Technical Summary: Aligning Language Models with Selective Prediction

Problem Statement

Large Language Models (LLMs) are increasingly deployed in high-stakes domains such as finance, healthcare, and law. In these settings, the cost of an erroneous decision often significantly outweighs the cost of deferring to a human. Standard LLM training paradigms require models to answer every query, regardless of their internal certainty, leading to "confident mistakes" and hallucinations.

While Confidence Calibration (CC) is often viewed as a prerequisite for Selective Prediction (SP), the authors argue that CC is neither sufficient nor necessary for effective SP.

• CC ensures that assigned confidence scores match the true probability of correctness (e.g., an 80% confidence implies an 80% chance of being right).
• SP requires that correct predictions are ranked higher than incorrect ones based on confidence, allowing the model to abstain from low-confidence queries to reduce the error rate on the selected subset.

The paper demonstrates that perfect calibration does not guarantee perfect selective prediction (as shown in Fig. 1), and conversely, perfect SP can exist with ill-calibrated confidence if the ranking is correct. Existing alignment methods, such as Reinforcement Learning with Verifiable Rewards (RLVR) and Reinforcement Learning with Calibration Rewards (RLCR), focus on accuracy or calibration metrics, respectively, but fail to directly optimize the risk-coverage trade-off required for SP.

Methodology: Reinforcement Learning for Selection Reward (RLSR)

The authors propose RLSR, a novel alignment framework that directly optimizes LLMs for Selective Prediction performance using the Group Relative Policy Optimization (GRPO) algorithm. The core objective is to minimize the Area Under the Risk-Coverage Curve (AURC), a standard metric summarizing the trade-off between the error rate (risk) and the proportion of inputs answered (coverage).

Key Technical Challenges and Solutions

Integrating AURC into a reinforcement learning framework presents two primary challenges, which RLSR addresses through specific innovations:

1. Sparse and One-Sided Rewards:

   • Challenge: The canonical AURC formulation only penalizes errors (incorrect predictions) and provides no reward signal for correct predictions. This leads to sparse learning signals and fails to enforce a margin between the confidence of correct and incorrect answers.
   • Solution: The authors introduce a Lifted AURC objective. They reformulate the binary error indicator $I\{f(x) \neq y\}$ into a signed variable $(2I\{f(x) \neq y\} - 1) \in \{-1, +1\}$. This creates a "push-pull" dynamic where correct predictions receive positive rewards and incorrect ones receive negative rewards, actively encouraging the model to separate the confidence distributions of correct and incorrect answers.

2. Global Ranking Dependency:

   • Challenge: AURC is a population-level metric dependent on the global ranking of all samples' confidence scores. Standard mini-batch stochastic gradient descent (SGD) cannot compute global rankings without iterating over the entire dataset, which defeats the purpose of efficient batch training.
   • Solution: RLSR employs a Batch Training Strategy. For each training step, the model generates $G$ rollouts for each of $B$ prompts. These $B \times G$ samples are pooled into an "effective batch." Ranks and weights are computed within this pooled batch to approximate the global AURC. This allows the algorithm to remain compatible with standard GRPO while approximating the population-level metric.

Algorithm Overview

The reward function $R_{RLSR}$ for a sample $i$ is defined as:
$$R_{RLSR}(y_i, \hat{y}_i, s_i) = \begin{cases} +\hat{\alpha}_i & \text{if } \hat{y}_i = y_i \\ -\hat{\alpha}_i & \text{if } \hat{y}_i \neq y_i \end{cases}$$
where $\hat{\alpha}_i$ is a weight derived from the sample's rank within the effective batch (based on the harmonic number formulation of AURC). The model is then updated using the GRPO advantage function, which normalizes rewards within groups of rollouts for the same prompt.

Key Contributions

1. Direct SP Alignment: This is the first work to propose aligning pretrained LLMs directly with Selective Prediction metrics, rather than relying on intermediate objectives like calibration.
2. RLSR Framework: The introduction of a new alignment framework that utilizes a Lifted AURC objective and a batch-based training strategy to overcome the technical hurdles of optimizing AURC via RL.
3. Empirical Superiority: Comprehensive experiments demonstrating that RLSR outperforms strong baselines (RLVR and RLCR) in both in-domain and out-of-domain settings.

Experimental Results

The authors evaluated RLSR on Qwen2.5-7B and Llama-3.1-8B models across multiple datasets (HotPotQA, BigMath, and various OOD benchmarks like GSM8K, GPQA, and MedQA).

• Risk-Coverage Trade-off: RLSR consistently achieved lower AURC scores (indicating better performance) compared to RLVR and RLCR. For example, on HotPotQA (in-domain), RLSR achieved an AURC of 0.44, significantly outperforming RLVR (0.56) and RLCR (0.51).
• Accuracy at Low Coverage: In high-stakes scenarios where low coverage is acceptable for high accuracy, RLSR showed superior performance. At 10% coverage on HotPotQA, RLSR reached 66.0% accuracy, compared to 45.5% for RLVR and 53.5% for RLCR.
• Out-of-Domain (OOD) Generalization: RLSR maintained strong performance on OOD tasks, whereas RLVR and RLCR often degraded or showed inconsistent improvements. This suggests RLSR learns a more robust ranking of confidence rather than overfitting to specific distributional properties.
• Risk-Controlled Deployment: In a simulated medical deployment (MedQA) with a target accuracy of 75%, RLSR was the only method capable of achieving the target accuracy while maintaining a reasonable coverage rate. RLVR failed to reach the target accuracy regardless of the threshold, and RLCR's performance dropped significantly on the test set compared to validation.
• Confidence Separation: Analysis of the confidence distributions (Fig. 3) confirmed that the Lifted AURC reward successfully created a larger margin (gap) between the confidence scores of correct and incorrect predictions compared to standard AURC or calibration-based rewards.

Significance and Claims

The paper claims that RLSR provides a principled approach to enhancing LLM reliability for real-world deployment. By directly optimizing the risk-coverage trade-off, RLSR enables LLMs to effectively "abstain" from answering when uncertain, thereby reducing error rates in critical applications without requiring human intervention for every query.

The authors emphasize that while calibration (ECE) is a useful metric, it is not a proxy for selective prediction performance. Optimizing for calibration does not necessarily improve the ranking of correct vs. incorrect answers, which is the fundamental requirement for SP. Consequently, RLSR represents a more direct and effective path to deploying LLMs in high-stakes environments where error control is paramount.

Limitations: The authors note that their current alignment focuses on the holistic AURC metric. They acknowledge that in practical high-stakes settings, a specific target risk level is often known in advance, suggesting that future work could formulate risk-constrained coverage maximization objectives.