Supervised Optimism Correction: Be Confident When LLM s Are Sure (2025)

Abstract: In this work, we establish a novel theoretical connection between supervised fine-tuning and offline reinforcement learning under the token-level Markov decision process, revealing that large language models indeed learn an implicit Q -function for inference.Through this theoretical lens, we demonstrate that the widely used beam search method suffers from unacceptable over-optimism, where inference errors are inevitably amplified due to inflated Q -value estimations of suboptimal steps. To address this limitation, we propose **S**upervised **O**ptimism **C**orrection (SOC), which introduces a simple yet effective auxiliary loss for token-level Q -value estimations during supervised fine-tuning. Specifically, the auxiliary loss employs implicit value regularizationto boost model confidence in expert-demonstrated responses, thereby suppressing over-optimism toward insufficiently supervised responses.Extensive experiments on mathematical reasoning benchmarks, including GSM8K, MATH, and GAOKAO, showcase the superiority of the proposed SOC with beam search across a series of open-source models.

## introduction recent advances in large language models (llms) have demonstrated remarkable success across diverse tasks such as instruction following @xcite @xcite @xcite , code generation @xcite @xcite , and medical diagnosis @xcite . within these developments, complex reasoning capabilities have attracted increasing attention from research communities, attributed to their capability of enabling llms to tackle intricate problemsolving @xcite @xcite . despite these achievements, learning to reason remains a critical yet challenging task for llms, particularly for smaller models with limited parameters @xcite . the inherent complexity stems from inefficient exploration of llms, as the combinatorial nature of the vocabulary space results in an exponential growth of potential reasoning paths @xcite . a prevalent paradigm for enhancing reasoning capabilities involves supervised fine-tuning (sft) on high-quality demonstration data @xcite @xcite , where models learn to imitate expert reasoning patterns through next-token prediction. the complementary strategies employ search-based decoding techniques during inference @xcite @xcite , such as beam search, which aims to enhance reasoning by exploring multiple candidate pathways. however, there is often an overlooked disconnect between the local token-level optimization of sft and the global sequence-level objectives pursued by searchbased decoding. this disconnect results in a critical misalignment: while sft focuses on maximizing the likelihood of individual token predictions, search-based decoding operates by scoring entire sequences, which may not directly align with the local training objectives of llms. by addressing this gap, we can potentially unlock more robust reasoning capabilities in llms, aligning training objectives more closely with inference-time goals. in this work, we employ the token-level markov decision process to establish a novel theoretical connection between sft and offline reinforcement learning (rl) @xcite @xcite . we theoretically demonstrate that llms indeed learn an implicit q-function to estimate the expected utility of token sequences during sft. through this lens, we further explore the over-optimism problem in the widely used beam search method, re- figure 1 : illustration of supervised optimism correction. guided by an implicit q-function, beam search suffers from the over-optimism problem during the decoding process, which confuses llm for reliable response choice. in particular, the over-optimism can amplify errors during beam search, leading to the selection of incorrect trajectories with higher q-values. to alleviate this problem, soc introduces an auxiliary loss during sft to boost model confidence in expert-demonstrated responses, about which llm should be sure. to alleviate this problem, we propose supervised optimism correction (soc), which introduces a simple auxiliary loss during sft to give supervised responses a state-value bonus. unlike prior rl-based methods that require explicit reward modeling, soc operates purely within the sft paradigm by imposing implicit value regularization. this regularization boosts confidence in expert-demonstrated data while potentially penalizing high q-value estimations for insufficiently supervised reasoning steps. as a result, the model learns to autonomously prune low-quality reasoning paths during inference without relying on external verifiers or reward models. for instance, when encountering erroneous intermediate steps, the corrected q-values suppress further exploration of those branches, mirroring human-like error recognition and recovery patterns, requiring no architectural modifications or additional inference-time computations. our core contributions are summarized as follows: ## supervised optimism correction in subsection 3.1, we first discuss the overoptimism problem observed in beam search during inference. next, subsection 3.2 investigates a potential cause of this issue by analyzing how inflated q-value estimation errors, particularly those arising from insufficiently supervised states, can amplify over-optimism through the maximization operation of beam search. finally, subsection 3.3 introduces our proposed supervised optimism correction (soc) method, which incorporates an auxiliary v loss to better align the response selection process with expert-demonstrated responses. ## theoretical analysis in this section, we provide a theoretical analysis of our objective. theorem 4.1 (contraction of value differences). let v θ (s) be the approximate value function of state s. suppose that the value function is always positive for any state s. if the objective in sft includes a additional term in equation (17), then after one step of gradient descent with learning rate α ∈ [0, 1], the gap between adjacent states' values contracts, i.e., we provide a detailed proof in appendix a. theorem 4.1 states that the auxiliary loss encourages the implicit value function of neighboring states to become closer. to further illustrate the benefit of this regularization, consider the policy evaluation problem in a token-level mdp with sparse rewards, where the agent only receives a reward at the end of the trajectory while all intermediate rewards are zero. in this setting, the optimal value function satisfies the bellman equation: since the reward is zero for all intermediate states, it follows that: ## experiments in this section, we empirically demonstrate the effectiveness of soc on the mathematical reasoning ability of llms. in subsection 5.1 and 5.2, we introduce the datasets and setup of our experiments. we present the results of llms scaling from 1.5b to 7b parameters in subsection 5.3 and provide an ablation study in subsection 5.4.