Walk me through the transformer encoder architecture block by block.
A transformer encoder turns token embeddings into contextual representations by repeating self-attention, residual connections, layer normalisation, and a position-wise feed-forward network. The original block is post-normalised, while many deeper implementations use pre-normalisation, so the order is an important implementation detail.
How to think about it
An encoder is a stack of layers that repeatedly lets every input token look at every other non-masked token, then transforms each token independently. The canonical layer is:
self-attention → residual add and LayerNorm → feed-forward network → residual add and LayerNorm
The sequence length and hidden width stay the same through every layer. If the input has shape L × d_model, the encoder output has that shape too.
Start with tokens and positions
The input begins as token IDs. An embedding lookup turns each ID into a learned vector. For a sequence of length L, this gives a matrix X with shape L × d_model.
The model also needs to know order. Self-attention by itself does not know that “the dog chased the cat” differs from “the cat chased the dog.” Without positional information, permuting the input tokens simply permutes the outputs in the same way. That property is called permutation equivariance: the operation respects rearrangement but cannot identify a particular position.
So the model adds a positional representation to each token embedding:
X_0 = token embeddings + positional representations
The original Transformer used sinusoidal positional encodings. Other encoder models use learned absolute positions, relative position biases, or rotary position methods. The exact position mechanism varies, but the purpose is constant: make order available to the attention calculation.
Block 1: multi-head self-attention
Take the representation entering one encoder layer, called X.
The layer creates three projections:
Q = XW_Q
K = XW_K
V = XW_V
Q means queries, K means keys, and V means values. A useful intuition is that a query describes what a token is looking for, a key describes what another token offers for matching, and a value contains the information that gets passed along. These are not three different input sequences. In self-attention, all three come from the same X, using different learned matrices.
The attention scores are computed as:
QK^T / sqrt(d_k)
This compares every query with every key. A row of the resulting matrix contains one token’s compatibility scores with all tokens. A row-wise softmax converts those scores into weights, and the weighted values are summed:
Attention(Q, K, V) = softmax(QK^T / sqrt(d_k))V
A padding mask is applied before softmax so that padding positions receive effectively zero attention. An encoder normally allows a token to attend both left and right. It is therefore bidirectional. It usually needs a padding mask, not the causal mask used by a decoder-only language model.
Why multiple heads?
Multi-head attention runs several attention calculations in parallel. Each head gets its own query, key, and value projections. One head may learn a subject–verb relationship, another may track a nearby phrase, and another may connect a pronoun to a noun. This is a useful capability, not a guarantee that every head has a neat human-readable job.
The head outputs are concatenated and passed through an output projection. The result returns to width d_model.
Here is a concrete shape walkthrough. Assume four input positions, d_model = 512, and eight heads.
| Object | Shape |
|---|---|
Input X | 4 × 512 |
Q, K, V inside one head | 4 × 64 |
| One head’s score matrix | 4 × 4 |
| One head’s output | 4 × 64 |
| Concatenated heads | 4 × 512 |
| Output projection | 4 × 512 |
For the four positions in The bank approved it, an illustrative attention row for it might be [0.10, 0.60, 0.20, 0.10]. That head would take a weighted combination of the value vectors for all four positions, placing most weight on bank. The numbers are learned and input-dependent; they are examples of the mechanism, not fixed linguistic rules.
The important boundary is this: self-attention is where tokens exchange information. Before attention, each token has only its own embedding and position. After attention, the representation for it can contain information about bank, approved, and the surrounding context.
Residual connection and layer normalisation
The attention result is not allowed to replace the incoming representation outright. It is added back to the input through a residual connection:
Y = LayerNorm(X + Dropout(SelfAttention(X)))
Dropout is used during training in the canonical architecture and disabled during inference.
The residual path gives information and gradients a shorter route through the network. Without it, every layer would have to learn a complete transformation from scratch, and deep stacks would be much harder to optimise. With it, a layer can learn a useful adjustment to the existing representation rather than reconstructing everything.
Layer normalisation, or LayerNorm, normalises the feature values for each token independently. With d_model = 512, it computes statistics over that token’s 512 features, then applies learned scale and shift parameters. It does not compute statistics across other tokens in the sequence and does not depend on the batch containing a particular number of examples.
That distinction matters. Batch normalisation can behave awkwardly with changing batch sizes or batch size one. LayerNorm works naturally when a request contains one sequence, 32 sequences, or sequences with different lengths.
Common trap: LayerNorm is not across the sequence
LayerNorm does not make all tokens share one mean and variance. Each token is normalised across its own feature dimension. The attention operation mixes tokens; LayerNorm itself does not.
Block 2: the position-wise feed-forward network
The second sub-layer is a feed-forward network, usually a small multilayer perceptron applied separately to every position:
FFN(h) = W_2 σ(W_1h + b_1) + b_2
For the original Transformer example, the dimensions are:
512 → 2048 → 512
The inner dimension is often called d_ff. The common four-times rule means d_ff is approximately four times d_model, not that this ratio is mandatory.
“Position-wise” means the same network is applied independently to every token. It does not mean every position has separate parameters. For our four-token example, the FFN processes four vectors of width 512, but it never compares the vector for bank with the vector for it.
This gives the two sub-layers different jobs:
- Self-attention mixes information between positions.
- The FFN mixes and transforms features within each position.
The original paper used ReLU for σ. BERT uses GELU, and many newer architectures use gated variants such as SwiGLU-style feed-forward layers. Those changes affect capacity and optimisation, but the attention-versus-position-wise-transformation division remains.
The FFN is also larger than many beginners expect. Ignoring biases, the standard 512 → 2048 → 512 FFN has about 2.10 million weights. The four attention projections together have about 1.05 million weights. So in this common configuration, the FFN carries roughly twice as many parameters as the attention projections.
The second residual and normalisation step is:
Z = LayerNorm(Y + Dropout(FFN(Y)))
Z is the output of one encoder layer, still shaped 4 × 512 in our example.
Stack the layers
The encoder repeats this layer N times. The original Transformer used six encoder layers. BERT-base uses 12 layers, a hidden width of 768, 12 attention heads, and an intermediate FFN width of 3072.
The layers have the same architecture and usually different parameters. They are not twelve copies sharing one set of weights.
As representations pass upward, later layers can combine and refine information produced by earlier layers. The final matrix H contains one contextualised vector per input position. For classification, a model may use a special classification token or pool the token vectors. For token labelling, it can classify every row separately. In the original encoder–decoder Transformer, the decoder uses H as its keys and values during cross-attention.
The encoder does not itself generate a sentence one token at a time. It reads the supplied sequence and produces representations for it.
The senior-level nuance: post-LN versus pre-LN
The equations above describe the original post-LayerNorm layout: the residual addition happens first, and LayerNorm follows it.
Many later deep Transformer implementations use pre-LayerNorm instead:
Y = X + SelfAttention(LayerNorm(X))
Z = Y + FFN(LayerNorm(Y))
A final normalisation layer is often added after the stack. Pre-LN commonly makes very deep models easier to optimise because the residual stream has a cleaner path through the layers. Post-LN can still work well, and the two versions are not interchangeable. If an interviewer asks about a specific implementation, state which order it uses rather than claiming that LayerNorm always comes after the sub-layer.
The other important nuance is the attention mask. Bidirectional attention is excellent when the whole input is available, as in classification, retrieval, or masked-language understanding. It is wrong for ordinary next-token prediction if target tokens are included in the input, because a token could look at information from its future. For that job, use a causal decoder or apply a causal mask, accepting that the model loses the encoder’s full-context behaviour.
Cost and a practical failure mode
Self-attention compares every query with every key, so its score matrix grows quadratically with sequence length. At length 4096, one head’s score matrix contains about 16.8 million values. In float16, that is roughly 32 MiB if materialised, before counting gradients and other activations. With 16 heads, storing all those matrices would be about 512 MiB.
FlashAttention reduces memory traffic and can avoid materialising the full matrix, but it does not remove the underlying quadratic number of pairwise comparisons. For very long documents, chunking, sparse attention, or retrieval may be a better engineering choice.
A common production failure appears when the padding mask is missing. Suppose the same four-word sentence is batched once with three padding positions and once with 124. Padding tokens then participate in attention, and positional encodings can give those positions distinct representations. The first symptom may be that the same sentence receives different logits depending on the batch’s maximum padded length. The fix is to construct and verify the padding mask, not to tune the learning rate.
What they’ll ask next
Why divide the attention scores by the square root of d_k?
If query and key components have roughly unit variance, their dot product accumulates variance as d_k grows. Large logits push softmax toward nearly one-hot weights, which produces small gradients. Dividing by sqrt(d_k) keeps the scores in a more useful range.
Why not use one large attention head?
Multiple heads provide separate learned projection spaces and can represent different relationships at the same layer. More heads are not automatically better, though: very small head dimensions can reduce the usefulness of each head, and some heads may become redundant.
Why does an encoder usually not use a causal mask?
Its usual purpose is to understand a complete input, so each token should be able to use both left and right context. A causal mask is necessary when looking at future tokens would leak the answer, such as next-token generation.
Say this in the interview
“A Transformer encoder adds positional information, then repeats multi-head bidirectional self-attention and a position-wise feed-forward network, with residual connections and LayerNorm around both sub-layers; attention mixes tokens, the FFN transforms each token’s features, and stacking the blocks produces contextual representations.”