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

Thank you very much for your insightful and detailed comments. Below, we address each concern and hope that our responses sufficiently clarify your questions.

Weakness

W1. Critical Token Estimation

In 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.

We 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.

W2. Ablation Study for Hyperparameter Sensitivity

Thank 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.

As 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.

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.

Weakness W1. Baseline results

Thank 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:

• "Problem:\n" (line break)
• versus "Problem: " (single space)

This 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:

We will update and clarify the results accordingly in the revised manuscript.

W2. More experiments

We 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:

Additionally, 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:

These results demonstrate clearly that:

• cDPO consistently surpasses the baseline model performance.
• cDPO maintains stability and robustness across diverse temperature settings.

W3. Ablation of the proposed methods

Thank 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:

These 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.

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.

Concerns:

C1. Comparison Between Contrastive Estimation (CE) and the Rollout Algorithm

The AUC values of CE compared to rollout sampling on LLaMA-3-8B are as follows:

1. GSM8K: 0.77
2. MATH: 0.84

We will include these results in the next revision.

C2. Use of Base Model Instead of Instruct-Tuned Model

We 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:

1. To ensure controlled and consistent model evaluations.
2. 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.

C3. Detailed Information on Evaluation Setup (COT, k-shot, pass@k, etc.)

Thank you for pointing out the need for clarification. Our experimental configurations strictly follow established practices:

• 8-shot prompting for GSM8K as in [4], and 4-shot prompting for MATH500 as described in [5];
• Temperature fixed at 0 for main evaluations.

Furthermore, we perform additional analyses by sampling each question 10 times at varying temperatures and measure Pass@1 accuracy as shown below:

Key observations:

• cDPO consistently outperforms baseline models across all temperatures.
• cDPO demonstrates stable and robust performance under varying sampling conditions.

C4. Computational Overhead of Training the Negative Model

We 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.

Looking 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.

---

Questions:

Q1. Derivation of Equation (1)

• 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.
• 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.

Q2. Timing for Training the Negative Model

As 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:

1. Critical token estimation using CE, involving prior training of both positive and negative models.
2. Integration of CE-derived scores into cDPO training.

Q3. Implementation Details of the Rollout Method

We perform rollout sampling exactly as described in Lines 194–202:

• Given an incorrect response T = {t_1, t_2, ..., t_n}, we traverse each token t_i.
• 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.
• The accuracy averaged over these k continuations forms a score for token t_i.

---

References:

[1] ARGS: Alignment as Reward-Guided Search. ICLR 2024.

[2] Token-level direct preference optimization. ICML 2024.

[3] Enhancing llm reasoning via critique models with test-time and training-time supervision. arXiv 2024.

[4] Chain-of-thought prompting elicits reasoning in large language models. NeurIPS 2022.

[5] Solving quantitative reasoning problems with language models. NeurIPS 2022.