Aligning Language Models with Selective Prediction
For practitioners deploying LLMs in high-stakes applications, this work provides a method to improve reliability by allowing models to abstain from uncertain predictions, enabling safer human-AI collaboration.
The paper addresses the reliability of large language models (LLMs) by integrating selective prediction (SP) during post-training alignment. The proposed framework, RLSR, optimizes the area under the risk-coverage curve (AURC) and achieves substantially better risk-coverage trade-off than multiple alignment baselines on both in-domain and out-of-domain tasks.
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.