Research map / Large language models / Reasoning and chain of thought
Critical Tokens Matter: Token-Level Contrastive Estimation Enhances LLM’s Reasoning Capability
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.
Defines critical tokens in incorrect math reasoning trajectories via rollout sampling, proposes contrastive estimation with positive/negative LoRA models to score them cheaply, and uses these scores to token-weight DPO (cDPO) for better math reasoning.
Problem
LLMs make token-level errors in mathematical reasoning, and standard DPO can reduce the likelihood of correct trajectories due to lexical overlap with incorrect ones; identifying which tokens cause failure is costly via rollout sampling.
Approach
Defines critical tokens as the first tokens in an incorrect trajectory whose rollout correctness drops to zero and stays low afterward; estimates them cheaply via contrastive estimation between a positive-trained and a negative-trained model, and penalizes them in DPO by weighting each negative token's implicit reward by (1 - s_t).
Key results
- cDPO averages 77.2 on GSM8K and 33.4 on MATH500, versus 71.3 for DPO-from-base... (best DPO baseline RPO averages 75.4 on GSM8K and 32.3 on MATH500)
- On GSM8K, cDPO reaches 90.8 with Llama-3-70B versus 80.4 for the baseline model.
- Contrastive estimation is claimed to need about 0.002% of rollout-sampling compute on GSM8K.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use cDPO instead of RPO | Llama-3-8B · GSM8K | accuracy: 67.5 -> 67.9* | Helped |
| Use cDPO instead of RPO | Llama-3-70B · GSM8K | accuracy: 89.7 -> 90.8* | Helped |
| Use cDPO instead of TokenDPO | DeepSeek-math-7B · GSM8K | accuracy: 69.6 -> 72.9* | Helped |
| Use cDPO instead of RPO | Llama-3-8B · MATH500 | accuracy: 18.4 -> 19.6* | Helped |
| Use cDPO instead of RPO | Llama-3-70B · MATH500 | accuracy: 43.8 -> 45.6* | Helped |
| Use cDPO instead of RPO | DeepSeek-math-7B · MATH500 | accuracy: 34.8 -> 35.0* | Helped |
| Replace the critical token with an alternative token during decoding (ablation) | Llama3-8B · GSM8K and MATH500 training instances | Pass@1 accuracy: approximately 30% | Helped |
| Replace the critical token with an alternative token during decoding (ablation) | Llama3-8B · GSM8K and MATH500 training instances | Pass@64: over 90% | Helped |
Surprising findings
- Replacing a single critical token with an alternative raises Pass@1 to about 30% and Pass@64 above 90% on sampled incorrect trajectories.
- Critical tokens often differ from human-annotated error tokens (65% on GSM8K, 87% on MATH500).
Limitations
- Critical token labels are heuristic, based on rollout correctness thresholds and a small analysis set of 100 trajectories per dataset.
- Gaussian distribution analysis of contrastive estimation is a simplifying assumption.
- Gains on MATH500 are small in absolute terms (base accuracy around 17-42).
Experimental setup
Datasets: GSM8K MATH
Benchmarks: GSM8K MATH500
Metrics: accuracy Pass@1 Pass@k average accuracy
Baselines: SFT DPO TokenDPO RPO
Builds on: DPO TokenDPO RPO contrastive decoding LoRA
Closest papers by meaning
- Token Assorted: Mixing Latent and Text Tokens for Improved Language Model Reasoning (ICML 2025 · similarity 0.90)
- Think Deep, Not Just Long: Measuring LLM Reasoning Effort via Deep-Thinking Tokens (ICML 2026 · similarity 0.89)
- Turning Dust into Gold: Distilling Complex Reasoning Capabilities from LLMs by Leveraging Negative Data (AAAI 2024 · similarity 0.89)
- Sample Smart, Not Hard: Correctness-First Decoding for Better Reasoning in LLMs (ICLR 2026 · similarity 0.89)
- TokUR: Token-Level Uncertainty Estimation for Large Language Model Reasoning (ICLR 2026 · similarity 0.89)
- Reasoning Can Be Restored by Correcting a Few Decision Tokens (ICML 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.
