Attention from Scratch

Every act below runs on one tiny example: the six tokens "Your journey starts with one step" as 3-dimensional vectors. Small enough to compute every number in your browser — and the exact numbers the reference Python prints.

That average had no dials to turn. Learning starts the moment we insert weight matrices between the tokens and their comparison.

One problem remains for a text generator: this attends to everything — including words that haven't been written yet.

A single masked head reads the sentence one way. Transformers want several, in parallel, without paying several times over.

The lab

Drive the attention matrix yourself. Toggle the causal mask, warm or cool the softmax, and click any cell to watch one query meet one key and become a weight.

And the scaling trick from Act 2, made tangible — slide the key dimension up and watch the un-scaled softmax collapse into a spike.

The real code

The demos run a TypeScript port; here is the Python they mirror. The whole of self-attention is a dozen lines — three projections, scaled dot products, softmax, a weighted sum of values:

class SelfAttention_v2(nn.Module):
    def __init__(self, d_in, d_out, qkv_bias=False):
        super().__init__()
        self.W_query = nn.Linear(d_in, d_out, bias=qkv_bias)
        self.W_key   = nn.Linear(d_in, d_out, bias=qkv_bias)
        self.W_value = nn.Linear(d_in, d_out, bias=qkv_bias)
 
    def forward(self, x):
        keys    = self.W_key(x)
        queries = self.W_query(x)
        values  = self.W_value(x)
 
        attn_scores  = queries @ keys.T                       # ω
        attn_weights = torch.softmax(
            attn_scores / keys.shape[-1]**0.5, dim=-1)        # ÷ √dₖ
        return attn_weights @ values                          # context vectors

Causal masking is two extra lines — set the upper triangle to -inf before the softmax, so a token can never attend to its future:

mask = torch.triu(torch.ones(context_length, context_length), diagonal=1)
masked = attn_scores.masked_fill(mask.bool(), -torch.inf)
attn_weights = torch.softmax(masked / keys.shape[-1]**0.5, dim=-1)

class MultiHeadAttention(nn.Module):
    def __init__(self, d_in, d_out, context_length, dropout, num_heads, qkv_bias=False):
        super().__init__()
        assert d_out % num_heads == 0, "d_out must be divisible by num_heads"
        self.d_out = d_out
        self.num_heads = num_heads
        self.head_dim = d_out // num_heads
        self.W_query = nn.Linear(d_in, d_out, bias=qkv_bias)
        self.W_key   = nn.Linear(d_in, d_out, bias=qkv_bias)
        self.W_value = nn.Linear(d_in, d_out, bias=qkv_bias)
        self.out_proj = nn.Linear(d_out, d_out)
        self.dropout = nn.Dropout(dropout)
        self.register_buffer(
            "mask",
            torch.triu(torch.ones(context_length, context_length), diagonal=1))
 
    def forward(self, x):
        b, num_tokens, d_in = x.shape
        keys    = self.W_key(x)      # (b, num_tokens, d_out)
        queries = self.W_query(x)
        values  = self.W_value(x)
 
        # split d_out into heads → (b, num_tokens, num_heads, head_dim)
        keys    = keys.view(b, num_tokens, self.num_heads, self.head_dim)
        values  = values.view(b, num_tokens, self.num_heads, self.head_dim)
        queries = queries.view(b, num_tokens, self.num_heads, self.head_dim)
 
        # heads become a batch dim → (b, num_heads, num_tokens, head_dim)
        keys, queries, values = keys.transpose(1, 2), queries.transpose(1, 2), values.transpose(1, 2)
 
        attn_scores = queries @ keys.transpose(2, 3)          # per-head dot products
        mask_bool = self.mask.bool()[:num_tokens, :num_tokens]
        attn_scores.masked_fill_(mask_bool, -torch.inf)
 
        attn_weights = torch.softmax(attn_scores / keys.shape[-1]**0.5, dim=-1)
        attn_weights = self.dropout(attn_weights)
 
        context_vec = (attn_weights @ values).transpose(1, 2) # (b, num_tokens, num_heads, head_dim)
        context_vec = context_vec.contiguous().view(b, num_tokens, self.d_out)
        return self.out_proj(context_vec)                     # combine heads

Four things to remember

1. Attention is a weighted average. Each token's output is a blend of all tokens' value vectors; the weights are a softmax over query·key scores. Acts 1–2 are that idea, first without and then with learnable weights.
2. Query, key, value are three views of the same token. What it looks for, how it's found, and what it passes on — each a learned linear projection.
3. Causal masking makes it a language model. Setting the upper triangle to −∞ before softmax (so those weights come out zero) stops a token from seeing the future, which is what lets the model generate left to right.
4. Multi-head attention is a split, not a stack. One projection sliced into heads gives several attention patterns for the price of one — the block GPT repeats a dozen times.

Adapted from Build a Large Language Model (From Scratch) by Sebastian Raschka — original notebook (Apache 2.0). The attention weights and context vectors on this page are torch's real seeded outputs, verified number-for-number against the Python with npm run verify:attention.