Research map / Large language models / Instruction and fine-tuning
On the Inductive Bias of Stacking Towards Improving Reasoning
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes MIDAS, a middle-block gradual stacking variant that speeds up UL2 language model pretraining by up to ~40% and, at matched perplexity, improves reasoning-oriented downstream tasks, supported by synthetic "reasoning primitives" and a looped-model connection.
Problem
Gradual stacking speeds up training, but its effect on generalization and the inductive biases it induces in language model pretraining is largely unexplored.
Approach
Instead of stacking the last block of layers as in gradual stacking, MIDAS duplicates the middle block when growing depth, which yields layer similarity structure resembling looped (ALBERT-style) models; this gives faster training and, at matched validation perplexity, better reasoning-oriented downstream performance.
Key results
- MIDAS with the Prop-2 schedule matches baseline validation log perplexity (2.009 vs 1.996 at 1B) while being 24% faster, and gets 26.8% average downstream accuracy vs 24.0% for baseline.
- At 8B, MIDAS reaches 43.1% on math word problems vs 34.9% for baseline, with similar loss (1.844 vs 1.841).
- On variable assignment depth-1 code primitive after fine-tuning on 64 examples, MIDAS achieves 68.54% vs 43.75% for baseline.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Stack the middle block rather than the last block | UL2 1B with schedule 2 and block size 4 · UL2 pretraining | All Tasks Average: 20.4 -> 26.8 | Helped |
| Stack the middle block rather than the last block | UL2 2B with schedule 2 · UL2 pretraining | All Tasks Average: 25.9 -> 32.9 | Helped |
| Stack the middle block rather than the last block | UL2 pretraining at matching speedup levels · Validation log perplexity | log perplexity | Helped |
| Use schedule 2 instead of schedule 1 (ablation) | MIDAS, UL2 1B, block size 4 · UL2 pretraining | All Tasks Average: 26.7 -> 26.8 | Mixed |
| Use block size 3 instead of block size 4 (ablation) | MIDAS, UL2 1B, schedule 1 · UL2 pretraining | All Tasks Average: 26.7 -> 25.6 | Hurt |
| Use block size 3 instead of block size 4 (ablation) | MIDAS, UL2 1B, schedule 2 · UL2 pretraining | All Tasks Average: 26.8 -> 27.5 | Helped |
| Use schedule 3 instead of schedule 2 (ablation) | MIDAS, UL2 1B, block size 4 · UL2 pretraining | All Tasks Average: 26.8 -> 28.3 | Helped |
| Pretrain with MIDAS instead of standard training | UL2 1B · GSM8K test, fine-tuned | accuracy: 5.3% to 10.4% | Helped |
| Pretrain with MIDAS instead of standard training | UL2 8B · GSM8K test, fine-tuned | accuracy: 12.3% to 15.2% | Helped |
| Pretrain with MIDAS instead of standard training | UL2 1B · Reasoning primitives, 5-shot copying random-letter words | accuracy: 14.9 -> 24.3 | Helped |
| Pretrain with MIDAS instead of standard training | UL2 1B · Reasoning primitives, 5-shot variable assignment depth 0 (math) | accuracy: 50.1 -> 92.8 | Helped |
| Pretrain with MIDAS instead of standard training | UL2 1B · Reasoning primitives, 5-shot variable assignment depth 1 (basic) | accuracy: 21.9 -> 20.6 | Hurt |
Surprising findings
- Models with equal or slightly worse validation perplexity show substantially better downstream reasoning and math performance, suggesting an inductive bias from stacking.
- The gap between MIDAS and baseline widens when a calculator is used for math word problems, which the authors attribute to reasoning vs memorization interplay.
Limitations
- The inductive bias is conjectured to stem from connection to looped models but not established causally.
- Downstream gains are measured on a small set of tasks and absolute accuracies at 1B-8B are low (e.g., GSM8K around 5-15%).
Experimental setup
Datasets: C4 Wikipedia Arxiv GitHub TriviaQA TydiQA Natural Questions WebQuestions SquadV2 DROP QuAC CoQA
Benchmarks: TriviaQA TydiQA Natural Questions WebQuestions SquadV2 DROP QuAC CoQA ASDiv MAWPS SVAMP GSM8K
Metrics: validation log perplexity accuracy speedup
Baselines: Baseline standard training gradual stacking
Builds on: gradual stacking progressive stacking UL2 ALBERT Universal Transformers
Closest papers by meaning
- An Analysis for Reasoning Bias of Language Models with Small Initialization (ICML 2025 · similarity 0.90)
- Reasoning with Latent Thoughts: On the Power of Looped Transformers (ICLR 2025 · similarity 0.89)
- Rethinking Fine-Tuning when Scaling Test-Time Compute: Limiting Confidence Improves Mathematical Reasoning (NeurIPS 2025 · similarity 0.88)
- Phenomenal Yet Puzzling: Testing Inductive Reasoning Capabilities of Language Models with Hypothesis Refinement (ICLR 2024 · similarity 0.88)
- MIRAGE: Evaluating and Explaining Inductive Reasoning Process in Language Models (ICLR 2025 · similarity 0.88)
- Selection, Reflection and Self-Refinement: Revisit Reasoning Tasks via a Causal Lens (ICLR 2026 · similarity 0.88)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.