← Concept library

NLP Foundations

Depth, Width, and Aspect Ratio

Given a fixed parameter budget, why pretraining loss barely cares whether you go deep or wide, and the hardware and capability reasons architects still don't pick the ratio at random.

intermediate · 8 min read · Premium

Fix a transformer's total parameter count and you still have one degree of freedom left: how many blocks to stack (depth, N) versus how wide to make each block (d_model). A block's parameter count scales roughly as 12 * d_model^2 (anatomy-of-a-transformer-block works out the 4d^2 attention plus 8d^2 FFN breakdown), so total parameters scale as N * 12 * d_model^2, and the same total can be reached by many different (N, d_model) pairs: deep and narrow, shallow and wide, or anything between.

The loss is nearly indifferent

Kaplan et al., 2020 measured pretraining loss across a wide range of depth-width combinations at fixed parameter count and found the effect of shape on loss surprisingly small; across a broad middle range of reasonable aspect ratios, loss is close to flat. This is a genuinely useful empirical result: it means architects are not required to find one precise "correct" ratio to avoid leaving loss on the table, and it explains why different model families settled on visibly different ratios without any of them being obviously wrong.

But the flat region has real edges elsewhere

Flat loss does not mean shape is irrelevant, it means shape's effect on this particular metric, at this particular range of ratios, is small. Two other axes push architects toward specific choices anyway.

Hardware. Width parallelises cleanly, a wide matmul splits trivially across accelerators with minimal communication. Depth adds sequential dependency, block N+1 cannot begin until block N finishes, which hurts pipeline latency and complicates pipeline parallelism, since a deep model run across many devices in a pipeline pays a "bubble" cost every time the pipeline fills and drains. For inference-heavy deployment specifically, shallower-and-wider is often favoured even at matched loss, purely because it is faster per forward pass on parallel hardware.

Capability, not just loss. Levine et al., 2020 argue that self-attention has an effective depth efficiency distinct from what aggregate pretraining loss reveals: certain kinds of compositional, multi-step computation benefit from additional depth in a way that a single flat loss number, averaged over an entire corpus, does not clearly surface. The claim is not that Kaplan's finding is wrong, it is that "loss is flat across shape" and "every downstream capability is indifferent to shape" are different claims, and only the first has strong support.

What real models actually chose

Chowdhery et al., 2022 report PaLM's 540B configuration used an aspect ratio (d_model / n_layers) of roughly 128, deliberately narrower per layer and deeper than some contemporary models of similar total size. Part of that choice tracked the parallel attention-and-FFN block formulation PaLM used (see parallel-attention-and-ffn) and the communication pattern of the TPU pods it trained on, a reminder that aspect ratio in practice gets set as much by the target hardware's topology as by any pure loss-optimal number derived in the abstract.

Keep reading with Pro.

You're reading the preview. Unlock the full concept plus the library, study plans, the AI mentor, and daily emails.

Sign in to save and react.
Share Copied