Inside Transformer Blocks

Suppose you have a working attention layer. It lets every token gather information from every other token. That is powerful, so the obvious move is to stack forty of them and see what happens.

What happens is that the model does not train. The loss drops for a few hundred steps, then plateaus at a value barely better than predicting the most common token. Sometimes it produces NaN in the third epoch and dies outright. Meanwhile a four-layer version of the same model trains fine.

Two separate things are going wrong, and both have to be fixed before depth becomes an asset rather than a liability. On top of that, attention on its own has a structural gap: it moves information between positions but does almost no computation within a position. A transformer block is the small set of components that solves all three problems. Understanding it means understanding why each piece is there — and what breaks when it is not.

Problem one: gradients cannot survive forty layers

Backpropagation computes a gradient at the output and pushes it backwards through every layer by repeated multiplication. For a stack of $L$ layers, the gradient reaching layer 1 is a product of $L$ Jacobian matrices.

Products of many numbers behave badly. If each layer's Jacobian shrinks the gradient by a factor of 0.9, then after 40 layers the surviving fraction is $0.9^{40} = 0.0148$ — the early layers receive about 1.5% of the signal. If each layer amplifies by 1.1 instead, $1.1^{40} = 45.3$, and the gradients explode. The window between "vanishing" and "exploding" is razor-thin, and nothing keeps a stack of learned matrices inside it.

The fix is a residual connection: instead of replacing the input with the layer's output, add the output to the input.

Without residual:   x -> [layer] -> y          y = F(x)With residual:      x -> [layer] -> (+) -> y   y = x + F(x)                    |______________|

The derivative of $y = x + F(x)$ with respect to $x$ is $1 + F'(x)$. That leading $1$ is everything. Even if $F'(x)$ is close to zero, the gradient still passes through at full strength via the identity path. Forty layers of $(1 + \text{something small})$ multiply to roughly 1, not to 0.0148.

The residual connection turns a deep network from a chain, where every link can break the signal, into a highway with optional exits. Each layer proposes an adjustment; it does not get to overwrite everything upstream.

This reframing is worth carrying: modern practice describes the vector running through the network as the residual stream. Every attention layer and every feed-forward layer reads from the stream, computes something, and writes its result back by addition. Nothing is ever deleted, only added to. A model with 32 layers is 64 sublayers all reading and writing one shared bus.

Problem two: the scale of activations drifts

Even with residual connections, values change magnitude as they pass through layers. If layer 5 produces activations around ±3 and layer 6 produces ±30, the layers that follow face a moving target, learning rates that suited one scale fail at another, and in low precision the large values overflow.

Layer normalisation fixes the scale at each position independently. For a vector $\mathbf{x}$ of length $d$:

Worked on a 4-dimensional vector:

x = [2, 4, 6, 8]mean:      (2+4+6+8)/4 = 5variance:  ((2-5)^2 + (4-5)^2 + (6-5)^2 + (8-5)^2)/4         = (9 + 1 + 1 + 9)/4 = 5std:       sqrt(5) = 2.236normalised: [(2-5)/2.236, (4-5)/2.236, (6-5)/2.236, (8-5)/2.236]          = [-1.342, -0.447, 0.447, 1.342]        mean 0, variance 1

The learned parameters $\gamma$ (scale) and $\beta$ (shift) let the model undo the normalisation if it is genuinely unhelpful for some dimension, so nothing expressive is lost.

Note the axis carefully, because this is the most common confusion. Layer norm computes statistics across the features of a single token. Batch norm computes them across the batch for a single feature. The difference is decisive for language models: batch statistics depend on which other examples happen to be in the batch, which makes them meaningless at inference time when the batch is one sequence, and worse with variable-length padded sequences. Layer norm has no such dependency — each token normalises itself, batch size 1 behaves identically to batch size 512.

RMSNorm: the same job, two-thirds of the work

Most modern large models use a simplification. It turns out that subtracting the mean contributes little; what matters is controlling the magnitude. RMSNorm keeps only the rescaling:

Text
Copy
x = [2, 4, 6, 8]mean of squares: (4 + 16 + 36 + 64)/4 = 30RMS:             sqrt(30) = 5.477output:          [0.365, 0.730, 1.095, 1.461]
No mean, no subtraction, no $\beta$ bias term. It is cheaper to compute and, in practice, costs no measurable quality, which is why it is now the default.
Where the norm goes: the choice that decides whether training works
The original transformer applied normalisation after the residual addition. Modern models apply it before the sublayer. The difference looks cosmetic and is not.
Post-norm (original)
Pre-norm (modern standard)
Formula
x = LayerNorm(x + Attn(x))
x = x + Attn(LayerNorm(x))
Is there a clean identity path?
No — every residual passes through a norm
Yes — the residual stream is never normalised
Learning-rate warmup
Essential; training diverges without it
Helpful but far less critical
Behaviour at 40+ layers
Frequently unstable
Stable
Final quality when it does train
Sometimes marginally better
Marginally worse in theory, vastly more reliable in practice
The mechanism behind the table: in post-norm, the gradient flowing backwards through the residual stream must pass through a normalisation at every single layer, and each one rescales it. Forty rescalings compound. In pre-norm the residual stream itself is untouched — the norm sits on the branch, not the highway — so gradients reach layer 1 essentially intact. Pre-norm models are also why you can train a 70-layer model without an elaborate warmup schedule.
One real cost of pre-norm: because nothing normalises the stream itself, its magnitude grows steadily with depth. Models therefore apply one final normalisation after the last block, before the output projection.
Problem three: attention does not actually compute anything
Look again at what attention does. It produces a weighted average of value vectors. Averaging is a linear operation, and the projections around it are linear too. Stack pure attention layers and, aside from the softmax weighting, you have something close to a repeatedly-averaged linear map. It moves information around brilliantly. It transforms information hardly at all.
The feed-forward network supplies the missing computation. It runs on each position independently — no mixing between tokens — and it is where nonlinear processing happens:
$$\text{FFN}(\mathbf{x}) = W_2\,\phi(W_1\mathbf{x} + \mathbf{b}_1) + \mathbf{b}_2$$
The hidden layer is conventionally 4× wider than the model dimension: with $d_{\text{model}} = 4096$, the hidden size is 16,384. Expand, apply a nonlinearity, project back down.
What the widening is for
A useful way to read the FFN is as a key–value memory. Each of the 16,384 rows of $W_1$ acts as a detector: it dots against the incoming vector and fires if it finds its pattern. The corresponding column of $W_2$ is what gets written back into the residual stream when it fires. Interpretability work has found individual FFN neurons that activate for remarkably specific things — a particular syntactic construction, a factual relation, a programming idiom.
Under this reading, the wide hidden layer is the model's storage. It holds an enormous number of learned "if you see this pattern, add this to the stream" rules, and the 4× expansion is simply how many such rules fit. That is consistent with a finding that surprises people the first time they see it: most of a language model's parameters live in the feed-forward layers, not in attention.
Activation functions, with numbers
The nonlinearity $\phi$ has evolved across model generations.
Input $x$
ReLU
GELU
Swish (SiLU)
−2.0
0.000
−0.046
−0.238
−0.5
0.000
−0.154
−0.189
0.0
0.000
0.000
0.000
0.5
0.500
0.346
0.311
2.0
2.000
1.954
1.762
ReLU is $\max(0,x)$. Simple and fast, but its derivative is exactly 0 for all negative inputs. A neuron pushed into that region receives no gradient and can never recover — the "dying ReLU" problem.
GELU is $x \cdot \Phi(x)$, where $\Phi$ is the standard normal cumulative distribution — a soft, probabilistic gate. At $x = -0.5$: $\Phi(-0.5) = 0.3085$, so $\text{GELU}(-0.5) = -0.5 \times 0.3085 = -0.154$. The small negative output means a nonzero derivative, so no neuron is permanently dead.
SwiGLU is the current standard in large models and is structurally different: it uses a gate. Two projections are computed from the same input, one is passed through Swish, and the two are multiplied elementwise before the down-projection:
$$\text{SwiGLU}(\mathbf{x}) = W_3\Big(\text{Swish}(W_1\mathbf{x}) \odot (W_2\mathbf{x})\Big)$$
That is three weight matrices where the classic FFN has two. To keep the parameter count comparable, implementations shrink the hidden dimension to roughly $\tfrac{2}{3}$ of the usual $4d$ — which is why real model configs show odd hidden sizes such as 11,008 rather than a round 16,384.
The whole block, and where the parameters actually go
Python
Copy
1import torch.nn as nn23class TransformerBlock(nn.Module):4    def __init__(self, d_model, n_heads, d_ff, dropout=0.0):5        super().__init__()6        self.norm1 = nn.LayerNorm(d_model)7        self.attn  = nn.MultiheadAttention(d_model, n_heads,8                                           dropout=dropout, batch_first=True)9        self.norm2 = nn.LayerNorm(d_model)10        self.ffn = nn.Sequential(11            nn.Linear(d_model, d_ff),12            nn.GELU(),13            nn.Linear(d_ff, d_model),14            nn.Dropout(dropout),15        )1617    def forward(self, x, causal_mask):18        # pre-norm: the residual stream (x) is never normalised in place19        h = self.norm1(x)20        x = x + self.attn(h, h, h, attn_mask=causal_mask, need_weights=False)[0]21        x = x + self.ffn(self.norm2(x))22        return x
Now count parameters for a realistic 7-billion-parameter configuration: $d_{\text{model}} = 4096$, 32 layers, 32 heads of dimension 128, SwiGLU hidden size 11,008, vocabulary 32,000.
Text
Copy
Attention, per layer  W_Q, W_K, W_V, W_O : 4 x (4096 x 4096) = 67,108,864Feed-forward (SwiGLU), per layer  gate: 4096 x 11008 = 45,088,768  up:   4096 x 11008 = 45,088,768  down: 11008 x 4096 = 45,088,768                        135,266,304Norms, per layer      : 2 x 4096          =       8,192                        ---------------------------------Per layer total       :                       202,383,360x 32 layers           :                     6,476,267,520Token embedding       : 32000 x 4096     =    131,072,000Output projection     : 32000 x 4096     =    131,072,000Final norm            :                             4,096                        ---------------------------------TOTAL                 :                     6,738,415,616
Component
Parameters
Share of model
Feed-forward layers
4.33 B
64.2%
Attention layers
2.15 B
31.9%
Embedding + output projection
0.26 B
3.9%
Attention gets all the attention, and accounts for under a third of the weights. Almost two-thirds of a language model is feed-forward layers — the per-position pattern-matching memory. If you are thinking about where a model's factual knowledge is stored, that ratio is a strong hint.
The same arithmetic gives you the memory footprint directly: 6.74 billion parameters at 2 bytes each in 16-bit precision is 13.5 GB just to hold the weights, before any activations or KV cache. At 4-bit quantisation it is about 3.4 GB, which is the difference between "needs a data-centre GPU" and "runs on a laptop".
Mixture of experts: a bigger feed-forward layer you only partly run
Since the feed-forward layers hold most of the weights, they are the obvious place to add capacity. A mixture-of-experts (MoE) layer replaces the single FFN with several independent FFNs — the experts — plus a small router, a linear layer that scores every expert for each token. Only the top-scoring few experts run for that token, and their outputs are added together, weighted by the router's scores. Attention, norms and embeddings are unchanged.

Work it through for Mixtral 8x7B: the same 32 layers and $d_{\text{model}} = 4096$ as above, but with 8 SwiGLU experts of hidden size 14,336 per layer, the router choosing 2 of them per token, and grouped-query attention with 8 KV heads (so $W_K$ and $W_V$ are $4096 \times 1024$).
Text
Copy
Attention, per layer  W_Q, W_O : 2 x (4096 x 4096) = 33,554,432  W_K, W_V : 2 x (4096 x 1024) =  8,388,608                                 41,943,040One expert (SwiGLU): 3 x 4096 x 14336 = 176,160,768Router:              4096 x 8         =      32,768Stored:   32 x (41.9M + 8 x 176.2M + router + norms) + embeddings        = about 46.7 billion parametersActive:   32 x (41.9M + 2 x 176.2M + router + norms) + embeddings        = about 12.9 billion parameters per token
Those match the figures Mistral published for the model: 46.7B total, 12.9B used per token. The trade is plain. Compute per token is that of a 13B dense model, but the memory bill is that of a 47B one, because any token may be routed to any expert. MoE buys capacity cheaply in FLOPs and not at all in memory, which is why it suits large serving clusters better than a single small GPU. Many of the largest recent open models are MoE, with far more experts per layer and a smaller fraction of them active per token.
Using this when something goes wrong
If a model you are training produces NaN losses, work through the block components in order. Check that you are using pre-norm. Check that the norm epsilon is not so small it underflows in mixed precision — $10^{-5}$ is typical, and $10^{-8}$ in fp16 can divide by an effective zero. Check for a missing residual connection, which is a one-character bug that silently caps how deep your model can usefully be.
If you are choosing a model to deploy, the parameter breakdown tells you what quantisation will and will not save. Quantising the feed-forward weights recovers most of the memory because that is where most of the weights are. Quantising the KV cache is a separate problem entirely — it scales with context length and users, not with parameter count, and no amount of weight compression touches it.
If you are reading a new architecture paper, the block is the right unit of comparison. Nearly every "new transformer variant" is a change to one of four slots: what normalisation is used, where it sits, what the attention pattern is, and what the feed-forward layer looks like — a mixture of experts, for instance, is a change to the last slot only. Identify which slot changed and you have understood the contribution in a sentence.