PyTorch A Single Head
A head is an nn.Module that packages one full attention pass, from three projections to one weighted mix of values. The Head class holds the query, key, and value layers plus the causal mask, and mixes one window per call.
The mask lives inside the module
The Head constructor builds the triangular mask once and registers it as a buffer. The printed corner rows show the staircase, with row one holding a single visible position.
import torch
from torch import nn
from torch.nn import functional as F
batch_size = 32
block_size = 64
n_embd = 128
n_head = 4
n_layer = 4
class Head(nn.Module):
def __init__(self, head_size):
super().__init__()
self.query = nn.Linear(n_embd, head_size, bias=False)
self.key = nn.Linear(n_embd, head_size, bias=False)
self.value = nn.Linear(n_embd, head_size, bias=False)
self.register_buffer("mask", torch.tril(torch.ones(block_size, block_size)))
def forward(self, x):
B, T, C = x.shape
q = self.query(x)
k = self.key(x)
v = self.value(x)
scores = q @ k.transpose(-2, -1) * k.shape[-1] ** -0.5
scores = scores.masked_fill(self.mask[:T, :T] == 0, float("-inf"))
weights = F.softmax(scores, dim=-1)
return weights @ v
head = Head(32)
print("mask shape:", tuple(head.mask.shape))
print("mask corner row 1:", head.mask[1, :4].tolist())
print("mask corner row 3:", head.mask[3, :4].tolist())One call, one mixed window
A random window of four vectors enters the head at width 128. The head returns four mixed vectors at the head size of 32.
import torch
from torch import nn
from torch.nn import functional as F
batch_size = 32
block_size = 64
n_embd = 128
n_head = 4
n_layer = 4
class Head(nn.Module):
def __init__(self, head_size):
super().__init__()
self.query = nn.Linear(n_embd, head_size, bias=False)
self.key = nn.Linear(n_embd, head_size, bias=False)
self.value = nn.Linear(n_embd, head_size, bias=False)
self.register_buffer("mask", torch.tril(torch.ones(block_size, block_size)))
def forward(self, x):
B, T, C = x.shape
q = self.query(x)
k = self.key(x)
v = self.value(x)
scores = q @ k.transpose(-2, -1) * k.shape[-1] ** -0.5
scores = scores.masked_fill(self.mask[:T, :T] == 0, float("-inf"))
weights = F.softmax(scores, dim=-1)
return weights @ v
torch.manual_seed(42)
head = Head(32)
x = torch.randn(1, 4, n_embd)
out = head(x)
print("input shape:", tuple(x.shape))
print("output shape:", tuple(out.shape))Example
import torch
from torch import nn
from torch.nn import functional as F
batch_size = 32
block_size = 64
n_embd = 128
n_head = 4
n_layer = 4
class Head(nn.Module):
def __init__(self, head_size):
super().__init__()
self.query = nn.Linear(n_embd, head_size, bias=False)
self.key = nn.Linear(n_embd, head_size, bias=False)
self.value = nn.Linear(n_embd, head_size, bias=False)
self.register_buffer("mask", torch.tril(torch.ones(block_size, block_size)))
def forward(self, x):
B, T, C = x.shape
q = self.query(x)
k = self.key(x)
v = self.value(x)
scores = q @ k.transpose(-2, -1) * k.shape[-1] ** -0.5
scores = scores.masked_fill(self.mask[:T, :T] == 0, float("-inf"))
weights = F.softmax(scores, dim=-1)
return weights @ v
torch.manual_seed(42)
text = "Moss the cat woke in the warm house"
chars = sorted(set(text))
stoi = {c: i for i, c in enumerate(chars)}
window = text[:12]
ids = torch.tensor([[stoi[c] for c in window]])
token_embedding = nn.Embedding(len(chars), n_embd)
position_embedding = nn.Embedding(block_size, n_embd)
x = token_embedding(ids) + position_embedding(torch.arange(len(window)))
head = Head(32)
out = head(x)
# The same computation again, written out step by step.
q = head.query(x)
k = head.key(x)
v = head.value(x)
scores = q @ k.transpose(-2, -1) / 32 ** 0.5
scores = scores.masked_fill(head.mask[:12, :12] == 0, float("-inf"))
weights = F.softmax(scores, dim=-1)
by_hand = weights @ v
diff = (out - by_hand).abs().max().item()
print("window:", window)
print("output shape:", tuple(out.shape))
print("out row 0, first 4:", [round(n, 3) for n in out[0, 0, :4].tolist()])
print("max difference:", f"{diff:.2e}")Key points
- nn.Module: the base class for PyTorch parts that own weights
- register_buffer stores a tensor inside a module without training it
- The mask is a buffer, so no gradient ever changes the triangle
- Head size 32 gives each of the four planned heads one quarter of the width
- The forward pass runs projections, scores, mask, softmax, and mix in that order
In this example, the window is the first twelve characters of a real corpus sentence about Moss the cat. Token and position embeddings turn the window into vectors at width 128. One Head with head size 32 mixes them. The program then repeats the whole computation step by step with the same layers and the same mask. The two results are compared entry by entry, and the largest difference prints in scientific notation. A difference at that scale is floating-point rounding, not a different computation. The module and the hand-written steps agree, so the class can be trusted as the earlier lessons in one package.
Note
The class reads n_embd and block_size from the module constants above it
Try it Yourself
Edit the code, then press Run.
Running is simulated · Python Playground
Exercise
Why is the mask stored with register_buffer instead of as a plain tensor attribute?