Critical Tokens Matter: Token-Level Contrastive Estimation Enhances LLM's Reasoning Capability

TL;DR

This paper introduces critical tokens in reasoning trajectories and demonstrates that identifying and replacing them with contrastive estimation significantly improves large language models' reasoning accuracy on benchmarks like GSM8K and MATH500.

Contribution

The work presents a novel framework for identifying critical tokens and extends it with a contrastive estimation method to enhance model training and reasoning capabilities.

Findings

Identifying critical tokens improves model accuracy on reasoning tasks.

Replacing critical tokens reduces reasoning errors.

The proposed cDPO method outperforms baseline training approaches.

Abstract

Mathematical reasoning tasks pose significant challenges for large language models (LLMs) because they require precise logical deduction and sequence analysis. In this work, we introduce the concept of critical tokens -- elements within reasoning trajectories that significantly influence incorrect outcomes. We present a novel framework for identifying these tokens through rollout sampling and demonstrate their substantial divergence from traditional error tokens. Through extensive experiments on datasets such as GSM8K and MATH500, we show that identifying and replacing critical tokens significantly improves model accuracy. We propose an efficient methodology for pinpointing these tokens in large-scale datasets using contrastive estimation and extend this framework to enhance model training processes with direct preference optimization (DPO). Experimental results on GSM8K and MATH500 benchmarks with the widely used models Llama-3 (8B and 70B) and Deepseek-math (7B) demonstrate the effectiveness of the proposed approach, cDPO. Our results underscore the potential of leveraging critical tokens to reduce errors in reasoning tasks, advancing the development of AI systems capable of robust logical deduction. Our code, annotated datasets, and trained models are available at https://github.com/chenzhiling9954/Critical-Tokens-Matter to support and encourage future research in this promising field.