Attention Is All You Need, in depth: a modern summary that follows one sentence through the encoder-decoder, the three masks, teacher forcing versus decoding, and a toy model you can train

The paper in five claims

1. Attention alone is enough. A stack of self-attention and position-wise feed-forward layers can replace recurrent encoders and decoders for translation.
2. Scaled dot-product attention. Attention(Q, K, V) = softmax(Q Kᵀ / √d_k) V, where the division by √d_k keeps the softmax out of its saturated, small-gradient region as the key dimension grows.
3. Multi-head attention. Instead of one attention over 512 dimensions, the base model runs eight heads of 64 dimensions each in parallel and concatenates them, so different heads can attend to different relationships.
4. Order comes from position encodings. Attention is permutation-invariant, so fixed sinusoidal vectors are added to the embeddings. The authors report that learned position embeddings gave nearly identical results.
5. It trains faster and translates better. The big model reached 28.4 BLEU on WMT 2014 English-German and 41.8 on English-French, the latter after 3.5 days on eight P100 GPUs, at a small fraction of the training cost of the best earlier models, by the paper's estimates.

Why recurrence was the bottleneck

An RNN reads token t only after it has finished token t-1, so the work inside one sequence cannot be spread across a GPU's parallel units, and information from the first word must survive n state updates to reach the last. The paper's Table 1 compares layer types on exactly these points, with n the sequence length and d the representation width:

Self-attention does more arithmetic per layer when n exceeds d, but all positions are computed at once and any two tokens are one step apart. For sentences of a few dozen tokens with d = 512, that was a clear win. The n² term is also the reason long-context work since 2017 has focused on attention's cost.

Following one sentence through the model

Take an English sentence to be translated into German. It is first split into subword tokens; the paper used byte-pair encoding with a shared English-German vocabulary of about 37,000 tokens. Each token id becomes a 512-dimensional embedding, scaled by √512, plus the sinusoidal position vector for its index. For a source sentence of S tokens the encoder input is an S by 512 matrix.

The encoder applies six identical layers. Each has a self-attention sublayer, in which every source token builds its new representation as a weighted mix of all source tokens, and a feed-forward sublayer of width 2048 applied to each position separately. Each sublayer is wrapped as LayerNorm(x + Sublayer(x)), the post-norm arrangement. The output, still S by 512, is called the memory, and it is computed once per sentence.

The decoder produces the German one token at a time. Its input is the German tokens produced so far, shifted right behind a start token. Each of its six layers has three sublayers: masked self-attention over the German prefix; cross-attention, where queries come from the decoder and keys and values come from the encoder memory; and a feed-forward network. A final linear layer, whose weights the paper shares with the embedding matrices, turns the last position into scores over the vocabulary, and a softmax turns those into the distribution for the next token.

Cross-attention is how translation happens: at every decoder layer, each German position asks which English positions are relevant right now. Cross-attention covers that mechanism and its later uses beyond translation.

The three masks

Masks decide which keys each query may see. A masked score is set to a very large negative number before the softmax, so its weight becomes zero. The original model needs three:

• Encoder padding mask. Sentences in a batch have different lengths and are padded to the same length. No query may attend to a padding key.
• Decoder causal mask. In training the whole German target is fed at once, so position t must be prevented from seeing positions after t, or it would learn to copy the answer. Combined with a padding mask for the target.
• Cross-attention mask. Decoder queries may attend to every real source token, but not to source padding.

The numpy demo below builds the causal and source-padding masks for a tiny example and prints the resulting attention weights.

import numpy as np

PAD = 0
src = np.array([[5, 9, 7, PAD]])          # batch=1, source length 4, last is padding
tgt = np.array([[1, 4, 8]])               # decoder input: <bos> y1 y2

src_pad = (src != PAD)[:, None, None, :]  # (B, 1, 1, S): which keys are real
causal = np.tril(np.ones((3, 3), dtype=bool))[None, None]  # (1, 1, T, T)

def attend(q, k, v, mask):
    d_k = q.shape[-1]
    scores = q @ k.swapaxes(-1, -2) / np.sqrt(d_k)
    scores = np.where(mask, scores, -1e9)
    w = np.exp(scores - scores.max(-1, keepdims=True))
    return (w / w.sum(-1, keepdims=True)) @ v, w / w.sum(-1, keepdims=True)

rng = np.random.default_rng(0)
x = rng.normal(size=(1, 1, 3, 4))         # 3 decoder positions, d_k=4
mem = rng.normal(size=(1, 1, 4, 4))       # 4 encoder positions

_, w_self = attend(x, x, x, causal)
_, w_cross = attend(x, mem, mem, src_pad)
np.set_printoptions(precision=2, suppress=True)
print("decoder self-attention weights (rows = query position):")
print(w_self[0, 0])
print("cross-attention weights (column 4 is the padded source token):")
print(w_cross[0, 0])

Running it prints:

decoder self-attention weights (rows = query position):
[[1.   0.   0.  ]
 [0.25 0.75 0.  ]
 [0.17 0.13 0.7 ]]
cross-attention weights (column 4 is the padded source token):
[[0.23 0.48 0.3  0.  ]
 [0.14 0.6  0.26 0.  ]
 [0.67 0.24 0.1  0.  ]]

The self-attention matrix is lower-triangular: position 0 can only see itself, position 2 sees all three. The cross-attention matrix has a zero column for the padded source token. Two practical notes. The demo uses -1e9 rather than negative infinity, because a row in which every key is masked would otherwise compute 0/0 and produce NaN. And this demo uses True to mean may attend, while PyTorch's nn.Transformer uses True to mean blocked; mixing the two conventions is the most common bug in reimplementations. Causal versus bidirectional attention explains why the causal mask also defines the training objective.

Teacher forcing in training, one token at a time in inference

Training and inference run the decoder differently. In training, the model sees the correct target prefix at every position, shifted right by one, and predicts the next token at all positions in a single parallel pass. This is teacher forcing: with the causal mask, one forward pass yields a loss for every target token.

At inference there is no correct prefix. The encoder runs once, then the decoder runs repeatedly: predict a token, append it, run again, until an end token appears. The paper used beam search with a beam of four and a length penalty of 0.6 rather than picking the single best token each step. A naive loop recomputes attention for the whole prefix at every step; implementations cache each layer's keys and values for earlier positions, which is the KV cache every LLM server relies on today.

Training used label smoothing of 0.1, which spreads a little probability over wrong tokens. The authors note it hurts perplexity, because the model becomes less sure, but improves accuracy and BLEU.

A toy model you can train

The script below uses PyTorch's nn.Transformer, which implements the paper's encoder-decoder with post-norm layers by default, at a much smaller size: two layers each, d_model 128, four heads. The task is to reverse a sequence of digits, which a model can only do by using cross-attention to look up the right source position for each output position. It keeps the paper's sinusoidal positions, embedding scaling, weight tying, Adam settings, warmup schedule and label smoothing.

import math
import torch
import torch.nn as nn

V, PAD, BOS, EOS = 13, 0, 1, 2          # ids 3..12 are the digits 0..9
D, MAXLEN = 128, 16

class SinusoidalPositions(nn.Module):
    def __init__(self, d, max_len=64):
        super().__init__()
        pos = torch.arange(max_len).unsqueeze(1)
        div = torch.exp(torch.arange(0, d, 2) * (-math.log(10000.0) / d))
        pe = torch.zeros(max_len, d)
        pe[:, 0::2] = torch.sin(pos * div)
        pe[:, 1::2] = torch.cos(pos * div)
        self.register_buffer("pe", pe)

    def forward(self, x):                # x: (B, T, d)
        return x + self.pe[: x.size(1)]

def causal_mask(t):
    """Bool mask, True = blocked: position i may not see j > i."""
    return torch.triu(torch.ones(t, t, dtype=torch.bool), diagonal=1)

class Seq2Seq(nn.Module):
    def __init__(self):
        super().__init__()
        self.emb = nn.Embedding(V, D, padding_idx=PAD)   # shared: encoder, decoder, output
        self.pos = SinusoidalPositions(D)
        self.tf = nn.Transformer(d_model=D, nhead=4, num_encoder_layers=2,
                                 num_decoder_layers=2, dim_feedforward=256,
                                 dropout=0.1, batch_first=True)

    def embed(self, ids):
        return self.pos(self.emb(ids) * math.sqrt(D))

    def forward(self, src, tgt_in):
        causal = causal_mask(tgt_in.size(1))
        h = self.tf(self.embed(src), self.embed(tgt_in), tgt_mask=causal,
                    src_key_padding_mask=(src == PAD),      # True = ignore
                    tgt_key_padding_mask=(tgt_in == PAD),
                    memory_key_padding_mask=(src == PAD))
        return h @ self.emb.weight.T                         # tied output projection

def batch(n=64):
    """Task: output the input digits reversed."""
    src = torch.full((n, MAXLEN), PAD)
    tgt = torch.full((n, MAXLEN), PAD)
    for i in range(n):
        L = int(torch.randint(3, 11, ()))
        d = torch.randint(3, V, (L,))
        src[i, :L] = d
        tgt[i, 0], tgt[i, 1:L + 1], tgt[i, L + 1] = BOS, d.flip(0), EOS
    return src, tgt

model = Seq2Seq()
opt = torch.optim.Adam(model.parameters(), lr=1.0, betas=(0.9, 0.98), eps=1e-9)
warmup = 400   # the paper used 4000 steps for a much larger model
sched = torch.optim.lr_scheduler.LambdaLR(
    opt, lambda s: D ** -0.5 * min((s + 1) ** -0.5, (s + 1) * warmup ** -1.5))
loss_fn = nn.CrossEntropyLoss(ignore_index=PAD, label_smoothing=0.1)

for step in range(3000):
    src, tgt = batch()
    tgt_in, tgt_out = tgt[:, :-1], tgt[:, 1:]               # teacher forcing: shift right
    logits = model(src, tgt_in)
    loss = loss_fn(logits.reshape(-1, V), tgt_out.reshape(-1))
    opt.zero_grad(); loss.backward(); opt.step(); sched.step()
    if step % 500 == 0:
        print(step, round(loss.item(), 3))

@torch.no_grad()
def greedy(src):
    model.eval()
    memory = model.tf.encoder(model.embed(src), src_key_padding_mask=(src == PAD))
    out = torch.full((src.size(0), 1), BOS)
    for _ in range(MAXLEN - 1):                              # one decoder pass per token
        causal = causal_mask(out.size(1))
        h = model.tf.decoder(model.embed(out), memory, tgt_mask=causal,
                             memory_key_padding_mask=(src == PAD))
        nxt = (h[:, -1] @ model.emb.weight.T).argmax(-1, keepdim=True)
        out = torch.cat([out, nxt], dim=1)
        if (nxt == EOS).all():
            break
    return out

The script is small enough to train on a laptop CPU; its output is not shown here. Read the loss curve diagnostically. Label smoothing keeps it above zero. A loss that drops suspiciously fast means a leak: an unshifted target or an inverted causal mask lets positions see their own answers. A loss that stays near log(13), about 2.56, means the model cannot use the source, so check the polarity of the padding masks. When it works, greedy() returns BOS, the reversed digits and EOS.

What a modern reader should update

The core is unchanged in today's models: scaled dot-product attention, multiple heads, residual connections around attention and feed-forward sublayers. Most of the rest has moved on. Large language models are usually decoder-only, dropping the encoder and cross-attention. LayerNorm moved before each sublayer (pre-norm), which trains deep stacks more stably than the paper's post-norm. Sinusoidal absolute positions gave way to relative schemes such as rotary position embeddings. Feed-forward layers use gated activations, and attention kernels are rewritten to avoid materialising the n² score matrix. Transformer block anatomy walks through a modern block with shapes and parameter counts.

Bugs people hit when reimplementing it

• Inverted masks. Pass a True-means-keep mask to an API that expects True-means-blocked and the effect flips: an inverted padding mask hides real tokens and exposes padding, and an inverted causal mask exposes only the future, so training loss looks excellent while generation fails. Its last row allows no keys at all, which can produce NaN.
• No shift. Feeding the target unshifted makes position t predict the token it already sees.
• Missing scale. Forgetting the √d_k divisor lets dot products grow with d_k and pushes softmax towards saturation, the effect the paper's scaling was meant to prevent.
• Softmax on the wrong axis. Normalise over keys, the last dimension of the score matrix, not over queries.
• Unscaled embeddings with tied weights. The paper multiplies embeddings by √d_model; a common explanation is that without it the position signal can swamp the token signal.
• Rows with every key masked. Negative infinity in the mask yields NaN; use a large finite negative value or make sure every query has at least one valid key.

What to do next

1. Read the paper's Sections 3.1 to 3.5 with the diagram above beside it and match every sublayer to a box.
2. Run the numpy mask demo, then change the padding position and confirm the zero column moves.
3. Train the toy model, then deliberately break one mask and watch how the loss and the greedy outputs change.
4. Add beam search with a beam of four to the toy model and compare it with greedy decoding.
5. Convert the toy model to pre-norm with norm_first=True and compare training stability at more layers.
6. Read the revisited article for the full training recipe, then the block anatomy article for what modern decoder-only models changed.