CarBoN: Calibrated Best-of-N Sampling Improves Test-time Reasoning

TL;DR

The paper addresses the inefficiency of test-time scaling with Best-of-$N$ sampling by introducing a general test-time calibration framework that steers generation toward high-reward reasoning paths without retraining the LLM. It instantiates CarBoN as a two-phase method: first explore with $N_1$ samples, then calibrate logits via an input-specific additive shift $\delta$ and temperature $T$ to guide the remaining $N_2$ rollouts, with the final answer drawn from all $N$ candidates. The authors prove the existence of an improving joint calibration $(\delta^*, T^*)$ and show that such calibration strictly increases the lower bound on the expected Best-of-$N$ reward, providing a theoretical foundation for the approach. Empirically, CarBoN delivers consistent gains on MATH-500 and AIME-2024 across multiple models, reducing the required rollout budget by up to $4\times$ and, in some cases, surpassing larger LLMs under fixed budgets; the work also demonstrates that temperature and delta impart complementary benefits and generalizes to beam-search style decoding. Overall, the framework offers a practical, provably beneficial path to more reliable test-time reasoning and broadens the applicability of test-time calibration beyond Best-of-$N$ to other step-level strategies.

Abstract

Allocating more computation during inference time (test-time scaling) improves language model performance, especially for reasoning tasks. However, popular methods like Best-of-$N$ sampling often show diminishing returns as $N$ increases. To address this inefficiency, we introduce a general test-time calibration framework that adaptively modifies the model toward high-reward reasoning paths, with theoretical guarantees of improving the lower bound of expected reward under finite sampling, all without large language model (LLM) retraining. Within this framework, we propose CarBoN (Calibrated Best-of-$N$), a two-phase method that first explores the solution space and then learns a calibration of the logits via an input-specific temperature $T$ and additive shift vector $δ$, guiding generation toward more reliable reasoning. Experiments on MATH-500 and AIME-2024 show that CarBoN improves efficiency, with up to $4\times$ fewer rollouts to reach the same accuracy, while often achieving higher accuracy under fixed budgets. We also analyze the complementary roles of $T$ and $δ$ in balancing output diversity and correctness, and demonstrate that the framework also generalizes to step-level sampling strategies such as beam search. For more information, please refer to our project page at huggingface.co/spaces/TrustSafeAI/Test-Time-Calibration.

Figures (5)

• Figure 1: Reward‑guided calibration accelerates binary search. Left: Increasing per‑step noisy reward (inverse‑distance signal + noise) lowers average search steps versus vanilla. Right: Example showing reward guidance converges early; vanilla keeps oscillating. See Appendix \ref{['appendix:toy-example']} for details.
• Figure 2: (a) Test-time calibration framework. With a rollout budget $N = N_1 + N_2$, the model first explores by generating and scoring $N_1$ candidate responses. The model then learns calibration parameters $(\delta, T)$ from high-scoring responses, , using them to adjust the logits for the remaining $N_2$ generations. The final answer is selected from all $N$ candidates. (b) MATH-500 Results. CarBoN improves weighted Best-of-N accuracy across four models. For all models, calibrated accuracy at $N=64$ (orange dash line) matches or exceeds uncalibrated accuracy at $N=256$, corresponding to up to a $4\times$ reduction in rollout budgets. Notably, with Qwen2.5-Math-1.5B-Instruct at $N=64$, CarBoN surpasses GPT-4o (red dashed line), while uncalibrated Best-of-N with $N=256$ does not.
• Figure 3: Correlation between problem difficulty, calibrated temperature, and top-$k$ completion entropy on MATH-500. Bars (left y-axis) show the average learned temperature across five difficulty levels, while the line plot (right y-axis) shows the normalized entropy of the top-$k$ completions used for calibration. Both temperature and entropy strongly increase with problem difficulty, indicating that harder problems require higher temperatures to capture the more diverse top-$k$ token distributions. See Appendix \ref{['appendix:cor_level_temp']} for full Spearman correlation statistics.
• Figure 4: Learned calibration temperatures across different rollout budget $N$ and problem difficulty levels. Both larger $N$ and higher difficulty consistently lead to higher $T$, reflecting the increased diversity of top-$k$ completions. This adaptive scaling complements our earlier grid-search results in Appendix \ref{['appendix:llama-temp']}, where small $N$ favored lower temperatures but larger $N$ required higher temperatures to fully leverage the broader exploration.
• Figure 5: Results of Llama-3.2-1B-Instruct on MATH-500 with different temperatures $T$ for best-of-$N$ inference. From left to right: majority voting, naive, and weighted selection. Blue curves indicate lower temperatures, red curves indicate higher temperatures. Lower temperatures improve accuracy, but overly low temperatures limit diversity and the benefit of increasing $N$.

Theorems & Definitions (6)

• Lemma 1: Existence of an Improving Joint Solution $(\delta, T)$
• Theorem 2: Joint Calibration $(\delta, T)$ Improves Expected Reward from Best-of-$N$ Sampling
• Corollary 3: Sub-optimality of Exploitation Alone
• proof
• proof
• proof