Attention Is All You Need

A single attention head has a critical limitation: it produces a single set of attention weights per position. This means each token can only focus on one "aspect" of the input at a time.

Consider the sentence: "The animal didn't cross the street because it was too tired."

To understand this sentence, we need to resolve multiple relationships simultaneously:

• Coreference: "it" → "animal" (not "street") — requires semantic reasoning
• Syntax: "didn't" modifies "cross" — requires syntactic structure
• Proximity: "too tired" relates to "it" — requires positional awareness
• Negation scope: "didn't cross" — requires understanding modifier scope

A single attention head would be forced to compromise between these patterns, mixing syntactic and semantic signals into one weighted average. The result is a blurred, less useful representation.

Furthermore, with $d_k = d_{model} = 512$, a single head's attention scores are computed from 512-dimensional dot products. This high dimensionality means that many unrelated features contribute to every attention decision, making it harder to learn focused patterns.

Multi-Head Attention solves this by splitting the representation space into $h$ independent subspaces, running attention in each, then recombining.

Dimension splitting — the key insight:

Instead of one attention with $d_k = 512$, use 8 heads each with $d_k = 64$:

$d_k = d_v = d_{model} / h = 512 / 8 = 64$

Step-by-step with concrete dimensions (base model):

1. Project into $h$ subspaces: For each head $i = 1, ..., 8$:

• $Q_i = XW^Q_i$ where $W^Q_i \in \mathbb{R}^{512 \times 64}$
• $K_i = XW^K_i$ where $W^K_i \in \mathbb{R}^{512 \times 64}$
• $V_i = XW^V_i$ where $W^V_i \in \mathbb{R}^{512 \times 64}$

Each head has its own $W^Q_i$, $W^K_i$, $W^V_i$ — 24 separate weight matrices total (8 heads × 3 matrices).

2. Compute attention independently: $\text{head}_i = \text{Attention}(Q_i, K_i, V_i) \in \mathbb{R}^{n \times 64}$

Each head produces a 64-dimensional output per token.

3. Concatenate all heads: $\text{Concat}(\text{head}_1, ..., \text{head}_8) \in \mathbb{R}^{n \times 512}$

Stacking 8 heads of 64 dims = 512 dims = $d_{model}$. No information is lost.

4. Final linear projection: $\text{MultiHead}(Q, K, V) = \text{Concat}(...)W^O$ where $W^O \in \mathbb{R}^{512 \times 512}$

$W^O$ mixes information across heads, allowing the model to learn combinations of the different attention patterns.

Computational cost:

Total parameters per multi-head attention layer:

• Projections: $h \times 3 \times (d_{model} \times d_k)$ = $8 \times 3 \times (512 \times 64)$ = 786,432
• Output: $d_{model} \times d_{model}$ = $512 \times 512$ = 262,144
• Total: ~1M parameters — same as a single-head with $d_k = d_{model}$

The magic: multi-head costs the same as single-head attention (since $h \times d_k = d_{model}$), but captures richer patterns.

What different heads learn (empirically observed):

Research has shown that different heads tend to specialize:

• Positional heads: Attend to adjacent tokens (local context)
• Syntactic heads: Follow dependency parse structure (subject → verb)
• Rare token heads: Pay attention to infrequent/important words
• Separator heads: Attend to punctuation and sentence boundaries
• Coreference heads: Link pronouns to their referents

Self-attention has no notion of position. Formally, if you permute the input sequence, the output is permuted in exactly the same way — the attention weights don't change.

This means:

• "The cat sat on the mat" and "mat the on sat cat The" would produce identical attention patterns (just reordered)
• The model cannot distinguish "dog bites man" from "man bites dog"
• Word order, which is fundamental to meaning in most languages, is completely invisible

Why this happens: The attention score $q_i \cdot k_j$ depends only on the content of tokens $i$ and $j$, not their positions. Permuting the input just permutes which dot products are computed — the values remain the same.

This is by design (it makes attention flexible), but we need a mechanism to inject positional information without breaking the architecture.

Sinusoidal Positional Encodings add fixed, position-dependent vectors to token embeddings:

$\text{input}_i = \text{embedding}_i + PE_i$

where $PE \in \mathbb{R}^{n \times d_{model}}$ encodes each position as a unique vector.

The encoding formula:

• Even dimensions: $PE_{(pos, 2i)} = \sin(pos / 10000^{2i/d_{model}})$
• Odd dimensions: $PE_{(pos, 2i+1)} = \cos(pos / 10000^{2i/d_{model}})$

Understanding the frequencies:

Each pair of dimensions $(2i, 2i+1)$ forms a sinusoidal wave with a specific wavelength (period):

$\lambda_i = 2\pi \cdot 10000^{2i/d_{model}}$

• Dimension 0-1: wavelength = $2\pi \approx 6.28$ positions (high frequency, changes every ~3 positions)
• Dimension 254-255: wavelength = $2\pi \cdot 10000^{254/512} \approx 61,600$ positions (very low frequency)
• This creates a spectrum from fine-grained to coarse-grained position signals

Why sinusoids enable relative position attention:

The key mathematical property: for any fixed offset $k$, $PE_{pos+k}$ can be expressed as a linear transformation of $PE_{pos}$:

$PE_{pos+k} = M_k \cdot PE_{pos}$

where $M_k$ is a rotation matrix that depends only on the offset $k$, not on $pos$. This means the model can learn to attend to "3 positions to the right" through a simple linear function — regardless of absolute position.

Why not learned embeddings?

The paper tested learned positional embeddings and found nearly identical results to sinusoidal. However, sinusoidal has two advantages:

1. Extrapolation: Can handle sequences longer than those seen in training (the waves continue naturally)
2. No extra parameters: The positional encoding requires zero learned parameters

Modern alternatives:

• RoPE (Rotary Positional Embedding): Used in LLaMA, Mistral. Encodes relative position directly into the attention computation via rotation matrices, rather than adding to embeddings. Better for extrapolation.
• ALiBi (Attention with Linear Biases): Used in BLOOM. Adds a linear distance penalty to attention scores: closer tokens get higher scores.
• Learned absolute embeddings: Used in GPT-2, BERT. Simple but cannot extrapolate beyond training length.

Sequence-to-sequence tasks like translation require two fundamentally different capabilities:

1. Encoding (understanding): Process the full input sentence and build a representation that captures all relationships. The encoder can (and should) look at the entire input — knowing the last word helps understand the first word.

2. Decoding (generation): Produce output tokens one at a time, left-to-right. At step $t$, the decoder can only see tokens $1, ..., t-1$ (the tokens it has already generated). It must also access the encoder's representation of the input.

RNN seq2seq limitations:

• Encoder compresses the entire input into a single fixed-size vector (the final hidden state). For long sentences, this bottleneck loses information catastrophically.
• Bahdanau attention (2014) partially addressed this by letting the decoder attend to all encoder hidden states, but the encoder itself was still sequential.
• Training was inherently sequential: each timestep depends on the previous one, making GPU parallelization impossible.

The Transformer architecture cleanly separates encoding and decoding into two stacks of layers, connected by cross-attention.

Complete data flow through the architecture (base model dimensions):

Input Processing:

1. Tokenize input text → token IDs (vocabulary size ~37,000 for WMT English-German)
2. Embed token IDs → $X_{embed} \in \mathbb{R}^{n \times 512}$ via learned embedding matrix $E \in \mathbb{R}^{V \times 512}$
3. Scale embeddings by $\sqrt{d_{model}}$ (prevents positional encodings from dominating small embedding values)
4. Add positional encoding: $X_0 = \sqrt{d_{model}} \cdot X_{embed} + PE \in \mathbb{R}^{n \times 512}$
5. Apply dropout ($p = 0.1$) to the sum

Encoder (6 identical layers):

Each encoder layer $l$ applies two sub-layers:

$\text{Encoder Layer}(X) = \text{FFN}(\text{SelfAttn}(X))$

With residual connections and layer normalization around each:

1. $Z_l = \text{LayerNorm}(X_{l-1} + \text{MultiHeadSelfAttn}(X_{l-1}))$
   • Self-attention: $Q, K, V$ all come from the previous layer's output
   • Bidirectional: every token attends to every other token (no masking)
2. $X_l = \text{LayerNorm}(Z_l + \text{FFN}(Z_l))$
   • Position-wise FFN: $512 \to 2048 \to 512$

The output of the final encoder layer is $\text{EncoderOut} \in \mathbb{R}^{n \times 512}$, which is passed to every decoder layer.

Decoder (6 identical layers):

Each decoder layer applies three sub-layers:

1. Masked Self-Attention: $Q, K, V$ from previous decoder output

   • Uses causal mask to prevent attending to future positions
   • At training: processes all positions in parallel (teacher forcing)
   • At inference: processes one new token at a time (autoregressive)

2. Cross-Attention: $Q$ from decoder, $K, V$ from encoder output

   • This is how the decoder "reads" the input sequence
   • The decoder queries the encoder: "given what I've generated so far, what part of the input is relevant for the next token?"
   • $K$ and $V$ are the same encoder output for all decoder layers — computed once, reused 6 times
   • Not masked: the decoder can attend to the entire input

3. Feed-Forward Network: Same structure as encoder (512 → 2048 → 512)

Output Processing:

1. Final decoder output $\in \mathbb{R}^{m \times 512}$ (where $m$ = output sequence length)
2. Linear projection $W_{out} \in \mathbb{R}^{512 \times V}$ maps to vocabulary logits
3. Softmax → probability distribution over vocabulary
4. The output embedding matrix $W_{out}$ is shared with the input embedding matrix — this weight tying reduces parameters and improves performance

Training vs. Inference:

• Training: Both encoder and decoder process full sequences in parallel (teacher forcing — ground truth is fed as decoder input). The causal mask simulates autoregressive behavior.
• Inference: Encoder runs once. Decoder runs $m$ times (one per output token), each time extending the output sequence by one token. Each step reuses cached $K, V$ from previous steps (KV-cache).

Three types of attention in the Transformer:

1. Encoder self-attention: Q, K, V all from encoder. Bidirectional. Builds rich input representations.
2. Decoder masked self-attention: Q, K, V all from decoder. Causal mask prevents future peeking. Models the output sequence structure.
3. Encoder-decoder cross-attention: Q from decoder, K and V from encoder output. This is the only connection between encoder and decoder — it's how generated tokens access input information.

How these connect across layers:

• Encoder layer 1 processes raw embeddings → encoder layer 2 processes layer 1's output → ... → encoder layer 6 produces final representations
• Decoder layer 1 processes output embeddings (with mask) → then cross-attends to encoder layer 6's output → then FFN → decoder layer 2 takes this output → cross-attends to the same encoder output → ... → decoder layer 6 → linear → softmax

Each encoder and decoder layer maintains the same dimensionality: $\mathbb{R}^{n \times 512}$ (encoder) and $\mathbb{R}^{m \times 512}$ (decoder). The residual connections enforce this — you can't change the dimension within a residual block.

Self-attention computes context-aware representations by mixing information across positions, but it has a subtle limitation: the mixing is linear.

The attention output for each token is a weighted average of value vectors: $\text{out}_i = \sum_j \alpha_{ij} v_j$. While the weights $\alpha_{ij}$ are non-linear functions of the input (through softmax), the combination itself is linear — it's just a weighted sum.

This means attention alone can only represent functions that are linear combinations of the input values. Complex per-token transformations — like detecting whether a token is a noun, computing whether a number is greater than a threshold, or mapping between representation spaces — require non-linear processing.

Additionally, attention has relatively few parameters per layer (~1M for base model). The model needs a larger parameter capacity to store the knowledge required for language understanding.

Feed-Forward Network (FFN) adds non-linear, per-token processing capacity:

$FFN(x) = \text{ReLU}(xW_1 + b_1)W_2 + b_2$

Concrete dimensions (base model):

For each token independently (same weights, applied to each position):

1. Expand: $xW_1 + b_1$ where $W_1 \in \mathbb{R}^{512 \times 2048}$, $b_1 \in \mathbb{R}^{2048}$ — projects from 512 to 2048 dimensions (4x expansion)
2. Activate: ReLU (or GELU/SwiGLU in modern variants) — introduces non-linearity. After ReLU, approximately half the neurons are zeroed out, creating a sparse intermediate representation
3. Contract: $\cdot W_2 + b_2$ where $W_2 \in \mathbb{R}^{2048 \times 512}$, $b_2 \in \mathbb{R}^{512}$ — projects back to 512 dimensions

Parameter count per FFN layer:

• $W_1$: $512 \times 2048$ = 1,048,576
• $b_1$: 2,048
• $W_2$: $2048 \times 512$ = 1,048,576
• $b_2$: 512
• Total: ~2.1M parameters per FFN layer
• Compare to attention: ~1M parameters per attention layer
• FFN contains ~2/3 of each layer's parameters

Across all 6 encoder + 6 decoder layers: $12 \times 2.1M \approx$ 25M FFN parameters (of 65M total in the base model)

Why the 4x expansion?

The expansion creates a higher-dimensional intermediate space where the model can represent more complex features. Recent research ("Transformer Feed-Forward Layers Are Key-Value Memories", Geva et al. 2021) suggests that each row of $W_1$ acts as a pattern detector (key) and the corresponding column of $W_2$ stores the associated information (value). The FFN essentially functions as a learned lookup table — a compressed knowledge base.

Modern FFN variants:

• GELU (used in GPT, BERT): Smoother than ReLU, $\text{GELU}(x) = x \cdot \Phi(x)$ where $\Phi$ is the Gaussian CDF
• SwiGLU (used in LLaMA, PaLM): $\text{SwiGLU}(x) = (xW_1 \odot \text{Swish}(xW_g))W_2$ — adds a gating mechanism with a third weight matrix $W_g$. Typically uses $d_{ff} = \frac{8}{3}d_{model}$ (instead of 4x) to keep parameter count comparable
• MoE (Mixture of Experts): Replaces single FFN with multiple "expert" FFNs + a router. Only 2-4 experts activate per token, enabling massive parameter count with fixed compute (used in Mixtral, Switch Transformer)

The Transformer base model has 12 layers (6 encoder + 6 decoder), each containing attention and FFN sub-layers. Training such a deep network faces serious challenges:

Vanishing gradients: During backpropagation, gradients are multiplied by the Jacobian of each layer. If each layer's Jacobian has eigenvalues < 1, gradients shrink exponentially with depth: after $L$ layers, gradients can be $O(\lambda^L)$ where $\lambda < 1$. For $L = 12$ with $\lambda = 0.9$: $0.9^{12} = 0.28$ — gradients reach the first layer at only 28% strength.

Exploding activations: Without normalization, the variance of activations can grow unbounded through layers. If each layer multiplies the variance by a factor $c > 1$: after 12 layers, variance is $c^{12}$ — potentially causing numerical overflow.

Internal covariate shift: Each layer's input distribution changes as parameters in previous layers update during training, forcing each layer to continuously re-adapt.

Loss landscape complexity: Deep networks without residual connections have highly non-convex loss landscapes with many saddle points, making optimization extremely difficult.

Two complementary techniques solve these problems:

1. Residual (Skip) Connections:

$x_{out} = x + \text{SubLayer}(x)$

Instead of learning the full mapping $f(x)$, the network learns the residual $\Delta x = f(x) - x$. This is strictly easier — if the optimal transformation is close to identity, the network just needs to learn a small residual.

Why they fix gradient flow:

The gradient of the output w.r.t. the input is:

$\frac{\partial x_{out}}{\partial x} = I + \frac{\partial \text{SubLayer}(x)}{\partial x}$

The identity matrix $I$ guarantees that gradients flow directly from the loss to any layer. Even if the sub-layer's gradient vanishes, the skip connection preserves the signal. Through 12 layers, the gradient always has a direct path with no multiplicative decay.

2. Layer Normalization:

$\text{LayerNorm}(x) = \gamma \odot \frac{x - \mu}{\sqrt{\sigma^2 + \epsilon}} + \beta$

For each token independently (across the $d_{model} = 512$ features):

• $\mu = \frac{1}{d_{model}} \sum_{i=1}^{d_{model}} x_i$ (mean of the 512 features)
• $\sigma^2 = \frac{1}{d_{model}} \sum_{i=1}^{d_{model}} (x_i - \mu)^2$ (variance)
• $\gamma, \beta \in \mathbb{R}^{d_{model}}$ are learnable scale and shift (512 params each)
• $\epsilon = 10^{-5}$ prevents division by zero

LayerNorm vs. BatchNorm:

• BatchNorm normalizes across the batch dimension (statistics over $B$ examples for each feature) — depends on batch size, doesn't work well for variable-length sequences
• LayerNorm normalizes across the feature dimension (statistics over $d_{model}$ features for each token) — independent of batch size and sequence length

Parameters per LayerNorm: Just $2 \times d_{model} = 1024$ (γ and β). There are 2 LayerNorms per layer × 12 layers = 24 LayerNorms total = 24,576 parameters — negligible compared to attention and FFN.

Pre-LN vs Post-LN placement:

The original paper used Post-LN: $x_{out} = \text{LN}(x + \text{SubLayer}(x))$

Modern Transformers use Pre-LN: $x_{out} = x + \text{SubLayer}(\text{LN}(x))$

Why Pre-LN is better:

• Post-LN: The residual path passes through LayerNorm, which can distort the gradient signal. Requires learning rate warmup (the paper used 4000 warmup steps). Without warmup, training diverges.
• Pre-LN: The residual path is clean (no normalization on the skip connection). Gradients flow unimpeded. Can often train without warmup. Used in GPT-2, GPT-3, LLaMA, and most modern Transformers.

Complete sub-layer pattern (Pre-LN): $x_{out} = x + \text{Dropout}(\text{SubLayer}(\text{LayerNorm}(x)))$

During generation, we face a fundamental causality constraint:

At time step $t$, the decoder is generating token $t$. Tokens $t+1, t+2, ...$ don't exist yet — they haven't been generated. Standard self-attention would attend to all positions, including these non-existent future positions.

During training, the situation is even more subtle. We use teacher forcing: the decoder receives the ground-truth output sequence as input. All output tokens are available simultaneously, so the decoder could trivially "cheat" by looking at the next token and copying it.

Without masking, the model would learn to copy rather than predict, and would be useless at inference time (when future tokens are truly unavailable).

The problem is that we want the parallelism benefits of processing all positions simultaneously (unlike RNNs), but we also need the causal property of never looking ahead.

Causal Masking applies a lower-triangular mask to the attention scores before softmax, setting future positions to $-\infty$:

$\text{MaskedAttention}(Q, K, V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}} + M\right)V$

where $M_{ij} = 0$ if $j \leq i$ (can attend), $M_{ij} = -\infty$ if $j > i$ (blocked).

Step-by-step with a 5-token sequence:

$M = \begin{bmatrix} 0 & -\infty & -\infty & -\infty & -\infty \\ 0 & 0 & -\infty & -\infty & -\infty \\ 0 & 0 & 0 & -\infty & -\infty \\ 0 & 0 & 0 & 0 & -\infty \\ 0 & 0 & 0 & 0 & 0 \end{bmatrix}$

Why $-\infty$ works:

• $\text{softmax}(-\infty) = \frac{e^{-\infty}}{\sum e^{\cdot}} = \frac{0}{\sum e^{\cdot}} = 0$
• The masked positions get exactly zero attention weight
• The remaining weights are renormalized to sum to 1
• Gradients through masked positions are also zero — the model cannot learn to use future information

Effective attention pattern:

Token 1 attends to: [1]              (only itself)
Token 2 attends to: [1, 2]           (itself + one prior)
Token 3 attends to: [1, 2, 3]        (itself + two prior)
Token 4 attends to: [1, 2, 3, 4]     (itself + three prior)
Token 5 attends to: [1, 2, 3, 4, 5]  (all tokens)

Where masking is used (and not used):

Component	Masked?	Why
Encoder self-attention	No	Input is fully available; bidirectional helps
Decoder self-attention	Yes	Must not see future outputs
Cross-attention (dec→enc)	No	Decoder can see the entire input

KV-Cache optimization at inference:

During autoregressive generation, the decoder processes one new token at a time. Recomputing attention over all previous tokens at every step would be $O(n^2)$ per token. The KV-cache stores previously computed keys and values, so only the new token's $Q, K, V$ need to be computed:

• Step $t$: compute $Q_t, K_t, V_t$ for the new token
• Append $K_t$ to cached keys, $V_t$ to cached values
• Attend $Q_t$ to all cached $K$ values → $O(n)$ per step instead of $O(n^2)$

This reduces total generation complexity from $O(n^3)$ to $O(n^2)$.

Training a deep attention-based model from scratch presented several open questions in 2017:

• Learning rate: Too high → divergence. Too low → slow convergence. The model is sensitive to LR especially in early training, before LayerNorm statistics stabilize.
• Regularization: With 65M parameters and training sets of ~4.5M sentence pairs, overfitting is a real risk. How much dropout, where to apply it?
• Label smoothing: Standard cross-entropy training encourages the model to be "overconfident" — predicting one token with probability ~1.0 and all others ~0.0. This hurts generalization.
• Optimizer: Adam was standard, but the Transformer required specific hyperparameter choices to train well.

Learning Rate Schedule (Warmup + Inverse Square Root Decay):

$lr = d_{model}^{-0.5} \cdot \min(step^{-0.5}, \, step \cdot warmup^{-1.5})$

This schedule has two phases:

1. Linear warmup (first 4000 steps): LR increases linearly from 0 to peak. This is critical for Post-LN Transformers — without it, early gradients through LayerNorm are too large and training diverges. Peak LR at step 4000: $512^{-0.5} \cdot 4000^{-0.5} \approx 0.0014$
2. Inverse square root decay (after 4000 steps): LR decreases proportionally to $1/\sqrt{step}$. This is gentler than cosine or step decay, giving the model a long, slow annealing phase.

Optimizer:

• Adam with $\beta_1 = 0.9$, $\beta_2 = 0.98$, $\epsilon = 10^{-9}$
• The unusually high $\beta_2 = 0.98$ (vs typical 0.999) makes the second moment estimate adapt faster, which works well with the warmup schedule

Regularization (three types):

1. Residual Dropout ($p = 0.1$): Applied to:

   • The output of each sub-layer (attention, FFN) before adding to the residual
   • The sum of embeddings + positional encodings
   • Total: applied at every sub-layer in every layer

2. Attention Dropout ($p = 0.1$): Applied to the attention weights (after softmax). This randomly drops attention connections, forcing the model to build redundant attention patterns.

3. Label Smoothing ($\epsilon_{ls} = 0.1$): Instead of training with one-hot targets $[0, 0, 1, 0, ...]$, uses smoothed targets: $[\epsilon/V, \epsilon/V, 1-\epsilon, \epsilon/V, ...]$ where $V$ is vocabulary size. This:

   • Hurts perplexity (the model is less confident, so perplexity increases by ~0.5)
   • Improves BLEU and translation quality — the model learns that multiple tokens can be reasonable next predictions, leading to better beam search diversity

Training Setup:

• Hardware: 8 NVIDIA P100 GPUs
• Training time: Base model 12 hours (100K steps), Big model 3.5 days (300K steps)
• Batch size: ~25,000 source + target tokens per batch (sentences batched by approximate length)
• Gradient accumulation used to achieve effective large batch sizes

Key Results (WMT 2014):

Model	EN-DE BLEU	EN-FR BLEU	Training Cost
Previous SOTA (ensemble)	26.4	41.0	—
Transformer Base	27.3	38.1	12 hours × 8 GPUs
Transformer Big	28.4	41.0	3.5 days × 8 GPUs
Transformer Big (ensemble)	—	41.8	—

The Transformer Big achieved SOTA on both benchmarks as a single model, while training in a fraction of the time of competing approaches (some required weeks of training).

Ablation insights from the paper:

• Reducing $d_k$ without compensating heads hurts quality (total attention dimension matters)
• Bigger models don't always help — quality saturates around 6 layers for this dataset
• Learned positional embeddings perform nearly identically to sinusoidal
• Replacing attention with average pooling or convolutions significantly degrades performance — attention is genuinely necessary