ΒΆPaper Feed

Issue 22 Β· Pick 09 AI / ML βœ“ read

Why Larger Models Learn More: Effects of Capacity, Interference, and Rare-Task Retention

Jing Huang, Daniel Wurgaft, Rachit Bansal, Laura Ruis, Naomi Saphra, David Alvarez-Melis, Andrew Kyle Lampinen, Christopher Potts, Ekdeep Singh Lubana

TL;DR: Larger models don't just learn faster β€” there are parts of the data distribution a small model will never learn, even with infinite data, and this paper proposes a concrete mechanism for why. In a mixture of tasks, a small model spends all its neurons on frequent tasks, and the gradients from those tasks keep overwriting whatever fragile progress the rare tasks make between their sparse appearances. A large model can fully absorb the frequent tasks, which makes their gradients nearly silent, leaving spare capacity where rare-task signal can accumulate undisturbed. The authors prove this in a linear multi-task toy model and then reproduce every signature of it β€” behavioral, representational, and gradient-level β€” by injecting controlled synthetic tasks into real OLMo pretraining runs from 4M to 4B parameters.

The scaling-law hint that data can't always substitute for parameters

Start from the Chinchilla-style scaling law:

L(N, D) = L_0 + \frac{A}{N^{\alpha}} + \frac{B}{D^{\beta}},

where N is parameter count, D is data, L_0 is irreducible loss, and \alpha \approx 0.46, \beta \approx 0.51 empirically. Under compute-optimal training, loss falls as L_C(N) \propto N^{-\gamma} with \gamma = 0.34. But if you give a model infinite data, its loss floor is L_\infty(N) \propto N^{-\alpha}.

Here is the observation the paper builds on: since \alpha > \gamma, the asymptotic loss curve is steeper than the compute-optimal one. That means there exist model-size pairs where a small model trained forever still can't match a larger model trained under a finite budget. The paper formalizes this as two regimes: loss gaps that are learnable via data scaling (the small model is just undertrained and can catch up) versus learnable via model scaling (no amount of data closes the gap). The second regime is the interesting one β€” it says a fraction of the data distribution is genuinely locked behind parameter count. The question is what is locked away, and mechanistically why, given that the small model can typically express the missing tasks just fine (you can teach them via fine-tuning).

A toy universe where "what width buys" is exactly computable

The authors build the simplest setting that reproduces monotonic scaling curves: a mixture of K linear regression tasks. Task k appears with frequency \pi_k (power-law: \pi_k \propto k^{-\beta}) and has a covariance C_k with eigenvalues \lambda_{k,1} \geq \lambda_{k,2} \geq \dots (power-law: \lambda_{k,j} \propto j^{-\alpha}). Slowly decaying spectra mean the task needs many directions β€” that's their proxy for complexity. Different tasks occupy orthogonal subspaces, so there is no representational conflict baked in; any conflict must come from learning dynamics. The student is a shared width-N linear encoder U (the bottleneck β€” this is "model size") plus per-task decoders.

For this model, the asymptotically optimal solution is fully characterized (Theorem 3). Define the utility of feature j of task k as

u_{k,j} = \pi_k \lambda_{k,j},

frequency times variance-explained. The optimal width-N encoder keeps exactly the N features with the highest utility. That's the entire answer to "what does width buy": each extra neuron admits the next feature down the global utility ranking. Rare tasks and complex tasks (whose signal is spread across many low-\lambda directions) have low-utility features, so they only make the cut at large N β€” even with infinite data.

features ranked by utility u = frequency Γ— variance utility small model: N neurons large model frequent task features medium task features rare task features
The optimal width-N model keeps the N highest-utility features. Rare tasks' features (and frequent tasks' deep, low-variance modes) sit in the tail β€” below the small model's cutoff but within the large model's. Note the interleaving: a rare task's top mode can outrank a frequent task's fifth mode, so learning order is set by utility, not frequency alone.

They verify this experimentally with K=32 tasks: the empirical phase diagram of which task gets learned at which width matches the analytic staircase N^*(m) almost exactly (Fig. 2 of the paper, and a sharp rank-1 version in Appendix E where the predicted critical width is literally N_{\text{crit}}(k) = k). Crucially, training 10Γ— longer (1M steps, Appendix E.4) doesn't bend the curves: above-capacity tasks sit at the random baseline indefinitely. That's the model-scaling regime made concrete.

The mechanism: quiet gradients and the update-and-forget loop

The utility ordering says what the optimum looks like, but not why gradient descent gets there β€” especially for a task that's 0.25% of the data. A rare task delivers a statistically real learning signal every time it appears. Why does the small model fail to compound it?

The answer is the paper's core contribution. The gradient contributed by the frequent tasks obeys (Theorem 4):

\|G_{\mathsf{F}}(U)\|_F \leq 2\sqrt{\lambda_1(M_{\mathsf{F}})\,\delta_{\mathsf{F}}(U)},

where M_{\mathsf{F}} is the frequent tasks' weighted covariance and \delta_{\mathsf{F}}(U) is the residual β€” the part of the frequent tasks not yet explained by the current representation. Tasks only push on the model through what it hasn't learned about them. A wide model can drive the frequent-task residual near zero, at which point frequent-task gradients go quiet. A narrow model can't: it lacks the neurons to fully absorb even the common tasks, so their gradients stay loud forever, and every frequent-task batch keeps churning the parameters.

Proposition 6 sharpens this into a threshold: a rare feature with utility \pi_r \lambda_r can stably occupy a neuron only if it outranks the weakest occupied frequent-task mode, \pi_r \lambda_r > \mu_N^{\mathsf{F}}. Below the critical width, you get a swinging update-and-forget dynamic: the rare batch pulls a neuron toward the rare feature, then thousands of frequent batches pull it back. The one-neuron caricature (Appendix C.4) makes the rate explicit β€” after a rare-task update, alignment decays like \theta_G \approx e^{-2\eta G}\theta_0 across the gap of G frequent-task steps. In a wide model, the frequent gradient is quiet, so the rare-task update persists until the next rare batch arrives, and the signal compounds.

training steps rare-task signal rare-task batches large model: retains, accumulates small model: update-and-forget
The matched-frequency injection experiment (paper Fig. 4). Both models see the same total amount of rare-task data; only the gap structure varies. The small model gains signal at each injection but frequent-task gradients erase it before the next one β€” an update-and-forget loop. The large model, whose frequent-task gradients are quiet, retains each increment and staircases up.

The clean test is the matched-frequency injection protocol: withhold the rare task for G steps, then inject a batch enlarged by exactly the right factor so that long-run frequency is identical across all settings. If learning depended only on total exposure, all runs would match. Instead, rare-task signal decays with G at all widths, but catastrophically in narrow models and barely at all in wide ones β€” isolating retention between observations, not sample count, as the bottleneck.

Does this survive contact with real pretraining?

This is where the paper earns its keep. They pretrain OLMo models at 4M, 20M, 300M, 1B, and 4B parameters on Dolma v1.7 (up to 210B tokens), injecting instances of two synthetic tasks at controlled frequencies from 7.8 \times 10^{-3} down to 2.4 \times 10^{-8} β€” roughly 1,000 instances per batch down to one instance every 10 batches. The tasks are three-token sequences over 100 randomly chosen vocabulary tokens: a comparison task (is val(\texttt{TOK1}) < val(\texttt{TOK2})?) and modular addition mod 100. With 5K held-out test instances, generalization requires learning the structure β€” the hidden order or arithmetic over arbitrary tokens β€” not memorizing pairs.

All three levels of the toy story replicate:

Behavior. Larger models learn lower-frequency injected tasks; smaller ones plateau at chance (Fig. 5). Tasks are learned in frequency order relative to naturally occurring reference tasks (Fig. 6a). The matched-frequency injection result reproduces too: injecting N instances every N batches, larger gaps degrade task loss despite identical global frequency (Fig. 6b). On modular addition, only large models at high frequency exhibit grokking β€” suggestive evidence that grokking-like memorize-then-generalize dynamics occur inside genuine LLM pretraining, not just in isolated toy runs.

Representations. Using distributed alignment search, they find a 1-D subspace in the first layer's residual stream that causally encodes global token order for the comparison task (96% interchange-intervention success), and Fourier modes for modular addition β€” the same features grokked transformers are known to use. Larger models and higher injection frequencies yield more of these task features, and feature presence tracks test accuracy (Fig. 7). This is the LLM analogue of the toy model's \mathrm{Tr}(P_U C_k) signal measurements.

Gradients. They localize the task circuit (first-layer MLP neurons), define a task reference direction g_r from the aggregated task gradient, and decompose each batch gradient into task-token and non-task-token parts, g = g_t + g_{nt}. The non-task part is the interference channel β€” general language modeling gradients pushing on the task circuit's neurons.

chart error: Expecting ',' delimiter: line 2 column 1 (char 646)
{"type": "bar", "title": "Cosine similarity with task reference direction g_r (T_CMP, task neurons)", "ylabel": "cosine similarity", "categories": ["non-task grad, 20M", "non-task grad, 1B", "full batch @ injection, 300M", "full batch @ injection, 1B"], "series": [{"name": "mean", "values": [0.10, 0.0000758, 0.04, 0.08]}], "note": "Sec. 4.4 / Fig. 9. The 20M model's non-task gradient collides with the task direction (0.10 Β± 0.09) β€” heavy interference; the 1B model's is essentially orthogonal (7.6e-5 Β± 0.02). Meanwhile the 1B model's batch gradient carries more genuine task signal at injection steps (0.08 Β± 0.02 vs 0.04 Β± 0.04 for 300M)."

This is the mechanism caught in the act at LLM scale: in the 20M model, general language-modeling gradients randomly collide with the task direction and scramble it between injections; in the 1B model they are almost exactly orthogonal, so the task circuit is left in peace to accumulate signal.

What it means, and where to push back

If this account holds, several things follow. First, an interpretation shift: "emergent" abilities of large models may often be rare-task retention effects β€” the capability was always learnable in principle, but only large models create the quiet parameter subspace where sparse signal can compound. Second, a practical lever: the paper explicitly notes that upweighting a target task's frequency in the data mixture may be a cheaper substitute for scaling parameters β€” the bottleneck is \pi_r \lambda_r vs. \mu_N^{\mathsf{F}}, and you can raise the left side instead of lowering the right. Gradient-surgery-style multi-task methods are another lever the framework naturally suggests. Third, a reframing of memorization: retaining traces of individual rare batches is the prerequisite for eventually abstracting the task β€” memorization as scaffolding for generalization, in line with Feldman's argument, here given a dynamical mechanism.

Now the skepticism, which the authors are unusually candid about. The theory lives entirely in a linear encoder with orthogonal task blocks β€” no depth, no attention, no feature superposition or reuse across tasks, which is exactly where real LLMs get interesting (real tasks share features; orthogonality assumes away transfer). The OLMo evidence, while genuinely multi-level, concerns artificial injected tasks with a controlled frequency knob; there is no demonstration that this mechanism quantitatively accounts for real-world scaling curves versus plain capacity or sample-efficiency effects, and the authors explicitly frame it as complementary to expressivity and sample-efficiency accounts, not a replacement. Scaling here is almost purely width (all OLMo models 300M+ have depth 16); whether depth scaling behaves the same is open. They did not test over-trained models or extreme frequency regimes, and the gradient analysis covers one task on a handful of first-layer neurons. And the phenomenological setup piece (Sec. 2) leans on power laws holding asymptotically, which they concede need not be true β€” it's motivation, not proof.

Still, the direction of the argument is hard to dismiss: three independent measurement modalities (loss, causal features, gradient geometry) tell the same story in both a solvable model and a real pretraining pipeline, with a matched-frequency control that specifically rules out "more exposure" explanations.

Where to spend your time: Section 3.2 with Figures 3–4 is the heart β€” the residual-controls-learning theorem and the injection experiment together are the "aha." Then Section 4.4 (Fig. 9) for the most striking empirical fact: non-task gradients orthogonal to the task circuit at 1B, colliding with it at 20M. Appendix C.4's one-neuron model is a five-minute read that makes the whole update-and-forget dynamic obvious.