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

Training-Trajectory-Aware Token Selection

ICML 2026 · Large language models › Reasoning and chain of thought · approach: Efficient and adaptive thinking · 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.

Identifies "Imitation Shock" in continual reasoning distillation, where performance crashes at an Imitation Bottleneck despite decreasing loss, and proposes T3S, which masks early-learned anchor tokens from the AR loss (and union-masks yet-to-learn tokens for dLLMs) to improve distillation from a few hundred examples.

Problem

Continual reasoning distillation into a strong student often yields limited gains or degradation even while training loss falls monotonically; the token-level cause of this failure is unclear.

Approach

Use the training trajectory around the lowest-training-accuracy checkpoint (Imitation Bottleneck) to split tokens by confidence change: tokens whose confidence rises early (Imitation-Anchor Tokens) are masked from the AR loss so yet-to-learn tokens get optimization priority; for dLLMs, yet-to-learn tokens are force-included in the random mask via union masking.

Key results

• T3S (R1) on BOBA-200 reaches 77.30 AIME24/25 average vs 71.46 base and 63.13 SFT for Qwen3-8B.
• With 200 examples, Qwen3-8B with T3S (R1+235B+QWQ mixing) reaches 80.42 AIME average, versus 74.90 for DeepSeek-R1 and 83.60 for Qwen3-235B.
• T3S-trained LLaDA-2.0-Mini reaches 66.01 average on AIME25/MATH500/TheoremQA vs 55.95 base and 56.19 SFT.

Surprising findings

• Discarding pre-bottleneck updates (RRT) converts SFT degradation into gains, suggesting early distillation updates can be harmful.
• Inverting the T3S mask (training only on anchor tokens) collapses performance well below SFT, e.g., 28.13 average on BOBA-200 with R1.

Limitations

• Requires training accuracy with gold answers or a verifier to locate the bottleneck.
• Main experiments use only 200 examples and a single fixed recipe (50 steps, batch 64, lr 1e-5), with AIME-heavy evaluation on small test sets.

Experimental setup

Datasets: BOBA-200 S1K-200 OpenThought3 RL-Code-Math-v5 S1K

Benchmarks: AIME24 AIME25 MMLU-Pro PhyBench GPQA-Diamond LiveCodeBench SuperGPQA CEval IFEval MATH500 TheoremQA

Metrics: accuracy (averaged over 16 runs) delta average vs base

Baselines: SFT RRT T3S-simple Ling-mini Qwen3-14B Ministral-14B-Instruct-2512 Qwen3-30B-A3B-Instruct-2507 Qwen3-235B-A22B

Builds on: RHO-1 LIMA block diffusion language models

Closest papers by meaning

• IAPO: Information-Aware Policy Optimization for Token-Efficient Reasoning (ICML 2026 · similarity 0.89)
• On the Token Value Inequality in Efficient Reasoning (NeurIPS 2026 · similarity 0.89)
• In-Token Rationality Optimization: Towards Accurate and Concise LLM Reasoning via Self-Feedback (AAAI 2026 · similarity 0.88)
• Critical Tokens Matter: Token-Level Contrastive Estimation Enhances LLM’s Reasoning Capability (ICML 2025 · similarity 0.88)
• Cornerstones or Stumbling Blocks? Deciphering the Rock Tokens in On-Policy Distillation (NeurIPS 2026 · similarity 0.88)
• Reasoning Can Be Restored by Correcting a Few Decision Tokens (ICML 2026 · similarity 0.88)