How Depth Growth and Looping Shape In-Context Learning in Language Models
Authors: Ferdinand Kapl, Kaitlin Maile, Stefan Bauer, Johannes von Oswald
Organizations: Technical University of Munich · Helmholtz AI, Munich · Munich Center for Machine Learning (MCML) · Google, Paradigms of Intelligence Team
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.
Figures & tables
Figure 1: Grown and looped models benefit more from context. In the 3B comparison, grown and looped models show greater context gains in (a) NLL relative to the dense baseline and (b) early-to-late loss reduction relative to the baseline per token position, although the looped model retains higher NLL. Both panels aggregate token losses in the same power-of-two position bins: (a) compares the binned loss directly with the baseline, whereas (b) first computes each model’s improvement from a fixed early token position (50) to the bin ending at x and then subtracts the baseline’s improvement. Dashed horizontal lines show mean Δ NLL over all positions. (c) Both models exceed the baseline’s few-shot GSM8K ( Cobbe et al., 2021 ) accuracy, with larger gains for LIDAS . Results use 10,000 FineWeb-Edu documents and three demonstration seeds, respectively. Additional model sizes and exact experimental setup in Appendix B .
Figure 2: Grown and looped models predict unseen formal languages more accurately. In-context language learning on RegBench ( Akyürek et al., 2024 ) . (a) Test total variation distance (TVD) for the dense 12-layer, \LoopOp2×6 , and Grown b=3 model (mean ± standard error over four training seeds). (b) TVD relative to the dense model against in-context string index. (c) The same comparison by symbol index, separately for the first string and strings 2–10. Lower is better for all panels. Bands show standard errors. The looped model has width 256, versus 128 for dense and grown models: hyperparameters are selected on validation data, rather than matched in width. Panel (a) uses the full test set. Panels (b,c) retain only string indices present in every test sequence (strings 1–10). Experimental setup and full details in Appendix C .
Figure 3: Growing and looping favor stronger retention of ICL, at the cost of in-weights learning (IWL), and earlier acquisition of ICL when IWL is not possible. Solution preference and emergence of in-context learning on Omniglot-based tasks ( Singh et al., 2023 ) . (a) Training uses fixed class labels and P(bursty)=1 , so both context and in-weights associations can solve each training query. ICL is measured on held-out classes with labels reassigned per sequence. IWL is measured on training classes whose query class is absent from context. (b) Training follows the ICL probe setup using images from training classes, with labels reassigned in every sequence to remove a fixed class–label solution. Curves show mean accuracy, with bands indicating standard error across four (a) and eight (b) smoothed seed trajectories. Growth occurs at 50k and 100k steps. Experimental details in Appendix D .
Figure 4: Induction behavior during training. Top-2% mean PS (3B), with heads reselected at each checkpoint. Vertical lines mark growth. Hollow circles show immediate post-growth scores.
Figure 5: Depth growth propagates early strong induction heads across layers. Development and final location of individual induction heads. Top: smoothed prefix-matching scores (PS) of each 3B model’s twelve final top-2% heads over training (logarithmic step axis). LIDAS lineages follow layer duplication, with shared ancestral paths drawn once. In LIDAS , the twelve final heads trace back to four ancestors in the initial four-layer model. Two of these, L2H9 and L2H10, already have high PS before growth and account for ten of the twelve final heads. Bottom: final PS for every layer and head (zero-based indices). Among the final top-2% PS heads, orange highlights descendants of the early L2H9/L2H10 heads in LIDAS and heads within the repeated block in \LoopOp10-8×2-10 . Other heads are blue. Black outlines identify final top-2% heads, except the original L2H9 in LIDAS , which shares a red outline with L2H10. Vertical lines mark growth events (top), four-layer intervals in LIDAS and loop-block boundaries (bottom). Selection, plotting, and evaluation details in Section E.2 . FV heads show analogous repetition but emerge later (Appendix Figure 14 ).
Figure 6: Grown and looped models rely more on induction heads than the baseline for few-shot ICL and later-token prediction. Mean ablation in the 3B baseline, LIDAS , and \LoopOp10-8×2-10 . Solid lines ablate heads ranked by PS after excluding the top-2% FV heads. Similarly, dashed lines show FV-head ablations. (a) Change in few-shot ICL (F1) on 19 tasks. (b) Change in ICL-Score on FineWeb-Edu. Both changes are ablated minus the same model’s unablated score (lower is worse). At larger ablation fractions, induction-head ablation reduces both scores more in grown and looped models than in the baseline, although FV-head ablation hurts few-shot ICL more in all three models. (c) Position-wise NLL increase from each model’s unablated reference at 9% ablation over the 4096-token context (higher is worse). Curves are log-binned and smoothed. FV-ablation effects largely plateau, while induction-ablation effects grow with position and eventually exceed FV-ablation effects in LIDAS and \LoopOp10-8×2-10 . Additional ablation fractions appear in Appendix Figure 21 with experimental details in Appendix F .
Appendix figures & tables18 assets
Supplementary material from the paper’s appendix.
Appendix
360M
1.7B
3B
Dense layers L
32
24
36
Hidden / MLP width
960 / 2560
2048 / 8192
2048 / 11008
Query / key-value heads
15 / 5
32 / 32
16 / 4
Vocabulary size
49,152
49,152
128,256
Context length C
2048
2048
4096
Rotary base
104
104
5×104
Appendix
Table 1: Language-model training configurations.
Model
Strategy
9-task MC ↑
MMLU ↑
360M
Baseline
54.69
34.67
MIDAS
56.00
34.11
LIDAS
57.10
34.09
\LoopOp4-12×2-4
54.21
33.69
1.7B
Baseline
62.97
40.12
MIDAS
64.61
40.43
Appendix
Table 2: General benchmark performance across model sizes. Accuracy (%), with higher values better. 9-task MC is the unweighted mean over ARC-Easy, ARC-Challenge, BoolQ, CommonsenseQA, HellaSwag, OpenBookQA, PIQA, Social IQa, and WinoGrande; MMLU is the unweighted mean over its 57 subjects. Both columns use five-shot ranked-classification (RC) scoring of answer continuations with OLMES task-specific normalization and a 2048-token context limit. Here MC denotes multiple-choice benchmarks; scores use the RC formulation throughout, rather than answer-letter scoring. Model sizes denote dense reference configurations; loop layouts follow Section A.2 . Best values within each size are bold.
Figure 7: Token loss and context gains across model sizes. Columns show 360M, 1.7B, and 3B, using the position bins of Figure 1 . Top: final-checkpoint NLL relative to the same-size dense model in power-of-two position bins plotted at their upper endpoints, with matching-color dashed full-context means computed over raw positions. Bottom: the early-to-late loss reduction relative to dense ( Equation 4 ) for later bins ending at x∈{256,…,C} . All curves use 10,000 long FineWeb-Edu documents. Native contexts are C=2048/2048/4096 positions; the evaluation protocol is given in Section B.2 . Loop configurations are specified in Section A.3 .
Figure 8: Few-shot GSM8K performance across model sizes. Panels show mean strict exact-match accuracy for 360M, 1.7B, and 3B, with standard-error bands over three demonstration seeds. The shot range is k∈{0,1,2,4} for 360M and 1.7B and extends to eight shots for 3B. Details in Section B.3 for the evaluation and loop configurations are specified in Section A.3 .
Figure 9: Few-shot performance on additional multiple-choice and generation tasks. Columns show 360M, 1.7B, and 3B. The top row averages chance-centered accuracy over five multiple-choice tasks; the bottom row averages exact-match accuracy over three generation tasks, for k∈{0,1,2,4,8} . Tasks have equal weight within each subset; bands show standard error across three demonstration-seed subset averages. Individual tasks and normalization are given in Figures 10 and B.3 . Loop configurations are specified in Section A.3 .
Figure 10: Few-shot performance on individual multiple-choice and generation tasks. Columns show 360M, 1.7B, and 3B. The top five rows show chance-centered accuracy for multiple-choice tasks; the bottom three show exact-match accuracy for generation tasks. Bands show standard error across three demonstration seeds at k∈{0,1,2,4,8} . Subset averages and normalization are given in Figures 9 and B.3 . Loop configurations are specified in Section A.3 .
Figure 11: Word-grid tracing with frequent-word labels. Columns show 360M, 1.7B, and 3B; rows show 3×3 and 4×4 grids. Curves give full-vocabulary probability mass assigned to valid neighbors, with standard error over 200 trials. Black curves are one- and two-observation edge-memorization references. All panels share the probability scale and use base-2 logarithmic context axes. Both settings include the same selected methods at each model size. See Section B.4 for the protocol. Loop configurations are specified in Section A.3 .
Figure 12: RegBench across all configurations evaluated with four seeds. Columns show TVD (left) and DFA accuracy (right). Rows show full-test means (top), differences from dense by string index (middle), and differences by symbol index (bottom) for string 1 (dashed) and strings 2–19 (dotted). All available strings and symbols are retained; indices start at 1 and symbol positions use bins of five. The positive mean \LoopOp2×6 TVD differences at strings 18 and 19 concern only 102 (20.4%) and 54 (10.8%) of the 500 test sequences, respectively. Accuracy differences are in percentage points. Error bars and difference bands show standard errors across training seeds. All models have effective depth twelve; validation-selected widths are 128 except for \LoopOp2×6 (256). For additional details see Section C.3 .
Figure 13: Training trajectories across model sizes. Columns show 360M, 1.7B, and 3B. Rows report PS and FV averaged over their respective checkpoint-specific top-2% head sets, nineteen-task ten-shot F1, and the window-averaged ICL-Score ( Equation 3 ). Curves connect raw trained checkpoints, with regular spacing of 10,000 steps at 360M and 2000 steps at 1.7B and 3B, plus additional growth points. At 3B, all models use the shared regular grid. The first three rows use linear steps; the last uses logarithmic steps. Vertical rules mark growth and hollow markers show immediate post-growth evaluations. The downward marker in the 1.7B ICL-Score panel denotes MIDAS at −0.723 , below the plotted range. Loop configurations are specified in Section A.3 .
Figure 14: FV heads recur at the grown or looped block size but emerge later than induction heads. FV counterpart of Figure 5 for the 3B models. Top: trajectories of the twelve final top-2% FV heads per model, selected independently of PS. In LIDAS , these heads trace to eight ancestors; two, L3H13 and L3H15, account for six final heads (two and four, respectively; orange). Bottom: final FV scores for all heads; black outlines mark the selected set. Orange in \LoopOp10-8×2-10 marks selected heads in the repeated block; other heads are blue. Scores share a common scale that retains negative values. Axes, growth and block guides, lineage tracing, and smoothing follow Figure 5 and this section.
Figure 15: Trajectories of the final top-2% prefix-matching heads. Rows show 360M, 1.7B, and 3B; columns show Baseline, MIDAS , LIDAS , and Loop. Each panel follows the final 10, 16, or 12 heads ranked by PS, respectively, through their actual growth ancestry; sets are not reselected at intermediate checkpoints. Shared ancestral paths are drawn once, and vertical rules mark growth. Lines use the within-stage smoothing and boundary policy in Section E.2 ; step ticks marked k denote thousands. Scores share a common vertical range. There is no 3B MIDAS run. Loop configurations are specified in Section A.3 .
Figure 16: Final prefix-matching scores for every attention head. Rows show 360M, 1.7B, and 3B; columns show Baseline, MIDAS , LIDAS , and Loop. All panels share one PS color scale. Black outlines identify the final top-2% heads traced in Figure 15 . Layer and head indices are zero-based; loop applications occupy separate layer positions. The empty 3B MIDAS slot denotes an unavailable model. Loop configurations are specified in Section A.3 .
Figure 17: Trajectories of the final top-2% function-vector heads. Rows show 360M, 1.7B, and 3B; columns show Baseline, MIDAS , LIDAS , and Loop. Each panel follows the final 10, 16, or 12 heads ranked by FV, respectively, through their actual growth ancestry; sets are not reselected at intermediate checkpoints. Shared ancestral paths are drawn once, and vertical rules mark growth. Lines use the within-stage smoothing and boundary policy in Section E.2 ; step ticks marked k denote thousands. Scores share a common vertical range. There is no 3B MIDAS run. Loop configurations are specified in Section A.3 .
Figure 18: Final function-vector scores for every attention head. Rows show 360M, 1.7B, and 3B; columns show Baseline, MIDAS , LIDAS , and Loop. All panels share one FV color scale. Black outlines identify the final top-2% heads traced in Figure 17 . Layer and head indices are zero-based; loop applications occupy separate layer positions. Negative FV scores, where present, are retained. The empty 3B MIDAS slot denotes an unavailable model. Loop configurations are specified in Section A.3 .
Figure 19: Intermediate few-shot readouts across model sizes. Columns show 360M, 1.7B, and 3B; rows use the logit and tuned lenses. Each curve runs the nineteen-task, ten-shot F1 evaluation after the indicated fraction of layers. The tuned lens is trained separately for each frozen language model by next-token distillation ( Section E.3 ). These readouts measure accessibility of predictions, not the causal importance of a layer. Loop configurations are specified in Section A.3 .
Figure 20: Exclusive head ablations across model sizes. Columns show 360M, 1.7B, and 3B. Top: change in nineteen-task ten-shot F1. Bottom: change in windowed ICL-Score ( Equation 7 ). Solid lines ablate induction-ranked heads after excluding the top-2% FV heads; dashed lines reverse these roles. Markers distinguish models. Fractions are 1%, 3%, 5%, 7%, 9%, 15%, and 20%; zero denotes the corresponding unablated result. Loop configurations are specified in Section A.3 .
Figure 21: Position-wise head ablation at higher fractions. Columns show 360M, 1.7B, and 3B; row pairs show exclusive induction and FV ablation at 7%, 9%, 15%, and 20%. Curves give the mean NLL increase on 10,000 FineWeb-Edu documents, with log-binned positions on a base-2 axis. Each induction/FV pair within a size shares an uncropped vertical range. MIDAS is unavailable at 3B. See Appendix F . Loop configurations are specified in Section A.3 .
Figure 22: Head ablation across model sizes. Row pairs show 360M, 1.7B, and 3B, each with absolute few-shot F1 above the windowed ICL-Score ( Equation 3 ); columns show dense, MIDAS , LIDAS , and looped models. Dashed rankings exclude the other measure’s top-2% heads. Gray bands show the mean ± standard error over random head sets: three for F1 and sixteen for the ICL-Score at every model size. Dotted lines give unablated values. Fractions count applications in looped models. Panel scales vary; the empty 3B MIDAS column reflects the missing model. See Appendix F . Loop configurations are specified in Section A.3 .
Looping, reusing a block of layers across depth, and depth growing, training shallow-to-deep models by duplicating middle layers, have both been linked to stronger reasoning, but their relationship remains unclear. We provide a mechanistic unification: looped and depth-grown models exhibit convergent depth-wise signatures, including increased reliance on late layers and recurring patterns aligned with the looped or grown block. These shared signatures support the view that their gains stem from a common form of iterative computation. Building on this connection, we show that the two techniques are adaptable and composable: applying inference-time looping to the middle blocks of a depth-grown model improves accuracy on some reasoning primitives by up to 2×, despite the model never being trained to loop. Both approaches also adapt better than the baseline when given more in-context examples or additional supervised fine-tuning data. Additionally, depth-grown models achieve the largest reasoning gains when using higher-quality, math-heavy cooldown mixtures, which can be further boosted by adapting a middle block to loop. Overall, our results position depth growth and looping as complementary, practical methods for inducing and scaling iterative computation to improve reasoning.
Ferdinand Kapl, Emmanouil Angelis, Kaitlin Maile +2
Technical University of Munich · Helmholtz AI, Munich · Google, Paradigms of Intelligence Team
In-context learning (ICL) allows language models to adapt their behavior from instructions and demonstrations in a prompt, but accuracy alone does not reveal how contextual information controls their predictions. We study the acquisition and trainability of ICL with a controlled diagnostic that separates copying a nearby demonstration from integrating information across the prompt. Across 17 models, about a third prefer the copied answer, a few are undecided, and the rest prefer the rule-consistent answer, which goes with higher accuracy and with control by the requested value rather than the closest conflicting example. Checkpoint analyses show that models begin prioritizing task-relevant prompt tokens over conflicting examples long before their accuracy improves beyond chance level. In matched training experiments, synthetic documents that define endings and then query them shift models toward compositional responses on the controlled diagnostic, whereas fine-tuning on definition exercises built from rewritten real documents improves accuracy when prompts reassign label meanings. The results suggest that ICL is not acquired as a single, uniform capability. Instead, prompt control, answer preference, construction without candidates, and transfer can develop along distinct trajectories.
Chain-of-thought (CoT) reasoning has become a widely used mechanism for eliciting multi-step reasoning in large language models by generating intermediate reasoning steps at inference time. Yet the scaling behavior of generalization with CoT depth remains poorly understood. To address this question, we study a theoretically solvable model of CoT for in-context weight prediction in linear regression, where test-time reasoning is represented as an iterative refinement of the weight-parameter estimate. Using tools from random matrix theory under high-dimensional asymptotics, we derive an exact formula for the generalization error as a function of reasoning depth, pretraining data amount, and context length. Our analysis reveals a sharp phase transition separating exponential and polynomial improvement, saturation, and overthinking, and characterizes how the optimal reasoning depth scales. We further show that deeper reasoning is most effective with sufficiently rich pretraining and in-context information, whereas limited pretraining or context makes longer reasoning prone to error amplification or saturation. We also validate these predictions through experiments on fully learned linear attention and softmax attention models. Our results provide a unified theoretical account of how test-time CoT depth affects generalization.
Kaito Takanami, Cengiz Pehlevan
Department of Physics, Graduate School of Science, The University of Tokyo · John A. Paulson School of Engineering and Applied Sciences, Harvard University · Kempner Institute for the Study of Natural and Artificial Intelligence, Harvard University +1