Beyond Binary Rewards: Training LMs to Reason About Their Uncertainty

TL;DR

<3-5 sentence high-level summary> RLCR addresses the instability in calibration caused by standard RL training with binary rewards for reasoning tasks. It jointly optimizes correctness and calibrated confidence by augmenting the reward with a bounded proper scoring rule (Brier score) and requires LMs to output both an answer and a verbalized confidence. Across HotPotQA, Math, and related datasets, RLCR maintains or matches accuracy while substantially improving calibration, with test-time confidence enabling further gains through scaling and ensembling. The work demonstrates a practical path toward uncertainly-aware, more trustworthy reasoning in language models, including better cross-domain robustness.</3-5 sentence high-level summary>

Abstract

When language models (LMs) are trained via reinforcement learning (RL) to generate natural language "reasoning chains", their performance improves on a variety of difficult question answering tasks. Today, almost all successful applications of RL for reasoning use binary reward functions that evaluate the correctness of LM outputs. Because such reward functions do not penalize guessing or low-confidence outputs, they often have the unintended side-effect of degrading calibration and increasing the rate at which LMs generate incorrect responses (or "hallucinate") in other problem domains. This paper describes RLCR (Reinforcement Learning with Calibration Rewards), an approach to training reasoning models that jointly improves accuracy and calibrated confidence estimation. During RLCR, LMs generate both predictions and numerical confidence estimates after reasoning. They are trained to optimize a reward function that augments a binary correctness score with a Brier score -- a scoring rule for confidence estimates that incentivizes calibrated prediction. We first prove that this reward function (or any analogous reward function that uses a bounded, proper scoring rule) yields models whose predictions are both accurate and well-calibrated. We next show that across diverse datasets, RLCR substantially improves calibration with no loss in accuracy, on both in-domain and out-of-domain evaluations -- outperforming both ordinary RL training and classifiers trained to assign post-hoc confidence scores. While ordinary RL hurts calibration, RLCR improves it. Finally, we demonstrate that verbalized confidence can be leveraged at test time to improve accuracy and calibration via confidence-weighted scaling methods. Our results show that explicitly optimizing for calibration can produce more generally reliable reasoning models.

Figures (6)

• Figure 1: (a): Sample chain-of-thought from a model trained with RLCR, using <think>, <answer>, <analysis>, and <confidence> tags. (b) On in-domain evaluation tasks, RLCR improves on standard reasoning training (RLVR) and even slightly outperforms a combination of RLVR and a dedicated classifier trained to predict RLVR correctness. (c) When evaluating generalization to novel tasks, RLCR improves both accuracy and calibration, while other methods leave accuracy unchanged and sometimes harm calibration. All results shown are for HotpotQA, see \ref{['sec:experiments']} for results on math tasks and additional baselines.
• Figure 2: (a): RLVR focuses solely on correctness, which can incentivize guessing. (b): RLCR uses a calibrated reward that jointly optimizes for correctness and calibration.
• Figure 3: (a)Reward curves for RLCR (ours) and RLVR. Both correctness and calibration rewards improve under our method, demonstrating simultaneous gains in correctness and calibration. The Brier reward is shifted upward by 1 for clarity. (b) Completion lengths during training. The completion lengths of our method gradually increase during training as uncertainty reasoning improves.
• Figure 4: Test-time scaling curves.(a)Accuracy vs Number of Samples (N). Accuracy improves for all methods with increasing compute. Confidence-weighted majority vote outperforms both vanilla majority vote and max-confidence, highlighting complementary benefits of combining voting with confidence scores. (b) Brier Scores vs Ensemble Size (K). Here we evalute the effect of applying test-time scaling to confidence estimation alone, resampling multiple analyses (blue text in \ref{['fig:formatted-tags-ex']}). Calibration improves as the size of the analysis ensemble grows.
• Figure 5: Brier scores (a) and ECE (b) of baseline / analysis classifiers on HotPotQA-Modified across three model sizes. Analysis classifiers outperform baselines at smaller sizes, suggesting that uncertainty CoT is essential for better calibration when capacity is limited.

Theorems & Definitions (8)

• Theorem 1
• Lemma 1: Calibration incentive
• proof
• Lemma 2: Correctness incentive
• proof
• proof
• proof : Proof of \ref{['thm:main']}
• Corollary 1