Sample Smart, Not Hard: Correctness-First Decoding for Better Reasoning in LLMs

TL;DR

This work reframes decoding for reasoning in LLMs as a correctness-driven problem rather than a pure exploration or confidence-driven task. By defining rank-wise accuracy and constructing a calibration grid that links token probability to expected correctness, the authors develop Greedy-Threshold, Calibrated-TopK, and Calibrated-$\varepsilon$ strategies that constrain sampling to high-correctness tokens. Across GSM8K, MMLU-Pro, BBH, and AIME, these correctness-first samplers yield consistent gains, with Calibrated-TopK and Calibrated-$\varepsilon$ providing the strongest improvements and Greedy-Threshold offering notable benefits for smaller models, all with negligible decoding overhead. The results support a shift in decoding practice toward correctness-aware truncation, showing that reducing sampling in low-confidence, high-risk steps can prevent error amplification while preserving useful diversity when appropriate. Overall, the work demonstrates that rank-based calibration signals offer a robust, transferable approach to improving reasoning capabilities in advanced LLMs and prompts a broader rethinking of uncertainty handling in autoregressive decoding.

Abstract

Large Language Models (LLMs) are increasingly applied to complex tasks that require extended reasoning. In such settings, models often benefit from diverse chains-of-thought to arrive at multiple candidate solutions. This requires two competing objectives: to inject enough stochasticity to explore multiple reasoning chains, and to ensure sufficient accuracy and quality in each path. Existing works pursue the first objective by increasing exploration at highly uncertain steps with higher temperature or larger candidate token sets, while others improve reliability by rejecting samples with low confidence post-generation, implying that low confidence correlates with low answer quality. These two lines of thought are in conflict, as they conflate different sources of uncertainty. To resolve this, we argue that the decoding rule should be calibrated by correctness, not confidence alone. We should sample from tokens with higher estimated correctness, and reduce sampling where expected correctness is low. We propose simple strategies that achieve this goal: Greedy-Threshold makes sampling greedy at very low confidence steps. Calibrated-TopK and Calibrated-epsilon set truncation threshold based on estimated rank-wise correctness. Together, our findings challenge prevailing heuristics about decoding under uncertainty and show gains across math and general reasoning benchmarks.

Figures (15)

• Figure 1: Comparison of common and our proposed truncation strategies. Each panel shows which tokens remain available for sampling, with tokens ordered from highest to lowest model-assigned probability (left to right). The $y$-axis represents the max token probability (“confidence”). Our methods explicitly suppress low-confidence tail tokens.
• Figure 2: Calibration grid of Qwen2.5-1.5B-Instruct with 5 bins shows the average probability $\hat{p}$ and correctness $\hat{c}$ for each confidence-bin and rank. Correctness is notably low in the lower-confidence bins, and decreases as rank increases. Percentages indicate frequency of occurrence of this bin.
• Figure 3: Expected accuracy increases with confidence across all model sizes. In the lowest confidence bin, expected accuracy drops regardless of model size. Frequency refers to the proportion of decoding steps whose maximum probability falls into each confidence bin. Larger models assign more predictions to the 0.9–1.0 confidence range, where both accuracy and frequency are highest, reflecting stronger benchmark performance. In contrast, smaller models place more probability mass in low-confidence bins, where accuracy is poor.
• Figure 4: Plot of the majority voted accuracy and the number of unique answers as the number of samples increase. Sampling is greedy unless the maximum probability falls into a certain confidence bin, in which case sample from the full token distribution. Sampling at the lowest confidence bin results in no gain in accuracy while contributing little to diversity in terms of number of unique answers.
• Figure 5: Effect of low-confidence events on sequence-level accuracy. (a) Accuracy decreases both when the model directly samples low-probability tokens ($p < 0.1$, blue) and when it is in a low-confidence state regardless of the sampled token ($p_{\max}<0.3$, orange). (b) Accuracy also drops as the mean rank of sampled tokens increases, showing that drifting into lower-ranked tokens degrades sequence quality.