The Transformer

The whole idea in one sentence: turn tokens into vectors, then repeat "let every position gather information from the others, then think about it on its own", with a shortcut around each step so nothing gets lost. Everything else on this page is detail on that sentence.

Problem

Before 2017, the best systems for translation and other sequence tasks read a sentence one word at a time, carrying a running summary from word to word (see neurons and layers for the building block). That has two costs. Each step waits for the previous one, so the work cannot be spread across many chips at once; and a fact from the start of a long sentence must survive many re-summarizations to reach the end. The Transformer asks whether a model can drop the one-at-a-time reading and let every word look directly at every other word.

Intuition

Picture a room where every word is a person holding a note. In one round of attention, each person asks a question, glances at everyone's note, decides whose matter most, and copies a weighted blend of them. Then each person thinks privately about what they gathered (the feed-forward step). Do a round, then another: after six, each word's note carries the sentence's context. Because everyone looks at everyone in the same round, all the words can be processed at once, and any two words are one step apart however far apart they sit.

Follow one sentence through it

Translating "the cat sat down" into French, with a batch of one and the base sizes. Drag the slider (or read straight down with JavaScript off). Shapes are written batch-first, which is what you get with batch_first=True.

Mechanics

The pieces, in the order a token meets them. Each one has its own shapes and parameters; the tables give the PyTorch classes that implement them.

The big picture

Click any block to jump to its section. Solid boxes are what nn.Transformer builds; dashed boxes are the parts you add yourself, because PyTorch's module starts at vectors and stops at vectors. The paper's base model uses exactly these defaults: d_model=512, 8 heads, 6 encoder and 6 decoder layers, a 2048-wide feed-forward block and dropout 0.1 (constructor defaults).

Mixing and thinking Tokens in Residual shortcut You add this

Embeddings and positions

An embedding is just a learnable lookup table: weight[token_id]. Attention, the next step, treats its input as an unordered set: swap two words and every output swaps with them, nothing else changes. So the model must be told the order. The paper adds a fixed pattern of sines and cosines of different frequencies to each position:

PE(pos, 2i) = sin(pos / 10000^2i/d)    PE(pos, 2i+1) = cos(pos / 10000^2i/d)

Low dimensions wiggle fast, high dimensions slowly, like the hands of a clock, so every position gets a unique fingerprint and nearby positions get similar ones. The paper found learned position embeddings worked about as well. Modern LLMs mostly use rotary embeddings (RoPE) instead; the Positional encoding page covers all three. nn.Transformer includes neither the embedding nor any positional encoding.

Attention, by hand

Attention answers one question for every position: "which other positions should I read from, and how much?" Each position makes three vectors from its own embedding with three learned matrices:

• Query q: what am I looking for?
• Key k: what do I contain, as an address?
• Value v: what do I hand over if you pick me?

Score every query against every key with a dot product, divide by √d_k so the numbers don't grow with the width, softmax each row into weights that sum to 1, and take that weighted average of the values:

Attention(Q, K, V) = softmax(Q·Kᵀ / √d_k) · V

That's the whole mechanism. PyTorch's reference version of it (the docstring of F.scaled_dot_product_attention) is a dozen lines: a 1/√d scale factor, an optional -inf mask, q @ kᵀ, softmax, then @ value.

Multi-head attention

One attention pattern per layer is limiting: a word might need its subject and the word right before it. So the layer runs 8 attentions in parallel, each on its own 64-wide slice (512 / 8), then concatenates the 8 outputs back to 512 and mixes them with one more matrix. Total cost is about the same as one 512-wide head; you just get 8 different patterns. PyTorch asserts that the width divides evenly (head_dim = embed_dim // num_heads).

The feed-forward block

Attention moves information between positions. The feed-forward block then works on each position alone: widen to 2048, apply ReLU, narrow back to 512. Same weights at every position. It holds two thirds of each encoder layer's parameters and is where much of a model's stored knowledge is thought to live. In PyTorch it is linear2(dropout(activation(linear1(x)))); activation is ReLU by default, or GELU if you pass "gelu" (_get_activation_fn). Modern LLMs use a gated variant, SwiGLU.

Residuals and LayerNorm

Every sublayer is wrapped the same way: keep a copy of the input, add the sublayer's output to it (the residual, from ResNet), and normalize. The residual means each layer only learns a small edit to a running vector, the residual stream, and gradients get a straight path back to the first layer. LayerNorm rescales each position's 512 numbers to mean 0 and variance 1, then applies a learned scale and shift.

Where the norm goes turned out to matter. The paper normalizes after the add (post-norm, PyTorch's default). Almost every model since normalizes before the sublayer (pre-norm), which trains stably without a careful warmup. PyTorch supports both with norm_first, in twelve lines that cite the pre-norm paper:

The encoder stack

An encoder layer is self-attention then feed-forward, each wrapped as above. TransformerEncoder deep-copies one layer N times (_get_clones, so the copies start identical and train apart), runs them in a plain loop, and applies the optional final norm. At inference time, under the right conditions, the layer skips the Python code entirely and calls one fused kernel (torch._transformer_encoder_layer_fwd).

The decoder stack

A decoder layer has three sublayers. Masked self-attention lets each output position read the earlier outputs only. Cross-attention is the bridge: its queries come from the decoder, but its keys and values come from the encoder's memory, so each French word can look up the English words it needs (_mha_block calls multihead_attn(x, mem, mem)). Then the usual feed-forward.

Masks: who may look at whom

A mask is added to the scores before softmax. -inf becomes a weight of exactly 0, so a blocked key contributes nothing. The decoder needs a causal mask during training: it sees the whole French sentence at once, and without the mask position 2 would simply copy the answer from position 3. Transformer.generate_square_subsequent_mask(n) builds it as torch.triu(torch.full((n, n), -inf), diagonal=1) (source):

tensor([[0., -inf, -inf, -inf],
        [0.,   0., -inf, -inf],
        [0.,   0.,   0., -inf],
        [0.,   0.,   0.,   0.]])

Watch out for the boolean form: in PyTorch a True means blocked (docs; converted to -inf by _canonical_mask). That's the opposite of F.scaled_dot_product_attention's attn_mask, where True means "may attend". Padding masks (*_key_padding_mask) work the same way, one row per batch item.

The output head

The decoder's final vectors become predictions with one Linear(512, vocab) and a softmax. The paper shares one matrix between the source embedding, target embedding and this output layer (weight tying), which saves parameters and helps. None of it is inside nn.Transformer: its forward returns the decoder output, [batch, tgt_len, d_model]. You add the head, and train with cross-entropy on the logits (softmax is folded into the loss).

Worked example

With JavaScript on, the table above is live: toggle the causal mask, turn off the √d scaling, or sharpen and flatten the softmax with a temperature, and watch each row re-weight. Try a temperature of 0.2: every row collapses onto its single best key, which is why the scale factor matters. Without it, dot products in 64 dimensions get large, softmax saturates, and gradients vanish (paper, §3.2.1).

Counting parameters

Implementation

In PyTorch

The library version, with the pieces nn.Transformer leaves out:

import math
import torch
from torch import nn

def sinusoidal(n, d):
    pos = torch.arange(n).unsqueeze(1)
    angle = pos / 10000 ** (torch.arange(0, d, 2) / d)
    pe = torch.zeros(n, d)
    pe[:, 0::2], pe[:, 1::2] = angle.sin(), angle.cos()
    return pe

d_model, vocab = 512, 37000
embed = nn.Embedding(vocab, d_model)
model = nn.Transformer(d_model=d_model, nhead=8, batch_first=True)
head = nn.Linear(d_model, vocab, bias=False)
head.weight = embed.weight                      # weight tying

def encode(ids):                                # [batch, seq] -> [batch, seq, d_model]
    return embed(ids) * math.sqrt(d_model) + sinusoidal(ids.size(1), d_model)

src = torch.randint(vocab, (2, 10))
tgt = torch.randint(vocab, (2, 7))
mask = nn.Transformer.generate_square_subsequent_mask(7)
out = model(encode(src), encode(tgt), tgt_mask=mask, tgt_is_causal=True)
logits = head(out)                              # [2, 7, vocab]

And attention plus one encoder layer from scratch, matching PyTorch's post-norm layer parameter for parameter:

def attention(q, k, v, causal=False):
    scores = q @ k.transpose(-2, -1) / math.sqrt(q.size(-1))
    if causal:
        n = scores.size(-1)
        blocked = torch.ones(n, n, dtype=torch.bool).triu(1)
        scores = scores.masked_fill(blocked, float("-inf"))
    return scores.softmax(dim=-1) @ v

class EncoderLayer(nn.Module):
    def __init__(self, d=512, heads=8, ff=2048, p=0.1):
        super().__init__()
        self.attn = nn.MultiheadAttention(d, heads, dropout=p, batch_first=True)
        self.ff = nn.Sequential(nn.Linear(d, ff), nn.ReLU(), nn.Dropout(p), nn.Linear(ff, d))
        self.norm1, self.norm2 = nn.LayerNorm(d), nn.LayerNorm(d)
        self.drop1, self.drop2 = nn.Dropout(p), nn.Dropout(p)

    def forward(self, x, mask=None):
        a, _ = self.attn(x, x, x, attn_mask=mask, need_weights=False)
        x = self.norm1(x + self.drop1(a))
        return self.norm2(x + self.drop2(self.ff(x)))

sum(p.numel() for p in EncoderLayer().parameters())   # 3,152,384

Tradeoffs

What it costs

Attention compares every position with every other, so its score matrix is seq × seq per head: doubling the context quadruples that work and memory, the O(n²·d) term. The feed-forward block is O(n·d²): linear in length. For short sequences the feed-forward block dominates; for long ones attention does. That trade-off is why the paper could drop recurrence (every position is computed in parallel, with a path of length 1 between any two words), and why long-context models fight so hard over attention. Optimizing it without changing the math is what fused kernels such as FlashAttention do, and what F.scaled_dot_product_attention dispatches to.

Training versus inference

Pitfalls

• batch_first defaults to False. Without it, PyTorch reads [seq, batch, d]; a [batch, seq, d] tensor runs without error and silently attends across the batch (defaults).
• Boolean masks are inverted between APIs. True blocks in nn.MultiheadAttention and nn.Transformer, but allows in F.scaled_dot_product_attention.
• No positions, no order. Forget the positional encoding and the encoder treats the input as a bag of words.
• need_weights=True is the default for nn.MultiheadAttention, which takes the slower path and averages weights over heads (signature). Pass need_weights=False when you only need the output; the transformer layers already do.
• Post-norm needs warmup. With the default norm_first=False, a large learning rate early on diverges; use warmup or norm_first=True.
• Don't forget model.eval() at inference, or dropout stays on.

Check yourself

Connections

References

• Vaswani, Shazeer, Parmar, Uszkoreit, Jones, Gomez, Kaiser and Polosukhin. Attention Is All You Need. 2017. arXiv:1706.03762
• Ba, Kiros and Hinton. Layer Normalization. 2016. arXiv:1607.06450
• He, Zhang, Ren and Sun. Deep Residual Learning for Image Recognition. 2015. arXiv:1512.03385
• Xiong et al. On Layer Normalization in the Transformer Architecture. 2020. arXiv:2002.04745
• Press and Wolf. Using the Output Embedding to Improve Language Models. 2016. arXiv:1608.05859
• Bahdanau, Cho and Bengio. Neural Machine Translation by Jointly Learning to Align and Translate. 2014. arXiv:1409.0473
• Dao, Fu, Ermon, Rudra and Ré. FlashAttention. 2022. arXiv:2205.14135