Critical Tokens Matter: Token-Level Contrastive Estimation Enhances LLM’s Reasoning Capability
Authors: Zicheng Lin
Tian Liang
Jiahao Xu
Qiuzhi Liu
Xing Wang
Ruilin Luo
Chufan Shi
Siheng Li
Yujiu Yang
Zhaopeng Tu
Venue: ICML 2025 poster (2025)
Keywords: Large Language Models
Reasoning
Primary area: deep_learning->large_language_models
TL;DR: We introduce critical tokens -- elements within reasoning trajectories that influence incorrect outcomes. We use contrastive estimation to pinpoint these tokens and extend this approach to enhance the direct preference optimization (DPO) method.

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.

Critical Tokens Matter:
Token-Level Contrastive Estimation Enhances LLM's Reasoning Capability
Zicheng Lin
Tian Liang" 2* Jiahao Xu 2 Qiuzhi Liu Xing Wang? 2
Ruilin Luo'
Chufan Shi
Siheng Li Yujiu Yang' Zhaopeng Tu
Abstract
2024; Yu et al.; Collins et al., 2024). These tasks neces- sitate logical and sequential deduction to derive solutions, making them challenging for models trained primarily on general language processing. The chain of thought (COT) method (Wei al., 2022) has significantly improved the reasoning capability of L.LMs by employing series of in- termediate reasoning steps, reasoning trajectories. Prior research has categorized trajectory errors based on mod- ifications required to correct the COT, such as calculator errors, with the objective of identifying avenues for model improvement (Wei et al., 2022; Wang et al., 2023).
Mathematical reasoning tasks pose significant challenges for large language models (LLMs) be- cause 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 novel frame- work for identifying these tokens through rollout sampling and demonstrate their substantial diver- gence from traditional error tokens. Through ex- tensive experiments on datasets such as GSM8K and MATH500, we show that identifying and replacing critical tokens significantly improves model accuracy. We propose an efficient method- ology for pinpointing these tokens in large-scale datasets using contrastive estimation and extend this framework to enhance model training pro- cesses with direct preference optimization (DPO). Experimental results on GSM8K and MATH500 henchmarks with the widely used models Llama-3 (8B and 70B) and Deepseek-mate (7B) demon- strate the effectiveness of the proposed approach, cDPO. Our results underscore the potential of leveraging critical tokens to reduce errors in rea- soning tasks, advancing the development of Al systems capable of robust logical deduction.
Despite these advancements, the token-level discrepancies within mathematical reasoning contexts have not been sys- tematically explored. Our study seeks bridge this gap by introducing novel framework for identifying and quanti- fying the impact of critical tokens on model accuracy. We define critical tokens in mathematical reasoning as crucial components within an incorrect trajectory that significantly alter the final outcome. We utilize rollout sampling iden- tify tokens that substantially influence the correctness of rea- soning trajectories. Our findings reveal that critical tokens often diverge from annoat-d error tokens, highlight- ing their unique role in disrupting logical coherence and computational accuracy. By analyzing the characteristics of critical tokens through word type and positional analysis, we provide novel insights into their nature and influence mech- anisms. Furthermore, manipulating a single critical token in incorrect trajectories can significantly enhance accuracy, underscoring their pivotal role in error mitigation.
1. Introduction
Building on these insights, we illustrate how critical tokens can enhance reasoning capabilities within Direct Preference Optimization (DPO), a commonly used reinforcement learn- ing algorithm. Although DPO proves effective for general tasks, it encounters difficulties in mathematical reasoning because it may reduce the generation likelihood of positive examples due to lexical similarities with negative examples. Our proposed method, cDPO, addresses this issue by tar- geting critical tokens unique to negative examples, thereby improving the model's ability to differentiate between pos- itive and negative instances. cDPO involves the efficient identification and penalization of critical tokens predom-
In the domain of artificial intelligence, mathematical rea- soning tasks are seen crucible for evaluating the pro- ficiency of large language models (LLMs) (Cobbe et al., 2021; Hendrycks et al., 2021; Yuan et al., 2023; Ahn et al.,
'Equal contribution 'Tsinghua University "Tencett. Correspon- dence to: Yujiu Yang Zyang-yujiu wuiiuyyjjiuut... wwt.siiguua..duu...... Zhaopeng Tu Minze223@milss hain:e33eemmssst.nnuuu......... cdd.......cu..nn>
tenuueeeeneet.mm>... First Author
Proceedings of the 42nd International Conference Machine Learning, Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

Critical Tokens Matter: Token-Level Contrastive Estimation Enhances LLM's Reasoning Capability
$125 initially, She still owes $125 0.75 593.75. The answer is $93.75.
x
Question:
owed
125 dollars. 75% 3/4, 3/4 125 3/4 125 93.75. The answer is 93.75.
x
Mathilda is determined to pay back the money she owes a friend so she decides pay initial installment of
125 dollars initially. If she still owes 75% of that, that is 0.75 x 125 93.75 dollars. The unswer is 93.75.
x
$125. If she still has 75% left to pay. how much did she owe originally?
Response:
an initial installment $125, which represents 25% of the total amount she owed. the total amount she is: 125 / 0.25 500. The answer 500.
Let's think step by step. She owed $125 initially, She still has 75% left pay, she owes $125** 0.75- $93.75. The answer is $93.75.
125 dollars. She owes 75% of the original amount. 100% minus 75% is 25%. 125 0.25 is 500. The answer is 500.
paid
$125 an initial installment with 75% left pay. So, the total owed is 125 wes is: 125 / 0.25 = 500. The answer is 500.
Figure 1. An illustration of the critical token "owed" shows that it fails to lead to the correct answer in any case. Replacing it with an alternative significantly increase model accuracy.
inantly found in negative examples, refining the model's learning process and enhancing its understanding of positive outcomes. Experimental results on GSM8K and MATH500 benchmarks using widely recognized models like Llama-3 (8B and 70B) and Deepseek-math (7B) demonstrate that our approach surpasses strong DPO baselines, such as Tok- enDPO (Zeng et al., 2024) and RPO (Liu et al., 2024; Pang et al., 2024), across all evaluated scenarios.
specific conditions in sequence analysis. Our findings reveal that critical tokens differ from human-annotated error to- kens in substantial proportion of cases, emphasizing their unique role in reasoning failure (Table 1). Further experi- mentation shows that replacing critical tokens boosts model accuracy significantly, highlighting their importance in er- ror reduction (Figure 2). We conclude with an analysis of critical tokens based on word types and positional attributes (Table 2), providing insights into their characteristics.
In summary, contributions are fhree-fold:
We introduce the concept of critical tokens in mathemat- ical reasoning tasks and empirically validate their exis- tence through extensive rollout sampling, distinguishing them from traditional error tokens.
Intuition Mathematical reasoning tasks require logical and sequential deduction to find solutions. We have ob- served that within incorrect reasoning trajectories, certain tokens are pivotal in leading to incorrect outcomes. These to- kens disrupt the logical flow, misrepresent relationships, introduce computational errors, thus significantly affecting the final result. Unlike other tokens that may have negligible effects on the reasoning process, these "critical tokens" are crucial points of failure. Identifying these tokens is essen- tial because avoiding or correcting them can often result in correct outcomes, even within an incorrect trajectory. As illustrated in Figure the token "owed" is predominantly re- sponsible for incorrect reasoning trajectories as it misguides the logical deduction process. In contrast, prompting the model to decode alternative tokens like "paid" significantly increases the likelihood producing a correct final result.
We propose an efficient approach using contrastive esti- mation to practically identify critical tokens in large-scale training data, requiring only as little as 0.002% of the computational cost of rollout sampling for GSM8K.
We develop cDPO, an innovative approach that leverages critical tokens within DPO, enhancing the algorithm's ability to distinguish between positive and negative exam- ples in mathematical reasoning.
2. Critical Tokens in Mathematic Reasoning
In this section, we explore the concept of "critical tokens" in mathematical reasoning tasks and their impact model accuracy. We begin by defining critical tokens as pivotal points in incorrect reasoning trajectories that significantly in- fluence outcomes using the example in Figure 1. To validate the presence of critical tokens, we perform rollout sampling, identifying tokens with zero correctness scores that meet
Empirical Validation with Rollout Sampling To empiri- cally validate the existence of critical tokens, we conducted 64 rollout samplings for each token within an incorrect tra- jectory. We calculated a score for each token based on the correctness ratio of the generated completions, to quantify influence on the trajectory. The first token that meets the
2

---

Critical Tokens Matter: Token-Level Contrastive Estimation Enhances LLM's Reasoning Capability
Table 1. The ratio of incorrect trajectories where critical tokens different from tokens across the error types.
GSM8K Training Data
MATH500 Training Data
Error Types
#Count
Critical Error
#Count
Critical Error
Calculation Error
60
56.7%
54
79.6%
One step Missing
17
72.7%
8
87.5%
Semantic Misund.
22
82.3%
17
100.0%
Degeneration
1
100.0%
21
95.2%
Total
100
65.0%
100
87.0%
Table 2. Analysis of identified critical tokens.
Word Types
Relative Positions
Error Types
Function
Content
Number
Operator
Punct.
Before
After
100 Instances from GSM8K Training Data
Calculation Error
14
19
17
17
One step Missing
5
2
0
5
Semantic Misund.
3
6
3
3
9
Degeneration Total
0
0
0
0
1
18
22
32
13
15
33
32
100 Instances from MATH500 Training Data
Calculation Error One step Missing
30
13
3
0
2
3
1
Semantic Misund.
4
3
2
6
10
7
Degeneration
6
3
6
5
1
10
10
Total
23
16
25
31
5
56
31
following two conditions identified as the critical token:
to 87% in the MATH500 training dataset. This variance further nuanced when examining specific error categories. GSM8K data, calculation errors show a 56.7% mismatch, which suggests that nearly half the time the critical tokens identified do not correspond to the actual error tokens. One-step missing error demonstrates higher 72. 72.7% discrepancy, and semantic misunderstandings show an even greater divergence of 82.3%. Notably, degeneration errors though based on a single occurrence exhibit a complete 100% discrepancy with error tokens.
The token's ' correctness score is O;
For example,
The scores of all subsequent tokens are below 5%.
We analyzed 100 incorrect trajectories from Llama3-8B. base, randomly selected from the MATH training dataset, and successfully identified the critical token in all cases. In addition, by examining 100 random incorrect trajectories from the GSM8K training data, identified the critical token in 99 cases out of 100. For the outlier case, we iden- tified critical token that only satisfied the first condition. These results demonstrate the existence of critical tokens.
In the MATH500 dataset, similar pattern observed. Cal- culation errors exhibit 79.6% discrepancy, one-step miss- ing errors an 87.5% mismatch, semantic misunderstandings 100% divergence, and degeneration errors a significant 95.2% discrepancy. The results in MATH500, particularly the large differences in degeneration errors, underscore the complexity involved in high-precision domains.
Critical Tokens Are Not Necessarily Error Tokens Re- searchers might hypothesize that critical tokens tend to co- incide with еитог tokens, given their definition (i.c., the correctness ratio of rollout samplings is 0). However, Ta- ble 1 demonstrates that critical tokens frequently differ from human-annotated enror tokens across various error types (Wei et al., 2022; Wang et al., 2023).
These findings underscore critical insight: while critical to- kens valuable for flagging potential issues in trajectory, there isn't always direct correlation with human-annotadl error tokens. This divergence emphasizes the intricate na- ture of error detection and correction in algorithmic analyses,
In the GSM8K training dataset, 65% of the critical tokens do match the error tokens, while this disparity increases
3

---

Critical Tokens Matter: Token-Level Contrastive Estimation Enhances LLM's Reasoning Capability
suggesting that critical tokens capture a broader context of
100%
underlying issues, possibly before error becomes evident.
80%
Analysis of Critical Tokens We analyze critical tokens from the following perspectives, as detailed in Table 2.
66.21
Pass@@k
60%
24.11
Tokens are categorized into five types based on their lin- guistic roles: function words, content words, and numbers. Additionally, operators (e.g., mathematical symbols) and punctuation marks are identified. Our analysis reveals differ- ing patterns across eπrоr types and datasets, highlighting the complexity of math reasoning tasks. In the GSM8K dataset, calculation errors predominantly occur with numbers (19 instances) and are fairly distributed among other word types. On the other hand, one step missing errors primarily involve numbers and function words. This suggests an emphasis on numerical manipulation where precision in number usage and function words indicating operations are critical. The MATH500 dataset follows a somewhat similar trend, though it exhibits a higher occurrence of related calcula- tion errors (17 instances), indicating a more frequent use of complex mathematical operations in this dataset. This accentuates the need for careful handling of operators in mathematical computations.
42.00
40%
MATHS000: w/o Critical
30
MATH500: w/ Critical Takens
20%
GSM8K: wio Criticl Tokens
GSM8K: Critical Tokens
0%
0.0%
00%
0.0%
0.0%
0.0%
2
16
32
64
Number of Samples (k)
Figure 2. Impact critical tokens on reasoning accuracy. Re- placing critical tokens with alternatives ("w/o Critical Tokens") can significantly increase model accuracy on both GSM8K and MATH500, highlighting the importance of these tokens.
observed a significant improvement in the model's perfor- mance on these datasets, with Pass@1 I accuracy reaching approximately 30% and Pass 64 increasing to over 90%. These results underscore the erucial role that critical tokens play potential stumbling blocks in the reasoning process. By preventing the model from proceeding with these criti- cal tokens and suggesting alternatives, we enable higher likelihood of reaching accurate conclusions.
The position of critical tokens is analyzed relative to the corresponding error tokens, categorizing them as occurring before or after the error tokens. If the critical token is the error token itself, it is not counted. In GSM8K, critical tokens are almost evenly distributed: 33 occur before and 32 after the error. However, in MATH500, more critical tokens occur before error tokens, suggesting that critical tokens capture broader context of underlying issues in the complex MATH500 problems, potentially before an error becomes evident.
The findings emphasize the importance of understanding and manipulating critical tokens to enhance model perfor- mance, especially in complex reasoning tasks. By identi- rying these critical tokens in reasoning trajectories, we can mitigate the risk of errors, thereby significantly improving the model's effectiveness. Such insights could be instru- mental in refining training processes and increasing the reliability of Al systems in real-world applications.
Impact of Critical Tokens We investigate the effect of critical tokens by replacing them with alternative tokens during model decoding. Specifically, let t. be a critical token in a given trajectory, and let T<i = {4,,...,4: 4-13 represent the preceding tokens. We perform & rollout samples based on two different prefixes and calculate the Pass k metric:
3. Enhancing Reasoning Capability with Critical Tokens
In this section, we demonstrate how reasoning capabili- ties, trained with commonly-used DPO, can be enhanced by using critical tokens. Despite the success of DPO in general instruction tuning tasks, challenges persist when it is applied to reasoning and mathematical tasks. Studies have delved into this issue from an optimization perspective, identifying that these algorithms often diminish the genera- tion likelihood of positive examples in reasoning tasks due to the lexical similarity between positive and negative exam- ples. This overlap can lead to a situation where the model struggles to effectively prioritize and generate the correct tra- jectory during reasoning or mathematical problem-solving tasks. As a result, the adopted approach is to optimize pref- erences while ensuring that high generation likelihoods are maintained exclusively for positive examples (Liu et al., 2024; Pang al., 2024; Pal al., 2024; Feng et al., 2024).
w/ critical tokens:
The prefix is (t1...., ..). Based on the definition critical token, all Pass@ @ scores are zero.
w/o critical tokens: The prefix is (4I,..., 4-14,, 1), exclud- ing the critical token 4, We force the model to decode an alternative token at the same position, using the model's 's probability distribution while masking out the critical to- ken. This substitution allows us to explore it leads to improved outcomes in model predictions.
Figure 2 displays the results on 100 instances sampled indi- vidually from the GSMSK and MATH500 training datasets. By replacing critical tokens with alternative tokens, we
4

---

Critical Tokens Matter: Token-Level Contrastive Estimation Enhances LLM's Reasoning Capability
However, these approaches still attribute high likelihood to tokens that appear in both negative and positive examples, thereby failing to adequately distinguish between truly ben- eficial features and those that are merely ubiquitous across positive and negative examples.
For training the positive model, we randomly selected single correct trajectory because we expect the model to develop decisiveness using its own accurate reasoning paths. For training the negative model, we chose the incorrect tra- jectories that most frequently occur and account for 50% of all incorrect cases. This approach ensures both variety and representativeness, allowing us to accurately identify critical tokens. For instance, if there are 10 incorrect tra- jectories comprising 3 cases with the incorrect answer a, 2 cases with the incorrect answer b, and other cases with answers {c, d, e, f.g] occurring only once each, we would randomly select one incorrect trajectory from those with an- swers a and b, as they appear in 5 cases in total. Finally, we train the negative model on this example using two incorrect trajectories: one with answer a and the other with b.
In this work, we mitigate this problem by leveraging critical tokens that occur only in negative examples. By focusing on these tokens, we aim to enhance the model's ability to effec- tively differentiate between positive and negative examples. Our approach involves the identification and penalization of critical tokens that are prevalent exclusively in negative examples. This allows us to adjust the model's learning process by reducing the likelihood of these negative-specifie tokens, thereby refining the model's understanding of what constitutes positive outcome. By explicitly incorporating mechanism that penalizes frequent negative-example tokens, we ensure that the model learns to prioritize features that truly contribute to successful task completion.
Contrastive Estimation With both the positive model and the negative model available, we can automatically annotate the likelihood of each token in an incorrect trajectory being a critical token using contrastive estimation. Let x be query, and y = -u,,..... ................................ Bt.................... length T used in DPO training. We compute the likelihood of token y/ being a critical token, denoted as 8t, with the following equation:
3.1. Efficient Identification of Critical Tokens
yr) } be a negative example of
While it is straightforward to identify critical tokens using rollout sampling as described in Section 2, such methods incur prohibitively high sampling costs and face significant scalability challenges. Moreover, existing methods (Guo et al., 2023; Yoon et al., 2024) depend on external models for token-level annotations, which, although providing effective supervision signals, are costly and limited by the capabili- ties of the external models. To efficiently identify critical tokens, we propose a method called contrastive estimation, which leverages models trained to learn patterns from both correct and incorrect reasoning trajectories. Figure 3 depicts the framework. By comparing the token-level likelihoods produced by two separately trained models, contrastive esti- mation can effectively pinpoint critical tokens contributing to incorrect outcomes. The contrastive estimation probabil- ity naturally highlights tokens (e.g., "'wwed") that lead to incorrect reasoning outcomes. We provide additional details throughout the remainder of this section.
log st =(1 + B) log PP
Bllog PM (ux yet) - log Z
(1)
Here, B is a scaling hyperparameter, while (P(-) and P'(() represent the probabilities from the positive and negative models, respectively. The term log Z is the partition func- tion used in the softmax computation. A low st indicates low likelihood under the correct pattern and high likeli- hood under the incorrect pattern, signaling the presence of critical tokens. We further provide theoretical and empirical analysis in Appendix A.
Efficiency Analysis of Contrastive Estimation To eval- uate the computational efficiency of contrastive estimation compared to rollout sampling, we estimate the number of forward passes required. Rollout sampling, used to identify critical tokens, incurs substantial inference costs as it relies on sampling from the base model. For GSM8K, obtaining critical tokens through rollout sampling for 100 incorrect examples (64 samples per token) results in an average of 581,422 1,25 additional tokens per response (7,613,942 tokens for MATH). Therefore, for n examples, rollout sampling requires approximately 581, 425 n forward passes.
Training Positive and Negative Models To implement the contrastive estimation, we need to develop models that can effectively estimate a wide range of both correct and incorrect reasoning distributions. To this end, we collect reasoning trajectories based on the sampling strategy: given dataset of M instances D {z,, ) we utilize pre-trained LLM to sample N reasoning trajectories. Then, we verify the outcome results based on the golden labels ys. which yields k; positive reasoning trajectories and N - k; negative reasoning trajectorics, which is denoted as:
In contrast, contrastive estimation involves both training and inference costs. On GSM8K, the dataset for training the positive and negative models contains 26,131 examples (68,391 examples for MATH), and SFT on this dataset re- quires approximately 3 x 26, 131 = 78, 393 forward passes, assuming a batch size of 1. For inference on n examples,
DP = =((x,, 15
D (x., (ut,, 1/y ) M
5

---

Critical Tokens Matter: Token-Level Contrastive Estimation Enhances LLM's Reasoning Capability
Base Model
D ((,
Sampling with Correctness Validatiion
$125 inittally.....
$93.75:
((x, Dipl=-)
{(x, 1
Positive Model
Negative Model
SFT
SFT
Preslction
Positive Model
Negative Model
Contrastive Estimation
Base Model
= ((x.)))))
eDPO
Trained Model
Figure 3. Pipeline of the proposed cDPO that involves the efficient identification and penalization of critical tokens that are prevalent exclusively in negative examples.
contrastive estimation only requires each of the positive and negative models to perform one forward pass per example, totaling 2n. Consequently, the total cost for contrastive 783 estimation 78, 393 + 2n forward passes for GSM8K, which significantly lower than the cost of rollout sam- pling. For example, for GSM8K that consists of 7,500 examples, contrastive estimation requires only as little as 01004-+2777000(58814555777777700
Formulation
Given the pairwise preference dataset D (I,,y ( M/ the original DPO loss formulated as:
M
EDPO =
log (o - 6(2,, v())
1=1
Here, o(x, y) is an implicit reward function, given by:
of the computa- tional cost of rollout sampling. Despite its efficiency, con- trastive estimation demonstrates reasonable alignment with rollout-based supervision. Specifically, we measured the Area Under the Curve (AUC) of CE with respect rollout on L.L.MA3-8BB, obtaining values of 0.77 GSM8K and 0.84 MATH.
d(z, v) = 7 log T x)) Tee( x))
where no(- and < represent the policy model and the reference model, respectively, and is the coefficient for the KL divergence penalty.
We extend the sample-level DPO to token-level DPO with critical rewards (i.e., cDPO). First, modify the reward function d(x, y) to include token-level scores sf (Equa- tion 1) follows:
3.2. cDPO: Explicitly Penalizing Critical Tokens in DPO
Intuition Critical tokens in incorrect trajectories signif- icantly contribute to errors, even when other tokens may be correctly placed. By assigning teken-level scores to incorrect trajectories, we can specifically penalize critical tokens without adversely affecting correct ones. Conversely, scoring correct trajectories to encourage certain tokens can inadvertently penalize other valid tokens, resulting in unde- sired distribution shifts in DPO (Rafailov et al., 2024; Xu scoring
Σ (1 st) log To vee) 1=1
T
6s (.z.
Tiet a M<ce)
where T is the total length of the response y. and s; rep- resents the token-level reward score in cDPO for the 1-th token. Accordingly, the objective of cDPO is formulated as: LDDO
et al., 2024). Therefore, we focus exclusively tokens within incorrect trajectories and extend DPO from the example level to the token level by utilizing token-level rewards for preference optimization.
M
Σ
loga(ó(zu.
(2,,, up: 8())
Note that only the reward function for the negative example y" modified. Intuitively, lower values of 8; suggest a
6

---

Critical Tokens Matter: Token-Level Contrastive Estimation Enhances LLM's Reasoning Capability
higher likelihood of being critical tokens, which more prone to result in incorrect outcomes. By weighting each token's contribution with 1 St, the model effectively pe- nalizes generating these critical tokens. This token-level approach helps ensure that the model reduces the likelihond of generating critical tokens, thus improving the overall accuracy of the responses.
Log probubilltty
-100
-150
Chosen in cDRO
-000
Rejected in cDPV
3.3. Experimental Results
Chosco in RPO
Rajatta in RPO
Experimental Setup We used two widely recognized math reasoning datasets: GSMSK (Cobbe et al., 2021) and MATH (Hendrycks al., 2021). For training, we sampled from all questions in the training set to generate the data. For evaluation, we utilized the MATH500 subset, which is uniformly sampled and has a distribution of difficulty levels and subjects that matches the full MATH test set, as demonstrated in Lightman et al. (2023). Additionally, for both training sampling and evaluation, we adopted few-shot prompting: 8-shot prompting for GSM8K, following Wei et al. (2022), and 4-shot prompting for MATH500, as de- scribed in Lewkowycz al. (2022). For all main evaluations, the temperature was fixed at O.
-250
Chosco in DPO
Rajetal in DPQ
1500
Training Steps
Figure 4. Log probabilities of chosen and rejected sequences dur- ing training on the GSM8K dataset. Solid lines represent chosen sequences, while dashed lines denote rejected sequences. The figure illustrates how cDPO achieves a better separation between chosen and rejected sequences compared to DPO and RPO.
with a learning rate of 2e-5 for all baseline methods. For our cDPO approach, since the token-level scores range between and (whereas DPO, the scores were all 1), we simply increased the learning rate to 4c-5.
We conducted experiments on a range of models, including the general-purgee purpose models Llama-3-8B 3-seeb and Llama- 3-70B-base (Dubey et al., 2024), as well as the domain- specific model DeepSeek-math-7BB- ma-bee (Shao et al., 2024). For comparison, evaluated multiple baseline methods us- ing the data generated from the process described in Section 3.1. For Supervised Fine-Tuning (SFT), we fine-tuned the model using the positive response set DD, For preference optimization (PO) methods, we utilized the token-level an- notated pair-wise preference dataset D. The baselines we compared include:
For proposed cDPO, each problem was sampled N 64 times, selecting the top-p 50% of incorrect trajectories to train the negative model q(-). During estimation, the hyperparameter / was set to 1.0.
Learning Curves We begin by examining the impact of cDPO on training dynamics. Figure 4 shows the log proba- bility trends for selected and non-selected sequences over training steps on the GSM8K dataset using the Llama-3-8B model with DPO, RPO, and cDPO.
DPO (Rafailov al., 2024): We tested two different starting points for training: based on the base model and on the SFT model.
The proposed cDPO method successfully differentiates be- tween chosen and rejected sequences by significantly in- creasing the log probability of correct sequences while sharply decreasing that of incorrect ones. In comparison, RPO (DPO with an additional NLL term) increases the probability of correct sequences, but its impact on reducing incorrect response probabilities is less significant. On the other hand, DPO notably decreases the probability of gener- ating incorrect sequences, but also reduces the probability of correct sequences. This suggests that cDPO strikes balanced approach, effectively enhancing the probability of correct outputs while minimizing critical errors, exceeding the performance of both DPO and RPO.
ToKenDPO (Zeng et al., 2024), which is token-level ap- proach that enhances Kullback-Leible (KL) divergence regulation by incorporating forward KI. divergence con- straints at the token level. The SFT model is used as the starting point for training. We implemented TDPO using the publicly available implementation.
RPO (Liu et al., 2024; Pang et al., 2024) introduces an ex- negative log- likelihoooo term to improve performance reasoning tasks. We implemented it using Hugging- Face's implementation and starting with the base model.
Main Results Table 3 presents the experimental results for various methods across the GSM8K and MATH500, Our proposed method consistently outperforms all baselines and other methods, achieving the highest scores across both datasets. For the GSM8K dataset, our method achieves a
We used LoRA adapters (Hu et al., 2022) to train all the models. We trained both positive and negative models for 1 epoch with a learning rate of 3e-4. For preference opti- mization training, we set ? 1.0 and trained for epochs
7

---

Critical Tokens Matter: Token-Level Contrastive Estimation Enhances LLM's Reasoning Capability
GSM8K
MATH500
Method
Llama-3
DeepSeek
Llama-3
DeepSeek
Avg.
Avg.
8B
70B
math-7B
8B
70B
math-7B
Baseline
56.4
80.4
64.1
67.0
18.6
43.6
34.0
32.1
+ SFT
61.2
82.1
67.1
70.1
16.8
43.4
34.2
31.5
+ DPO
59.7
87.8
66.5
71.3
17.6
40.8
34.6
31.0
+ ToKenDPO
62.3
83.3
69.6
71.7
19.2
43.6
34.6
32.5
+ DPO
59.6
88.9
63.1
70.5
14.6
42.6
35.0
30.7
+ RPO
67.5
89.7
68.9
75.4
19.4
44.0
33.6
32.3
+ cDPO (Ours)
67.9*
90.8
72.9*
77.2
19.6
45.0*
35.2*
33.3*
Table 3. Experimental results on GSM8K and MATH500 datasets. Our proposed method surpasses all the strong baselines at a large margin individual settings and average performance. denotes the significance where p 0.005.
remarkable average score of 77.2, surpassing the Baseline and notable improvements such as those incorporating SFT and DPO. Specifically, our approach reaches the highest scores with Llama-3 (90.8 for 70B) and DeepSeek (72.9). These results highlight the effectiveness of our method in leveraging the strengths of both large-scale models (Llama- 3) and task-specific models (DeepSeek).
2010; Bose et al., 2018; He et al., 2020; Denize et al., 2023). Specifically, our work is closely related contrastive de- coding (CD), an application of contrastive estimation in downstream tasks. CD (Li et al., 2023) involves contrast- ing token distribution likelihoods between expert and am- ateur models during decoding. As described by O' O'Brien & Lewis (2023), this technique avoids high- sish-probability but low-quality tokens, ensuring text fluency and coherence.
Similarly, on the MATH500 dataset, our method attains an average score of 33.3, marking an improvement over the baseline (32.1) and other enhanced methods such SFT and RPO. Notably, our approach yields the highest individual score with Llama-3 (45.0 for 70B) and performs robustly across all model configurations.
Subsequent research has emphasized CD's potential im- prove factuality (Zhang et al., 2023; Yang et al., 2024), knowledge editing (Bi et al., 2024), safety (Zhao et al., 2024), and reasoning (O'Brieq & Lewis, 2023; Shi et al., 2024; Yuan et al., 2024; Yang et al., 2025). Notably, O' O'Brien & Lewis (2023) showed that CD enhances reasoning tasks and mitigates typical errors. Shi et al. (2024) demonstrated that unchosen experts in Mixture-of-Experts models could be applied for CD, thereby improving model reasoning ca- pacities. Different from those works that focus the infer- ence of contrastive decoding, our research primarily uses contrastive estimation to identify "critical tokens" that sig- nificantly affect the correctness of the reasoning process.
Since models exhibit sensitivity to prompt formatting when evaluated with temperature 0 on MATH5000 we further perform additional analyses by sampling each question 10 times under varying temperatures and report Pass@ in the Table 5. As shown, cDPO consistently outperforms the corresponding base models across all temperature settings and exhibits stable, robust performance.
To further verify the effectiveness of our method on stronger models, we conduct experiments using Qwen-2.5-7B and Qwen-2.5-32B both GSM8K and MATH500. The results presented in Table 6. As shown, cDPO consistently im- proves performance across both datasets and model scales.
Reinforcement Learning Among various alignment al- gorithms (Christiano al., 2017: Schulman et al., 2017; Ziegler al., 2019; Ouyang et al., 2022; Bai al., 2022), DPO (Rafailov et al., 2024) is one of the most representa- tive algorithmss DPO uses the LLM itself as a secret reward model and conducts preference optimization on a preference pair of positive and negative examples. Since then, various contributions have been made to further advance the DPO development of LLM alignment (Pal et al., 2024; Amini et al., 2024; Azar et al., 2024).
The consistent performance improvements observed across various settings underscore the superiority of our method compared to existing techniques. The significance tests, which were conducted to verify the statistical reliability of these results, confirm the competitive advantage of our proposed approach.
A line of research aims to improve the performance of DPO in reasoning tasks (Liu et al., 2024; Pang et al., 2024; Pal et al., 2024; Feng et al., 2024). Furthermore, Lai et al. (2024) leverage human or GPT-4 validation to pinpoint incorrect reasoning steps; Guo al. (2023); Yoon al. (2024) utilize
4. Related Work
Contrastive Estimation Notable works have refined con- trastive estimation techniques (Gutmann & Hyvärinen,
8

---

Critical Tokens Matter: Token-Level Contrastive Estimation Enhances LLM's Reasoning Capability
external LLMs to refine responses, deriving token-level preferences from pre- and post-revision comparisons. In contrast, our study seeks to establish an automatic process supervision strategy devoid of human annotation and easy to scale. Specifically, we harness contrastive estimation to identify critical tokens, providing token-level signals for preference optimization that significantly enhance LLM reasoning capabilities.
References
Ahn, J., Verma, R., Lou, R., Liu, D., Zhang, R., and Yin, W. Large language models for mathematical reasoning: Progresses and challenges. In Proceedings of the 18th Conference of the European Chapter of the Association for Computational Linguistics: Student Research Work- shop, pp. 225-237, 2024.
Amini, A., Vieira, T., and Cotterell, R. Direct pref- erence optimization with an offset. arXiv preprint 200i::2402.05777 2024.
5. Conclusion
Our work contributes a significant framework for under- standing and enhancing mathematical reasoning in LLMs through critical token analysis. By defining and identifying critical tokens, we provide valuable insights into token-level discrepancies that disrupt logical reasoning. Our approach, cDPO, successfully integrates this analysis into DPO, im- proving model performance on mathematical tasks by im- proving the model's differentiation between positive and negative examples. Experimental results from GSM8K and MATH500 benchmarks have shown that our method outper- forms existing DPO baselines, underscoring the potential of critical token interventions in enhancing model accuracy. Our research opens new doors for further exploration of token-level influences in complex reasoning tasks, which could lead to more refined and effective LLMs.
Azar, M. G., Guo. Z. D.. Piot, B., Munos, R., Rowland, M., Valko, M., and Calandriello, D. A general theoretical paradigm to understand learning from human preferences. In International Conference on Artificial Intelligence and Statistics, pp. 4447-4455, 2024.
Bai, Y., Jones, A., Ndousse, K., Askell, A., Chen, A., Das- Sarma, N., Drain, D., Fort, S.. Ganguli, D., Henighan, T., et al. Training a helpful and harmless assistant with rein- forcement learning from human feedback. arXiv preprint arXiv:2204.0588 2022.
Bi, B., Liu, S., Mei, L.,, Wang, Y., coding by contrasting knowledge: Enhancing Ilms' confi- dence on edited facts. arXiv preprint arXiv:2405.11613. 2024.
Ji, P., and Cheng, X. De-
Future work should explore the integration of cDPO with other reasoning frameworks and extend its application to diverse logical reasoning domains, contributing towards the broader aim of developing more robust and reliable LLMss.
Bose, A. J., Ling, H., and Cao, Adversarial contrastive estimation. arXiv preprint arXiv: (005: 03645, 2018.
Christiano, P. F., Leike, J., Brown, T., Martic, M., Legg, S., and Amodei, D. Deep reinforcement learning from human preferences. Advances in neural information pro- cessing systems, 30, 2017.
Impact Statement
In our research, we focus exclusively on developing models for solving mathematical problems, which inherently mini- mizes common ethical concerns typically associated with AI applications in broader domains. The primary function of our models is to enhance computational accuracy and efficiency in mathematical tasks, ensuring that any potential bias, privacy issues, or harmful outputs are effectively non- existent. Since our dataset consists solely of mathematical questions, it does not involve personal, sensitive, or contro- versial information that could lead to ethical dilemmas.
Cobbe, K., Kosaraju, V., Bavarian, M., Chen, M., Jun, H., Kaiser, L., Plappert, M., Tworek, J., Hilton, J., Nakano, R., et al. Training verifiers to solve math word problems. arXiv preprint arXiv: 10. 14168, 2021.
Collins, K. M., Jiang, A. Q., Frieder, S., Wong, L., Zilka, M., Bhatt, U., Lukasiewicz, T., Wu, Y., Tenenbaum, J. B., Hart, W., et al. Evaluating language models for mathe- maties through interactions. Proceedings of the National Academy of Sciences, 121(24):2218124122 2024.
Acknowledgments
Denize, J., Rabarisoa, J., Orcesi, A., Hérault, R., and Canu, S. Similarity contrastive estimation for self- self-supervised soft contrastive learning. In Proceedings of the IEEF/CVF Winter Conference on Applications of Computer Vision, pp. 2706 2023.
This work was partly supported by the National Key Research and Development Program of China (No. 2024YFB2808900 the Shenzhen Science and Technol- ogy Program 00X1120220181010140000 and the research fund of Tsinghua University- Tencent Joint Laboratory for Internet Innovation Technology.
Dubey, A., Jauhri, A., Pandey, A., Kadian, A., Al-Dahle, A., Letman, A., Mathur, A., Schelten, A., Yang, A., Fan, A., et al. The llama 3 herd of models. arXiv preprint arXi::2407.2178. 2024.
9

---

Critical Tokens Matter: Token-Level Contrastive Estimation Enhances LLM's Reasoning Capability
Feng, D., Qin, B., Huang, C., Zhang, Z., and Lei, W. Towards analyzing and understanding the limitations of dpo: A theoretical perspective. arXiv preprint mrXiv:2404.04626, 2024.
the Association for Computational Linguistics (Volume I: Long Papers), pp. 12286-12312, 2023.
Lightman, H.,, Kosaraju, V., Burda, Y., Edwards, H., Baker, B., Lee, T., Leike, J., Schulman, J., Sutskever, I., and Cobbe, Let's verify step by step. arXiv preprint arXiv:2305.cm 20050, 2023.
Guo, G., Zhao, R., Tang, T., Zhao, W. X., and Wen, J.- R. Beyond imitation: Leveraging fine- une-grained quality signals for alignment. arXiv preprint arXiv:231..0407) 2023.
Liu, Z., Lu, M., Zhang, S., Liu, B., Guo. H., Yang, Y., Blanchet, J., and Wang, Provably mitigating overopti- mization in rlhf: Your sft loss implicitly an adversarial regularizer. arXiv preprint arXiv:2405. 16436, 2024.
Gutmann, M. and Hyvärinen, A. Noise-contrastive estima- tion: new estimation principle for unnormalized statisti- cal models. In Proceedings of the thirteenth international conference on artificial intelligence and statistics, pp. 297-304. JMLR Workshop and Conference Proceedings, 2010.
O' O'Brien, S. and Lewis, M. Contrastive decoding improves reasoning large language models. arXiv preprint arXi::2309.09/ 2023.
Ouyang, L.,. Wu, J., Jiang, X., Almeida, D., Wainwright, C., Mishkin, P., Zhang, C., Agarwal, S., Slama, K., Ray, A., et al. Training language models to follow instructions with human feedback. Advances neural information processing systems, 35:27730-27744 2022.
He, K., Fan, H., Wu, Y., Xic, S., and Girshick, R. Mo- mentum contrast for unsupervised visual representation learning. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pp. 9729-9738, 2020.
Pal, A., Karkhanis, D., Dooley, S., Roberts, M., Naidu, S., and White, C. Smaug: Fixing failure modes of pref- erence optimisation with dpo-positive. arXiv preprint arXiv:2402.13222 2024.
Hendrycks, D., Burns, C., Kadavath, S., Arora, A., Basart, S., Tang, E., Song, D., and Steinhardt, Measuring mathematical problem solving with the math dataset. In Thirty-fifth Conference on Neural Information P Process- ing Systems Datasets and Benchmarks Track (Round 2), 2021.
Pang, R. Y., Yuan, W., Cho, K., He, H., Sukhbaatar, S., and Weston, J. Iterative reasoning preference optimization. arXiv preprint arXiv:2404. 19733,, 2024.
Hu, E. J., Wallis, P., Allen-Zhu, Z., Li, Y., Wang, S., Wang, L.. Chen, W., et al. Lora: Low-rank adaptation of large language models. In International Conference on Learn- ing Representations, 2022.
Rafailov, R., Sharma, A., Mitchell, E., Manning, C. D., Er- mon. S., and Finn, C. Direct preference optimization: Your language model is secretly a reward model. Ad- vances in Neural Information Processing Systems, 36, 2024.
Lai, X., Tian, Z., Chen, Y., Yang, S., Peng, X., and Jia, J. Step-dpo: Step-wise preference optimiza- tion for long-chain reasoning of llms. arXiv preprint arXiv:2406. 18629, 2024.
Schulman, J., Wolski, F., Dhariwal, P., Radford, A., and Klimov, O. Proximal policy optimization algorithms. arXiv preprint vVXiv1107.00347, 2017.
Lambert, N., Pyatkin, V., Morrison, J., Miranda, L., Lin, B. Y., Chandu, K., Dziri, N., Kumar, S., Zick, T., Choi, Y., et al. Rewardbench: Evaluating reward models for language modeling. arXiv preprint arXiv:2403.1 13787, 2024.
Shao, Z., Wang. P., Zhu, Q., Xu, R., Song. J., Zhang, M., Li, Y., Wu, Y., and Guo, D. Deepseekmath: Pushing the limits of mathematical reasoning in open language models. arXiv preprint arXiv:2402.033000 2024.
Lewkowycz, A., Andreassen, A., Dohan, D., Dyer, E., Michalewski, H., Ramasesh, V., Slone, A., Anil, C., Schlag, I., Gutman-Solo, T., et al. Solving quantitative reasoning problems with language models. Advances in Neural Information Processing Systems, 35:3843-3857, 2022.
Shi, C.,, Yang, C.,, Zhu, X., Wang, J., Wu, T., Li, S., Cai, D., Yang, Y., and Meng, Y. Unchosen experts can contribute too: Unleashing moe models' power by self-contrast. arXiv preprint arXiv:2405. 2024.
Wang. B., Min, S., Deng, X., Shen, J., Wu, Y., Zettlemoyer, I.,, and Sun, H. Towards understanding chain-of-thouge prompting: An empirical study of what matters. In Pro- ceedings of the 61st Annual Meeting of the Association for Computational Linguistics (Volume Long Papers), pp. 2717-2739, 2023.
Li, X. L., Holtzman, A., Fried, D., Liang, P., Eisner, J., Hashimoto, T. B., Zettlemoyer, L., and Lewis, M. Con- trastive decoding: Open-ender text generation as opti- mization. In Proceedings of the 61 st Annual Meeting of
10

---

Critical Tokens Matter: Token-Level Contrastive Estimation Enhances LLM's Reasoning Capability
Wei, J., Wang, X., Schuurmans, D., Bosma, M., Xia, F,, Chi, E., Le, Q. V., Zhou, D., al. Chain-of- prompting elicits reasoning in large language models. Advances in neural information processing systems, 11124824-24887 2022.
Zhao, Z.., Zhang, X., Xu, K., Hu, X., Zhang, R., Du, Z., Guo, Q-. and Chen, Y. Adversarial contrastive decod- ing: Boosting safety alignment of large language mod- els via opposite prompt optimization. arXiv preprint arXiv:2406. 16743, 2024.
Xu, S., Fu, W. Gao, J., Ye, W., Liu, W. Mei, Z., Wang. G., Yu, C., and Wu, Y. Is dpo superior to ppo for Ilm align- ment? a comprehensive study. In Forty-fort International Conference on Machine Learning, 2024.
Ziegler, D. M., Stiennon, N., Wu, Brown, T. B., Radford, A., Amodei, D., Christiano, P., and Irving, G. Fine-tuning language models from human preferences, arXiv preprint arXiv: 119999900999. 2019.
Yang, C., Shi, C., Li, S., Shui, B., Yang, Y., and Lam, W. LLM2: Let large language models harness sys- tem 2 reasoning. In Chiruzzo, L., Ritter, A., and Wang, L. (eds.), Proceedings of the 2025 Conference of the Nations of the Americas Chapter of the Asso- ciation for Computational Linguisties: Human Lan- guage Technologies ( (Volume 2: Short Papers), pp. 168- 177, Albuquerque, New Mexico, April 2025. Associa- tion for Computational Linguistics. ISBN 979-8-89176- 190-2. URL www.hhtpss//acllneeee naacl www.thhoottt.. 15/.
org.orggg//// 2025.
Yang, D., Xiao, D., Wei, J., Li, M., Chen, Z., Li, K., and Zhang, L. Improving factuality in large language models via decoding time hallucinatory and truthful comparators. arXiv preprint ar arXiv:2408. 12325, 2024.
Yoon, E., Yoon, H. S., Eom, S., Han, G., Nam, D., Jo, D., On, K.-W., Hascgawa-Johnson, M., Kim, S., and Yoo, C. Tier: Token-level continuous reward for fine- grained reinforcement learning from human feedback. In Findings of the Association for Computational Linguistics ACL 2024, pp. 14969-14981, 2024.
Yu, L., Jiang, W., Shi, H., Jincheng, Y., Liu, Z., Zhang, Y., Kwok, J., Li, Z., Weller, A., and Liu, W. Metamath: Bootstrap your own mathematical questions for large lan- guage models. In The Twelfth International Conference Learning Representations.
Yuan, H., Lu, K., Huang, F., Yuan, Z., and Zhou, C. Specu- lative contrastive decoding. In Proceedings of the 62nd Annual Meeting of the Association for Computational Linguistics ( (Volume 2: Short Papers), pp. 56 64, 2024.
Yuan, Z., Yuan, H., Li, C., Dong, G., Lu, K., Tan, C., Zhou, C., and Zhou, J. Scaling relationship on learning math- ematical reasoning with large language models. arXiv preprint ar Xiv:22308..0 01825, 2023.
Zeng, Y., Liu, G., Ma, W., Yang, N., Zhang, H., and Wang, Token-level direct preference optimization. arXiv preprint 4004v:2404, 1999, 2024.
Zhang, Y., Cui, L., Bi, W., and Shi, S. Alleviating hal- lucinations of large language models through induced hallucinations, arXiv preprint arXiv:2312.157100 2023.

---

Critical Tokens Matter: Token-Level Contrastive Estimation Enhances LLM's Reasoning Capability
A. Contrastive Estimation: Theoretical and Empirical Analysis
A.1. Distribution Analysis of Contrastive Estimation
demonstrate that contrastive estimation does not fundamentally alter the nature of the trajectory distribution. According (Guo al., 2023; Lambert et al., 2024), the trajectory distribution can be modeled as Gaussian distributions based on correctness. Consequently, define the probability density functions for the correct and incorrect distributions, denoted Pe and pm respectively, as follows:
1
(x Up)
(2(x)
exp
(2)
V2TOO
20²
m4n
pa (x(x)
exp
(3)
V2TT
20²
where the means satisfy MP > Mm. and both distributions share the same standard deviation ø, facilitating straightforward comparison between them.
Therefore, according to Eq. I, the probability density function pe of the CE distribution can be calculated as follows:
log( (x)) = (1 ()) log (x))) PM (x:) Blog(P"(x))
(1))
(4)
where Zi is the partition function.
Substituting the definitions of P2(x) and P"(x), we obtain:
(1 1)x pp)
3(x Ma)2
Iog (x)) log V2TO
log (1))
(5)
20
Thus, the CE distribution pce is:
1
(
(1+ B)(x Hpp)
B(x Mn)2
(x) Zi
exp
(6)
V
20²
The term B)xx - up) - B(x - pn) can be expressed as:
(1+ - 2p)² - fz -m)2 (x 1 + Zz.
(7)
where ========== 14p B(1p - Mo). and Z3 is constant independent of .r.
Substituting this result into (6) gives:
(- (x - 11093
Zz
(=( = Zi V2TO
exp
exp
(8)
20
20
where Z1 exp (-2g) (3). Finally, the CE distribution be written as:
exp (- (x - 10000222
pee (x) = /2mo
(9)
20²
Hence, (x) is also a Gaussian distribution with mean n* = Up + 3(pp Pa) and variance o², Here, contrast factor 3 controls the extent to which the distribution shifts away from the negative mean f4 towards the positive mean
12

---

Critical Tokens Matter: Token-Level Contrastive Estimation Enhances LLM's Reasoning Capability
A.2. Effect of the Contrast Factor |
To better understand the influence of the scaling factor contrastive estimation, we conduct ablation study using the -LaMA-3-8B model on the GSM8K. The results shown in Table 4.
Table 4. Ablation study on contrast factor on GSM8K using LLaMA-3-8B.
0.5
0.75
1.0
1.25
1.5
1.75
2.0
2.25
2.5
cDPO (%)
66.5
69.2
67.9
67.2
69.5
70.7
68.4
66.8
68.9
As discussed in Appendix 1, the scaling factor shifts the mean of the modified distribution pee away from the negative mean Mn toward the positive mean μę, following μ Mp 3(4p Ma). The ablation results show that setting 3 in the range of 1.5 to 1.75 leads to better performance, indicating that pe aligns more closely with the optimal distribution when the contrast is moderately emphasized.
B. Futhrer Experiments
B.1. Pass@1 accuracy on MATH500
Table Pass@ accuracy on MATH500, Each result is averaged over 10 samplings per question.
Temperature
0
0.25
0.5
0.75
1.0
1.25
1.5
LLaMA-3-8B
18.6
16.4
15.3
13.0
9.5
3.3
1.2
cDPO
19.6
20.3
20.1
19.6
18.7
19.7
18.3
DeepScck-Math-7B
34.0
31.8
30.5
26.2
21.1
11.3
3.0
cDPO
35.2
34.5
34.5
34.5
33.9
32.9
32.8
Qwen-2.5-7B
49.2
46.9
45.1
41.4
34.0
20.1
2.8
+ cDPO
54.0
54.2
53.6
52.8
52.9
53,4
51.9
B.2, Accuracy of Qwen-2.5 models
Table 6. Accuracy of Qwen-2.5 models GSM8K and MATH500.
Model
GSM8K
MATH500
Qwen-2.5-7B
85.5
49.2
cDPO
87.5
54.0
Qwen-2.5-32B
93.0
58.8
cDPO
93.5
64.8
13


Reviews:
[{"review_id": "Ph77QfTUef", "reviewer": "Reviewer_Qri4", "summary": "This paper introduces the concept of critical tokens, which are tokens that significantly influence the reasoning trajectories, leading to incorrect outcomes. They propose to use rollout algorithm to identify critical tokens, then study the difference between critical tokens and wrong tokens. They further propose an efficient method to detect the critical tokens. Based on these findings, they propose a cDPO method that assign more weights on critical tokens, and show that it improves the performance.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your thoughtful and constructive review. We sincerely appreciate your recognition of our work's novelty and contributions. Your detailed comments are very helpful; we respond to each of your concerns below.\n\n**Concerns:**\n\n**C1. Comparison Between Contrastive Estimation (CE) and the Rollout Algorithm** \n\nThe AUC values of CE compared to rollout sampling on LLaMA-3-8B are as follows:\n1. GSM8K: 0.77\n2. MATH: 0.84\n\nWe will include these results in the next revision. \n\n**C2. Use of Base Model Instead of Instruct-Tuned Model**\n\nWe primarily adhere to setting protocols established by previous studies [1, 2, 3] by using base models rather than instruct-tuned models. This choice serves two main purposes:\n1. To ensure controlled and consistent model evaluations.\n2. To isolate the effects specific to the post-training phase, thereby clearly assessing our method’s direct contribution without interference from any prior fine-tuning effects.\n\n**C3. Detailed Information on Evaluation Setup (COT, k-shot, pass@k, etc.)**\n\nThank you for pointing out the need for clarification. Our experimental configurations strictly follow established practices:\n- 8-shot prompting for GSM8K as in [4], and 4-shot prompting for MATH500 as described in [5];\n- Temperature fixed at 0 for main evaluations.\n\nFurthermore, we perform additional analyses by sampling each question 10 times at varying temperatures and measure Pass@1 accuracy as shown below:\n\n| Temperature (T) | 0 | 0.25 | 0.5 | 0.75 | 1.0 | 1.25 | 1.5 |\n|-----------------|---|------|-----|------|-----|------|-----|\n| LLaMA-3-8B      |18.6|16.4  |15.3 |13.0  |9.5  |3.3   |1.2  |\n| + cDPO          |19.6|20.3  |20.1 |19.6  |18.7 |19.7  |18.3 |\n| DeepSeek-math-7B|34.0|31.8  |30.5 |26.2  |21.1 |11.3  |3.0  |\n| + cDPO          |35.2|34.5  |34.5 |34.5  |33.9 |32.9  |32.8 |\n| Qwen-2.5-7B     |49.2|46.9  |45.1 |41.4  |34.0 |20.1  |2.8  |\n| + cDPO          |54.0|54.2  |53.6 |52.8  |52.9 |53.4  |51.9 |\n\nKey observations:\n- cDPO consistently outperforms baseline models across all temperatures.\n- cDPO demonstrates stable and robust performance under varying sampling conditions.\n\n**C4. Computational Overhead of Training the Negative Model**\n\nWe appreciate your insightful observation. Identifying critical tokens via rollout sampling incurs significant computational overhead. To address this limitation, we introduce contrastive estimation (CE), a computationally efficient alternative utilizing trained positive and negative models to estimate critical tokens. To validate this approach further, we incorporate CE-based scoring into our cDPO strategy and demonstrate its effectiveness experimentally.\n\nLooking ahead, if sufficiently extensive datasets annotated with critical tokens become available, training a dedicated token-level reward predictor model can deliver an even more scalable and lightweight alternative solution.\n\n---\n\n**Questions:**\n\n**Q1. Derivation of Equation (1)**\n\n- *Clarification of Derivation*: As thoroughly discussed in the Appendix \"Distribution Analysis of Contrastive Estimation,\" the term s_t presented in Equation (1) is derived directly as a combination of the probability distributions P^p (positive model) and P^n (negative model). This formulation results in an improved distribution, efficiently suppressing incorrect token likelihoods while promoting correct ones.\n- *Use of Logits as Weights in cDPO*: Further, as detailed in Section 3.2 \"Formulation\", cDPO refines the negative portion of the DPO loss into a token-level loss, using s_t as weighting factors. This mechanism explicitly guides training towards avoiding generation of critical tokens.\n\n**Q2. Timing for Training the Negative Model**\n\nAs depicted explicitly in Figure 3, the negative model is trained prior to initiating cDPO training. Specifically, our full experimental pipeline consists of the following two distinct phases:\n1. Critical token estimation using CE, involving prior training of both positive and negative models.\n2. Integration of CE-derived scores into cDPO training.\n\n**Q3. Implementation Details of the Rollout Method**\n\nWe perform rollout sampling exactly as described in Lines 194–202:\n- Given an incorrect response T = {t_1, t_2, ..., t_n}, we traverse each token t_i.\n- At each position t_i, we generate k=64 sampled continuations from the prefix T_{≤i} , employing the identical sampling configuration used for original generation.\n- The accuracy averaged over these k continuations forms a score for token t_i.\n\n---\n\n**References:** \n\n[1] ARGS: Alignment as Reward-Guided Search. ICLR 2024.\n\n [2] Token-level direct preference optimization. ICML 2024.\n\n [3] Enhancing llm reasoning via critique models with test-time and training-time supervision. arXiv 2024.\n\n [4] Chain-of-thought prompting elicits reasoning in large language models. NeurIPS 2022.\n\n [5] Solving quantitative reasoning problems with language models. NeurIPS 2022.", "claims_and_evidence": "The claims for critical tokens are convincing: The authors demonstrate the existence of critical tokens. The criteria chosen by the authors are very stringent. Nevertheless, they successfully identify a large amount of critical tokens, showing that critical tokens widely appear in LLM generations.\n\nHowever, I find the claims for the efficient identification method less convincing. There lacks comparison between the efficient identification approach and the \"golden standard\", the rollout algorithm. I can only find the efficiency comparison. There should be a table showing the correct identification probability of critical tokens for the efficient algorithm. Even if the two methods do not lead to coherent approach, it will be important to understand the cause of the difference.", "methods_and_evaluation_criteria": "There are several concerns:\n1. The use of base model instead of instruct fine-tuned model: Can the authors comment on the choice of the model? It would make more sense to use the instruct finetuned model as the baseline and for further experiment.\n2. There should also be some details for evaluation (COT, k-shot, pass@k, etc). \n3. Since the negative model serves only as an intermediate step and is discarded afterward, the current approach introduces significant computational overhead through the negative model used for critical token identification. \n\n[1] L. Team and A. Meta, “The Llama 3 Herd of Models,” Jul. 2024. Available: https://arxiv.org/pdf/2407.21783\n‌\n\n‌", "theoretical_claims": "I did not check the theoretical claims as I think the paper mostly focus on empirical applications.", "experimental_designs_or_analyses": "1. Can the authors explain more about the derivation of Equation (1)? For example, why does s_t can be realised with this form? How does this form fit into the weights in the proposed cDPO?\n2. When is the negative model trained? Is it prior to the cDPO training? \n3. There should be more details on the implementation of the rollout method.", "relation_to_broader_scientific_literature": "The critical tokens fit with previous intuitions that some tokens are more important than others, both in the general setting [1], or in mathematical reasoning setting [2]. This may also relate to COT compression problems.\n\n[1] Lin, Zhenghao, Zhibin Gou, Yeyun Gong, Xiao Liu, Ruochen Xu, Chen Lin, Yujiu Yang, Jian Jiao, Nan Duan, and Weizhu Chen. \"Not all tokens are what you need for pretraining.\" Advances in Neural Information Processing Systems 37 (2024): 29029-29063.\n\n[2] Xia, Heming, Yongqi Li, Chak Tou Leong, Wenjie Wang, and Wenjie Li. \"Tokenskip: Controllable chain-of-thought compression in llms.\" arXiv preprint arXiv:2502.12067 (2025).", "essential_references_not_discussed": "None", "other_strengths_and_weaknesses": "None.", "other_comments_or_suggestions": "None", "questions_for_authors": "None", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "eqjdrrH0PY", "reviewer": "Reviewer_KGMH", "summary": "In this paper, the authors start from the observation that the existence of critical tokens will influence the model performance and propose a contrastive estimation method to identify the critical tokens. At last, the authors proposed the cDPO to improve the model performance.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you very much for your constructive feedback and your willingness to engage in further discussions with us. We have responded to each of the issues you raised below and have carefully addressed all your concerns.\n\n**Weakness**\n **W1. Baseline results**\n\nThank you for pointing out this important issue. Upon careful investigation, we discovered that this disparity stems from a formatting inconsistency in the few-shot prompting examples. Specifically, the prompts differed slightly in spacing:\n  - \"Problem:\\n\" (line break)  \n  - versus \"Problem: \" (single space)\n\nThis seemingly minor formatting difference unfortunately resulted in underestimated baseline performance on the MATH500 dataset. We have since corrected this issue. The revised performance numbers are reported below:\n\n| Model               | MATH500 |\n|---------------------|---------|\n| DeepSeek-math-7B    | 34.0    |\n| + cDPO              | 35.2    |\n| Llama-3-8B          | 18.6    |\n| + cDPO              | 19.6    |\n| Llama-3-70B         | 44.4    |\n| + cDPO              | 45.0    |\n| Qwen-2.5-7B         | 49.2    |\n| + cDPO              | 54.0    |\n| Qwen-2.5-32B        | 58.8    |\n| + cDPO              | 64.8    |\n\nWe will update and clarify the results accordingly in the revised manuscript. \n\n **W2. More experiments**\n\nWe agree with your comment on the limited scope of experiments. To address this, we have conducted additional experiments using Qwen-2.5-7B and Qwen-2.5-32B on both GSM8K and MATH500 datasets. The expanded experimental results are detailed below:\n\n| Model           | GSM8K | MATH500 |\n|-----------------|-------|---------|\n| Qwen-2.5-7B     | 85.5  | 49.2    |\n| + cDPO          | 87.5  | 54.0    |\n| Qwen-2.5-32B    | 93.0  | 58.8    |\n| + cDPO          | 93.5  | 64.8    |\n\nAdditionally, we have also evaluated Pass@1 accuracy under various temperature settings. We sampled each question 10 times for each temperature setting and report average Pass@1 accuracy:\n\n| Temperature (T)        | 0    | 0.25 | 0.5  | 0.75 | 1.0  | 1.25 | 1.5  |\n|------------------------|------|------|------|------|------|------|------|\n| Llama-3-8B             | 18.6 | 16.4 | 15.3 | 13.0 | 9.5  | 3.3  | 1.2  |\n| + cDPO                 | 19.6 | 20.3 | 20.1 | 19.6 | 18.7 | 19.7 | 18.3 |\n| DeepSeek-math-7B       | 34.0 | 31.8 | 30.5 | 26.2 | 21.1 | 11.3 | 3.0  |\n| + cDPO                 | 35.2 | 34.5 | 34.5 | 34.5 | 33.9 | 32.9 | 32.8 |\n| Qwen-2.5-7B            | 49.2 | 46.9 | 45.1 | 41.4 | 34.0 | 20.1 | 2.8  |\n| + cDPO                 | 54.0 | 54.2 | 53.6 | 52.8 | 52.9 | 53.4 | 51.9 |\n\nThese results demonstrate clearly that:\n\n- cDPO consistently surpasses the baseline model performance.\n- cDPO maintains stability and robustness across diverse temperature settings.\n\n**W3. Ablation of the proposed methods**\n\nThank you for the valuable suggestion. To strengthen our analysis, we have conducted an ablation study on the hyperparameter β controlling the mean of the contrastive distribution P^{ce}, as explained in the Appendix section “Distribution Analysis of Contrastive Estimation”. Using LLaMA-3-8B on GSM8K, we observed the performance effects of β values, shown below:\n\n| β-value | 0.5  | 0.75 | 1.0  | 1.25 | 1.5  | 1.75 | 2.0  | 2.25 | 2.5  |\n|---------|------|------|------|------|------|------|------|------|------|\n| cDPO    | 66.5 | 69.2 | 67.9 | 67.2 | 69.5 | 70.7 | 68.4 | 66.8 | 68.9 |\n\nThese results indicate that optimal performance occurs when β is set within a range of 1.5–1.75, highlighting the need to appropriately balance the amplification and suppression of token likelihoods during training. We will present a more detailed discussion and expanded analysis on this topic in the next revision of the manuscript.", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "Yes", "theoretical_claims": "NaN", "experimental_designs_or_analyses": "yes", "relation_to_broader_scientific_literature": "NaN", "essential_references_not_discussed": "NaN", "other_strengths_and_weaknesses": "### Strength\n1. The paper is well-written and easy to follow.\n2. The observation and analysis are novel and insightful.\n3. Results of the proposed methods work well.\n\n### Weakness. \n1. My main concern is the correctness of the baseline results. I find the baseline results are lower than those reported in the original paper. For example, in Table 2 of DeepSeekMath, its performance on MATH is 36.2, not 31.4 as reported in the paper. Also, from the Table. 12 in LLaMA3, the LLaMA 3-8B gets 20.3 on MATH, but not 16.8. This is a serious problem, and I think the authors should clarify it in the rebuttal.\n\n2. The experiment is somewhat simple, with only three models and two benchmarks. This also hinders the convincness of the proposed methods. \n\n3. Similar to 2, there are no ablation and few analyses of the proposed methods. \n\nIn a word, I think the paper starts from an interesting observation and convincing analyses. But the results of the experiment part have some problems. So, I give a weak reject to the current version, and I'll adjust the final rating based on the rebuttal.\n\n## update after rebuttal\nThe rebuttal partly solved my concerns. So I raise my score to weak accept.", "other_comments_or_suggestions": "NaN", "questions_for_authors": "NaN", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "fNmA9QceYd", "reviewer": "Reviewer_W2W8", "summary": "The paper introduces the concept of critical tokens in mathematical reasoning tasks, which are pivotal points in incorrect reasoning trajectories that significantly influence the final outcome. The authors propose a novel framework for identifying these tokens through contrastive estimation and they further introduce cDPO. Reducing the occurrence probability of critic tokens through the DPO is straightforward.\n\n## update after rebuttal\nI will keep my positive score because my concerns have been partially addressed.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you very much for your insightful and detailed comments. Below, we address each concern and hope that our responses sufficiently clarify your questions.\n\n**Weakness**\n\n **W1. Critical Token Estimation**\n\nIn our approach, the distribution for critic tokens is approximated using a model trained with incorrect trajectories. Although these incorrect trajectories do not necessarily satisfy the two conditions specified on page 2, our experimental results nevertheless indicate robust and competitive performance.\n\nWe measured the Area Under the Curve (AUC) of Contrastive Estimation (CE) with respect to rollout on LLaMA-3-8B, obtaining values of: (1) 0.77 on GSM8K, and (2) 0.84 on MATH. These metrics and corresponding analyses will be explicitly included in our revised manuscript.\n\n**W2. Ablation Study for Hyperparameter Sensitivity**\n\nThank you for pointing out this important issue. We fully agree that examining hyperparameter sensitivity will strengthen the paper, enhancing its usefulness for practitioners. Accordingly, we have now included an ablation study investigating different values of the scaling factor β using the LLaMA-3-8B model on GSM8K, as summarized in the table below.\n\n| β-value | 0.5  | 0.75 | 1.0  | 1.25 | 1.5  | 1.75 | 2.0  | 2.25 | 2.5  |\n|---------|------|------|------|------|------|------|------|------|------|\n| cDPO    | 66.5 | 69.2 | 67.9 | 67.2 | 69.5 | 70.7 | 68.4 | 66.8 | 68.9 |\n\nAs further elaborated in Appendix section 'Distribution Analysis of Contrastive Estimation,' the hyperparameter β affects the mean of the modified distribution P^{ce}. The above ablation results demonstrate that selecting β in a suitable range (approximately 1.5 to 1.75) shifts P^{ce} toward a more accurate distribution, thereby improving overall model performance.", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "Yes", "theoretical_claims": "Yes", "experimental_designs_or_analyses": "Yes", "relation_to_broader_scientific_literature": "I am not familiar enough with this area to seriously recommend anything.", "essential_references_not_discussed": "I am not familiar enough with this area to seriously recommend anything.", "other_strengths_and_weaknesses": "**Strength**\n\n1. The introduction of critical tokens is a novel and insightful contribution to the field of mathematical reasoning in LLMs. The authors provide a clear definition and empirical validation of these tokens, showing their significant impact on model accuracy.\n\n2. This paper is well-written and easy to follow. The paper provides sufficient technical details for readers to understand.\n\n\n**Weakness**\n\n1. A model trained with incorrect trajectories is used to simulate the probability distribution of critic tokens, but the incorrect trajectories do not necessarily satisfy the two conditions on page 2.\n\n2. Lack the sensitivity of the proposed method to hyperparameters, such as the scaling factor β in contrastive estimation. Understanding how sensitive the results are to these parameters would be valuable for practitioners looking to implement the method.", "other_comments_or_suggestions": "This is discussed in the Strengths And Weaknesses.", "questions_for_authors": "The Question is discussed in the Weaknesses.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: The paper is well-written and easy to follow. The authors propose a simple yet reasonable definition of critical tokens and successfully identify a large number of them in their experiments, demonstrating that critical tokens are prevalent in LLM outputs. They also leverage this method to stabilize DPO training, which adds practical value. Overall, this is a solid piece of work, and all reviewers are in favor of acceptance.
Decision: Accept (poster)