Chapter 11 // Architectures

Attention & Transformers

Every architecture so far had a fixed wiring diagram. Attention lets a model decide, for each output, which inputs to read — and that one idea, scaled up, became the Transformer that now underlies nearly all of modern AI.

Adapted from the original chapter of Dive into Deep Learning · CC BY-SA 4.0 · the Python shown is the book's; the figures run a faithful TypeScript port in your browser

For thirty years, the recipe was to bake the connections into the architecture: convolutions wire each output to a fixed local window, recurrences pass a single state down a fixed chain. Attention throws that out. It lets the network compute, on the fly, how much each piece of the input should matter to each piece of the output. Start with the mechanism itself — a differentiable lookup.

Act 1 · Attention is a soft lookup

01/05

Start with a database: keys on the left, values on the right. A last name is the key, a first name the value. To look something up, you bring a query.

The whole thing is one equation: given a query and a database of key–value pairs, the output is a weighted sum of the values, with weights from a softmax over query–key scores.

\text{Attention}(\mathbf{q}, \mathcal{D}) = \sum_{i} \alpha(\mathbf{q}, \mathbf{k}_i)\,\mathbf{v}_i, \qquad \alpha(\mathbf{q}, \mathbf{k}_i) = \frac{\exp(a(\mathbf{q}, \mathbf{k}_i))}{\sum_j \exp(a(\mathbf{q}, \mathbf{k}_j))}

This is not a new idea. The oldest version of it predates deep learning by half a century — and it's the clearest way to see attention at work.

Act 2 · Attention, fifty years early

01/05

Forty noisy samples of y = 2·sin(x) + x. We want to predict y at a new x — with no training at all.

Kernel regression hand-picks the similarity function. The leap is to learn it — but first, the scoring function $a(\mathbf{q}, \mathbf{k})$ needs to be cheap to compute and safe to feed a softmax.

Act 3 · Scoring, cheaply and safely

01/05

The cheapest way to score a query against a key is their dot product. But there's a catch: for random vectors of length d, the dot product's variance grows in lockstep with d — past 120 by d=128.

The dot product $\mathbf{q}^\top\mathbf{k}/\sqrt{d}$ is the scoring function of every modern Transformer. Now make its parameters learnable, and attention becomes something a network can train end to end — which is exactly how it entered deep learning, as a fix for machine translation.

Act 4 · Learning to align

01/05

Bahdanau's 2014 idea launched the whole field: instead of a fixed similarity kernel, learn what queries and keys should match. Translate the red car into la voiture rouge.

Bahdanau attention made one RNN translator better. The bigger bet came next: throw the RNN away entirely and build a model out of nothing but attention. Two ingredients make that possible — several attention heads working in parallel, and a way to smuggle word order back in.

Act 5 · Many heads, and a sense of order

01/05

In self-attention, a sequence attends to itself: every token is a query, every token a key. Here is one attention head of a real trained Transformer — each row shows where that position looks.

Self-attention with multiple heads and positional encodings is the entire substrate. Stack it into a block, wrap each sublayer in a residual connection and layer normalization, and you have the Transformer.

Act 6 · The Transformer, assembled

01/05

Now stack it into a block: multi-head self-attention, then a position-wise feed-forward network, each wrapped in a residual connection and layer norm. Repeat N times for the encoder; add cross-attention for the decoder. No recurrence, anywhere.

That is the architecture — encoder, decoder, and the vision variant — running on toy tasks small enough to train in your browser but structurally identical to the models behind machine translation and image recognition.

The lab

Three ideas from the acts, yours to drive: attention pooling and how a kernel's width trades smoothing for overfitting, the positional code and why it encodes relative distance, and the trained Transformer reversing whatever digits you hand it.

▸ lab

Pool by similarity

Choose a kernel — and for the Gaussian, how wide. Watch attention interpolate, smooth, or overfit.

kernel

width σ = 0.50

mean error vs truth: 0.51

Nearby points count more. No training happened — the attention weights do all the work.

▸ lab

The shape of a position

Slide two positions. Their sinusoidal codes differ, but their similarity depends only on the gap between them.

oppositesimilarity 0.71identical

position i = 5position j = 12

distance |i − j| = 7
same gap at 25→32: sim 0.71

Move both sliders keeping the gap fixed: the similarity barely changes. Position codes encode relative distance.

▸ lab

Reverse it yourself

Set six digits — the trained Transformer flips them, and the cross-attention shows it reading the input backwards.

set the input

Reversed perfectly — 3 8 0 5 2 7. The bright anti-diagonal is the model reading the input back to front.

target 3 8 0 5 2 7 · output 3 8 0 5 2 7

Where this goes next: pretraining

The chapter's last section isn't an architecture — it's the shift in practice that made Transformers dominant. Rather than train from scratch for each task, you pretrain one large Transformer on a mountain of unlabeled text, then adapt it.

[three modes, and the scaling law]

The same block runs in three configurations. Encoder-only (BERT) sees the whole sequence at once and is pretrained by masking random tokens and predicting them — great for understanding, classification, tagging. Encoder–decoder (T5) casts every task as text-to-text, reconstructing corrupted spans. Decoder-only (the GPT series) uses causal attention — each token sees only its past — and simply predicts the next token, which turns out to be enough to learn translation, summarization, and reasoning as a side effect. The reason to build them ever larger is empirical: performance follows a power law in model size, data, and compute, and large models are also more sample-efficient. Scale the three together and the loss keeps falling — the observation that launched the era of foundation models, and the bridge to everything on the sister site about training LLMs from scratch.

The real code

Every act above carries a </> chip with its section's exact PyTorch. The encoder block is the whole architecture in four lines — self-attention and a feed-forward network, each wrapped in a residual connection and layer norm:

transformer.md (PyTorch)

class TransformerEncoderBlock(nn.Module):
    def __init__(self, num_hiddens, ffn_num_hiddens, num_heads, dropout, use_bias=False):
        super().__init__()
        self.attention = d2l.MultiHeadAttention(num_hiddens, num_heads, dropout, use_bias)
        self.addnorm1 = AddNorm(num_hiddens, dropout)
        self.ffn = PositionWiseFFN(ffn_num_hiddens, num_hiddens)
        self.addnorm2 = AddNorm(num_hiddens, dropout)
 
    def forward(self, X, valid_lens):
        Y = self.addnorm1(X, self.attention(X, X, X, valid_lens))
        return self.addnorm2(Y, self.ffn(Y))

Scaled dot-product attention + multi-head, the core (the full tour)

attention-scoring-functions.md · multihead-attention.md (PyTorch)

class DotProductAttention(nn.Module):
    def forward(self, queries, keys, values, valid_lens=None):
        d = queries.shape[-1]
        scores = torch.bmm(queries, keys.transpose(1, 2)) / math.sqrt(d)
        self.attention_weights = masked_softmax(scores, valid_lens)
        return torch.bmm(self.dropout(self.attention_weights), values)
 
class MultiHeadAttention(d2l.Module):
    def forward(self, queries, keys, values, valid_lens):
        queries = self.transpose_qkv(self.W_q(queries))   # split into heads
        keys = self.transpose_qkv(self.W_k(keys))
        values = self.transpose_qkv(self.W_v(values))
        output = self.attention(queries, keys, values, valid_lens)
        return self.W_o(self.transpose_output(output))    # concat + mix

[from the book]

Everything above runs a TypeScript port of these operations, verified against a Python implementation of the same maths (npm run verify:attention — the forward primitives and the alignment/ViT training match to 1e-6, every trained decode matches exactly, and both backprop passes are proven against finite differences). Honest deviations, all so the runs stay checkable and the figures legible: the book's Transformer translates English→French (d 256, 2+2 blocks, 4 heads); the site trains the identical architecture — sinusoidal positions, multi-head scaled dot-product attention, residual add-and-norm, a position-wise feed-forward net, a causal decoder mask, and Adam with the book's gradient clipping — shrunk (d 16, 1+1 blocks, 2 heads) on a reverse-the-digits task, so a full run fits in seconds and the cross-attention is a clean anti-diagonal. The Vision Transformer is the book's pre-norm ViT — patch embedding, <cls> token, learnable positions, GELU MLP — on chapter 4's 8×8 garments. The Bahdanau act trains only the additive-attention scorer on a four-token alignment (the book trains a full GRU translator; that chapter isn't on this site yet). Because 600 Adam steps amplify last-bit floating-point differences between engines, the two languages' late loss values diverge even though both converge to the same near-perfect model — the acts replay one baked, verified run. The book goes much further than six acts can: BERT/T5/GPT pretraining, the scaling laws, and the full derivations are all worth reading in the original.

Three things to remember

Attention is a differentiable lookup. Score a query against every key, softmax the scores into weights, return the weighted sum of the values. Every variant — kernel regression, Bahdanau, the Transformer — is that one equation with a different way of computing the scores.
The Transformer is attention without recurrence. Multi-head self-attention lets every token read every other in one parallel step; positional encodings restore word order; residual connections and layer norm make the stack trainable. No convolutions, no recurrence — and the shortest possible path between any two positions.
One architecture, scaled, ate the field. The same block reverses digits, classifies image patches, translates, and — pretrained on enough text — writes. Performance follows a power law in size, data, and compute, which is why the frontier keeps getting bigger.

Adapted from Chapter 11 of Dive into Deep Learning by Zhang, Lipton, Li, and Smola (CC BY-SA 4.0). The figures run this site's verified TypeScript port; all computation happens live in your browser.