Rethinking Fine-Tuning when Scaling Test-Time Compute: Limiting Confidence Improves Mathematical Reasoning
Authors: Feng Chen
Allan Raventos
Nan Cheng
Surya Ganguli
Shaul Druckmann
Venue: NeurIPS 2025 poster ()
Keywords: Test-time compute
inference-time compute
coverage
pass@N
reasoning
large language model
formal math
finetuning
overfitting
overconfidence
Primary area: deep_learning
TL;DR: We show that limiting a model's confidence during training can improve test-time scaling in mathematical reasoning.

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.

Rethinking Fine-Tuning when Scaling Test-Time Compute: Limiting Confidence Improves Mathematical Reasoning
Feng Chen"
Allan Raventós" Stanford University Stanford, CA 94305 aravento@stanford.ccom
Nan Cheng University of Michigan Ann Arbor, MI 48 48109 nancheng@umico ch. edu
Stanford University Stanford, CA 94305 fengc@stannfrre ord. edu
Surya Ganguli Stanford University Stanford, CA 94305
Shaul Druckmann
Stanford University
Stanford, CA 94305
gganguli@stanfi edu
shauld@stanforms edu
Abstract
Recent progress large language models (LLMs) highlights the power of scaling test-time compute to achieve strong performance on complex tasks, such as mathe- matical reasoning and code generation. This raises 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, simple test-time strategy that searches for a correct answer in N independent samples. We show, surprisingly, that training with cross-entropy (CE) loss can be misaligned with pass @N N in that pass @ N accuracy decreases with longer training. We explain the origins of this misalignment in terms of model overconfidence induced by CE, and experimentally verify our prediction of overconfidence scaling test-time compute via pass@N N. Furthermore we suggest a principled, modi- fied training loss that is better aligned to pass @N by limiting model confidence and rescuing pass @ test performance. Our algorithm demonstrates improved math- ematical reasoning on MATH and MiniF2F benchmarks under several scenarios: (1) providing answers to math questions; 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.
impediment
Introduction
Scaling test-time compute has been integral to unprecedented improvements in LLMs' reasoning skills for complex tasks such math and coding. Thus, test-time compute has emerged new dimension for improving LLMs, leading to key tradeoff between allocating additional compute to inference versus pretraining 11. Diverse test-time strategies include Chain-of-Thougee (CoT) (2). thc-outth B. self- consistency [4) self-reflection (5), self-critique (6), verfficriiinntoon 7 and Monte-Cari tree search [8). These have shown great success in boosting model performance in the post-train pest-training phase or at inference time. More recently. OpenAl's s O1 model [9] and DeepSeek's 's R. model 10 have combined some of these strategies with reinforcement learning to generate high-
Equal Contribution
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

quality reasoning traces for problems of various difficulty levels, demonstrating clear performance improvements as more test-time compute allocated.
These successes fit into a broader paradigm in which a frontier model is first fine-tuned on a reasoning task with supervised fine-tuning (SFT) 2 12], and then test-time algorithm is applied to model outputs or reasoning traces to improve performance [3 [3. Many test-time algorithms are independent of the fine-tuning process. As result, the fine-tuning is agnostic and thus decoupled from the test-time algorithm [14.. However, for a given choice of test-time strategy and compute budget, is not a priori clear which fine-tuning approach, including the loss objective, would be best aligned with the test-time strategy so as to maximize the test accuracy under the overall strategy.
Our work studies the problem of aligning fine-tuning with test-time algorithms. We consider what is perhaps the simplest setting, SFT with CE loss and pass@ N as the test-time strategy. This setting reveals case of misalignment: standard SFT is not maximizing performance under pass @ N. We believe that this misalignment presents itself in several combinations of fine-tunin/test-timee approaches, motivating our thorough study in this paper. Our main contributions are,
We identify a misalignment between standard fine-tuning with CE loss and the pass@N coverage metric test time (Section 4.1p.
We develop and experimentally verify a framework that suggests this misalignment arises from overconfidence induced by training on CE loss (Sections and
We propose a loss function that directly optimizes the pass @N coverage metric, demon- strating consistent improvement over the CE loss objective and achieving superior accuracy frontiers on MATH and MiniF2F (Sections 5.1 to 5.3
We extend our algorithm to more complex test-time scenarios including searching over proof-trees of varying shapes to improve automated theorem proving, and answering math questions using Chain-of-Thouglh reasoning traces, demonstrating improved mathematical reasoning in both cases (Sections 5.3 and 5.4).
2
Related Works
Test-time compute and P strategy. A growing body of work explores how training and test-time strategies interact to shape model performance. Jones [15] highlight tradeoff between train and test-time compute in a controlled board game setting. Brown et al. [16) identify an exponentiated power law between coverage and number of in-context samples, captured phenomennlooically in Levi
[17 and Schaeffer et al. [18] while Snell et al. [1] explore compute-optimal strategies for scaling test-time inference. Our paper focuses primarily on the pass@) N test-time strategy. Gui et al. [19 show that, in the best-of-N setting, the sample distribution is well-aligned with the reward model, and alignment methods have been proposed to distill this distribution and reduce sampling costs [19] 20
However, their work focuses on improving the pass@1 I performance from the Best-of-N policy, while we focus on directly improving the pass @ N performance. Closely related our work. Chow et al. [4]] derives a similar training objective for reinforcement learning (RL) to directly optimize for the best-of-N test-time strategy. Li et al. (2 argue that fine-tuning with cross-entropy loss limits output diversity and propose a maximum entropy-based method to address this issue. Our work similarly identifies limitations of cross- coosopy when scaling test sst-time compute, but focuses on the alignment of training and test-time strategies. In addition, Yue et al. [22] argue that fine-tuning with Group Relative Policy Optimization (GRPO) does not improve reasoning capabilities beyond the base model when evaluated with pass @N N strategy. We will show that GRPO training may also lead
to overconfidence in Figure 3
Post-training for mathematical reasoning, Various post-training techniques have been developed to improve mathematical reasoning in LLMs. Instruction-tume and reinforcement learning with human feedback boost model performance on math tasks (3325 25 while continued training on domain-
specific math or code data also boosts downstream reasoning (3). Rejection-sampe
B1
and seff-improvement (2) methods augment training data for SFT. Recent approaches ( 10 have incorporated RL to achieve exceptional reasoning capabilities. GRPO 10 29 is newly proposed RL algorithm demonstrating strong improvements in mathematical reasoning. Several recent studies have further addressed response length and problem difficulty biases [3), as well as conducted extensive analyses of accuracy and response length across wide variety of base models [33, 34).

---

While our primary focus is SFT, our loss function can be applied to other settings that train under CE loss, such as continual training, pnrrucctiontuninng and data augmentation.
Data pruning and hard example mining. The loss function we derive below can be viewed from the lens of data pruning and hard example mining. Data selection is used to curate high-quality datasets for pretraining [35], where Sorscher et al. (36) show that pruning easy samples can improve loss scaling as a function of dataset size. Zhou al. (37), Ye et al. [38] demonstrate that, in SFT, even small datasets can yield strong alignment and reasoning performance. Hard example mining focuses on identifying and emphasizing challenging samples to improve model performance [3]] In the domain of mathematical reasoning, Tong et al. [40] find a difficulty imbalance in rejectionnsamplee datasets and show that more extensive training on difficult samples improves model performance.
3 Problem setup
Given a vocabulary set W. consider dataset D = {(x(0, y )
where xx()) E W is prompt, y E W is its ground-truth completion, and n, and m, are the prompt and completion lengths. In the context of math, (x(9) is the problem statement and y() is its solution. To model the conditional distribution ((x())((()) we use autoregressive transformer model [41], which is traditionally trained by minimizing the cross-entropy loss
Co==
E log ((yl())
(1)
(x
where ý denotes the model's distribution.
To use and evaluate the model at test time, we assume the existence of an efficient oracle verifier V which takes input an (zz, y) pair and returns Vzz, 1 if y is correct completion of z and otherwise. Practical examples of verifiers include compilers or pre-defined unit tests for coding problems, automatic proof checkers in mathematical theorem proving. In such applications, simple method, known as pass@N, for trading test-time compute for accuracy involves sampling N completions from P given the test prompt z and applying the verifier V to all of them to search for a correct solution. The probability of a correct answer is then no longer the probability that 1 completion is correct, but rather the probability that at least one of N is correct. This probability, for a dataset D, is given by the pass @ N coverage metric
Cº =
E H4 p( x))
P(3j E [N] V(, = =1).
(2)
JZ D.g.
Minimizing the CE loss in Equation is equivalent to maximizing the pass@1 metric Cp on training set D. But if we scale up test-time compute so that the pass pass@N metric CD on test set 'D for N > is the relevant performance metric, Lee still good training loss, or can we do better?
4
Misalignment between CE loss and pass @N
4.1 The CE loss induces overfitting for pass@N
To understand the impact of training with CE loss on pass@1 @N test performance, we fine-tune Llama-3-8B-bass 42 on the MATH 43] dataset. We start from the base model rather than Llama- 3 InSt-Inst avoid potential leakage of the MATH dataset into Llama-3-8B 3-8B-Instruect through post-training. We follow Lightman et al. (24 and use 12, 000 problems for training and the remaining 500 for testing. Here train the model to provide a direct answer without reasoning trace. We will discuss training with CoT in Section 5.4 We confirm that the results presented in this section also hold for larger model, Llama-3-708 base, and direct readers to Appendix C.33 for additional results.
Table reveals that the pass N performance C. on test set monotonically increases with the number of training epochs only for N = 1, when minimizing CE loss is equivalent maximizing CD- However, for N 2 16, minimizing CE loss during training does not monotonically increase CN at test; indeed for N 2 256, @ test performance, remarkably, monotonically decreuses with the number of training epochs, despite the fact that pass@) I performance monotonically increuses. This effectively corresponds to a novel type of overfitting, in which test, performance degrades over training, likely due to mismatch between the test time (pass @ N) and training time (pass @ I) strategies.
3

---

4.2 Overfitting, confidence, and explore-exploit tradeoff
What are the origins of this overfitting? First we show that in the simple case of a single problem s, overfitting cannot occur. Let pox)) denote the probability assigned by the model to all correct answers for problem I. Then the pass 1 coverage is C f(z)) while the pass@N N coverage is N 1 (1 - C This formula for C in the single problem case obeys:
CN =1 -
Lemma 4.1. VN, N' > 0, C.N is monotonic in CN"
Thus increasing pass N' coverage implies increasing pass @ N coverage for any N and N'. When N' 1, this implies minimizing CE loss maximizes pass @ N coverage.
However, Lemma 4.1 can fail when there is more than one problem. To understand this in a simple setting, consider a test set with two problems and I2 with unique correct answers yı and 1/2- Consider two models. The first model assigns probabilities P1 (11|=)) - and Pn (W2lzz) = 0. If we think of the probability model assigns to answer its confidence in that answer, this model is highly confident and correct for problem I1, but highly confident and wrong for æ2- Its pass @1 coverage is thus 50%. Moreover, since it always gets 21 right and 22 wrong, its pass @ N coverage remains at 50%. In contrast, consider a second model which assigns probabiliites (2 (m1) (12 (22) 0.1. This model has high confidence (0.9) but unfortunately answers. Therefore its pass@1 accuracy is only 10%. However, is willing to explore or hedge by placing some low confidence (0.1) on the other answer, which happens to be correct. Thus if this model samples N times, the pass @ N coverage increases with N, eventually approaching 100% as N 0o. Thus the first model outperforms the second in terms of pass I but not pass @ N for large N. This indicates tradeoff amongst policies: those that do better on pass @ may not do better at pass @. N. Of course the best one can do is be confident and correct, corresponding to the optimal model with 6- (y1 (x1) = 5" (82)22) = 1. However, this toy example reveals that if one cannot guarantee correctness, it can be beneficial limit confidence and explore more solutions.
incorrect
To demonstrate more generally the existence of tradeoffs between confident exploitation of a few answers versus unconfident exploration of many answers, function of the number of passes N, we prove two lemmas. To set up these lemmas, given any problem, let f. denote the model probability or confidence assigned answer i, and assume answers are sorted from highest lowest confidence so that P; 2 pi+1, Vi 2 1. Thus ÿi is the model's maximal confidence across all answers. Moreover, let pt be the probability across all the problems that the 7th ranked answer is actually correct. Assume the model policy is approximately well calibrated so that higher confidence implies higher or equal probability of being correct, i.e. p. 2 Pi+1, Vi 2 We empirically verify that the model trained with CE loss on the MATH dataset approximately satisfies this assumption (Figure [3]. Indeed, optimal policies maximizing pass @ N coverage in Equation (2) are approximately well calibrated (See Lemma I.I). P1 is then the model's maximal accuracy across all answers. We prove:
Lemma 4.2 (Upper bound on max confidence). Assume max accuracy Pt, and assume El_ 1 - E for some 0 < e < 1--1-- P1. Then optimal policies maximizing pass@1 @N coverage in Equa- tion (2), subject to the above accuracy and calibration constraints, have confidence upper bound u( N, P1, e, k), which is a non-increasing function of N.
We proved this upper bound as well as the proof in Appendix A decreasing in N (Figure 12p. implying that at large N, an optimal policy for pass @ N must limit its
This upper bound is monotonically
Table 1:
Pass @N coverage on the MATH test set for a Llama-3-8B-Bl model fine-tuned with CE loss on direct answers from the MATH training set. While pass@) @ 1 improves with continued training, for large N, surprisingly, P decreases with number of training epochs.
pass@1
pass@ 16
pass @ 256
pass @ 4k
Epoch 1
4.4%
30.0%
65.2%
82.5 %
Epoch
5.3%
31.4%
64.5%
80.0%
Epoch 3
6.5%
28.7%
54.5%
79.2%
Epoch
7.4%
22.9%
44.5%
63.0%

---

(8)
CE Loss
(b)
CE Loss
DCO (N')
DCO
150
E
Incorrect
Epech
Epoch
Correct
Epech
I
Epoch
Epech
M-= 256
Counts
100
Epoch
50
-8
-8
pressy)
prea((
logptymmeay/
prepo(preooyx))
Figure 1: A model trained with CE loss becomes overconfident in its greedy completions, which harms its pass@N N coverage; our proposed DCO objective limits this overconfidence. We fine- tune -lam--3-BB base model the MATH dataset produce direct answers. Ogreedy is the model's greedy completion when choosing the most likely token at each sampling step. (a) The model trained with CE loss assigns progressively larger confidences p P(Ugresdy (2)) to its greedy completions over the course of training.
(b) At the end of training, only small portion of the model's highly confident completions correct. This will harm the model's pass @ N performance when scaling up N. (c) Same as (a) but shown for the DCO loss with N' 256. The model trained on DCO shows a much milder overconfidence effect. (d) The confidence distribution greedy completions after four epochs with DCO for various choices of N'. As N' increases, the model's 's confidence in its greedy completion more stringently limited, directly as consequence of the overconfidence regularizer F'.
max confidence to be low, thereby favoring exploration and discouraging exploitation. Furthermore, we prove:
Lemma 4.3 (Lower bound on max confidence). Assume top two accuracies pı > P2. Then optimal policies maximizing pass@N N coverage in Equation subject to above accuracy and calibration constraints, must have max confidence obeying
(1 PI
+
p2
P1
P2 --T
11 21
1
P1 P2
This lower bound a monotonically decreasing function of N (Figure 12 implying that at small N, optimal policies must have high max confidence, thereby favoring exploitation and discouraging exploration. Note in the limit N + 1+, the lower bound always 1, recovering the intuitive result that the optimal policy for pass @ is to place 100% confidence the highest accuracy answer.
4.3 Overconfidence prevents performance gains from scaling test-time compute
Lemma suggests that at large N it is beneficial to unconfidently explore by assigning low model confidence to many answers, while Lemma 4.3 suggests that at small N is beneficial confidently exploit by assigning high model confidence to one or few answers. These lemmas make a prediction that could explain the empirical observation in Table T namely, maximization of pass@1 I makes the @..
model overconfident, thereby preventing good performance
To test this theoretical prediction of model overconfidence, estimated the confidence of the model as p(Ugreedy (r) where Vggredy is the greedy completion to obtained by sampling autoregres- sively from the model p. and at each step selecting the token with the highest probability. p(ygreedy (x) approximates the max confidence $1 in Lemmas 4.2 and 4.3 For the model fine-tuned on MATH with CE loss, we plot the distribution of p(ygreedyy x over the test set in Figure I (a). This demon- strates the model becomes progressively more confident about its greedy completion over training, thereby confirming our theoretical prediction. In (b) we see why this a problem: only some of the model's 's highly confident answers are correct. Thus the model becomes overconfident and largely wrong. Scaling test-time compute cannot easily rescue such a model, as over multiple samples, it is likely to confidently provide the same wrong answers, explaining the origins of poor p test performance. In Appendix C.4 we observe similar trend of overconfidence on the out-of -ut-o--distribution test set AIME24, further highlighting the generality of this failure mode.
5

---

5 Direct Coverage Optimization
5.1 A solution that naturally prevents overconfidence
The misalignment between CE training loss and pass@N coverage suggests a simple solution: directly optimize pass @N coverage for each training example at training time, whenever the pass @N strategy is to be used at test-time. We thus propose the Direct Coverage Optimization (DCO) objective, Le B(z,y)~D DCo (2, y), where
(1, y) = log (1 (1- (()=)))))))))
the log probability that the model produces y at least once N samples given x. Thus for each prompt and correct completion y in the training set, we maximize the probability y is found in N passes. This loss naturally prevents model overconfidence, as can be seen via its gradient
V Vol (x,,) v) = F (N.l()) VolCR(I, v)
(4)
where cE is the standard CE loss on a single example, and F(N, p(y]2)) = N ((1-))))))"-11 Maula)
1- 1-(1-3 Plu [()z)
is an extra overconfidence regularization factor that multiplies the standard CE gradient. Note that F(N. ((N,)(y)))) = for N = 1, so DCO reduces to CE for pass@1. 1. Furthermore F(N, p(y]=)) monotonically decreases in the model confidence ( (see Figure 2 (b)). Thus gradients for examples (ar, y) on which the model is more confident are attenuated, and this attenuation is stronger for larger N. This justifies the interpretation of F(N, p(y(=)) as a regularizer that prevents model overconfidence. Indeed, for large N. F(N, p(y) model confidence on an example exceeds 1/N, its gradient becomes negligible. Thus interestingly. aligning training and test better through DCO naturally yields a simple emergent regularization of model overconfidence, which was itself identified as an impediment to performance gains through scaling test-time compute in Figure I (a). Indeed, F(N, ((yy))) is smaller for easier examples (Figure 14, right), effectively regularizing their contribution. This regularization mitigates the potential detrimental effects of overconfidence from easy examples as discussed Appendix [C.]
for confidence p(y]x) Z L/N. As soon as
In practice the introduction of F can lead to some samples in a batch contributing minimally to the current gradient step. To maintain stable effective batch size, we introduce threshold e, and if for a given example, (I, ((yl))) < e, we replace the example with a new one.
5.2 DCO can prevent overconfidence and rescue test-time scaling
We next test whether DCO can rescue test-time scaling by preventing model overconfidence. We first perform experiments on the MATH dataset in this section, and then in the next section we perform experiments on the LeanDojo automated theorem proving benchmark 44 We fine-tune Llama-3-88 base model for epochs on the MATH training set using DCO for N' 256 and confirm that the model is much less confident in its greedy completion than when trained with CE loss after multiple epochs (compare Figure] | (a) and (c)). Moreover, training with DCO at larger N yields lower model greedy confidences at the end of training [ consistent with Lemma 4.2
We next assess pass @ test performance as function of N for models trained by minimizing 4Dco for different values of N' (Figure 2 (a)). For any given N, there is an optimal N' for the training loss LoCO that maximizes pass@1 test coverage, yielding a Pareto optimal performance frontier (black curve) that is achieved when N' is close to N. In particular the model trained with CE loss (equivalent to pass@1 maximization at training) performs poorly relative to the Pareto frontier at large N (red curve below black at large N). Conversely, models trained with DCO at large N' perform poorly relative to the Pareto frontier for small N (green curves below black at small N).
Together these results indicate that the alignment of the training loss CN with the test time strategy pass@N, N, with N' close to N, is crucial for obtaining Pareto optimal performance. Moreover, these results once again confirm the tradeoff between exploration and exploitation, with good performance using pass @ N test-time strategies requiring high (low) exploration with low (high) confidence at large (small) N.
5.3 Improved theorem proving via ensembled tree search through modified step-wise DCO
To further test our method in realistic settings with a verifier, we conduct experiments theorem proving using an interactive proof assistant on the LeanDojo benchmark [44] extracted from the math

---

(n)
219
(b)
1.00
(c)
Small Neff
Large Ness
Coverage
0.8
FIM, ply/x))
Theorem
State
Theorem
0.6
Tactic Edge
Concet proof
I Trused
0.25
0.2
CE
0.0
DCO Fronter
0.00
10a
10a
10a
10"
10
10P
0.0
0.5
1.0
p(y(k)
Figure 2: (a) DCO improves on CE loss for pass@N N test coverage over a broad range of N and traces a Pareto-optimal frontier. We fine-tune Llama-3-8B base models on MATH to produce direct answers: one with CE loss and others using CDco for various N' (color-coded). Each curve shows pass@N coverage for single fine-tuned model. Note that no N' is optimal for all N. The black curve is a Pareto-optimal performance frontier traced by the max of coverage curves for DCO over all N'. (See F which highlights N' = 256, and Figures & and 10 10 Llama-3-70B and AIME24 results.) (b) DCO limits overconfidence by attenuating gradients for examples on which the model is highly confident. We plot the confidence regularization factor F in Equation (4). For CE loss (N- 1). F = regardless of confidence p(y) (z). For N > 1, F decreases with p(y (z: and drops to zero around 1/N 1/N-once confidence on an example reaches (11/N), FF vanishes, preventing it from further increasing. (c) Inverse confidence Nar at training time controls search-tree exploration shapes at test-time. We provide schematic proof search trees for models trained with DCOstep for small (left) and large (right) Nett Circles are proof states; edges are tacties; solid edges are explored, while dashed edges are not. Small New yields a narrow search tree that may miss correct proof' (red, dashed), while larger Natr expands the tree to include the correct proof. However, if the tree becomes too wide (N pass Mass@N N coverage decreases as CN x N/N eff Thus Neft chosen at training time, is a powerful knob to control the tradcoff between exploitation and exploration at test time.
off >> N for a k-step proof), sampling the correct path becomes unlikely and
In this task, at each step i of the proof, the model is prompted with the current proof state x =[i], and outputs a proof tactic y with a trainable model probability p[y[i] p[y[i]|=[i]) The proof assistant takes the sampled tactic yli). verifies whether it is a valid tactic, and if so, returns the next proof state æ/i + 1] after the tactic y|i] is applied. At test time this interactive process between the model and the assistant continues until either: (1) it terminates in an invalid tactic; (2) hits a maximal allowed search depth set by computational constraints; or (3) terminates a successful proof of the initial proof goal as verified by the proof assistant.
library of LEAN4 45
LEAN4 (46) as the formal proof language. We start from Qwen2.5- M5th-1.5B (28). and fine-tune it on the LeanDojo benchmark [4] In contrast to solving math problems above where the model first autoregressively generates an entire answer y according to p(ylz) and then the entire y is verified for correctness, in theorem proving we must obtain correct tactics ylo]. as verified by the proof assistant, at every proof step i. A straightforward application of DCO fine-tuning in this setting then involves replacing the model confidence p()(x; in Equation (3) on a single answer y. with the model confidence on an entire successful, complete
In our experiments,
k-step proof p(y[0], .... y/k - []]=(0)) = II [(yl]|==(i)) Because of the chain rule, the DCO gradient single proof has the same form as Equation (4) but with F(N, p(y)==) replaced with F(N. I 1-1 p (y]:)(=(())) We naively applied this DCO Linettnnig method with N' 4k and evaluated the resulting model, as well as base model trained with CE loss, on MiniF2F, using pass@44 4k. CE loss achieves a proof success rate of 37.4%, while fine-tuning with DCO achieves a 38.7% success rate. Thus, a naive application of DCO to theorem proving by matching the parameter N' in DCO to the parameter N in pass @N achieves only a modest improvement.
In theorem proving, since every single intermediate proof step tactic yli] must be valid, it is natural to consider a step-wise generalization of DCO. For example, at each step i, the model chooses a tactic vlt with confidence P(ylū)
[ (x(t)) Our proposed step-wise DCO optimizes this single-step confidence according to Equation 44) as before, except now the step-wise confidence regularizer is given by F(Nen, P()[i]](=()) [ Here one can think of Nen as a step- wise DCO hyperparameter that controls the exploration width at every proof step. In essence, the confidence regularizer prevents the confidence of any chosen tactic from becoming much higher than 1/Nan. Since the sum of the confidences over all tacties at each step must be 1, this means that at test time, model exploration at each step corresponds to searching on a search tree in which each proof state æ[i] allows the exploration of
7

---

Table Proof success rates on Mathlib and MiniF2F using pass @ 4k. DCOstP outperforms CE loss (Neit = 1) on both datasets for well-chosen Not. Ensembling models trained with different Nott leverages complementary strengths, with substantial gains over CE at matched test-time compute.
New
I (CE loss)
4
8
16
32
Ensemble of Net'S
Natf = 1
(5x test compute)
Mathlib
55.6%
56.4%
56.1% %
56.5%
55.8%
62.2%
57.0%
MiniF2F
37.4%
39.0%
39.5 %
37.0%
37.4%
43.6%
39.5%
approximately only New tactics. Thus using Net in step-wise DCO during fine-tuning selects the approximate branching factor New of the proof search tree at test time (see Figure F
In particular, valid (partial) proof of length k will have probability of order Now Thus small Neef limits the exploration at test time to a tree with small branching factor, but allows longer proofs with larger numbers of steps k to have higher sampling probability. This limited width search strategy should work well at test time using pass@1 @N as long as two conditions hold: (1) a successful proof of length l is present in the search tree of branching factor Nets and (2) the N in pass @ N at test time is large enough that this proof of probability O(NoT*) is selected with high probability in N passes (which starts to occur soon as N Z N Conversely, larger Nen allows for more exploration at test time using a wider search tree of larger branching factor, but makes finding longer proofs with large & harder, since the approximate success condition for pass @N of N > Ner is harder to satisfy at fixed N and large Net and k. However, this wide exploration strategy can work well for short proofs with small i if successful short proof does not lie in the limited width search tree obtained at small Nen, but does lie in a wider search tree at larger New (Figure 2 (c) right).
The mean and median lengths of the proofs in the training set are 4.2 and 2, respectively. Thus we explore fine-tuning with the step-wise DCO algorithm, DCOSle,, with Neff ranging from 1 (CE loss) to Nor= = 32. We evaluate pass 4k performance on the Mathlib and MiniF2F test sets and find that larger Netf improves accuracy over the CE baseline (Table Interestingly, the optimal Neff increases with more passes test time (Table O similar to the Pareto frontier in Figure 2 (a). Since different choices of Nar correspond to different search strategies at test time larger Net favors short proofs in wide trees, smaller Nen favors long proofs in narrow trees
hypothesized that ensembling these methods could significantly boost performance. This was indeed the case (ensemble entry in Table 2 The ensemble strategy samples 5 4k times, proper baseline is CE loss with the same test-time compute of pass@N N with N = 20k. The ensemble strategy outperforms this stringent baseline with significant excesses of 5 5.2% on Mathlib and 4.1% on MiniF2F.
These results indicate that varying Nar in DCOS at training time allows a diversity of tree search strategies trading depth and breadth that can be exploited by scaling test-time compute via pass @N
5.4 Approximate DCO also improves performance with chain-of-though reasoning
In Section 5.2 the model is trained to directly give an answer y to a problem x. In this case, one can implement DCO at training time by explicitly computing the model confidence ((ylær) and using it in Equation 4 However, in a Cfain-of-Thie of-Thhtht (CoT) training paradigm, the training data consists of triplets (I. e, y) where v and y are the problem and answer before, but now e is a CoT reasoning trace that explains how to derive y from æ. The model trained on the triplet (x,, c, y), and at test time, given a new problem x', it generates a reasoning trace e' followed by an answer y Importantly, the model is evaluated on whether its answer y is correct independent of its reasoning trace d'.
Thus, to estimate the probability P(y(z) that the model assigns to any answer y, one can no longer compute directly as in the direct answer case. Instead, one must marginalize over all possible rea- soning traces e to obtain p(y]z) = Σ. P(y, c/z). Exact marginalization is intractable, so we replace the marginalization with Monte Carlo estimate of p(ylar), which we insert into the overconfidence regularization factor F in Equation 4.. We call this algorithm approximate DCO, denoted by DCO*. We also compute Monte Carlo estimates of the probability the model assigns to its most likely final answer, mode (z),, as proxy for the model's 's max confidence. Note that in the CoT setting the greedy answer would correspond to single most likely reasoning trace, but we want the probability of the most likely final answer marginalizing over all reasoning traces.
8

---

CE Loss
(b)
DCO" (N4=64)
GRPO
GRPO
80
Booch
Epoch 1
Incorrect
3
200
60
Correct
Counts
Epoch 5
40
100
20
50
-2
log pogply
1X
log
[x]
log ply
(x]
log.pty
[x)
Figure 3: (a, b) Overconfidence persists in CoT fine-tuning with CE loss. We fine-tune 3 8B base model on CoT traces in the MATH dataset and plot the distribution of the estimated model confidences p(ymodela (z) over samples in the test set at various points in training. (a) For CE loss, the model becomes more confident in its most likely answers as training progresses- confidence distribution shifts to the right. This effect is milder than in the direct answer setting (Figure IT. (b) DCO* limits this shift of the confidence distribution over training. (c, d) Overconfidence is also present in GRPO fine-tuning and can result in model that is overconfident and wrong. (c) The model trained with GRPO assigns progressively larger confidences to ( moyde |z) over the course of training. (d) At the end of training, only a small portion of the model's highly confident completions are correct, which will harm the model's pass@) N performance when scaling up N.
Table Pass @ coverage on MATH (#: indicates standard error of the mean). We fine-tune Llama- 3-8B-base on CoT traces. Training with DCO* using N' = 64 underperforms CE loss for pass @ 1, but outperforms for higher N, with the largest improvement when N' = N. The GRPO model has the strongest @ass@1 I performance, but its performance degrades at larger N due to overconfidence.
Epoch 1
Epoch
Epoch 5
Pass@N
CE
GRPO
DCO*
CE
GRPO
DCOa
CE
GRPO
DCO*
Pass@1 @1
4.3
3.0
5.0 5.0
9.0 9.0-1.11
12.6
8.3
9.9
14.6
7.8 7,8
Pass @ 16 30.2
27.1
34.2
41.6+0.4
35.8+0033
42.6+0.6
40.9*0.4
29.2
42.8
Pass 64 51.2#1.2 50.6
55.3
61.740.4
63.1 46.3
60.9
35.4 t..3
We conduct experiments the MATH dataset. As baseline, we first fine-tune Allama-3-8B-bass model using CE loss on the golden CoT solutions. We find Table pass@1 coverage robustly improves over the course of CE loss training. Intriguingly, in the CoT setting, overfitting induced by the misalignment of CE loss at training time with the pass@N strategy at test time, appears less severe compare with the direct answer setting (cf. Table Ib. Nevertheless, we still find clear increase in the model's 's max confidence through training (Figure (a)). In contrast, training with DCO mitigates the rise in the model's 's max confidence (Figure (b)). As with the direct answer setup, model trained with DCO" (N' = 64) exhibits lower pass @1 performance compared to the CE baseline but demonstrates superior performance when evaluated with pass @ N for larger N values (Table 3).
To provide a comprehensive comparison, we also evaluate a model trained with reinforcement learning. Specifically, we use GRPO (29 to train the base model with the rl [10] prompt template. Interestingly, we observe that the GRPO-trained model exhibits pronounced overconfidence (Figure (c)). yet only small portion of the model's 's highly confident answers are correct (Figure 3 (d)). This negatively impacts the model's 's pass@y performance when scaling up N (Table 3}. Although GRPO demonstrates superior pass @1 performance, the model trained with DCO* achieves significantly stronger performance when evaluated with pass @ N for larger N values (Table B..
Overall these results once again validate our theory and algorithmic approaches, now in a CoT reasoning setting. In essence, improved performance by scaling test time compute via pass N strategy at large N can be best obtained by aligning the training strategy via choosing DCO* with the same N, and the origin of the performance improvement comes from limiting model overconfidence.
6 Discussion
In summary, all our results suggest the need for a tight co-design of two traditionally separate phases of LLM development: (1) model training or fine-tuning and (2) test-time search/reasoning strategies

---

and budget. If the former is misaligned with the latter, then more training can actually impair performance gains from scaling test-time compute. But if they are properly co-designed, end-to-end training and test time performance can be far more effective. We have shown how to modify standard cross-entropy loss for training be better aligned to a pass@ N strategy for large N test-time. Moreover, we have suggested and empirically confirmed why this co-design of training loss and Pass@N @ strategy is essential, because optimal policies for pass N at large N should unconfidently explore while the optimal policies for pass@N at small N should confidently exploit.
Limitations. Our study primarily investigates SFT with the pass@N test-time strategy for its simplicity. However, this strategy requires a verifier, restricting its applicability to problems where verification is feasible. Extending our findings other test-time strategies, such as best-of-N, an important direction for future research. Additionally, while we have noted similar overconfidence phenomena in GRPO, more comprehensive study of RI. methods beyond the current scope. Lastly, our work only focuses on mathematical reasoning; generalizing our findings to other domains, such as coding tasks where unit tests serve as natural verifiers, left for future work.
Acknowledgments and Disclosure of F Funding
We would like to thank Zhinan Cheng, Clémentine Dominé, Marco Fumero, David Klindt, Daniel Kunin, Ben Sorscher and Atsushi Yamamura for helpful discussions. S.G. thanks the James S. McDonnell Foundation, Simons Foundation, NTT Research, Schmidt Foundation and an NSF CAREER Award for support. S.D. and F.C. were partially supported by the McKnight Foundation, the Simons Foundation and an NIH CRCNS grant R01DC020874, We are also grateful to Adam Ousherovitch and Hongda Yuan for identifying error in Lemma 4.2 which has been corrected in the revised version of this paper.
References
[1] Charlie Snell, Jaehoon Lee, Kelvin Xu, and Aviral Kumar. Scaling Ilm test-time compute optimally can be more effective than scaling model parameters, 2024. URL htttts://aaxiiii org/abs/224 03314
[2] Jason Wei, Xuezhi Wang, Dale Schuurmans, Maarten Bosma, brian ichter, Fei Xia, Ed Chi, Quoc V Le, and Denny Zhou. thhhugh prompting elicits reasoning in large language models. In S. Koyejo, S. Mohamed, A. Agarwal, D. Belgrave, K. Cho, and A. Oh, editors, Advances in Neural Information Processing Systems, volume 35, pages 14824-24833 Curran Associates, Inc., 2022.
[3] Shunyu Yao, Dian Yu, Jeffrey Zhao, Izhak Shafran, Tom Griffiths, Yuan Cao, and Karthik Narasimhan. Tree of thoughts: Deliberate problem solving with large language models. In A. Oh, T. Naumann, A. Globerson, K. Saenko, M. Hardt, and S. Levine, editors, Advances in Neural Information Processing Systems, volume 36, pages 1 118809-1880 11822. Curran Associates, Inc., 2023.
[4] Xuezhi Wang, Jason Wei, Dale Schuurmans, Quoc V Le, Ed H. Chi, Sharan Narang, Aakanksha Chowdhery, and Denny Zhou. Sonfscoon improves chain of thought reasoning in language models. In The Eleventh International Conference on Learning Representations, 2023. URL https://openreviex. httprrrreee..ttteoooommmdddLP.L
[5] Noah Shinn, Federico Cassano, Ashwin Gopinath, Karthik Narasimhan, and Shunyu Yao. Reflex- ion: language agents with verbal reinforcement learning. In A. Oh, T. Naumann, A. Globerson, K. Saenko, M. Hardt, and S. Levine, editors, Advances in Neural Information Processing Systems, volume 36, pages 8634-8652, Curran Associates, Inc., 2023.
[6] William Saunders, Catherine Yeh, Jeff Wu, Steven Bills, Long Ouyang, Jonathan Ward, and Jan Leike. Self-critiquing models for assisting human evaluators. CoRR, abs/2206.05802, 2022.
https://://dot. org/ 10. 48550/ arxiv 2206. 05802
[7] Yixuan Weng, Minjun Zhu, Fei Xia, Bin Li, Shizhu He, Shengping Liu, Bin Sun, Kang Liu, and Jun Zhao. Large language models are better reasoners with self verfficriiic...nn.. In
10

---

The 2023 Conference on Empirical Methods in Natural Language Processing, 2023. URL http:////ppoeeeeeeeeeeeorrrrrmm www.oormimidddddoeemmm
[8] Zirui Zhao, Wee Sun Lee, and David Hsu. Large language models as commonsense knowledge for large-scale task planning. In A. Oh, T. Naumann, A. Globerson, Saenko, M. Hardt, and S. Levine, editors, Advances in Neural Information Processing Systems, volume 36, pages 31967-31987. Curran Associates, Inc., 2023.
[9] OpenAI. Learning to reason with llms, 2024. URL. http::://opeaaa.ccmmmeeeeee
learning- to reason-withe withhlmsstt
[10] Haowei Zhang DeepSeck-AI: Daya Guo, Dejian Yang et al. Deepseek-rl: Incentivizing reasoning capability in Ilms via reinforcement learning, 2025, URL httth:/t/aaaiiii org/
abs/2501. 12948
[11]
Long Ouyang, Jeffrey Wu, Xu Jiang, Diogo Almeida, Carroll Wainwright, Pamela Mishkin, Chong Zhang, Sandhini Agarwal, Katarina Slama, Alex Ray, John Schulman, Jacob Hilton, Fraser Kelton, Luke Miller, Maddie Simens, Amanda Askell, Peter Welinder, Paul F Christiano, Jan Leike, and Ryan Lowe. Training language models to follow instructions with human feedback. In S. Koyejo, S. Mohamed, A. Agarwal, D. Belgrave, K. Cho, and A. Oh, editors, Advances in Neural Information Processing Systemx, volume 35, pages 27730-27744. Curran Associates, Inc., 2022. URL 200p://iiaaaaaassaee444233555522 hrr.a//00eecetpppeeer-ooeoeeeeee Cowfeerrooooeeeceeee... pdf
[12]
Hyung Won Chung, Le Hou, Shayne Longpre, Barret Zoph, Yi Tay, William Fedus, Yunxuan Li, Xuezhi Wang. Mostafa Dehghami, Siddhartha Brahma, Albert Webson, Shixiang Shane Gu, Zhuyun Dai, Mirac Suzgun, Xinyun Chen, Aakanksha Chowdhery, Alex Castro-Ros. Marie Pellat, Kevin Robinson, Dasha Valter, Sharan Narang, Gaurav Mishra, Adams Yu, Vincent Zhao, Yanping Huang. Andrew Dai, Hongkun Yu, Slav Petrov, Ed H. Chi, Jeff Dean, Jacob Devlin, Adam Roberts, Denny Zhou, Quoc V. Le, and Jason Wei. Scaling instruction- finetumed language models. Journal of Machine Learning Research, 25(70): 53, 2024. URL http://jmlr. wrggpappps/v22222
0870 htmml
[13]
Mark Chen, Jerry Tworek, Heewoo Jun, Qiming Yuan, Henrique Ponde de Oliveira Pinto, Jared Kaplan, Harri Edwards, Yuri Burda, Nicholas Joseph, Greg Brockman, Alex Ray, Raul Puri, Gretchen Krueger, Michael Petrov, Heidy Khlaaf, Girish Sastry, Pamela Mishkin, Brooke Chan, Scott Gray, Nick Ryder, Mikhail Pavlov, Alethea Power, Lukasz Kaiser, Mohammad Bavarian, Clemens Winter, Philippe Tillet, Felipe Petroski Such, Dave Cummings, Matthias Plappert, Fotios Chantzis, Elizabeth Barnes, Ariel Herbert- Voss, William Hebgen Guss, Alex Nichol, Alex Paino, Nikolas Tezak, Jie Tang, Igor Babuschkin, Suchir Balaji, Shantanu Jain, William Saunders, Christopher Hesse, Andrew N. Carr, Jan Leike, Josh Achiam, Vedant Misra, Evan Morikawa, Alec Radford, Matthew Knight, Miles Brundage, Mira Murati, Katie Mayer, Peter Welinder, Bob McGrew, Dario Amodei, Sam McCandlish, Ilya Sutskever, and Wojciech Zaremba. Evaluating large language models trained on code, 2021, URL https www./b///////////////o/oeaeeb// abs/ 221//2377773337
//arxiv. hwwp:/orggaahle
[14] Yinlam Chow, Guy Tennenholtz, Izzeddin Gur, Vincent Zhuang, Bo Dai, Aviral Kumar, Rishabh Agarwal, Sridhar Thiagarajan, Craig Boutilier, and Aleksandra Faust. Inference-aware fine- tuning for best-of-n sampling in large language models. In The Thirteenth International Conference on Learning Representations, 2025. URL https:://pperevieeee www.arreeiiii.neeoornn..e nwt/ffrum? id=77gQUdQhE7
[15] Andy L. Jones. Scaling scaling laws with board games, 2021. URL httts:://t//at http://axxvv...o org ftss/204.003111
[16) Bradley Brown, Jordan Juravsky, Ryan Ehrlich, Ronald Clark, Quoe V. Le, Christopher Ré, and Azalia Mirhoseini. Large language monkeys: Scaling inference compute with repeated sampling, 2024. URL. https://:///rt https/://axxwvv.or./// abs/a4s///20 21787
[17] Noam Levi. A simple model of inference scaling laws, 2024. URL, https: abs/ has//241..0 16377
11

---

[18] Rylan Schaeffer, Joshua Kazdan, John Hughes, Jordan Juravsky, Sara Price, Aengus Lynch, Erik Jones, Robert Kirk, Azalia Mirhoseini, and Sanmi Koyejo. How do large language monkeys get their power (laws)?, 2025. URL htt.b:t//p//attttt org warg/aba/25000 2502. 17578
[19] Lin Gui, Cristina Gårbacea, and Victor Veitch. Bonbon alignment for large language models and the sweetness of best-of-n sampling. In A. Globerson, L. Mackey, D. Bel- grave, A. Fan, U. Paquet, J. Tomezak, and C. Zhang. editors, Advances in Neu- ral Information Processing Systems, volume 37, pages 2851-2885. Curran Associates, Inc., 2024. URL httppptps:/pprreeeednnee www.ips...iiipeee wwwppcccpppppp..eeeeeeeeeeee..ee saappee_filleeppppppeee /iappr__.1ss/ppppr//22/////// 066228886cc622227776676666468882 http:///aaaerrroaoooooaaaaeaaaar Conferenceeeee.e. pdf
[20]
Leonard Hussenot, Johan Ferret, Nino Vieillard, Alexandre Rame, Bobak Shahriari, Sarah Perrin, Abram L. Friesen, Geoffrey Cideron, Sertan Girgin, Piotr Stanczyk, Andrea Michi, Danila Sinopalnikov, Sabela Ramos Garea, Amélic Héliou, Aliaksei Severyn, Matthew Hoffman, Nikola Momchev, and Olivier Bachem. BOND: Aligning LLMs with best-of-n distillation. In The Thirteenth International Conference Learning Representations, 2025. URL pactps://openrettt
Pier Giuseppe Sessa, Robert
wwt maet/formm?i id idt.ummmd-d0ttAuuuttttt
[21] Ziniu Li, Congliang Chen, Tian Xu, Zeyu Qin, Jiancong Xiao, Ruoyu Sun, and Zhi-Quan Luo. Entropic distribution matching in supervised fine-tuning of llms: Less overfitting and better diversity. CoRR, abs/2408. 16673, 2024. Uttpttps: org/ 10. 48550/ arxiv 2408 16673
[22] Yang Yue, Zhiqi Chen, Rui Lu, Andrew Zhao, Zhaokai Wang, Yang Yue, Shiji Song, and Gao Huang. Does reinforcement learning really incentivize reasoning capacity in Ilms beyond the base model?, 2025. URL httthhttp:/aaara http:/axxiv..ooaaaba strgg/bbs/5500 13837
[23] Xiang Yue, Xingwei Qu, Ge Zhang, Yao Fu, Wenhao Huang, Huan Sun, Yu Su, and Wenhu Chen. MAmmoTH: Building math generalist models through hybrid instruction tuning. In The Twelfth International Conference on Learning Representations, 2024. URL https:// openrevies
yCC1G77770
[24] Hunter Lightman, Vineet Kosaraju, Yuri Burda, Harrison Edwards, Bowen Baker, Teddy Lee, Jan Leike, John Schulman, Ilya Sutskever, and Karl Cobbe. Let's verify step by step. In The Twelfth International Conference on Learning Representations, 2024. URL https
//openreview wwwpeevvviwi.eefffoooooe
[25] Jonathan Uesato, Nate Kushman, Ramana Kumar, Francis Song, Noah Siegel, Lisa Wang, Antonia Creswell, Geoffrey Irving, and Irina Higgins. Solving math word problems with process- and outcome-based feedback, 2022.
org/abs/ wttpaaaa..aaaaaaaeeeeeaaaaaaaeea 14275
[26] Aitor Lewkowycz, Anders Andreassen, David Dohan, Ethan Dyer, Henryk Michalewski, Vinay Ramasesh, Ambrose Slone, Cem Anil, Imanol Schlag, Theo Gutman-Solo, Yuhuai Wu, Behnam Neyshabur, Guy Gur-Ari, and Vedant Misra. Solving quantitative reasoning problems with language models. In Koyejo, S. Mohamed, A. Agarwal, D. Belgrave, K. Cho, and Oh, editors, Advances in Neural Information Processing Systems, volume 35, pages 3843-3857. Curran Associates, Inc., 2022.
[27] Zhangir Azerbayev, Hailey Schoelkopf, Keiran Paster, Marco Dos Santos, Stephen Marcus McAleer. Albert Q. Jiang, Jia Deng. Stella Biderman, and Sean Welleck. Llemma: An open language model for mathematics. In The Twelfth International Conference on Learning Representations, 2024. URL HRRLhttps: httttp:://openeviic..
[28] An Yang, Beichen Zhang, Binyuan Hui, Bofei Gao, Bowen Yu, Chengpeng Li, Dayiheng Liu, Jianhong Tu, Jingren Zhou, Junyang Lin, Keming Lu, Mingfeng Xue, Runji Lin, Tianyu Liu, Xingzhang Ren, and Zhenru Zhang. Owenn.55-aath technical report: Toward mathematical expert model via self-improvement, 2024. URL https: http://rwivvoorr//bb /aob//244... 2200.111222..
[29] Zhihong Shao, Peiyi Wang, Qihao Zhu, Runxin Xu, Junxiao Song, Xiao Bi, Haowei Zhang, Mingchuan Zhang, Y. K. Li, Y. Wu, and Daya Guo. Deepseekmath: Pushing the limits of mathematical reasoning in open language models, 2024. URL https:: htapp:///arae //arxiv... org/ abs/
2402.03300!
12

---

[30] Huaiyuan Ying, Shuo Zhang, Linyang Li. Zhejian Zhou, Yunfan Shao, Zhaoye Fei. Yichuan Ma, Jiawei Hong. Kuikun Liu, Ziyi Wang, Yudong Wang, Zijian Wu, Shuaibin Li, Fengzhe Zhou, Hongwei Liu, Songyang Zhang, Wenwei Zhang, Hang Yan, Xipeng Qiu, Jiayu Wang, Kai Chen, and Dahua Lin. Intermlm-math: Open math large language models toward verifiable reasoning. CoRR, abs/2402.06332, 2024. URLhttps: //doi..://d.... org/ 2/ngg/100 10. 48550/arXiv.
[31] Eric Zelikman, Yuhuai Wu, Jesse Mu, and Noah Goodman. Star: Bootstrapping reasoning with reasoning. In S. Koyejo, S. Mohamed. A. Agarwal, D. Belgrave, K. Cho, and A. Oh, editors, Advances in Neural Information Processing Systems, volume 35, pages 15476-15488. Curran Associates, Inc., 2022.
[32] Zhenting Qi, Mingyuan MA, Jiahang Xu, Li Lyna Zhang, Fan Yang, and Mao Yang. Mutual reasoning makes smaller LLMs stronger problem-solver. In The Thirteenth International Conference on Learning Representations, 2025. htt.::ttttsss/eetee http://wwwweeeffrruue net/forum? id=6aHUmotXaw
[33] Zichen Liu, Changyu Chen, Wenjun Li, Penghui Qi, Tianyu Pang, Chao Du Wee Sun Lee, and Min Lin. Understanding rl-zero-like training: A critical perspective, 2025. URL https 20783
//arxiv. www./ogg//////ooooooooooo.oeoo
[34] Weihao Zeng. Yuzhen Huang, Qian Liu, Wei Liu, Keqing He, Zejun Ma, and Junxian He. Simplerl-zoo: Investigating and taming zero reinforcement learning for open base models in the http//bb//200... 2503. 18892.
wild, 2025. httttt:ts://axii
[35] Max Marion, Ahmet Üstün, Luiza Pozzobon, Alex Wang, Marzich Fadace, and Sara Hooker. When less is more: Investigating data pruming for pretraining Ilms at scale. CoRR, abs/2309.04564, 2023. URL BRLpstss:://do1 00t.:////rr//00./0/////////////0
48550/ arxiv. 2309. .04664
[36] Ben Sorscher, Robert Geirhos, Shashank Shekhar, Surya Ganguli, and Ari Morcos. Beyond neural scaling laws: beating power law scaling via data pruning. In S. Koyejo, S. Mohamed, A. Agarwal, D. Belgrave, K. Cho, and A. Oh, editors, Advances in Neural Information Process- ing Systems, volume 35, pages 19523-19536 Curran Associates, Inc., 2022.
[37] Chunting Zhou, Pengfei Liu, Puxin Xu, Srinivasan Iyer, Jiao Sun, Yuning Mao, Xuezhe Ma, Avia Efrat, Ping Yu, LILI YU, Susan Zhang, Gargi Ghosh, Mike Lewis, Luke Zettlemoyer, and Omer Levy. Lima: Less is more for alignment. In A. Oh, T. Nau- mann, A. Globerson, K. Saenko, M. Hardt, and S. Levine, editors, Advances in Neu- ral Information Processing Systems, volume 36, pages Inc., 2023. URL https: htttps:// www.:ooooedinnnnnoooooeeeeeeeeee httpppppcccppppppafepeeepp/ wwt.:ffieeeeseeeeeeeeeeeeeeee paaprr/228//flee/ 4c662877822944477cca http://00440000.11110677777444/0 4666
Curran Associates,
[38] Yixin Ye, Zhen Huang, Yang Xiao, Ethan Chern, Shijie Xia, and Pengfei Liu. Limo: Less is more for reasoning, 2025. URL Porraaaeee777aaaaaaaaaoao-orarar pdf
www.orgg///bb.aoeeeeeeee 03387
[39] Abhinay Shrivastava, Abhinav Gupta, and Ross Girshick. Training region-based object detectors with online hard example mining. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition (CVPR), June 2016
[40] Yuxuan Tong, Xiwen Zhang, Rui Wang, Ruidong Wu, and Junxian He. DART math: Difficulty- aware rejection tuning for mathematical problem-solvingg In The Thirty-eighth Annual Con- ference on Neural Information Processing Systems, 2024. URL StLhhttps:/// /ttp::///oeeeeeeeeeeee net/
forum/id=zzU221oojDS
[41] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In I. Guyon, U. Von Luxburg, S. Bengio, H. Wallach, R. Fergus, S. Vishwanathan, and R. Garnett, ed- itors, Advances in Neural Information Processing Systems, volume 30. Curran Associates, Inc., 2017. URL http:l/ttttt::/ppeeeeeiiii neurips cc/ wwtpcccrcppcceeeeeeeeeeee www.les/ppae/2011/////// pdf
[42] Aaron Grattafiori, Abhimanyu Duberyyyyyuuuu....aaaaa Abhinav Jauhri, Abhinav Pandey, et al. The llama 3 herd of models, 2024. http::/:: org/ ab///////...... .21783
13

---

[43] Dan Hendrycks, Collin Burns, Saurav Kadavath, Akul Arora, Steven Basart, Eric Tang, Dawn Song. and Jacob Steinhardt. Measuring mathematical problem solving with the MATH dataset. In Thirty-fifth Conference on Neural Information Processing Systems Datasets and Benchmarks Track (Round 2), 2021. URL
[44] Kaiyu Yang, Aidan Swope, Alex Gu, Rahul Chalamala, Peiyang Song, Shixing Yu. Saad Godil, Ryan J Prenger, and Animashree Anandkumar. Leandojo: Theorem proving with retrieval- augmented language models. In A. Oh, T. Naumann, A. Globerson, K. Saenko, M. Hardt, and S. Levine, editors, Advances in Neural Information Processing Systems, volume 36, pages 21573-21612. Curran Associates, Inc., 2023.
[45] The mathlib Community, The lean mathematical library. In Proceedings of the 9th ACM SIGPLAN International Conference on Certified Programs and Proofs, POPL '20. ACM, Jan- uary 2020. doi: 10. 10.1145/33728853333322 URLhttp::/oo doi org/10 wttg/..../0000001 1145/33728888 3373824
[46] Leonardo de Moura and Sebastian Ullrich. The lean 4 theorem prover and programming language. In André Platzer and Geoff Sutcliffe, editors, Automated Deduction CADE 28, pages 625-635, Cham, 2021. Springer International Publishing. ISBN 978-3-030-79876-3
[47] Jason Ansel, Edward Yang. Horace He, Natalia Gimelshein, Animesh Jain, Michael Voznesen- sky, Bin Bao, Peter Bell. David Berard, Evgeni Burovski, Geeta Chauhan, Anjali Chourdia, Will Constable, Alban Desmaison, Zachary De Vito, Elias Ellison, Will Feng, Jiong Gong. Michael Gschwind, Brian Hirsh, Sherlock Huang, Kshiteej Kalambarkar, Laurent Kirsch, Michael Lazos, Mario Lezcano, Yanbo Liang, Jason Liang, Yinghai Lu, CK Luk, Bert Maher, Yunjie Pan, Christian Puhrsch, Matthias Reso, Mark Saroufim, Marcos Yukio Siraichi, Helen Suk, Michael Suo, Phil Tillet, Eikan Wang. Xiaodong Wang, William Wen, Shunting Zhang, Xu Zhao, Keren Zhou, Richard Zou, Ajit Mathews, Gregory Chanan, Peng Wu, and Soumith Chintala. PyTorch 2: Faster Machine Learning Through Dynamic Python Bytecode Transformation and Graph Compilation. In 29th ACM International Conference on Architectural Support for Program- ming Languages and Operating Systems. Volume 2 (ASPLOS '24). ACM, April 2024. doi: pdf
00.1145/3620665.6443333 URL URL
[48] Sylvain Gugger, Lysandre Debut, Thomas Wolf, Philipp Schmid, Zachary Mueller, Sourab Mangrulkar, Marc Sun, and Benjamin Bossan, Accelerate: Training and inference at scale made simple, efficient and adaptable. battps://githin htwp://www.ccom//gggggggg/////ee
2022.
[49] Woosuk Kwon, Zhuohan Li, Siyuan Zhuang, Ying Sheng, Lianmin Zheng. Cody Hao Yu, Joseph Gonzalez, Hao Zhang, and Ion Stoica. Efficient memory management for large language model serving with pagedattention.
Proceedings of the 29th Symposium on Operating Systems Principles, SOSP '23, page 611-626, New York, NY, USA, 2023. Association for Computing Machinery, ISBN 9798400702297. //doi. 1145/3600006 00155/33000006361166555
doi: 10. 00.1145//60000066661166 URL https
[50] Tsung-Yi Lin, Priya Goyal, Ross Girshick, Kaiming He, and Piotr Dollar. Focal loss for dense object detection. In Proceedings of the IEEE International Conference on Computer Vision (ICCV), Oct 2017.
[51] Thomas Anthony, Zheng Tian, and David Barber. Thinking fast and slow with deep learning and tree search. In I. Guyon, U. Von Luxburg, S. Bengio. H. Wallach, R. Fergus, S. Vishwanathan, and R. Garnett, editors, Advances Neural Information Processing Systems, volume 30. Curran Associates, Inc., 2017. URL http::::/ neurips. wttpk/pp..cccppeeeee wwwpprrflleee//aaapaapaapeeeeele / 0t57752i.68833342227766666666662
pdf
[52] Stanislas Polu, Jesse Michael Han, Kunhao Zheng, Mantas Baksys, Igor Babuschkin, and Ilya Sutskever. Formal mathematics statement curriculum learning. In The Eleventh International Conference on Learning Representations, 2023. URL htttps://operreeeeee www.rrvvvieeeeee/fommrm net/forum? 1Lz-P7G-8dmSh4
14

---

A Proofs
A.1 Proof of approximately well calibration under optimal policy
Lemma A.1. Given any problem, let p. denote the model probability or confidence assigned to answer and assume answers are sorted from highest to lowest confidence that Pi 1i+1, 1. Let ps be the probability that answer is actually correct across all the problems. Then optimal policies maximizing pass @ coverage in Equation 2 (demoted by P is approximately well calibrated, i.
higher confidence implies higher or equal probability of being correct, i.e. pi 2 Pc+1, 1.
Proof. Let R be the total number of answers the model generate. Here R could possibly equal infinity. Let 5 denote the model probability or confidence assigned to answer i under optimal policy maximizing pass @ N coverage in Equation ( Let A; be the event that the 7h ranked answer is not chosen in the N passes and B; be the event that 5th ranked answer is chosen in the N passes. Let 1s, be the indicator function of B,. Then, given an event w, the probability of getting the correct answer is Ex=
Pt PclB, (w). Hence, the expected probability of getting the correct answer is
R
-
R - Ep((1-((A)) 1- 1 _ P.P(A,) = 1- > P1((--))) (5)
R
IE
Σ PiLB.
the current strategy is already optimal, any small change (fi.... P7) would not increase the expected probability of getting the correct answer. Now suppose we don't always have pt P1+1, in other words, < Pj+1 for some 2 1. If fī < f5411 we have
j some
R
1-1
R
Σ P: P:lB.
< 1-
Σ
Σ imj+1
p: p.(1-1))*
- 1,(1-8))-P-HA-C PD+1(1103)) - -D5(1 - -j+1)), (6)
where the last inequality says that increase the expected probability of getting the correct an- swer by swapping the confidence on answer originally labeled as j and j 1, which contradiction. pj+1: then for all 5 > 0 sufficiently small, always have
If 63
E
Σ
PLB,
< 11- Σ
())
Σ p+(1-;)) P5 - Pp+1 (1-11 +-))
which also contradict with the current policy being optimal.
A.2 Proof of Lemma 4.2
Assuming is known, an upper bound on p is given by,
Epi(L - -p))m() ~ 1)
A
max Ti
argmin
(A.)
(8)
k
s.t. Σ Ps =1 - pi 21211 p2.. 2 >.
s=1
Suppose (p:)+= solves Equation 8 and let (i be the optimizer of
&
argmin P:(1 P.) s.t. Σ p. = 12 (A)) ==1
Pr. 20 0.
(9)
==1
By definition, u = pj. Since pn 2 p2... 2 Pk. we also have ग 2 P.... 2 Pt. There are k possible cases.
1) / and 5 for all >
1) / 0 but p = 0. For 2 ls 1. There are k 2 cases.
k) P 0.
We now characterize when each case can occur and derive upper bound in each case.
15

---

A.2.1 Case k
Lemma A.2. Pf O, then p1 = P2 Ph
Proof. In this case, we have
Pkc = (k - 1)Pk a
0, p7 1- (k 1)p1" a
(10)
E Ps
Σs s=1 Ps
where a = N- Suppose, for contradiction, that p1 > 1k-1- We construct a perturbed configuration
{qe fi=1 such that the corresponding optimizer {ū. ) satisfies of > pr. Choose small 5 > 0 and define
P1.
Qk- = Pk-2: Q = Ph-1 + 6, g = Pk
For > 0 sufficiently small, we still have 0% > (by continuity), and moreover oी > pt. contradict- ing the optimality of the original configuration {pi) Hence, p1 p2 1--1-
By Lemma A.2 we have
(k- 1)p1 Pk: = - e, 1)pT (k 2)pp Pi 2 Pk 0.
(11)
This implies the following range for pl in Case k:
1
1-6
S PI
(12)
N-1'
Under condition Equation
12
We have
(k: 1)p
(13)
(k: 1)p1 + (1 + pi - kp:)-a
Moreover, the right hand side of Equation (13 is an increasing function of pt.
A.2.2 Case j
Lemma A.3. If p3 > 0 but P]+1 for some s j S k 1, then P1 Pj-1 and
Pj-1 = = Pk >>0
Proof. We proceed in three steps.
Step 1: we prove Pk 0. Suppose, by contradiction, Pt: 0. Let > 2 be the smallest index such that =
P = 0. If s > j. change Pj to p, 6. Pj+1 to ö for some small 6 >> If s S j. change Pa-1 1 to Pa-1 p, for some small ò > In either case, for sufficiently small, the new system remains in Case j but has strictly larger Pi. a contradiction.
Step 2: we prove Pj+1
Pk. Being in Case j implies
Cj >
0, Vj + +lsrs k.
(14)
where = 1 (r - 1)p
Σt Pt
If fo441 0, change P, to Pj 5, Pj+1 to Pj+1 8 for some small > The new system still in Case j but with strictly larger P7. which leads to contradiction. Hence Cj+1 0.
If c2422 < then Pj+2 < 1 and small 8 > O. Under this change, Cs+1 decreases, cj+2 increases, C.. s 2 + 3 decreases, so the configuration remains Case but with equal ji for > sufficiently enough. But then cj+1 contradicting Cj+1 = 0. Iterating this argument yields cj+1 = Pý-1 Pk-
change Pj+1 to Pj+1 d, PJ+2 P5+2 + 8 for some
= 0, which implies
Step 3: We prove P1 P--1- This follows from the same argument as in the proof of Lemma A.2
D
16

---

As direct consequence, we obtain the following corollary.
Corollary A.4. Let Pj = 9. Pj+1 = =...
= Pk = p. The upper bound u is given by
(j 1-1p N
q N-
(15)
(j - 1)p1 N-1 +9 N-I
(j- 1)p1
N-
+ 9
x=1
Moreover,
1
j j u
(16)
3-1
and p. satisfies
Pi 2 2 0,
(j - 1)pi q (k - j}p 1-6,
(j () - PP )
a N N-1
(17)
j 2 2 q
P1.
Proof.
This is an immediate consequence of Lemma A.3 and its proof (using cj+1 = 0, cj > 0).
Once j is determined, sense.
compute and hence Equation (1 determines j in the following
Lemma A.5. If {go.Po Pol satisfies Equation then jo 2
Proof. Let un be the corresponding maximal value associated with {40, Jo.. Then uo
Combining Equation (17) and Equation (15) we have
S ug
(18)
ju
ju 1 1
j < jo- then up u, which contradiction.
Lemma A.6. For any fixed e, Pr and k, = j(N) is non-decreasing function of N. Since 1/j 2 u 1/(j 1), u u(pr) is a non-decreasing function of N.
Proof. Suppose, for contradiction, that j(N + 1) < j(N). Then
j(N + 1) -
j(N)
- N-1
(19)
J(N +1)-1
J(N)
Let (j(N + (N + 1), P(N + 1)} be the corresponding configuration at N + 1. Using the inequality above, this configuration be used construct larger value of N with index j = -(N + 1) < j(N), contradicting the minimality of j in Lemma A.5
A.2.3 Case 1
This case is impossible.
A.2.4
Summary of the Proof of Lemma 4.2
If
1-c
1-c
SPI
(20)
N-I
k
k-1
then Case k applies. In this regime, the upper bound u is increasing function of P1 and is given by
(k 1)P x
(21)
- (k-1)pp N-1 + (1 + -pp)) xL.
17

---

Otherwise, Case applies, and the upper bound is
(j - 1)p
(22)
(j 1) + N-1
For any fixed Pı and k, let j be the smallest positive integer such that the following equations (viewing {4, p} as variables) admits a solution:
P1 2 q p 0,
(j 1)pi + q (k f)p==p=----- =1-6
(j 1)(p pp
a N
(23)
j 1-2 N-1 q >
P
The index j is non-decreasing with respect to N, and u is further bounded by
5 su< 3-1
(24)
A.3 Proof of Lemma 4.3
Proof. The following proof is motivated by the following intuition: given the values of p1 and pi attains its smallest possible value when the Pj (j 2 3) are as large as possible. Let be the smallest integer such that P1
sp2 2 1. We have,
a+1 > 1=2
p = π1 argmin (
Σ P. -.) 2 m1 1=1
argmin(pi(1- D1) +
p2(1 - -1)0)))))) (25)
(A-)
Second, for fixed pl, Jensen's 's inequality tells us that
1 1 P1 yN
Pi D1)
Σ i=2
p2(1 D.) 2 pi - (1) spp2(1
(26)
Therefore, we have
argmin (p1(1-1)) ( +
Σ (1 (=(11)))))
++1
P1) argmin (1)-))) (1- + sp2(1 Pr
yN
(A:)
2=
(27)
Computing the right hand side of Equation (27) by taking derivatives we have
api 21)N + sp2(1 1-
sp
p2
p
(28)
f1
Pi
-- N-1
P2
Since is the smallest integer such that p1 sp2 1, have
+ > P1 + sp2,
(29)
Combining Equation
29 and Equation
28 we arrive at
11 PI + P2 )
N
P2
> l
(30)
1 - P1 P2 + P
P2
18

---

B Experimental details
Our codebase uses PyTorch
Accelerate 48
and deepspeed /tttts://github.. com/
microsoft/Deepspeed to enable efficient training with memory constraints, and vllm [49 for efficient inference. Code will be made available here https///ggttub..comllllaaaaaeee
MATH. For experiments with MATH dataset
43 (license: MIT), we fine-tune the Llama-3-8B-
Instruct to through post- ttt-trrining process. We follow Lightman et al. (24 and use 12, 000 problems for training and the remaining 500 for testing. In Sections 4 and 5.2 Figure
base 42 the MATH [3] dataset. We start from the base model rather than Llama-3-8B --n-a------nnstructt
avoid potential leakage of the MATH dataset into L
Figure[ic (a) and (b), we fine-tune the model for 4 epochs with learning rate of 2e-5 and batch size 64. We adopt linear learning rate warmup in the first 20 steps. For experiments with DCO, some of the data may have an extreme confidence regularizer value, if the model is already quite confident in answers to certain problems. To maintain an approximately fixed batch size, we set a threshold of 0.3 on the confidence regularizer F. Any training data examples with an F' lower than this threshold will be replaced with new training examples to construct the batch. For the CoT experiments in Section 5.4 same learning rate warmup.
we use learning rate 2e-5 and batch size 128, with the
Theorem proving. We use LeanDojo benchmark 44 (license: CC-BY and adopt the random train and test split as introduced in Yang et al. [44] The random test set includes 2,000 theorems. We fine-tune the model 10wen2.5-Math-i
(28] on the training set for 3 epochs with learning rate le-5 and batch size 64. We adopt linear learning rate warmup in the first 20 steps. To evaluate the model, we use LojnDojo [44] to interact with the proof assistant. We impose maximum wall clock time for each theorem, in addition to limiting the number of passes per problem. For experiments with 4k passes, the time budget is fixed at 5,000 seconds. In order to avoid the model going infinitely deep in the search tree, we limit the number of proof steps to be most 50.
DCO objective. The DCO" introduced in Section 5.4 is an approximation for DCO. To construct batch of size B with DCO", we process batches of samples sequentially: for each batch, we run online inference on each of the samples and discard all samples with probability of success rate larger than thrrehh We choose discard samples which have DCO confidence reguarizer F lower than 0.01, corresponding to thresh = 0.1 for N' = 64. This process continues until we have enough training data for single batch. In Figure 4 we plot the number of discarded samples as a function of training step for our DCO" experiments (same ones as in Table 3).
GRPO. We perform GRPO fine-tuning training set. We follow the modifications to GRPO proposed in Dr. GRPO ((33),, that is, we do not normalize advantages by the standard deviation of the group rewards, do not normalize a roll-out's contribution to the loss by its length, and do not include a KL divergence term with respect to the reference model. Furthermore, we perform only one policy gradient update for cach set of model roll-outs, so there is no need for elipping. This leads to the surrogate objective taking the form,
10 29 on problems of all difficulty levels in the MATH
.V'
lodl
1 (()) = N
To(or, q. 01,4e)
Ac.
stop_grad (To(o:, 9, (.or,.e))
where we adopt notation similar to [33] N' is the number of roll-outs for a single problem statement q. o, is the ith roll-out in the group, on, the ith token o, and Tp(0i, q, Vi,ce) is the probability assigned by the model to token oi,r conditioned on context 0, The stop_grad operator prevents
the flow of gradients when computing VoL. The advantages, in turn, are computed Åi,
N" r(g,o,) 0:) N Ey= >>"ir"rr r(q, o.)), where r(q, o:) is the binary reward assigned by the verifier to the solution o, for the problem g.
For our experiments, we use N' = 8 and batch size 128, which results in a total of 1,024 roll-outs per batch. We also perform batch balaneing such that each batch contains no problems for which either for all i, since these samples do not provide any gradient signal.
r(q, o,) = 0 for all i or r(g. o;) We use constant learning rate of le-6, a 20-step linear warmup: we train for epochs. Note that batch balancing results in fewer gradient steps. Additionally, we use the rl template from [10 for our experiments,
A conversation between User and Assistant. The User asks question, and the Assistant solves it. The Assistant first thinks about the reasoning process in
19

---

the mind and then provides the User with the answer. The reasoning process is enclosed within <think> </think> and the answer is enclosed within <answer> </ answer? tags, respectively, i.e., <think> reasoning process here </think> <answer > answer here </answer>. User: (question)
Assistant: <think
We strictly require solutions follow the format specified in the prompt.
200
Number of Discarded Samples
150
How!
50
300
Step
Figure 4: Number of discarded samples as function of training step for the DCO* experiments. The step structure reflects the model revisiting examples has seen previously training, where each step closely matches the start of a new epoch.
Implementing online inference. Throughout our experiments and analysis, we extensively use the open-source vllm package [49] for efficient inference. Integrating inference into the training loop to enable training under the DCO* objective, as Section 5.4, poses a challenging implementation problem. We solve this problem by placing vllm worker processes, which together perform inference on all GPUs, in separate process group and use Ray / com/ray- pro to isolate them from the training loop. This enables running concurrent training and inference on the same set of GPUs. We believe this inference in the loop setup will be useful to the community, as enables straightforward implementation of online data filtering approaches for LLM training.
Computational cost. All experiments with model size smaller than 10B are performed on machines with 8 NVIDIA H100 GPUs or 8 NVIDIA A100 GPUs. Experiments with model size larger than 10B are performed on machines with NVIDIA H200 GPUs or NVIDIA B200 GPUs. Specifically, for DCO on the MATH dataset, each run required approximately NVIDIA H100 GPU-hour when using the L.lama3-8BB-B Llama-3-70B-B -asee model. For DCO on the LeanDojo benchmark with the Qwen2.5-Mait model, each run required around 12 NVIDIA H100 GPU-hours. For DCO" on the MATH dataset with CoT, each run consumed about 60 NVIDIA H100 GPU-hours for the Llama-3-8B-ban model.
model, and approximately 4 NVIDÍA B2000 GPU-hours when using the
For both DCO and DCOStep, the additional computational cost minimal limited to computing the DCO factor, which is negligible compared to the cost of forward and backward passes. However, DCO", when applied to fine-tuning with CoT traces, introduces significant overhead due the need for online Monte Carlo estimation. In the setting of Llama-3 8B, with 32 layers, hidden dimension 4096, context length 1024, and batch size 128, we estimate petaflops for each forward- backward pass. Generating NMc 64 CoT's for each of 128 examples, in turn, requires approximately 50 petaflops. The actual inference compute cost depends on model performance: stronger models, which will discard more samples, will need process more samples in order to construct full batch for the next training step. Therefore there is tradeoff: constructing a batch of 128 hard samples requires running inference more samples, but also means fewer steps are required to make full pass through the dataset.
limitedd
The memory overhead is approximately 30% of GPU memory, when running one inference engine per pair of H100 80GB GPUs. We note that while the flop overhead is significant, costly online inference is also required in RL frameworks for CoT training. Furthermore, the choice of N can be made smaller, leading to less precise estimates the model's probability of success on given problem, while speeding up training.
20

---

Additional results
C.1 Easy data drives overconfidence
Early in 1st epoch
Late in 1st epoch
2
gest
0
-2
-2
4
1
2
1
2
3
5
Difficulty Level
Difficulty Level
Figure 5: Easy data drives overconfidence and degrades performance when scaling test-time compute. For the experiments in Figure [] compute the directional derivative V yg two points during the first training epoch: 11% (carly) and 86% (late). At each stage, the gradient direction 9 L(,,,)) (fc 34) is evaluated batches of previously unseen training data of each difficulty level defined in the MATH dataset. Early in training, data from all difficulty levels contribute to decreasing the test loss (left). However, later in training, easier examples (difficulty level 2) provide further benefit, while the easiest examples (difficulty level 1) actively degrade test performance (right). The plotted V g S is an average over batches of unseen data, and we use N 256, corresponding pass @ 256.
We have shown in Section that overconfidence limits performance gains from scaling test-time compute. Here we investigate whether all training data contribute equally to the drop in pass @ N test performance. To quantify this effect, we define the test loss as the negative rrg-rrbbability of coverage, C.N
(i.))EDe log (z.g)) where (,.,)) denotes the coverage for single test example (Jr, y). To measure how training batch B influences the final test loss, compute the directional derivative V, g C for the gradient direction 9 = "m V V1), where fcE
the standard CE loss on single example. A negative value, V CN O, indicates that gradient step in the direction of g decreases the test loss.
To analyze the impact of data difficulty, the MATH dataset) and use the directional derivative to examine how a batch of previously unseen data from given difficulty level would affect Liy at different stages of training. Early in training, we observe that CN is negative across all difficulty levels, indicating that data from all difficulty levels contribute positively to reducing the test loss (Figure 5 left). However, near the end of the first epoch, training easier examples (difficulty levels 1 and 2) no longer improves performance, and the easiest examples (level 1) actively degrade pass @ test performance (Figure right). This shows how continued training easy data can harm test performance when scaling test- test-time compute.
group training data by difficulty level (as specified in
21

---

C.2 Optimality of DCO is achieved when the training objective aligns with the test objective.
1.0
21
0.8
N' used in DCO
Coverage
0.6
2
0.4
24
N 256
DCO Frontier
0.0
2
100
10a
102
103
10a
10s
N
given test strategy pass @ N for some N obtained by DCO training for some N' close to N. Same as Figure but with N' = 256 highlighted. We fine-tune Llama-3-8B base models on the MATH dataset produce direct answers. We fine-tune for 4 epochs one model using CE loss and several models under the C objective, for choices of N indicated by color. Note that there is no choice of N' that is optimal across all N (different colors higher at different N). The black curve is Pareto-optimal performance frontier traced by the max of coverage curves for DCO over all N'. Fine-tuning with the DCO loss (N' 256) achieves Pareto-optimality for 180 < N < 385 at test-time, further confirming that optimalyp at a given test strategy pass @ N for some N on the x-axis is obtained by DCO training for some N' close to N.
Figure 6: Pareto-optimallty
C.3 Additional results with Llama-3-70B-base
In this section, we present additional results using the larger Llama-3-70B-baas model, which are consistent with the findings reported for Llama-3-8B-baln
Table 4: Same Table 1 but with -lama-3-70B-haar model. Pass@N N coverage metric the MATH test set for L-lam-3-700--a model fine-tuned with CE loss on direct answers from the MATH training set. Surprisingly, @ass@N test accuracy at large N decreases with number of training epochs.
pass @
pass 16
pass @256
pass @ 4k
Epoch
5.9%
34.3%
68.8%
84.2 %
Epoch 2
6.9%
32.0%
63.0%
81.5%
Epoch
8.2%
24.2%
48.8%
72.2%
Epoch 4
7.9%
22.2%
44.2%
65.9%
22

---

150
Epoch
Epoch
Fpoch
Incorrect
Correct
Epoch
Fpoch
Epoch
256
Counts
100
Ezoch
50
-8
poep()
logp(@mmea/k))
logply
Figure 7: Larger-sized model trained with CE loss also becomes overconfident in its greedy completions; our proposed DCO objective limits this overconfidence on larger-sized models. Same as Figure but with -lama-3-70B-ball dataset to produce direct answers without reasoning trace. Igreedy is the model's greedy completion when sampling autoregressively and choosing the most likely token at each step. Leftmost: The model trained with CE loss assigns progressively larger confidences pi ggreddy (z) to greedy completions over the course of training. Left: At the end of the training, only a small portion of the model's s highly confident completions are correct. This will harm the model's pass performance when scaling up N. Right: Same leftmost but shown for the DCO loss with N 256. Relative to the CE loss, the model trained on DCO shows much milder overconfidence effect. Rightmost: The confidence distribution of the greedy completions after four epochs with DCO for various choices of N. As increases, the model's 's confidence on the greedy completion is more stringently limited, directly consequence of the overconfidence regularizer F.
We fine-tune a Llama-3-70B-Be (0B-base model on the MATH
1.0
1
0.8
N' used in DCO
Coverage
28
0.6
0.4
24
0.2
CELoss
DCO Frontier
0.0
100
10
102
103
104
105
Figure 8: DCO improves CE for pass @ N test coverage Pareto-optimal frontier with Llama-3-70B-base Same as Figure 2 (a) but with Llama-3-70B- model. We fine-tune Llama-3-70B base models the MATH dataset to produce direct answers. We fine-tune for epochs choices of N' indicated by color. We plot pass @ N test coverage as a function N, with each curve (solid red or faint blue-green) corresponding to one fine-tuned model. Note that there is no choice of N' that is optimal across all N (different colors are higher at different N). The black curve is Pareto-optimal performance frontier traced by the max of coverage curves for DCO over all N Pareto- optitootttimlltty at given test strategy @ for some N on the x-axis is obtained by DCO training for some N' close N.
broad range of N and traces
model using CE loss and several models under the Lico objective, for
C.4 Out-of-distributiio evaluation
In this section, we present additional evaluation results AIME24. We confirm that overconfidence persists even on the nuttr--distribution test set, and that the DCO algorithm saves test-time scaling by regularizing the max confidence.
23

---

Table Same as Table but evaluate on AIME24 dataset. Pass@N coverage metric on the AIME24
dataset for a Llama- 8B base model fine-tuned with CE loss on direct answers from the MATH
training set. Surprisingly, Pass@ test accuracy at large N decreases with number of training epochs.
pass @
pass @ 16
pass 2266
pass @ 4k
Epoch 0.2%
3.4%
30.5%
80.5 %
Epoch
0.3%
4.1 %
31.7%
75.6%
Epoch 3
0.2%
2.9%
16.0%
42.0%
Epoch
1.9%
4.9%
15.2%
35.8%
Loss
Loss
DCO (N-256)
DCQ
Epoch 1
Epoch
#1
Epach 3
Tpoen 2
Correct
Epoch
Epach
0 #12255
#1
Counts
Epach 4
Epach 4
#444
-3
-
-4
-
-8
-8
1ogp [X]
loqpt/grwa// [K]
logp5/pweyxx)
ioget/graay [K]
Figure 9: A model trained with CE loss becomes overconfident in its greedy completions oitor--distributinn test set; our proposed DCO objective limits this overconfidence on out-of- distribution test set. Same Figure I but evaluated on AIME24. We fine-tune (lama-3-8B-bar model on the MATH dataset to produce direct answers without a reasoning trace and evaluate the model on AIME24. Fgreddy is the model's greedy completion when sampling autoregressively and choosing the most likely token at each step. Leftmost: The model trained with CE loss assigns progressively larger confidences [[grredyy (z:) its greedy completions over the course of training. Left: At the end of the training, only a small portion of the model's highly confident completions correct. This will harm the model's pass @N performance when scaling up N. Right: Same as leftmost but shown for the DCO loss with N = 256. Relative to the CE loss, the model trained on DCO shows a much milder overconfidence effect. Rightmost: The confidence distribution of the greedy completions after four epochs with DCO for various choices of N. As N increases, the model's s confidence on the greedy completion is more stringently limited, directly as consequence of the overconfidence regularizer F.
1.0
22
CELoSS
0.8
DCO Frontier
N Used in DCO
Coverage
0.6
2
0.4
24
0.2
0.0
2
109
10
102
10
10a
105
N
Figure 10: DCO improves on CE for pass N test coverage over broad range of N and traces Pareto-optimal frontier on out-of-disttibution test set. Same as Figure (a) but evaluated on AIME24. We fine-tune Llama-3-8B base models on the MATH dataset to produce direct answers and evaluate on AIME24. We fine-tune for epochs one model using CE loss and several models under the L objective, for choices of N' indicated by color. We plot @ test coverage as function of N, with each curve (solid red or faint blue-green) corresponding to one fine-tuned model. Note that there is no choice
(different colors are higher different N). The black curve is Pareto-optimal performance frontier traced by the max of coverage curves for DCO over all N'. Pareto- optimalll at a given test strategy pass@ for some on the x-axis is obtained by DCO training for some N' close to N.
that is optimal across

---

C.5
Comparing with Focal Loss
MATH
AIME24
1.0
20
1.0
20
used in Focal Loss
V used Focal Loss
0.8
0.8
15
15
Coverage
Coverage
0.6
0.6
10
10
0.4
0.4
0.2
5
0.2
0.0
1
0.0
10
10a
102
103
104
10
10"
10a
102
10
104 10
N
Figure 11: DCO outperforms Focal Loss (50) for pass N test coverage over a broad range both for in- and out-of-distribution test sets. We fine-tune Llama-3-8B base models on the MATH dataset to produce direct answers, using the Focal Loss [50]: C (I. y) -1 -p x)7 log Pog at. We fine-tune several models for 4 epochs under L for choices of y indicated by color. We plot pass@N test coverage as a function of for MATH (left) and AIME24 (right), with each curve corresponding to one fine-tuned model. The black curve shows the Pareto-optimal performance frontier traced by the max of coverage curves for DCO over all N'. The CE loss curves (red) and the DCO frontier curves (black) are the same as in Figure 2 (a) for MATH and Figure 10 for AIME24.
C.6 Theorem proving
In 6 we show additional results for model performance under the DCOSTEP objective. We find that the optimal Nar grows with increasing passes, agreeing with results in Sections 5.2 and 5.4 We also conduct expert iteration [51] [2) on Mathlib with theorems that do not have proof traces in the training set. We pass @ 1k to prove those theorems. We find that algorithm achieves stronger improvement over the baseline for pass @4k after the iteration. This improvement might result from the fact that the models can prove easy theorems where the model has higher confidence. As result, we believe our method will perform better with expert iteration.
Table 6: Success rate on lean-dojo benchmark random test set trained with DCO
EXPERT ITERATION 0
EXPERT ITERATION 1
DCOSTEP
@ 16 PASS @64 PASS @ 256 PASS @ 1K PASS @ 4K PASS @ 16 PASS @ 256 PASS @ 4K
Nat (CE)
30.0% %
38.75%
46.05
50.75%
55.55%
40.3%
52.65%
58.55%
Nat
30.15%
39.5%
47.2%
52.95%
56.35%
40.8%
53.05 %
59.45%
Net
30.2%
38.9%
47.15%
52.7%
56.1%
40.1%
53.25%
59.5%
Nat
28.65%
46.7%
46.45%
52.9%
56.5%
39.05%
52.8%
60.05
Nan = 32
26.05%
46.7%
45.6%
51.5%
55.8%
37.05%
52.2%
59.15
ENSEMBLE
40.6%
49.15%
54.6%
59.0%
62.15%
49.05%
59.3%
64.8%
25

---

C.7 Plot of the upper bound and the lower bound
1.0
Upper Bound
Lower Bound
0.8
2 0.6
0.4
100
101
102
103
N
Figure 12: Plot of the upper bound (Appendix A.2.4) and the lower bound (Equation (28). We plot the upper bound (blue) and lower bound (orange) for p1 = 1.2 P2 1. c 1 and 2 Both
upper bound and the lower decrease monotonically in and they both tends to as + 1+
C.8 Empirical evaluation the approximately well calibrated assumption
0.08
0.20
Accuracy (pi)
0.06
0.15
0.04
0.10
0.02
0.05
10 15
10
15
Model confidence rank
Model confidence rank
Figure 13: We empirically verify the assumption that models are approximately well-calibrated. Left: We fine-tune Llama-3-8B-bas model with CE loss the MATH dataset without CoT and perform beam search with width of 256 to obtain the top most probable completions. We then measure the accuracy of these completions at each confidence rank. Right: We fine-tune Llama-3-8B--art model with CE loss on the MATH dataset with CoT and perform estimate the top 16 most frequent answers after tracing out the reasoning traces. We then measure the accuracy of these answers at each frequency rank. The results demonstrate that test accuracy decreases approximately monotonically with model confidence rank, supporting the assumption.
26

---

C.9 The data dependency of confidence and factor F
Late in 1st epoch
Late 1st epoch
0.10
0.8
0.08
F(N, p(y[x))
0.6
plyxx)
0.06
0.4
0.04
0.2
0.00
0.0
Difficulty Level
Difficulty Level
Figure 14: Model is more confident on easy problems; DCO improves test-time scaling by regularizing the easy examples. We fine-tune Llama-3-8B-bass model on the MATH dataset with CE loss and plot the model confidence ply(r) and the factor P(N, [(y]()) 86% of the first training epoch. Both model confidence and factor evaluated the unseen data grouped by difficulty level from the MATH dataset. Model confidence decreases with increasing difficulty, whereas the regularization factor increases with problem difficulty. As result, DCO effectively regularizes contribution from easy examples. This regularization mitigates the potential detrimental effects of overconfidence from easy examples discussed in Appendix C.T to pass @ 256 for the plots.
We use N 256 corresponding
C.10
Quantifying the output diversity of models trained with DCO
AIME24
N'
MATH AIME24
150
N' 256
Counts
100
10
N' AK
50
0.0
2.5
5.0
7.5
0.0
2.5
5.0
7.5
2 24 2
2
5
N'
Figure 15: DCO enhances output diversity in fine-tuned models. Increasing N' leads to greater diversity in model completions. We fine-tune Llama-3-8B-bar model the MATH dataset generate direct answers without explicit reasoning steps. Shannon entropy of model completions, conditioned on the input prefix, estimated using 4096 samples per test example. Left and Middle: Histograms depicting the estimated entropy distributions for the test of MATH (left) and AIME24 (middle) respectively. Higher values of N' shift the entropy distribution the right, reflecting increased diversity of model outputs. Right: Mean entropy values plotted against N' with standard errors of the mean. The estimated entropy increases with N'.
27

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
The abstract and introduction clearly and accurately reflect the main claims made in the paper, including (I) identifying a misalignment between the CE loss used in SFT and the Pass@N coverage metric at test time; (2) developing and verifying theoretical framework that suggests the misalignment arises from overconfidence induced by CE loss; (3) proposing loss function that directly optimizes for pass@1 N and demonstrating its superior performance on MATH and MiniF2F; (4) extending the proposed algorithm to more complex scenarios with CoTs and theorem proving. We explicitly reference the relevant sections supporting each contribution in the introduction.
Justification:
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
is fine to include aspirational goals as motivation as long is clear that these goals are not attained by the paper.
2. Limitations
Question:
Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification: We explicitly discuss the limitations of our findings, including the reliance on a verifier for the pass @ strategy, the limited scope mathematical reasoning tasks, and the unexplored generalization to alternative test-time strategies and reinforcement learning methods, in Section
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section in their paper.
The paper should point any strong assumptions and how robust the results are to violations of these assumptions (e.g.. independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested few datasets with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they seale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and faimess.
28

---

While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory assumptions and proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?
Answer: [Yes]
Justification: We explicitly state all assumptions required for each lemma in our paper. Complete proofs for these theoretical results are provided in Appendix
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless whether the code and data are provided or not)?
Question:
Answer: [Yes]
Justification: We provide sufficient details for reproducing each experiment in Appendix B
including hyperparameter settings and required compute. Additionally, we include code to reproduce all reported results as supplementary material, and we intend to publicly open-source our codebase upon publication.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is specific model and empirical evaluation, it may be necessary to either make possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., s., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution is primarily a new algorithm, the paper should make clear how reproduce that algorithm.
29

---

(b) If the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution is new model (e.g.. a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors welcome describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, described in supplemental material?
Answer: [Yes]
Justification: We plan to release our code publicly upon publication. Comprchensive docu- mentation describing the usage, installation instructions and dependencies of our codebase is included and will be made available alongside the code. Some of the experimental details are included in Appendix B
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://niis.cc)
for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https
y //nips.ccom ctips.cc/public/uuii hwwubligguuide//////eeeee htt////ooSSuiisseeeeeee pttsSSummmssiooooooooo ssionnolicy for more details.
The authors should provide instructions on data access and preparation, including how access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (c.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
We clearly describe our experimental settings and provide comprehensive details, including hyperparameters and dataset splits, for each experiment in Appendix B Additionally, we include code as supplementary material to enable full reproducibility of all reported results, and we plan to publicly open-source our codebase upon publication.
Justification:
Guidelines:
The answer NA means that the paper does include experiments.
The experimental setting should be presented the core the paper to level of detail that is necessary to appreciate the results and make sense of them.
30

---

The full details can be provided either with the code, in appendix, or as supplemental material.
7. Experiment statistical significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [No]
Justification: We generally do not report error bars due the substantial computational cost associated with performing multiple repeated runs. However, for experiments exhibiting smaller performance margins (e.g.. Table 3 we explicitly include error bars. We believe our main findings are robust, reproducible with the provided experimental details and code, and not attributable to superficial noise.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.
The factors of variability that the error bars capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
should be clear whether the error bar is the standard deviation or the standard error of the mean.
It is OK to report I-sigma error bars, but should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful figures symmetric error bars that would yield results that are out of range (c.g. negative error rates).
show in tables
error bars are reported in tables or plots, The authors should explain the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification: We provide detailed information regarding computational resources and ciated costs in Appendix
Guidelines:
The answer NA means that the paper does include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't 't make into the paper).
9. Code of ethics
Question: Does the research conducted the paper conform, in every respect, with the NeurlPS Code of Ethics htttp:::/neuripssccc/// httppp////aarrp..cpppaaaaaaaaaaa
Answer: [Yes]
31

---

Justification:
We make sure to conform with the NeurIPS Code of Ethics.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special cireumstances that require a deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., if there special consid- eration due to laws or regulations in their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification: Our research foundational nature, focusing on understanding factors that limit performance gains when scaling test-time compute, and designing algorithms better aligned with test-time strategies to enhance mathematical reasoning capabilitiess in LLMs. While work could have broad impacts in education, research, and automated problem-solving we do not see a direct or immediate path toward specific negative societal consequences beyond those already associated generally with large language model research.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g.. disinformation, generating fake profiles, surveillance), faireess considerations (e.g., dep.lymenn of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer[[NA]
Our paper does not release data or models and thus the paper poses no such risks.
Justification:
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
32

---

Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification: We have explicitly cited all papers whose open-source code or datasets were utilized in our experiments.
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset is used and, if possible, include URI..
The name of the license (e.g.. CC-BY 4.0) should be included for each asset. For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperssi whrode.ehcc...... com datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that re-packaged, both the original license and the license of the derived asset (if has changed) should be provided.
If this information is not available online, the authors are encouraged to reach out to the asset's creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [Yes]
Justification: We plan to release our code publicly upon publication. Comprehensive documentation describing the usage, installation instructions and dependencies of our codebase is included and will be made available alongside the code. We have carefully cited all papers whose open-source code or datasets were utilized in our experiments. Some of the details are included in Appendix B
Guidelines:
The answer NA means that the paper does not release new assets. Researchers should communicate the details the datasct/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification: The paper does not involve crowdsoureing nor research with human subjects.
33

---

Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included the main paper.
According to the NeurIPS Code Ethics, workers involved in data collection, curation, other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals or equivalent for research with human subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: The paper does not require crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent)
may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
16. Declaration of LLM usage
Question: Does the paper describe the usage of LLMs if it is an important, original, or non-standard component of the core methods this research? Note that if the L.M used only for writing, editing, or formatting purposes and does not impact the core methodology. scientific rigorousness, or originality the research, declaration not required.
Answer: [NA]
Justification: The core method development in this research does not involve LLMs any important, original, or non-standard components.
Guidelines:
The answer NA
that the core method development in this research does involve LLMs as any important, original, or non-standard components.
Please refer to our LLM policy /ettps://neurr cc/ //oofffrencee/22255555 for what should or should be described.
34


Reviews:
[{"review_id": "0TWgm9DbWf", "reviewer": "Reviewer_4Nd6", "summary": "The 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.", "questions": "- 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?", "limitations": "yes", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strengths:\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?", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely thank the reviewer for their thorough and insightful comments. We are encouraged that the reviewer found the theoretical analysis solid, our writing clear, and our topic timely and relevant. Below, we directly address the reviewer's concerns and questions, with additional empirical results and clarifications to strengthen our contributions.\n## Comparison to Focal Loss\nThe reviewer correctly notes that various works have also proposed loss functions aiming to reduce overconfidence in the cross entropy (CE) loss. However, we would like to emphasize that the overconfidence we identified arises specifically from a train-test mismatch in sampling strategies. While there are other methods/loss functions to alleviate overconfidence, our DCO loss is directly designed to address this misalignment.\nTo concretely demonstrate this difference, we followed the reviewer’s suggestion and conducted a thorough comparison with Focal Loss (FL) [Lin et al., 2017], $\\ell^\\gamma_\\text{FL}(x, y) = -(1-\\hat{p}(y \\mid x))^\\gamma \\log \\hat{p}(y \\mid x)$, optimizing over its hyperparameter $\\gamma$ for each value of $N$. Our results are as follows:\n\n\n**Table 1**. We fine-tune a Llama-3-8B-base model on the MATH dataset to generate direct answers with CE loss, FL with different $\\gamma$ and the DCO losses with different $N$, and test the model on the test set of MATH. We have reported pass@N performance for the CE baseline and DCO frontier (as plotted in Figure 2a). Here, we include the FL frontier, which selects the best performing $\\gamma$ for each $N$. The FL frontier consistently outperforms the CE baseline but underperforms the DCO frontier across all values of $N$. This further highlights the advantage of DCO in improving coverage under heavy sampling.\n| | N=1 | N=8 | N=64 | N=512 | N=1024 | N=2048 | N=4096 |\n|----------|----------|----------|----------|----------|----------|----------|----------|\n| CE   | **0.07**     | 0.18     | 0.33     | 0.49     |  0.54   |  0.59     |   0.63    |\n| Focal Loss frontier    | **0.07** | 0.22  | 0.44 | 0.63  | 0.68  | 0.72  | 0.75  |\n| DCO frontier    | **0.07**     | **0.24**     | **0.54**     | **0.75**     |**0.80** | **0.84**    |  **0.87**     |\n\n**Table 2**. Same as Table 1, but test on the out-of-distribution AIME24.\n\n|                      | N=1  | N=8  | N=64 | N=512 | N=1024 | N=2048 | N=4096 |\n|----------------------|------|------|------|-------|--------|--------|--------|\n| CE                   | **0.02** | 0.04 | 0.09 | 0.19  | 0.23   | 0.29   | 0.36   |\n| Focal Loss frontier  | **0.02** | 0.04 | **0.12** | 0.24  | 0.33   | 0.45   | 0.55   |\n| DCO frontier         | **0.02** | **0.04** | 0.11 | **0.48** | **0.64** | **0.79** | **0.89** |\n\nThese results clearly show that while heuristically designed objectives like FL can indeed outperform vanilla CE loss, DCO loss consistently achieves significantly superior performance, particularly at larger $N$ values critical for practical applications. We will include these comparisons explicitly in the revised paper. Specifically, we will add a corresponding plot illustrating that the FL frontier consistently lies above the CE baseline but below the DCO frontier across nearly all values of $N$ on both in-distribution MATH test set and out-of-distribution AIME24.\n\nWe fully agree that providing clear guidelines for choosing $N’$ greatly enhances the practicality of our approach, given that training a separate model for each choice of $N'$ is computationally expensive. The main argument of this paper is that one should align the training objective with the test-time strategy. That said, we study the calibration of the DCO objective and find that DCO can be robust to some extent to mismatches between $N'$ (the DCO objective parameter) and $N$ (the $N$ used to evaluate pass@N at test time). We define an $N'$ as optimal for a given $N$ if its pass@N performance is within 1% of the DCO frontier at that $N$ (see Figure 2a). Based on experiments fine-tuning Llama-3-8B on direct-answer generation for the MATH dataset, we observe:\n- $N’=1$ is optimal for $1\\leq N\\leq 3$.\n- $N’=16$ is optimal between $2\\leq N\\leq 20$.\n- $N’=256$ is optimal between $180\\leq N\\leq 385$.\n- $N’=4096$ is optimal for $N\\geq 1151$.\n\nWe also remind the reviewer that we proposed a rule of thumb for selecting $N_{\\text{eff}}$ in the case of theorem proving (see the paragraph starting at line 269). If one has an estimate of task difficulty (e.g., proof length $k$), we suggest using $N_{\\text{eff}} = N^{-1/k}$, where $N$ is the inference-time sampling budget.\nAdditionally, if an inference-time budget (in FLOPs or wall-clock time) is available, a practical strategy is to estimate the distribution of cost per sample and then choose the largest $N'$ such that the expected total cost remains within budget.\n\n\n## Clarification of the role of $\\epsilon$ in the lemmas\nWe thank the reviewer for raising this important point. We clarify that $\\epsilon$ captures uncertainty regarding exact probability values for lower-ranked terms. The aim of Lemma 4.2 is to establish an upper bound for the model confidence $\\hat{p_1}$* associated with the highest-ranked answer under the optimal policy that maximizes coverage. The precise value of $\\hat{p_1}$* inherently depends on all probabilities $p_{i}$. The introduction of the $\\epsilon$ term explicitly acknowledges uncertainty regarding the exact values of probabilities $p_{i}$ for indices $i>k$, capturing our incomplete knowledge of these lower-ranked terms. Thus, $\\epsilon$ is a theoretically motivated component ensuring robustness rather than a practically burdensome hyperparameter.\n\n## Comparison to other inference-aware tuning methods\nWe appreciate the suggestion to expand our discussion regarding BonBon and BOND. Both methods elegantly address the problem of distilling the Best-of-$N$ (BoN) policy, thereby avoiding the cost of $N$ inference passes at test-time, closely related to RLHF objectives. However, their work focuses on improving the pass@1 performance from the BoN policy, while we focus on directly improving the pass@N performance. We will incorporate these distinctions clearly in our revised related-work section, highlighting the complementary contexts in which each method applies.\n## Effect of sequence length on the probability\nWe recognize the reviewer’s concern regarding longer sequences naturally receiving lower probabilities. Importantly, our method is designed such that sequence length does not inherently penalize coverage:\nIn direct-answer tasks, DCO optimizes directly for the correct final answer, independent of sequence length.\nFor formal math, we apply DCO at the step-level granularity, avoiding any length-related bias.\nIn Chain-of-Thought scenarios, our method marginalizes over reasoning traces, ensuring no direct penalization of longer sequences.\nThus, sequence length inherently does not disadvantage our proposed approach.\n\n\nWe greatly appreciate Reviewer 4Nd6's valuable feedback, which significantly strengthens our paper. We hope our detailed clarifications and additional experiments address all concerns and increase the reviewer’s confidence in our work. We would be happy to clarify any points that are still unclear or any additional questions.\n\nReference:\n\n[Lin et al., 2017] Lin, Tsung-Yi, et al. \"Focal loss for dense object detection.\" Proceedings of the IEEE international conference on computer vision. 2017."}, {"review_id": "frBgrTSYAJ", "reviewer": "Reviewer_asva", "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.", "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", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 4, "clarity": 3, "significance": 4, "originality": 3, "strengths_and_weaknesses": "Strengths: \n - 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.\n - The performance of DCO and its variants are good on MATH, AIME, MiniF2F, and Mathlib, cross multiple domains.\n - The method applies to both step-wise proof search and chain-of-thought, showing robustness.\n\nWeaknesses:\n - Training with DCO may need more sampling and GPU time, which is acceptable.", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely thank Reviewer asva for their thorough and insightful feedback. We are pleased that the reviewer found our identification of the gap between training and inference insightful, and recognized the robustness and strong empirical performance of our proposed DCO. Below, we address the reviewer's main concerns and questions directly, providing additional clarifications and evidence to strengthen our contribution.\n\n## Robustness of DCO \n\nWe appreciate the reviewer’s observation that training a separate model for each choice of $N'$ is computationally expensive. Indeed, the central theme for this paper is that we cannot simply train a single model to be optimal for every $N$ without conditioning on $N$ at test time. \n\nThe reviewer raises a question about the robustness of DCO to mismatches between $N'$ (the DCO objective parameter) and $N$ (the $N$ used to evaluate pass@N at test time). We would like to remind the reviewer that there is an example plot in Figure 6, where we highlight and show how robust a model trained with a specific $N=256$ is with respect to the Pareto frontier. To further investigate the robustness, we define an $N'$ as optimal for a given $N$ if its pass@N performance is within 1% of the DCO frontier at that $N$ (see Figure 2a). Based on experiments fine-tuning Llama-3-8B on direct-answer generation for the MATH dataset, we observe:\n\n- $N’=1$ is optimal for $1\\leq N\\leq 3$.\n- $N’=16$ is optimal between $2\\leq N\\leq 20$.\n- $N’=256$ is optimal between $180\\leq N\\leq 385$.\n- $N’=4096$ is optimal for $N\\geq 1151$.\n\nThese results strongly indicate that while matching $N'=N$ is a standard and intuitive approach, a fixed choice of $N'$ can result in near-frontier performance for a range of test-time $N$. Selecting a few strategic $N'$ values can yield near-optimal performance across broad ranges, reducing computational overhead.\n\n\nMoreover, as demonstrated in our theorem proving results, an ensemble of models trained with different $N'$ values can actually surpass the performance of a single model using larger $N$. Thus, training multiple models and forming an ensemble can be an advantageous strategy rather than purely a limitation.\n\n\n## Strategies for Supporting Multiple Test-time Budgets\n\nThe reviewer raises an important question regarding potential strategies to support multiple test-time budgets without complete retraining. Indeed, a promising approach is to explicitly condition the model on $N$ during training and testing (i. e.  $p(y|x,N)$), enabling a single model to flexibly adapt to various inference budgets. While exploring this direction in depth remains beyond the scope of our current work, we acknowledge it as a valuable direction for future research.\n\n\n\n## Computational overhead in the CoT setting\nWe agree with the reviewer regarding the significant computational overhead involved in Monte Carlo estimation in CoT settings. While acknowledging this limitation, we briefly note here that, although this overhead is significant relative to standard SFT, costly online inference is also required in RL frameworks for CoT training, as discussed in Appendix B. Importantly, practitioners can significantly reduce training costs by using fewer Monte Carlo samples ($N_{MC}$), trading precision in estimating success probability for increased computational efficiency.\n\n\nWe thank the reviewer again for their valuable suggestions and insightful questions, which help improve our paper substantially. We hope our responses clearly address their concerns and further strengthen their confidence in our findings. Please don’t hesitate to reach out if anything remains unclear or if additional questions arise."}, {"review_id": "3RSA0u5Uay", "reviewer": "Reviewer_52g7", "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.", "questions": "See the Weakness.", "limitations": "yes", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strength\n\n- 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).\n\n- 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.\n\n- Theoretical Depth: Lemmas 4.2–4.3 formally characterize the confidence-exploration tradeoff, and Appendix A provides complete proofs.\n\nWeakness\n\n- Approximate DCO for CoT requires online Monte Carlo estimation, increasing compute by 10× vs. standard fine-tuning.\n\n- The paper links overconfidence to reduced pass@N performance but does not quantify how DCO affects output diversity.", "was_revised": true, "final_justification": "The detailed response solves my concern.", "rebuttal": "We thank the reviewer for their thorough review and insightful feedback and suggestions. We are pleased that the reviewer recognizes the novelty and practical relevance of our findings and appreciates the depth of our theoretical analysis. Below, we directly address the reviewer's concerns and provide additional evidence and clarifications.\n\n## The effect of DCO on output diversity\nWe thank the reviewer for suggesting this point. We agree with the reviewer that the effect of the DCO objective on output diversity is important to understand. To quantify this, we estimated the mean Shannon entropy of the model outputs using Monte Carlo sampling with 4096 samples per problem. Specifically, we fine-tuned a Llama-3-8B-base model on the MATH dataset to generate direct answers using DCO losses with different values of N’ and compute the mean Shannon entropy of the model distribution on the test set of MATH and AIME24. We will explicitly incorporate these diversity results into our revised paper. We summarize the results in the table below:\n\n**Table 1**.  Mean and standard deviation of Shannon entropy for models trained with DCO of varying N’ values.\n| Dataset | $N’=1$ (CE loss)  | $N’=8$  | $N’=64$  | $N’=512$  | $N’=4096$ |\n|----------|----------|----------|----------|----------|----------|\n| MATH   | $2.27\\pm1.43$    | $3.23\\pm1.16$    | $4.22\\pm0.91$     | $5.52\\pm0.71$     |  $6.71\\pm0.58$  |\n| AIME24    | $2.74\\pm1.27$    | $4.11\\pm1.17$    | $4.39\\pm0.91$     | $5.89\\pm0.41$     |  $7.10\\pm0.22$  |\n\nThe table clearly demonstrates that the entropy consistently increases with larger $N’$. This indicates that models trained with higher $N’$ indeed produce more diverse outputs, thus directly confirming that DCO naturally encourages diversity. Intuitively, superior pass@N performance at larger $N$ requires distributing probability mass across multiple candidate solutions, reducing overconfidence and enhancing coverage. \n\n## Computational overhead in the CoT setting\n\nWe agree with the reviewer regarding the significant computational overhead involved in Monte Carlo estimation in CoT settings. While acknowledging this limitation, we briefly note here that, although this overhead is significant relative to standard SFT, costly online inference is also required in RL frameworks for CoT training, as discussed in Appendix B. Importantly, practitioners can significantly reduce training costs by using fewer Monte Carlo samples, trading precision in estimating success probability for increased computational efficiency.\n\n\nWe thank the review again for their thoughtful questions and constructive feedback. We hope our detailed responses address all concerns, and we remain available to provide further clarifications if needed."}, {"review_id": "M4bxM8BTWh", "reviewer": "Reviewer_j14W", "summary": "This 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.", "questions": "1. 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?", "limitations": "yes", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 4, "clarity": 4, "significance": 2, "originality": 3, "strengths_and_weaknesses": "**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.", "was_revised": true, "final_justification": "Overall this is a good paper addressing the misalignment between cross-entropy based training and pass@N metrics during inference. The rebuttal addressed my concerns regarding the choice of N' and how the method can be used with soft verification.", "rebuttal": "We sincerely thank Reviewer j14W for the thorough and insightful review. We appreciate that the reviewer found the paper to be exceptionally well-written, the DCO formulation elegant, and the empirical validation strong and thorough. Below, we directly address the reviewer’s important questions and provide additional clarification and evidence to strengthen our contributions.\n\n## Overconfidence Severity at Scale (Larger Models)\n\nWe thank the reviewer for raising this important problem. Indeed, one of our motivations for investigating the Llama-3-70B model was precisely to explore whether the severity of overconfidence diminishes or intensifies as models grow larger and more capable. Our empirical findings indicate that overconfidence may become more severe at larger scales: For Llama-3-8B model, the maximum improvement of DCO frontier over CE loss happens at $N=541$ with an improvement of $26.3$%, while for Llama-3-70B model, the maximum improvement over CE loss occurs at $N=258$ with a larger improvement of $29.3$%. Interestingly, the performance gap between DCO and CE not only persists but grows slightly larger at the 70B scale. We suspect that the larger model, because it has much more parameters, may be more prone to overfitting and thus exhibit greater overconfidence under CE training. We acknowledge that this observation is based specifically on the Llama-3 series, and broader validation across more model families would further strengthen generalizability. To clarify, we are not suggesting that larger models inherently exhibit overconfidence without fine-tuning; rather, we aim to highlight a critical failure mode resulting from the misalignment between CE training and pass@N evaluation. We believe this phenomenon of training-test misalignment and the resulting overconfidence is likely a general issue.\n\nWe also recognize that the pass@1 performance of the model on the benchmark is not near state-of-the-art, because we use the base model instead of an instruction-tuned model. We intentionally selected base models for two reasons: (1) to avoid potential data contamination or interference from post-training and (2) because lower pass@1 accuracy in the base models provides more meaningful room for improvement to demonstrate the effectiveness of scaling test-time compute.\n\nWe will integrate these key results and comparisons in the revised paper to better address scalability concerns.\n\n\n## Robust choices of N’ under unknown test-time compute vs. ensembling\n\nWe appreciate the reviewer’s observation that the test-time $N$ might not be known at training time. To address the possibility of having “sweet-spot” choices of $N’$ for training, we evaluate the robustness of DCO to mismatches between $N′$ (the DCO objective parameter) and $N$ (the $N$ used to evaluate pass@N at test time). We define a choice of $N’$ to be optimal for a given test-time $N$, if its pass@N performance is within 1% of the DCO frontier at that $N$ (see Figure 2a for the frontier and individual DCO curves). Based on experiments fine-tuning Llama-3-8B on direct-answer generation for the MATH dataset, we observe:\n\n- $N’=1$ is optimal for $1\\leq N\\leq 3$.\n- $N’=16$ is optimal between $2\\leq N\\leq 20$.\n- $N’=256$ is optimal between $180\\leq N\\leq 385$.\n- $N’=4096$ is optimal for $N\\geq 1151$.\n\nThese results indicate that while matching $N’=N$ is a standard and intuitive approach, a fixed choice of $N’$ can result in near-optimal performance for a range of test-time $N$. Selecting a few strategic $N’$ values can thus yield near-optimal performance across broad ranges.\n\nWe also believe that ensemble may be beneficial for certain use cases such as theorem proving. As highlighted by the reviewer, an ensemble of models trained with different $N_{eff}$ values can actually surpass the performance of a single model when using larger $N$ in theorem proving. Thus, training multiple specialized models can be an advantageous strategy rather than purely a limitation. Of course, both the robustness of DCO to the test-time $N$ and the benefits due to ensembling models at various choices of $N’$ are likely task-dependent.\n\nRegarding the question of whether a single model could effectively handle multiple inference budgets at test time, we would like to suggest a potential strategy: explicitly conditioning the model on the target sampling budget $N$ during training and testing (i. e. learning the distribution of $p(y|x,N)$), thus enabling a single model to flexibly adapt to various inference budgets. While exploring this direction in depth remains beyond the scope of our current work, we acknowledge it as a valuable direction for future research.\n\n\n## Settings with “soft” verification\n\nWe agree with the reviewer regarding the importance of exploring extensions of DCO to domains where verification is not strictly binary but rather “soft”. We acknowledge that restricting the scope to verifiable problems is one of the limitations as we stated at the end of our paper. Extending DCO to handle continuous verification signals represents a promising direction for future research. Given the substantial density and content of our current manuscript, we leave such extensions to future work and hope the reviewer could understand.\n\nWe want to thank the reviewer again for their thoughtful comments and constructive suggestions. We remain eager to address any additional questions or concerns the reviewer might have."}]

Author rebuttal:

Meta-review: In this paper, the authors note that large language models are often evaluated on pass@k evaluations (i.e. whether one sample is correct from k samples conditioned on a prompt) but they are trained on pass@1.  They observe that models trained on standard objectives actually degrade in performance on pass@k.  As such, they propose a new training loss they call Direct Coverage Optimization (DCO).  They show empirically that DCO improves on pass@k evaluations compared to cross-entropy and the recent GRPO method.  They also show theoretically why pass@k performance degrades with existing strategies.

Overall the reviews are borderline leaning towards accept (4, 4, 4, 5).  The reviewers seemed to find the paper "exceptionally" well written and clear.  They found the work well-motivated, topically "very relevant" and theoretically "solid".  They noted that the method presented was simple and seemed empirically convincing.  The reviewers who gave scores of 4, noted very little in terms of weaknesses and overal seemed positive.  Unfortunately, that didn't give the authors a significant amount to address and, in-addition, those reviewers didn't seem to engage meaningfully in discussion despite what seems like an in-depth response.  Two of these reviewers noted that the proposed algorithm required additional sampling and thus was computationally less efficient (however, GRPO does as well).  The reviewer that recommended accept gave a thorough review and engaged with the authors.  Given the lack of any significant concerns raised by some of the reviewers, I would deduce that their scores are perhaps lower than they should be.  Therefore, the recommendation is to accept (i.e. up-weighing the 5 score).  

The Accept review noted that there remains a question whether the results generalize to the current largest SOTA models.  Naturally, that is difficult to empirically verify, but seems like an important potential limitation.
Decision: Accept (poster)