Bayes-Consistent Losses for Adaptive Inference in Classification

Shirley Xiaoqi Liu; Florence Regol; Antonios Valkanas; Patrick Rebeschini; Mark Coates

stat.MLOct 29, 2026

Bayes-Consistent Losses for Adaptive Inference in Classification

Authors: Shirley Xiaoqi Liu, Florence Regol, Antonios Valkanas, Patrick Rebeschini, Mark Coates

Organizations: University of Oxford · McGill University

Abstract

Adaptive inference systems balance accuracy and inference cost by routing inputs to models with varying capabilities. Although jointly training models and routers is empirically more sample-efficient than training them separately, principled losses for joint training remain underexplored. We address this gap in a classification setting by proposing cost-sensitive hinge losses for jointly training classifiers with single-stage or cascade routers. Single-stage routers directly select a model per input, whereas cascade routers sequentially decide at each stage whether to predict with the current model or defer to the next stage, which uses a more accurate but more costly model. For both settings, we prove linear calibration bounds establishing Bayes consistency: minimizing the hinge risk recovers the Bayes-optimal risk under the target cost-sensitive zero-one loss. Moreover, since routers can require more labeled data to train than classifiers, we propose a novel post-training prediction-powered inference (PPI) procedure that uses unlabeled data to refine a trained single-stage router. Experiments on synthetic and real datasets demonstrate that the proposed hinge losses are effective for joint training and that PPI can correct routing errors when labeled data is scarce.

Figures & tables

Panel of Figure 1 — Figure 1: Illustration of flat (left) and cascade (right) routing models. A flat router selects a single classifier $f_{t}$ based on $X$ , while a cascade router sequentially invokes classifiers $f_{1},\dots,f_{T}$ until the first exit decision is made based on $(X,Z_{1:t})$ , or until the last classifier $f_{T}$ is reached.

Figure 2: (Top) Monte-Carlo estimate of the Bayes-optimal conditional risk and Bayes-optimal routing decisions for the flat router

r^{*}

(left) and for the first stage of the cascade router

d^{*}_{1}

(right). (Bottom) Learned routing decisions on the test set, where points are colored according to the routing probabilities learned by the model.

Figure 3: Comparison of the oracle router

r^{\dagger}_{\textrm{flat}}

, the router

r_{\textrm{flat}}

trained on labeled data, and the PPI-corrected router

r^{\textrm{PP}}_{\textrm{flat}}

leveraging unlabeled data. The fixes and errors introduced by PPI are computed with respect to

r^{\dagger}_{\textrm{flat}}

Figure 4: Classifier accuracy and routing fractions on GSM8K across three cost configurations (costs increasing from left panel to right panel).

Appendix figures & tables9 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 6: Empirical risk of the learned model trained with surrogate loss in the flat case (two left figures) and the cascade case (two right figures), with varying costs per classifier. In all cases, the empirical risk evaluated on the test set of the learned model converges to the empirical risk of the Bayes-optimal solution.

Figure 7: Classifier and routing-function accuracy during training in the flat router setting with

\kappa=0.005

. The router is slower to learn compared to the classifiers.

Panel of Figure 8 — Figure 8: Comparison of the Bayes-optimal routing function $r_{\textrm{flat}}^{*}$ , the router learned on labeled data $r_{\textrm{flat}}$ , and the PPI-corrected router $r^{\textrm{PP}}_{\textrm{flat}}$ , for increasing numbers $n=l$ of labeled training and PPI samples. The error of the oracle w.r.t. the Bayes-optimal $r^{*}_{\textrm{flat}}$ is shown in blue shaded regions, the error of the trained router w.r.t. the oracle router is shown in gray shaded regions, and the fixes and errors introduced by PPI are computed w.r.t. the oracle router and shown by the green and red shaded regions. Note that these are computed pointwise for every input $x$ , whereas the PPI correction is applied at the bin level, so the errors and fixes also reflect the effect of binning.

Panel of Figure 9 — Figure 9: Comparison of the Bayes-optimal routing function $r_{\textrm{flat}}^{*}$ , the router learned on labeled data $r_{\textrm{flat}}$ , and the PPI-corrected router $r^{\textrm{PP}}_{\textrm{flat}}$ , for increasing numbers $n$ of labeled samples, when the same labeled data is used both for training and for the PPI correction.

Figure 10: Classifier accuracy and deferral fractions of the cascade router with LLaMA classifiers on GSM8K across three cost configurations (costs increasing from left panel to right panel). The top row shows the accuracy of each classifier on all test prompts and on the prompts it answers, and the bottom row shows the fraction of prompts deferred at each classifier.

Figure 11: Classifier accuracy and routing fractions of the flat router with Qwen classifiers on CommonsenseQA across three cost configurations (costs increasing from left panel to right panel). The top row shows the accuracy of each classifier on all test prompts and on the prompts routed to it, and the bottom row shows the fraction of prompts routed to each classifier.

Method	Error Rate	Avg Cost	Test Empirical Risk
Hinge (ours)	0.1626	0.0344	0.1970
Joint CE	0.1870	0.0334	0.2204
FrugalGPT [ Chen et al., 2024a ]	0.2195	0.0351	0.2546
TREACLE [ Zhang et al., 2024 ]	0.1951	0.0328	0.2279

Appendix

Table 1: Error rate, average cost and test empirical risk of each method with Qwen classifiers on the test set of CommonsenseQA. FrugalGPT and TREACLE are run with a target cost budget of 0.03, which acts as a soft constraint, so their average cost can exceed it slightly. Our method attains the lowest test empirical risk. Bold indicates statistically significant difference between the two best models via pairwise Wilcoxon test with 10 trials (our method achieves statistical significance at

p<0.05

level).

Figure 12: Test

0

1

c

risk during training. Our method is compared with the separate-training baselines (CT, CT-c, Soft Deferral and L2D from Mozannar and Sontag 2020) and the joint-training baselines (Joint CE and Regol et al. 2025), and it reaches the lowest risk at the end of training.

Figure 13: Test

0

1

c

empirical risk on CommonsenseQA as a function of the number of classifiers

T

, jointly-trained using our surrogate hinge loss. Starting from a naive single-model setting with the Qwen-1.5B model (

T=1

), where no routing rule needs to be learned, we incrementally add more models: Qwen-32B (

T=2

), Qwen-72B (

T=3

), LLaMA-8B (

T=4

), and LLaMA-405B (

T=5

) alongside a trainable routing rule. Risk decreases monotonically from

T=1

T=5

(with diminishing returns beyond

T=3

), indicating that the proposed framework scales well with the number of models and consistently benefits from additional routing options.

Adaptive Model Routing

Adaptive model routing selects models or specialized experts per input to balance prediction quality, computational cost, latency, and resource constraints.

See all 92 papers

Realizable Bayes-Consistency for General Metric Losses
May 5, 2026Dan Tsir Cohen, Steve Hanneke, Aryeh KontorovichStatistical Learning Theory
1Ben-Gurion University of the Negev · 2Purdue University, USA.
Flexible Routing via Uncertainty Decomposition
May 8, 2026Charlotte Peale, Siddartha Devic, Parikshit Gopalan +2Uncertainty Quantification Adaptive Model Routing
1Stanford University · University of Southern California · 3Apple
Coupled Training with Privileged Information and Unlabeled Data
May 22, 2026Jiahao Shi, Omar Hagrass, Jason M. KlusowskiLearning Using Privileged Information Semi-Supervised Learning
Department of Electrical and Computer Engineering, Princeton University · Department of Operations Research and Financial Engineering, Princeton University

May 5, 2026J/K move · Enter open · S save

Realizable Bayes-Consistency for General Metric Losses

Dan Tsir Cohen, Steve Hanneke, Aryeh Kontorovich

1Ben-Gurion University of the Negev · 2Purdue University, USA.

Statistical Learning Theory

Open paper AskPDF arXiv

We study strong universal Bayes-consistency in the realizable setting for learning with general metric losses, extending classical characterizations beyond

0

1

classification (Bousquet et al., 2020; Hanneke et al., 2021) and real-valued regression (Attias et al., 2024). Given an instance space

(X,ρ)

, a label space

(Y,\ell)

with possibly unbounded loss, and a hypothesis class

H \subseteq Y^{X}

, we resolve the realizable case of an open problem presented in Tsir Cohen and Kontorovich (2022). Specifically, we find the necessary and sufficient conditions on the hypothesis class

H

under which there exists a distribution-free learning rule whose risk converges almost surely to the best-in-class risk (which is zero) for every realizable data-generating distribution. Our main contribution is this sharp characterization in terms of a combinatorial obstruction: Similarly to Attias et al. (2024), we introduce the notion of an infinite non-decreasing

(γ_k)

-Littlestone tree, where

γ_k \to \infty

. This extends the Littlestone tree structure used in Bousquet et al. (2020) to the metric loss setting.

Abstract

Adaptive Model Routing

Explore similar work

Realizable Bayes-Consistency for General Metric Losses

Flexible Routing via Uncertainty Decomposition

Coupled Training with Privileged Information and Unlabeled Data