Skip to content
datarekha

Describe the components of a transformer block and the difference between pre-norm and post-norm.

The short answer

A transformer block combines multi-head self-attention and a position-wise feed-forward network, with a residual addition and normalization around each sublayer. Post-norm normalizes after each residual addition; pre-norm normalizes the sublayer input before adding its output, usually making deep training more stable, though neither placement is universally best.

How to think about it

A transformer block combines multi-head self-attention with a position-wise feed-forward network. Each sublayer is connected to the running representation through a residual addition and a normalization layer. In post-norm, normalization happens after the residual addition; in pre-norm, it happens before the sublayer, inside the residual branch. Pre-norm usually makes very deep models easier to optimize, but it is not a universal quality guarantee.

What is inside a transformer block?

Suppose a block receives X, a tensor containing one vector for every token. Its shape is usually [batch, sequence length, d_model], where d_model is the width of each token representation. That representation is often called the hidden state or residual stream.

A standard encoder block has two main sublayers:

  1. Multi-head self-attention
  2. A position-wise feed-forward network

A decoder block in the original encoder-decoder Transformer has three: masked self-attention, cross-attention over the encoder output, and the feed-forward network. A decoder-only LLM normally has masked self-attention and the feed-forward network, repeated many times.

1. Multi-head self-attention

Self-attention lets each token build a weighted mixture of information from other tokens in the same sequence.

The block projects X into three tensors:

  • Q, the queries that ask what information each token wants
  • K, the keys that describe what each token offers
  • V, the values containing the information to pass along

The projections are learned matrix multiplications:

Q = X W_Q, K = X W_K, and V = X W_V.

For a query at position i and a key at position j, the attention score is their dot product, divided by sqrt(d_head). Here, d_head is the width of one attention head. A softmax converts the scores into weights that sum to one, and those weights are used to average the value vectors.

The scaling matters. Without it, dot products tend to grow in magnitude as the head gets wider. The softmax then becomes almost one-hot, so most gradients become tiny. The scaling keeps the scores in a useful range.

Multi-head attention runs several smaller attentions in parallel. Each head has its own projections, so different heads can learn different relationships. One may focus on nearby syntax, another on a long-range reference, and another on a delimiter pattern. That description is a useful intuition, not a promise: heads do not come with assigned human-readable jobs.

The attention output is concatenated across heads and projected back to width d_model. A causal mask is applied in a decoder-only model so that token i cannot use information from a future token. The mask must affect the scores before softmax. Otherwise, future positions can receive probability mass.

2. The residual connection

A residual connection, also called a skip connection, adds a sublayer’s output to its input:

output = x + F(x).

This gives the network an identity path. A layer can learn a useful transformation, but it can also learn to make a small correction to the existing representation. More importantly for optimization, gradients can travel through the addition without being forced entirely through the learned transformation.

The residual stream is not a second copy of the input that gets thrown away later. It is the main running state of the model. Attention writes context into it. The feed-forward network then transforms each token’s features.

3. The position-wise feed-forward network

The feed-forward network is a small multilayer perceptron applied independently to every token position, using the same weights at every position. A typical version expands the vector, applies a nonlinear activation, and projects it back:

d_model → d_ff → d_model.

The original Transformer used a ReLU activation. Modern models commonly use GELU or a gated variant such as SwiGLU, where one part of the expanded representation controls another part.

“Position-wise” is a common source of confusion. It does not mean the token has no context. Attention has already mixed information from other positions. It means that once the attention result exists, the same feed-forward computation is applied separately to each token vector.

The two sublayers are therefore doing different jobs:

  • Attention mixes information across positions.
  • The feed-forward network performs richer feature transformation at each position.

Normalization keeps activation scales under control. LayerNorm normalizes each token across its feature dimensions, using that token’s mean and variance, followed by learned scale and offset parameters. It does not normalize across the batch or mix information between tokens. Many modern LLMs use RMSNorm instead, which scales by root mean square and does not subtract the mean. Pre-norm and post-norm describe placement, not whether the model uses LayerNorm or RMSNorm.

Pre-norm versus post-norm

The difference is easiest to see in pseudocode. Let A be attention, F be the feed-forward network, and N be LayerNorm or RMSNorm.

Post-norm places normalization after each residual addition:

h = N(x + A(x))
y = N(h + F(h))

Pre-norm normalizes the input to each sublayer, then adds the unnormalized residual stream back:

h = x + A(N(x))
y = h + F(N(h))

The names are literal. Post-norm normalizes the post-residual result. Pre-norm normalizes the pre-sublayer input.

That small rearrangement changes the gradient path. In pre-norm, every block retains a direct identity route from its input to its output through the addition. The gradient can pass along that route even if the attention or feed-forward branch is poorly behaved. In post-norm, the gradient must pass through the normalization after the addition, so the normalization’s derivative sits between successive residual states.

This is why pre-norm is generally easier to train as depth increases. It tends to reduce sensitivity to initialization and learning-rate choices, and it often needs less delicate warmup tuning. Modern deep decoder-only LLMs therefore commonly use a pre-norm arrangement, frequently with RMSNorm and a final normalization after the last block.

A concrete shape calculation

Take a model with batch size 2, sequence length 512, d_model = 768, 12 attention heads, and d_head = 64.

The input has shape [2, 512, 768]. The projected query, key, and value tensors initially have the same shape. After splitting into heads, each has shape [2, 12, 512, 64].

The attention score tensor has shape [2, 12, 512, 512]. That is:

2 × 12 × 512 × 512 = 6,291,456 scores.

If stored as two-byte floating-point values, that single tensor occupies about 12.6 MB. A naive implementation may also hold attention probabilities and intermediate buffers, which is why attention memory grows quickly. At sequence length 1,024, the score count becomes 25,165,824, exactly four times larger. The sequence doubled; the pairwise attention work quadrupled.

After attention, the result returns to shape [2, 512, 768] and enters the residual addition. An ordinary feed-forward network might expand every one of the 1,024 token vectors from 768 features to 3,072, apply its activation, and project back to 768.

The senior-level nuance

Pre-norm is not “better normalization.” It is a different placement of the same general control mechanism.

Pre-norm usually improves optimization stability, but its residual stream is not normalized after every addition. Across many layers, activation magnitudes can therefore grow unless the architecture, initialization, residual scaling, and final normalization are chosen carefully. Pre-norm can also produce different optimization and representation behavior from post-norm, so “more stable to train” does not mean “always higher final accuracy.”

Post-norm gives each sublayer output a normalized scale, which can be attractive. It was used in the original Transformer and remains viable when depth, initialization, learning rate, warmup, and residual scaling are tuned together. Some very deep architectures use explicit scaling or other normalization designs to recover stable post-norm-like behavior.

A practical failure mode appears when someone copies hyperparameters from a pre-norm model into a deep post-norm model. The first symptom may be sharply spiking gradient norms, a training loss that becomes NaN, or a loss curve that fails during the first few thousand updates. The fix is not automatically “move the norm.” Check the exact block equations, learning-rate schedule, initialization, residual scaling, and whether the checkpoint was trained with the same architecture. A post-norm checkpoint is not interchangeable with a pre-norm checkpoint merely because the layer names look similar.

What they’ll ask next

Why divide attention scores by sqrt(d_head)?
If query and key components have roughly unit variance, their dot product accumulates variance proportional to d_head. Larger heads therefore create larger logits. Softmax saturates on those logits, concentrating almost all weight on one position and weakening gradients. Dividing by sqrt(d_head) keeps the distribution trainable.

Is pre-norm always the right choice?
No. It is usually the safer default for a deep Transformer, but final quality depends on the whole recipe. Post-norm can work well with suitable initialization, residual scaling, warmup, and depth. The choice should be validated with training stability and downstream quality, not selected by slogan.

How does a decoder block differ from an encoder block?
A decoder-only block uses causal self-attention, so a token cannot see future tokens. An encoder-decoder decoder also has cross-attention: its queries come from the decoder state, while its keys and values come from the encoder output. The feed-forward sublayer then transforms the result.

Say this in the interview

“A Transformer block alternates multi-head self-attention, which mixes information across tokens, with a position-wise feed-forward network, which transforms each token’s features; residual connections and normalization surround those sublayers, with post-norm applying normalization after the add and pre-norm applying it before the sublayer, usually giving more stable gradients in deep models.”

Learn it properly Inside the transformer block

Keep practising

All Deep Learning questions

Explore further