How Depth Growth and Looping Shape In-Context Learning in Language Models
Ferdinand Kapl, Kaitlin Maile, Stefan Bauer, Johannes von Oswald
cs.CL, cs.AI, cs.LG
2026-10-22
On matched 3B runs (100B tokens), grown and looped models learn from context earlier than a same-depth dense model and rely more on induction heads. LIDAS is about +7 on GSM8K.
Two shortcuts add depth without paying for a full stack of independent layers. Depth growing trains a shallow network, then duplicates a middle block, copying weights and Adam moments. The copies train separately after that. Looping runs one shared block several times inside a single forward pass, with the weights tied throughout. Both have been used for reasoning, and earlier results suggested they are also better at in-context learning (ICL): using the current sequence without a parameter update.
A better reasoning score does not show that the model is reading the context. The next token can come from a fixed association stored in the weights, or from the sequence in front of it. A drop in loss on later tokens is a different measurement from answering a new question after a few demonstrations. On matched data and optimizer steps, the question is whether growing and looping still pick the context solution when memorization would also work, whether that choice lasts, and which attention heads carry it.
The main runs are decoder-only Transformers trained from scratch, with dense references from 360M to 3B. The 3B setup follows SmolLM3: 36 layers, the same 100B tokens in the same order, and matched inference FLOPs at the end. Growth uses LIDAS, which copies the middle of the network and ignores earlier block boundaries. Block size is 4, so the model expands from 4 layers to 36 without resetting the learning-rate schedule. Early training is shallower, which cuts estimated training FLOPs by about 26% at 3B and 22% at 360M and 1.7B.
A loop is written Loop p-b×r-c. The 3B loop is Loop 10-8×2-10: 36 layer applications, 28 unique layers, about 2.45B parameters.
Three probes pull the behaviors apart. FineWeb-Edu gives position-wise negative log-likelihood. The ICL-Score subtracts mean loss at positions 490 to 510 from mean loss at positions 40 to 60. GSM8K uses 0 to 8 demonstrations and strict exact match on the extracted answer. RegBench concatenates strings from an unseen finite automaton and scores total variation distance (TVD) between the predicted next-symbol distribution and the uniform distribution over legal continuations. Omniglot with fixed labels lets either context or a stored class-label map answer the query. The ICL probe switches to new classes and reassigns labels every sequence. The IWL probe uses training classes whose query class is missing from context.
Prefix-matching score (PS) marks induction heads: in a random token sequence played twice, how much attention lands on the token that followed the earlier copy of the current prefix. Function-vector heads are scored by swapping a head's output for its mean activation on a clean task and checking whether the right answer returns after demonstration labels are shuffled. Ablations replace outputs with that mean, after dropping the top 2% of the other head type, then remove 1% to 20% of heads.
On 10,000 long documents at 3B, grown and looped models are worse at the start. The gap is largest at token index 8 or earlier, then their loss falls faster than the dense baseline. LIDAS catches the baseline near position 1000 and is lower after that, including the average over all positions. The loop has fewer unique parameters and still has higher loss at the end of the sequence.
On GSM8K the loop beats the baseline at every shot count, but the baseline nearly closes the gap at 8 shots. LIDAS's lead widens as demonstrations are added and finishes about 7 points ahead. On five multiple-choice tasks and three generation tasks at 3B, LIDAS is ahead on both, and the loop's edge is mostly in generation.
| Scale | Model | 9-task | MMLU |
| 3B | Dense | 61.33 | 39.15 |
| 3B | LIDAS | 62.89 | 39.28 |
| 3B | Loop | 60.53 | 38.32 |
| 1.7B | Dense | 62.97 | 40.12 |
| 1.7B | LIDAS | 64.69 | 40.86 |
| 1.7B | Loop | 63.25 | 39.89 |
| 360M | LIDAS | 57.10 | 34.09 |
| 360M | Dense | 54.69 | 34.67 |
MMLU barely moves. At 360M, LIDAS is slightly worse, 34.09 against 34.67. On the nine-task suite LIDAS is higher. The loop is often no better than the dense model.
On RegBench every model applies twelve layers. Grown and looped models reach lower test TVD. The gap peaks near the fourth string. Inside the first string it shows up after about 15 symbols. The loop has width 256 against 128 for the others. Hyperparameters were chosen on validation loss, so width and compute are not matched.
With fixed Omniglot labels, ICL rises and then falls while IWL climbs. At the end of the shared run, grown and looped models keep higher ICL and lower IWL than the dense 12-layer model. The grown model stays close to the 4-layer network it started from. When labels are reassigned every sequence, the loop, the grown model, and the 4-layer dense model all develop ICL earlier than the 12-layer dense model. By 1M steps all four are near 100%.
During 3B training, mean PS over the top 2% of heads stays higher for LIDAS than for the dense baseline. The loop finishes near the baseline. Each duplication briefly disrupts prefix matching. The twelve strongest final heads trace back to four ancestors in the initial 4-layer model. L2H9 and L2H10 are already strong before the first copy and account for ten of those twelve. In the loop, high-PS heads appear in both passes of the shared 8-layer block, with the weights still tied.
Ablations line up across the three 3B models. Removing function-vector heads hurts few-shot F1 on 19 tasks more. Removing induction heads hurts the token-loss ICL-Score more. At larger fractions, induction-head ablation hurts grown and looped models more than the dense baseline. At 9%, function-vector ablation mostly hurts early positions, while induction-head damage grows with position and overtakes it later in LIDAS and the loop. In the dense baseline, function-vector heads remain more important for later tokens.
Building depth by reusing parameters pushes the model toward context, and it makes that solution appear earlier and last longer. Growth also spends less on training: about 26% fewer training FLOPs at 3B, with inference FLOPs still matched to a 36-layer dense model. For few-shot math and later-token loss, LIDAS is the run to follow. The loop buys the same effective depth with fewer unique parameters. It can win on formal languages and few-shot generation, and it often loses on absolute loss and multiple choice.
Layer copying mostly stamps a few early prefix-matching heads into later layers. Few-shot tasks still depend more on function-vector heads. Later-token prediction shifts toward the induction heads.
Each language-model arm has one training seed. The 360M loop also skips an initialization rescaling, so that comparison is not purely architectural. On RegBench the loop is width 256 against 128, and compute is not matched. In the Omniglot setting that only measures when ICL appears, all four models are near 100% by 1M steps. Keeping ICL means weaker IWL. If the job is to store stable labels, the bias is a cost.
The GSM8K gap of about 7 points sits on low absolute accuracy. The 3B panel's axis tops out at 18%. LIDAS gains 1.56 points on the nine-task suite and 0.13 on MMLU, while the loop is below the dense baseline on both. The 3B models see 100B tokens. The 360M and 1.7B runs see about 210B and 419B, with different tokenizers and context lengths.
Copied early context computations being harder for memorization to displace is one explanation, and no intervention tests it. How function-vector ablation compares with the baseline also changes across sizes. Induction heads contribute less to few-shot accuracy than to later-token loss.