Can you walk me through a transformer?
I am interviewing at ML companies this fall, and there is one question that shows up in every loop in some form: “walk me through a transformer.” It sounds like a warm-up. It is not. It is the interviewer checking whether you can hold the entire machine in your head at once, and whether you know why each piece is there rather than just that it is there.
Let me be honest about the setting. In an interview, with a whiteboard and a clock, it is the nuances that slip: where exactly the nonlinearity lives, why the FFN is four times wider than the model, whether attention is really the expensive part.
Every one of those has a crisp answer. And as a PhD student, or any serious ML engineer, I hold myself to a simple standard: these answers should be on my fingertips, not reconstructed under pressure. The transformer is the machine everything else is built on. Knowing it cold is the baseline, not a bonus.
There is a second reason to write this down, and it is worth being candid about. It is September 2026. Coding agents are not a forecast, they are how most of us work: I hand Claude Code a task, it plans the steps, writes the attention block, runs the tests, fixes what fails, and opens the pull request while I read something else.
Agents run for hours on their own now. That is a genuine gift, and I am not going to pretend I do not lean on it every day. It also makes it very easy to let the internal workings of things fade, because nothing in that loop ever asks you why the square root is there.
The basics are not something to outsource. They are exactly what lets you read what the agent wrote and know whether it is right, and they are what an interviewer is really probing when the warm-up question arrives.
So this series is the refresher I want on hand, written the way I would want to answer. Six parts, each built around the questions interviewers actually ask, each answered from first principles rather than memorized. This part is the map: one pass through the block, every shape accounted for.
Parts 2 through 5 interrogate the components one at a time: attention and its square root, residuals and normalization, positional encoding, the feed-forward network. Part 6 is the diff between the 2017 paper and the blocks people ship today.
The 30-second version
Tokens become rows of a $(T, d)$ matrix through an embedding lookup, with position information added. That matrix is the residual stream. Each of $L$ identical blocks reads it, computes an update, and adds the update back: first an attention sublayer, where every row becomes a weighted mix of other rows, then a feed-forward sublayer, where every row is transformed on its own. Each sublayer is preceded by a normalization. After the last block, a final norm and an unembedding matrix turn each row into logits over the vocabulary. Shape in, shape out: $(T, d)$ the whole way through.What is the one picture that makes the whole block easy?
The mental model that finally made the whole thing hold together is this: a transformer is not a pipeline of layers that transform their input into something new. It is a shared bus, and layers are devices attached to it.
The bus is a matrix $H \in \mathbb{R}^{T \times d}$: one row per token, $d$ numbers per row. This is the residual streamResidual streamThe running $(T, d)$ state that every block reads from and adds to.
It is called a stream because it flows unchanged from the embedding to the output, with each block contributing an additive update:</span>\(H_{l+1} = H_l + f_l(H_l)\)Unroll it and the final state is the embedding plus the sum of every block’s contribution. Nothing is ever overwritten.</span></span>. Every block does exactly one thing to it:
\[H \leftarrow H + \text{Sublayer}(\text{Norm}(H))\]Read the stream, normalize what you read, compute something, add it back. The stream is never overwritten, only nudged. A block with all-zero weights is a no-op, and the network still works with that block removed. This is the property Part 3 is about, and it is why a hundred of these can be stacked.
A transformer block has two such sublayers, and they differ in one crucial way. Attention mixes across rows: token 4 gets to read tokens 1 through 3. The feed-forward network mixes across columns: each row is transformed by itself, and no row sees any other. The block alternates these two, and the whole architecture is that alternation repeated $L$ times.
What goes into the block?
A tokenizer has already turned the text into integer ids, $T$ of them. An embedding matrix $E \in \mathbb{R}^{V \times d}$ has one learned row per vocabulary entry, and the input is just a lookup:
\[H_0 = E[\text{ids}] + P \qquad H_0 \in \mathbb{R}^{T \times d}\]The $P$ term is position information. Attention on its own has no idea which token came first, a fact that is surprising the first time you prove it and obvious ever after: Part 4 does the proof. The original paper adds a fixed sinusoidal table here, once, at the input. Modern models mostly do not add anything here at all and instead inject position inside every attention layer, which is also Part 4’s story.
What does the attention sublayer do to the stream?
The stream is normalized row by row (Part 3 covers why row-wise and not batch-wise), and three linear maps produce queries, keys, and values:
\[Q = \tilde{H} W_Q, \qquad K = \tilde{H} W_K, \qquad V = \tilde{H} W_V\]with $\tilde{H} = \text{Norm}(H)$ and each $W \in \mathbb{R}^{d \times d}$. Every token now has a query (what am I looking for), a key (what do I contain), and a value (what I will hand over if someone looks at me). The score between token $i$ and token $j$ is the dot product of $i$’s query with $j$’s key, scaled, and turned into a distribution over $j$ with a softmax:
\[A = \text{softmax}\!\left(\frac{QK^\top}{\sqrt{d_k}}\right) \qquad A \in \mathbb{R}^{T \times T}\]Row $i$ of $A$ is a probability distribution over which tokens $i$ should read from. The output for token $i$ is the average of the values under that distribution:
\[\text{Attn}(\tilde{H}) = A V \qquad \in \mathbb{R}^{T \times d}\]That is the entire mechanism. A soft dictionary lookup: each query retrieves a weighted blend of values, with weights from key similarity. Why the dot product, why the softmax, and why that square root are the whole of Part 2, because the square root question in particular is asked in almost every interview and most people give half the answer.
In practice attention is multi-headed. Instead of one $d$-dimensional lookup, split $d$ into $h$ heads of width $d_k = d / h$, run the lookup $h$ times in parallel with separate projections, and concatenate:
\[\text{MultiHead}(\tilde{H}) = \text{Concat}(\text{head}_1, \ldots, \text{head}_h)\, W_O\]The concatenation is $(T, d)$ again, and $W_O \in \mathbb{R}^{d \times d}$ mixes the heads. Nothing about the shape changed; the model just gets $h$ independent lookups per token instead of one.
For a language model there is one more ingredient: a causal mask. Before the softmax, every score where $j > i$ is set to $-\infty$, so the softmax assigns it exactly zero weight and token $i$ can only read from itself and its past. Without the mask (BERT, the original encoder) every token reads every other.
Then the add:
\[H \leftarrow H + \text{MultiHead}(\text{Norm}(H))\]What does the feed-forward sublayer do?
Normalize the stream again, and apply the same two-layer MLP to every row independently:
\[\text{FFN}(x) = W_2\, \sigma(W_1 x + b_1) + b_2 \qquad W_1 \in \mathbb{R}^{d_{ff} \times d},\; W_2 \in \mathbb{R}^{d \times d_{ff}}\]The hidden width is $d_{ff} = 4d$ in the original paper and in most models since. The nonlinearity $\sigma$ was ReLU, then GELU, and today it is usually a gated variant called SwiGLU. Because the FFN acts on each row with no reference to any other row, it is sometimes called position-wise, and it is exactly a $1 \times 1$ convolution along the sequence.
\[H \leftarrow H + \text{FFN}(\text{Norm}(H))\]That is one block. Two thirds of its parameters are in the FFN, and if you ask what the FFN is for when attention already exists, the answer is more interesting than “it adds capacity.” That is Part 5.
What comes out at the top?
After $L$ blocks, a final normalization, and an unembedding matrix $W_U \in \mathbb{R}^{d \times V}$:
\[\text{logits} = \text{Norm}(H_L)\, W_U \qquad \in \mathbb{R}^{T \times V}\]Row $i$ is a score for every vocabulary entry, and softmax over it gives the model’s prediction for the token after position $i$. In training, all $T$ rows are used at once: the loss is the cross-entropy of each row against the actual next token, and the causal mask is what makes this legal, since row $i$ was computed without seeing token $i+1$.
That is why one forward pass gives $T$ training signals for the price of one, and why the decoder-only design won.
At inference you only need the last row to pick the next token. The rest of the rows are not wasted: their keys and values are cached so that the next step does not recompute them. The KV cache is the reason inference-time architecture choices look different from training-time ones, which is a Part 6 topic.
How many parameters and FLOPs, and where do they sit?
Interviewers like this one because it separates people who have looked at a config file from people who have derived it.
The 30-second version
Attention has four $d \times d$ matrices ($W_Q, W_K, W_V, W_O$): $4d^2$. The FFN has $d \times 4d$ and $4d \times d$: $8d^2$. So a block is about $12d^2$, two thirds of it in the FFN. Norm parameters are $O(d)$ and negligible. Add the embedding matrix $V \times d$ (and the unembedding if untied) on top. GPT-2 small has $d = 768$ and 12 layers: $12 \times 12 \times 768^2 \approx 85$M in blocks, plus $50257 \times 768 \approx 39$M in embeddings, which is the 124M everyone quotes.Per block, ignoring biases and norms:
| Component | Matrices | Parameters |
|---|---|---|
| Attention | $W_Q, W_K, W_V, W_O$, each $d \times d$ | $4d^2$ |
| FFN | $W_1$: $d \times 4d$, $W_2$: $4d \times d$ | $8d^2$ |
| Block total | $12d^2$ |
The compute question has a subtlety that most people miss, and interviewers know it.
The 30-second version
Two kinds of cost per layer. The weight matmuls cost $O(T d^2)$: every token goes through $12d^2$ parameters. The attention scores cost $O(T^2 d)$: every token dots against every other. Their ratio is $T / d$ up to a constant, so attention FLOPs only dominate once the context is several times longer than the model width. For a 4096-wide model that is tens of thousands of tokens. Below that, the linear layers dominate compute, and what attention actually strains is memory, because the $T \times T$ score matrix per head has to live somewhere. That is the problem FlashAttention solves.The clean way to remember it: forward FLOPs per token are about $2N$ for a model with $N$ non-embedding parameters (one multiply and one add per weight), plus an attention term of about $4Td$ per layer that grows with context. The training rule of thumb, $6N$ FLOPs per token (forward plus a backward that costs twice the forward), comes straight from this.
Where is the nonlinearity?
The 30-second version
Two places, and they are different in kind. The softmax in attention makes the mixing weights a nonlinear function of the input, but the values are still combined linearly. The activation inside the FFN is the only elementwise nonlinearity applied to the token's own representation. Everything else (projections, residual adds, unembedding) is linear, and the norms are a mild nonlinearity of scale. Take the FFN out and the model can still route information, but it can barely transform it.I put this question here because it is the hinge between the parts of this series. Attention decides where to look. The FFN decides what to do with what it found. Both have a nonlinearity but in different places, and knowing which is which is the difference between describing the architecture and understanding it.
Encoder, decoder, or both: which stack does a modern model keep?
The 2017 paper had two stacks. An encoder read the source sentence with unmasked attention, and a decoder generated the target with masked self-attention plus a third sublayer, cross-attention, whose queries came from the decoder and whose keys and values came from the encoder’s output. That is the right shape for translation, where the input is fully known before you write a word.
BERT kept only the encoder: bidirectional attention, trained to fill in masked tokens, good for classification and retrieval where you want a representation of a whole input. GPT kept only the decoder: causal attention, trained to predict the next token, and it turned out that this one objective on enough text subsumes almost everything else.
Nearly every model you would be asked about in 2026 is decoder-only, and the encoder’s job has been absorbed by the prefix: the “source” is just the beginning of the same sequence.
Can you write the block in code?
The last thing an interviewer may do is hand you a laptop. This is the block in the fewest lines I trust, with the pre-norm placement that modern models use:
import torch, torch.nn as nn, torch.nn.functional as F
class Block(nn.Module):
def __init__(self, d, n_heads, d_ff=None):
super().__init__()
self.h, self.dk = n_heads, d // n_heads
self.norm1, self.norm2 = nn.LayerNorm(d), nn.LayerNorm(d)
self.qkv = nn.Linear(d, 3 * d, bias=False) # W_Q, W_K, W_V stacked
self.out = nn.Linear(d, d, bias=False) # W_O
d_ff = d_ff or 4 * d
self.ffn = nn.Sequential(nn.Linear(d, d_ff), nn.GELU(), nn.Linear(d_ff, d))
def attention(self, x):
B, T, d = x.shape
q, k, v = self.qkv(x).split(d, dim=-1) # each (B, T, d)
q, k, v = (t.view(B, T, self.h, self.dk).transpose(1, 2) # (B, h, T, dk)
for t in (q, k, v))
scores = q @ k.transpose(-2, -1) / self.dk ** 0.5 # (B, h, T, T)
mask = torch.triu(torch.ones(T, T, dtype=torch.bool, device=x.device), 1)
scores = scores.masked_fill(mask, float("-inf")) # causal
y = F.softmax(scores, dim=-1) @ v # (B, h, T, dk)
return self.out(y.transpose(1, 2).reshape(B, T, d))
def forward(self, x): # x: (B, T, d), the residual stream
x = x + self.attention(self.norm1(x)) # tokens talk to each other
x = x + self.ffn(self.norm2(x)) # each token thinks on its own
return x
Every shape comment in there is something I have been asked to state out loud. The transposes exist only so that the matmul broadcasts over heads; the split, view, and reshape are the multi-head bookkeeping and nothing more.
Wrap it in an embedding, $L$ of these, a final norm, and a linear layer to $V$, and you have GPT-2. What you do not yet have is the reason any of it works, and that is what the next four parts are for.
Rapid fire: can you do these from memory?
Close the tab and answer these out loud. If one of them takes more than a minute, that is the part to reread.
- What is the shape of the residual stream, and does it ever change inside the network?
- Which sublayer moves information between tokens, and which one does not?
- Write the attention equation and name the shape of every intermediate.
- Why does one forward pass over a sequence give $T$ training signals rather than one?
- Count the parameters in a block in terms of $d$. What fraction is FFN?
- When does the $O(T^2 d)$ attention term dominate the $O(T d^2)$ term?
- What is cross-attention, and why do decoder-only models not need it?
- What is cached between decoding steps, and why is that possible?
Part 2 starts with the question I have been asked most: why divide by $\sqrt{d_k}$, and what happens if you do not.