Self-Attention from Scratch

Self-attention is a mechanism that lets each position in a sequence attend to every other position in a single parallel step, using weights computed from the data itself.

It exists because RNNs process a sequence one token at a time. By token 50, the information from token 1 has been squeezed through 50 compression steps. Long-range dependencies get crushed into a fixed-size hidden state, a bottleneck that LSTM gating softens but doesn't fully solve.

The 2014 Bahdanau attention paper offered a fix: let the decoder look back at every encoder position and decide which ones matter for the current step. It was still bolted onto an RNN. The 2017 "Attention Is All You Need" paper asked a sharper question: what if attention is the sole mechanism? No recurrence, no convolution, just attention.

That single parallel step is what makes transformers fast, scalable, and dominant.

The Database Lookup Analogy

Self-attention is a function that maps a sequence of n vectors to a sequence of n vectors. Each output is a weighted sum of value vectors, and the weights come from how well that token's query matches every token's key.

The easiest way to hold that in your head is a soft database lookup:

Traditional database:
  Query: "capital of France"  -->  exact match  -->  "Paris"

Attention:
  Query: "capital of France"  -->  similarity to ALL keys  -->  weighted blend of ALL values

Each token produces three vectors:

• Query (Q): "What am I looking for?"
• Key (K): "What do I contain?"
• Value (V): "What information do I provide if selected?"

A query's dot product with every key gives one score per token, and a high score means "this key matches my query." After softmax (covered below), those scores weight the values, and the output is their weighted sum.

Q, K, V Computation

Each token embedding is multiplied by three learned weight matrices:

Input embeddings (sequence of n tokens, each d-dimensional):

  X = [x1, x2, x3, ..., xn]       shape: (n, d)

Three weight matrices:

  Wq  shape: (d, dk)
  Wk  shape: (d, dk)
  Wv  shape: (d, dv)

Projections:

  Q = X @ Wq    shape: (n, dk)      each token's query
  K = X @ Wk    shape: (n, dk)      each token's key
  V = X @ Wv    shape: (n, dv)      each token's value

Visually, for one token:

             Wq
  x_i ------[*]------> q_i    "What am I looking for?"
       |
       |     Wk
       +----[*]------> k_i    "What do I contain?"
       |
       |     Wv
       +----[*]------> v_i    "What do I offer?"

The Attention Matrix

With Q, K, V in hand for all tokens, the attention scores form a matrix:

Scores = Q @ K^T    shape: (n, n)

              k1    k2    k3    k4    k5
        +-----+-----+-----+-----+-----+
   q1   | 2.1 | 0.3 | 0.1 | 0.8 | 0.2 |   <- how much q1 attends to each key
        +-----+-----+-----+-----+-----+
   q2   | 0.4 | 1.9 | 0.7 | 0.1 | 0.3 |
        +-----+-----+-----+-----+-----+
   q3   | 0.2 | 0.6 | 2.3 | 0.5 | 0.1 |
        +-----+-----+-----+-----+-----+
   q4   | 0.9 | 0.1 | 0.4 | 1.7 | 0.6 |
        +-----+-----+-----+-----+-----+
   q5   | 0.1 | 0.3 | 0.2 | 0.5 | 2.0 |
        +-----+-----+-----+-----+-----+

Each row: one token's attention over the entire sequence

Watch one query at a time sweep the keys: each row scores every token, softmax turns the scores into weights, and the context vector is the weighted blend of values.

Why Scale?

Dot products grow with the dimension dk. If dk = 64, they can land in the tens, which pushes softmax into regions where gradients vanish. Dividing by sqrt(dk) fixes it:

Scaled scores = (Q @ K^T) / sqrt(dk)

That keeps the scores in a range where softmax still gives useful gradients.

Softmax Turns Scores into Weights

Softmax converts raw scores into a probability distribution across each row:

Raw scores for q1:   [2.1, 0.3, 0.1, 0.8, 0.2]   (before the 1/sqrt(dk) scaling, to keep the arithmetic simple)
                            |
                         softmax
                            |
Attention weights:   [0.58, 0.10, 0.08, 0.16, 0.09]   (sums to 1.0, up to rounding)

Each token now has a row of weights saying how much to attend to every other token.

Weighted Sum of Values

The final output for each token is a weighted sum of all value vectors:

output_i = sum( attention_weight[i][j] * v_j  for all j )

For token 1:
  output_1 = 0.58 * v1 + 0.10 * v2 + 0.08 * v3 + 0.16 * v4 + 0.09 * v5

Full Pipeline

The whole pipeline in one line:

Attention(Q, K, V) = softmax( Q @ K^T / sqrt(dk) ) @ V

Causal Masking

Causal masking is the change that turns bidirectional self-attention (every token sees every token) into the decoder-style attention used for text generation. Before softmax, every score whose key position j comes after the query position i is set to negative infinity. Since exp(-inf) = 0, those weights come out exactly zero, and row i spreads its weight over positions 0 through i only. A token can no longer peek at the tokens it is supposed to predict. The code in Build It has no mask parameter, so adding one is Exercise 1.

Step 1: Softmax from scratch

Softmax converts raw logits into probabilities. Subtracting the max first keeps the exponentials from overflowing.

import numpy as np
 
def softmax(x):
    shifted = x - np.max(x, axis=-1, keepdims=True)
    exp_x = np.exp(shifted)
    return exp_x / np.sum(exp_x, axis=-1, keepdims=True)
 
logits = np.array([2.0, 1.0, 0.1])
print(f"logits:  {logits}")
print(f"softmax: {softmax(logits)}")
print(f"sum:     {softmax(logits).sum():.4f}")

The full self_attention.py also runs softmax on [100, 200, 300]. Because of the max shift it returns a finite result, roughly [0, 0, 1], and the sum is still 1.

Step 2: Scaled dot-product attention

The core function. It takes Q, K, V and returns the attention output plus the weight matrix.

def scaled_dot_product_attention(Q, K, V):
    dk = Q.shape[-1]
    scores = Q @ K.T / np.sqrt(dk)
    weights = softmax(scores)
    output = weights @ V
    return output, weights

Step 3: Self-attention class with learned projections

A self-attention module that holds Wq, Wk, Wv, initialized with Xavier-like scaling. The Q/K scale is named scale_qk so it doesn't collide with scale_v, and forward returns the (output, weights) tuple straight from scaled_dot_product_attention.

class SelfAttention:
    def __init__(self, d_model, dk, dv, seed=42):
        rng = np.random.default_rng(seed)
        scale_qk = np.sqrt(2.0 / (d_model + dk))
        self.Wq = rng.normal(0, scale_qk, (d_model, dk))
        self.Wk = rng.normal(0, scale_qk, (d_model, dk))
        scale_v = np.sqrt(2.0 / (d_model + dv))
        self.Wv = rng.normal(0, scale_v, (d_model, dv))
        self.dk = dk
 
    def forward(self, X):
        Q = X @ self.Wq
        K = X @ self.Wk
        V = X @ self.Wv
        return scaled_dot_product_attention(Q, K, V)

Step 4: Run it on a sentence

Make random embeddings for a sentence and print the attention weights.

sentence = ["The", "cat", "sat", "on", "the", "mat"]
n_tokens = len(sentence)
d_model = 16
dk = 8
dv = 8
 
rng = np.random.default_rng(42)
X = rng.normal(0, 1, (n_tokens, d_model))
 
attn = SelfAttention(d_model, dk, dv, seed=42)
output, weights = attn.forward(X)
 
print("Attention weights (each row: where that token looks):\n")
print(f"{'':>6}", end="")
for token in sentence:
    print(f"{token:>6}", end="")
print()
 
for i, token in enumerate(sentence):
    print(f"{token:>6}", end="")
    for j in range(n_tokens):
        w = weights[i][j]
        print(f"{w:6.3f}", end="")
    print()

Step 5: Visualize attention with ASCII heatmap

Map each weight to a shade character for a quick look at the matrix.

def ascii_heatmap(weights, tokens, chars=" ░▒▓█"):
    n = len(tokens)
    print(f"\n{'':>6}", end="")
    for t in tokens:
        print(f"{t:>6}", end="")
    print()
 
    for i in range(n):
        print(f"{tokens[i]:>6}", end="")
        for j in range(n):
            level = int(weights[i][j] * (len(chars) - 1) / weights.max())
            level = min(level, len(chars) - 1)
            print(f"{'  ' + chars[level] + '   '}", end="")
        print()
 
ascii_heatmap(weights, sentence)

Step 6: Multi-head attention

One attention function can only score tokens one way. The file's MultiHeadSelfAttention runs several in parallel. Each head is a SelfAttention with its own seed and dk = dv = d_model // n_heads, the head outputs are concatenated, and a final matrix Wo projects the result back to d_model.

class MultiHeadSelfAttention:
    def __init__(self, d_model, n_heads, seed=42):
        assert d_model % n_heads == 0
        self.n_heads = n_heads
        self.dk = d_model // n_heads
        self.dv = d_model // n_heads
        self.heads = [
            SelfAttention(d_model, self.dk, self.dv, seed=seed + i)
            for i in range(n_heads)
        ]
        rng = np.random.default_rng(seed + n_heads)
        scale = np.sqrt(2.0 / (d_model + d_model))
        self.Wo = rng.normal(0, scale, (n_heads * self.dv, d_model))
 
    def forward(self, X):
        head_outputs = []
        all_weights = []
        for head in self.heads:
            out, w = head.forward(X)
            head_outputs.append(out)
            all_weights.append(w)
        concatenated = np.concatenate(head_outputs, axis=-1)
        output = concatenated @ self.Wo
        return output, all_weights

Run python self_attention.py to see the single-head matrix, the heatmap, and both heads' weights for the same sentence. Two things are left to you: the file has no mask parameter (causal masking is Exercise 1), and it gives each head its own projections rather than splitting one Q, K, V into chunks (Exercise 2).

PyTorch's nn.MultiheadAttention runs the same computation we built, with multi-head splitting and an output projection added:

import torch
import torch.nn as nn
 
d_model = 8
n_heads = 2
seq_len = 6
 
mha = nn.MultiheadAttention(embed_dim=d_model, num_heads=n_heads, batch_first=True)
 
X_torch = torch.randn(1, seq_len, d_model)
 
output, attn_weights = mha(X_torch, X_torch, X_torch)
 
print(f"Input shape:            {X_torch.shape}")
print(f"Output shape:           {output.shape}")
print(f"Attention weight shape: {attn_weights.shape}")
print(f"\nAttn weights (averaged over heads):")
print(attn_weights[0].detach().numpy().round(3))

The difference is the multi-head part. It runs several attention functions in parallel, each with its own Q, K, V projections of size dk = d_model / n_heads, then concatenates the results. That lets the model attend to different relationship types at once.

What you ship:

• outputs/prompt-attention-explainer.md: a prompt for explaining attention through the database lookup analogy

1. Modify scaled_dot_product_attention to accept an optional mask matrix that sets certain positions to negative infinity before softmax (this is how causal/decoder masking works). self_attention.py has no mask parameter, so this one is entirely yours: build a lower-triangular mask and check that every weight above the diagonal comes out zero
2. Implement multi-head attention from scratch: split Q, K, V into n_heads chunks, run attention on each, concatenate, and project through a final weight matrix Wo
3. Take two different sentences of the same length, feed them through the same SelfAttention instance, and compare their attention patterns. What changes? What stays the same?

• Attention Is All You Need (Vaswani et al., 2017) - the original transformer paper
• The Illustrated Transformer (Jay Alammar) - a strong visual walkthrough of the full architecture
• The Annotated Transformer (Harvard NLP) - line-by-line PyTorch implementation with explanations