Self Attention from Scratch

Self-attention is a module where queries, keys, and values all come from the same tensor. One module, tensor in and same-shape tensor out. Stack it with a feed-forward block and residuals and you have a transformer.

The idea

The module has two learned matrices: a qkv projection and an output projection.

• One projection, three views. Linear(C, 3C) produces queries, keys, and values together, then splits them: one matmul instead of three.
• Heads. Split C channels into H heads of size C/H, so H attention patterns run in parallel over separate channel groups. Each head can learn a different relation — the previous token, a repeated word — at the same total cost as one full-width head.
• Scale. Each head divides its scores by √head_dim, not √C.
• Mask. torch.tril builds the lower-triangular mask; the forbidden scores are set to −∞, so the softmax gives them exactly 0.

The shape walk with B = 1, T = 4, C = 8, H = 2:

Parameters are 3C² for qkv plus C² for proj, so 4C²: 256 at C = 8, about 16.8M at C = 2048.

Worked example

Causality is the property worth testing. Give the module a random (1, 4, 8) input, then change only position 2 and run it again.

• Outputs at positions 0 and 1 are identical: max absolute difference 0.0.
• Outputs at positions 2 and 3 move: max difference 0.3175.

Without the mask, positions 0 and 1 would move and training would silently cheat on next-token labels. The same test catches mask off-by-ones: tril keeps the diagonal, so a position attends to itself and to everything before it.

In code

import torch, torch.nn as nn, torch.nn.functional as F

class SelfAttention(nn.Module):
    def __init__(self, n_embd, n_head):
        super().__init__()
        assert n_embd % n_head == 0
        self.n_head, self.head_dim = n_head, n_embd // n_head
        self.qkv = nn.Linear(n_embd, 3 * n_embd, bias=False)
        self.proj = nn.Linear(n_embd, n_embd, bias=False)

    def forward(self, x):
        B, T, C = x.shape
        q, k, v = self.qkv(x).split(C, dim=2)
        q, k, v = (t.view(B, T, self.n_head, self.head_dim).transpose(1, 2) for t in (q, k, v))
        att = (q @ k.transpose(-2, -1)) * self.head_dim ** -0.5
        mask = torch.tril(torch.ones(T, T, dtype=torch.bool))
        att = F.softmax(att.masked_fill(~mask, float("-inf")), dim=-1)
        y = (att @ v).transpose(1, 2).contiguous().view(B, T, C)
        return self.proj(y)

torch.manual_seed(0)
attn = SelfAttention(8, 2)
x = torch.randn(1, 4, 8)
y1 = attn(x)
x2 = x.clone()
x2[0, 2] = torch.randn(8)
y2 = attn(x2)

print(tuple(y1.shape))                                  # (1, 4, 8)
print(sum(p.numel() for p in attn.parameters()))        # 256 == 4 * 8 * 8
print((y1[0, :2] - y2[0, :2]).abs().max().item())       # 0.0
print(round((y1[0, 2:] - y2[0, 2:]).abs().max().item(), 4))   # 0.3175

Check yourself

1. Why does changing position 2's input leave the outputs at positions 0 and 1 unchanged?
2. What do multiple heads buy over a single head of the same total width?
3. Why scale the scores by head_dim rather than C?

Key takeaways

• Self-attention is one qkv projection, a scaled q·kᵀ, a causal mask, and an output projection.
• Multi-head splits channels, not compute: heads run in parallel over the same budget.
• Test causality by perturbing one position and checking earlier outputs.