Is the acquisition worth the cost? Surrogate losses for   Consistent Two-stage Classifiers
Authors: Florence Regol
Joseph Cotnareanu
Theodore Glavas
Mark Coates
Venue: NeurIPS 2025 spotlight ()
Keywords: consistency
dynamic networks
learning to defer
Primary area: theory

Abstract: Recent years have witnessed the emergence of a spectrum of foundation models, covering a broad range of capabilities and costs. Often, we effectively use foundation models as feature generators and train classifiers that use the outputs of these models to make decisions. In this paper, we consider an increasingly relevant setting where we have two classifier stages. The first stage has access to features $x$ and has the option to make a classification decision or defer, while incurring a cost, to a second classifier that has access to features $x$ and $z$.  This is similar to the ``learning to defer'' setting, with the important difference that we train both classifiers jointly, and the second classifier has access to more information. The natural loss for this setting is an $\ell_{01c}$ loss, where a penalty is paid
for incorrect classification, as in $\ell_{01}$, but an additional penalty $c$ is paid
for consulting the second classifier. The $\ell_{01c}$ loss is unwieldy
for training. Our primary contribution in this paper is the derivation
of a hinge-based surrogate loss $\ell^c_{hinge}$ that is much more
amenable to training but also   satisfies the property that
$\ell^c_{hinge}$-consistency implies $\ell_{01c}$-consistency.

Is the acquisition worth the cost? Surrogate losses for Consistent Two-stage Classifiers
Florence Regol
Joseph Cotnareanu
Block, Toronto, Canada and McGill University, ILLS, Mila florence. florence.robertttee. robert-regol @mail.mmal......m megill.ca
McGill University, ILLS, Mila
Theodore Glavas
Mark Coates
McGill University, ILLS, Mila
McGill University, ILLS, Mila
Abstract
Recent years have witnessed the emergence of a spectrum of foundation models, covering a broad range capabilities and costs. Often, we effectively use foun- dation models as feature generators and train classifiers that the outputs of these models to make decisions. In this paper, we consider an increasingly relevant setting where we have two classifier stages. The first stage has access features x and has the option to make a classification decision or defer, while incurring a cost, second classifier that has access to features and 2. This is similar the "learning to defer" setting, with the important difference that we train both classifiers jointly, and the second classifier has access more information. The natural loss for this setting is an lnte loss, where a penalty is paid for incorrect classification, in for, but additional penalty c is paid for consulting the second classifier. The foic loss is unwieldy for training. Our primary contribution in this paper is the derivation of hinge-based surrogate loss Ehange that is much more amenable training but also satisfies the property that Chinge connsiseeyy implies fone-consis..ncy.
Introduction
With the emergence of a spectrum of foundation models, covering a broad range of capabilities and costs, we are increasingly faced with decision as to which model to use. For example, can we make a decision locally, on an edge device, or should we incur the additional communication and computational cost of sending the query to more powerful remote model? In many cases, we use pre-trained foundation model essentially as feature generator, and strive to train a classifier that uses the output of the foundation model as its input. In this setting, then face a task of training two classifiers, while simultancously learning when to defer to the more powerful model.
One approach to solve this problem is to I) train the more powerful classifier first; and 2) train the decision module with the smaller classifier afterwards (either jointly or separately). This strategy has proven successful and benefits from strong theoretical foundations [Keswani et al., 2021 Wilder al. 202T Verma et al. 2022, Mao al., 2023 2024b), but intuitively, this appears inefficient. Indeed, with this approach, both classifiers expend effort exploring regions of the input space where their predictions will ultimately not used. Because of this, important consider and formalize the problem where the classifiers and the module deciding which model to use are trained jointly.
While there has been significant body of work establishing consistent losses for the related problem of classification with learning to defer to experts or with reject [Verma et al.) 2022 Herbei and Wegkamp 2006], these losses do not cover the case where we jointly train multiple inference
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

classifiers along with the decision module. In these well-studied settings, the task is to learn one classifier and either defer an oracle (learning to reject) [Chow] 1970], an expert (learning to defer) Madras et al., 2018 or to multiple experts (learning to defer to multiple experts) [Verma et al. 2022]. However, these frameworks assume that the experts are external to the problem setting. They do not address how to train the experts alongside the base classifier.
In this work, we address the problem of jointly training two classifiers with a decision module. The two classifiers incur different costs, with the implication that the more expensive classifier offers better performance. We model this problem by introducing an additional information variable, Z, which represents the extra information available to the more powerful classifier. The decision module and the base classifier both have access to the same base input variable X. We refer this setup as the two-stage classification problem.
We provide the optimal solution to this problem, as well as surrogate loss function that is more suitable for training. The surrogate loss, which is based on the hinge loss, aligns with the standard cost-aware 0 loss formulation commonly used in classification tasks. We validate our theoretical findings on synthetic datasets and demonstrate the practical relevance of the problem by presenting results on a standard large language model (LLM) task, where two LLMs of varying sizes are used to answer multi-question math problems. Additionally, we provide a proof that the cross-entropy loss, which is sometimes used heuristically in existing literature, is not Bayes consistent with the natural 0 - ic loss, further justifying the need to explore this problem at theoretical level.
Our main contributions are follows:
1. We formulate a problem setting for learning a model that integrates two classifiers, where one has access to additional information but comes with a cost The goal is to train the models and simultaneously learn the decision function to determine whether to consult the more powerful classifier for a given sample.
2. We present surrogate loss function based on the hinge loss, which is suitable for training with cost-aware classification tasks, We show that it is consistent with respect to the 0 - Ie loss that is natural for the considered problem.
3. We validate the theoretical findings, which are the primary contribution of the work, on synthetic datasets and provide practical insights through experiments on standard LLM task.
2 Related Work
Loss consistency is an important topic that has been widely explored, as it serves as the fundamental link between the loss we optimize in practice and the actual loss we aim minimize. Foundational results have been established for classical risks [Steinwart] 2007 Tewari and Bartleu, 20077 Bartlett ct al. 2006]. and the emergence of new target losses has prompted the development of new consistency results, Learning to defer (L2D) is wide category of settings in which the task is to learn a classifier and a deferral rule, either reject (learning to abstain) [Chow] 1957 1970 Herbei and Wegkamp 2006 Cao et al., 2022, Wiener and El- Yaniy, 2011 Geilman and El- Yaniv 2019 or to defer to one
or more experts of varying costs [Madras et al. 2018 Keswani et al. 2021 Wilder et al., 2021
Mozannar and Sontag [2020] were the first to provide Bayes- consis-cn results for their proposed
generalized cross entropy loss for learning to defer, followed by Verma al. [2022], who used
a one-vs-all loss. AAaatth et al.] [2022] explored stronger guarantees than Bayes-consistency by
introducing H-consisteney bounds [Long and Servedio 2013 Mozannar al. (2023) prove that
earlier approaches, such as [Mozannar and Sontag] 2020 Verma et al. 2022] J. fall short of realizable
H-consistency, and propose a new algorithm without Bayssscon proof. This shortcoming is addressed by Mao et al. [2024a]. who recently published a unifying work, introducing new family of surrogate losses for the learning to defer problem with a single expert, and providing Bayes- consistency, realizable H--conistency, and H-consistencyy bounds. Verma et al. [2022] extended the work of |Mozannar and Sontagl [2020). Verma et al. (2022] to the multi-expert setting with Bayes-consistency guarantees. Mao et al. [202bb] introduced general cost functions and surrogate
losses, extending the results of Mozannar and Sontag (2020) with H-consistency bounds for joint training, and offering stronger guarantees than Bayes-consistence Mao et al. (2023) also provided H-consistency bounds for a slightly different setting where training the classifier and the deferral rule is done separately.

---

It appears that consistency results have been thoroughly studied in the case where only a single classifier trained. However, the setting where multiple classifiers are trained jointly with decision module has been largely neglected. This type of architecture is used in adaptive computation or dynamic networks [Han et al., 2022a], a branch of research focused on developing architectures that adaptively allocate computation. Since the main objective is to improve average inference efficiency, such dynamic network architectures have attracted significant interest in the development of scalable
LLM inference [Liu et al., 20233 Elbayad et al. 2020 Zeng et al. 2024
Xia et al. 2024 Leviathan
ct al.
2023
Chen et al. 20244].
Although there is growing interest in these types of networks, the losses used to train such models are mostly heuristic and lack strong theoretical foundations. In practice, these models typically train the classifiers separately and rely on threshold beeseedd-baeedd decisions [Han et al.] [202aa] [Schuster et al.
2022 Some theoretical research has been conducted on this separated training approach: Jitkrittum
et al. [023] explored the connection between threshold decisions and risk under principled 0-1 loss, identifying conditions under which the two coincide. However, this separate training is not guaranteed to be the best approach. In fact, the importance of jointly learning the classifiers and the decision module has been empirically demonstrated [Han et al.) 2022b Yu et al., 2022 Regol et al.
2024| Krzepkowski et al., 2024), motivating the development of joint learning approaches [Regol et al. 2024] and classifier seffffer-deeerallawwaaal -wwrre training methods [Han et al. 2022b Yu et al.. [022]. These
works lack a connection to surrogate losses and well-defined risk framework, which the gap
aim to address in this work.
3 The two-stage classification problem
We consider setting of two-stage classification where there are classifiers: and ʃz. The second classifier, /2, has access to additional information z, but it also incurs an additional cost, denoted as c. We therefore have the choice between using the prediction of the first classifier, to pay the additional cost and then use the more informed second classifier.
In practice, can be explicitly modeled as additional input signal or feature, which may come with higher access costs. For instance, in recommendation systems, different types of user data queries can vary significantly in terms of latency and infrastructure expense. A common approach is to first run a lightweight model for initial inference, and then selectively identify instances that would benefit from a more complex model with access to richer features. This tiered architecture is notably used by Youtube's recommendation system [Covington et al. 2016), for instance.
Alternatively, z can conceptually represent the augmented modeling capability of a larger model that has more parameters and/or was trained on a larger data set.
Denote by X' the feature space, Z the additional information space, and y = {L, K} the label space. We are given mnstance-labee-infrmaaioo triples {(x,, Zi, v1)) independently and identically drawn from an underlying distribution D with probability function p(X, Z, Y). We additionally introduce the decision module r X -> 0, 1, which indicates whether we are using, for the final decision, the first classifier fi if r(x) = 0 or to defer to the second classifier f2 if r(z) =1
The goal of two-stage classification is train two-stage classifier f X'x x Z + y that encompasses both the classifiers /1 : X' + f2 X' x Z - y, and the decision module r. The set H. of two-stage classifiers is therefore defined as follows:
[1((s,, r(x)= = 0 (e(a, r(x) = 1.
H = (f f(,
(1)
The loss associated with such a setting is the zero-one-exit loss lotes which can be expressed as a variant of the traditional zero-one loss for (f()),)) = (J() + v):
[1lf((z) z [1[f((z)) + v| +c
r(æ) = 0, r(z) =
(xr ().))
where 1[:] is the indicator function. The cost e can be an instance-specific function, i.e., e(x), provided it is known and deterministic. Since the additional information is only accessible at a cost e, the first classifier and the decision function do not have access it; the classifiers fi (zr) and rax) take only æ as input.
3

---

Our task is to train two-stage classifier f E H. as defined by ① that can minimize the expectation of loie over the data distribution. The risk is:
Roie Bo( (f 2), v)],
(3)
and its optimal value Roic (f))(f)) is obtained by the Bayes-optimal classifier:
arg min Rote JEH
(4)
The Ole loss discrete, and thus difficult to work with. We would like to be able identify surrogate loss la such that l cg-ooniittnccy implies fo1e-Consistency. This is our main contribution in this work. We specify surrogate loss function that satisfies this property, and show that other heuristic surrogate losses that are used in the literature for joint training (Regol et al. 2024 Ding jet al. 2024) do not. Taking a step beyond this, we specify how to construct and train a two-stage classifier using the posited surrogate loss and present empirical results to validate our result.
3.1 The solution
We start by providing the solution the optimization problem specified by 4 We first define compact notation for the posteriors:
(x(x)) = p(Y = y|æ),
(x A p(Y y/a,
(5)
Lemma 3.1. The optimal solution f arg minsew (f) is the following:
arg max, Fl (2), if max, ny 2 |maxx Ca 2)] C,
(6)
arg maxy (y(x, 2) else.
See Appendix A.4 for the proof of the lemma.
Using our previous definition of two-stage classifier, this would correspond to:
fi()) arg max (z) if (x)r"(x) = 0}
(7)
(I,, arg max Cy (I, 2) [x; ( -
(8)
"
(x) 0 if max, (x) Ep(ziz)(max, (, Z)] c, o.w.
The optimal solution is interesting. It hints towards model that is slightly different from most existing methods. Yes, the optimal decision should depend on Pmax = max 17(), but the threshold for Pmax should be based on the expected future gain: т [x [max (2x, Z)] In that setting, r(x:) should identify the set which max 7)x) 2 Ep(s|2) max maxç( 2) only, it should select the class with the largest probability according the posterior 17(zz) This result is similar to the solution for the decision rule given fixed classifiers first provided by |Jikkittum et al. (2023), which would read as max, 7/y (x) 2 max, (2, Z) c. However, the two-tiered information available to fi, r and fa provides a more practical and detailed so solution, as it allows to integrate the constraint that cannot fully access the information available to f2- This modeling choice leads to a decision based on the expected future gain.
and, for these elements
explicit modeling of
Unsurprisingly, the first and second classifiers simply predict the class with the highest probability according to their respective posteriors, but only for the samples assigned to them.
The proposed hinge-based surrogate loss
A common strategy to develop consistent loss for more complex risk functions is to propose surrogate loss and verify its consistency. This strategy was employed by early work for the learning defer problem [Mozannar and Sontag 2020 Verma and Nalisnick 2022].
Our proposed surrogate loss built on a multiclass version of the hinge loss Tarigan and van de
Geer 2008 We chose this version because it is Bayes-consistent, unlike other multiclass hinge

---

losses. We use a hinge loss rather than the more popular cross entropy is because of its linear scaling, which allows account for the cost an additive way in Eqn. 2 Following the definition of the multiclass hinge loss from [Tarigan and van de Geer) 2008), the classifiers are based on K -d-dmensional real valued outputs t(2), ((z 2) E R with the constraints that Σ t²(2) =
and Σi= Vi 2) = 0. The label prediction is obtained by returning the max element the vector. The decision function F(x) returns a real value bounded between 0 and 1. For brevity, we omit the dependence on the inputs and only write t. V. We can therefore introduce the link function that connects the real valued output and soft decision function f(x:) two-stage classifier function
f {fi. f2, r}
pt, f) {max t (z), max V 2), [f(x) 2 0.5]].
(10)
yEY
wEY
Letting [z]+ max(z, 0). our proposed hinge-baser based surrogate loss is given by:
Eivo + K-1lt + Ke V'+4
Phinge (t, v, v, = (1 - -x)))
Σ It + K- K-1 + f(x) v'
K-1
(11)
The loss is composed of a sum of two terms: the first trains the first classifier, and the second trains the second classifier. The balance or weight assigned to each term per-sample basis is intuitively controlled by the learned soft decision F. If f(z) indicates that a sample should be inferred by fi,
then fi will receive more weight during training at that point. In the second term, corresponding to
Ke fz. we include an additional fixed term K-1 that encodes the penalty of using the second classifier. This encourages f(-) to favor the first term unless the benefit of using f2 outweighs the cost. We can then define the associated risk
Rhinge (t, v) Bp I v, 2, v)l,
(12)
and consider the triplet of minimizers t" (x), (a, 2), F* (z) of such a risk:
t²
arg min t.veR rE(0,1]
R Rainge (t, v, V)
(13)
Shinge pp(tt, v",r".
(14)
In the following theorem, we establish the consistency of our proposed surrogate loss w.r.t. loie- meaning that if a learned two-stage classifier f converges to the optimal surrogate risk Riinge it also converges to the optimal target risk Role.
Theorem 4.1. There exists link function &p s.t. for any distribution p(æ, 2, v), have that:
Rainge (t, v, v) + Riingg
Roie(pp(t, v, v,)) + Riie
(15)
i.e. the surrogate loss C hinge (v, tt,, lote(sp(t, v,,),
is consistent with respect to the loss of interest
The proof is provided in Appendix A.5 The proof is built by showing that the minimizers of both risks are unique and coincide:
fhinge fr (Lemma A.1 with proof provided in Appendix A.5.I y and 2) that for some increasing function W with v (0) 0, the following holds:
(16)
(y F)) Rite S v RRhinge (t, v, v) Riingge)
(17)
(Lemma A.2 with proof included in Appendix A.5.2). consistency. We can actually establish that
Taken together, these results guarantee
2(K- - 1) (Rnänge (t, v, Bhinge)
(o v. F)) -Roc I
(18)
K
The bound on the risk gap provided by (18) allows us to further quantify the relationship between the two optimization problems, showing that the consistency not merely asymptotic. This upper bound is tight and attainable for some cases of and Ç. The K-1 term comes from the scaling of "K" the hinge loss, while the factor of accounts for corner cases where the routing decision uncertain (f(z) 0.5) and the model perfectly estimates the posteriors 1) and Ç.
5

---

To validate our findings, we present a synthetic experiment in which the ground-truth posteriors are known. We design simple K-clast classification task with oie-dimeniional inputs X and Z to enable visualization of the learned functions. Our primary interest lies in visualizing the decision boundary F(x:) of model f trained with the proposed surrogate loss. This boundary should closely 4.1 Cross entropy version
One might be tempted to build similar formulation using the widely used cross-entropy loss - -Py)) Some heuristics in the literature for training two-stage or early exit models are built around a similar version of this loss [Regol al. 2024]. Interestingly. we can prove that such a loss is in fact not Bayes consistent with the 0 - 1e loss that we presented. To build a cross entropy version of the proposed loss, we now need to assume that the model outputs predicted class probabilities p' E A for fi and p² E A K for fz. where AK is the K didimensional simplex and yo is the same link function that was previously defined. The cross entropy version of the loss that we consider adds an arbitrary function of the cost g(c) and is given by:
galc) ,pp,p'' F. 2, v)
((1 - f(z)) log(p!)) + F(x) (log(()))) + +(()))
(19)
We again consider the associated risk:
Rox Rp''p', r) = E Epl,,,,) (ga (p' pp. F,
(20)
and the minimizing function:
f& = arg min Ree (p' .P) pt FEjo.
(21)
The following lemma shows that this coss-entropy surrogate loss cannot be Bayes-consistent. Lemma 4.2. There is no function g(r)for which the solution föe to the associated problem in Eqn. is equal the Bayes-classifer f" defined in Eqn. for all distributions p(X, Z.
21
The proof is included in Appendix A.6
5 Experiments
5.1 Synthetic Experiments
approximate the optimal decision rule (x)) as defined in Eqn. 9 For completeness, we additionally include an experiment using the related cearningtto-defe baselines, which we adapt to this particular setting in Appendix A.3
Task Description The inputs X and Z are drawn uniformly from the interval [-1, 1). The label Y sampled from a categorical distribution with parameter O = =00, 00, 0z.......... 0x] I E 1]K, where
>x1 0; = 1 and p(Y = i) = 0,. The function (zx, =) is defined piecewise by partitioning the domain of x, z into K slanted regions. Full details of the construction of the synthetic dataset are provided in Appendix A.I.I The random variables are distributed as:
X ~ Uniform[-1, x(X)
(22)
Z ~ [oiformm-11,11 = p(Z)
(23)
Y Categorical (O(z, =)) = ((((x,,))
(24)
The constructed task can be visualized in Figure L where we show the class distribution in terms of most likely class and the samples x;, Vi. ~ p X: Z, Y)) for K = 5. For this example, we can see that at 0, the value of provides essentially no additional information
estimate the correct posterior, which should translate into no deferral to f2 (r (x 0) = 0). At = 0.25, the variable becomes informative. Therefore, the optimal decision function r" (ax) will alternate as vertical strips along the x-axis, with width of size that varies based on the cost parameter c.
Given this construction, the exact posterior probabilities can be computed in closed form, allowing us to derive the optimal decision rule r" (z)) To approximate the expectations Ep(z|z)) we use Monte Carlo estimation by sampling from p( / :)):
0 if max /M MM (x,
2:) 2 M EM max (y(x,
(x)
(25)
1 otherwise
where 2; 23 Uniform( - 1)
(26)

---

1.0
5
5
0.5
0.5
Classes
Classes
0.0
3
0.5
-0.5
-1.0
1
-1.0
1
-1.0
-0.5
0.0
0.5
1.0
-1.0
-0.5
0.0
0.5
1.0
X
X
Figure 1: Visualization of multi class synthetic dataset with k 5. Left) Max probability labels
arg max, p(Y x|x, Z). The black shaded region indicates where the optimal decision rule is to defer to f2 (x) = 1), given a cost 0.03. Right) Samples 2, 2, y p(X, Z,Y) of the synthetic experiment.
In our experiments, we use 1000 samples to approximate the expectations.
Training details We build simple 3-layer neural networks (NN) for t, v, and F. Following the requirements for F(x), the corresponding network takes x as input and ends with sigmoid activation. For t and v. the NNs take r and (xr 2) as inputs, respectively, and output a real-valued vector of dimension K-1. Appendix [A.1.2] provides parameter size and layer details. We use the Adam optimizer with learning rate Ir 0.001, a batch size of 512 and train for 50 epochs using our surrogate loss defined Eqn. The training set size is Ntr 10, 000 and the test
is
Nte 000.
Result and discussion Figure illustrates the ground truth and predicted decision boundaries for cost values c = 0.03, 0.07 and number of classes 3, 5. We observe that the model trained with the proposed surrogate loss successfully learns the correct decision boundary across different cost values and numbers of classes K. The learned decision function f(x:) perfectly tracks with the ground truth r (I)) Additionally, although the trained model output any value in the range F(x:) E 1] due to the sigmoid activation, it learns to produce sharp values near O or 1, which is the optimal behavior.
Now if we turn to a model trained with the additive version of the cross-entropy- surrogate loss introduced in Eqn 19 using the identity function g(c) with K 5, observe in Figure that the behavior of the learned decision function F(x) differs significantly. First, we note that since consistency cannot be established for this surrogate loss, it is not possible to precisely target desired cost level in the loic loss, unlike with the hinge-based surrogate. Looking at the results, the learned decision boundaries can be observed in the tip-right plot of Figure where we see that four regions should evenly spaced and deferred f2 (regardless of c). This pattern is adequately learned in Figure 3 For instance, we see that the right-most region of x that should be deferred is slowly erased as the cost increases.
generally unstable and uneven. The correct pattern of deferral for K
Lastly, we visualize the behavior of the learned model fuinge during training in Figure We track the empirical target risk estimated from sampling Rote(f) N Eim1 lote( 21), mi) and observe that it converges to the (empirical) optimal risk Rije as expected.
5.2 Large Language Model Experiment
To illustrate a practical setting of the problem we consider, we present experiment based on large language models (LLMs). In this experiment, we use two LLMs of different sizes, which correspond to different inference costs. The additional inference cost used by the larger model corresponds to our setup. The task involves solving multi-answer math questions. The intuition behind this setting is that some test questions should be more difficult than others. Therefore, it would be desirable to efficiently dispatch simpler questions the smaller L.LM and more challenging ones to the larger LLM. This allows us to achieve strong performance at reasonable inference cost.

---

c 0.03, k
== 0.07, k=
C 0.03,
C 0.07, k=5
1.0
Groung truth
-0.6
Deferral
0.4
Predicted
0
-0.2
0.0
-1
1 -1 0
0
-1
1
X
Figure 2: Top) Ground truth decision boundary 8 (a:) with costs values 0.03, 0.07 and number of classes 3, 5. Bottom) Learned f(z) of the model that was trained with our surrogate loss. In all cases, the decision boundaries perfectly aligned, confirming our result that the model trained with the proposed surrogate loss successfully learns the optimal decision function. As the cost increases, the black region which represents points deferred f2 shrinks.
c=0.03, k=5
C=0.035, k=5
0=0,336, k=5
C=0 ==0.037, k=5
Ground truth
-0.8
0.6
Deferral
N
g(c)=0.00 0.03
9(c)=0.035
g(c)=0.036 0.066
g(c)=0.037
00.4
Predicted
-0.2
-1
X
Figure Top) Ground truth decision boundary f (x) with costs values and number of classes K 5. Bottom) Learned F(z) of a model trained with the additive cross-entropy surrogate loss with g(c) = c, for varying e values. Unlike the model trained with the hinge-based surrogate, the learned decision patterns are generally wrong and not consistent.
Task description We use the Instruction- Pre-trained models LLaMA 8B and LLaMA 70B Grattafioni et al. 2024] solve multi-ane answer math questions from the AQUA dataset [Zhong et al. 2024). The AQUA dataset is composed of multiple-choie math reasoning questions, each with choices. We frame the task as 5-class classification problem, where the model must select the correct option from a fixed set. The inputs r and z are formed by extracting the hidden-states from the final tokens of the LLM and the 70B L.M,, respectively. We use the first 1000 AQUA [Zhong datapoints from the test split our dataset, and use 80/ 10/ 10 train/val/test split.
et al. 2024
Training details We use similar architecture to the one previously presented. The model is trained for 1000 epochs using a learning rate of 0.001, batch size of 32, and early stopping with patience of 20 epochs. Additional details provided in Appendix A.1.3

---

c=0.03
c=0.05
c=0.07
Roic(
0.5
Ro
Roie
0.4
0
10
20
30
10
20
30
10
20
30
Epoch
Epoch
Epoch
Figure 4: Empirical 0 - 1c risks of the learned function trained with the surrogate loss and of the We can see that for varying cost values, function trained with the surrogate loss converges to the optimal solution.
optimal solution for
Results and discussion Although we do not have access to the ground truth decision function in this setting, we can examine the accuracy evaluated on the selected samples vs. all samples of model trained with the surrogate loss. The selected samples of fi or f2 are the samples routed to these functions by F(zr. Ideally, the two-stage classifier model f should learn to route "hard" examples to f2. and "easier" examples to S1. practice, surrogate loss can have two effects: I) Si and f2 are additionally trained their respective selected samples; and 2) fi and f2 may receive smaller gradient updates depending on the average deferral rates.
These two effects can be observed in Figure 5 In the left figure tracking f1, we see that the average accuracy slightly increases as the cost increases (and consequently the deferral rate decreases), and the inverse behavior can be seen for f2 in the right figure. fi should, in principle, be given an easier task, so we can expect its selected accuracy be higher than the average accuracy, which we observe in the left panel of Figure 5 The two values are closest when the selected samples comprise almost all the data (i.e., a deferral rate of 90%). For f2, selected accuracy closer to the average accuracy. This could suggest that training only on the samples deferred to f2 does not result in better performance-positiee because these consist of "harder" instances.
LLaMA 3 8B f
LLaMA 3 70B f2
1.0
1.0
deferral
0.8
0.8
Accuracy
all samples
proportinn
selected samples
0.4
0.4
0.2
0.2
0.0
0.0
0.01
0.1 Cost C
0.3
0.01
0.1
0.3
Cost
Figure 5: Deferral rate and average accuracy on all samples and selected samples by left) LLaMA 3 8B fi and by right) LLaMA 70B f2. The confidence intervals are computed on trials.
In addition to aggregate performance, we can also inspect which types of queries are routed to each model. Figure 6 shows examples of math questions that were consistently routed to the smaller model (f1) and the larger model (fz) across various cost settings. From the presented examples, it appears that the "easy" questions that were consistently routed the small LLM (f1) generally involve basic arithmetic or proportions. In contrast, the labeled "hard" questions that were consistently routed to the large LLM (f2) seem to require more comprehensive knowledge (such as motion or number theory). This suggests that the routing function aligns with our perceived notion of complexity and the type of reasoning required. See Appendix A.2 for the complete list of questions that were consistently routed to fi and f2-

---

Example questions routed to the small LLM (f)
Question 1: The cost of 10 of mangos is equal to the cost of 24 kg of rice. The cost of kg of flour equals
the cost of kg of ricc. The cost of each kg of flour $22. Find the total of kg mangos, kg of rice and 5 kg of flour?
Question 2: A man buys an article and sells it at a profit of 20%. If he had bought it at 20% less and sold it for Rs.75 less, he could have gained 25%. What is the cost price?
Example questions routed to the larger LLM (f2)
Question Two trains 140 m and 160 m long run the speed 60 km/hr and km/hr respectively opposite directions on parallel tracks. The time which they take to cross each other is? Question 2: If the product of two numbers is 17820 and their H.C.F. is 12, find their I.C.M.
Figure 6: Sampled questions that are consistently being routed to f2 across different costs.
6 Conclusion and Limitations
In conclusion, this work aims to solidify the theoretical foundation behind the design and use of loss functions for the increasingly relevant problem of training multiple models with different costs, while also learning which model to use. We formalized this problem using a principled 0 - 1 cost- based loss formulation and proposed a surrogate loss based on the hinge loss, showing its consistency.
Limitations A clear limitation of our work is that we only consider two models in our setup, whereas dynamic networks often require more than two. Extending our approach to the multi-stage setting would be valuable direction for future research. Appendix A.7 provides sketch of how our method can be generalized to the multi-stage setting with T. classifiers. Moreover, although the theoretical results guarantee loss consistency, the hinge loss is less commonly used in practice. While we have presented a simple proposal of cross-entropy surrogate loss and shown that it is insufficient for this setting, exploring alternative, more stable losses would be an important next step to ensure the development of practical and principled methods.
7 Social Impact
Although we believe that this theoretical paper poses minimal direct societal impact, the broader problem of cost-sensitive deferral systems may raise concerns related to faireess and access. In such systems, the model determines whether a query is "simple" and can be handled by a smaller model, or "difficult" and should be deferred to a more powerful model, which may involve higher computational cost or latency. This can introduce bias in how different users' queries are treated. For instance, if particular user group systematically submits queries that the system deems "hard", they may consistently experience greater latency, potentially leading to unfair treatment or limited access. Additionally, this introduces deferral rule itself can be biased. This could further exacerbate disparities in user experience and overall system fairesss.
potential pathways for bias to enter the system, as the
Acknowledgements
We thank Prof. Rui Pires da Silva Castro for his valuable insights and suggestions, which greatly contributed the development of our solution. This research was partially funded by the Natural Sciences and Engineering Research Council of Canada (NSERC), [reference number 260250]. Cette recherche été partiellement financée par Conseil recherches en sciences naturelles et en génie du Canada (CRSNG), [numéro de référence 260250J. Ce projet de recherche nº 324302 est rendu possible grâce au financement du Fonds de recherche du Québec.
References
Pranjal Awasthi, Anqi Mao, Mehryar Mohri, and Yutao Zhong. H-consistency bounds for surrogate loss minimizers. In Proc. Int. Conf. on Machine Learning (ICML), 2022.
10

---

Peter L Bartlett, Michael I Jordan, and Jon D McAuliffe. Convexity, classification, and risk bounds. Journal of the American Statistical Association, 101(473):138- 156, 2006.
Yuzhou Cao, Tianchi Cai, Lei Feng, Lihong Gu, Jinjie GU, Bo An, Gang Niu, and Masashi Sugiyama. Generalizing consistent multi-class classification with rejection to be compatible with arbitrary losses. In Proc. Adv. Neural Info. Process. Syst. (NeurIPS), 2022.
Yanxi Chen, Xuchen Pan, Yaliang Li, Bolin Ding, and Jingren Zhou. Ec-llm: Large-scale training and inference of early-exit large language models with 3d parallelism. In Proc. Int. Conf. Mach. Learn. (ICML), 2024.
C. Chow. On optimum recognition error and reject tradeoff. IEEE Transactions on Information Theory, 16(1):41-46, 1970.
C. K. Chow. An optimum character recognition system using decision functions. IRE Transactions on Electronic Computers, EC-6(4):247-22 1957.
Paul Covington, Jay Adams, and Emre Sargin. Deep neural networks for youtube recommendations. In Proc. ACM Conf. on Recommender Systems, page 191-198, 2016.
Dujian Ding, Ankur Mallick, Chi Wang, Robert Sim, Subhabrata Mukherjee, Vietor Rühle, Laks Lakshmanan, and Ahmed Hassan Awadallah. Hybrid LLM: Cost-efficient and quality-award query routing. In Proc. Int. Conf. Learning Representations (ICLR), 2024.
Maha Elbayad, Jiatao Gu, Edouard Grave, and Michael Auli. Depth-adaptive transformer. In Proc. Int. Conf. Learn. Representations (ICLR), 2020.
Yonatan Geifman and Ran El- Ea-ya.iv. SelectiveNet: A deep neural network with an integrated reject option. In Proc. Int. Conf. on Machine Learning (ICML), 2019.
Aaron Grattafiori, Abhimanyu Dubey, Abhinav Jauhri, Abhinav Pandey, Abhishek Kadian, Ahmad Al-Dahle, Aiesha Letman, Akhil Mathur, Alan Schelten, Alex Vaughan, Amy Yang, Angela Fan, Anirudh Goyal, Anthony Hartshorn, Aobo Yang. Archi Mitra, Archie Sravankumar, Artem Korenev, Arthur Hinsvark, Arun Rao, Aston Zhang, Aurelien Rodriguez, Austen Gregerson, Ava Spataru, Baptiste Roziere, Bethany Biron, Binh Tang, Bobbie Chern, Charlotte Caucheteux, Chaya Nayak, Chloe Bi, Chris Marra, Chris McConnell, Christian Keller, Christophe Touret, Chunyang Wu, Corinne Wong, Cristian Canton Ferrer, and Zhiyu Ma. The Llama herd of models, 2024.
Y. Han, G. Huang, Song, Yang. H. Wang, and Y. Wang. Dynamic neural networks: A survey. IEEE Trans. on Pattern Analysis; Mach. Intell., 44(11):7436-7456 2022a.
Yizeng Han, Yifan Pu, Zihang Lai, Chaofei Wang, Shiji Song, Junfen Cao, Wenhui Huang, Chao Deng, and Gao Huang. Learning to weight samples for dynamic early-exiting networks. In Proc. European Conf. on Computer Vision (ECCV), 2022b.
Radu Herbei and Marten H. Wegkamp. Classification with reject option. The Canadian Journal of Statistics La Revue Canadienne de Statistique, 34(4))709-721, 2006.
Wittawat Jitkrittum, Neha Gupta, Aditya K Menon, Harikrishna Narasimhan, Ankit Rawat, and Sanjiv Kumar. When does confidence-based cascade deferral suffice? In Proc. Adv. Neural Info. Process. Syst. (NeurIPS), 2023.
Vijay Keswani. Matthew Lease, and Krishnaram Kenthapadi. Towards unbiased and accurate deferral to multiple experts. In Proc. AAAI/ACM Conf. on AI, Ethics, and Society, 2021.
Bartlomiej Krzepkowski, Monika Michaluk, Franciszek Szarwacki, Piotr Kubaty, Jary Pomponi, Tomasz Trzciáski, Bartosz Wójeik, and Kamil Adamczewski. Joint or disjoint: Mixing training regimes for early-exit models, 2024.
Yaniv Leviathan, Matan Kalman, and Yossi Matias. Fast inference from transformers via speculative decoding. In Proc. Int. Conf. Mach. Learn. (ICML), 2023.
Zichang Liu, Jue Wang, Tri Dao, Tianyi Zhou, Binhang Yuan, Zhao Song, Anshumali Shrivastava, Ce Zhang, Yuandong Tian, Christopher Ré, and Beidi Chen. Deja vu: Contextual sparsity for efficient Ilms inference time. In Proc. Int. Conf. Mach. Learn. (ICML), 2023.
11

---

Phil Long and Rocco Servedio. Consistency versus realizable h-consistency for multiclass classifica- tion. In Proc. Int. Conf. on Machine Learning (ICML), 2013.
David Madras, Elliot Creager, Toniann Pitassi, and Richard S. Zemel. Learning adversarially fair and transferable representations. arXiv preprint: arXiv 1802.06309, 2018.
Anqi Mao, Christopher Mohri, Mehryar Mohri, and Yutao Zhong. Two-stage learning to defer with multiple experts. In Proc. Adv. Neural Info. Process. Syst. (NeurIPS), 2023.
Anqi Mao, Mehryar Mohri, and Yutao Zhong. Realizable h-consistent and bayes-consistent loss functions for learning to defer. In Proc. Adv. Neural Info. Process. Syst. (NeurIPS), 2024a.
Anqi Mao, Mehryar Mohri, and Yutao Zhong. Principled approaches for learning defer with multiple experts. In Artificial Intelligence and Image Analysis, 2024b.
Hussein Mozannar and David Sontag. Consistent estimators for learning to defer to an expert. In Proc. Int. Conf. on Machine Learning (ICML), 2020.
Hussein Mozannar, Hunter Lang, Dennis Wei, Prasanna Sattigeri, Subhro Das, and David Sontag. Who should predict? exact algorithms for learning to defer to humans. In Proc. Int. Conf. on Artificial Intelligence and Statistics (AISTAT), 2023.
Florence Regol, Joud Chataoui, and Mark Coates. Jointly-learedd exit and inference for dynamic neural network. In Proc. Int. Conf. Learn. Representations (ICLR), 2024.
Tal Schuster, Adam Fisch, Jai Gupta, Mostafa Dehghani, Dara Bahri, Vinh Q. Tran, Yi Tay, and Donald Metzler. Confident adaptive language modeling. In Proc. Adv. in Neural Inf. Proces. Syst. (NeurIPS), 2022.
Ingo Steinwart. How to compare different loss functions and their risks. Constructive Approximation, 26(2):225 287, Aug 2007.
Bernadetta Tarigan and Sara van de Geer. A moment bound for multi-hinge classifiers. Journal of Machine Learning Research, 9(71): 2008.
Ambuj Tewari and Peter L. Bartlett. On the consistency of multiclass classification methods. Journal of Machine Learning Rexearch, 8(36):1007-1022 2007.
Rajeev Verma and Eric T. Nalisnick. Calibrated learning to defer with one-vs-all classifiers. In Proc. Int. Conf. on Machine Learning (ICML), 2022.
Rajeev Verma, Daniel Barrej' on, and Eric Nalisnick. Learning to defer to multiple experts: Consistent surrogate losses, confidence calibration, and conformal ensembles. In Proc. Int. Conf. on Artificial Intelligence and Statistics (AISTAT), 2022.
Yair Wiener and Ran El-Yaniv. Agnostic selective classification. In Proc. Adv. Neural Info. Process. Syst. (NeurIPS), 2011.
Bryan Wilder, Eric Horvitz, and Ece Kamar. Learning to complement humans. In Proc. Inf. Joint Conf. on Artificial Intelligence, 2021.
Heming Xia, Zhe Yang, Qingxiu Dong, Peiyi Wang, Yongqi Li, Tao Ge, Tianyu Liu, Wenjie Li, and Zhifang Sui. Unlocking efficiency in large language model inference: A comprehensive survey of speculative decoding. arXiv preprint: arXiv 2401.07851, 2024.
Haichao Yu, Haoxiang Li, Gang Hua, Gao Huang, and Humphrey Shi. Boosted dynamic neural networks. In Proc. AAAI Conf. on Artif. Intell., 2022.
Ziqian Zeng, Yihuai Hong, Hongliang Dai, Huiping Zhuang, and Cen Chen. ConsistentEE: A consistent and hardness-guice guided early exiting method for accelerating language models inference. In Proc. AAAI Conf. Artif. Intell., pages 19506-1951 19514, Mar. 2024.
Wanjun Zhong, Ruixiang Cui, Yiduo Guo, Yaobo Liang. Shuai Lu, Yanlín Wang, Amin Saied, Weizhu Chen, and Nan Duan. AGIEval: A human-centric benchmark for evaluating foundation models. In Findings of the Association for Computational Linguistics, June 2024.
12

---

A
Appendix
NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification:
The contributions are stated in the introduction are supported by theoritical
results and experiments.
2. Limitations
Question:
Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification: We include discussion on the limitation of our work in the conclusion.
3. Theory assumptions and proofs
Question: For each theoretical result, does the paper provide the full of assumptions and complete (and correct) proof?
Answer: [Yes]
Justification: Proofs are included the Appendix.
4. Experimental result reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Answer: [Yes]
Justification: Implementation details are included in the main text and in the Appendix.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, described in supplemental material?
Answer: [Yes]
Justification: n/a
6. Experimental setting/details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: The details are included in the main text and in the Appendix.
7. Experiment statistical significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [Yes]
Justification: We report statistical significance on our real data experiment. res tables in the text.
8. Experiments compute resources
Question: For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Answer: [NA]
13

---

Justification: The required computational resources were minimal.
9. Code of ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics
Answer: [Yes]
Justification: This research inline with NeurIPS Code of Ethics.
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification: answerNA
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: [NA]
Licenses existing assets
Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [NA]
Justification: [NA]
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification: [NA]
14. Crowdsoureing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification: [NA]
15. Institutional review board (IRB) approvals or equivalent for research with human subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: [NA]
16. Declaration of LLM usage
Question: Does the paper describe the usage of LLMs if is an important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, formatting purposes and does not impact core methodology, scientific rigorousness, or originality of the research, declaration is not required.
Answer: [NA]
Justification: [NA]
14

---

A.1 Additional experimental details
A.1.1 Synthetic task
In this section, we provide additional details of the synthetic task. We model Y by a categorical
distribution with parameter 0 = [01, .2.....00 0z]² E [0, 1j satisfying Σi= 0; = and p(Y = piecewise function by partitioning the range of X into k 1 equally
04. 0 is modeled sized bins {B1, B2...,, Bg-1). Assuming the range of X is [a, b). a, E R, we define B, = a ++ (i-1)(b-o) + following scaled and shifted sigmoid:
(6)(baa Within every bin Bi, only 0, and ds+1 take on non-zero values,
R-1
R-1
X ~ Uni 1,
(27)
Z ~ Unit 1,
(28)
Y Categorical(()
(29)
0 [0, 02 ...., 0x] E [0, 1]k
(30)
B 1 + 2(i 1)
2i
k: -1
Vi E (1, 1]
(31)
h: 1
(sX + + +2)
if E B:
0, (-1x x (sX - c--1 + 2))
if > 1 and X E Bi-1
(32)
0
else
k-1
(33)
6 = - 21
(34)
(35)
Using this model, we arrive at the closed form for the posteriors:
(x) / p(y = 1|2, Z 2)p(2 =p)() [0,]
(36)
(sx + 4 + 2)
i E B;
E2() Z)
o x (sx c + Z))
if > 1 and E Bi-1
(37)
else
(888 + 4: +2) }
if E B if > 1 and E Bi-1
Sy=i(x, =)= 0;
o o x (sx + c-11 + 2))
(38)
else
and we can derive the optimal decision function r r'(x):
(0
if max; [0:] 22()) |max, 0,] - C, o.w..
r' (x)
(39)
A.1.2 Synthetic Model details
We describe the architecture of the model used in the synthetic experiment. The hidden size of all networks is 64. Each neural network is defined as:
(x) sigmoid (BatchNorm (BRee (Ox))) to θ OBL.II1BBBcaAAoom@((AAAALAACBB
(40)
(41)
(0 2) ORIlU///accaaommmRRRRRaACCCCaaad
(42)
A.1.3 LLM Model details
We describe the architecture of the model used in the LLM experiment. The hidden size of all networks is 128. We performed a grid search for the hidden size across the values {32, 64, 128, 256) and for the learning rate across the values {0.01, 0.001, 0.0001}.
15

---

Each neural network is defined as:
re (z) = sigmoid (Batch ( (O)))))))))
(43)
to (2: O
(44)
Ve( =) (
=\)))))
(45)
A.2 Complete list of deferred questions (LLM experiments)
In this section, provide comprehensive list of the questions that were consistently deferred or sent to fi in the LLM experiment.
Example questions sent to the small LLM (f1):
1. A man buys an article and sells it at profit of 20%. he had bought it at 20% less and sold it for Rs. 75 less, he could have gained 25%. What is the cost price?
2. The cost of 10 kg of mangos is equal to the cost of 24 kg of rice. The cost of 6 kg of flour equals the cost of kg of rice. The cost of each kg of flour $22. Find the total cost of kg of mangos, 3 kg of rice and 5 kg of flour?
3. The speed of a boat in upstream is 100 kmph and the speed of the boat downstream is 180 kmph. Find the speed of the boat still water and the speed of the stream.
4. A and working together could mow field in 28 days and with the help of C they could have mowed it in 21 days. How long would C take by himself?
5. Evaluate the expression:
2 + 42 62
22
6. In examination, 60% failed in Math and 40% failed in French. If 15% failed in both subjects, what percentage of students passed in both?
7. One train crosses bridge of length 340 42 seconds, and the same train crosses another bridge of length 500 m in 50 seconds. What is the approximate speed of the train in km/hr?
8. Eshan and Mary each wrote two or three poems every day over period of time. Eshan wrote 43 poems while Mary wrote 61. What is the number of days in this period?
9. Find the value of in the sequence of numbers 5, 1, 4. 8, x. if the sum of the first numbers is 30 and the average is 4.
10. Roja and Pooja start moving in opposite directions from a pole. They are moving at speeds of 7 km/hr and 3 km/hr respectively. After 4 hours, what will be the distance between them?
Example questions deferred to the larger LLM (f2):
1. If the product of two numbers 17820 and their H.C.F. 12, find their L.C.M.
2. Two passenger trains start at the same hour in the day from different stations and move towards each other at the rate of 14 kmph and 21 kmph respectively. When they meet, it is found that one train has traveled 60 km more than the other one. What is the distance between the two stations?
3. Which is the odd one: 10, 25, 54, 60, 75, 80?
4. Two trains 140 m and 160 m long run at the speed of 60 km/hr and 40 km/hr respectively in opposite directions on parallel tracks. The time which they take to cross each other is?
5. A ladder 100 feet long leaning against vertical wall. Its lower end is 60 feet from the bottom of the wall. The side of the largest cubical box that can be placed between the wall and the ladder without disturbing the ladder is (to the nearest foot)?
6. On dividing certain number by 5, 7 and 8 successively, the remainders obtained are 2, 3 and 4 respectively. When the order of division is reversed and the number is successively divided by 8, and 5, what will be the respective remainders?
7. A tour group of 25 people paid total of $670 for entrance to museum. If this price included 5% sales tax, all the tickets cost the same amount, what was the face value of each ticket price without the sales tax?
16

---

8. A rectangular floor is covered by rug except for a strip 4 meters wide along each of the four edges. If the floor is 25 meters by 20 meters, what the area of the rug in square meters?
9. In each of the following questions number series is given with one term missing. Choose the correct alternative that will continue the same pattern and fill in the blank space.
2, 7, 14, ?, 34, 47
10. In a game of 500 points there are three participants A. B, and C. A gives to B 80 points and to 101 points. Then how many points can B give C?
11. When magnified 1,000 times by electron microscope, the image of certain circular piece of tissue has a diameter of centimeters. The actual diameter of the tissue, in centimeters,
is:
12. From the given equation, find the value of z:
2x + 9x 0
13. The sum of the non-prime numbers between 50 and 60, non-inclusive, is:
14. Solve the system of equations to find the values of and d:
I. 3 - 988 = 343
(46)
II. - 72 49
(47)
15. How many minutes does Aditya take to cover distance of 400 meters, if he runs at speed of 20 km/hr?
16. An engineer designed a ball so that when it was dropped, it rose with each bounce exactly one-half as high as it had fallen. The engineer dropped the ball from 18 -eterr platform and caught it after it had traveled 53.4 meters. How many times did the ball bounce?
A.3 Additional baselines
We include additional baselines that assume different settings from ours, such as the learning-tor to-eefer- to-expert setting, which trains a classifier alongside deferral function that defers to a fixed expert, as in [Mozannar and Sontag, 2020). We also consider simple thresholding methods that only produce results for deferral, assuming fixed classifiers, including the popular confidence-basee thresholding rule investigated by Jitkrittum et al. [2023]. These methods generally rely on training cross-entropy models for f1. and obtain the associated predicted probabilities P1. To adapt these basclines to our setting, we train the second classifier fa separately using standard eross-entropy loss, and then follow the scheme of the basclines to train the first classifier fi and obtain the decision function F(z). Let acc, denote the empirical accuracy of a model evaluated validation set.
We include the following rules:
CT-c: We pretrain fi using cross-entropy. For given cost c, we set the threshold 7 using the empirical accuracy of the second classifier minus the cost ačCfz c: and define the deferral rule as:
0
if max P/ 2 T, otherwise
F(x)
(48)
Soft deferral: We pretrain f1 using cross-entropy, and sample the deferral decision from a Bernouilii distribution with maxxey P,, parameter:
f(x) Bernoulli (1 max pT)
(49)
CT: We pretrain fi using cross coosseentropy. and search for the optimal threshold 7 that yields the smallest empirical risk on a validation set.
F(x) - -
if maxyey pt 2 T,
(50)
otherwise
17

---

L2D: We use our the pretrained f2 as the expert for the method of Mozannar and Sontag 2020]. We the confidence parameter in our expert, a, to be the average accuracy of f2: a accg.
Figure 7 shows the empirical risk during training for setting with a cost = 0.3 and number of classes k 5. We observe that proposed approach converges the fastest and achieves the lowest empirical risk. In Table | we report the empirical risk Rote along with the standard deviation computed across 10 trials for varying costs. Our proposed surrogate loss consistently attaining the lowest empirical risk across all cost settings.
0.3800
Hinge (ours)
0.3775
CT
CT-c
Roic(f)
0.3750
Soft Deferral
L2D
0.3725
0.3700
0.3675
0.0
2.5
5.0
7.5
10.0
12.5
15.0
Epoch
Figure 7: Empirical risk Roie of the learned function trained with our proposed surrogate loss and other baselines for K- = 5 and e = 0.3.
Table 1: Empirical risk Ruie for different baselines and deferral costs with K 5. The mean and standard deviation computed across 10 trials. The bolded entry denotes the lowest value.
Baseline
0.03
0.05
0.07
CT-c
0.3701 + 0.0013
0.3842
0.0007
0.4031 + 0.0014
CT
0.3700 + 0.0008
0.3784 + 0.0016
0.3853 + 0.0036
Soft deferral
0.3700 + 0.0012
0.3789
0.0008
0.3855 + 0.0014
L2D
0.3712 + 0.0015
0.3826
+ 0.0022
0.3906 + 0.0025
Hinge (ours)
0.3695 + 0.0004
0.3777
0.0004
0.3842 +0.0005
A.4 Proof of the solution f
In this section, we provide the proof of Lemma 3.1
which states that the optimal solution
f" = arg min Roie(f)
(51)
JEN
is the following:
arg max, My (x), arg max, Cy(x, 2), else.
if max, 1,((x) 2 Ep(zie) (max, (x(x, Z)] c,
ft
(52)
18

---

Or alternatively:
f [x)))
(x) =0
(53)
(z 2), r(x) =1
where f1(2) = arg max (x), V.s.t..""a =0 DEY
(54)
(2 z) = arg max (1(222),
Vx....r =1
(55)
DEY
(x)x) = 1 ma (x) S ( (o(x, Z)] c
(56)
Proof. We start by evaluating the risk w.r.t to the function:
f arg maxy 7y
if max 2 (p(l)) max ((x, () -c
(57)
arg max, Cy else,
(foses and prove the result by showing that any function fo # f results in higher risk, therefore showing that f is the optimal solution.
Rf Bp(a,2.)) [fo1e (f (2,
(58)
The risk is given by:
Rfic(f*) Ep(r,,.)) looe (z, 2),
(59)
J. (f Eiooc (z, 2), u)ll,, ())
Roie(f
(60)
We can partition X in two regions based on the decision function of i.. A {x; max ()x)) 2 Ep(z(z) max ((x,=)=00 and B {æ; max I Ep(a(z) max max{(xx =xx, and split the expectation two terms:
Raie(5*) = RA + RBB
(61)
where RS J /
-Ilooi 2), w)lp(ul 2)
(62)
Looking both terms separately, starting with R where f does not use (or corresponds to fi):
Ra = /
ZEA J. Σ (onees 2), M)l(l, =)
(63)
I DEA Σ( ([f((z) # vip(elæ, y)dz p(z: d) by def of ft and lote
(64)
-f JZEA
Σ
Eglza [[((z)) * u) p(., y)d
(65)
Σ
[[f((x) # vlp(,, y)dx nothing depends on
(66)
JUEA
Σ
aarg max my (x) # y)n (x)p(x)de by def of fi and p(yl=)
(67)
9
RA
max my (z) (())((()e) (x)dz.
(68)
DEA
We obtain more straightforwardly R²:
R =
(1+c c- max (r, z)p(x, d)ddr
(69)
J.EB
Hence, the risk of f is given by;
RoIe(5e) 1 - max max +
1 + max (o 2)p(1, d2ddd.
(70)
DEA
8
JEEB
19

---

Now, we consider different sto-stage classifier f H. and f + fo. We show that any fo E H. will lead higher risk Rf I Roia(f"), therefore proving that == We further partition the space X' by splitting and where f + f and f =
A _ ( E and f = })),
A A A and I
(71)
B A ( B and f = J
B - (x: and f + ln)
(72)
We then use those new partitions to further decompose the risks of f and f in 4 terms:
Roie RA + RA + RR R
(73)
and (f R + Rdd + R R
(74)
We can then write the difference between the risks as:
Ro) Roio(5)) = 1AA RA
+ +BB
-R
(75)
Roie(f*) Roie = 8:14 + 8T
(76)
where dA" BA* RAa
(77)
8 R.
R
(78)
In the following, we prove that
84
1 and 88" S.
(79)
which would imply that
Roic(f%) - Roic(f9) 5
Vf 1
(80)
= RoIe(f") 2 Ro #
(81)
f =f.
(82)
Showing that 8-14
< and 88
00
We first consider 8-1
8-12
= R&1
(83)
max my (x)p(())d R. using 68
(84)
JIEAU
y
We can (once again) further partition the space based on fa. We divide A in two based on decision function of f
de A {a; E Ad and f = f((x22)
(85)
d {a; E Ad and fo = ff(z))
(86)
20

---

Continuing our development of 8-1:
84"
- max (z)pp(<)) R
(87)
WEAA
/
1 - 1-maxx (z)p(z)dd (=)(xx(())(2)dd
RA* +
ZEAAS 1 max (ax)g())p(x)d Ro
(88)
JuEAdr
-
(1-max 1--max By (2))((())
pheelll) do
(89)
JoEAi
(1- 1-xaxx,())pp
do using
74
ZEAdS
(90)
- JSEAUL
(1-max max (1))
Σ Coid [onc///Ro)oo000ppb::::::::ttta p(x))) p(x)de
(91)
(1-maxnyle
> ( app
(92)
UREAdS
by def. of Ant Adz
(93)
= -JSAv (1 (1- max my (x))
Σ (1[5f()) (1))))()))))(())))))) (ol [*]
(94)
/ SEEACE (1- (1-maxny(2)
E + + c)p(zlx)
)c)p(()))))())////)))))))))))))) ((n(e)))))(zdddddd p(z)dadz [**].
(95)
(96)
At this stage we can focus on one term at the time, starting with the part of A where f is not using which the integral over Adz:
ZEAUe (1- (1-max max (2)) Σ (1[Sf(s) # vy)
) p(o)dr
[+]
(97)
JaEAdz (-maxny (nax))((x) (sf (x) p(z)dr
(98)
matter what value f2) takes, max (y()) 2 fly Therefore:
(99)
y
[+] $
(100)
Going the second term [**], when fo uses We start by restating the definition of the set A:
A {a; max 77x) Eppcta) [max Cy(2,, 2)] c).
(101)
[+*]
(1- max
my 11-manny(2
(z)
115 M (1 c) ((()))))))(((( p(x)dr (102)
TEAds
M
(1 (1- (1-max maxmy (x) Ep(ata)11 - Seg(n.e) (I,, 2]] p())
(103)
< 1 (- maxx (x) + (pccta) [max (x, =)] c) pl )d
(104)
A
[**] < 0 by def. of A.
(105)
Combining both results together, we have that
8 8
[$]
** <0.
(106)
Following similar steps, we also have that
8 nd
< 0.
(107)
21

---

Final step Since we have both that 8 3 0 and 8'14 3 0, we can conclude that
Roie(5%) Rote(f (fo) =8A + 68"
(108)
Roie (f*) I Roic (f) #
(109)
f= =-=
(110)
This concludes the proof.
A.5 Proof of consistency for the multi class surrogate hinge loss
We start by restating our surrogate loss:
Chinge (t. v. V. =(1 - F(x))
Σ (x)
+ 1
K 1-1
(111)
V'Fy
Ke
+ F(x)
Σ V
(x, + K 1111+ K -
(112)
where have t. e RRK and TRK the real-vectored outputs for fi and f2 respectively with the constraints that ||v|| = 0 and ||t|l = 0, and F E [0, 1] as soft decision output.
Since our surrogate optimization provides with the triplet t, F. map these to a two-stage classifier f H using the following a link function RR x RK x [0, 1] +
arg max yEy ty (x) arg max yEY (2,
if < 0.5
f = (t, v,
E H,
(113)
O.W.
where f1x) = arg max ty (x)) yEy
(114)
2x( = arg v 2x. vEY
(115)
r(x)
[f) < 0.5].
(116)
We can then define our risk as usual:
Rhinge (t, v, Ex ( (t, v, f, 2, x. 2,
(117)
and consider the triplet of minimizers t" (z), v (z, 2), (x) of such a risk, which correspond to two-stage solution:
arg min Rninge (t, v, P). viteR FE(0,11
(118)
Shinge ppt", v P*) H.
(119)
We prove that our surrogate loss is Bayes-consistent w.r. to the lote loss by showing that 1)
fhinge:
(120)
and 2)
Rotelyp(t, v.)) - Roic s v (Rninge v, Risinge)
(121)
Taken together, those results guarantee that for any distribution p(z, v). we have that:
Rhinge (t, v, f) Riinnge
Roie(y((t, v) R
(122)
which defines Bayes-consistency
(Steinwart 2007
22

---

A.5.1 The solutions fr of Rf and fhinge coincide
Restating the definitions of the solution of the target and surrogate problems:
f arg min Ruie(f)
(123)
JEH
Thinge ptt", v (T)
(124)
arg min Rasing (t, v, f). t.vER SE(O.1]
(125)
In this section, we show that the two solutions f Shänge coincide.
Lemma A.1. For any distribution p(X,Z,Y) Y)
Jhinge
(126)
Proof. We have previously shown Appendix A.4 that the solution our targeted problem
f arg min Roic(f)
(127)
SEH
is given by:
S [f((), (=)== =0
(128)
(z) r(z) =1
where (z) arg max (7 Vrs.t. r" (z) 0 DEY
(129)
(2 = arg max (o( Va r r (z)
(130)
yEY
r" (x) =
mazy"lyy)) (x) S [ (x( 2)
(131)
Next, we show that the two-stage classifier obtained from the triplet minimizer of our surrogate loss
Shinge = stt", f*) corresponds this solution.
arg min Rhinge (t, F)
(132)
tvERR ,F((01]
mrg Ep(x, Chinge (t, V,,,,, t.VER FE(o.1)
(133)
arg min Epc.a)) -(xx) Σ
It +
1 K- 11+
(134)
f(z)
Σ J'
Vi
1 --
Ke
(135)
-1
We can push the optimization problem inside the expectation w.r.t p(x) as t(x), (xx, F(x) are all functions of (and the inner expectation term is guaranteed to be bounded):
arg min L.VERK,AE(ON
Ep 11 f(x))
Σ
Ity
K-12 11+
(136)
K-
1/ fy
1 K
Ke
f(x)
Σ デ'
(vg'
K-1
K-1
(137)
Since the loss is linear combination of two terms that respectively depend t and v, that for any F, the minimizers for t and v will always equal to the minimizer of the individual terms:
(z) = arg min Ep(y
Σ Ity
1 i
1 (for any v()()) < 1)
(138)
tERK
K
y'
Ke
(x. 2) arg min Epsinn)
+
K |
(for any t, f(xx) 2 0).
(139)
VERK
K
23

---

For the multi class hinge loss that we are considering, it is known that the minimizing functions are given by the following (Tarigan and van de Geer 2008
if y = arg max vEy (z(
(0(x)
(140)
O.W.
y = arg max JEY (o( 2)
(2)
(141)
O.W.
which gets converted into fi and f2 by the link function y (see Equ 10
Shinge, (x)) = arg max (T)
vo..1.f(x)) <1
(142)
yEY
arg max (z)
Vx f(x) < 1 by Eqn. 140
(143)
MEy
fhänge, (I) (i)
(zstt.(a)) 1 by Eqn. 54
(144)
fhinge, (xx) (x)
Vx.st.F(aa) 2 0 by Eqn.
53
(145)
Next, turn to the decision function F(z). Using i as shorthands for F(z)):
F arg min (p(a.)))) Chinge (2, 2, t², v², :)] #E[0,1]
(146)
arg min(1 F(x:)) FE[o,]]
[t.jg + K 1+1
(147)
V'
1
Kc
+ E(x))p(v.)
Elve
+ K -1+
+
(148)
K
Since it is linear combination of 1 - f(x) and f(z). it is clear that the minimizer 7 (a:) will either be O or We can therefore rewrite the optimization problem as the following:
A(z, f) (1 F(Ir) (p(la))
Σ ity v'4y
+
--
(149)
K
Ke
+ K x-111))
+ F(x
K K-1 +
Σ +lvg
(150)
(1-r((x) Σ
12 2)
Σ1 Iti +
(151)
K +
Eivi Iv K - 11 11
Kc
+ f(x)
K 1- + Ep(z]e)
Σ
(o( 2)
(152)
DEY
V'
(*(x) arg min A(x, F).
(153)
FE(0.1)
24

---

We consider the two cases for A(x, R).
A(x,
Σ yey
The
(x)
Σ
(154)
K K-1+
V'Zy
Σ
(1)
Σ
lty
+
(155)
K
ydarg max, (x)
V'vu
+ larg max, (o(x)
[ +
(156)
K- 1-
y'targ max, (x)
E
nu(x)
(K-2) 2)1
-1
+
+ [1
K --1++)
1
(157)
K
1
K 1
y/arg max, (y(x)
+ larg max, no(o)
Σ
K-1 +
1
by def of tt" Eqn 140
K - 1
y'rarg max, (g(())
(158)
(x) (x))
K
A(r,
K-1 1 max
(159)
For the second case, using similar steps:
Ke
A(x, 1) = K -
E Ep(e/2)
Σ
Sa 2)
Σ
|+
(160)
30'
K
NEY
K
K - (Ep(e
(1
11
1 HEY
max Sa(xr) | |+ c).
(161)
This allows us to write the solution as
0
if F(r = 0, < A(r = 1,x)
(x) =
(162)
1 o.w.
if स (1 maxxyy My (x)) S k- (1 maxyey (y(2))] + c)
(163)
1
o.w.
0
if max, [v(x)) 2 Bp(ata) max Cy(x, 2)]
(x))
(164)
1 o.w.
We have therefore shown that
(())=r
(165)
Since we now have that (x) is restricted to the binary values F* (*(x) = {[0, 1}, we can rewrite the optimal classifiers that we previously obtained:
fhinge, (x) = (x) Vzs.t. <
(166)
Ihinge.2( (x) = (z) Vxs.t. f(x) 2 0
(167)
as
fhinge,1 (x) (x) Ves.t.r* (x) 0,
(168)
fh fhinge,2
(x) f2(2) Vrs.t. 1.
(169)
f fhinge, (x) f1(x),
(170)
fhinge,2 (x) f2(x).
(171)
Therefore, we can see that the optimal and v leads to the same solution for the internal classifiers of f". We have therefore shown that:
Sisnge.1 (x) (i),,
(172)
Jisnge.2(3 (x),
(173)
and 7 ('(x) r- (x), S Shinge
(174)
(175)
This concludes the proof.
25

---

A.5.2 Gap of the hinge loss
Next, we aim to show that for some increasing function with (0()) = 0, we can upper bound the risk gap of our loss of interest (tte v. Rote with the risk gap of our surrogate hinge loss Rhinge (t, f)
R Kinge
Lemma A.2. For any distribution p(X,Z,Y)
Rote(e(t, v, - Roie S v (Rringe(t; v. Rhinge)
(176)
Proof. We start by developing the hinge risk Bringe (t, v)
Rhsnge (t, v, = E Ep(e,2 (Chinge (t, v. f. 2, v)]
(177)
Bp((,))
So (2,, 2) - f(x)) y'ø8
Σ
ty +
K- 1
(178)
K
vE
+ Σ (1 2)((x)
Σ
Nv +
1
+ |
K K-1 }.
Ke
(179)
K
VEY
V'Zy
Next develop the term associated the optimal hinge risk Rhange
Rhinge Ep(x,))
Σ (x(x, F1
Σ Ity K-1
(180)
yEy
+ Σ (,,
Σirg +
K 1- +
1 - 11++
Ke
(181)
K -
yEY
EEpc (x) 0)A(1, 0)] E Ep(r,)) [1 (z) 1JA(x, 1)] reusing Eqn 16] 159
(182)
K
Rhinge
K Ep(e) [1 (x) 0j1
max my (x)] yEy
(183)
K
+ Bple.)) [1F 1](1 - max (x( + c)].
(184)
yEy
Bringing both Risinge and Rasngg (f) to evaluate the gap G A Ryange (S)nge(s)-1 we can decompose the gap by sum of 4 terms that are driven by the ground truth decision cases, i.e. r" 0 or and the decision of the model f(z):
G A Rhinge (t, v, - Riginge
(185)
G Ep(x,el Σ C,(2, =) f(x))) Σ
Ity' +
(186)
K +
yEY
8'
Ke
+ Σ (1,
Σ v'zy
M +
1+ +
(187)
K
K --1
DEy
K
K Bp(ay [1] 17 (x) 0](1 max 7 (x)) Ep(r. (1)7
1][c
max Cy (1 2))]
yEy
yEy
(188)
26

---

We can define the corresponding gap to each case as follows:
s G1 a ) 0.5]
Σ
-11=))
my 2)
Σ
bt
+
1
(189)
K K-1
DEY
V'
+ Σ 2)(x2)
Σ lvy +
1
| +
Ke
K
(1
- max fl (x))
(190)
K -
K -
K
yEy
NEY
G2 1 (z) = 0]1(F(z)
0.5) Σ
(0(x22)
- -((x))
Σ
[ty
+
K 1
1
(191)
DEY
V'vy
+ Σ 2)(xx)
Σ
Ny +
1
|+ +
Ke
K - max (1-max max7 (x))
K
K -
K K-1
yEy
VEY
(192)
G3 a 1lFt = 1]1[F(z) S 0.5] Σ
no(a; 2) - (1-2))11foo) -(x)))
Σ
Ity +
1
(193)
K +
DEY
+ Σ (f(x))
Eivo
+ K- - |+ +
1
Kc
K
(c + 1
max (o( z)),
K -1
K
PEy
NEY
8'48
(194)
G4 _ (f) = ()](F(x) 2 0.5] Σ
mo =) ()(1)(())) -((z))
Σ V/VV
[ty +
1
(195)
K 1-14
BEY
1
Kc
K
Σ 2)(1)))((())
Eivo
Vo + K-1
1+ +
1
(c
max Cy (x, 2)),
K-1
K 11 HEY
VEY
V'48
(196)
and rewrite the total gap as:
G G Ept,,))) G2 + G3 + Ga).
(197)
We can obtain a similar decomposition for the risk gap of our target risk Roie(f). We recall the definition:
(f) = 1p(
(198)
= -p(e,.))
[|((x) - 0]][(( ul) 1][1f2(xx))) 2) # v) + c]
(199)
= -Epe) 1 [[((x)) = =0]Ep(/a) 1 (x) + vl + [f(())) = Ep mn) [f 2) + v) + c]
(200)
(f) = Ep(r) [|((x) = 0j(1 - (z)) + 1x)(x) = 1|Bp(=|0][1 Staln.)) +c]
(201)
and decompose the gap risk with terms based on similar cases:
F A 1/(x)) 0, (x) (0) (z)
(202)
FB A 1/r(x) = 0, (<x) = 1[(ppGin)\(3(enn))) - (z) c)
(203)
F3 A 1/((x) = rr"(()) = 0](1)- (w) - -Bptele)SS/s(eee) +c)
(204)
F4 A (f((x) = 1,r (x) = 1](Ep(ele) (rz((.)) (t((n,)))
(205)
Rf - R = Ep(,.)) Fi + F2 + F3 + F4
(206)
This makes sense: if the optimal decision is defer and we don't, the risk is diminished by the saved computation e (the second case). If the optimal decision is not to defer and do, the risk is increased by the computation cost c (the third case).
27

---

To prove the result, we show that
Ep(n.s)
Fi
S 4 p(zz,,)) (1)))
(207)
Bp(a.s)
F2
(Rppp G2))
(208)
Ep(x,=)
F3
S <V =] (Ga)
(209)
E (p(as)
F4
1 11 p(,,)) (Ga))
(210)
G inequality Starting with G1 and F1:
G - 11=1 S 0.5)
Σ
Thy (2, 2) F(x))
Σ
1
(211)
K-1
DEy
V'Fy
+ Σ ()(12
Σlvo
Vv +
K -111- - + K -
Kc
K
(1
max fly (z))
(212)
K -
DEy
NEY
0'Fy
K
K (1-1(z))(1 777) f(x)(1 Cr2 + +c] (1 (1-1)) by def. of the hinge loss
(213)
K
(1 1111 f(z) + f(z)n) F(x) - f(x){f2 + F(z)c 1 m.)
(214)
K-1
K
G 2 1[(x) S 0.5]
(n. 75 +f S52 c)).
(215)
K
We need to find mapping v such that the following holds:
Ep(a.e) Fi W VV (1))
(216)
Ep(a.s) (mnp- -m1)
<V <V (Ept,,e))
(217)
Using the simple scaling function ((xx) 2(K-1) x, we can see that the previous inequality holds:
2(K - 1)G1
Ep(x,))
2 2E (z.z))
1l((x) < 0.5) - my1 +()) S52 + c))
(218)
K
2 ((..))
[1F(x) < 0.5
- -1 - +c]]
G.
(219)
2 2E (c.ee)
[1F(x) < 0.5)(n. - //1 + (c n.) + c))
(220)
r" 0
(221)
2 ((.)))
[1F(x) s 0.5] 15- - (n.))))-
(222)
[1F(x) s 0.5 f(z))(n- m/11)
(223)
2 2R ((.ee))
0.5(1). 1 -n
(224)
(n. - m/1)] - with ((x) 2(K 1)
((Btn.s)/Gil) 2 E p(,
K 2.
(225)
28

---

G2 inequality For the next inequality with G2 and E2, again using the same function (x)))
2(K
2, the inequality holds:
R
2(K 1)G₂
E Ep(o.s)
2 2E sp)ss
[[((x) 2 0.5](n. + )n( - 1) - +
+
(226)
K
2 2B p(r,z)
[[F(x)) 2 0.5[(n. + (f(x))) f(z){f2 F(z)c))
(227)
((a.))
[(f(x)) 2 -
- +
(228)
2 B((,,
0.5(n- - Sr +c)
(229)
V ((p(r,,))\2)) > E s))
F2 =.
(230)
(231)
Following similar steps, we obtain
V (B ((ap( (Ga)) 2 Bp(a,2)
F3
(232)
v ( (Bppar)))Gaa) 2 Ep(rrz))
(233)
Putting all results together, we obtain
Roiolp((t, v.)) - Roic - 2(K - 1)
PRisnge (t, v. Rhänge)
(234)
K
Roioly( F)) Boic S v (Rhinge (t, v, Rhinge))
O.
(235)
with V (0) and is increasing. This concludes the proof.
A.6 Proof of the failure of the cross entropy version
In this section, we prove that the cross entropy version of the surrogate loss we presented cannot be Bayes-consistent.
We recall the entropy version, with p' A K for fi and p² e A K for fa, with the same link function sp. The cross entropy version of the loss that we consider is given by:
(g(c)
(p' p², F, x 2, y) = ((1 F(x))
log(p)) + f(x) (log(p?) + g(c)))
(236)
ce
where g(c) is an arbitrary function. We again consider its associated risk:
Red (p' F² Ep(r, [69(c)
((9(e) (p' p² F, x, z, 9)]
(237)
and minimizing function:
(oe Sop,, p2 P)) Ice arg min
(238)
p FE(O,
Lemma 4.2. Cross-entropy surrogate loss not Bayes Consistent. There is no g(c) for which:
foe-
(239)
Proof. Following a similar reasoning as the proof of Lemma A.1. we can readily find that the optimal predicted probability vectors p' and p² in f should match the posteriors:
J (n(x), C(x, F%e.
(240)
Now, to find the optimal decision function of the cross entropy risk Foe. we can again obtain the solution as
(x) arg min B(F, where
(241)
FE(O.1)
B x) Ep(n.a) - - f(x)) log(n, (x)) (()))))(((((((((((())))))+)++
(242)
29

---

by following the same steps that were taking to obtain Eqn. 153
We again consider the two cases:
B(0 = - log(n, (z))
(243)
(1, (p() Dog (C. (Ir, =))] g(c)-
(244)
The solution for the cross-entropy decision function is hence given by:
[
if B(0,,) B(1,a),
(ze
(x)
(245)
[
O.W.
0 if log(n. (x)) 2 E Pett) (og p(c)
(z)
(246)
1 O.w.
If we recall the solution decision of our target problem;
r r=(x) n. (2) (2tt)))//((0,22))
(247)
are searching for function (c) for which
[log(n, (a:) < [ott If we define the decision sets of x:
=r( Vx,
(248)
D1 == fo r (z) ( 1)
(249)
Do fo r" (x) 0),
(250)
we can rewrite the condition Eqn. 248 as:
lo)(n-(x)) Ep(a (o, 2)]] < g(c)
for all D1,
(251)
(x)( [m()))) (og((S 2)]] 2 g(c) for all x E Do-
(252)
This implics that the function g(c) should satisfy:
sup log(n, (x)) - Ep(x(=) [og((S (x)] < g(c) x inf log(n.()) - Ep(e(e) (I, S.
rED;
=ED.
(253)
However, our condition for r" (x) is on the absolute scale, not on the log scale r r'(x) = 17 (x))
<
(st 2)] + c. We will therefore have that for some n,5
sup log(n=()) Ep(ala) [og((G,(, =))] > inf log(n(())) - Ep(la)) ( (x) =ED,
(254)
ZED
This implies that there is no g(c) that can satisfy
1 [log(n (x)) (ogg) s(c) r*))..
(255)
This concludes the proof.
A.7 Extension to the multi-classifier setting
In this section, we provide high-level view of how could extend results to the more general case of L + 1 classifiers. In this setting, there are still two stages, but there are L classifiers to choose from the second stage (with associated costs CI,...Cr). We would need introduce additional random variables multiple classifiers and costs:
and the loss in Eqn would need to be generalized to
[[f((z) + vl 1[f2(1,21) # +
r(x) 0
r(x)
lf (x.))).))))
(256)
(2, # vl +
r(x) L.
30

---

The corresponding solution for the decision boundaries (Eqn. 2 would become more complex. Instead of comparing the maximum posterior probability 77()) single expected future gain minus cost, the comparison would now need to be made against the best potential expected future gain for each classifier:
if maxy my (2))) > maxiel,I (Bp(z) [max, p(y12, 2i)1 ci)]
if maxiel,L (Ep22112) maxx, plyl x 2:)] - ci)
(*())=
(257)
if maxiel.L (pp (max, plyla, 21) - a)
The term m (pp()))))) [max, p(u1æ, Zi) a) returns the index of the best model that can be used at the second stage.
Then, we could propose multi-classifier surrogate loss (replacing Eqn. IT by using soft decision function that is now multiclass probability vector f(a:) [0. 1]_+1 with Ei Fr(z) 1. We would also need to introduce I. new learnable vectors (0)) ..... t(L), and index-dependen cost term 1(, ci,..., CL, K), where g(1, ,1,...,,c K" is some linear function of the costs that would need to be derived and obtained from the proof. The hinge surrogate loss could (potentially) have the following form:
Chinge (t F,
( 1 E
R-
g(4, C1, Ct,
K)). (258)
8'
remains to be seen whether we can verify the consistency of a loss of this form.
31


Reviews:
[{"review_id": "nbldxKlG8u", "reviewer": "Reviewer_uhgW", "summary": "This paper explores a two-stage classification framework in which the first classifier, operating on a basic feature set, can either issue a prediction or defer the decision to a second, more informed classifier that has access to additional features, at an associated cost. To enable effective joint training of both classifiers, the authors propose a hinge-based surrogate loss function that approximates the original cost-sensitive objective while remaining amenable to optimization.", "questions": "1. The loss function presented in Equation (11) could benefit from a clearer explanation. Adding more intuition or context around its components would help improve the clarity and readability of the paper.\n2. The paper introduces a cost c for deferring to the second-stage classifier. Could the authors clarify whether this represents computational inference cost, monetary cost, or some other type of penalty?\n3. While the paper references related methods, such as learning to defer and learning to abstain, it does not include empirical comparisons to these approaches. Including such comparisons would significantly strengthen the evaluation and highlight the advantages or trade-offs of the proposed method.", "limitations": "The authors acknowledge the main limitation that the work is restricted to a two stage setting and suggest extending to multi stage systems as a future direction. While the societal impact is not explicitly discussed, I believe the work does not pose immediate risks and focuses on improving efficiency and consistency in model deployment. Adding a brief discussion on how cost sensitive deferral strategies might affect fairness or access (e.g., systematically deferring difficult cases) would strengthen this section.", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "1. The paper discusses a practically important of jointly training a deferral mechanism between classifiers with different costs and capabilities. It introduces a hinge-based surrogate loss and rigorously proves its Bayes consistency with respect to the cost-aware 0-1 loss.\n2. Unlike previous works on learning to defer or reject, this approach handles joint training of multiple classifiers, rather than assuming access to fixed experts. The proof that no cross-entropy variant can be consistent for this loss is a useful result.\n3. The paper includes both synthetic and realistic LLM-based experiments that demonstrate the practical relevance of the proposed approach.\n4.  The framework and analysis are restricted to two classifiers; many real-world adaptive systems involve more than two stages or models.\n5.  The synthetic experiment is clean but limited in scope, and the LLM experiment has no ground-truth optimal decision for deeper evaluation.\n6. There’s a lack of comprehensive comparisons to other recent surrogate-based or dynamic routing methods (e.g., soft deferral, confidence thresholding).", "was_revised": true, "final_justification": "The authors addressed the main concerns raised in the initial review. They added relevant baselines, clarified the proposed loss function, and acknowledged fairness implications and multi-stage extensions. While evaluation on LLMs lacks ground truth, this is a known challenge and does not significantly weaken the contribution.", "rebuttal": "Thank you for your positive review and for highlighting the strong theoretical foundation of our work. We found your feedback valuable and have addressed the points you raised in our response.\n\n## 1. The framework and analysis are restricted to two classifiers; many real-world adaptive systems involve more than two stages or models.\n\n We acknowledge that extending our results to more than two classifiers is a very interesting and worthwhile direction. While developing a full set of new theoretical results would require substantial work and is better suited for a follow-up paper, we will include in the appendix a more thorough discussion and concrete pointers on how one might extend our theoretical to the multi-stage classification setting. (See our response to R2.5 for more details.)\n\n## 2. The synthetic experiment is clean but limited in scope, and the LLM experiment has no ground-truth optimal decision for deeper evaluation. \n\nUnfortunately, for any experiment involving real data, it is not possible to derive the ground-truth optimal decision boundary. To approximate this optimal decision boundary within our framework, we would need access to the true posterior distributions $p(Y|x), p(Y|x,z)$. However, these quantities are not available unless we assume that our learned model is perfectly calibrated. If we do make this assumption, the evaluation becomes circular (effectively assuming the model is correct in order to evaluate that same model).\nAs a result, the accuracy and deferral rates we report are the only reliable evaluation metrics in this setting.\n\n\n## 3. There’s a lack of comprehensive comparisons to other recent surrogate-based or dynamic routing methods (e.g., soft deferral, confidence thresholding).\n\nConfidence thresholding (CT) and soft deferral are both methods designed for the **fixed classifier** setting (not the joint learning setting we target). However, we agree that including them in the experiments would strengthen our presentation. We will therefore include both a confidence-based thresholding method and a soft deferral method as baselines.\n\nWe describe below the additional baselines that will be added to the paper.\nLet $acc(f)$ denote the empirical accuracy of a model $f$ evaluated on a validation set.\n\n- **CT-c:** We pretrain both $f_1$ and $f_2$ using cross-entropy (CE). For a given $c$, we set the threshold $\\tau = acc(f_2) - c$, so the deferral rule is defined as $r(x) = 1[\\max_y \\eta_y(x) < \\tau]$.\n\n- **CT:** We pretrain both $f_1$ and $f_2$ using cross-entropy (CE), and search for the optimal $\\tau$ using a validation set.\n\n- **Soft deferral:** We pretrain both $f_1$ and $f_2$ using cross-entropy (CE). The deferral rule is sampled per example using $p = 1 - \\max_y \\eta_y(x)$ as the Bernoulli parameter:\n  $r(x) \\sim \\text{Bernoulli}(p = 1 - \\max_y \\eta_y(x))$.\n  \n - **L2D:** We pretrain $f_2$ using cross-entropy (CE), and set it as an expert for the method [Mozannar and Sontag, 2020]. \n\n---\n\nWe present preliminary results on the synthetic experiment by reporting the empirical risk $\\hat{R}_{01c}(f)$. Since the rebuttal format limits us to tabular display, we report only the empirical risk of the four baselines. These results will be added to Figure 4, and Figure 6 will show corresponding results for the LLM experiment.\n\n#### Table 1. $\\hat{R}_{01c}(f)$ with $K=5$\n\n| Baseline      | c=0.03     | c=0.05     | c=0.07     |\n|---------------|------------|------------|------------|\n| CT-c          | 0.3701     | 0.3842     | 0.4031     |\n| CT            | 0.3700     | 0.3784     | 0.3853     |\n| Soft deferral | 0.3700     | 0.3789     | 0.3855     |\n| L2D |  0.3700  |  0.3788   |   0.3857   |\n| Hinge (ours)         | **0.3695** | **0.3777** | **0.3842** |\n\n---\n\nAs shown in the table, across different values of $c$, the proposed surrogate loss consistently achieves the lowest empirical risk. Notably, it even outperforms the confidence-threshold baseline, which uses a validation set to tune its threshold. Although the difference is small, the goal is to supplement and illustrate the theoretical results.\n\n\n## 4. The loss function presented in Equation (11) could benefit \nfrom a clearer explanation. Adding more intuition or context around its components would help improve the clarity and readability of the paper.\n\n  Thank you for that suggestion. We will add the following contextualization text in line 175 :  `The loss is composed of a sum of two terms: the first trains the first classifier, and the second trains the second classifier. The balance or weight assigned to each term on a per-sample basis is intuitively controlled by the learned soft decision $\\tilde{r}$. If $\\tilde{r}(x)$ indicates that a sample should be inferred by $f_1$, then $f_1$ will receive more weight during training at that point. In the second term, corresponding to $f_2$, we include an additional fixed term $\\frac{Kc}{K-1}$ that encodes the penalty of using the second classifier. This encourages $\\tilde{r}(\\cdot)$ to favor the first term unless the benefit of using $f_2$ outweighs the cost. `\n\n## 5. The paper introduces a cost c for deferring to the second-stage classifier. Could the authors clarify whether this represents computational inference cost, monetary cost, or some other type of penalty?\n\nWe don't assume any particular form of cost, as it can represent any of the examples you listed. This cost is highly task-dependent and should be determined by the practitioner. This is illustrated in the example provided on line 111: `For instance, in recommendation systems, different types of user data queries can vary significantly in terms of latency and infrastructure expense.`\nUltimately, the user should decide how much an error ($f(x) \\neq y$) costs to the system compared to the cost of using a more powerful architecture (or acquiring/using more information), as modeled by $f_2$. As an example, consider a fraud detection system. The first classifier can respond in 10ms. The second classifier requires 150 ms (e.g. round-trip time to a more powerful remote classifier). A practitioner may decide that this additional 140ms latency penalty is only worthwhile for a given transaction if the expected improvement in accuracy exceeds 10 percent, motivating a setting of c=0.1.\n\n## 6. While the paper references related methods, such as learning to defer and learning to abstain, it does not include empirical comparisons to these approaches. Including such comparisons would significantly strengthen the evaluation and highlight the advantages or trade-offs of the proposed method.\n\nThank you for this suggestion. We will include additional baselines, as suggested by other reviewers as well. In particular, we will compare against learning-to-defer methods by pretraining the $f_2$ function and then treating it as an expert. (See our response to your point 3).\n    \n ## 7. The authors acknowledge the main limitation that the work is restricted to a two stage setting and suggest extending to multi stage systems as a future direction. While the societal impact is not explicitly discussed, I believe the work does not pose immediate risks and focuses on improving efficiency and consistency in model deployment. Adding a brief discussion on how cost sensitive deferral strategies might affect fairness or access (e.g., systematically deferring difficult cases) would strengthen this section.\n\n Thank you for this suggestion, this is a fair point and we will include additional text to address the societal impact of how cost sensitive deferral strategies can affect fairness or access. In particular we will add a social impact section with the following text:\n    \n     Although this is a theoretical paper and we believe poses minimal direct societal impact, the broader problem of cost-sensitive deferral systems may raise concerns related to fairness and access. In such systems, the model determines whether a query is ``simple'' and can be handled by a smaller model, or ``difficult'' and should be deferred to a more powerful model, which may involve higher computational cost or latency. This can introduce bias in how different users' queries are treated. For instance, if a particular user or group systematically submits queries that the system deems ``hard'', they may consistently experience greater latency, potentially leading to unfair treatment or limited access. Additionally, this introduces new potential pathways for bias to enter the system, as the deferral rule itself can be biased. This could further exacerbate disparities in user experience and overall system fairness."}, {"review_id": "sJKCEy1MxL", "reviewer": "Reviewer_C48H", "summary": "The problem this paper tries to solve is as follows: we have a classifier f1 that uses input x and a classifier f2 that uses input z, however the classifier f2 incurs a cost c>0 every time it is used, how can we learn the pair (f1,f2) to minimize the misclassification loss taking into account the cost c. This translates to a loss function l_{01c} which is non-differentiable, this paper aims to create a surrogate loss that is consistent with this loss and easy to optimize. They first outline the optimal bayes solution to this problem. They then propose a hinge based surrogate loss function that is consistent with the l_{01c} loss function. Moreover, they prove that a cross-entropy like surrogate loss would not be consistent. They evaluate their new surrogate on synthetic experiments and an real experiment using features from 7b (x) and 70b (z) llama models on a multiple choice math dataset.", "questions": "- can you add baselines to the experiments? For instance simple confidence based thresholding without training the models. \n\n- Is it possible to conduct more comprehensive experiments in settings that mimic where this loss might be useful for? For instance an NLP experiment with small bert models (f1) vs llama models (f2) that are fine-tuned with Llora instead of taking the embeddings only? Other experiments are welcome.\n\n- Can you comment on how to extend to defer to multiple models each with different costs c_i ?", "limitations": "yes", "rating": 3, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 3, "significance": 2, "originality": 2, "strengths_and_weaknesses": "Clarity: The paper is easy to read and clear.\n\nQuality: The theoretical results of this paper are sound. The empirical results however fall short of providing sufficient proof that the method works better than simple baselines. The LLM experiments in Section 5.2 is quite limited: it lacks baselines, it evaluates a single setup on a single dataset and is with a small sample size. Given the limited empirical evidence, it is hard to judge whether this loss function is effective.\n\nOriginality: The loss function derived is novel. The technique to derive the loss function extends existing working on learning to defer to derive this hinge based loss. \n\nSignificance: The problem this paper is trying to solve has many real world examples. For instance, routing between different LLM models or model providers is likely to be one real world application. However, the empirical experiments do not provide enough evidence to use this loss function in practice.", "was_revised": true, "final_justification": "I have read the author's rebuttal. I remain on my original score of borderline reject. The authors note that the main contribution is theoretical one, however, as a researcher in this area I need to see  some convincing empirical evidence that this new loss function is worth it compared to baselines before I invest time implementing it for my own work. I don't think the loss function on it's own is differentiated in terms of theory compared to prior work in this area to stand on it's own without empirical validation. I think this papers fails to provide that empirical evidence, many other papers in this area do have theoretic contributions supplemented with empirical validation.", "rebuttal": "Thank you for engaging with our work and for your thoughtful feedback.\nMost of the concerns raised relate to the empirical results. While we appreciate these insights and have done our best to incorporate the suggestions where possible, we would like to emphasize that the primary contribution of this paper is theoretical. The empirical results are included mainly to provide insight into the nature of the theoretical findings, rather than to offer exhaustive experimental proof of the practical superiority of the proposed loss function. Even without extensive empirical comparisons, specifying a consistent surrogate loss function is an important step that can lay the groundwork for other researchers to design training schemes and loss adaptations or approximations that improve practical performance. There are similar works in this line of research that are primarily theoretical and also provide relatively limited empirical studies (e.g., Long and Servedio, 2013; Mao et al., 2023; Cao et al., 2022). Our paper follows a similar approach, aiming to advance theoretical understanding while including illustrative experiments to support key insights.\n\nWe believe this type of theoretical paper is as valuable as more empirical works and constitutes a meaningful contribution to the NeurIPS conference.\n<!-- Most of the criticisms below pertain to the empirical results. While we appreciate the insights, and have done our best to accommodate the suggestions, and believe that many of your points are very valuable indicators for future work, we would like to emphasize that the contribution of this paper is primarily theoretical. Empirical results are included to provide insights into the nature of the theoretical results, rather than to demonstrate thorough experimental proof that the proposed loss function is practically superior. Even without an exhaustive empirical comparison, the specification of a consistent surrogate loss function is an important step that can provide a foundation for other researchers to design training schemes and loss adaptations/approximations that yield improved practical performance. \n\nIn numerous settings in the machine learning literature, some papers make primarily theoretical contributions while others are primarily experimental. For example, in the learning to defer setting, there are multiple papers that provide theoretical results concerning consistency without including substantial experimental studies (e.g., Long and Servedio, 2013, Mao et al., 2023, and Cao et al., 2022). \nIn most of these papers, the experimental results are considerably more limited than what we report. The expectation that a single paper do both is somewhat unreasonable - the vast majority of papers at NeurIPS would be rejected according to these criteria, as very few provide both extensive novel theoretical contributions and expansive empirical investigations.  -->\n\n## 1.  The theoretical results of this paper are sound. The empirical results however fall short of providing sufficient proof that the method works better than simple baselines. The LLM experiments in Section 5.2 is quite limited: it lacks baselines, it evaluates a single setup on a single dataset and is with a small sample size. Given the limited empirical evidence, it is hard to judge whether this loss function is effective.\n\n We acknowledge that the scope of our empirical results is limited, and we will include an additional baseline as suggested (see preliminary results in our response R3.1). That said, we would like to emphasize that the main contribution of this paper is theoretical. The experiment involving LLMs is intended to demonstrate that the proposed loss can be applied in practice and leads to an intuitive separation of inputs.  We recognize that this is not a comprehensive empirical study comparing our method to a wide range of baselines.\n    It certainly is of interest to perform an experimental study of our proposed loss function, exploring practical scenarios, but we believe that it is more appropriately left for future work.\n    \nThe goal of our included experiments is to provide insight into the behavior of the learned functions $f_1, f_2, r$ when trained with our surrogate loss. One of the insights provided by the experiment we include can be seen in our results in Figure 5, where we observe a ''specialization'' of the classifiers. This behavior cannot arise for example  when using the simple baseline of confidence-based thresholding, since in that setting the models are trained independently and lack any mechanism for joint specialization/coordination. \n\n    \n## 2. Significance: The problem this paper is trying to solve has many real world examples. For instance, routing between different LLM models or model providers is likely to be one real world application. However, the empirical experiments do not provide enough evidence to use this loss function in practice.\n  See our response 1.\n    \n## 3 can you add baselines to the experiments? For instance simple confidence based thresholding without training the models.\n\nYes, thank you for this suggestion. We can add this as a baseline for the LLM experiment. We have included preliminary results of the suggested baselines on the synthetic experiment in our response to R1.3\n    \n## 4 Is it possible to conduct more comprehensive experiments in settings that mimic where this loss might be useful for? For instance an NLP experiment with small bert models (f1) vs llama models (f2) that are fine-tuned with Llora instead of taking the embeddings only? Other experiments are welcome.\n\nSee response to 2.1. The outlined experiment is indeed an excellent suggestion, but we don't think it is practical to embark on a completely new experiment of this relatively extensive nature during the rebuttal period. \n\n## 5 Can you comment on how to extend to defer to multiple models each with different costs $c_i$ ?\n\nWe believe that this result can indeed be obtained, but it would require a non-trivial amount of additional work to prove formally. It represents a valuable direction for future research, and we will add more details in the paper to outline how our results could  be extended in that direction. We will add an Appendix containing the following details:\n      \nTo extend our result to $L+1$ classifiers, where there are still two stages, but there are $L$ classifiers to choose from at the second stage, we would need to introduce additional random variables $Z_1, \\dots, Z_L$, and the loss in Eqn.2 would need to be generalized to multiple classifiers and costs:\n\n\\begin{align}\n\\ell_{01c}(f(x,z_{1:L}), y) = \\begin{cases}\n1[f_1(x)\\neq y] & r(x) = 0 \\\\\\\\\n1[f_2(x,z_1)\\neq y] + c_1 & r(x) = 1 \\\\\\\\\n\\dots, \\\\\\\\\n1[f_L(x,z_L)\\neq y] + c_L & r(x) = L. \n\\end{cases}\n\\end{align}\n\nThe corresponding solution for the decision boundaries (Eqn. 9) would become more complex. Instead of comparing the maximum posterior probability $\\eta(x)$ to a single expected future gain minus cost, the comparison would now need to be made against **the best** potential expected future gain for each classifier:\n\n\\begin{align}\nr^*(x) = \\begin{cases} 0 &\\text{ if } 1[\\max_y \\eta_y(x) > \\max_{l \\in {1,L}} (E_{p(Z_l|x)}[ \\max_y p(y|x,Z_l)] - c_l ) ] \\\\\\\\\n1 &\\text{ if } 1 = \\max_{l \\in {1,L}} (E_{p(Z_l|x)}[ \\max_y p(y|x,Z_l)] - c_l )  \\\\\\\\\n\\dots \\\\\\\\\nL &\\text{ if } L = \\max_{l \\in {1,L}} (E_{p(Z_l|x)}[ \\max_y p(y|x,Z_l)] - c_l ) \n\\end{cases}\n\\end{align}\n\nThe term $\\max_{l \\in {1,L}} \\Big(\\mathbb{E}_{p(Z_l|x)}[ \\max_y p(y|x,Z_l)] - c_l \\Big)$ returns the index of the best model that can be used at the second stage.\n\nThen, we could propose a multi-classifier surrogate loss (replacing Eqn. 11) by using a soft decision function that is now a multiclass probability vector $\\tilde{r}(x) \\in [0,1]^{L+1}$ with $\\sum^L_{l=0} \\tilde{r}^l(x) = 1$. We would also need to introduce $L$ new learnable vectors $\\mathbf{t}^{(0)} ,  \\dots, \\mathbf{t}^{(L)}$, and an index-dependent cost term $\\frac{K}{K-1}g(l, c_1, \\dots, c_L, K)$, where $g(l, c_1, \\dots, c_L, K)$ is some linear function of the costs that would need to be derived and obtained from the proof. The hinge surrogate loss could (potentially) have the following form:\n\n\\begin{align}\n\\ell_{hinge}^{c}( \\mathbf{t}^{(0)}, \\dots, \\mathbf{t}^{(L)}, \\tilde{r}, x, z_1, \\dots, z_L, y)  = \\sum^L_{l=0} \\tilde{r}^l(x) ( \\sum_{y' \\neq y} [t^{(l)}{y'} + \\frac{1}{K{-}1} ]_+ + g(l, c_1, \\dots, c_L, K) ).\n\\end{align}\nWe would then follow similar steps to those in the paper to verify its consistency.\n\nThese formulations would differ if the multiple models were used in a cascade of L+1 stages, so that we have multiple choices to make at different levels in the cascade, with decisions being made based on cumulative information. In such a setting, there would be a set of soft decision functions, with each $\\tilde{r}_i(\\cdot) \\in [0,1]$. Insights into this extension was requested by R3.3, so we provide the decision function there."}, {"review_id": "UO6CG2OU5S", "reviewer": "Reviewer_nQwJ", "summary": "This paper studies a two-stage classification setting where the first-stage model can defer to a second, more informed model at a cost. The authors introduce the $\\ell_{01c}$ loss to model the trade-off between classification errors and deferral cost, and propose a hinge-based surrogate loss, $\\ell^c_{\\rm{hinge}}$, that is more tractable for training while ensuring consistency with the original objective. The approach enables joint training of both classifiers and is well-suited for applications leveraging foundation model features.", "questions": "1. How is the deferral cost to the second-stage (larger) model approximated  in the experiments?\n\n2. Did the authors consider using smoother surrogates (e.g., softplus or  log-barrier variants) instead of the hinge loss to address optimization challenges with neural networks?\n\n3. How does the proposed method empirically compare with models trained using cross-entropy-based loss functions?\n\n4. How does the proposed methodology differ from, and compare to, the cascade framework presented in Jitkrittum et al., “When does confidence-based cascade deferral suffice?” NeurIPS 2023, which is applicable in the multi-stage setting? What advantages, if any, does the proposed method offer over this approach?", "limitations": "Yes.", "rating": 3, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 3, "significance": 2, "originality": 3, "strengths_and_weaknesses": "Strengths:\n\n1. The paper clearly defines a two-stage classification problem where a lightweight model can defer to a more costly model with access to richer features.\n\n2. The paper introduces the $\\ell_{01c}$ loss to explicitly encode trade-offs between prediction error and deferral cost, and derives a hinge-based surrogate loss ($\\ell^c_{\\rm{hinge}}$) that preserves consistency while being more tractable than the original discrete loss.\n\n3. The paper offers a nuanced extension to the standard learning-to-defer paradigm by allowing richer second-stage inputs and unfreezing both classifiers.\n\nWeaknesses:\n\n1. The hinge-based surrogate is non-smooth, which may hinder optimization with neural networks. The paper does not discuss this practical issue or whether smooth approximations (e.g., log-barrier or softplus variants) were considered. The negative results reported in Section 4.1 further underscore this limitation.\n\n2. The experiments do not include comparisons with more common baselines trained using cross-entropy loss, which would help contextualize the practical benefits of $\\ell^c_{\\rm{hinge}}$.\n\n3. The results are restricted to the two-stage setting, whereas extending the approach to the more practical and potentially impactful multi-stage setting would be of greater interest.", "was_revised": true, "final_justification": "I have carefully considered the author's rebuttal, but my assessment remains a borderline reject.\n\nMy primary reservation is the practical applicability of the proposed hinge-based surrogate loss. The core issue stems from the non-differentiability of the hinge function, a property known to introduce significant computational inefficiencies during optimization, especially within deep learning frameworks. This is a well-established challenge that poses a major barrier to real-world adoption.\n\nWhile I acknowledge the promising results in the additional experiments conducted under a simple setting, their limited scope does not sufficiently allay my concerns about generalizability and performance at scale. The fundamental question of whether this approach is viable for complex, real-world scenarios remains unanswered.\n\nTherefore, I stand by my initial score. As also noted by another reviewer, the paper requires more extensive and rigorous empirical validation to substantiate the claims of its practical utility.", "rebuttal": "Thank you for engaging with our work and for highlighting the theoretical strengths of our paper. We have integrated the suggested baselines, added further discussion on the relevant topics and points you raised, and provided pointers toward the requested extension.\n\n## 1.  The hinge-based surrogate is non-smooth, which may hinder optimization with neural networks. The paper does not discuss this practical issue or whether smooth approximations (e.g., log-barrier or softplus variants) were considered. The negative results reported in Section 4.1 further underscore this limitation.\n   \nThis is a valuable point. We will include additional discussion around the limitations of using the hinge loss, and suggest that smoother alternatives such as log-barrier or softplus [Dugas et al. (2001), Nock and Nielsen (2009), Zhang (2004)] variants could be explored as approximations if needed.\n\nWe did provide experiments showing that the proposed loss can indeed be optimized in practice with neural networks. Both the LLM-based and the synthetic experiments used neural network models trained end-to-end with our loss, following architectures similar to those commonly used in real-world applications. \n[1] C. Dugas, Y. Bengio, P. Bélisle, C. Nadeau, and R. Garcia, “Incorporating second-order functional knowledge for better option pricing,” in Advances in Neural Information Processing Systems, vol. 13, 2001, pp. 472–478.\n\n[2] R. Nock and F. Nielsen, “A real log barrier primal-dual algorithm for large-scale nonlinear programming,” Optimization Methods and Software, vol. 24, no. 5, pp. 771–799, 2009.\n\n[3] T. Zhang, “Statistical behavior and consistency of classification methods based on convex risk minimization,” Annals of Statistics, vol. 32, no. 1, pp. 56–85, 2004.\n\n\n ## 2. The experiments do not include comparisons with more common baselines trained using cross-entropy loss, which would help contextualize the practical benefits of lhinge\n\nIn Section 4.1, we did present a cross-entropy–based version that directly targets our problem formulation, and we showed both theoretically and empirically that it fails to recover the correct decision boundary.\n\nIf you are instead referring to the standard confidence-based approach (i.e., where a threshold is applied to the predicted probabilities of a separately trained model), we agree that this is a useful baseline for comparison, but it is also not consistent. We will integrate this method into our experiments as a baseline to better highlight the differences and limitations of such decoupled strategies in our setting. We've included preliminary results in our response to R1.3.\n\n \n## 3. The results are restricted to the two-stage setting, whereas extending the approach to the more practical and potentially impactful multi-stage setting would be of greater interest.\n\nOther reviewers have also requested a similar extension. We will integrate into the paper more explicit guidance on how our results could be extended, by adding an Appendix. In our response to R2.5, we've explained how the result could be extended to defer to multi-models at the second stage. For the multi-stage case, which corresponds to a cascade setting, the extension would be similar but with some modifications.\n\nThe decision function would become multiple decision functions, each making a decision at its respective stage and having access to increasingly more information. The classifiers would also have access to an increasing amount of data. This would be modeled as follows:\n\n\\begin{align}\n\\ell_{01c}(f(x,z_{1:L}), y) = \\begin{cases}\n1[f_1(x)\\neq y] & r_0(x) = 0 \\\\\\\\\n1[f_2(x,z_1)\\neq y] + c_1 & r_1(x,x_1) = 1 \\\\\\\\\n\\dots, \\\\\\\\\n1[f_L(x,z_1, \\dots, z_L)\\neq y] + c_L & r_L(x,z_1, \\dots, z_L) = L. \n\\end{cases}\n\\end{align}\n\n\nThat said, we would like to emphasize two points: 1) Even with just two models, our framework already supports meaningful and practical applications. For example, scenarios involving edge devices and servers often naturally involve two models (one lightweight and one more powerful);  and  2) our work already represents a significant step forward, as it introduces the first consistency-based loss tailored for this setting. To the best of our knowledge, no prior work has addressed this challenge in a theoretically grounded way. This mimics how the learning to defer literature developed: early works first established the consistency of deferring to a single expert (Mozannar and Sontag [2020] , Verma et al. [2022]), and only later were extensions to multiple experts presented in Verma et al. [2022].\n\n## 4. How is the deferral cost to the second-stage (larger) model approximated in the experiments?\n   \n The deferral cost is not approximated in our experiments; we present multiple experimental settings with various cost values. The reason for this choice is that the value of the cost parameter $c$ is inherently task-dependent. This cost can represent latency, energy consumption, or other computational resource metrics. It is therefore important to evaluate the method across a range of cost values to ensure that it performs consistently and does not fail in specific cost regimes.\n\nIn the LLM experiment, we report results for several values of $c$, as shown on the $x$-axis of Figure 5, including both low and high values: $c = 0.01$, $0.1$, $0.3$. \n\n  \n    \n## 5 Did the authors consider using smoother surrogates (e.g., softplus or log-barrier variants) instead of the hinge loss to address optimization challenges with neural networks?\n\n  No, as our focus was on providing a consistency surrogate loss. Using those approximations would break the consistency proof. We will include this as an approximation that could be used in practice if our hinge proposal is too troublesome, as we mentioned in our response to your first point.\n\n## 6 How does the proposed method empirically compare with models trained using cross-entropy-based loss functions?\n\n This result was provided in Section 4.1, where we explicitly showed that in a controlled environment, the cross-entropy version of the loss cannot recover the optimal boundary. If you are referring to other confidence-based methods, we have included preliminary results in our response to R1.3.\n\n## 7 How does the proposed methodology differ from, and compare to, the cascade framework presented in Jitkrittum et al., ''When does confidence-based cascade deferral suffice?'' NeurIPS 2023, which is applicable in the multi-stage setting? What advantages, if any, does the proposed method offer over this approach?\n\nIn short, the work by Jitkrittum et al. considers a fundamentally different setting and does not propose an approach for our setting.\n\n**Setting**: Jitkrittum et al. consider a cascade setting where the classifiers are **pre-trained and fixed**, and only the deferral rule is learned: \n    `we focus on a setting where the base models are pre-trained and fixed, and the goal is to train only the deferral rule.`\n    \nIn our setting, the classifiers are **not fixed** and are **jointly trained** with the deferral rule. \n\n**Approach:**  Jitkrittum et al. investigate scenarios where the confidence-based deferral rule is not optimal. To analyze this, they use post hoc rules designed to approximate the oracle deferral rule. However, these post hoc rules are not implementable in practice because they rely on access to information unavailable at decision time (such as the confidence of $f_2$). In our setting, this would correspond to having access to $z$ at no cost in order to make the decision to defer $r(x,z)$, which would render the problem nonsensical. Their results show that, indeed, in some settings, the confidence-based rule is suboptimal. \n\nTheir work can be seen as a useful justification and motivation for the problem setting we address, while also revealing that in some cases joint training is not likely to provide substantial benefits.\n\nWe directly contextualize the result from this work at line 153: \n`This result is similar to the solution for the decision rule given fixed classifiers first provided by Jitkrittum et al, which would read as $\\max_y\\eta_y(x) \\geq \\max_y\\zeta_y(x,Z) -c$. However, our explicit modeling of the two-tiered information available to $f_1, r$ and $f_2$ provides a more practical and detailed solution, as it allows us to integrate the constraint that $r$ cannot fully access the information available to $f_2$. This modeling choice leads to a decision based on the \\textbf{expected} future gain.`"}, {"review_id": "A0OFUF9cPa", "reviewer": "Reviewer_ysUj", "summary": "This paper addresses a very natural and increasingly common setting: we often have access to a cheap model that can handle most predictions, and a more expensive, better-informed model we'd rather only use when necessary. The authors formalize this as a two-stage classification setup, where the first classifier can either predict or defer to a second, stronger model at a fixed cost. They propose a new surrogate loss based on a hinge formulation that is proven to be consistent with the cost-aware 0–1–c loss. The theoretical results are backed by both synthetic experiments and an application involving routing math questions to either a small or large LLM, depending on difficulty.", "questions": "See the 'weaknesses' section.", "limitations": "Yes", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 4, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strengths\n\nThe paper is technically solid and conceptually thorough. It’s also very well written. Not only are the motivations clear and compelling, but the notations, assumptions, and formulations are all introduced with careful justification and helpful intuition.\nOne of the nicest things about the paper is that it’s quite self-aware — it anticipates natural questions (like “why not use cross-entropy?”) and addresses them directly, both with formal analysis and experiments. The related work section is especially good at situating this work within the broader literature on deferral, dynamic inference, and loss consistency.\nThe theoretical contribution is well crafted. The surrogate loss is not only consistent, but the paper goes further and gives a negative result for the commonly used cross-entropy approach, showing that it cannot be Bayes-consistent with respect to the 0–1–c loss. This contrast makes the case for the proposed surrogate even stronger.\n\nThe LLM experiment, while not deeply benchmarked, is a smart choice that makes the whole problem setting feel concrete and relevant. Routing questions between small and large models based on difficulty is a task that many practitioners are actively thinking about.\n\nWeaknesses\n\nThe visualization section could be improved. Figures 2, 3, and 4 require a fair amount of explanation to be fully understood, and they don’t clearly illustrate the practical difference between the hinge-based surrogate and the cross-entropy baseline. Visually comparing decision boundaries is tricky, and the figures don’t really speak for themselves. Since this is where the empirical case is made, it would help to make the differences more visually intuitive.\n\nRelated to that, it would be helpful to include the empirical 0–1–c risk of the cross-entropy model alongside the hinge-based one in Figure 4, even if the former is theoretically inconsistent. A quantitative comparison would help drive home the point about practical performance.\nA few small presentation issues:\nLine 122: the notation r: x → 0,1 is missing curly braces. Should be r: x → {0,1}.\nLine 126: the definition of H is missing a closing brace.\nThe phrase “0–1 c loss” (e.g., in lines 198 and 302) can be hard to parse visually — it reads like “closs.” Consider reformatting it more clearly, perhaps using hyphens or math mode if applicable.", "was_revised": true, "final_justification": "The paper is technically solid, clearly written, and presents a well-motivated contribution with strong theoretical support. While the visualization section and a few presentation details can be improved, the planned revisions directly address these concerns. Overall, the work is of good quality and significance, and I remain supportive of its acceptance.", "rebuttal": "Thank you for this thoughtful review, and for highlighting the clarity of our presentation and the contextualization of our work, as this was an important point for us. \n## 1. The visualization section could be improved. Figures 2, 3, and 4 require a fair amount of explanation to be fully understood, and they don’t clearly illustrate the practical difference between the hinge-based surrogate and the cross-entropy baseline. Visually comparing decision boundaries is tricky, and the figures don’t really speak for themselves. Since this is where the empirical case is made, it would help to make the differences more visually intuitive.\n\nYes, we agree that this could be made clearer. While we cannot provide new figures at this stage, we can describe the improvements we will make based on your comments. You are right that these points are central to our paper and should be presented as clearly and intuitively as possible.\n\nFirst, we will add a dashed line in Figure 1, labeled as the true decision boundary at a specific cost, to help readers better understand the optimal decision boundary and introduce it before Figures 2 and 3.\n\nNext, we will align and visually compare the true decision boundary, the boundary induced by our hinge loss, and the boundary produced by the cross-entropy loss. This will make it clearer that the cross-entropy loss fails to capture the optimal boundary, whereas our hinge loss closely matches it.\n\n## 2.  Related to that, it would be helpful to include the empirical 0–1–c risk of the cross-entropy model alongside the hinge-based one in Figure 4, even if the former is theoretically inconsistent. A quantitative comparison would help drive home the point about practical performance. A few small presentation issues: Line 122: the notation r: x → 0,1 is missing curly braces. Should be r: x → {0,1}. Line 126: the definition of H is missing a closing brace. The phrase `0–1 c loss` (e.g., in lines 198 and 302) can be hard to parse visually — it reads like `closs.` Consider reformatting it more clearly, perhaps using hyphens or math mode if applicable.\n   \n Thank you for the suggestions and pointed out typos, we will integrated them. We will change the $0-1$ c loss to $01c$ loss."}]

Author rebuttal:

Meta-review: This paper addresses a natural and increasingly common setting: we often have access to a cheap model that can handle most predictions, and a more expensive, better-informed model we'd rather only use when necessary. The authors formalize this as a two-stage classification setup, where the first classifier can either predict or defer to a second, stronger model at a fixed cost. They propose a new surrogate loss based on a hinge formulation that is proven to be consistent with the cost-aware 0–1–c loss. The theoretical results are backed by both synthetic experiments and an application involving routing math questions to either a small or large LLM, depending on difficulty.

The paper is theoretical in its approach but aims to address a practical problem. As such, the criticism of the experimental results and the potential hurdles in trying to implement its suggested paradigm are more relevant than in purely theoretical papers. However, overall, I think that the novelty of formulating the problem in a reality-relevant way and providing clear theoretical analysis outweighs those concerns.
Decision: Accept (spotlight)