Title: Is the acquisition worth the cost? Surrogate losses for   Consistent Two-stage Classifiers
Venue: NeurIPS 2025
Accepted: yes
Keywords: consistency

dynamic networks

learning to defer
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

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 ℓ 01 c loss, where a penalty is paid for incorrect classification, as in ℓ 01 , but an additional penalty c is paid for consulting the second classifier. The ℓ 01 c loss is unwieldy for training. Our primary contribution in this paper is the derivation of a hinge-based surrogate loss ℓ c hinge that is much more amenable to training but also satisfies the property that ℓ c hinge -consistency implies ℓ 01 c -consistency.

1 Introduction

With the emergence of a spectrum of foundation models, covering a broad range of capabilities and costs, we are increasingly faced with a 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 a more powerful remote model? In many cases, we use the pre-trained foundation model essentially as a feature generator, and strive to train a classifier that uses the output of the foundation model as its input. In this setting, we then face a task of training two classifiers, while simultaneously learning when to defer to the more powerful model.

One approach to solve this problem is to 1) 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 et al., 2021, Verma et al., 2022, Mao et 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 be used. Because of this, it is important to consider and formalize the problem where the classifiers and the module deciding which model to use are trained jointly.

While there has been a 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

classifiers along with the decision module. In these well-studied settings, the task is to learn one classifier and to either defer to an oracle (learning to reject) [Chow, 1970], to 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 to this setup as the two-stage classification problem.

We provide the optimal solution to this problem, as well as a 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 -1 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 -1 c loss, further justifying the need to explore this problem at a theoretical level.

Our main contributions are as follows:

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 c . 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.

We present a 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 -1 c loss that is natural for the considered problem.

We validate the theoretical findings, which are the primary contribution of the work, on synthetic datasets and provide practical insights through experiments on a 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 to minimize. Foundational results have been established for classical risks [Steinwart, 2007, Tewari and Bartlett, 2007, Bartlett et al., 2006], and the emergence of new target losses has prompted the development of new consistency results. Learning to defer (L2D) is a wide category of settings in which the task is to learn a classifier and a deferral rule, either to reject (learning to abstain) [Chow, 1957, 1970, Herbei and Wegkamp, 2006, Cao et al., 2022, Wiener and El-Yaniv, 2011, Geifman 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-consistency results for their proposed generalized cross entropy loss for learning to defer, followed by Verma et al. [2022], who used a one-vs-all loss. Awasthi et al. [2022] explored stronger guarantees than Bayes-consistency by introducing H -consistency bounds [Long and Servedio, 2013]. Mozannar et al. [2023] prove that earlier approaches, such as [Mozannar and Sontag, 2020, Verma et al., 2022], fall short of realizable H -consistency, and propose a new algorithm without a Bayes-consistency proof. This shortcoming is addressed by Mao et al. [2024a], who recently published a unifying work, introducing a new family of surrogate losses for the learning to defer problem with a single expert, and providing Bayesconsistency, realizable H -consistency, and H -consistency bounds. Verma et al. [2022] extended the work of Mozannar and Sontag [2020], Verma et al. [2022] to the multi-expert setting with Bayes-consistency guarantees. Mao et al. [2024b] 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-consistency. 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 is trained. However, the setting where multiple classifiers are trained jointly with a 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., 2023, Elbayad et al., 2020, Zeng et al., 2024, Xia et al., 2024, Leviathan et al., 2023, Chen et al., 2024].

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-based decisions [Han et al., 2022a, Schuster et al., 2022]. Some theoretical research has been conducted on this separated training approach: Jitkrittum et al. [2023] explored the connection between threshold decisions and risk under a 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-deferral-aware training methods [Han et al., 2022b, Yu et al., 2022]. These works lack a connection to surrogate losses and a well-defined risk framework, which is the gap we aim to address in this work.

3 The two-stage classification problem

We consider a setting of two-stage classification where there are two classifiers: f 1 and f 2 . The second classifier, f 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, or to pay the additional cost and then use the more informed second classifier.

In practice, z can be explicitly modeled as an 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 = { 1 , . . . , K } the label space. We are given instance-label-information triples { ( x i , z i , y i ) } n i =1 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 f 1 if r ( x ) = 0 or to defer to the second classifier f 2 if r ( x ) = 1 .

The goal of two-stage classification is to train a two-stage classifier f : X ×Z → Y that encompasses both the classifiers f 1 : X → Y , f 2 : X × Z → Y , and the decision module r . The set H of two-stage classifiers is therefore defined as follows:

̸

The loss associated with such a setting is the zero-one-exit loss ℓ 01 c , which can be expressed as a variant of the traditional zero-one loss ℓ 01 ( f ( · ) , y ) = [ f ( · ) = y ] :

̸

✶ where ✶ [ · ] is the indicator function. The cost c can be an instance-specific function, i.e., c ( x ) , provided it is known and deterministic. Since the additional information z is only accessible at a cost c , the first classifier and the decision function do not have access to it; the classifiers f 1 ( x ) and r ( x ) take only x as input.

̸

Our task is to train a two-stage classifier f ∈ H , as defined by (1), that can minimize the expectation of ℓ 01 c over the data distribution. The risk is:

and its optimal value R ∗ 01 c = R 01 c ( f ∗ ) is obtained by the Bayes-optimal classifier:

The 01 c loss is discrete, and thus difficult to work with. We would like to be able to identify a surrogate loss ℓ ϕ such that ℓ ϕ -consistency implies ℓ 01 c -consistency. This is our main contribution in this work. We specify a 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 et 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 to the optimization problem specified by (4). We first define a compact notation for the posteriors:

Lemma 3.1. The optimal solution f ∗ = arg min f ∈H R 01 c ( f ) is the following:

See Appendix A.4 for the proof of the lemma.

Using our previous definition of a two-stage classifier, this would correspond to:

The optimal solution is interesting. It hints towards a model that is slightly different from most existing methods. Yes, the optimal decision should depend on p max = max η ( x ) , but the threshold for p max should be set based on the expected future gain : τ < E p ( Z | x ) [max ζ ( x, Z )] -c . In that setting, r ( x ) should identify the set on which max η ( x ) ≥ E p ( z | x ) max ζ ( x, z ) -c , and, for these elements only, it should select the class with the largest probability according to the posterior η ( x ) . This result is similar to the solution for the decision rule given fixed classifiers first provided by Jitkrittum et al. [2023], which would read as max y η y ( x ) ≥ max y ζ 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 expected future gain.

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.

4 The proposed hinge-based surrogate loss

A common strategy to develop a consistent loss for more complex risk functions is to propose a surrogate loss and verify its consistency. This strategy was employed by early work for the learning to defer problem [Mozannar and Sontag, 2020, Verma and Nalisnick, 2022].

Our proposed surrogate loss is 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 to account for the cost in an additive way as in Eqn. 2. Following the definition of the multiclass hinge loss from [Tarigan and van de Geer, 2008], the classifiers are based on K -dimensional real valued outputs t ( x ) , v ( x, z ) ∈ R K , with the constraints that ∑ K i =1 t i ( x ) = 0 , and ∑ K i =1 v i ( x, z ) = 0 . The label prediction is obtained by returning the max element of the vector. The decision function ˜ r ( 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 a soft decision function ˜ r ( x ) to a two-stage classifier function :

̸

̸

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 ˜ r . If ˜ 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 Kc K -1 that encodes the penalty of using the second classifier. This encourages ˜ r ( · ) to favor the first term unless the benefit of using f 2 outweighs the cost. We can then define the associated risk as:

and consider the triplet of minimizers t ∗ ( x ) , v ∗ ( x, z ) , ˜ r ∗ ( x ) of such a risk:

In the following theorem, we establish the consistency of our proposed surrogate loss w.r.t. ℓ 01 c , meaning that if a learned two-stage classifier f converges to the optimal surrogate risk R ∗ hinge , it also converges to the optimal target risk R ∗ 01 c .

Theorem 4.1. There exists a link function φ s.t. for any distribution p ( x, z, y ) , we have that:

i.e., the surrogate loss ℓ c hinge ( v , t , r, x, z, y ) 11 is consistent with respect to the loss of interest ℓ 01 c ( φ ( t , v , r ) , x, z, y ) .

The proof is provided in Appendix A.5. The proof is built by showing that 1) the minimizers of both risks are unique and coincide:

(Lemma A.1, with proof provided in Appendix A.5.1); and 2) that for some increasing function Ψ with Ψ(0) = 0 , the following holds:

(Lemma A.2, with proof included in Appendix A.5.2). Taken together, these results guarantee consistency. We can actually establish that

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 is not merely asymptotic. This upper bound is tight and attainable for some cases of η and ζ . The K -1 K term comes from the scaling of the multi-hinge loss, while the factor of 2 accounts for corner cases where the routing decision is uncertain ( ˜ r ( x ) = 0 . 5 ) and the model perfectly estimates the posteriors η and ζ .

4.1 Cross entropy version

One might be tempted to build a similar formulation using the widely used cross-entropy loss -log( p y ) . Some heuristics in the literature for training two-stage or early exit models are built around a similar version of this loss [Regol et al., 2024]. Interestingly, we can prove that such a loss is in fact not Bayes consistent with the 0 -1 c 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 1 ∈ ∆ K for f 1 and p 2 ∈ ∆ K for f 2 , where ∆ K is the K -dimensional simplex and φ 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:

We again consider the associated risk:

and the minimizing function:

The following lemma shows that this coss-entropy surrogate loss cannot be Bayes-consistent. Lemma 4.2. There is no function g ( · ) for which the solution f ∗ ce to the associated problem in Eqn. 21 is equal to the Bayes-classifier f ∗ defined in Eqn. 6 for all distributions p ( X,Z,Y ) .

5 Experiments

5.1 Synthetic Experiments

To validate our findings, we present a synthetic experiment in which the ground-truth posteriors are known. We design a simple K -class classification task with one-dimensional inputs X and Z to enable visualization of the learned functions. Our primary interest lies in visualizing the decision boundary ˜ r ( x ) of a model f trained with the proposed surrogate loss. This boundary should closely approximate the optimal decision rule r ∗ ( x ) , as defined in Eqn. 9. For completeness, we additionally include an experiment using the related learning-to-defer 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 is sampled from a categorical distribution with parameter θ = [ θ 1 , θ 2 , . . . , θ K ] T ∈ [0 , 1] K , where ∑ K i =1 θ i = 1 and p ( Y = i ) = θ i . The function θ ( x, z ) is defined piecewise by partitioning the domain of x, z into K -1 slanted regions. Full details of the construction of the synthetic dataset are provided in Appendix A.1.1. The random variables are distributed as:

The constructed task can be visualized in Figure 1, where we show the class distribution in terms of most likely class and the samples x i , y i , z i ∼ p ( X,Z,Y ) for K = 5 . For this example, we can see that at x = 0 , the value of z provides essentially no additional information to estimate the correct posterior, which should translate into no deferral to f 2 ( r ∗ ( x = 0) = 0 ). At x = 0 . 25 , the variable z becomes informative. Therefore, the optimal decision function r ∗ ( x ) 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 ∗ ( x ) . To approximate the expectations E p ( Z | x ) , we use Monte Carlo estimation by sampling from p ( Z | x ) :

1.0

0.5

0.0

0.5

1.0

X

1.0

0.5

0.0

0.5

1.0

Z

1

2

3

4

5

Classes

1.0

0.5

0.0

0.5

1.0

X

1.0

0.5

0.0

0.5

1.0

Z

1

2

3

4

5

Classes

In our experiments, we use M = 1000 samples to approximate the expectations.

Training details We build simple 3-layer neural networks (NN) for t , v , and ˜ r . Following the requirements for ˜ r ( x ) , the corresponding network takes x as input and ends with a sigmoid activation. For t and v , the NNs take x and ( x, z ) 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 lr = 0 . 001 , a batch size of 512 and train for 50 epochs using our surrogate loss defined in Eqn. 11. The training set size is N tr = 10 , 000 and the test set size is N te = 1 , 000 .

Result and discussion Figure 2 illustrates the ground truth and predicted decision boundaries for cost values c = 0 . 03 , 0 . 07 and number of classes K = 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 ˜ r ( x ) perfectly tracks with the ground truth r ∗ ( x ) . Additionally, although the trained model can output any value in the range ˜ r ( x ) ∈ [0 , 1] due to the sigmoid activation, it learns to produce sharp values near 0 or 1, which is the optimal behavior.

Now if we turn to a model trained with the additive version of the cross-entropy-based surrogate loss introduced in Eqn 19, using the identity function g ( c ) = c with K = 5 , we observe in Figure 3 that the behavior of the learned decision function ˜ r ( x ) differs significantly. First, we note that since consistency cannot be established for this surrogate loss, it is not possible to precisely target a desired cost level in the ℓ 01 c loss, unlike with the hinge-based surrogate. Looking at the results, the learned decision boundaries are generally unstable and uneven. The correct pattern of deferral for K = 5 can be observed in the top-right plot of Figure 2, where we see that four regions should be evenly spaced out and deferred to f 2 (regardless of c ). This pattern is not 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.

Lastly, we visualize the behavior of the learned model f hinge during training in Figure 4. We track the empirical target risk estimated from sampling ˆ R 01 c ( f ) = 1 N ∑ N i =1 ℓ 01 c ( f ( x i , z i ) , y i ) and observe that it converges to the (empirical) optimal risk ˆ R ∗ 01 c as expected.

5.2 Large Language Model Experiment

To illustrate a practical setting of the problem we consider, we present an 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 c in 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 to the smaller LLM and more challenging ones to the larger LLM. This allows us to achieve strong performance at a reasonable inference cost.

-1

0

1

c = 0.03, k = 3

c = 0.07, k = 3

c = 0.03, k = 5

c = 0.07, k = 5

-1

0

1

-1

0

1

-1

0

1

-1

0

1

-1

0

1

0.0

0.2

0.4

0.6

0.8

1.0

Deferral

X

Z

Ground

truth

r

*

Predicted

r

-1

0

1

c=0.03, k=5

c=0.036, k=5

c=0.037, k=5

-1

0

1

g(c)=0.03

g(c)=0.035

g(c)=0.036

g(c)=0.037

0.2

0.4

0.6

0.8

Deferral

X

Z

Ground

truth

r

*

Predicted

r

Task description We use the Instruction-Tuned Pre-trained models LLaMA 3 8B and LLaMA 3 70B [Grattafiori et al., 2024] to solve multi-answer math questions from the AQUA dataset [Zhong et al., 2024]. The AQUA dataset is composed of multiple-choice math reasoning questions, each with 5 choices. We frame the task as a 5-class classification problem, where the model must select the correct option from a fixed set. The inputs x and z are formed by extracting the hidden-states from the final tokens of the 8B LLM and the 70B LLM, respectively. We use the first 1000 AQUA [Zhong et al., 2024] datapoints from the test split as our dataset, and use a 80 / 10 / 10 train/val/test split.

Training details We use a similar architecture to the one previously presented. The model is trained for 1000 epochs using a learning rate of 0.001, a batch size of 32, and early stopping with a patience of 20 epochs. Additional details are provided in Appendix A.1.3.

0

10

20

30

Epoch

0.4

0.5

R

01

c

(

f

)

c=0.03

0

10

20

30

Epoch

c=0.05

0

10

20

30

Epoch

c=0.07

R

01

c

(

f

)

R

*

01

c

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 the model trained with the surrogate loss. The selected samples of f 1 or f 2 are the samples routed to these functions by ˜ r ( x ) . Ideally, the two-stage classifier model f should learn to route 'hard' examples to f 2 , and 'easier' examples to f 1 . In practice, the surrogate loss can have two effects: 1) f 1 and f 2 are additionally trained on their respective selected samples; and 2) f 1 and f 2 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 f 1 , 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 f 2 in the right figure. f 1 should, in principle, be given an easier task, so we can expect its selected accuracy to 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 f 2 , the selected accuracy is closer to the average accuracy. This could suggest that training only on the samples deferred to f 2 does not result in better performance-possibly because these consist of 'harder' instances.

0.01

0.1

0.3

Cost

c

0.0

0.2

0.4

0.6

0.8

1.0

LLaMA 3 8B

f

1

0.01

0.1

0.3

Cost

c

LLaMA 3 70B

f

2

deferral rate

all samples

selected samples

0.0

0.2

0.4

0.6

0.8

1.0

Proportion

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 ( f 1 ) and the larger model ( f 2 ) across various cost settings. From the presented examples, it appears that the 'easy' questions that were consistently routed to the small LLM ( f 1 ) generally involve basic arithmetic or proportions. In contrast, the labeled 'hard' questions that were consistently routed to the large LLM ( f 2 ) 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 f 1 and f 2 .

Example questions routed to the small LLM ( f 1 )

Question 1: 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 2 kg of rice. The cost of each kg of flour is $22. Find the total cost of 4 kg of mangos, 3 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 ( f 2 )

Question 1: 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?

Question 2: If the product of two numbers is 17820 and their H.C.F. is 12, find their L.C.M.

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 a valuable direction for future research. Appendix A.7 provides a sketch of how our method can be generalized to the multi-stage setting with L 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 a 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 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.

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.

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. Ee-llm: Large-scale training and inference of early-exit large language models with 3d parallelism. In Proc. Int. Conf. Mach. Learn. (ICML) , 2024.

Chow. On optimum recognition error and reject tradeoff. IEEE Transactions on Information Theory , 16(1):41-46, 1970.

K. Chow. An optimum character recognition system using decision functions. IRE Transactions on Electronic Computers , EC-6(4):247-254, 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, Victor Rühle, Laks V. S. Lakshmanan, and Ahmed Hassan Awadallah. Hybrid LLM: Cost-efficient and quality-aware 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-Yaniv. 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 3 herd of models, 2024.

Han, G. Huang, S. Song, L. 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.

Bartłomiej Krzepkowski, Monika Michaluk, Franciszek Szarwacki, Piotr Kubaty, Jary Pomponi, Tomasz Trzci´ nski, Bartosz Wójcik, 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 llms at inference time. In Proc. Int. Conf. Mach. Learn. (ICML) , 2023.

Phil Long and Rocco Servedio. Consistency versus realizable h-consistency for multiclass classification. 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 to 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-learned exit and inference for a 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 A. van de Geer. A moment bound for multi-hinge classifiers. Journal of Machine Learning Research , 9(71):2171-2185, 2008.

Ambuj Tewari and Peter L. Bartlett. On the consistency of multiclass classification methods. Journal of Machine Learning Research , 8(36):1007-1025, 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. Int. 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-guided early exiting method for accelerating language models inference. In Proc. AAAI Conf. Artif. Intell. , pages 19506-19514, Mar. 2024.

Wanjun Zhong, Ruixiang Cui, Yiduo Guo, Yaobo Liang, Shuai Lu, Yanlin 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.

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 a 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 set of assumptions and a complete (and correct) proof?

Answer: [Yes]

Justification: Proofs are included in the Appendix.

4. Experimental result reproducibility

Question: Does the paper fully disclose all the information needed to reproduce the main experimental 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 instructions to faithfully reproduce the main experimental results, as 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, hyperparameters, 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 or tables in the text.

8. Experiments compute resources

Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?

Answer: [NA]

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 https://neurips.cc/public/EthicsGuidelines ?

Answer: [Yes]

Justification: This research is 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

Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?

Answer: [NA]

Justification: [NA]

12. Licenses for 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. Crowdsourcing 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 it 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, or formatting purposes and does not impact the core methodology, scientific rigorousness, or originality of the research, declaration is not required.

Answer: [NA]

Justification: [NA]

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 θ = [ θ 1 , θ 2 , . . . , θ k ] T ∈ [0 , 1] k satisfying ∑ k i =1 θ i = 1 and p ( Y = i ) = θ i . θ is modeled as a piecewise function by partitioning the range of X into k -1 equally sized bins { B 1 , B 2 , . . . , B k -1 } . Assuming the range of X is [ a, b ) , a, b ∈ R , we define B i = [ a + ( i -1)( b -a ) k -1 , a + ( i )( b -a ) k -1 ) . Within every bin B i , only θ i and θ i +1 take on non-zero values, following a scaled and shifted sigmoid:

(35)

Using this model, we arrive at the closed form for the posteriors:

and we can derive the optimal decision function r ∗ ( x ) :

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:

A.1.3 LLMModel 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 } .

Each neural network is defined as:

A.2 Complete list of deferred questions (LLM experiments)

In this section, we provide a comprehensive list of the questions that were consistently deferred to f 2 or sent to f 1 in the LLM experiment.

Example questions sent to the small LLM ( f 1 ):

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?

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 2 kg of rice. The cost of each kg of flour is $22. Find the total cost of 4 kg of mangos, 3 kg of rice and 5 kg of flour?

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 in still water and the speed of the stream.

A and B working together could mow a 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?

Evaluate the expression:

In an examination, 60% failed in Math and 40% failed in French. If 15% failed in both subjects, what percentage of students passed in both?

One train crosses a bridge of length 340 m in 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?

Eshan and Mary each wrote two or three poems every day over a period of time. Eshan wrote 43 poems while Mary wrote 61. What is the number of days in this period?

Find the value of x in the sequence of numbers 5, 1, 6, 0, 4, 8, x , 2 if the sum of the first 7 numbers is 30 and the average is 4.

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 ( f 2 ):

If the product of two numbers is 17820 and their H.C.F. is 12, find their L.C.M.

Two passenger trains start at the same hour in the day from two 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?

Which is the odd one: 10, 25, 45, 54, 60, 75, 80?

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?

A ladder 100 feet long is leaning against a 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)?

On dividing a 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, 7 and 5, what will be the respective remainders?

A tour group of 25 people paid a total of $670 for entrance to a museum. If this price included a 5% sales tax, and all the tickets cost the same amount, what was the face value of each ticket price without the sales tax?

Arectangular floor is covered by a rug except for a strip 4 meters wide along each of the four edges. If the floor is 25 meters by 20 meters, what is the area of the rug in square meters?

In each of the following questions a number series is given with one term missing. Choose the correct alternative that will continue the same pattern and fill in the blank space.

In a game of 500 points there are three participants A, B, and C. A gives to B 80 points and to C 101 points. Then how many points can B give to C?

When magnified 1,000 times by an electron microscope, the image of a certain circular piece of tissue has a diameter of 2 centimeters. The actual diameter of the tissue, in centimeters, is:

From the given equation, find the value of x :

The sum of the non-prime numbers between 50 and 60, non-inclusive, is:

Solve the system of equations to find the values of c and d :

How many minutes does Aditya take to cover a distance of 400 meters, if he runs at a speed of 20 km/hr?

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 an 18-meter 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-to-deferto-expert setting, which trains a classifier alongside a 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-based thresholding rule investigated by Jitkrittum et al. [2023]. These methods generally rely on training cross-entropy models for f 1 , and obtain the associated predicted probabilities p 1 . To adapt these baselines to our setting, we train the second classifier f 2 separately using a standard cross-entropy loss, and then follow the scheme of the baselines to train the first classifier f 1 and obtain the decision function ˜ r ( x ) . Let ˆ acc f denote the empirical accuracy of a model f evaluated on a validation set.

We include the following rules:

CT-c: We pretrain f 1 using cross-entropy. For a given cost c , we set the threshold τ using the empirical accuracy of the second classifier minus the cost τ = ˆ acc f 2 -c : and define the deferral rule as:

Soft deferral: We pretrain f 1 using cross-entropy, and sample the deferral decision from a Bernouilli distribution with 1 -max y ∈Y p 1 y as a parameter:

CT: We pretrain f 1 using cross-entropy, and search for the optimal threshold τ that yields the smallest empirical risk on a validation set.

L2D: We use our the pretrained f 2 as the expert for the method of Mozannar and Sontag [2020]. We set the confidence parameter in our expert, α , to be the average accuracy of f 2 : α = ˆ acc f .

Figure 7 shows the empirical risk during training for a setting with a cost of c = 0 . 3 and number of classes k = 5 . We observe that our proposed approach converges the fastest and achieves the lowest empirical risk. In Table 1, we report the empirical risk ˆ R 01 c 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.0

2.5

5.0

7.5

10.0

12.5

15.0

Epoch

0.3675

0.3700

0.3725

0.3750

0.3775

0.3800

R

01

c

(

f

)

Hinge (ours)

CT

CT-c

Soft Deferral

L2D

A.4 Proof of the solution f ∗

In this section, we provide the proof of Lemma 3.1, which states that the optimal solution

is the following:

Or alternatively:

where

R

≜

Proof. We start by evaluating the risk w.r.t to the function:

̸

and prove the result by showing that any function f o = f s results in a higher risk, therefore showing that f s is the optimal solution.

The risk is given by:

We can partition X in two regions based on the decision function of f s , i.e. : A = { x ; max η ( x ) ≥ E p ( z | x ) max ζ ( x, z ) -c } and B = { x ; max η ( x ) ≤ E p ( z | x ) max ζ ( x, z ) -c } and split the expectation in two terms:

S

∫

∫

∑

y

x

∈

S

z

Looking at both terms separately, starting with R A where f s does not use z (or corresponds to f 1 ):

̸

̸

̸

̸

We can obtain more straightforwardly R B :

Hence, the risk of f s is given by;

s

[

ℓ

01

c

(

f

(

x, z

)

, y

)]

p

(

y

|

x, z

)

p

(

z

|

x

)

p

(

x

)

(62)

In the following, we prove that which would imply that

Showing that δ A d ≤ 0 and δ B d ≤ 0 We first consider δ A d :

We can (once again) further partition the space based on f o . We divide A d in two based on the decision function of f o :

̸

Now, we consider a different two-stage classifier f o ∈ H and f s = f o . We show that any f o ∈ H will lead to a higher risk R 01 c ( f s ) ≤ R 01 c ( f o ) , therefore proving that f s = f ∗ .

̸

We further partition the space X by splitting A and B where f s = f o and f s = f o

̸

̸

We then use those new partitions to further decompose the risks of f s and f o in 4 terms:

We can then write the difference between the risks as:

Continuing our development of δ A d :

̸

(96)

At this stage we can focus on one term at the time, starting with the part of A d where f o is not using z which is the integral over A dx :

̸

Going to the second term [**], when f 0 uses z . We start by restating the definition of the set A : A = { x ; max η ( x ) ≥ E p ( z | x ) [max y ζ y ( x, z )] -c } . (101)

̸

Combining both results together, we have that

Following similar steps, we also have that

̸

Final step Since we have both that δ B d ≤ 0 and δ A d ≤ 0 , we can conclude that

̸

This concludes the proof.

A.5 Proof of consistency for the multi class surrogate hinge loss

We start by restating our surrogate loss:

̸

̸

where we have t ∈ R K and v ∈ R K as the real-vectored outputs for f 1 and f 2 respectively with the constraints that || v || 1 = 0 and || t || 1 = 0 , and ˜ r ∈ [0 , 1] as a soft decision output.

Since our surrogate optimization provides us with the triplet t , v , ˜ r , we map these to a two-stage classifier f ∈ H using the following a link function φ : R K × R K × [0 , 1] →Y :

We can then define our risk as usual:

and consider the triplet of minimizers t ∗ ( x ) , v ∗ ( x, z ) , ˜ r ∗ ( x ) of such a risk, which correspond to a two-stage solution:

We prove that our surrogate loss is Bayes-consistent w.r.t to the ℓ 01 c loss by showing that 1)

and 2)

Taken together, those results guarantee that for any distribution p ( x, z, y ) , we have that:

which defines Bayes-consistency [Steinwart, 2007].

A.5.1 The solutions f ∗ of R 01 c ( f ) and f ∗ hinge coincide

Restating the definitions of the solution of the target and surrogate problems;

In this section, we show that the two solutions f ∗ , f ∗ hinge coincide.

Lemma A.1. For any distribution p ( X,Z,Y ) ;

Proof. We have previously shown in Appendix A.4 that the solution of our targeted problem

Next, we show that the two-stage classifier obtained from the triplet minimizer of our surrogate loss f ∗ hinge = φ ( t ∗ , v ∗ , ˜ r ∗ ) corresponds to this solution.

̸

̸

We can push the optimization problem inside the expectation w.r.t p ( x ) as t ( x ) , v ( x, z ) , ˜ r ( x ) are all functions of x (and the inner expectation term is guaranteed to be bounded):

̸

̸

Since the loss is a linear combination of two terms that respectively depend on t and v , we can see that for any ˜ r , the minimizers for t and v will always be equal to the minimizer of the individual terms:

̸

̸

is given by:

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]:

which gets converted into f ∗ 1 and f ∗ 2 by the link function φ (see Eqn 10):

Next, we turn to the decision function ˜ r ( x ) . Using ˜ r as shorthands for ˜ r ( x ) :

̸

̸

Since it is a linear combination of 1 -˜ r ( x ) and ˜ r ( x ) , it is clear that the minimizer ˜ r ∗ ( x ) will either be at 0 or 1. We can therefore rewrite the optimization problem as the following:

̸

̸

̸

̸

We consider the two cases for A ( x, ˜ r ) .

̸

̸

̸

̸

̸

̸

For the second case, using similar steps:

̸

This allows us to write the solution as

We have therefore shown that as

This concludes the proof.

Since we now have that ˜ r ∗ ( x ) is restricted to the binary values ˜ r ∗ ( x ) = { 0 , 1 } , we can rewrite the optimal classifiers that we previously obtained:

Therefore, we can see that the optimal t ∗ and v ∗ leads to the same solution for the internal classifiers of f ∗ . We have therefore shown that:

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 R 01 c ( φ ( t , v , ˜ r )) -R ∗ 01 c with the risk gap of our surrogate hinge loss R hinge ( t , v , ˜ r ) -R ∗ hinge .

Lemma A.2. For any distribution p ( X,Z,Y ) :

Proof. We start by developing the hinge risk R hinge ( t , v , ˜ r ) :

̸

̸

Next we develop the term associated to the optimal hinge risk R ∗ hinge :

̸

̸

Bringing both R ∗ hinge and R hinge ( f ) to evaluate the gap G ≜ R hinge ( f ) -R ∗ hinge , we can decompose the gap by a sum of 4 terms that are driven by the ground truth decision cases, i.e. r ∗ = 0 or ∗ = 1 and the decision of the model ˜ r ( x ) :

̸

̸

We can define the corresponding gap to each case as follows:

̸

̸

̸

̸

̸

̸

̸

̸

We can obtain a similar decomposition for the risk gap of our target risk R 01 c ( f ) . We recall the definition:

̸

̸

̸

̸

R 01 c ( f ) = E p ( x ) [ ✶ [ r ( x ) = 0](1 -η f 1 ( x ) ) + ✶ [ r ( x ) = 1] E p ( z | x ) [1 -ζ f 2 ( x,z ) + c ] ] (201) and decompose the gap risk with terms based on similar cases:

This makes sense: if the optimal decision is defer and we don't, the risk is diminished by the saved computation c (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).

and rewrite the total gap as:

To prove the result, we show that :

G 1 inequality Starting with G 1 and F 1 :

̸

We need to find a mapping Ψ such that the following holds:

Using the simple scaling function Ψ( x ) = 2( K -1) K x , we can see that the previous inequality holds:

̸

G 2 inequality For the next inequality with G 2 and F 2 , again using the same function Ψ( x ) = 2( K -1) K x , the inequality holds:

(231)

Following similar steps, we obtain

Putting all results together, we obtain

with Ψ(0) = 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 1 ∈ ∆ K for f 1 and p 2 ∈ ∆ K for f 2 , with the same link function φ . The cross entropy version of the loss that we consider is given by:

where g ( c ) is an arbitrary function. We again consider its associated risk:

and minimizing function:

Lemma 4.2. Cross-entropy surrogate loss is not Bayes Consistent.

There is no g ( c ) for which:

Proof. Following a similar reasoning as the proof of Lemma A.1 , we can readily find that the optimal predicted probability vectors p 1 and p 2 in f ∗ ce should match the posteriors:

Now, to find the optimal decision function of the cross entropy risk ˜ r ∗ ce , we can again obtain the solution as

by following the same steps that were taking to obtain Eqn. 153. We again consider the two cases:

The solution for the cross-entropy decision function is hence given by:

If we recall the solution decision of our target problem;

we are searching for a function g ( c ) for which

If we define the decision sets of x :

we can rewrite the condition Eqn. 248 as:

This implies that the function g ( c ) should satisfy:

However, our condition for r ∗ ( x ) is on the absolute scale, not on the log scale r ∗ ( x ) = η ∗ ( x ) < E p ( z | x ) [ ζ ∗ ( x, z )] + c . We will therefore have that for some η, ζ

This implies that there is no g ( c ) that can satisfy

This concludes the proof.

A.7 Extension to the multi-classifier setting

In this section, we provide a high-level view of how we could extend our 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 at the second stage (with associated costs c 1 , . . . c L ). We would need to introduce additional random variables Z 1 , . . . , Z L , and the loss in Eqn 2 would need to be generalized to multiple classifiers and costs:

̸

̸

̸

The corresponding solution for the decision boundaries (Eqn. 9) would become more complex. Instead of comparing the maximum posterior probability η ( 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:

The term max l ∈ 1 ,L ( E p ( Z l | x ) [max y p ( y | x, Z l )] -c l ) returns the index of the best model that can be used at the second stage.

Then, we could propose a multi-classifier surrogate loss (replacing Eqn. 11) by using a soft decision function that is now a multiclass probability vector ˜ r ( x ) ∈ [0 , 1] L +1 with ∑ L l =0 ˜ r l ( x ) = 1 . We would also need to introduce L new learnable vectors t (0) , . . . , t ( L ) , and an index-dependent cost term K K -1 g ( l, c 1 , . . . , c L , K ) , where g ( l, c 1 , . . . , 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:

̸

It remains to be seen whether we can verify the consistency of a loss of this form.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths\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.\n\nQUESTIONS:\nSee the 'weaknesses' section.\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-06-28"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths:\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.\n\nQUESTIONS:\n1. 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?\n\nLIMITATIONS:\nYes.", "TITLE": "Official Review", "RECOMMENDATION": 3, "REVIEWER_CONFIDENCE": 4, "CLARITY": 3, "IMPACT": 2, "ORIGINALITY": 3, "DATE": "2025-07-03"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS AND WEAKNESSES:\nClarity: 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.\n\nQUESTIONS:\n- 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 ?\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 3, "REVIEWER_CONFIDENCE": 4, "CLARITY": 3, "IMPACT": 2, "ORIGINALITY": 2, "DATE": "2025-07-03"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS AND WEAKNESSES:\n1. 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).\n\nQUESTIONS:\n1. 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.\n\nLIMITATIONS:\nThe 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.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-07-03"}