Rethinking Fine-Tuning when Scaling Test-Time Compute: Limiting Confidence Improves Mathematical Reasoning

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.

Overconfidence Severity at Scale (Larger Models)

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

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

We will integrate these key results and comparisons in the revised paper to better address scalability concerns.

Robust choices of N’ under unknown test-time compute vs. ensembling

We 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’=1$ is optimal for $1\leq N\leq 3$.
• $N’=16$ is optimal between $2\leq N\leq 20$.
• $N’=256$ is optimal between $180\leq N\leq 385$.
• $N’=4096$ is optimal for $N\geq 1151$.

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

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

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

Settings with “soft” verification

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

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

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.

The effect of DCO on output diversity

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

Table 1. Mean and standard deviation of Shannon entropy for models trained with DCO of varying N’ values.

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

Computational overhead in the CoT setting

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

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

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.

Robustness of DCO

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

The 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’=1$ is optimal for $1\leq N\leq 3$.
• $N’=16$ is optimal between $2\leq N\leq 20$.
• $N’=256$ is optimal between $180\leq N\leq 385$.
• $N’=4096$ is optimal for $N\geq 1151$.

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

Moreover, 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.

Strategies for Supporting Multiple Test-time Budgets

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

Computational overhead in the CoT setting

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

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

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.

Comparison to Focal Loss

The 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. To 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:

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.

Table 2. Same as Table 1, but test on the out-of-distribution AIME24.

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

We 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’=1$ is optimal for $1\leq N\leq 3$.
• $N’=16$ is optimal between $2\leq N\leq 20$.
• $N’=256$ is optimal between $180\leq N\leq 385$.
• $N’=4096$ is optimal for $N\geq 1151$.

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

Clarification of the role of $\epsilon$ in the lemmas

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

Comparison to other inference-aware tuning methods

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

Effect of sequence length on the probability

We 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: In direct-answer tasks, DCO optimizes directly for the correct final answer, independent of sequence length. For formal math, we apply DCO at the step-level granularity, avoiding any length-related bias. In Chain-of-Thought scenarios, our method marginalizes over reasoning traces, ensuring no direct penalization of longer sequences. Thus, sequence length inherently does not disadvantage our proposed approach.

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

Reference:

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