LEC: Linear Expectation Constraints for Selection-Conditioned Risk Control in Selective Prediction and Routing Systems

Foundation models often generate unreliable answers, while heuristic uncertainty estimators fail to fully distinguish correct from incorrect outputs, causing users to accept erroneous answers without any statistical guarantee. We address this problem through selection-conditioned risk control, aiming to ensure that an accepted prediction has an error probability no larger than a user-specified risk level. To this end, we propose LEC, a principled framework that reframes selective prediction as a decision problem governed by a linear expectation constraint over selection and error indicators. This formulation directly controls the ratio between the expected number of accepted errors and the expected number of accepted predictions, which corresponds to the marginal error probability conditioned on selection. Under exchangeability, we derive a finite-sample sufficient condition that relies only on a held-out calibration set, enabling the computation of a risk-constrained, retention-maximizing threshold. Furthermore, we extend LEC to two-model routing systems: if the primary model's uncertainty exceeds its calibrated threshold, the input is delegated to a subsequent model, while maintaining system-level selection-conditioned error control. Experiments on both closed-ended and open-ended question answering (QA) and vision question answering (VQA) demonstrate that LEC maintains the prescribed risk level in accepted predictions and substantially improves sample retention compared to baselines.

AI systems such as chatbots and vision-language models are increasingly used to answer questions and support decisions, but they can still give wrong answers with high confidence. This can be risky when users do not know which answers to trust. Our paper proposes a simple way to help an AI system decide when to answer, when to pass the question to another model, and when to avoid answering. The main idea is to use a small set of checked examples to learn a rule for accepting answers. The user chooses an acceptable error limit, and our method selects a threshold so that, among the answers the system chooses to provide, the error rate is kept within that limit. We also extend this idea to systems with two models, where uncertain questions can be sent from a first model to a second model. Experiments on text question answering and visual question answering show that our method keeps the accepted answers within the desired error limit while allowing the system to answer more questions than several existing methods. This makes AI systems more useful and safer in settings where users need both good coverage and clearer control over mistakes.