How Depth Growth and Looping Shape In-Context Learning in Language Models
Ferdinand Kapl (TU Munich, Helmholtz AI) with Kaitlin Maile, Stefan Bauer and Johannes von Oswald (Google Paradigms of Intelligence) show that language models which reuse computation across depth, by growing depth during training or looping a shared block, learn more from context than a dense model of the same depth.
Ask this paper
Language-model comparison. In a 3B comparison, grown and looped models show larger early-to-late token-loss reductions than the dense baseline and higher GSM8K accuracy at every demonstration count. They also predict unseen formal languages on RegBench more accurately, with the advantage appearing within the first string.
Acquisition and retention. In controlled Omniglot tasks, both acquire in-context learning earlier than the dense reference when labels are reassigned per sequence. When fixed class-label associations could also be stored in the weights, they keep higher in-context accuracy and lower in-weights accuracy.
Mechanism. Tracing heads through training shows that depth growth copies early induction heads across layers, and most of the strongest final heads descend from a few early ancestors. Looped models show the same repeated pattern across applications of the shared block.
Ablations. Grown and looped models depend more on induction-ranked heads for few-shot performance and later-token prediction, while function-vector heads remain the most important heads for few-shot accuracy in all models.
Abstract
Looped and depth-grown language models reuse computations across depth and have attracted interest for their reasoning capabilities. Building on prior suggestive evidence of improved in-context-learning (ICL), we investigate how such reuse favors ICL, including how it shapes the acquisition, retention, and mechanisms of this ability. Our analysis spans pretrained language models up to 3B parameters and controlled tasks. The experiments provide evidence for a shared inductive bias toward learning from context: grown and looped models rely more on information in context and, compared with a full-depth dense reference, acquire ICL earlier in controlled tasks and retain it more strongly when learning fixed associations in the weights is also possible. Tracing attention heads through training reveals that depth growth propagates early induction heads across layers. Ablations show that grown and looped models rely more strongly on induction-ranked heads for both few-shot performance and later-token prediction. Together, these findings provide evidence for a shared inductive bias toward learning from context and suggest that (soft) parameter reuse shapes how this ability develops and which attention mechanisms support it.