Skip to content
AI.info

Deep architectures

Anatomy of a Transformer Block

Trace data and gradients through normalization, multi-head attention, feed-forward sublayers, residual connections, and dropout in modern Transformer blocks.

By the end you can

Visual

A pre-normalized Transformer block

The residual stream passes through two learned update branches.

Read the six steps as three pairs. Each pair normalizes the carried stream, computes one learned update from it, and adds that update back. Nothing in the block replaces the stream; every branch writes into it.

FigureProcess · 6 steps
  1. 1

    Normalize residual stream

    Prepare the current representation for attention.

  2. 2

    Multi-head attention

    Route information across permitted positions or sources.

  3. 3

    Residual addition

    Add the attention update back to the carried stream.

  4. 4

    Normalize again

    Prepare the updated stream for channel transformation.

  5. 5

    Feed-forward network

    Apply a position-wise nonlinear expansion and projection.

  6. 6

    Second residual addition

    Carry the transformed update into the next block.

Attention mixes positions; the feed-forward network mixes features

Self-attention lets one position pull information from other positions. The feed-forward sublayer then transforms each position on its own, through the same nonlinear channel map everywhere.

Neither component is optional. That is a theorem, not a design preference. An ICML paper in 2021 states the result: “Specifically, without skip connections or multi-layer perceptrons (MLPs), the output converges doubly exponentially to a rank-1 matrix. On the other hand, skip connections and MLPs stop the output from degeneration.” Read the rate. Not slow decay with depth, and not exponential decay either. Doubly exponential. A stack of pure attention does not gradually blur the differences between positions. It collapses them to one direction almost immediately. The paper is called Attention is not all you need.

The collapse shows up in training, not only in the forward pass. A different group carried it into optimization at NeurIPS the following year. Rank collapse makes the query and key gradients vanish at initialization, and depth-dependent scaling of the residual branches prevents it. A collapsed block does not report an error. It reports zero gradient on exactly the projections that were supposed to learn the routing.

So routing without transformation does not merely fail to add anything. Measured over depth, it destroys the value features that were already there.

Transformer blocks alternate communication across positions with computation within each position.

Comparison

Pre-norm and post-norm residual layouts

Where the normalization sits changes optimization and signal paths. The difference was derived before it was folded into practice. Ten authors put it precisely in 2020: “Specifically, we prove with mean field theory that at initialization, for the original-designed Post-LN Transformer, which places the layer normalization between the residual blocks, the expected gradients of the parameters near the output layer are large.” Large gradients at the output, at step zero. That is what the learning-rate warm-up stage exists to survive. In the Pre-LN layout the gradients are well-behaved at initialization, and the warm-up can go.

The same result has a price tag on the other side of the ledger, reported independently in 2019 by Nguyen and Salazar. Pre-norm residual connections together with smaller initializations enable warmup-free, validation-based training at large learning rates. Their three combined changes are PreNorm, ScaleNorm and FixNorm. Together they always converge on five low-resource translation pairs from TED Talks-based corpora, for an average of +1.1 BLEU over state-of-the-art bilingual baselines and a new 32.8 BLEU on IWSLT'15 English-Vietnamese. Then the high-resource case. On WMT'14 English-German, ScaleNorm and FixNorm stay competitive while PreNorm degrades performance. The layout that removes a training pathology in one regime loses ground in another.

The third column has a measured price too. RMSNorm, proposed by Zhang and Sennrich in 2019, drops the mean subtraction: they hypothesise “that re-centering invariance in LayerNorm is dispensable”. It “achieves comparable performance against LayerNorm”, they report, but “reduces the running time by 7%~64% on different models”. Read the range rather than the headline. On one architecture the swap is nearly free. On another it is worth most of the wall clock. A normalization choice is validated per architecture, not adopted as a rule.

FigureComparison · 3 columns

Pre-normalization

Normalize before each learned branch, then add its output.

  • Direct residual path through depth
  • Commonly stable for deep training
  • Final normalization often required
  • Branch scale still matters

Post-normalization

Apply the learned branch, add the shortcut, then normalize.

  • Matches the original Transformer layout
  • Normalized output after every sublayer
  • Can be harder to optimize deeply
  • May require careful schedules or initialization

Alternative norms and scaling

Use RMS-style norms, residual scaling, or gated updates.

  • Changes centering and scale behavior
  • Can improve efficiency or stability
  • Not interchangeable across checkpoints
  • Requires architecture-specific validation

Example

The feed-forward layer is often where most parameters live

A standard FFN expands width, applies a nonlinearity or gate, then projects back. Every choice below moves parameters, activation memory, or serving cost — usually all three.

  • Expansion ratio: a wider hidden layer increases channel transformation capacity and activation cost.
  • Activation choice: GELU, gated linear units, and related functions alter smoothness and parameterization.
  • Gated FFN: one branch modulates another, increasing parameters while changing feature selection.
  • Mixture-of-experts: selected FFNs replace the dense FFN for conditional capacity.
  • Quantization pressure: these matrices are what serving-side compression actually targets, and the savings are documented. A 2022 paper reports that its authors “develop a procedure for Int8 matrix multiplication for feed-forward and attention projection layers in transformers, which cut the memory needed for inference by half while retaining full precision performance”. The trick is a mixed-precision split. Outlier dimensions stay in 16-bit while more than 99.9% of values are multiplied in 8-bit. That is how a 175B checkpoint halves its inference memory with no performance degradation. GPTQ pushed the same weights further the next year: 3 or 4 bits, about four GPU hours for a 175B model, negligible accuracy loss. The model then fits on a single GPU, running around 3.25x faster than FP16 on an A100 and 4.5x on an A6000.

Case

Two thirds of a Transformer block's weights never look at another position

The original Transformer fixes the numbers. Its 2017 paper states that “the dimensionality of input and output is dmodel = 512, and the inner-layer has dimensionality dff = 2048” — a fourfold expansion. Table 3 lists the base configuration: six layers, eight heads, dk = dv = 64, 100K training steps, 65 million parameters. Do the arithmetic at those widths. The FFN's two matrices hold about 2.1 million parameters per block. The four attention projections hold roughly 1.05 million. Two thirds of a block's weights sit in the part that never looks at another position.

That ratio is not an artefact of one 2017 configuration, and it is not only arithmetic. An EMNLP paper in 2021 opens by stating it flatly: “Feed-forward layers constitute two-thirds of a transformer model’s parameters, yet their role in the network remains under-explored.”

A third paper counted the complementary share directly, on models its authors had trained themselves. Each head is 6.25% of an attention layer's parameters on WMT, about 8.34% on BERT. The paper gives the total: “roughly speaking, in both our models, approximately one third of the total number of parameters is devoted to MHA across all layers”. So one paper computes the split from published dimensions, a second states it as background, a third counts it on a WMT Transformer and on BERT. Omer Levy co-authored two of them, so those two are overlapping teams rather than fully independent ones. The measurement lands on the same third either way.

Figure

The part of the block people describe as “the attention layer” holds a third of its matrices. Dimensions and the 65-million total are Table 3 of “Attention is All you Need”.

Key idea

The residual stream is not a passive wire

Every block writes updates into the carried representation. Scale, normalization, dropout, and numerical precision determine whether those updates remain balanced across depth.

A model can train while some branches contribute little, or while residual norms drift. That is not a hypothetical failure mode of research prototypes. It is measurable in shipped weights. An ICLR paper in 2025 prunes whole blocks chosen by cross-layer similarity, then heals the model with QLoRA finetuning on a single 40GB A100: “Surprisingly, with this method we find minimal degradation of performance until after a large fraction (up to half) of the layers are removed for some common open-weight models.” Half the blocks. On question-answering benchmarks. On models people are running today.

A second group reached the same conclusion by a different route, in a 2025 paper called ShortGPT: “In this study, we identify notable redundancy across the layers of LLMs, where some layers contribute minimally to the overall network functionality.” Their measurement is the useful part for anyone instrumenting a stack. They score each layer by Block Influence, the similarity between a layer's input and its output. A block whose output closely resembles its input has written almost nothing into the stream. That is a number you can compute per layer, not a suspicion you carry.

So track update-to-stream ratios and layerwise statistics rather than assuming every block participates equally. Two independent groups have now found that the assumption fails by a wide margin.

Residual connectivity enables deep paths, but healthy contribution still requires measurement.

Analogy

A meeting followed by private desk work

In a team meeting each person gathers what they need from colleagues. The private desk work that follows transforms what each person learned.

Colleagues act on intentions they could state aloud, while attention and feed-forward layers are differentiable tensor operations. What carries over is the rhythm: communication, then local computation, then communication again.

Attention decides what to combine; the feed-forward sublayer decides how to transform the combined representation.

Steps

Instrument a Transformer block

A few probes can localize many architecture failures before expensive training completes.

Step 4 is usually written as a suggestion, so it is worth saying what an ablation actually returns. An ACL paper in 2019 reports: “For example, on the English-Russian WMT dataset, pruning 38 out of 48 encoder heads results in a drop of only 0.15 BLEU.” Ten heads of forty-eight were doing the work.

The same shape turns up on other models. In the head-counting paper from 2019, only 8 of 96 encoder self-attention heads changed performance significantly when removed. Half of those removals raised BLEU. Twenty per cent of WMT heads and 40% of BERT heads can be pruned with no noticeable impact, and pruning 50% of BERT's heads speeds up inference by up to about 17.5% at larger batch sizes. That is the expected outcome to compare your own numbers against. If ablating a head changes nothing, the finding is ordinary. Step 4 is how you tell which heads those are before you pay to serve them.

FigureProcess · 5 steps

  1. 1. Assert all head shapes

    Verify projections, head split, masks, concatenation, and output projection.

  2. 2. Record residual and update norms

    Measure each branch relative to the carried stream.

  3. 3. Track normalization statistics

    Look for collapse, extreme scale, or precision sensitivity.

  4. 4. Measure head diversity

    Compare routing patterns and ablate redundant heads.

  5. 5. Profile FFN activations

    Capture memory peaks and sparsity under realistic sequence lengths.

Activation memory scales with sequence length and retained intermediates

Training stores values needed for the backward pass: projections, normalization inputs, attention-related tensors. Checkpointing can trade recomputation for memory, and both sides of that trade have been measured.

The classic result is from 2016: O(sqrt(n)) memory for an n-layer network, at the price of one extra forward pass per mini-batch. Its authors put a number on it. A 1,000-layer deep residual network drops from 48G to 7G for about 30 percent extra running time on ImageNet. That exchange rate is not a constant of nature. An MLSys paper in 2023 showed most of the redundant compute is avoidable: “We evaluate our approach on language models up to one trillion parameters in scale and show that our method reduces activation memory by 5x, while reducing execution time overhead from activation recomputation by over 90%.” On a 530B GPT-style model on 2240 NVIDIA A100 GPUs, that is 54.2% Model FLOPs Utilization against 42.1% with recomputation.

How the kernel is written changes what gets stored. A fused attention kernel may never materialize a full score matrix, so anyone analysing an architecture has to separate the mathematical tensors from the ones the kernel actually keeps. FlashAttention made that distinction operational at NeurIPS in 2022. It is “an IO-aware exact attention algorithm that uses tiling to reduce the number of memory reads/writes between GPU high bandwidth memory (HBM) and GPU on-chip SRAM”. Exact, not approximate, and it never stores the full score matrix. On BERT-large at sequence length 512 it gave a 15% end-to-end wall-clock speedup against the MLPerf 1.1 training speed record. GPT-2 at length 1K gained 3×. Long Range Arena at 1K–4K gained 2.4×. These are also “the first Transformers to achieve better-than-chance performance on the Path-X challenge (seq. length 16K, 61.4% accuracy) and Path-256 (seq. length 64K, 63.1% accuracy)”.

The computational graph and the runtime memory graph are related but not identical.

Key takeaways