Title: Rethinking Fine-Tuning when Scaling Test-Time Compute: Limiting Confidence Improves Mathematical Reasoning
Venue: NeurIPS 2025
Accepted: yes
Keywords: Test-time compute

inference-time compute

coverage

pass@N

reasoning

large language model

formal math

finetuning

overfitting

overconfidence
Abstract: Recent progress in large language models (LLMs) highlights the power of scaling test-time compute to achieve strong performance on complex tasks, such as mathematical reasoning and code generation. This raises a critical question: how should model training be modified to optimize performance under a subsequent test-time compute strategy and budget? To explore this, we focus on pass@N, a simple test-time strategy that searches for a correct answer in N independent samples. We show, surprisingly, that training with cross-entropy (CE) can be _misaligned_ with pass@N in that pass@N accuracy _decreases_ with longer CE training. We explain the origins of this misalignment in terms of model overconfidence induced by CE, and experimentally verify our prediction of overconfidence as an impediment to scaling test-time compute via pass@N. Furthermore we suggest a principled, modified training loss that is better aligned to pass@N by limiting model confidence and rescuing pass@N test performance. Our algorithm demonstrates improved mathematical reasoning on MATH and MiniF2F benchmarks under several scenarios: (1) providing answers to math questions both with and without Chain-of-Thought reasoning traces; and (2) proving theorems by searching over proof trees of varying shapes. Overall our work underscores the importance of co-designing two traditionally separate phases of LLM development: training-time protocols and test-time search and reasoning strategies.


Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper investigates an misalignment between the standard fine-tuning objective for large language models (cross-entropy loss) and the widely used test-time strategy of pass@N sampling. The authors demonstrate, both theoretically and empirically that training with cross-entropy loss can lead to model overconfidence. This overconfidence, while improving pass@1 accuracy somewhat paradoxically harms performance when sampling a larger number of solutions (pass@N for large N), as the model becomes less exploratory and repeatedly generates the same incorrect high-confidence answers. To address this, the paper introduces Direct Coverage Optimization, a loss function that directly optimizes the pass@N metric during training. By acting as a confidence regularizer, DCO prevents the model from becoming overconfident in its solutions. The authors show that DCO not only rescues the degrading pass@N performance but also establishes a Pareto-optimal frontier, improving results on mathematical reasoning benchmarks like MATH and theorem-proving benchmarks like MiniF2F, for both direct-answer and chain-of-thought settings.\n\nSTRENGTHS AND WEAKNESSES:\n**Strengths**\n\n1. The paper addresses a highly relevant and overlooked problem in LLM training for test-time scaling. The core finding that more training via standard CE loss can actively harm performance for certain inference strategies is somewhat counter-intuitive at first but makes sense on further consideration. The paper compellingly argues for a tighter co-design of training and inference protocols.\n\n2. The paper is exceptionally well-written and clear. The authors do an excellent job of building intuition for the problem, starting with a simple two-problem example and generalizing with theoretical lemmas (4.2, 4.3). The proposed solution, DCO is elegant and directly motivated by the identified issue of overconfidence. The connection between the DCO gradient and the \"overconfidence regularizer\" (Figure 2b) provides a clear and satisfying explanation of why the method works.\n\n3. The empirical validation is good. The authors test their hypothesis across multiple scenarios including direct answers, chain-of-thought reasoning on the MATH dataset, and automated theorem proving on LeanDojo. The inclusion of comparisons to both standard CE loss and a reinforcement learning baseline (GRPO) is good. The visualizations of confidence distributions (Figure 1) and performance frontiers (Fig 2) are effective and strongly support the paper's claims.\n\n**Weaknesses**\n\n1. The central weakness of the paper is the question of whether the core findings and the effectiveness of DCO hold at the scale of state-of-the-art models. While the authors include results for a Llama-3-70B model in the appendix (Section C.3, Figures 7 and 8) even this is not near the scale of state-of-the-art. The bulk of the results, figures, and discussion are based on the 8B model. I wonder as models become significantly larger and more capable, does this overconfidence issue persist with the same severity? It is plausible that larger models might inherently have better-calibrated uncertainty, potentially diminishing the performance gap between CE and DCO. \n\n2. The entire premise of the pass@N strategy and the DCO objective relies on the existence of a binary verifier to determine the correctness of a generated sample. The authors acknowledge this limitation in their discussion. However, this is a major constraint that limits the applicability of the proposed method to specific domains like math, coding, and formal theorem proving. The approach is not directly usable for more open-ended tasks.\n\nQUESTIONS:\n1. The results for Llama-3-70B in the appendix are encouraging as they show a similar trend to the 8B model. Could you provide more insight into the magnitude of this effect at scale? Does the performance gap between DCO and CE loss widen, shrink, or remain constant when moving from 8B to 70B? To strengthen the paper's claim of generalizability, I strongly encourage you to integrate the 70B results more centrally into the main paper, perhaps directly compare the confidence distributions and performance frontiers across scales. \n\n2. The results show that the optimal DCO hyperparameter N' is closely tied to the test-time pass@ N. This suggests that one might need to train multiple models to be prepared for different inference-time compute budgets. What is your recommendation if the exact test-time N is not known in advance? Is there a \"sweet spot\" N' that provides robust performance across a wide range of test N values? Or is ensembling models trained with different N' values (as you did for theorem proving) the most practical approach?\n\n3. While the work focuses on domains with clear verifiers, have you considered how DCO might be adapted for tasks where verification is \"soft,\" such as using a reward model's score? How might the DCO loss function behave if the correctness signal p(y|x) is not binary (0 or 1) but a continuous score from a reward model?\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "CLARITY": 4, "IMPACT": 2, "ORIGINALITY": 3, "DATE": "2025-06-09"}

comments: SUMMARY:
The paper investigates the misalignment between cross-entropy (CE) fine-tuning and test-time compute scaling (e.g., pass@N). It shows that prolonged CE training induces overconfidence, degrading pass@N performance despite improving pass@1. The authors attribute this to a tradeoff between exploitation (high confidence) and exploration (diverse sampling). To realign training with pass@N, they propose Direct Coverage Optimization (DCO), a loss that limits model confidence via a regularization term attenuating gradients for high-confidence predictions. Experiments on MATH (direct answers/CoT) and LeanDojo (theorem proving) demonstrate that DCO rescues pass@N scaling and achieves Pareto-optimal frontiers. Extensions to step-wise DCO for proof trees and approximate DCO for CoT further validate the framework.

STRENGTHS AND WEAKNESSES:
Strength

- Novel Insight: Identifies a counterintuitive phenomenon—longer CE training harms pass@N performance—and traces it to overconfidence, supported by theory (Lemmas 4.2–4.3) and empirical validation (Table 1, Figure 1).

- Practical Impact: The ensemble of step-wise DCO models (Table 2) boosts theorem-proving success by +5.2% (MathLib) and +4.1% (MiniF2F) over CE with matched compute, highlighting real-world utility.

- Theoretical Depth: Lemmas 4.2–4.3 formally characterize the confidence-exploration tradeoff, and Appendix A provides complete proofs.

Weakness

- Approximate DCO for CoT requires online Monte Carlo estimation, increasing compute by 10× vs. standard fine-tuning.

- The paper links overconfidence to reduced pass@N performance but does not quantify how DCO affects output diversity.

QUESTIONS:
See the Weakness.

LIMITATIONS:
yes
TITLE: Official Review
CLARITY: 3
IMPACT: 3
ORIGINALITY: 3
DATE: 2025-06-30

comments: SUMMARY:
This paper focus on the problem that after few epochs of training, while the pass@1 result of model increases, pass@N decreases. It claims a mismatch between how we train LLMs (with cross-entropy loss that optimizes single best output) and how we test them (by pass@N), and the standard fine-tuning makes the model too confident at one guess and hurts performance when we sample many times. They propose a new loss objective, Direct Coverage Optimization (DCO), which directly maximizes the chance of getting at least one correct answer in N samples, showing gains in  pass@N on math problems and theorem proving benchmarks.

STRENGTHS AND WEAKNESSES:
Strengths: 
 - The paper provides a insightful observation and points out a real gap between training and inference and explains why it matters, with clear proofs show why cross-entropy leads to overconfidence and how DCO fixes it.
 - The performance of DCO and its variants are good on MATH, AIME, MiniF2F, and Mathlib, cross multiple domains.
 - The method applies to both step-wise proof search and chain-of-thought, showing robustness.

Weaknesses:
 - Training with DCO may need more sampling and GPU time, which is acceptable.

QUESTIONS:
In the paper, the authors train with N’ = N to match the test-time sampling budget. But retraining a separate model for every target N can be very costly. How robust is a model trained for one N’ when tested at different budgets (e.g., training at N’=16 but testing at N=64 or N=4)? Have you explored stragety to support multiple budgets without full retraining for each?

LIMITATIONS:
yes
TITLE: Official Review
CLARITY: 3
IMPACT: 4
ORIGINALITY: 3
DATE: 2025-07-01

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper investigates a subtle but important mismatch between how LLMs are fine-tuned and how they are later used with heavy test-time sampling (e.g. pass@N). The authors first show that training with standard cross-entropy (CE) pushes the model toward over-confident, low-diversity distributions that hurt coverage when many samples are drawn. On MATH, for instance, pass@256 accuracy falls from 65% after one epoch to 45% after four epochs despite pass@1 steadily improving. They trace the effect to a theory-backed exploit/explore trade-off and introduce Direct Coverage Optimization (DCO), a simple loss that directly maximizes the probability a correct answer appears in N samples while automatically regularizing over-confidence. Results on math are quite promising.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths:\n - Solid theoretical analysis explaining why CE induces over-confidence;\n - Writing is clear; figures/tables make trade-offs visually clear\n - Method is simple and easy to implement\n - Topic very relevant today\n\nWeakness:\n - Related work on inference-aware fine-tuning (e.g. BonBon, BOND) is acknowledged but could be contrasted more deeply, especially on diversity vs. confidence control\n - There are many previous loss functions that try to regularize cross entropy to avoid over confidence. Focal loss, which is commonly used in computer vision, is one such example. Could you discuss (or maybe even compare to) those losses?\n\nQUESTIONS:\n- Regarding N: Could you provide a rule-of-thumb (perhaps tied to available inference budget) to pick these hyper-parameters without sweeping? A calibration plot would make adoption easier.\n - The epsilon part is a bit sketchy. Is it really necessary? What's the rule of thumb for that? \n - As sequence length gets longer, the prob of the sequence naturally gets lower. How does that affect things?\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-07-04"}