Attention & Transformers

By the end of this lesson you will be able to explain self-attention in plain words, compute it step by step in NumPy, add the causal mask that text generators use, and see how attention, layer normalisation and skip connections combine into a transformer block. You will also build one in PyTorch.

The transformer was introduced in the 2017 paper "Attention Is All You Need" by Vaswani and colleagues, for translation. It now underlies large language models, many vision models and speech systems. The core idea fits on one line, and you will code it yourself.

Why attention?

An RNN reads a sentence one word at a time and squeezes everything into one hidden state. Two problems follow: information from far back fades, and the steps cannot run in parallel, so training is slow.

Attention takes a different approach. Every word looks directly at every other word in the sentence and decides how much to borrow from each. In "Meera put the bat in her bag because it was new", the word "it" needs to know whether it refers to the bat or the bag. Attention lets "it" look at both and take more information from the right one.

Think of a library in Lucknow:

• Your query is what you are looking for: "a book about cricket for beginners".
• Each book has a key: the label on its spine.
• Each book has a value: its actual content.

You compare your query with every key, give more weight to the books that match best, and read a weighted mix of their content. In self-attention, every word plays all three roles: it asks a query, offers a key, and provides a value.

Scaled dot-product attention

Each token's vector x is turned into three vectors by three learned matrices: q = x @ W_q, k = x @ W_k, v = x @ W_v. For a whole sentence, these become matrices Q, K and V, one row per token. Then:

1. Scores: Q @ K.T gives every pair's similarity (a dot product).
2. Scale: divide by sqrt(d), where d is the size of each key.
3. Weights: softmax across each row, so each token's weights add up to 1.
4. Output: weights @ V, a weighted mix of the values for each token.

In one line: attention(Q, K, V) = softmax(Q @ K.T / sqrt(d)) @ V.

Step 1: self-attention on a five-token sentence.

import numpy as np

def softmax(z):
    e = np.exp(z - z.max(axis=-1, keepdims=True))
    return e / e.sum(axis=-1, keepdims=True)

def attention(Q, K, V, mask=None):
    scores = Q @ K.T / np.sqrt(K.shape[-1])
    if mask is not None:
        scores = np.where(mask, scores, -np.inf)    # blocked pairs get zero weight
    weights = softmax(scores)
    return weights @ V, weights

rng = np.random.default_rng(0)
tokens = ["Meera", "bought", "a", "new", "bat"]
d_model, d_k = 8, 4
X = rng.normal(size=(len(tokens), d_model))        # token embeddings (random here)
W_q, W_k, W_v = (rng.normal(0, 0.5, (d_model, d_k)) for _ in range(3))

out, weights = attention(X @ W_q, X @ W_k, X @ W_v)
print("output shape:", out.shape, " weights shape:", weights.shape)
print("each row sums to 1:", weights.sum(axis=1).round(3))
for token, row in zip(tokens, weights):
    print(f"{token:<7}", row.round(2))

output shape: (5, 4)  weights shape: (5, 5)
each row sums to 1: [1. 1. 1. 1. 1.]
Meera   [0.03 0.32 0.12 0.51 0.02]
bought  [0.56 0.13 0.05 0.02 0.24]
a       [0.45 0.31 0.02 0.08 0.14]
new     [0.08 0.15 0.5  0.16 0.11]
bat     [0.02 0.38 0.2  0.39 0.02]

Each row shows how much one token draws from every token, itself included. With random weights the pattern means nothing yet; training adjusts W_q, W_k and W_v so that useful links get high weights. Notice that all five rows were computed at once with matrix multiplication. That parallelism is why transformers train so efficiently on GPUs.

Why divide by the square root of d?

The dot product of two random vectors grows with their length. With large d, the scores become large, softmax pushes almost all the weight onto one token, and its gradients become tiny, just like a saturated sigmoid. Dividing by sqrt(d) keeps the scores at a steady scale.

Step 2: see softmax saturate without scaling.

import numpy as np

def softmax(z):
    e = np.exp(z - z.max())
    return e / e.sum()

rng = np.random.default_rng(1)
for d in [4, 64, 512]:
    q = rng.normal(size=d)
    keys = rng.normal(size=(5, d))
    raw = keys @ q
    print(f"d = {d:>3}: raw score spread {raw.std():6.2f} | "
          f"max weight unscaled {softmax(raw).max():.2f}, scaled {softmax(raw / np.sqrt(d)).max():.2f}")

d =   4: raw score spread   0.48 | max weight unscaled 0.29, scaled 0.25
d =  64: raw score spread   5.73 | max weight unscaled 0.94, scaled 0.38
d = 512: raw score spread  14.71 | max weight unscaled 0.93, scaled 0.31

Without scaling, the larger the vectors, the more one token takes almost all the weight. With scaling, the weights stay spread out, so the model can still learn from several tokens at once.

The causal mask

A text generator such as GPT predicts the next word. While training, it must not peek at the words that come after the current position. A causal mask blocks those pairs: before softmax, their scores are set to minus infinity, which gives them a weight of exactly 0.

Step 3: add a causal mask. This repeats the setup from Step 1.

import numpy as np

def softmax(z):
    e = np.exp(z - z.max(axis=-1, keepdims=True))
    return e / e.sum(axis=-1, keepdims=True)

rng = np.random.default_rng(0)
tokens = ["Meera", "bought", "a", "new", "bat"]
X = rng.normal(size=(5, 8))
W_q, W_k, W_v = (rng.normal(0, 0.5, (8, 4)) for _ in range(3))
Q, K, V = X @ W_q, X @ W_k, X @ W_v

causal = np.tril(np.ones((5, 5), dtype=bool))      # True on and below the diagonal
print(causal.astype(int))
scores = np.where(causal, Q @ K.T / np.sqrt(4), -np.inf)
weights = softmax(scores)
for token, row in zip(tokens, weights):
    print(f"{token:<7}", row.round(2))

[[1 0 0 0 0]
 [1 1 0 0 0]
 [1 1 1 0 0]
 [1 1 1 1 0]
 [1 1 1 1 1]]
Meera   [1. 0. 0. 0. 0.]
bought  [0.81 0.19 0.   0.   0.  ]
a       [0.57 0.4  0.03 0.   0.  ]
new     [0.09 0.17 0.56 0.18 0.  ]
bat     [0.02 0.38 0.2  0.39 0.02]

The first token can only look at itself, so its weight is 1. The last token can look at all five. Every weight above the diagonal is exactly 0.

Word order: positional encoding

Attention on its own ignores order. If you shuffle the tokens, every token gets the same output as before, just in a shuffled position. "Dog bites man" and "man bites dog" would look identical. So transformers add positional information to each token's embedding before the first layer.

The original paper used fixed sinusoidal encodings: each position gets a vector of sine and cosine waves of different wavelengths. Many newer models instead learn position embeddings, or use methods that rotate queries and keys by position.

Step 4: order-blindness and the sinusoidal fix.

import numpy as np

def softmax(z):
    e = np.exp(z - z.max(axis=-1, keepdims=True))
    return e / e.sum(axis=-1, keepdims=True)

def self_attention(X, W):
    Q, K, V = X @ W[0], X @ W[1], X @ W[2]
    return softmax(Q @ K.T / np.sqrt(K.shape[1])) @ V

def positional_encoding(n_positions, d):
    pos = np.arange(n_positions)[:, None]
    i = np.arange(d // 2)[None, :]
    angles = pos / 10000 ** (2 * i / d)
    pe = np.zeros((n_positions, d))
    pe[:, 0::2], pe[:, 1::2] = np.sin(angles), np.cos(angles)
    return pe

rng = np.random.default_rng(3)
X = rng.normal(size=(4, 8))                          # "dog bites the man", say
W = [rng.normal(0, 0.5, (8, 4)) for _ in range(3)]
order = [3, 1, 2, 0]                                 # swap the first and last token

plain = self_attention(X, W)
shuffled = self_attention(X[order], W)
print("without positions, outputs just move with the tokens:", np.allclose(plain[order], shuffled))

pe = positional_encoding(4, 8)
with_pos = self_attention(X + pe, W)
shuffled_pos = self_attention(X[order] + pe, W)
print("with positions, swapping words changes the result:  ", not np.allclose(with_pos[order], shuffled_pos))
print("positional encoding for position 1:", pe[1].round(3))

without positions, outputs just move with the tokens: True
with positions, swapping words changes the result:   True
positional encoding for position 1: [0.841 0.54  0.1   0.995 0.01  1.    0.001 1.   ]

Without positions, swapping two words simply swaps their outputs, so the model cannot tell the sentences apart. Once each position adds its own pattern, the same words in a different order produce genuinely different results.

Multi-head attention and the transformer block

One attention pattern can capture one kind of relationship. Multi-head attention runs several attention computations ("heads") side by side, each with its own smaller W_q, W_k, W_v, and joins their outputs. One head might link pronouns to nouns, another adjectives to the nouns they describe.

A transformer block wraps attention with three ideas you already know:

1. Layer normalisation before each part (Batch Normalisation explained LayerNorm).
2. Multi-head self-attention, so tokens share information.
3. A small feed-forward network (two dense layers with an activation such as GELU) applied to each token separately.
4. Skip connections around both the attention and the feed-forward parts, as in ResNet.

The original paper put the normalisation after each part; most modern models put it before ("pre-norm"), which tends to train more stably. A full model stacks many such blocks.

Attention and transformers in PyTorch

First, PyTorch's built-in attention function gives the same answer as our NumPy code. Then nn.MultiheadAttention and a complete nn.TransformerEncoderLayer. This is plain code with its real output.

import math
import torch
import torch.nn.functional as F
from torch import nn

torch.manual_seed(0)
Q, K, V = torch.randn(5, 4), torch.randn(5, 4), torch.randn(5, 4)
manual = torch.softmax(Q @ K.T / math.sqrt(4), dim=-1) @ V
builtin = F.scaled_dot_product_attention(Q, K, V)
print("same as the formula:", torch.allclose(manual, builtin, atol=1e-6))
causal = F.scaled_dot_product_attention(Q, K, V, is_causal=True)
print("first row with causal mask equals V[0]:", torch.allclose(causal[0], V[0]))

x = torch.randn(2, 10, 32)                           # batch 2, 10 tokens, 32 features
mha = nn.MultiheadAttention(embed_dim=32, num_heads=4, batch_first=True)
out, attn = mha(x, x, x)                             # self-attention: query = key = value = x
print("multi-head output:", tuple(out.shape), " weights:", tuple(attn.shape))

block = nn.TransformerEncoderLayer(d_model=32, nhead=4, dim_feedforward=64,
                                   batch_first=True, norm_first=True, activation="gelu")
print("encoder block output:", tuple(block(x).shape))
print("block parameters:", sum(p.numel() for p in block.parameters()))

same as the formula: True
first row with causal mask equals V[0]: True
multi-head output: (2, 10, 32)  weights: (2, 10, 10)
encoder block output: (2, 10, 32)
block parameters: 8544

F.scaled_dot_product_attention also picks fast, memory-saving implementations on GPUs when it can. A transformer block keeps the shape (batch, tokens, features) unchanged, which is why blocks stack so easily. By default nn.MultiheadAttention returns the attention weights averaged over the heads.

Encoders, decoders and the cost of attention

Transformers come in three shapes:

Attention compares every token with every other token, so its cost grows with the square of the sequence length: twice as many tokens means about four times the work. This is why language models have a maximum context length and why much research goes into faster attention. To see how these ideas become chatbots, read How LLMs Work: Tokens & Training.

Common mistakes

• Forgetting the scale factor. Unscaled scores saturate the softmax and slow learning.
• Forgetting the causal mask in a text generator. The model "cheats" by looking at the answer during training and fails at generation time.
• Leaving out positional information. Attention alone cannot see word order.
• Mixing up batch_first. nn.MultiheadAttention defaults to (tokens, batch, features) unless you pass batch_first=True.
• Ignoring the quadratic cost. Doubling the sequence length roughly quadruples attention's time and memory.

Key points

• Self-attention lets every token take a weighted mix of every other token's value, weighted by query-key similarity.
• Attention is softmax(Q @ K.T / sqrt(d)) @ V; the scale keeps softmax from saturating.
• A causal mask stops tokens from seeing the future, which decoder models such as GPT need.
• Attention ignores order, so transformers add positional encodings.
• A transformer block combines layer norm, multi-head attention, a feed-forward network and skip connections; models stack many blocks.

import numpy as np

def softmax(z):
    e = np.exp(z - z.max(axis=-1, keepdims=True))
    return e / e.sum(axis=-1, keepdims=True)

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

rng = np.random.default_rng(0)
X = rng.normal(size=(5, 8))
outputs = []
for head in range(2):
    W_q, W_k, W_v = (rng.normal(0, 0.5, (8, 4)) for _ in range(3))
    out, weights = attention(X @ W_q, X @ W_k, X @ W_v)
    outputs.append(out)
    print(f"head {head} weights for token 0:", weights[0].round(2))
combined = np.concatenate(outputs, axis=1)
print("combined output shape:", combined.shape)

head 0 weights for token 0: [0.03 0.32 0.12 0.51 0.02]
head 1 weights for token 0: [0.19 0.16 0.3  0.27 0.08]
combined output shape: (5, 8)