Research map / Large language models / Reasoning and chain of thought

Understanding the Role of Training Data in Test-Time Scaling

ICLR 2026 · Large language models › Reasoning and chain of thought · approach: Routing, pruning and distillation for reasoning · 0 citations (OpenAlex)

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.

Theoretical analysis of test-time chain-of-thought scaling for transformers trained on in-context linear regression weight prediction, yielding context-length/test-compute trade-offs, overthinking conditions, and a quadratic task-selection objective favoring diverse, relevant, hard tasks.

Problem

When does increasing test-time compute (longer chain-of-thought) help, and what training data properties make test-time scaling effective or harmful for transformers?

Approach

Models a linear self-attention transformer trained on in-context weight prediction; shows CoT at test time implements pseudo-Newton iterations, derives a task hardness measure tr(Λ)/λmin(Λ), and formulates optimal multi-task selection probabilities as a quadratic program.

Key results

• Theory: for fixed test error, increasing CoT steps k permits reducing training context length n, with error bounded by d(1+n/(1+H))^{-2k}, where H=tr(Λ)/λmin(Λ).
• On Qwen2.5-7B-Instruct, Qwen-GCD reaches 75% accuracy on short-CoT GCD examples (0-1k chars) vs 30.39% for base, while Qwen-Poly drops to 20.83% on longer CoTs vs 27.2% for base.
• Task selection experiments show the quadratic program assigns higher probability to harder and more diverse task types.

Surprising findings

• Increasing test-time thinking can harm performance when task-relevant directions are underrepresented in training data (overthinking).
• Qwen-Poly, trained on polynomial roots, performs worse than base on longer GCD reasoning, showing training on a different task can make extra test-time thinking harmful.

Limitations

• Theory limited to linear regression and single-layer linear self-attention.
• Real-reasoning validation is a small experiment on two OMEGA tasks with one base model.
• Bounds rely on assumptions such as m = Omega(k^2 d) and bounded eigenvalues.

Experimental setup

Datasets: synthetic Gaussian linear regression prompts OMEGA

Benchmarks: OMEGA GCD test set

Metrics: test estimation error (squared norm) accuracy

Baselines: random task selection

Builds on: in-context weight prediction framework linear self-attention training dynamics chain-of-thought training for in-context learning

Closest papers by meaning

• Towards Thinking-Optimal Scaling of Test-Time Compute for LLM Reasoning (NeurIPS 2025 · similarity 0.90)
• Language models scale reliably with over-training and on downstream tasks (ICLR 2025 · similarity 0.90)
• Sample Complexity and Representation Ability of Test-time Scaling Paradigms (ICLR 2026 · similarity 0.90)
• Rethinking Fine-Tuning when Scaling Test-Time Compute: Limiting Confidence Improves Mathematical Reasoning (NeurIPS 2025 · similarity 0.89)
• Beyond What Seems Necessary: Hidden Gains from Scaling Training-Time Reasoning Length under Outcome Supervision (NeurIPS 2026 · similarity 0.89)
• Inference Scaling Laws: An Empirical Analysis of Compute-Optimal Inference for LLM Problem-Solving (ICLR 2025 · similarity 0.89)