Research map / Large language models / Reasoning and chain of thought
Multiple Token Divergence: Measuring and Steering In-Context Computation Density
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 Multiple Token Divergence (KL between a model's full next-token distribution and its shallow multi-token-prediction head) as a computational-effort measure, plus Divergence Steering, a geodesic-interpolation decoding method that biases generation toward or away from the MTP distribution.
Problem
Measuring in-context computational effort of language models, since next-token loss does not capture reasoning complexity, and prior latent-compressibility measures (PHi loss) are invasive and unstable to train.
Approach
Measure per-token computational effort as the KL divergence between the full model's output distribution and a shallow auxiliary MTP head's prediction, computable post hoc from existing MTP modules with no extra training; use the same geodesic (Fisher-Rao) interpolation between the two distributions as a decoding control knob alpha, complementary to temperature.
Key results
- MTD of MiMo-7B correlates positively with MATH difficulty (partial r=0.199 controlling for NLL on self-generated CoTs), while NLL correlates negatively (r=-0.158 uncontrolled).
- Selecting the lower-MTD CoT among correct/incorrect pairs yields 67.1% correct on MATH for MiMo, versus 73.3% for lower NLL and 80.4% when both agree.
- In creative writing, alpha=-0.1 gives the best Overall Impression score with MiMo; the optimum is negative with p<0.01 by quadratic regression.
Surprising findings
- Lower MTD is associated with correct reasoning, in contrast to prior PHi findings where correct answers had high PHi loss.
- MTD without access to the latest token embedding is negatively correlated with task complexity after controlling for NLL.
Limitations
- MTD may conflate genuine computation with simple patterns that exceed the shortcut head's capacity; results depend on MTP head size.
- Divergence steering gave no clear reasoning improvement for large pre-trained models in preliminary experiments.
- Correlations are modest (e.g., r around 0.2) and the Mistral results are quantitatively weaker than MiMo.
Experimental setup
Datasets: MATH GSM8K FineWeb
Benchmarks: MATH GSM8K creative writing benchmark
Metrics: partial correlation pairwise correct-CoT selection accuracy creativity score overall impression
Baselines: PHi loss next-token NLL
Builds on: Prediction of Hidden States (PHi) Multiple Token Prediction speculative decoding
Closest papers by meaning
- Think Deep, Not Just Long: Measuring LLM Reasoning Effort via Deep-Thinking Tokens (ICML 2026 · similarity 0.88)
- WASD: Wasserstein-based Knowledge Distillation for Large Language Models (NeurIPS 2026 · similarity 0.87)
- DiffAdapt: Difficulty-Adaptive Reasoning for Token-Efficient LLM Inference (ICLR 2026 · similarity 0.87)
- metabench - A Sparse Benchmark of Reasoning and Knowledge in Large Language Models (ICLR 2025 · similarity 0.87)
- On the Bias of Next-Token Predictors Toward Systematically Inefficient Reasoning: A Shortest-Path Case Study (NeurIPS 2025 · similarity 0.87)
- Token Assorted: Mixing Latent and Text Tokens for Improved Language Model Reasoning (ICML 2025 · similarity 0.87)
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.
