Aligning Language Models with Selective Prediction
2026-07-03 • Machine Learning
Machine LearningArtificial IntelligenceComputation and Language
AI summaryⓘ
The authors address making large language models (LLMs) more reliable by teaching them to only answer questions they are confident about, and let humans handle the rest. They introduce a new method called Reinforcement Learning for Selection Reward (RLSR), which helps the model decide when to respond or abstain, improving the balance between accuracy and coverage. Their approach focuses on optimizing a special measure called the area under the risk-coverage curve (AURC) instead of usual accuracy alone. Tests show RLSR works better than other methods both on familiar and new types of tasks.
large language modelsselective predictionrisk-coverage trade-offreinforcement learningmodel alignmentAURCpost-training alignmenthuman-AI collaborationerror rateout-of-domain tasks
Authors
Gaoxiang Luo, Yifan Wu, Sinian Zhang, Aryan Deshwal, Ju Sun
Abstract
Large language models (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.