Inverse Depth Scaling From Most Layers Being Similar

AI summary

Overview

Research area: Neural scaling laws, large language model (LLM) architecture, and mechanistic interpretability of Transformer depth.

Technical level: Advanced (assumes familiarity with neural scaling laws, residual networks, and neural ODE ideas), though the central claims are stated plainly enough for a motivated intermediate reader.

Scope: The paper quantifies how model depth (as opposed to total parameter count) affects loss in LLMs, argues that loss scales roughly inversely with depth, and attributes this to layers acting as an ensemble rather than as a composition of increasingly abstract features.

What This Paper Is About

Neural scaling laws describe how LLM loss falls predictably as models and datasets grow, but most work treats "model size" as a single lumped quantity rather than separating width from depth. The authors ask how LLMs actually use their depth and what quantitative law connects depth to performance. They compare three candidate pictures of depth use—compositional assembly, procedural assembly, and ensemble averaging—against measurements of real LLMs and controlled toy models.

Key Contributions

1. Measurement of inverse depth scaling in LLMs. The authors fit a decomposed scaling form to about 200 data points reconstructed from the Chinchilla models and find the depth-dependent loss component scales approximately inversely with depth (α_ℓ = 1.2 ± 0.3).

2. A hidden-state diagnostic for depth usage. They analyze the angle θ(h_l, h_{l+1}) between neighboring hidden states across layers in Pythia-410m evaluated on FineWeb, and use PCA of per-token angle trajectories to show that 99.6% of tokens are approximately evenly updated in the middle layers, while only 0.4% (mostly the first tokens of documents) stop updating early.

3. A mechanistic explanation via ensemble averaging. Combining LLM measurements with controlled toy-model teacher–student experiments, they conclude that most layers make similar incremental updates and reduce error by averaging, rather than by compositional hierarchy or by discretizing smooth dynamics.

4. Theory connecting smooth dynamics to depth exponents. They derive that procedural assembly (matching smooth dynamics) should yield a converged depth exponent α_ℓ = 3, which does not match LLM observations, whereas ensemble averaging yields α_ℓ typically equal to 1.

Main Findings

• Three regimes of depth use. The paper distinguishes compositional assembly (layers build increasingly abstract representations), procedural assembly (layers discretize a smooth dynamical system, like a neural ODE), and ensemble averaging (layers act as similar shallow subnetworks whose independent errors partly cancel).

• Most layers are not compositional. In Pythia-410m on FineWeb, the first and last layers rotate hidden states by large angles close to π/2, but most middle layers update hidden states by similarly small angles. PCA of per-token trajectories shows a small cluster containing 0.4% of tokens (typically document-initial tokens) that stop updating early, and a large cluster containing 99.6% that are approximately evenly updated in the middle layers, with middle-layer angles around 0.45 rad.

• Middle-layer updates shrink with depth. Averaging θ(h_l, h_{l+1}) over the middle layers l = 2, 3, ..., ℓ−1 gives a quantity that scales approximately inversely with depth. This supports procedural assembly or ensemble averaging rather than compositional assembly (Result 1: Not compositional assembly).

• Neighboring updates are weakly correlated. The angle θ(Δh_l, Δh_{l+1}) between neighboring layer updates is typically large, which is inconsistent with smooth dynamics, though the authors caution that LLM data alone lacks a calibrated reference.

• Fitted LLM scaling exponents. Using L = c_m/m^{α_m} + c_ℓ/ℓ^{α_ℓ} + c_D/D^{α_D} + L_0 (seven free parameters), they obtain α_m = 0.98 ± 0.08, α_ℓ = 1.2 ± 0.3, and α_D = 0.30 ± 0.01. The average relative error between empirical loss and the fitted prediction is 0.4%. The fitted α_D closely matches the original Chinchilla scaling exponent (Result 2: Empirical inverse depth scaling).

• Toy model setup. The student uses RMSNorm, ℓ residual blocks h_l = h_{l−1} + MLP_l(h_{l−1}) with MLP_l(v) = B_l ReLU²(A_l RMSNorm(v) + b_l), width m = 32, output dimension n = 128, teacher depth ℓ* = 128, and student depths ℓ from 6 to 48. Students are trained for 40,000 steps with Adam; the maximum training steps studied is t_max = 80,000.

• Two toy regimes reproduce inverse scaling. With independent teacher weights (weight correlation ρ = 0), the measured depth exponent is robustly close to α_ℓ ≈ 1 across temperatures. With tied teacher weights (ρ = 1), α_ℓ increases from 1 to 3 as teacher temperature increases.

• The α_ℓ ≈ 1 case with smooth dynamics is imperfect training, not a different mechanism. At late training with tied weights, α_ℓ approaches 3, matching the predicted discretization-dominated scaling; low teacher temperature slows convergence and produces an α_ℓ ≈ 1 that the authors attribute to incomplete optimization.

• Theory of ensemble averaging. Writing h(1) − h*(1) as (1/ℓ) Σ_l (f(l/ℓ) − ∫₀¹ f*(s) ds), each layer contributes an O(1/ℓ) error. Summation over ℓ layers gives a worst case ‖h(1) − h*(1)‖ = O(1), with typical behavior O(1/√ℓ) under independent layer-wise errors, implying L ∼ 1/ℓ with α_ℓ typically equal to 1.

• Toy hidden states match LLM hidden states. Trained students with independent teacher weights show approximately uniform layer-wise updates across depth, update magnitude scaling inversely with the number of layers, and neighboring-update angles large and close to π/2 — qualitatively similar to the LLM measurements (Result 3: Depth scaling from ensemble averaging).

• Width–depth balance implication. With α_m ≈ 1 and α_ℓ ≈ 1 under a fixed parameter budget N ∝ m²ℓ, the optimal width–depth relationship is m ∝ ℓ. At that ratio both terms scale as N^{−1/3}, close to the empirical 0.34 model-size exponent reported in Chinchilla scaling.

• Functional groups, not single layers. A causal tracing experiment on Pythia-12b (reproduced from Meng et al. 2022a) with the prompt "The Space Needle is in downtown" shows two distinct groups of layers — an earlier group integrating the tokens and a later group extracting "Seattle" conditioned on "downtown" — with recovery probability fluctuating across layers within each group, consistent with distributed rather than localized computation.

Methodology in Plain English

The authors work on two fronts.

First, they run inference-only probes on a pretrained LLM (Pythia-410m), streaming text from the FineWeb dataset. For each token they record the angle between the hidden state after layer l and after layer l+1, the norms of those states, and the angle between consecutive layer updates. They also take each token's full vector of angles across depth, apply PCA, and compare the resulting clusters to two constructed "ideal" trajectories: one that stops updating in the middle and one that updates evenly in the middle (0.45 rad). Separately, they fit a decomposed loss function with separate width, depth, and dataset-size power laws to roughly 200 Chinchilla model data points by minimizing the mean squared error of the logarithm of the loss.

Second, they build a teacher–student toy model that keeps the essential architectural ingredient of Transformers — residual connections — while stripping away attention, embedding training, and other complexity. The teacher is deeper than the student but shares the same width, so representational error is eliminated by construction and only transformation error remains. They manipulate two knobs: whether teacher MLP weights are tied across layers (which induces smooth, neural-ODE-like dynamics) or sampled independently (which induces random-walk-like, non-smooth dynamics), and the teacher temperature (which controls how peaked the target output distribution is). They then fit the student's loss against depth and compare the fitted exponent, plus the student's hidden-state statistics, against the LLM measurements.

Why This Matters

Impact on research. The paper moves neural scaling laws from a lumped "model size" picture to a decomposition of width- and depth-dependent contributions, and it supplies a quantitative target (α_ℓ) that future theory and architecture work can be tested against. It also reframes the well-documented redundancy of Transformer layers as a specific, measurable mechanism — averaging — rather than a vague inefficiency. The authors note that their hidden-state analysis characterizes layer behavior statistically and cannot reveal the mechanistic function of individual layers, and they cannot rigorously exclude other mechanisms that might produce the same inverse depth scaling.

Real-world applications (as grounded in the paper):

• Compute allocation and model design. Because the fitted exponents imply an optimal width–depth relationship of m ∝ ℓ under a fixed N ∝ m²ℓ, practitioners can use the decomposition to reason about how to split a parameter budget between width and depth.
• Architecture innovation. The findings suggest that improving LLM efficiency in depth may require architectural changes that encourage compositional use of depth; the authors point to recurrent depth as one direction that appears to better explore data hierarchy.
• Mechanistic interpretability and model editing. The functional-group picture, illustrated by causal tracing on factual recall in Pythia-12b, is consistent with prior work showing that editing memories across a group of layers works better than targeting a single layer.
• Scaling-law forecasting. The predicted model-size exponent of 1/3 arising from inverse depth scaling is close to the empirical 0.34 in Chinchilla scaling, which the authors suggest may mean empirical scaling laws are partly a consequence of inverse depth scaling.

Industry relevance. Teams that choose depth and width for large training runs, or that design post-training and editing interventions on specific layer groups, can use these results as a quantitative prior about where additional depth is likely to help and where it is likely to be absorbed as redundant averaging.

Future Directions

1. Derive the loss decomposition from first principles. The authors explicitly state that the decomposed scaling form was proposed from a combination of empirical findings and theoretical insights rather than derived rigorously, and is supported mainly by fit quality.

2. Study width–depth–dataset-size interactions. In particular, the interplay is flagged as potentially important for small models, where the fitted width-exponent ≈ 1 and the m ∝ ℓ relation are not rigorously true in practice.

3. Explain why ensemble averaging emerges. The paper calls for a more mechanistic account of the origins of ensemble averaging, pointing at the architectural bias of residual connections and at the possibility that next-token prediction cannot be modeled by a smooth dynamical system.

4. Pursue architectures that exploit composition. Since inverse depth scaling implies inefficient use of depth, the authors suggest encouraging compositional use of depth, including recurrent depth approaches, and note the possibility that other undiscovered mechanisms could produce the same inverse depth scaling and similar hidden-state signatures.

Target Audience

This paper is most useful to machine learning researchers working on scaling laws, Transformer architecture design, and mechanistic interpretability, as well as to engineers who decide how to allocate parameter budgets between width and depth for large training runs. Readers will get the most out of it if they are comfortable with residual networks, power-law scaling fits, and the neural-ODE view of deep networks; the core empirical claims about inverse depth scaling and layer similarity are accessible to a broader technically literate audience.

Authors’ abstract

Neural scaling laws relate loss to model size in large language models (LLMs), yet depth and width may contribute to performance differently, requiring more detailed studies. Here, we quantify how depth affects loss via analysis of LLMs and toy residual networks. We find loss scales inversely proportional to depth in LLMs, probably due to functionally similar layers reducing error through ensemble averaging rather than compositional learning or discretizing smooth dynamics. This regime is inefficient yet robust and may arise from the architectural bias of residual networks and target functions incompatible with smooth dynamics. The findings suggest that improving LLM efficiency may require architectural innovations to encourage compositional use of depth.