Aligning Language Models with Selective Prediction
Abstract
Large language models (LLMs) are increasingly deployed as decision making components in real-world systems for societal and scientific applications, creating a growing need for reliable predictions. In this paper, we study the problem of reliable decision making with LLMs via the lens of selective prediction, allowing the model to improve performance by trading-off coverage. The aim of selective prediction is to balance the key tradeoff of risk and coverage, where risk measures the predictive performance on selected inputs and coverage measures the fraction of abstentions. While existing LLM post-training approaches focus primarily on correctness or calibration, we propose to directly optimize for selective prediction performance by introducing reinforcement Learning for Selection Reward (RLSR), which targets the area under the risk-coverage curve (AURC) as its training objective. RLSR achieves substantially better risk-coverage tradeoff compared to multiple baselines on both in-domain and out-of-domain tasks.