Aligning Language Models with Selective Prediction

Abstract: LLMs are increasingly deployed as critical decision-making components in high-stakes real-world AI systems, rendering LLM reliability a foremost practical concern. In this paper, we focus on enhancing LLM reliability through selective prediction (SP), a strategy that allows an LLM to only predict for inputs where it is likely to be correct (i.e., coverage) and hence reduce the error rate (i.e., risk) on that portion of inputs -- flagging the remaining inputs for future human discretion. In other words, SP improves LLM reliability by balancing the risk-coverage trade-off and enabling seamless human-AI collaboration. To integrate SP into LLMs, we focus on the LLM post-training alignment stage and propose to align LLMs with SP performance metrics, in contrast with existing LLM alignment methods that focus primarily on correctness or calibration metrics. Specifically, we propose a novel alignment framework, Reinforcement Learning for Selection Reward (RLSR), which targets the area under the risk-coverage curve (AURC) -- a popular SP performance metric -- as its alignment objective. RLSR achieves substantially better risk-coverage trade-off compared to multiple alignment baselines on both in-domain and out-of-domain tasks.

• The paper introduces RLSR, a reinforcement learning framework that improves selective prediction by optimizing confidence ranking via Lifted AURC rewards.
• It employs a batch-wise ranking strategy that achieves lower AURC and superior risk-coverage trade-offs compared to conventional calibration and correctness-based methods.
• Empirical results on datasets like HotPotQA and BigMath confirm RLSR’s robustness and operational reliability in high-stakes applications.

Aligning LLMs with Selective Prediction: Formal Analysis and Empirical Insights

Motivation and Problem Formulation

Reliability is paramount in deploying LLMs in high-stakes environments such as medical diagnosis, finance, and law, where errors can have substantial consequences. Traditional LLM evaluation and alignment—focused on accuracy or confidence calibration (CC)—encourage prediction for all inputs, incentivizing confident but incorrect outputs and exacerbating hallucination phenomena. Selective Prediction (SP) offers an operationally critical alternative: the ability for an LLM to abstain from predicting on uncertain inputs, thereby controlling prediction risk and facilitating appropriate human-AI interaction.

The paper distinguishes SP fundamentally from CC, highlighting that calibration (the match between stated confidence and correctness probability) does not guarantee optimal selective prediction, nor does SP require perfect calibration. The SP paradigm requires models to maximally rank correct predictions above incorrect ones with respect to their confidence estimates, optimizing for area under risk-coverage curves (AURC) rather than traditional calibration error metrics. This dichotomy is visually and empirically demonstrated.

Figure 1: Perfect calibration does not necessarily yield optimal selective prediction, and vice versa; RLSR yields substantially improved SP metrics versus alternatives.

Reinforcement Alignment for Selective Prediction

AURC as an Objective

SP is formalized through the use of the risk-coverage curve, which quantifies the error rate (risk) as a function of the coverage (fraction of predictions not abstained). The area under this curve, AURC, is the canonical metric for SP and represents an integrated trade-off between risk reduction and prediction coverage. Importantly, AURC is driven by the ranking of prediction confidences, not their calibrated values.

A technical challenge arises because AURC is a population-level, ranking-based metric—difficult to optimize with mini-batch RL frameworks used for LLM fine-tuning. Additionally, the canonical AURC only penalizes errors, providing sparse and one-sided rewards, insufficient for gradient-based optimization.

Lifted AURC and RLSR Framework

To address these barriers, the paper introduces "Lifted AURC," a modified reward formulation assigning symmetric rewards: positive for correct, negative for incorrect, weighted according to confidence rank. The resulting reward, used within the Group Relative Policy Optimization (GRPO) RL framework, creates margin enforcement, actively encouraging separation between the confidence distributions for correct and incorrect predictions. This is empirically validated to produce larger confidence gaps than canonical AURC reward.

Figure 2: Lifted AURC induces greater separation between the normalized confidence distributions of correct and incorrect predictions, both on train and test sets.

To efficiently approximate population-level ranking in RL updates, a batch-wise pooling and ranking strategy is used: rollouts from multiple prompts are jointly ranked, and rewards are computed accordingly.

Comparison to Existing Alignment Paradigms

Alignment schemes such as RLVR (optimizing for correctness) and RLCR (joint correctness and calibration) are shown to be insufficient for optimal SP, neither targeting confidence ranking nor risk-coverage OPTIMA explicitly. RLVR's binary reward is sparse and fails to distinguish between highly confident and tentative responses, while RLCR focuses on calibration, which does not translate to improved selective prediction in practice.

Empirical Evaluation

Dataset and Model Setup

The evaluation uses Qwen2.5-7B and Llama-3.1-8B as base LLMs, aligned on HotPotQA-Modified (multi-hop reasoning with distractor evidence) and BigMath (multi-step math) datasets. SP performance is assessed in both in-domain and out-of-domain settings across reasoning, factuality, and ambiguity benchmarks, using exact-match, math-verify, and LLM-as-a-judge verifiers.

Risk-Coverage Analyses

RLSR demonstrates superior risk-coverage trade-off on all datasets and metrics:

Figure 3: Risk-coverage curves for Qwen2.5-7B aligned on HotPotQA; RLSR achieves sharper decline in risk as coverage is controlled, outperforming RLCR and RLVR.

Figure 4: Risk-coverage curves for Qwen2.5-7B aligned on BigMath; RLSR consistently yields lower risk across coverage levels.

Numerically, RLSR consistently achieves lower AURC and higher accuracy at fixed low coverages (10%, 25%, 50%), highlighting its value in operational scenarios where minimizing risk is prioritized over coverage.

Transfer and Robustness

RLSR's superiority persists in transfer to out-of-domain data, with minimal degradation, while RLVR and RLCR exhibit overfitting and reduced robustness to distributional shift. RLSR's batch-wise ranking approach ensures stable performance under varying batch sizes, and statistical significance is confirmed over multiple seeds and temperature settings.

Practical Deployment: Error Control

Deployment on MedQA (medical QA) demonstrates RLSR's ability to achieve stringent accuracy levels (75%) at meaningful coverage, outperforming RLVR and RLCR in both validation and test splits. Post-hoc calibration via temperature scaling fails to address ranking deficiencies, confirming the necessity of reward-driven confidence ranking optimization.

Related Work Context

The work situates itself within SP literature, differentiating its approach from post-hoc selector calibration, abstention-followup algorithms, and conformal prediction methods—all of which alter selection rules or confidence estimation outside end-to-end LLM training. The SP gap analysis highlights ranking error as the primary obstacle in attaining SP-optimality, which RLSR directly addresses.

Implications and Future Directions

RLSR provides a theoretically and empirically grounded mechanism for aligning LLMs to selective prediction, fundamentally improving reliability in risk-constrained deployments. By directly optimizing for confidence ranking, RLSR enables more robust abstention behavior, reduces confident errors, and enhances human-AI trust, particularly in sensitive domains.

Practically, this alignment paradigm can be combined with architectural advances, improved selectors, or mixed-objective RL schemes for further reliability gains. Theoretically, the paper suggests future modifications to focus on error-controlled coverage maximization, matching more closely real-world risk thresholds.

Conclusion

RLSR is a rigorously designed RL-based alignment framework for selective prediction in LLMs, leveraging lifted AURC rewards and batch-wise ranking to achieve substantially better risk-coverage trade-offs than prior correctness- or calibration-centric alternatives. Its demonstrated empirical superiority and theoretical soundness make it a preferred mechanism for reliability-driven LLM alignment, with broad applicability across domains and datasets.