PaperlayerView exact revision on arXiv

Source-linked AI summary

How Depth Growth and Looping Shape In-Context Learning in Language Models

Ferdinand Kapl, Kaitlin Maile, Stefan Bauer, Johannes von Oswald

arXiv:2610.31421v1cs.CLcs.AIcs.LG

TL;DR

The paper asks whether depth growth and looping promote learning from context, and how this develops when memorization can compete. It compares pretrained language models and controlled tasks, including analyses of attention heads. Grown and looped models show stronger context use, earlier ICL acquisition, and stronger retention, while growth propagates early induction heads across layers.

  • Problem

    How depth growth and looping affect context-based learning when memorization is possible, the persistence of that learning, and its attention-head mechanisms remains less understood.

    Source: §1 · pp. 1–2

  • Method

    The authors compare dense, depth-grown, and looped language models on language-model and controlled tasks, measuring context use and tracing and intervening on attention heads.

    Source: §1 · p. 2, §3.1 · p. 4

  • Results

    Grown and looped models show stronger context use, earlier ICL acquisition, and stronger ICL retention than dense references, while depth growth propagates early induction heads across layers.

    Source: §1 · p. 2, §4.1 · p. 7, §5 · p. 9

  • Takeaways & Limitations

    The experiments provide evidence that depth reuse favors learning from context and shapes the attention mechanisms supporting it.

    Source: §1 · p. 2

Abstract

from arXiv · show

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.

1 INTRODUCTION

The paper investigates whether depth growth and looping favor learning from context, how this develops alongside in-weights learning, and which attention mechanisms support it. Its experiments find stronger context use and earlier acquisition or retention of ICL in grown and looped models, alongside a role for induction heads.

1 INTRODUCTION

  • Depth growth duplicates blocks of a progressively expanded network, whereas looping repeatedly applies a shared block within each forward pass.The two approaches can produce related iterative-computation patterns despite different parameter sharing.Source: §1 · p. 1
  • Prior evidence links depth reuse with ICL, but how it affects context-based solutions when memorization is possible, their persistence, and attention-head mechanisms remains less understood.Improved reasoning alone does not establish ICL, since predictions may draw on weights, context, or both, and later-token prediction differs from learning from demonstrations.Source: §1 · pp. 1–2
  • The study compares dense, grown, and looped language models up to 3B parameters, using token loss, few-shot learning, formal-language learning, controlled classification tasks, and attention-head analyses.The experiments include tasks that separate learning from context from learning fixed associations in the weights.Source: §1 · p. 2
  • Grown and looped models show stronger context use than dense baselines across token-loss, GSM8K, and unseen-language evaluations.The grown model also attains lower absolute late-token loss in the main 3B language-model comparison.Source: §1 · p. 2
  • In controlled Omniglot tasks, grown and looped models acquire ICL earlier when labels vary by sequence and retain higher ICL accuracy when fixed associations can compete.In the latter setting, both show lower in-weights learning accuracy than the full-depth dense reference.Source: §1 · p. 2
  • Depth growth propagates early induction heads across layers, while grown and looped models rely more on induction-ranked heads for few-shot performance and later-token prediction.The interventions also find function-vector heads remain important.Source: §1 · p. 2
  • The experiments suggest that constructing depth through parameter reuse shapes both context use and the attention mechanisms supporting it.Source: §1 · p. 2

2 RELATED WORK

Prior work studies layer growth and parameter sharing as ways to build depth, alongside research on in-context learning and attention mechanisms. This paper extends that work by examining ICL acquisition, retention, and mechanisms across pretrained models and controlled tasks.

Grown Models

  • Growth approaches include layer copying, progressive stacking, learned transformations, and function-preserving expansions; growth operators have also been studied for their effects on model capabilities.MIDAS duplicates a central layer block, while LIDAS duplicates layers around the network center independently of existing block boundaries.Source: §2 · p. 2

Looped Models

  • Cross-layer parameter sharing provides computational depth with fewer weights, and related work studies recurrent stacking, layer-sharing designs, adapters, and mixture-of-experts extensions.Recurrence has also been explored for length generalization, latent reasoning, inference-time scaling, and iterative refinement.Source: §2 · pp. 2–3

In-Context Learning

  • Prior ICL work examines learning functions from examples, context-dependent meanings, and competition with in-weights associations; induction and function-vector heads are linked to pattern copying and task representations.Source: §2 · p. 3
  • Kapl et al. (2026b) report shared depth-wise patterns and improved adaptation from demonstrations, while this paper extends the evidence to ICL acquisition, retention, and mechanisms.The study tests context use when memorization is possible and assesses induction- and function-vector-head contributions to few-shot learning and token prediction.Source: §2 · p. 3

3 IN-CONTEXT LEARNING OF DEPTH-GROWN AND LOOPED MODELS

This section compares dense, depth-grown, and looped models to examine how they use context across language modeling and controlled tasks. It reports greater context use, stronger ICL retention when fixed associations are also learnable, and earlier ICL acquisition when they are not.

3.1 MODELS AND EVALUATION MEASURES

  • Dense models train independently parameterized layers; depth-grown models duplicate and independently train blocks, while looped models repeatedly apply a tied block.The main 3B comparison uses a 36-layer dense reference, a grown model reaching that depth, and a Loop 10-8×2-10 model with 36 layer applications and 28 unique layers; models share data and optimizer-step budgets and match final inference FLOPs.Source: §3.1 · pp. 3–4

3.2 CONTEXT USE IN LANGUAGE MODELS

  • Both grown and looped models gain more from context than the dense baseline; LIDAS reaches lower late-token NLL, while the looped model retains higher NLL.On GSM8K, both models exceed the baseline, and LIDAS’s advantage widens with additional demonstrations to around +7%.Source: Figure 1 · p. 4, §3.2 · pp. 4–5

3.1 MODELS AND EVALUATION MEASURES

  • The evaluation uses next-token NLL, an ICL-Score measuring early-to-late mean-NLL reduction, and few-shot benchmark accuracy or token-overlap F1.The ICL-Score compares mean NLL at positions 40–60 with positions 490–510; higher scores indicate greater improvement.Source: §3.1 · p. 4

3.3 IN-CONTEXT LEARNING LANGUAGES: REGBENCH DATASET

  • On RegBench, grown and looped models achieve lower TVD than the dense baseline on unseen formal languages, with their advantage emerging within the first string and peaking around the fourth.Within the first string, both predict symbols better after roughly 15 symbols; their advantage then remains approximately constant with more symbols.Source: §3.3 · p. 5, Figure 2 · p. 5

3.4 IN-CONTEXT AND IN-WEIGHTS LEARNING: OMNIGLOT

  • When fixed class–label associations are learnable, grown and looped models retain higher ICL accuracy but have lower IWL accuracy than the dense 12-layer model.The looped model has significantly lower IWL accuracy, while the grown model has ICL and IWL accuracies similar to its four-layer starting reference.Source: §3.4 · p. 6, Figure 3 · p. 6
  • When ICL is the only correct solution, looped, grown, and four-layer dense models acquire it earlier than the dense 12-layer model, though all four nearly reach 100% accuracy by 1M steps.This setting differentiates the timing and consistency of ICL emergence rather than its eventual attainment.Source: §3.4 · pp. 6–7

4 GROWING AND LOOPING INCREASE RELIANCE ON INDUCTION HEADS

Depth growth propagates early induction behavior across layers, while grown and looped models rely more on induction-ranked heads for few-shot performance and later-token prediction.

4 GROWING AND LOOPING INCREASE RELIANCE ON INDUCTION HEADS

  • The analysis tracks induction behavior during training, then ablates heads to test their contributions to context use.Prefix-matching scores characterize induction behavior; exclusive induction rankings exclude the top 2% of function-vector heads before ablation.Source: §4.1 · p. 7, §4 · p. 7, §4.2 · p. 8

4.1 DEPTH GROWTH PROPAGATES EARLY INDUCTION HEADS

  • LIDAS develops higher mean top-2% prefix-matching scores than the dense baseline and finishes higher, while the looped model finishes close to baseline.Growth temporarily disrupts prefix matching, followed by recovery during further training.Source: §4.1 · p. 7, Figure 4 · p. 7
  • 10 of the 12 final top-2% LIDAS induction heads descend from two ancestors that already had high prefix-matching scores before growth.The twelve final heads trace to four initial ancestors; the other two originated from initially low-scoring heads whose induction behavior developed later.Source: §4.1 · p. 7, Figure 5 · p. 8
  • Strong induction heads recur at four-layer intervals in LIDAS and across the shared block’s applications in Loop 10-8×2-10, unlike the dense baseline.Grown copies train independently after duplication, whereas looped heads are tied across applications.Source: §4.1 · p. 7, Figure 5 · p. 8
  • Function-vector heads also recur at grown or looped block size, but emerge later than induction heads.In LIDAS, two ancestors give rise to six of the twelve final top-2% function-vector heads.Source: §4.1 · p. 7, Figure 5 · p. 8

4.2 INDUCTION HEADS CONTRIBUTE MORE TO CONTEXT USE

  • At larger ablation fractions, induction-head ablation reduces few-shot performance and ICL-Score more in grown and looped models than in the baseline.Function-vector heads still cause greater few-shot ICL degradation in all three models.Source: §4.2 · p. 8, Figure 6 · p. 9
  • At 9% ablation, induction-head effects on NLL grow with token position and eventually exceed function-vector effects in LIDAS and Loop 10-8×2-10.Function-vector ablation hurts early-token prediction more and then largely plateaus; the baseline retains greater function-vector effects across most positions.Source: §4.2 · p. 8, Figure 6 · p. 9
  • Both head types support context use, with grown and looped models relying more on induction heads.Function-vector heads remain more important for few-shot performance in all three models.Source: §4.2 · pp. 8–9

5 CONCLUSION

The conclusion reports that depth-grown and looped models favor learning from context across language-model and controlled tasks. It links this pattern to early induction-head propagation and greater reliance on induction heads in grown and looped models.

5 CONCLUSION

  • Grown and looped models rely more on induction-ranked heads for few-shot performance and later-token prediction, while function-vector heads remain more important for few-shot performance.The conclusion also reports that these models achieve larger early-to-late token-loss reductions and higher GSM8K accuracy than the dense baseline, and predict unseen formal languages more accurately.Source: §5 · p. 9
  • Most of LIDAS’s strongest final induction heads descend from a small set of early ancestors, producing repeated patterns across depth.These grown layers train independently after copying, whereas looped models repeatedly apply shared layers.Source: §5 · p. 9
  • In controlled classification tasks, grown and looped models acquire ICL earlier and retain higher ICL accuracy when fixed associations can also be learned in the weights.They retain higher ICL accuracy while attaining lower in-weights-learning accuracy than the full-depth dense reference.Source: §5 · p. 9

A.1 TRAINING SETUP

This appendix describes model training configurations, growth and looping procedures, and checkpoint comparison scope, then introduces the language-model context-use evaluations.

A.1 TRAINING SETUP

  • The models are autoregressive decoder-only Transformers trained with next-token cross-entropy; sizes refer to dense references, while looped models use fewer unique parameters and grown models reach reference depth during training.Within each size, methods share data and optimization settings, but sizes differ in training budget; the 3B configuration also changes data, tokenizer, and context length.Source: §A.1 · p. 14

A.2 GROWING AND LOOPING

  • Grown models start with four layers and repeatedly duplicate four contiguous layers until reaching reference depth; copied weights and optimizer moments become independently trainable immediately.LIDAS selects the middle four layers without respecting prior block boundaries; layer indices identify copied parameter sources, not persistent sharing.Source: §A.2 · pp. 14–15
  • The Prop-1 schedule allocates training steps in proportion to stage depth, with learning-rate decay appended to the final full-depth stage.For 3B, 2,000 warm-up steps are excluded from the proportional split and prepended to the first stage; only LIDAS is evaluated at 3B.Source: §A.2 · p. 15
  • Loop configurations specify prelude layers, a repeated shared block, and coda layers; effective depth counts repeated applications, while unique parameters count each shared layer once.Loop 10-8×2-10 has 36 layer applications and 28 unique layers; comparisons match effective final depth and training tokens within a size.Source: §A.2 · p. 15

A.3 CHECKPOINTS AND COMPARISON SCOPE

  • Checkpoint comparisons distinguish trained pre-growth checkpoints from immediate post-growth evaluations, which are shown as hollow markers at the same training step.Vertical rules mark growth events; final evaluations use step 200,000 for 360M/1.7B and step 42,400 for 3B.Source: §A.3 · p. 15
  • Each language-model comparison arm uses one training seed, so evaluation bands do not estimate variation across training runs.There is no 3B MIDAS run, and the configuration differences and 360M loop initialization exception apply throughout.Source: §A.3 · p. 15

B CONTEXT USE IN LANGUAGE MODELS

  • The following section reports evaluations across sizes for general benchmarks, token loss, downstream few-shot performance, and graph tracing.Shared training configurations and size-specific loop definitions are in Appendix A.Source: B CONTEXT USE IN LANGUAGE MODELS · p. 15

B.1 GENERAL BENCHMARK EVALUATIONS

The section specifies final-checkpoint benchmark evaluation and the token-loss measures used to compare context-dependent prediction with dense baselines.

B.1 GENERAL BENCHMARK EVALUATIONS

  • Final benchmark scores use OLMES five-shot ranked-classification scoring, a 2048-token limit, and final checkpoints at step 200,000 for 360M/1.7B and 42,400 for 3B.The nine-task core suite and MMLU are evaluated for every model.Source: §B.1 · p. 16, Table 2 · p. 16

B.2 TOKEN LOSS AND THE ICL SCORE

  • Token-loss evaluation uses 10,000 long documents, scoring C predictions per document after truncating each to C + 1 tokens.The 2048- and 4096-position evaluations use each model’s native tokenizer and separately selected document sets.Source: Figure 1 · p. 16
  • Position-dependent loss compares each model’s power-of-two-binned NLL with the same-size dense baseline; negative differences indicate better prediction at those positions.The bins are plotted at their upper endpoints on a base-2 logarithmic axis.Source: §B.2 · pp. 16–17, Figure 7 · p. 18
  • The context-use score compares early and later mean NLL, but a change can arise from either window and need not mean absolute late-position prediction improved.The reported score uses positions 40–60 as the early window and later-position bins for the other window.Source: §B.2 · p. 17

B.3 FEW-SHOT SCALING ON DOWNSTREAM TASKS

This appendix details few-shot evaluations across downstream tasks and graph tracing, followed by the RegBench setup and its model-selection and position-wise scoring procedures.

B.3 FEW-SHOT SCALING ON DOWNSTREAM TASKS

  • Few-shot evaluations average results over three demonstration seeds, cap shot ranges to avoid truncation, and use an input limit equal to the training context.GSM8K uses strict exact match; additional task sets report multiple-choice and generation averages across k ∈ {0, 1, 2, 4, 8}.Source: §B.3 · p. 19

B.4 IN-CONTEXT GRAPH TRACING

  • Graph tracing scores the full-vocabulary probability mass assigned to valid neighboring node labels after each node in a random-walk sequence.Frequent-word experiments cover 3 × 3 and 4 × 4 grids and compare against one- and two-observation edge-memorization references.Source: §B.4 · p. 22, Figure 11 · p. 22

C.1 TASK AND DATA

  • RegBench uses DFA-generated sequences split by sampled automaton, with 40,000 training, 500 validation, and 500 test sequences shared across models.Training and validation cross-entropy mask separator targets and padding.Source: §C.1 · p. 23

C.2 MODELS, OPTIMIZATION, AND SELECTION

  • The RegBench comparison includes dense 12-layer, looped 2×6, and three grown configurations; selection uses mean validation loss, not test performance.Comparisons match data and optimizer steps but not width or compute.Source: §C.2 · p. 23

C.3 METRICS AND POSITION DECOMPOSITION

  • RegBench reports TVD and DFA accuracy, with negative TVD differences and positive accuracy differences favoring alternatives to dense.Appendix position views retain all available strings and symbols; position averages may differ from full-test statistics because aggregations differ.Source: §C.3 · pp. 23–24

C.4 RESULTS ACROSS CONFIGURATIONS

The appendix presents RegBench comparisons across configurations and describes the data and sequence-model architecture used for the Omniglot experiments.

C.4 RESULTS ACROSS CONFIGURATIONS

  • Figure 12 compares five RegBench configurations using TVD and DFA accuracy, showing full-test means and differences from dense by string and symbol index.The positive Loop 2×6 TVD differences at strings 18 and 19 concern 20.4% and 10.8% of test sequences, respectively.Source: Figure 12 · p. 24

D.1 DATA, SEQUENCE FORMAT, AND ARCHITECTURE

  • Omniglot sequences interleave eight image–label context pairs with a query image, and training loss is applied only at the query position.The data comprise 1,623 character classes with twenty grayscale exemplars each, with ten classes reserved for held-out evaluation.Source: §D.1 · p. 25
  • The Omniglot comparison uses dense depths 4 and 12, a Loop 4×3 model, and a LIDAS-grown model that expands from four to twelve layers.The grown model and four-layer reference share exactly the same pre-growth trajectory; the looped model repeats Transformer blocks only.Source: §D.1 · p. 25

D.2 TRAINING DISTRIBUTIONS AND OPTIMIZATION

The experiments contrast fixed-label sequences, where context and stored associations can both solve queries, with ICL-only sequences that remove persistent class-to-answer mappings.

D.2 TRAINING DISTRIBUTIONS AND OPTIMIZATION

  • In the fixed-label setting, bursty sequences include repeated examples of two classes, and the query label can be recovered from context or the fixed training association.The setup samples two burst classes, adds two other training items, then shuffles the eight context items.Source: §D.2 · p. 25
  • ICL-only training randomly reassigns labels 0 and 1 to two classes in each four-shot, two-way sequence, leaving no persistent class-to-answer mapping.The architecture, optimizer, sequence length, and 1,623-output head remain unchanged.Source: §D.2 · p. 25

D.3 EVALUATION: ICL AND IWL

  • The ICL probe tests held-out classes with new per-sequence labels, while the IWL probe tests training classes absent from context using their fixed labels.ICL accuracy restricts the argmax to outputs 0 and 1; IWL uses the unrestricted 1,623-output argmax and is only used in the fixed-label setting.Source: §D.3 · pp. 25–26
  • Evaluation reuses fixed probe sequences across checkpoints and averages curves across four fixed-label seeds or eight ICL-only seeds.Displayed curves are smoothed, but endpoint statistics and emergence thresholds use unsmoothed evaluations.Source: §D.3 · p. 26

E.1 MEASURES AND AGGREGATE TRAINING TRAJECTORIES

The section defines head, few-shot, and token-loss measurements, then describes how it tracks selected heads through growth and reads out predictions at intermediate depths.

E.1 MEASURES AND AGGREGATE TRAINING TRAJECTORIES

  • PS measures attention to the earlier matching token, while FV uses a task-mean head-output patch to measure its effect on the correct answer probability.The fixed ten-shot evaluation reports token-overlap F1 across nineteen held-out tasks; PS and FV summaries average checkpoint-specific top-2% head sets.Source: §E.1 · pp. 26–27

E.2 FINAL-HEAD IDENTITIES, LINEAGES, AND LOCATIONS

  • Final top-2% PS- and FV-ranked heads are traced backward through actual duplication maps, with shared ancestors drawn once and loop applications treated separately.The selections are fixed identities across training rather than sets reselected at intermediate checkpoints; no seed-head pair is assumed to recur across sizes.Source: §E.2 · p. 29

E.3 PER-LAYER READOUTS

  • Per-layer logit and tuned-lens readouts measure how accessible few-shot predictions are at each exit depth, not which layers causally compute them.The tuned lens learns a residual translator for each frozen model; the full-depth exit uses the original readout.Source: §E.3 · p. 34, Figure 19 · p. 34

F.1 HEAD SELECTION AND MEAN ABLATION

The section specifies exclusive head rankings and mean-activation ablations, then measures their effects on few-shot F1 and token prediction across positions.

F.1 HEAD SELECTION AND MEAN ABLATION

  • Exclusive induction rankings exclude top-2% FV heads, and exclusive FV rankings exclude induction heads before selected head applications are mean-ablated.Ablation fractions are counted against the full set before exclusion; looped models intervene on applications, not unique parameters.Source: §F.1 · pp. 34–35

F.2 MEASUREMENTS AND PLOTTED SUMMARIES

  • Few-shot effects are measured as changes in nineteen-task ten-shot F1, using clean task-mean head activations as replacements during prompt and generation passes.Figure 22 also reports absolute F1 alongside ablation results.Source: §F.2 · p. 35, Figure 20 · p. 36, Figure 22 · p. 38
  • Token-loss effects use NLL increases on 10,000 FineWeb-Edu documents, with position-wise curves and windowed ICL-Score summaries.Position curves use logarithmic bins; the plotted figure compares exclusive induction and FV ablations at higher fractions.Source: §F.2 · p. 35, Figure 21 · p. 37

Related papers

Loading 2610.31421v1…