Chapter 94 — Transformer Architecture Deep Dive

📖 Definition

A transformer is a neural network built from self-attention and feed-forward layers, with residual connections and normalisation. Every modern LLM (GPT, Claude, LLaMA, Gemini) is a decoder-only transformer trained to predict the next token.

Senior interviews expect you to walk through one forward pass, write attention from memory, and explain the design choices (RoPE, GQA, SwiGLU, RMSNorm).

🔍 One Forward Pass, Step by Step

1. Tokenise the text into token ids ("Hello world" → [15496, 995]).
2. Embed: look up a vector of size d_model for each id.
3. Repeat N transformer blocks (e.g. 32–120 layers). Each block:
   • h = x + Attention(Norm(x)) — tokens exchange information.
   • out = h + FeedForward(Norm(h)) — each token is processed independently.
4. Final norm, then the LM head (a linear layer to vocabulary size) produces one score (logit) per vocabulary token.
5. Softmax turns logits into probabilities for the next token; a sampling strategy picks one.
6. Append the token and repeat (autoregressive generation).

🔍 Self-Attention in Plain Words

Each token asks a question (Query), advertises what it contains (Key), and carries information (Value). A token's new representation is a weighted average of all Values, weighted by how well its Query matches each Key.

Q = X · Wq      K = X · Wk      V = X · Wv          (X: tokens × d_model)

scores  = Q · Kᵀ / sqrt(d_k)          tokens × tokens similarity
scores  = scores + causal_mask        −∞ above the diagonal: no peeking at future tokens
weights = softmax(scores)             each row sums to 1
output  = weights · V

Why divide by sqrt(d_k)? Dot products grow with dimension; large values push softmax into tiny gradients. Scaling keeps the variance around 1.

Why the causal mask? During training the model sees the whole sentence at once; the mask stops position t from using tokens after t, so it learns to predict the next token honestly.

💻 Code Example — Attention from Scratch (NumPy)

import numpy as np

def softmax(x, axis=-1):
    x = x - x.max(axis=axis, keepdims=True)     # numerical stability: avoid exp overflow
    e = np.exp(x)
    return e / e.sum(axis=axis, keepdims=True)

def causal_self_attention(X, Wq, Wk, Wv):
    """X: (T, d_model). Returns (T, d_k)."""
    Q, K, V = X @ Wq, X @ Wk, X @ Wv
    d_k = Q.shape[-1]
    scores = Q @ K.T / np.sqrt(d_k)             # (T, T)
    T = X.shape[0]
    mask = np.triu(np.ones((T, T), dtype=bool), k=1)
    scores = np.where(mask, -np.inf, scores)    # block the future
    return softmax(scores) @ V

rng = np.random.default_rng(0)
T, d_model, d_k = 5, 16, 8
X = rng.normal(size=(T, d_model))
out = causal_self_attention(X, *(rng.normal(size=(d_model, d_k)) for _ in range(3)))
print(out.shape)   # (5, 8)

💻 Code Example — Multi-Head Attention (PyTorch)

import torch
import torch.nn.functional as F
from torch import nn

class MultiHeadAttention(nn.Module):
    def __init__(self, d_model: int, n_heads: int):
        super().__init__()
        assert d_model % n_heads == 0
        self.h, self.d = n_heads, d_model // n_heads
        self.qkv = nn.Linear(d_model, 3 * d_model, bias=False)
        self.out = nn.Linear(d_model, d_model, bias=False)

    def forward(self, x):                                   # x: (B, T, d_model)
        B, T, C = x.shape
        q, k, v = self.qkv(x).split(C, dim=-1)
        # (B, T, C) → (B, heads, T, head_dim)
        q, k, v = (t.view(B, T, self.h, self.d).transpose(1, 2) for t in (q, k, v))
        y = F.scaled_dot_product_attention(q, k, v, is_causal=True)   # uses FlashAttention when available
        y = y.transpose(1, 2).contiguous().view(B, T, C)    # concat heads
        return self.out(y)

Why multiple heads? Each head can learn a different kind of relation (syntax, coreference, position) in its own sub-space; cost is the same as one big head.

📊 Complexity

For short sequences the FFN dominates; for long contexts attention's T² dominates — hence FlashAttention, sliding-window attention, and other long-context tricks.

🔍 Positional Information

Attention by itself ignores order ("dog bites man" = "man bites dog"). Options:

Context extension: RoPE scaling methods (position interpolation, NTK-aware, YaRN) stretch a model trained on 4k tokens to 32k+ with a little fine-tuning.

🔍 Feed-Forward Network (FFN / MLP)

FFN(x) = W2 · act(W1 · x)                     classic: hidden size ≈ 4 × d_model
SwiGLU(x) = W2 · (swish(W1 · x) ⊙ (W3 · x))   modern: gated, hidden ≈ 8/3 × d_model

The FFN holds about two-thirds of the parameters and is where much "knowledge" is thought to be stored.

🔍 MHA vs MQA vs GQA (Inference Memory)

The KV cache (Inference chapter) is often the memory bottleneck when serving, so fewer K/V heads = more concurrent users.

📊 Encoder vs Decoder Architectures

Why decoder-only won: one simple objective that scales, works for any task framed as text continuation, and is efficient to train on huge unlabelled data.

🔍 Tokenisation

Practical effects:

• ~0.75 English words per token; many other languages and code use more tokens per word → higher cost and shorter effective context.
• Numbers split oddly → arithmetic weaknesses; spelling/character tasks are hard because the model doesn't "see" letters.
• Vocabulary size trade-off: bigger vocab = shorter sequences but bigger embedding/LM-head matrices.

🔍 Mixture of Experts (MoE)

• Replace each FFN with many expert FFNs plus a small router that sends each token to the top-k (often 2) experts.
• Total parameters are huge, but active parameters per token are small → more capacity for the same compute (e.g. Mixtral 8x7B: ~47B total, ~13B active).
• Challenges: load balancing (auxiliary loss so experts get similar traffic), all experts must still fit in memory, more complex distributed training (expert parallelism).

🔍 Decoding / Sampling Strategies

🧠 What a Senior Answer Adds

• Connects architecture choices to serving cost: GQA and KV cache, MoE active params, context length and T².
• Knows tokenisation affects cost, multilingual quality and arithmetic.
• Can write attention, multi-head reshaping and the causal mask without notes.

🎯 Likely Interview Questions

1. Explain self-attention. — Queries compare with keys to make weights; output is a weighted sum of values; scaled by sqrt(d_k); causal mask for decoders.
2. Why scale by sqrt(d_k)? — Keeps dot-product variance stable so softmax doesn't saturate.
3. Why multi-head attention? — Different heads learn different relations in parallel sub-spaces.
4. What is the complexity of attention? — O(T²·d) compute, O(T²) memory naive; long contexts are expensive.
5. How do transformers know word order? — Positional encodings; modern LLMs use RoPE (relative positions via rotations).
6. Pre-LN vs post-LN? — Normalising before sub-layers is more stable for deep models.
7. Encoder-only vs decoder-only? — Bidirectional understanding (BERT) vs causal generation (GPT).
8. What is GQA and why use it? — Query heads share K/V heads → smaller KV cache, faster inference, small quality cost.
9. What is a Mixture of Experts? — Router sends tokens to a few expert FFNs: many total params, few active per token.
10. How does BPE work and why does tokenisation matter? — Merges frequent pairs; affects cost, multilingual fairness, arithmetic, spelling.
11. Temperature vs top-p? — Temperature reshapes the whole distribution; top-p truncates to the probable set.
12. Where are most parameters in a transformer? — Feed-forward layers (~2/3), then attention projections; embeddings matter for big vocabularies.

← Deep Learning & Training | Index | Next: Pretraining & Alignment →