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

AI summary

Overview

• Research area: Selective prediction and uncertainty quantification for foundation models (LLMs and large vision-language models), combined with distribution-free statistical risk control.
• Technical level: Advanced. The paper assumes familiarity with split conformal prediction, exchangeability, empirical risk, false discovery rate style metrics, and PAC-style high-probability bounds.
• Scope: The paper introduces LEC (Linear Expectation Constraints), a calibration framework that guarantees the error rate among accepted predictions stays below a user-specified risk level α, for both single-model selective prediction and two-model routing systems.

What This Paper Is About

Foundation models often produce wrong answers, and uncertainty scores drawn from model logits, self-consistency, or semantic entropy do not cleanly separate correct from incorrect outputs — so users can accept erroneous answers with no statistical guarantee. The paper reframes selective prediction as a decision problem governed by a constraint on the expectation of a linear functional of two binary indicators (whether a prediction is selected and whether it is wrong), so that the marginal error probability given selection is provably bounded by α. It then extends the same construction to a two-model routing system in which uncertain inputs are escalated to a second model, while preserving a system-level guarantee.

Key Contributions

1. A reformulation of selective prediction as a linear expectation constraint. The selection-conditioned error rate is written as E[Z^(a)(λ^(a))] / E[S^(a)(λ^(a))], where Z^(a) = S^(a) · err^(a); the requirement SCER ≤ α becomes equivalent to E[Z^(a)(λ^(a)) − α·S^(a)(λ^(a))] ≤ 0 (Eq. 4).

2. A finite-sample sufficient condition and retention-maximizing threshold (Theorem 3.1). Using only a held-out calibration set and a leave-one-out style correction, the paper requires Σ_{j=1}^{k(λ)} (err_(j) − α) ≤ −1 over the k(λ) calibration points accepted at threshold λ, then picks the largest feasible threshold. Under exchangeability, the calibrated threshold satisfies Pr(err_{n+1} = 1 | u_{n+1} ≤ λ̂) ≤ α, with the guarantee vacuous if nothing is accepted.

3. Extension to two-model routing with joint threshold calibration (Theorem 3.2). System-level selection and error indicators are defined so that routing selects at most one prediction, the same linear-constraint argument is applied to the pair (Z_i, S_i), and (λ̂^(a), λ̂^(b)) are chosen to maximize empirical acceptance subject to the constraint. When the feasible set is empty, the risk level is declared infeasible and the system abstains on all inputs. Appendix B outlines extension to routing with more than two models.

4. An empirical study across QA and VQA benchmarks and multiple uncertainty estimators. LEC was tested on CommonsenseQA and TriviaQA with eight LLMs (LLaMA, Qwen, Vicuna, OpenChat families) and on ScienceQA and MM-Vet v2 with four LVLMs (LLaVA1.5, LLaVA-NeXT, InternVL2 groups), with predictive entropy, black-box semantic entropy, graph-Laplacian eigenvalue, degree-matrix, eccentricity, and length-normalized PE estimators.

Main Findings

• The guarantee holds empirically on closed-ended QA. On CommonsenseQA at risk level 0.05 with OpenChat-3.5, LEC achieves an average empirical accepted error rate of 0.0497 across 500 random calibration–test splits.

• LEC controls risk more tightly than confidence-bound baselines. Confidence interval methods (UCB-CLP, UCB-HFD, and the routing variants UCB-CLP-Routing and UCB-HFD-Routing) stay well below the target risk level, whereas LEC operates near it. On TriviaQA with Qwen2.5-3B, LEC attains an empirical accepted error rate of 0.0987 versus 0.0878 for UCB-CLP, while UCB-HFD fails to identify feasible thresholds at all.

• Tighter control translates into higher power (more accepted predictions). Across all evaluated LLMs and risk levels, LEC's power exceeds the UCB baselines. In Table 1 on TriviaQA, at α = 0.05 with OpenChat-3.5 the powers are 0.7230 (LEC), 0.6684 (UCB-CLP) and 0.6091 (UCB-HFD); at α = 0.1 with LLaMA-3.1-70B they are 1.0 (LEC), 1.0 (UCB-CLP) and 1.0 (UCB-HFD); at α = 0.1 with Qwen2.5-14B they are 1.0 (LEC), 0.9987 (UCB-CLP) and 0.9718 (UCB-HFD). At α = 0.05 with Qwen2.5-14B, LEC reports a power of 0.7193 and retains 9.5% more admissible samples than UCB-CLP, while UCB-HFD yields no feasible threshold.

• The abstract reports a headline retention gain. The introduction summarizes the improvement over confidence-interval-based methods as "+9%" in accepted samples on TriviaQA.

• Joint calibration is necessary in the routing setting. Applying LEC independently to each model without joint threshold calibration does not achieve a valid system-level guarantee, whereas LEC-Routing with a single constraint over system-level selection and error indicators does.

• Routing increases coverage while respecting the risk budget. On CommonsenseQA at α = 0.05, Qwen2.5-3B alone accepts only 20.3% of test samples; LEC-Routing with Qwen2.5-3B as primary and LLaMA-3.1-8B as secondary accepts 33.9% in total, split as 17.44% handled by Qwen2.5-3B and 16.46% routed to LLaMA-3.1-8B, a 13.6% absolute increase.

• Routing is adaptive rather than blindly escalating. At the same risk level, Qwen2.5-7B alone accepts 50.7% and LLaMA-3.1-8B alone accepts 33.4%; LEC-Routing reaches a 51.0% acceptance rate with the majority (48.65%) still processed by the primary Qwen2.5-7B model.

• The system accepts more correct samples in absolute terms. Table 2 on CommonsenseQA reports accepted correct sample counts; for the Qwen2.5-7B & LLaMA-3.1-8B pair these are 2413, 3144, 3615, 3896, 3928 and 3928 at risk levels 0.05, 0.1, 0.15, 0.2, 0.25 and 0.3 respectively, versus 2392, 3078, 3523, 3858, 3924 and 3924 for Qwen2.5-7B alone.

• Advantages persist under alternative correctness judgments. With bi-entailment as the alignment criterion in the admission function A, LEC maintains tight selection-conditioned error control and higher power than UCB baselines across five LLMs and all tested risk levels.

• Robustness is reported across configuration choices. LEC is described as maintaining statistical rigor and higher power across different UQ methods, admission functions, calibration–test split ratios, and sampling sizes under black-box scenarios.

• Results for ScienceQA and MM-Vet v2 are not reported in the available content. The benchmarks and LVLMs are listed as part of the experimental setup, but no ScienceQA or MM-Vet v2 result figures or tables appear in the text available here.

Methodology in Plain English

The starting point is a standard split-calibration setup: a dataset is divided into a held-out calibration set and a test set. For every calibration example, the researcher records an uncertainty score (from entropy, semantic entropy, or another estimator) and whether the model's answer was actually wrong under an alignment check such as sentence similarity at a 0.6 threshold, bi-entailment, or LLM-as-a-Judge.

Instead of asking "is the uncertainty score good at ranking correct above incorrect?", the paper asks a single statistical question: over the calibration points that would be accepted if the threshold were set to λ, is the accumulated quantity (number of errors minus α times number of selections) at most −1? Because the threshold changes which points are counted, the set of λ satisfying this condition can be scanned directly, and the largest such λ is chosen because it accepts the most examples. A −1 leave-one-out style correction is what converts this empirical check into a finite-sample guarantee.

For two models, the same idea is applied at the system level. Each calibration example yields two uncertainty scores and two error indicators. An input is handled by the primary model if its score is at or below λ^(a), otherwise it is routed to the second model and accepted if that model's score is at or below λ^(b), and otherwise the system abstains. Because each input can be accepted by at most one model, the system-level selection indicator stays binary, and the pair (λ^(a), λ^(b)) can be searched jointly to maximize the number of accepted calibration examples while satisfying the same expectation constraint.

Evaluation follows prior protocols: 500 random calibration–test splits, reporting mean and standard deviation. Hyperparameters include beam search with num_beams = 5, sampling 10 answers per input for uncertainty quantification in open-domain QA, and a calibration–test split ratio fixed at 0.5. The plotted "FDR" axis is a shorthand for the observed proportion of erroneous predictions among accepted ones; the paper states the theoretical guarantee is always stated in terms of selection-conditioned error rate.

Why This Matters

• Impact on research: The paper shifts selective prediction away from ranking heuristics and worst-case confidence bounds toward a direct expectation constraint, which is both statistically valid and less conservative. Because the guarantee follows from a linear decomposition rather than model-specific assumptions, it applies to any deterministic routing policy where each example goes to at most one model, and it composes across routing depths.

• Real-world applications:

  • Clinical or legal question answering, where an LLM that answers directly must keep its error rate under a stated ceiling rather than returning a set of candidate answers.
  • Customer-facing assistants and agentic pipelines, where uncertain requests are escalated from a cheap model to a stronger one or handed to a human, with a system-level error budget.
  • Vision-language assistants for image or document understanding, where a small VLM handles routine inputs and uncertain ones are routed to a larger VLM.
  • Cost-aware deployment, where the calibration explicitly trades the use of an expensive secondary model against the risk level chosen by the operator.

• Industry relevance: The framework gives practitioners a knob (α) they can set, a calibration procedure that outputs thresholds, and a documented fallback — declaring a risk level infeasible and abstaining — rather than silently over-promising. That makes it easier to attach reliability claims to deployed models and to justify model-routing architectures on both accuracy and cost grounds.

Future Directions

• Pushing beyond two models. The paper sketches multi-model routing in Appendix B, but all experiments use at most two models; the authors note that more models add threshold parameters and nested searches during calibration, so scaling behavior remains to be characterized empirically.

• Reporting the VQA results. ScienceQA and MM-Vet v2 with LLaVA1.5, LLaVA-NeXT and InternVL2 are described as part of the setup but their outcomes are absent from the available content, leaving open how LEC behaves on vision-language tasks and open-ended VQA scoring.

• Reducing calibration cost. The threshold search scans candidate λ values and, for the routing case, candidate pairs; more efficient search or closed-form solutions for large routing systems are a natural next step.

• Behavior in low-data regimes and under distribution shift. The guarantee rests on exchangeability between calibration and test examples, and the paper does not report performance when that assumption is stressed or when the calibration set is small relative to the number of feasible thresholds.

Target Audience

Researchers and graduate students working on uncertainty quantification, conformal prediction, and trustworthy machine learning for foundation models; practitioners who deploy LLMs or LVLMs in pipelines that must accept, escalate, or abstain with a stated error budget; and engineers designing model-routing or cascading systems who need a calibration procedure with a finite-sample guarantee rather than a heuristic confidence cutoff.

Authors’ abstract

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.