LLM Atlas/Part 04

Attention Is All You Need (2017)

The 2017 architecture, component by component: QKV projections, multi-head attention, positional encoding, the feed-forward block, residuals and norms.

Published: 2 August 2026
Reading time: 15 min read
Figures: 9 figures
Equations: 17 equations

Encoder-Decoder Architecture Overview

The original Transformer (Vaswani et al., 2017) is an encoder-decoder designed for machine translation. The encoder reads the entire input simultaneously and builds rich contextual representations. The decoder generates output one token at a time, attending to both its own outputs and the encoder's representations.

Both encoder and decoder are stacks of identical layers. Each layer contains: (1) a multi-head (self-)attention sub-layer, and (2) a position-wise feed-forward network sub-layer. Residual connections and layer normalization wrap each sub-layer.

The encoder uses bidirectional self-attention (each token can attend to all others). The decoder uses masked self-attention (each token can only attend to earlier generated tokens — causal masking) plus cross-attention to the encoder.

Encoder Output

H = Encoder (X) \in R^{n \times d}

Encoder-Decoder Architecture Overview

Positional Encoding — Injecting Order

Self-attention processes all tokens simultaneously with no inherent notion of order. Without positional information, 'dog bites man' and 'man bites dog' look identical to the attention mechanism.

Solution: add a fixed positional signal to each token embedding before entering the network. The sinusoidal encoding uses a geometric progression of frequencies — low frequencies encode coarse position; high frequencies encode fine-grained position.

The key mathematical property: for any fixed offset k, PE(pos+k) is a linear function of PE(pos). This means the model can learn to detect relative offsets between positions from absolute encodings alone, without ever being explicitly taught what 'distance 5 apart' looks like.

Even dims

P E (p os, 2 i) = sin (\frac{p os}{1000 0 ^{2 i / d}})

Odd dims

P E (p os, 2 i + 1) = cos (\frac{p os}{1000 0 ^{2 i / d}})

Positional Encoding — Injecting Order

Scaled Dot-Product Self-Attention

Intuition: 'for this word's representation, which other words should influence it, and by how much?'

Step 1 — Project: each token embedding X becomes three vectors via learned matrices. Query Q = XW_Q asks 'what am I looking for?' Key K = XW_K says 'what do I advertise?' Value V = XW_V says 'what will I contribute if attended to?'

Step 2 — Score: compute raw relevance via dot products QK^T — an n×n matrix showing every pair's affinity.

Step 3 — Scale: divide by √d_k. Why? For random vectors, the dot product variance scales with d_k. Large variances push softmax into saturated flat regions where gradients vanish. Dividing by √d_k keeps variance ≈ 1.

Step 4 — Softmax: convert raw scores to a probability distribution over positions. The output is a convex combination of Value vectors — weighted average of what everyone offers, weighted by relevance.

Q, K, V projections

Q = X W_{Q}, K = X W_{K}, V = X W_{V}

Scaled dot-product attention

Attn (Q, K, V) = softmax (\frac{Q K ^{⊤}}{d _{k}}) V

Softmax

softmax (z_{i}) = \frac{e ^{z_{i}}}{\sum _{j} e ^{z_{j}}}

Scaled Dot-Product Self-Attention

Multi-Head Attention — Multiple Subspaces

A single attention computation captures one type of relationship at a time. Multi-head attention runs h independent attention operations in parallel, each in a lower-dimensional subspace of size d_k = d_model/h.

With d_model=512 and h=8, each head operates on 64 dimensions — the same total compute as single-head over the full 512 dimensions.

The outputs of all heads are concatenated and projected back to d_model via W_O. Empirically, different heads specialize: some track syntactic dependencies, some track long-range coreference, some focus on adjacent tokens, some encode semantic similarity. Multi-head attention gives the model multiple 'perspectives' simultaneously.

Each head

head_{i} = Attn (Q W_{Q}^{i}, K W_{K}^{i}, V W_{V}^{i})

Multi-head

MHA (Q, K, V) = Concat (head_{1}, \dots, head_{h}) W_{O}

Multi-Head Attention — Multiple Subspaces

python

def multi_head_attention(q, k, v, mask):
    # q, k, v shape: (batch, heads, seq_len, head_dim)
    scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(head_dim)
    if mask is not None:
        scores = scores + mask
    probs = torch.softmax(scores, dim=-1)
    return torch.matmul(probs, v)

KV-Cache: Inference Acceleration

During text generation, the model predicts one token at a time. To predict token $t$, the attention mechanism needs the Key and Value vectors of all $t-1$ previous tokens.

Recomputing these past vectors every step is O(N^3) over the sequence length. Instead, models cache the $K$ and $V$ vectors in memory. For the new token, it only computes the new $Q, K, V$, appends the new $K, V$ to the cache, and attends to the cache. This reduces generation to O(N^2) time, but introduces a massive memory bottleneck (the KV-Cache memory footprint).

KV-Cache: Inference Acceleration

python

def prefill(prompt, W):
    kv_cache = []
    for token in prompt:
        k, v = W.k(token), W.v(token)
        kv_cache.append((k, v))
    return kv_cache

def decode(new_token, kv_cache, W):
    k, v = W.k(new_token), W.v(new_token)
    kv_cache.append((k, v))
    # Attention over entire kv_cache
    return compute_attention(W.q(new_token), kv_cache)

The Causal Mask

In a decoder, tokens cannot look into the future during training, otherwise the task of 'predict the next word' becomes trivial cheating.

A causal mask is an upper-triangular matrix of $-\infty$ added to the $QK^\top$ attention scores before the softmax. Since $e^{-\infty} = 0$, the attention weights for future tokens become exactly zero. The model can only attend to itself and previous positions.

Masked Attention

softmax (\frac{Q K ^{⊤}}{d _{k}} + M) V

The Causal Mask

python

import torch

def get_causal_mask(seq_len):
    # Lower triangular matrix
    mask = torch.tril(torch.ones(seq_len, seq_len))
    # Convert 0s to -inf for softmax
    return mask.masked_fill(mask == 0, float('-inf'))

Feed-Forward Network & Residual Connections

After attention mixes information across tokens, each token position is independently processed by a feed-forward network — a two-layer MLP applied identically to every position. The FFN expands to 4×d_model (the 'expansion ratio') and projects back.

If attention asks 'what information from where?', the FFN asks 'how should I transform this information?' Research suggests FFN layers act as key-value memories — the first matrix matches patterns, the second matrix reads out stored values.

Residual connections: instead of just computing sublayer(x), compute x + sublayer(x). By the chain rule, ∂(x + f(x))/∂x = 1 + ∂f/∂x. The '1' term provides a direct gradient highway — gradients never vanish through the residual path, enabling networks of 100+ layers to train stably.

FFN (ReLU, original)

FFN (x) = max (0, x W_{1} + b_{1}) W_{2} + b_{2}

Residual connection

x \leftarrow x + sublayer (LN (x))

Feed-Forward Network & Residual Connections

python

class FFN(nn.Module):
    def __init__(self, dim, hidden_dim):
        super().__init__()
        self.up = nn.Linear(dim, hidden_dim)
        self.down = nn.Linear(hidden_dim, dim)
        self.act = nn.GELU()

    def forward(self, x):
        return self.down(self.act(self.up(x)))

Cross-Entropy Loss & Training Objective

From information theory: entropy H(p) = −Σ p·log p measures irreducible uncertainty. Cross-entropy H(y,p) = −Σ y·log p is the expected code length when using model distribution p to encode samples from true distribution y. Minimizing cross-entropy drives p toward y.

For causal LMs (GPT-style): L = −Σ_t log P(x_t | x_{<t}) — sum the negative log-probability of each token given its preceding context. Teacher forcing: always feed ground-truth previous tokens during training (not the model's own predictions), which speeds and stabilizes learning.

Perplexity = exp(L_CE) is the model's 'effective branching factor': how many equally likely choices it appears to be making at each step. Perplexity 10 means the model behaves as if choosing uniformly among 10 candidates — lower is better.

Cross-entropy

L_{C E} = - i \sum y_{i} lo g (p_{i})

Causal LM objective

L = - t \sum lo g P (x_{t} ∣ x_{< t})

Perplexity

PPL = exp (L_{C E})

Cross-Entropy Loss & Training Objective

python

def cross_entropy(logits, targets):
    # Standard language modeling objective
    # Maximize log probability of correct next token
    log_probs = F.log_softmax(logits, dim=-1)
    return -log_probs.gather(dim=-1, index=targets.unsqueeze(-1)).mean()

Parameter Count Analysis

For a Transformer with d-dimensional embeddings, h attention heads, and n layers:

• Attention block: 4 weight matrices W_Q, W_K, W_V, W_O, each d×d → 4d² params per layer • FFN block (inner dim 4d): two matrices d→4d→d → 8d² params per layer • Total per layer: ~12d² • Embedding table: V×d (vocabulary size × hidden dim)

Validation on GPT-3: d=12288, n=96, V≈50K → 96 × 12 × 12288² + 50000 × 12288 ≈ 174.5B + 0.6B ≈ 175B ✓

Key implication: widening the model (increasing d) is quadratically expensive. Doubling d quadruples the parameter count. This is why MoE architectures — which scale total parameters without scaling active compute — became so attractive.

Per-layer params

\approx 4 d^{2} + 8 d^{2} = 12 d_{model}^{2}

Total params

\approx n \times 12 d_{model}^{2} + V \times d_{model}

Parameter Count Analysis

Regularization: Dropout & Label Smoothing

To prevent the millions of parameters from simply memorizing the training data (overfitting), the Transformer uses two key regularizers:

Dropout randomly zeroes out a percentage of neuron activations during training, forcing the network to build redundant, robust representations. Label Smoothing prevents the model from becoming overly confident. Instead of aiming for 100% probability on the correct token and 0% on others, it aims for (e.g.) 90% on the correct token and distributes the remaining 10% uniformly across the rest of the vocabulary, preventing exploding gradients.

Smoothed Target

y_{k}^{L S} = (1 - α) y_{k} + α / K

Beam Search Decoding

Greedy decoding only picks the immediate best next token. But what if a slightly suboptimal token now leads to a much better sequence later?

Beam Search maintains the top-B most likely sequences (the 'beam width') at each step. It expands all B sequences, scores them, and keeps the new top-B. It is standard in tasks requiring strict correctness like translation or summarization, whereas stochastic sampling (Top-p) is preferred for open-ended creative generation.

Full Forward Pass Walkthrough

Putting it all together for one step of generation: 1. Tokenize the input prompt into integer IDs. 2. Lookup the embedding for each ID and add Positional Encoding. 3. Pass through $N$ layers: compute Multi-Head Attention, add residual, apply LayerNorm, compute FFN, add residual, apply LayerNorm. 4. Multiply the final hidden state by the unembedding matrix to get logits over the vocabulary. 5. Apply Temperature and Softmax to get probabilities. 6. Sample a token, append to prompt, and repeat.