Transformers

Why Transformers exist (motivation)

Before Transformers, sequence modeling was dominated by RNNs/LSTMs/GRUs and CNN-based sequence models. Those families have two persistent pain points:

1) Long-range dependencies are hard. Even with gating, recurrent models struggle to move information across hundreds or thousands of steps.

2) Parallelism is limited. Recurrence is inherently sequential: to compute step t, you need step t−1. That slows training.

Transformers address both by making the core operation all-to-all token interaction in a single layer: every token can “look at” every other token (subject to masking). This interaction is differentiable, learnable, and highly parallelizable on GPUs/TPUs.

The Transformer idea in one sentence

A Transformer layer repeatedly does:

• •Mix information across tokens using self-attention (content-based weighted averaging).
• •Transform each token independently using a position-wise feed-forward network (an MLP applied to each position).

Crucially, both steps are wrapped in residual connections and layer normalization for stable optimization.

What a token representation looks like

Assume a sequence length L and model width d_model.

• •Input embeddings form a matrix X ∈ ℝ^{L×d_model}
• •Row i, written xᵢ, is the vector for token i.

Because attention has no inherent sense of order, we add positional information:

• •X₀ = E + P, where E are token embeddings and P are positional encodings (sinusoidal or learned).

The two canonical Transformer families

This lesson focuses on the core mechanics that all of these share: scaled dot-product attention, multi-head attention, and the Transformer layer structure.

Why attention is “routing”

Suppose token i needs information from token j (e.g., a pronoun needs its antecedent). Attention lets token i compute a weighted mixture of other tokens’ information.

The key design choice is: weights depend on content (learned similarity), not just distance.

Queries, keys, values (the roles)

Each token representation xᵢ is linearly projected into three vectors:

• •qᵢ (query): what token i is looking for
• •kᵢ (key): what token i offers / how it can be matched
• •vᵢ (value): the information token i will contribute if attended to

In matrix form for a whole sequence X ∈ ℝ^{L×d_model}:

• •Q = XW_Q
• •K = XW_K
• •V = XW_V

where W_Q, W_K, W_V are learned matrices.

Typically:

• •W_Q ∈ ℝ^{d_model×d_k}
• •W_K ∈ ℝ^{d_model×d_k}
• •W_V ∈ ℝ^{d_model×d_v}

In most standard implementations, d_v = d_k = d_model / h per head.

Attention scores and why we scale by √d_k

The raw similarity between token i and token j is the dot product:

• •score(i, j) = qᵢ · kⱼ

In matrix form, the score matrix is:

• •S = QKᵀ where S ∈ ℝ^{L×L}

If q and k components have roughly unit variance, then the dot product grows with dimension:

• •Var(qᵢ · kⱼ) ≈ d_k

Large magnitudes push softmax into saturation, giving tiny gradients. To keep logits in a reasonable range, we scale:

• •Ŝ = (QKᵀ) / √d_k

This is “scaled dot-product attention.”

Softmax to turn scores into weights

For each query position i, we take a softmax over j:

• •A = softmax(Ŝ)

So Aᵢⱼ ≥ 0 and ∑ⱼ Aᵢⱼ = 1.

Interpretation: row i of A is a probability distribution over which tokens i will attend to.

Masking (padding and causal)

Masks are incorporated by adding −∞ (or a very negative number) to disallowed positions before softmax.

Let M ∈ ℝ^{L×L} where:

• •Mᵢⱼ = 0 if attention from i to j is allowed
• •Mᵢⱼ = −∞ if disallowed

Then:

• •A = softmax( (QKᵀ)/√d_k + M )

Two common masks:

1) Padding mask: disallow attending to padding tokens.

2) Causal mask (decoder): disallow attending to future tokens, i.e., j > i.

Weighted sum of values

Finally, the output at each position i is a weighted sum of values:

• •oᵢ = ∑ⱼ Aᵢⱼ vⱼ

In matrix form:

• •O = AV

Where O ∈ ℝ^{L×d_v}.

Putting it together (single-head self-attention)

The full formula:

• •Attention(X) = softmax( (XW_Q)(XW_K)ᵀ / √d_k + M ) (XW_V)

A geometric intuition

Dot products measure alignment. If qᵢ points in a similar direction to kⱼ, token i will attend more to token j. But because W_Q and W_K are learned, the model can invent similarity notions that match the task (syntax, coreference, topic, etc.).

Complexity note (why long context is expensive)

Self-attention builds an L×L matrix of scores. That’s:

• •Time: O(L² d_k)
• •Memory: O(L²)

This quadratic dependence is why long-context Transformers require approximations or architectural tricks—but the core mechanism remains the same.

Why one attention “view” isn’t enough

A single attention map must decide one set of weights per token. But language often requires multiple simultaneous relationships:

• •One head might track coreference (pronouns → nouns).
• •Another might track syntactic dependencies (verbs → objects).
• •Another might focus on local context (previous few tokens).

Multi-head attention lets the model compute multiple attention distributions in parallel, each in its own learned subspace.

The shape story: d_model, h, d_k

Let:

• •d_model = model width
• •h = number of heads
• •d_k = per-head query/key width

Commonly:

• •d_k = d_model / h

Example: d_model = 768, h = 12 ⇒ d_k = 64.

Per-head projections

For head r ∈ {1,…,h}, we have separate projection matrices:

• •W_Q^{(r)} ∈ ℝ^{d_model×d_k}
• •W_K^{(r)} ∈ ℝ^{d_model×d_k}
• •W_V^{(r)} ∈ ℝ^{d_model×d_v}

Compute:

• •Q^{(r)} = XW_Q^{(r)}
• •K^{(r)} = XW_K^{(r)}
• •V^{(r)} = XW_V^{(r)}

Then each head output:

• •O^{(r)} = softmax( Q^{(r)} (K^{(r)})ᵀ / √d_k + M ) V^{(r)}

Concatenate and mix

Concatenate head outputs along the feature dimension:

• •O_concat = [O^{(1)} | O^{(2)} | … | O^{(h)}]

If each O^{(r)} ∈ ℝ^{L×d_v} and d_v = d_k, then O_concat ∈ ℝ^{L×(h d_k)} = ℝ^{L×d_model}.

Then apply a final learned output projection:

• •Y = O_concat W_O where W_O ∈ ℝ^{d_model×d_model}

This final mixing matters: it lets the model combine information across heads.

Self-attention vs cross-attention inside MHA

Multi-head attention is a pattern, and it can be used in different places:

• •Self-attention: Q, K, V come from the same X.
• •Cross-attention: Q comes from decoder states X_dec; K and V come from encoder states X_enc.

Cross-attention formula:

• •Attention(X_dec, X_enc) = softmax( (X_dec W_Q)(X_enc W_K)ᵀ / √d_k + M ) (X_enc W_V)

Here the decoder learns to retrieve information from the encoded source sequence.

Practical interpretation of heads

A helpful mental model:

• •Each head defines its own “matching function” (via W_Q and W_K)
• •and its own “payload space” (via W_V)

So attention isn’t only where to look; it’s also what to bring back.

A note on head dimension choice

Holding d_model fixed, increasing h decreases d_k. There is a trade-off:

Empirically, standard settings (8–32 heads depending on width) work well, but variants exist (multi-query, grouped-query attention) to reduce memory/compute during decoding.

Why the Transformer is a stack of simple blocks

A single attention layer can mix tokens once, but deep language understanding requires multiple rounds of:

• •gathering context
• •transforming it
• •gathering again at a higher abstraction

So Transformers stack identical blocks. The stability of deep stacking relies on residual connections and layer normalization.

The canonical encoder layer

An encoder layer has two sublayers:

1) Multi-head self-attention (MHA)

2) Position-wise feed-forward network (FFN)

Each sublayer is wrapped by residual + layer norm. There are two common normalization conventions:

We’ll write pre-LN, since it’s widely used.

Pre-LN encoder layer equations

Let X ∈ ℝ^{L×d_model} be the layer input.

(1) Attention sublayer

• •U = LN(X)
• •A = MHA(U, U, U; M_padding)
• •X′ = X + Dropout(A)

(2) Feed-forward sublayer

• •V = LN(X′)
• •F = FFN(V)
• •Y = X′ + Dropout(F)

Output is Y.

What “position-wise FFN” means

The FFN is the same MLP applied independently to each position.

Typical form:

• •FFN(x) = W₂ σ(W₁ x + b₁) + b₂

Where:

• •W₁ ∈ ℝ^{d_model×d_ff}
• •W₂ ∈ ℝ^{d_ff×d_model}
• •σ is a nonlinearity (ReLU, GELU, SwiGLU variants)
• •d_ff is often 4× d_model (e.g., 3072 for 768)

Even though FFN doesn’t mix tokens, it adds substantial capacity: it can reshape and recombine features within each token vector.

The decoder layer

A decoder layer adds one more attention sublayer:

1) Masked multi-head self-attention (causal)

2) Cross-attention over encoder outputs (optional in encoder–decoder)

3) Feed-forward network

Pre-LN decoder (encoder–decoder) sketch:

• •X₁ = X + MHA(LN(X), LN(X), LN(X); M_causal)
• •X₂ = X₁ + MHA(LN(X₁), K_enc, V_enc; M_enc_padding)
• •Y = X₂ + FFN(LN(X₂))

If it’s decoder-only (GPT-style), you omit the cross-attention term.

Why residual connections matter (conceptual)

Residuals let the model learn modifications rather than complete rewrites.

If a sublayer initially does something unhelpful, the residual path preserves the input:

• •X_out = X_in + ε

This makes gradients flow more directly through many layers. In deep Transformers (dozens to hundreds of layers), residual pathways are essential.

Why layer normalization is placed around sublayers

Layer norm stabilizes the scale of activations, making training less sensitive to initialization and learning rate.

Layer norm operates per token vector xᵢ:

• •μ = (1/d_model) ∑ₖ xᵢₖ
• •σ² = (1/d_model) ∑ₖ (xᵢₖ − μ)²
• •LN(xᵢ) = γ ⊙ (xᵢ − μ)/√(σ² + ϵ) + β

(You already know LN; here it matters because attention logits and FFN activations can drift in magnitude as depth increases.)

Where positional encoding enters

Positional information is usually added once at the bottom:

• •X₀ = Emb(tokens) + PosEnc(positions)

Then all layers operate on X₀, X₁, …, X_N.

However, there are alternatives (relative position bias, rotary embeddings), but the principle remains: attention needs a way to distinguish positions.

Encoder-only Transformers (bidirectional)

Use case: classification, retrieval, tagging, masked language modeling.

Mechanics:

• •Full self-attention over the entire sequence (no causal mask).
• •Output is a sequence of contextual embeddings.

To classify, you might pool:

• •use a special [CLS] token representation y_cls
• •or mean-pool: y_pool = (1/L) ∑ᵢ yᵢ

Decoder-only Transformers (autoregressive)

Use case: text generation, code generation, next-token prediction.

Mechanics:

• •Causal mask enforces that token i can only attend to tokens ≤ i.

Training objective (next-token):

• •Given tokens t₁…t_L, model predicts t_{i+1} from context t₁…t_i.

If logits at position i are zᵢ ∈ ℝ^{|Vocab|}, then:

• •p(t_{i+1} | t_{≤i}) = softmax(zᵢ)

Loss is cross-entropy summed across positions.

Encoder–decoder Transformers (seq2seq)

Use case: translation, summarization, speech-to-text.

Mechanics:

• •Encoder produces memory H_enc.
• •Decoder uses masked self-attention for partial output and cross-attention to read from H_enc.

Practical issues that shape real implementations

1) KV caching during decoding

Autoregressive decoding generates one token at a time. Recomputing attention over the whole prefix is expensive. Instead, store previous keys/values.

At time step t:

• •compute K_t, V_t for the new token
• •append to cache
• •attend with Q_t against all cached K_{≤t}

This reduces per-step cost from O(t²) to roughly O(t) for attention score computation (still linear in context length per step).

2) Attention masks are not optional

If you forget masking:

• •the model may attend to padding and learn spurious correlations
• •a decoder may “cheat” by attending to future tokens during training

Mechanically, masks must be added before softmax.

3) Initialization and normalization choices

Deep Transformers are sensitive to:

• •pre-LN vs post-LN
• •residual scaling
• •dropout placement

These details often decide whether training is stable.

4) Why positional encoding is central

Self-attention is permutation-invariant without positional signals:

If you permute tokens, QKᵀ permutes correspondingly, producing the same pattern up to permutation. Positional encoding breaks this symmetry, enabling order-sensitive tasks.

Connecting the mechanics to behavior

When you inspect trained models, attention heads often develop recognizable patterns:

• •local heads: attend strongly to nearby tokens
• •delimiter heads: attend to separators (commas, periods)
• •induction heads: copy patterns like “A … A”

Not every head is interpretable, and attention weights are not the whole story (FFN and residual streams matter), but the routing intuition remains useful.

Summary of the full forward pass (encoder-only)

Given tokens → embeddings + positions:

• •X₀ = Emb(t) + PosEnc(pos)

For ℓ = 1…N layers:

• •X_ℓ = TransformerEncoderLayer(X_{ℓ−1})

Final outputs X_N feed task heads (classification, MLM head, etc.).

For decoder-only, the same structure holds, but with causal masking and an output projection to vocabulary logits.