Gatekeeper: Improving Model Cascades Through Confidence Tuning
Authors: Stephan Rabanser
Nathalie Rauschmayr
Achin Kulshrestha
Petra Poklukar
Wittawat Jitkrittum
Sean Augenstein
Congchao Wang
Federico Tombari
Venue: NeurIPS 2025 poster ()
Keywords: cascading
deferral
selective prediction
confidence tuning
Primary area: infrastructure
TL;DR: We introduce Gatekeeper, a novel loss function that calibrates smaller models in cascade setups to confidently handle easy tasks while deferring complex ones, significantly improving deferral performance across diverse architectures and tasks.

Abstract: Large-scale machine learning models deliver strong performance across a wide range of tasks but come with significant computational and resource constraints. To mitigate these challenges, local smaller models are often deployed alongside larger models, relying on routing and deferral mechanisms to offload complex tasks. However, existing approaches inadequately balance the capabilities of these models, often resulting in unnecessary deferrals or sub-optimal resource usage. In this work, we introduce a novel loss function called Gatekeeper for calibrating smaller models in cascade setups. Our approach fine-tunes the smaller model to confidently handle tasks it can perform correctly while deferring complex tasks to the larger model. Moreover, it incorporates a mechanism for managing the trade-off between model performance and deferral accuracy and is broadly applicable across various tasks and domains without any architectural changes. We evaluated our method on encoder-only, decoder-only, and encoder-decoder architectures. Experiments across image classification, language modeling, and vision-language tasks show that our approach substantially improves deferral performance.

GATEKEEPER: Improving Model Cascades Through Confidence Tuning
Stephan Rabanser"
Nathalie Rauschmayr Google rauschmayrogooglee.coo com
Achin Kulshrestha
Princeton University rabanserdprinceton.. edu
Google
kulac@google.cmm com
Petra Poklukar Google poklukar@google..oom com
Wittawat Jitkrittum
Sean Augenstein
Google
Google saugenst@google.com com
vittawat@google.com
Congchao Wang
Federico Tombari
Google
Google tombari@google. com
congchaovangggoolle..com
Abstract
Large-scale machine learning models deliver strong performance across wide range of tasks but come with significant computational and resource constraints. To mitigate these challenges, local smaller models are often deployed alongside larger models, relying on routing and deferral mechanisms to offload complex tasks. However, existing approaches inadequately balance the capabilities of these models, often resulting in unnecessary deferrals or sub-optimal resource usage. In this work we introduce novel loss function called GATEKEEPER for calibrating smaller models in cascade setups. Our approach fine-tunes the smaller model to confidently handle tasks it can perform correctly while deferring complex tasks to the larger model. Moreover, incorporates a mechanism for managing the trade-off between model performance and deferral accuracy, and is broadly applicable across various tasks and domains without any architectural changes. We evaluate our method encoder-only, decoder-only, and encoder-decoder architectures. Experiments across image classification, language modeling, and vision-language tasks show that our approach substantially improves deferral performance.
Introduction
In recent years, large-scale machine learning models such Gemini (GeminiTean et al. 2023). GPT-4 (Achiam et al. 2023) or Claude (Anthropic) 2024) have gained significant traction due to their remarkable ability address wide array of tasks. These tasks range from natural language understanding and generation, including machine translation, summarization, and conversational agents, computer vision applications like image recognition, object detection, and image captioning. The versatility and high performance of these expansive models make them invaluable tools across
diverse domains, including healthcare (Nazi & Peng 2024 finance (Li et al. 2023), education (Wang
et al. 2024b), and entertainment (Gallotta et al. 2024).
Deploying and operating such large models presents significant challenges in terms of latency, memory, compute and storage (Pope et al. 2023). Optimizing inference costs is an active research
Work done while a Student Researcher at Google.
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

Join: Join: Annuracy of and K5
0.950
(d)
Input
Response Ms(2)
Response Mz(
0.925
0.900
No
0,075
0.850
Small model Msg
Defer?
Large model Mr
0.825
Ideal
0.000
Random
0.775
(b)
(c)
Realized
0.50
0.75
1.00
Doterral Ratio
Figure 1: Overview of the cascading setup (left) and performance trade-off (right). Left: Cascading determines which inputs should be predicted by a small model Mg or routed to a large model Mr, Right: Performance measured as a trade-off between joint accuracy across Mg and Mr and deferral ratio. Ideal deferral strategies optimize this trade-off and push the realized deferral curve closer to the ideal deferral depicted in (d). (a) depicts full deferral: (b) depicts no deferral; and (c) depicts excessive deferral of requests that could have been correctly handled by Ms.
area which includes both techniques for reducing the size of the existing large model such as model compression (Hoefler et al. 202T), model pruning (Ma et al., 2023 Cheng et al. 2024) and distillation (Yang et 2024), and those aiming to leverage sequence of models such as specu speculative decoding (Leviathan ct al. (2023) and model cascades (Dohan et al.) 2022; Chen et al. 2023a Gupta
ct al. 2024; Chen et al. 2024a). However, due to scaling laws showing that the performance of
Large Language Model (LLM) increases with its size (Kaplan et al. 2020), the latter category of
methods leveraging a sequence of models is currently more promising direction to lower inference
costs without sacrificing the capabilities of large models.
Both speculative decoding and model cascading rely on the existence of a large performant model Mr. and small model MMs that cheap, fast, and less accurate. Speculative decoding leverages Ms for generating set of draft tokens that are then validated by Mr in parallel, a technique successfully de- ployed in industry applications (Leviathan] 2024). contrast, model cascades leverage deferral rule for selecting the most suitable model to process given request (see Figure Ieft). While the success of speculative decoding necessitates a highly performant Ms to generate quality draft tokens, model cascades allow the deployment of a less capable Ms by invoking Mr only for inference requests out- side the small model's scope. In this work, we contribute to the advancement of the model cascades.
Model cascades achieve efficient deferral by optimizing two objectives: compute budget and joint accuracy. We illustrate this trade-off in Figure (rightt) Assume we have x inference requests and a small model Ms that uses only 20% of the compute budget of the large model My, There are three worst-case scenarios: (a) Ms defers all requests to Mr, yielding the highest joint accuracy (equal to that of My.) but the worst compute budget (1. (1.2x), since both models process all inputs; (b) Ms never defers, achieving the lowest compute budget (0. (0.2x) but also the lowest joint accuracy (equal to that of Ms); (c) Ms defers only requests it would have answered correctly, increasing compute budget over (b) without improving joint accuracy. In contrast, the ideal case (d) occurs when Ms defers only the requests it would misclassify, achieving optimal joint accuracy for given compute budget (0.2-Ix). We define how closely a model approximates this ideal as its deferral performance.
In this paper, we address the following research question:
How can we optimize model cascades to maximize deferral performance?
In other words, we focus designing effective model cascades by making the small model more aware of what it does not know. We do so by introducing a general-purposs loss function, called GATEKEEPER, that calibrates the small model's confidence in its predictions. By fine-tuning Ms to output high confidence for correct predictions and low confidence for incorrect ones, we improve the reliability of its uncertainty estimates and facilitate learning of common tasks- proving deferral performance. Crucially, GATEKFEPER includes a built-in mechanism for managing the trade-off between model and deferral performance, and is applicable arbitrary architectures.
directly im-

---

We empirically demonstrate the efficacy of the GATEKEEPER loss on encoder-only vision models for image classification, decoder-only language models (LMs) for flosed-form text generation, and encoder-decoder models for vision-language (VI.) tasks such as open set classification and captioning. Our main results show that models trained with GATEKFEPER outperform an untuned baseline by 0.72x/2x on CIFAR-100/ 100/innlmaegeet and 7x/10x on ARC-e/c, respectively, in terms of deferral performance. As a result, GATEKEEPER paves the way for more scalable and efficient deployment strategies, leveraging the collaboration between local and large-scale models to deliver high-quality results in applications with real-time processing demands.
2
Related Work
Our proposed method improves model cascades through uncertainty- finetuning. Next, describe related work for both research areas. We provide an extended discussion in Appendix[B.]
Model Cascades. A cascade consists of series of models and a deferral rule which determines the appropriate model given an input request. The concept of model cascades has first been proposed
by Viola & Jones] (2001), where it is used to accelerate object detection models. Cascades have been extensively studied for classification-base computer Vision (Wang et al. [0177] (Tapezznik &
Saligrama| 2013 [Bolubbasi et al. [017a] Jitkrittum et al. (223) and in models for natural language
processing (Dohan al. 2022; Mamou et al. [022 Varshney & Baral 2022).
Cascades are particularly promising in the context of generative models such as LLMs and VLMs since they can significantly reduce inference costs. In contrast to speculative decoding (Leviathan itt al., (2033), they aim to invoke the large model only for difficult examples. However, the two approached can also be combined (Narasimhan et al., 2025). While Chen et al. (2024b) combine the deferral logic with speculative decoding to generate initial tokens using larger models and later tokens using a smaller model, the majority of research on model cascades has focused on using pre-trained
LLMs with a post-hoc deferral logic (Narasimhan et al. 2022 Jitkrittum et al. 2023 Yue et al. 2024). Kolawole et al. (2024) use agreement across multiple models to make deferral decisions, while Gupta et al. (2024) present a method to learn a deferral rule based on quantiles of per-token log probabilities.
Model cascades can be improved through training and fine-tuning, Wang et al. (2024a) train the small model only on easier examples by masking tokens where both large and small models are incorrect. Enomoro Eda (202T) extend the training objective of image classification models with confidence calibration. In contrast, our approach extends cascades to VLMs and boosts inference performance by making smaller models less confident when incorrect.
Uncertainty-Aware Models. Extensive research has been conducted in the field of uncertainty quantification in deep learning and refer to Abdar et al. (2021) for detailed survey. While many models, measuring uncertainty for generative
methods have been proposed for classification models is still an active area of research. Based on the assumed level of access model internals, existing methods can be summarized into three main categories:
Black box methods operate solely via the model's 's query interface by injecting tailored instructions into prompts. These modify the prompt x by appending instructions x' for the model to respond less confidently: +- x x'x'. Related methods are confidence quantification (Shrivastava et al. 2023). rejection and remote model awareness (Kadavath et al.) 20022), and self-critiquing (Gou et al. 2023).
Xiong al. (2024) show that LLMs express their confidence through prompting and sampling strategies and their experiments indicate that these models tend to be overconfident.
Gray box approaches employ confidence- based strategies centered on post- post-processing the model's log- its. Many uncertainty techniques such as ensembling (Lakshminarayanan et al. (2017) and Bayesian methods (Blundell et al. 2015)) are not scalable. Related techniques are max confidence (Hendrycks)
& (201el][011),, predictive entropy, and confidence reduction prompting. Malinin & Gales 2201
uses token-entropy as measure of uncertainty in auto-regressive models and Kuhn et al. (0233) leverages linguistic invariances via semantic entropy.
White box methods use uncertainty-awann aware fine-tuning to produce better-calibrated models. Chuang jet al. 2024) introduces Self-REF, framework that leverages confidence tokens during fine-tuning to improve downstream routing. Krishnan et al. (2024) proposes an uncertainty-aware causal language modeling loss that captures the trade-off between accuracy and calibration. contrast, our method calibrates the model so that correct predictions receive low uncertainty and incorrect ones high
3

---

uncertainty. We apply this uncertainty-aware model in cascade inference system, where it improves overall performance. Prior work by Rawat et al.] (202T) pre-partitions data into easy and hard examples, e.g. based on Mr,'s confidence, and trains M, with explicit labels. We improve on this static partitioning by dynamically assigning examples during training based on Ms's 's current state.
3 The GATEKEEPER Loss
3.1 Overview & Setup
Our framework consists of large, highly capable model Mr. and smaller, resource-efficiet model Mg. We assume that S E N and L. E N represent the parameter count of each model with S < L. Both models can either function as classifiers (i.e.. M ID' - [C] with IR denoting the input space and C' the number of total classes), or (multi-modal) sequence models (i.c., M : R + [V] where is the vocabulary and T is the sequence length). We include experiments on all of these model classes in Section 4 Furthermore, we do not require shared model family to be deployed on both Ms and MM; for example, MMs could be a custom convolutional neural network optimized for efficient inference and M vision transformer (Dosovitskiy) (2020). The primary objective is to design a deferral mechanism that enables Ms to decide when to return its predictions without the assistance of Mr and when to instead defer to it. We assume that Mr. is either outside of our control (e.g., an API endpoint) or too costly to modify, and that only Ms is subject to adaptation.
Deferral decisions are made using signals derived from the small model .Ms as this approach is typically more cost-effective than employing a separate routing mechanism (Tecrapitaayaon et al.. 2016). Approaches that involve querying the large model Mr to assist in making deferral decisions common at test time are excluded from our setup. Such methods-coommnn in domains like L.MMs are coun- terproductive to our goal since querying M defeats the purpose of making a deferral decision in the first place. Examples of these inapplicable methods include collaborative LLM frameworks (Mielke) et al.
2022) and techniques that rely on semantic entropy for uncertainty estimation (Kuhn et al. As part of our setup, we assume that Mr, dominates Ms as per the following assumption.
2023
Dominance Assumption. Let D denote the target deployment distribution defined over covariates X and labels V. We assume that Mr dominates the Ms with high probability under D; formally,
Pr Pr [ML(z) + F AAA.AA Ms(x) = v] < 5. (x.g)=
with < This "almost-always" dominance, supported by scaling-law trends (Kaplan et al. 2020). implies that deferring from Ms to Mr cannot hurt accuracy in expectation, while still allowing rare counter-examples where the small model outperforms the large model. Note that we empirically observe 5 = O across all tasks considered in this work, meaning that M, strictly dominates Ms.
As discussed in Section [] the choice of deferral strategy often depends the level of access available to Ms. We assume white box access with full access to Ms's 's internals. As such, deferral mechanisms can be directly integrated into the model's architecture and parameters. This involves fine- tuning Ms to predict deferral decisions or to incorporate rejection mechanisms within its predictive process. Our work falls into this category as it proposes new loss function to fine-tune Ms.
Our goal is to train a small model that can effectively distinguish between correct and incorrect predictions. While many past works have considered the question of whether it is possible to find proxy measures for prediction correctness, the central question we ask is:
Can optimize the small model to separate correct from incorrect predictions?
We show that this is indeed achievable through a carefully designed fine-tuning stage that does not require any architectural modifications. This ensures that the ability to separate correct from incorrect decisions is integrated seamlessly into Msss 's existing structure.
3.2 Confidence-Tunnigg for Deferral
Prerequisite: Standard Training. We begin with an Ms that has already been trained on the tasks it is intended to perform upon deployment. However, due to its limited capacity, Ms cannot achieve

---

the performance levels of My. Importantly we make no assumptions about the training process of Ms wwetther was trained from scratch without supervision from an external model or with the help of soft labels through a distillation approach.
Stage (Finetuning): Correctness-Awan Finetuning with GATEKEEPER. Next, introduce correctness-awar loss, dubbed GATEKEEPER, to fine-tune My for improved confidence calibration. Specifically, the model is trained make correct predictions with high confidence while reducing the confidence of incorrect predictions (see Figure 2). This loss can either rely on true labels or utilize the outputs of Mr with soft probabilities as targets.
In its canonical form, GATEKEEPER is defined as a hybrid loss L = aLCon + (1 - a Lincor with
N Coom
N
1
Σ 1=
1l 9.) (x)C 84) Lincon
Hy # UKL (p. (). (2)
N
N
Here, Y/ and Ys are the true and predicted labels for x,, respectively, p; is the predicted probabil- ity distribution of MMs over classes, represents the uniform distribution over all classes, N de- notes the number samples in the current batch, (0, 1) is tunable hyperparameter controlling the emphasis between correct and incorrect predic- tions, and the cross-entropy function and KL diver- gence are defined as CE(p, y) = E. Ve log Pe and KL(P 9) E. Pe log( Pc) respectively. We note that similar loss has previously been pro- posed in Outlier Exposure (OE) (Hendrycks et al.] 2018) for distibbtitiiiooonnn (OOD) sample detec- tion. Here, the goal is to make sure that OOD exam- ples assigned low confidence scores by tuning the confidence on an auxiliary outlier dataset, How- ever, to the best of our knowledge, this idea has not previously been used to improve deferral per- formance of smaller model cascading chain.
Batch
Minimize CB(ys,
Probability
0.5
Correct
C1
C4
C5
Probabiltt
Incorrect
C1
C2
C3
C4
C5
q.
Minimize (L. ()
Figure 2: GATEKEEPER Overview: We want correctly predicted samples to maintain their current prediction by ensuring that cross en- tropy is decreased. At the same time, we want incorrectly predicted samples to yield uniform confidence across classes, leading to low over- all confidence score (high predictive entropy).
We emphasize that the trade-off parameter a plays a critical role as part of this optimization setup it directly influences model utility and deferral performance. A lower value of emphasizes reducing confidence incorrect predictions by pushing them closer to the uniform distribution, making the model more cautious in regions where it may make mistakes. Conversely, higher value of a encourages the model to increase its confidence on correct predictions, sharpening its decision boundaries and enhancing accuracy where is already performing well. Thus, a serves as a crucial hyperparameter that balances the trade-off between improving calibration by mitigating overconfidence errors and reinforcing confidence in accurate classifications. By appropriately tuning practitioners can control the model's 's behavior to achieve desired balance between reliability uncertain regions and decisiveness in confident predictions, tailored to the specific requirements of their application.
We further generalize this loss token-based models (e.g., LMs and VLMs) where
Loom Σ 1(y., = de. (..p..CE...) Lincote
Σ
1 * 0.,) (pc, 1| 2)). (3)
Here, Vi,t and Di,t denote the true and predicted tokens at position for sample i, Pi, is the predicted token distribution at position for sample i, and T is the sequence length for the token-based model. The token-level loss ensures that correct token predictions are made confidently while incorrect tokens assigned smaller confidences.
Practical Computation of the Gatekeeper Loss. We evaluate Gatekeeper once per mini-baich within the standard training loop-no no auxiliary passes or data-set re-shuffling are required. Given a mini-batch B = {(1,, ) we perform single forward pass through Ms. This allows us to obtain (i) logit vector fo(x)), and (ii predicted labels y, arg max. 2i,0- Two
5

---

binary masks m. = H(: ds and m - are computed on-the-fly. The hybrid loss is then assembled in fully vectorized manner with components Coor = 인 Ei m CE(Ps, and Lincor N 2 mi (p.l1.) where p. softmax(2;). Because both masks and losses are computed inside the same tensor graph, back-propagation incurs only the cost of ONN x C) element-wist operations-identice
to vanilla cross-entropy step. This single-pass design keeps the computational overhead negligible while guaranteeing that in the full loss = a& + (1 0k) Lincorr every sample contributes either Leon or Lincor in the same optimization step.
Stage 2 (Inference): Confidence Computation & Thresholding. After fine-tuning Mg with GATEKEEPER, we apply standard confidence- and entropy-based based techniques for model uncertainty to obtain a deferral signal. We use the selective prediction framework to determine whether query point x E R.' should be accepted by .Ms or routed to MM. Selective prediction alters the model inference stage by introducing deferral state through a gating mechanism Yaniv & Wiener 2010). At its core, this mechanism relies on a deferral function g IR_ IR which determines if Ms should output prediction for a sample x or defer to Mr. Given a targeted acceptance threshold the resulting predictive model can be summarized as:
(L.)MM. = [Ms(x) g(x) TT (M otherwise.
(4)
Classification Models (Max Softmax). Let Ms produce a categorical distribution {p(y = |x) over C classes. Then we define the gating function as
ply c|x). (x) ISe5C
(5)
Token-based Models (Negative Predictive Entropy). Let Ms produce a sequence of categorical distributions {p(y= x))c=1 for each token index & E T. Then we define the gating function
T
c ply- =c|x) log = =<|x).
(NENT(x)=
(6)
where y, [C] is the predicted token time step L, p(u: x) is the (conditional) probability of token l at step i, and T is the total number of token positions for the sequence. Across both model classes, higher values of JCL or INENT indicate higher prediction confidence (i.e., lower uncertainty).
Experiments
In this section, we detail the experiments used to evaluate the effectiveness of GATEKEEPER across three distinct model classes: encoder-only classification models, decoder-only language models, and encoder-decoder vision-languas models. Each setup involves a cascade where a smaller model can defer inputs to larger, more capable model. Our method and all competing baselines are applied on top of the same training/fine-tnnigg protocol used for obtaining the initial Ms and My models.
4.1 Encoder-only Setup (Classification Models)
We comprehensively assess the performance of GATEKEEPER across different model architectures and task types, starting with image classification. We train both a large model and a small model on the following datasets: 1IFAR-10/100 (Krizhevsky et al. 2009), Food-101 (Bossard al., 2014). and TinyImageNe1200 (Le & Yang. 2015). For both CIFAR datasets we use a RRsNeet-1 18 (He et al.) 2016) as Mr and a custom CNN as Mg. For Food- 101 and TinylmageNe1200 we instead use a ResNet-50 (He et al. 2016) as Mr and a Mobilenet V3 Small (Howard et al. 2019) as Ms. where the latter is trained using knowledge distillation from the big model.
Evaluation Metrics. We measure the performance of GATEKEBPER and the resulting deferral function g(-) using the following performance metrics (see example in Figure for an overview):
1. The Distributional Overlap of Confidences of Correct and Incorrect Predictions 8, is defined as the integral of the minimum of the probability density functions (PDFs) of confidence scores

---

a) Distribational Overlap
Deferral Performance
Delezral (aCuzeai
Correct
0.95
Razdem Deferval (acemnel
Donsity
Overlap Area So
Incorrect
Accuracy
0.90
Performanco Area (Apai)
Headroom Area
0.05
Useful Azea (Ausea)
Realzed Deferral (acCwal)
0.80
No deferral acci
Full deferral accl.v))
0.0
0.2
0.4
0.6
0.8 1.0
0.0
0.2
0.5 Deferral Ratio
0.4
0.8
1.0
Confidence
Figure 3: Performance metries overview: (a) Distributional Overlap s.: the densities of confidence scores for correctly (green) and incorrectly classified (red) samples, with the overlap area shaded in blue. Smaller values are better (4). (b) Deferral Performance sd: how joint accuracy between MM and Mr varies with deferral ratio, showing random (red), ideal (green), and realized (black) deferral strategies. The blue region shows the realized performance gain, the hatched portion represents the range of useful deferral functions, and the green region indicates the potential headroom over the realized deferral. Larger values are better (1).
for correctly classified samples, (c), and incorrectly classified samples, Pincor (c) (see Figure 30). Formally, given the confidence sets Ceoi and Cincorr. the overlap 8, is computed as
8, / min (Pener (c), Prncom {c)} de.
where the PDFs are estimated using Kernel Density Estimation (KDE), If s, then MMs cannot distinguish the confidence distribution of correct and incorrect predictions; if s. = O, then Ms can perfectly separate correct and incorrect predictions. Note that related way of capturing the distributional separability is given by the Area Under the Receiver Operating Characteristic Curve (AUROC) which we discuss in Appendix C.3
2. Deferral Performance sd: To formally quantify how well MM defers difficult inputs to ML, examine the joint performance across all possible deferral ratios r E [0, 1], where r denotes the fraction of inputs sent Mr, based on particular threshold (recall Equation 4 Figure b) illustrates how, as r increases from to 1, the overall (joint) accuracy acc(r) increases from the accuracy of MMs (yellow circle, deferral) to the accuracy of Mr (purple square, full deferral). Useful deferral models are constrained to operate between random deferral (accrand- red dotted line) and ideal deferral (accideal, green dashed line). The ideal deferral accideal corresponds to the oracle solution that perfectly defers examples misclassified by MM and we discuss its exact functional form in Appendix [..3] We also define the realized deferral curve, acCreal. as the joint accuracy obtained under the learned deferral strategy g(-) employed by MM and Mr. The deferral performance metric sa is then given as:
Apert So (accreal(r) acCrand (r)) dr Auseful
(8)
So (acoideat(7) - (r)) dr
This ratio quantifies the fraction of the potential improvement over random deferral that has been realized by the achieved deferral strategy. Note that sa indicates perfect deferral, matching the idcal strategy, while an Sa = 0 implies no improvement over random deferral.
Accuracy of the small model (Mc(Ms): Finally, since GATEKEEPER emphasizes patterns for distinguishing correct/incorrect examples, the model is longer encouraged to minimize the classification loss over the full population. As a result, improving on the correct/incorreet separation task can lead to changes in utility over the full data distribution. Hence, practically useful deferral methods need to balance both deferral performance and the accuracy of Ms.
Our main results are shown in Figure 4 We report performance for both a bascline model (an instance of Ms not trained with GATEKEEPER) and small models trained with GATEKEEPER at various a values. We also compare against Narasimhan al. (2022), a common cascading baseline for supervised learning tasks. For all models, we compute deferral performance and correct/incorrect separation (center, left). The strongest performance occurs at low as, where the model pushes outputs of incorrect examples closer to uniform. However, this comes at cost: the small model's accuracy degrades at low as (right), highlighting that the model effectively "unlearns" performance on harder examples to focus casier ones. At larger as, accuracy remains stable
Results.
7

---

MC)(M))
0003 Baseline
e- 0.9
2-0.7
0.5
2-3
0.1
Nurssiman et sl., 2022
Distriuctional Overlap (
Deferral Porformsance (
Acouracy Ms (T)
0.1
3.0
0.8
0.8
04
0.0
0.4
500000
0.0
meelfn
cileta
bodie1
etartto
Emdd
dida
Figure 4: Performance on image classification tasks. We observe that lower levels of a lead to decreased distributional overlap between correct/incorreet predictions (left), increased deferral performance (center) and generally decreased performance over the full data distribution (right). These results support our conclusion that the small model MM learns to refocus on easier subsets of the distribution while understanding more reliably when it should defer to the large model ML.
IOC Dieceliee
0.5
0-0.3
Redure Conf
Atacuur
Gapta ml. 2004
Distributional Overiap
())
Daferral Performance
1
Accuracy
3.6
5.6
2.4
0.4
6.4
Figure 6: Performance on language modeling tasks. Similar as Figure 4 In addition a non-tuned baseline, we also add an uncertainty prompting baseline, an Answer "N" option, as well as the post-hoc confidence calibration method from Gupta al. (024). We observe that GATEKEBPER outperforms other methods at lower levels of a.
or improves slightly, as training emphasizes already well-predicted points. While the baseline from Narasimhan et al. (2022) preserves model accuracy, it requires explicit estimation of the expert model's correctness- necessitatiii either an architectural change (c.g., an added prediction head) or a separate prediction network. In contrast, GATEKEEPER relies solely on confidence tuning of the small model, making it operationally easier to deploy. This approach not only simplifies implementation but also leads to improved deferral performance by producing more reliable uncertainty estimates through a stronger correlation with correctness.
This result highlights a critical trade-off which is directly controlled by a:
How strongly do we want to degrade model performance over the full data distribution in order obtain better deferral model?
We note that this compromise between raw model utility and deferral performance is not surprising and similar trade-offs exist fairness (Dutta et al., 2020; Yaghini et al.) (023) and privacy (Abadi et al. 2016, Rabanser et al. 2023). We study this trade-off explicitly in Higure 5 showing (i) clear negative correlation between deferral performance and the small model's accuracy; and (ii) clear positive correlation between the overlap of correct/incorrect confidences and the accuracy of Ms.
4.2 Decoder-Only Setup (LLMs)
In the decoder-only setup, we explore the application of LLMS.. Our primary models of interest are the scalable LMs from the Gemma model class (GemmaTTean et al.) 2024). We choose Gemma2B as Ms and Gemma7B as My, Similar to the encoder-only setup, we employ smaller LMs as the initial classifiers to manage simpler next-token prediction tasks. The deferral strategy involves routing only those token sequences that exhibit high uncertainty-as determined by high predictive entropy-to the more powerful model ML-
8

---

Our experiments begin by taking the instruction-tune tuned checkpoints of Gemma2B and Gemma7B and fine-tuning both models on the training split of each dataset to ensure that the model (i) performs well on the task and (ii) is familiar with the desired response format. This step is performed using standard supervised fine-tuning. Next, we fine-tune My with GATEKEEPER on the same training split to reduce confidence on incorrect next-token predic- tions. Finally, evaluate the model trained with GATE- KEEPER on a validation split. The datasets used are ARC- e/c (Clark al., 2018), MMLU (Hendrycks al.) 2020).
Distributional Overlap So (1)
0.5
Baseline
0.4
Cp.9
0.3
0g.5
0o.3
0.2
0.2
0.4
0.6
0.8
and GSM8K (Cobbe et al. 2021). The evaluation metrics for our LLM experiments match those used in Section 4.1
Small model accuracy acc(MM) (1)
Deferral Performance So (t)
Results. We present our main results Figure 6 compar- ing the baseline model's deferral and correct/incorrect sep- aration ability to our fine-tuned model across different as. We observe similar trend as in the image classification set- ting: higher as maintain raw prediction performance closer to the baseline but offer limited gains in separation, while lower as improve deferral more substantially at the cost of overall accuracy. In addition to the baseline model (not fine- tuned with GATEKEEPER) ) we include results from Gupta et al. 2024) (an extension of Narasimhan et al. (2022) to
Ro.1
Co.33
0.8
doo.5
0.7
doo
0.6
Baseline
0.5
0.4
0.2
0.4
0.6
0.8
loken-based sequence models), as well as two uncertainty prompting baselines (described in Appendix C.2): (i) Re- duce Confidence, which appends instructions to encourage the model to lower confidence when uncertain; and (ii) An- swer "N", which instructs the model to respond with "N"
Small model accuracy acc(MMs ( (1)
tinyimagenet200
cifar10
food101
cifar100
uncertain. We find that GATEKEEPER outperforms Gupta itt al. (2024) in terms of correct/incorrect sepration and de- ferral at the cost of overall utility. Consistent with prior findings from Kadavath et al. (2022), the prompting base- lines do not reliably improve deferral.
Figure 5: Performance trade-off between small model accuracy acc(Ms) and deferral evaluation metries. The baseline model obtained without fine-tuning KEEPER is often the most accurate model over the full data distribution. With the introduction of GATEKEEPER we can improve distinguishability of correct/incorrect predictions (top) as well as deferral performance (bottom) at the expense of model utility. Success- ful cascading solutions in practice need to balance both maintaining high model accuracy and deferral performance.
using GATE-
4.3 Encoder-Decoder Setup (VLMs)
Finally, we examine models with both visual and textual processing capabilities, ideal for tasks requiring joint im- age understanding and language generation. We use the PaliGemma (Steiner et al. 2024) model family- decoder models designed for VL tasks such image cap- tioning, visual question answering, and descriptive image classification. The encoder processes input images into rich feature representations, while the decoder generates textual outputs. We use PaliGemmal IB as Ms and PaliGemma7B Mr,-. Our deferral strategy runs the smaller VLM on all inputs and only defers to the more bigger 7B model when Ms 's predictive entropy falls below a set threshold.
Similar to our experiments on LMs in Section 4.2 we employ two stages of fine-tuning. First, we take the instruction- tunee checkpoints of PaliGemmal and PaliGemma7B and then fine-tune both models the training split of a given dataset. Next, we fine-tune only MMs using GATEKEEPER before evaluating the model on a validation/test split of the dataset. The datasets we consider are two classification datasets (VQAv2 (Goyal et al. 2017). A12D (Hiippala et al. 202T) and two captioning datasets (Cococap (Lin et al. 2014), Screen2 Words (Wang et al. 2021)) This allows us to evaluate GATEKEEPER in closed-form lasion-language classification setups and open-form text generation.
Factuality Scoring. For classification tasks we apply our analysis in the same way as Section 4.2 However, for captioning datasets we need to evaluate the quality of caption generated by PaliGemma. To do that, we compute factuality score which judges whether the generated caption semantically coherent with respect to reference caption using the Gemini LLM (GeminiTeam et al. 2023

---

8558 Biseliw
G.1
Distriuttonal Overlap 4)
Deferral Performance
Poarson Correlation ( 11
Deferral
Accuracy (
Performance ( (T)
Fadiuality Mr
0.8
0.8
0.6
Socke 66
0.2
0.6
0.6
V.
0.4
011
0.0
0.0
Figure 7: Performance on VLM classification (left) and captioning tasks (right). Consistent with results in Figures 4 and 6 we see that smaller as lead to improved deferral performance.
Specifically, the Gemini LLM is prompted with an instruction of the form: "Are these captions semantically equivalent?", followed by both the candidate caption and the reference caption. The model then responds with either "Yes" or "No". Finally, we compute the log-likelihowe of each response and normalize it to probability, reflecting the LLM's confidence in the captions being factually aligned. We detail this process in Appendix [.4 and denote the factuality score for input point x; with candidate caption y; and ground truth caption y; as Spac (pac (F.,
Measuring Correlation Between Factuality and Negative Predictive Entropy. Since the result of evaluating $Fac ( is no longer binary, our evaluation metrics which previously relied on accuracy cannot be used directly evaluate deferral performance and the correct/incorrect entropy distribution separation. We address this issue by replacing the distributional overlap computation with the Pearson correlation ((9NENT(x:)) SPac y:)) between the negative predictive entropy of a caption INENT (x:) and its associated factuality score F (Y:,.))) y:)). We also adapt our deferral performance metric from Equation (8) to rely on factuality measures instead of accuracy.
Results. We present our results in Figure 7 comparing the baseline model's 's deferral ability to our fine-tuned models across different as. For classification (Figure 7 left), we observe the same trends as in previous classification and language modeling experiments. For captioning (Figure right), GATEKEEPER increases the correlation between factuality and negative predictive entropy, enabling better deferral from Ms to Mr. as a decreases. This shows that our method generalizes beyond closed-form classification to open-form sequence generation tasks. While we also benchmarked the prompting baselines from Section 4.2, PaliGemma did not return responses for the modified likely due to its rigid pretraining and prompting instructions (Beyer et al. 2024).
5 Conclusion
In this work we present a novel loss function called GATEKEEPER for improving confidence calibra- tion in a cascade between a small local and a larger remote model. Our loss is architecture and task agnostic, making it flexibly applicable across wide range of applications. Our results on encoder- only classification models, decoder-only language models, and encoder-decoder vision-language models demonstrate that our approach improves over standard confidence-brrt based deferral rules and effectively leads the small model to unlearn how to handle complex queries in favor of casier ones.
Limitations. While our approach demonstrates promising results, several limitations remain. First, we assume that only the smaller model can be fine-tuned, whereas in some applications the larger model could also be adjusted to improve deferral. Second, in language settings, GATEKHEPER may be overly aggressive: multiple token sequences can convey the same meaning, so penalizing deviations based on exact tokens rather than semantics may be suboptimal. Ideally, deferral should account for semantic correctness rather than surface-level mismatches. Third, we did not extensively evaluate across diverse model families in LLM and VLM settings, though we did include such comparisons for classification tasks. Fourth, our evaluation focusses on two-stage cascade consisting a small model and big model. However, we confident that GATEKEEPER can scale multi-stage setups (see Appendix [B.5). Finally, our use of generative model (c.g., Gemini) to evaluate captioning introduces the possibility of erroncous judgments, as LLMs are themselves imperfect evaluators.
10

---

References
Abadi, M., Chu, A., Goodfellow, I., McMahan, H. B., Mironov, I., Talwar, K., and Zhang. L. Deep learning with differential privacy. In Proceedings of the 2016 ACM SIGSAC conference computer and communications security, pp. 308-318, 2016.
Abdar, M., Pourpanah, F., Hussain, S., Rezazadegan, D., Liu, L.,, Ghavamzadeh, M., Fieguth, P., Cao, X., Khosravi, A., Acharya, U. R., et al. A review of uncertainty quantification in deep learning: Techniques, applications and challenges. Information fusion, 76:243-221 297.2297222 2021.
Achiam, J., Adler, S., Agarwal, S., Ahmad, L.. Akkaya, I., Aleman, L., Almeida, D., Altenschmidt, J., Altman, S., Anadkat, S., et al. Gpt-4 technical report. arXiv preprint arXiv:2303.08777, 2023.
Anthropic. Models overview anthropic. https: httte::/ddocs.aatrrppic..om wwt.apacc.ccoo/eee//aaooacaiacee muild-with-clan citiuiioeiitiiiinseeio 2024. [Online; accessed 24-Jamuary-2025),
Azaria, A. and Mitchell, The internal state of an Ilm knows when it's lying. arXiv preprint arXiv:2304. 13734, 2023.
Beyer, L.. Steiner, A., Pinto, A. S., Kolesnikov, A., Wang. X., Salz, D., Neumann, M.. Alabdulmohsin. I. Tschannen, M., Bugliarello, E., et al. Paligemma: A versatile 3b vlm for transfer. arXiv preprint arXiv:2407.07726. 2024.
Blundell, C., Cornebise. J.. Kavukcuoglu, K., and Wierstra, D. Weight uncertainty in neural network.
In International conference on machine learning, pp. 1613-1622. PMLR, 2015.
Bolukbasi, T., Wang, J., Dekel, O., and Saligrama, V. Adaptive neural networks for fast test-time prediction. 2017a
Bolukbasi, T., Wang, J., Dekel, O., and Saligrama, V. Adaptive neural networks for efficient inference. In International conference machine learning, pp. 527-536. PMLR, 20170 7b.
Bossard, L.., Guillaumin, M., and Van Gool, L. Food- 101 mining discriminative components with random forests. In European Conference Computer Vision, 2014.
Chen, B., Zhu, M., Dolan-Gavitt, B., Shafique, M., and Garg, Model cascading for code: Reducing inference costs with model cascading for LLM based code generation, 2024a. URL https://arxiv. org. abs////s///// 2405. 15842
Chen, L., Zaharia, M., and Zou, J. FrugalGPT: How to use large language models while reducing cost and improving performance. arXiv preprint arXiv:2305.cm 05/ 176, 2023a.
Chen, L.. Zaharia, M., and Zou, J. Frugalgpt: How to use large language models while reducing cost and improving performance. arXiv preprint arXiv:2305.05177 2023b.
Chen, Z., Yang, X., Lin, J., Sun, C., Chang, K., and Huang, J. Cascade speculative drafting for even faster LLM inference. In The Thirty-iighth Annual Conference Neural Information P Processing Systems, 2024b. URL https::///// /ttp::////ppeeiieeeeeee net/ forumm wit.mmmmddd-1-yiiiiii.
Cheng, H., Zhang, M., and Shi, J. Q. A survey deep neural network pruning: Taxonomy, comparison, analysis, and recommendations. IEEE Transactions on Pattern Analysis and Machine Intelligence, 46( (6(12):10) 100010558-0000
2024. doi: 10. 10.119//PAMM..224..4444
Chuang. Y.-N., Zhou, H., Sarma, K., Gopalan, P., Boccioc J., Bolouki, S., and Hu, X. Learning to route with confidence tokens, 2024. URL https://axxiv.. wttp:///..rraaa//////// 13284
Clark, P., Cowhey, I., Etzioni, O., Khot, T., Sabharwal, A., Schoenick, C., and Tafjord, Think you have solved question answering? try arc. the ai2 reasoning challenge. arXiv preprint 17X.v:1803.05457. 2018.
Cobbe, K., Kosaraju, V., Bavarian, M., Chen, M., Jun, H., Kaiser, L., Plappert, M., Tworek, J., Hilton, J., Nakano, R., et al. Training verifiers to solve math word problems. arXiv preprint arXiv:2110.1410 2021.

---

Ding, D., Mallick, A., Wang, C., Sim, R., Mukherjee, S., Ruhle, V., Lakshmanan, L. V., and Awadallah, A. H. Hybrid Ilm: Cost-efficient and quality-aware query routing, arXiv preprint arXiv:2404. 14618, 2024.
Dohan, D., Xu, W. Lewkowycz, A., Austin, J., Bieber, D., Lopes, R. G., Wu, Y., Michalewski, H., Saurous, R. A., Sohl-Dickstein, J., et al. Language model cascades. arXiv preprini arXiv:2207.10342, 2022.
Dosovitskiy, A. An image is worth 16x16 words: Transformers for image recognition at scale. arXiv preprint ar arXhv:2010.11992 2020.
Dutta, S., Wei, D., Yueksel, H., Chen, P.Y., Liu, S., and Varshney, K. there a trade-off between fairness and accuracy? a perspective using mismatched hypothesis testing. In International conference on machine learning, pp. 2803-2813. PMLR, 2020.
El- Yaniv, R. and Wiener, Y. On the foundations of noise-free selective classification. Journal of Machine Learning Research, 11(53):1055-1641, 2010. URL http::/jjmr. http://mmmr..rg//ppprr//////////
el- yl-yanivvia.. html
Enomoro, and Eda, T. Learning to cascade: Confidence calibration for improving the accuracy and computational cost of cascade inference systems. Proceedings of the AAAI Conference on Artificial Intelligence, 35(8):7331-7339 May 2021. doi: 10.1609/aaai.v333. v35i8. 16900. URL
battps:/ojj) aaal org/ iidex.ii/xxx..ee/ pttpp///////rrrtce//eeeeee////ee
Gallotta, R., Todd, G., Zammit, M., Earle, S., Liapis, A., Togelius, J., and Yannakakis, G. N. Large language models and games: A survey and roadmap. arXiv preprint arXiv:2402. 2024.
GeminiTeam, G., Anil, R., Borgeaud. S., Alayrac, J.-B., Yu, J., Soricut, R., Schalkwyk, J., Dai, A. M., Hauth, A., Millican, K., et Gemini: family of highly capable multimodal models. arXiv preprint arXiv:2312. 1805, 2023.
GemmaTeam, G., Riviere, M., Pathak, S., Sessa, P. G., Hardin, C., Bhupatiraju, S., Hussenot, L., Mesnard, T., Shahriari, B., Ramé, A., et al. Gemma 2: Improving open language models at practical size. arXiv preprint arXiv:2408.00111
2024.
Geng, J., Cai, F., Wang, Y., Koeppl, H., Nakov, P., and Gurevych, A survey of confidence estimation and calibration in large language models. arXiv preprint arXiv:23/1.088 1.08298. 2023.
Gou, Z., Shao, Z., Gong, Y., Shen, Y., Yang, Y., Duan, N., and Chen, W. Critic: Large language models can self-correct with tool-interactive critiquing. arXiv preprint arXiv:2305.11738 2023.
Goyal, Y., Khot, T., Summers-Suee D., Batra, D., and Parikh, D. Making the v in vqa matter: Elevating the role of image understanding in visual question answering. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp.
2017.
Guo, C., Pleiss, G., Sun, Y., and Weinberger, K. Q. On calibration of modern neural networks. In International conference on machine learning, pp. 1321-1330. PMLR, 2017.
Gupta, N., Narasimhan, H., Jitkrittum, W., Rawat, A. S., Menon, A. K., and Kumar, S. Language model cascades: Token-level uncertainty and beyond. 2024. URL https://opeerrevieu net/ Kggauzid=ggaScc44V1
He, K., Zhang. X., Ren, S., and Sun, J. Deep residual learning for image recognition. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 770-778, 2016.
Hendrycks, D. and Gimpel, K. A baseline for detecting misclassified and out-of-distribution examples in neural networks. arXiv preprint 3rXiv:1610.02136. 2016.
Hendrycks, D., Mazeika, M., and Dietterich, T. Deep anomaly detection with outlier exposure. arXiv preprint arxiv: arXiv:1812.. 012.0.000666..... 2018.
Hendrycks, D., Burns, C., Basart, S., Zou, A., Mazeika, M., Song, D., and Steinhardt, Measuring massive multitask language understanding. arXiv preprint arXiv:2009.03300 2020.
12

---

Hiippala, T., Alikhani, M.. Haverinen, J., Kalliokoski, T., Logacheva, E., Orekhova, S., Tuomainen, A., Stone, M., and Bateman, A. Ai2d-rst: multimodal corpus of 1000 primary school science diagrams. Language Resources and Evaluation, 55:661-688 2021.
Hoefler, T., Alistarh, D., Ben-Nun, T., Dryden, N., and Peste, A. Sparsity in deep learning: Pruning and growth for efficient inference and training in neural networks. Journal of Machine Learning Research, 22(241):
2021.
Howard, A., Sandler, M., Chu, G., Chen, L.-C., Chen, B., Tan, M., Wang, W., Zhu, Y., Pang, R., Vasudevan, V., et al. Searching for mobilenetv3. In Proceedings of the IEEE/CVF international conference on computer vision, pp. 1314-1324, 2019.
Jitkrittum, W., Gupta, N., Menon, A. K., Narasimhan, H., Rawat, A., and Kumar, When does confidence-based cascade deferral suffice? In Oh, A., Naumann, T., Globerson, A., Saenko, K., Hardt, M., and Levine, S. Advances pp. 9891-99909 Curran Associates, Inc., 2023.
Advances in Neural Information Processing Systems, volume 36,
Kadavath, S., Conerly, T., Askell, A., Henighan, T., Drain, D., Perez, E., Schiefer, N., Hatfield-Dodds, Z., DasSarma, N., Tran-Johnson, E., et al. Language models (mostly) know what they know. arXiv preprint arXiv:2207.0522 2022.
Kaplan, J., McCandlish, S., Henighan, T., Brown, B., Chess, B., Child, R., Gray, S., Radford, A., Wu, J., and Amodei, D. Scaling laws for neural language models. arXiv preprint arXiv:2001.08361, 2020.
Kolawole, S., Dennis, D., Talwalkar, A., and Smith, V. Agreement-basee cascading for efficient inference, 2024. htttp: wttp://www..rrg////////1/// 02348
Krishman, R., Khanna, P., and Tickoo, O. Enhancing trust in large language models with uncertainty- aware fine-tuning, 2024. URL ltttss//aaaa org/ abs/ abs///11 22////12..0000 02904
Krizhevsky, A., Hinton, G., et al. Learning multiple layers of features from tiny images. 2009.
Kuhn, L.,, Gal, Y., and Farquhar, S. Semantic uncertainty: Linguistic invariances for uncertainty estimation in natural language generation. In The Eleventh International Conference on Learning Representations, 2023. URL /tt////tt////peeeevieeer.r www.:nne//ooomdttttttoet
Lakshminarayanan, B., Pritzel, A., and Blundell, C. Simple and scalable predictive uncertainty estimation using deep ensembles. Advances in neural information processing systems, 30, 2017.
Le, Y. and Yang, X. S. Tiny imagenet visual recognition challenge. 2015.
Leviathan, Y. Looking back at speculative decoudinn 2024. URL httpp://research.gogo. blog/looking back-at- www.aappppecaaaeeeeeeeeeeoeeeeee
Leviathan, Y., Kalman, M., and Matias, Y. Fast inference from transformers via speculative decoding.
In International Conference on Machine Learning, pp. 19274-19286. PMLR, 2023.
Li, Y., Wang, S., Ding, H., and Chen, H. Large language models in finance: survey. In Proceedings of the fourth ACM international conference on Al finance, pp. 374-382, 2023.
Lin, S., Hilton, J., and Evans, O. Teaching models to express their uncertainty in words. arXiv preprint arXiv:2205.1433 2022.
Lin, T.-Y., Maire, M.. Belongie, S., Hays, J., Perona, P., Ramanan, D., Dollár, P., and Zitnick, C. L. Microsoft coco: Common objects in context. In Computer Vision-ECCV 2014: 13th European Conference, Zurich, Switzerland, September 6-12, 2014, Proceedings, Part V 13, pp. 740-755. Springer, 2014.
Lin, T.-Y., Goyal, P., Girshick, R., He, K., and Dollár, P. Focal loss for dense object detection. In Proceedings of the IEEE international conference on computer vision, pp. 2980-2988, 2017,
Liu, I.., Pan, Y., Li, X., and Chen, G. Uncertainty estimation and quantification for llms: simple supervised approach. arXiv preprint 4rXiv:2404.15900 2024.
13

---

Ma, X., Fang, G., and Wang, X. LLM- On the structural pruning of large language models. In Thirth-seventh Conference on Neural Information Processing Systems, 2023. URL https //openreview.co wwt.:nnutfrr/ee/.e (t ja
//
Mahaut, M., Aina, L.., Czamowska, P., Hardalov, M., Müller, T., and Märquez, L. Factual confidence of Ilms: On reliability and robustness of current estimators. arXiv preprint arXiv:2406. 13415, 2024.
Malinin, A. and Gales, M. Uncertainty estimation in autoregressive structured prediction. In International Conference on Learning Representations, 2021. URL httttpps///pppereeveeeeeee forum?id= Morumiid=-JU5y-22500
Mamou, J., Pereg, O., Wasserblat, M., and Schwartz, R. Tangobert: Reducing inference cost by using wwt.:/org//a.aa.eeeee 06271
cascaded architecture, 2022.
Mielke, J., Szlam, A., Dinan, E., and Boureau, Y.-.. Reducing conversational agents' overconfi- dence through linguistic calibration. Transactions of the Association for Computational Linguistics, 10:857-872, 2022.
Narasimhan, H., Jitkrittum, W., Menon, A. K., Rawat, A., and Kumar, S. Post-hoc estimators for learning to defer to am expert. In Koyejo, S., Mohamed, S., Agarwal, A., Belgrave, D., Cho, K., and Oh, A. (eds.), Advances in Neural Information Processing Systems, volume 35, pp. 29292 29304. Curran Associates, Inc., 2022.
Narasimhan, H., Jitkrittum, W., Rawat, A. S., Kim, S., Gupta, N., Menon, A. K., and Kumar, S. Faster cascades via speculative decoding. In The Thirteenth International Conference on Learning Representations, 2025. URRL
Nazi, Z. A. and Peng, W. Large language models in healthcare and medical domain: A review. Informatics, 11(3), 2024. ISSN 2227-9709, doi: 10.3399/infommaties/ 1030057. URL https // www. mdpi. www.comoo///// 0ttp//222229779111/////
Nie, L., Ding. Z., Hu, E., Jermaine, C.,, and Chaudhuri, S. Online cascade learning for efficient inference over streams, arXiv preprint arXiv:2402.0451/ 2024.
Pereyra, G., Tucker, G., Chorowski, J., Kaiser, L.,, and Hinton, G. Regularizing neural networks by penalizing confident output distributions. arXiv preprint arXiv:/701. 06548, 2017.
Pope, R., Douglas, S., Chowdhery, A., Devlin, Bradbury, J., Heek, J., Xiao, K., Agrawal, S., and Dean, J. Efficiently scaling transformer inference. In Song, D., Carbin, M., and Chen, T. (eds.), Proceedings of Machine Learning and Systems, volume 5, pp. 606-6244 Curan, 2023.
Rabanser, S., Thudi, A., Guha Thakurta, A., Dvijotham, K., and Papernot, N. Training private models that know what they don't know. Advances in Neural Information Processing Systems, 36: 5371 13711-53727, 2023.
Rawat, S., Zaheer, M., Menon, K., Ahmed, A., and Kumar, S. When in doubt, summon the titans: Efficient inference with large models, 2021. URL https: org abs 2210. 10305
Shnitzer, T., Ou, A., Silva, M., Soule, K., Sun. Y., Solomon, J., Thompson, N., and Yurochkin, M.
Large language model routing with benchmark datasets. arXiv preprint arXiv:2309. 2023.
Shrivastava, V., Liang. P., and Kumar, A. Llamas know what gpts don't show: Surrogate models for confidence estimation. arXiv preprint arXiv:2311.088 1.08877,777 2023.
Steiner, A., Pinto, S., Tschannen, M., Keysers, D., Wang, X., Bitton, Y., Gritsenko, A., Minderer, M., Sherbondy, A., Long, S., et al. Paligemma 2: A family of versatile vims for transfer. arXiv preprint arXiv:2412.0355: 2024.
Teerapittayanon, S., McDanel, B., and Kung, H.-T. Branchynet: Fast inference via early exiting from deep neural networks. In 2016 23rd international conference on pattern recognition (ICPR), pp- 2464-2469, IEEE, 2016.
14

---

Trapeznikov, K. and Saligrama, V. Supervised sequential classification under budget constraints. In Carvalho, C. M. and Ravikumar, P. (eds.). Proceedings of the Sixteenth International Conference on Artificial Intelligence and Statistics, volume 31 of Proceedings of Machine Learning Research, pp- 581-589, Scottsdale, Arizona, USA, 29 Apr-01 May 2013. PMLR. URL get..https:/prrocedeine mlr w..ppeesss/S///rraeeeeoooo/ html
Varshney, N. and Baral, C. Model cascading: Towards jointly improving efficiency and accuracy of NLP systems. In Goldberg, Y., Kozareva, Z., and Zhang, Y. (eds.), Proceedings of the 2022 Conference Empirical Methods in Natural Language Processing, pp. 11007-11021, Abu Dhabi, United Arab Emirates, December 2022. Association for Computational Linguistics. doi: 10. 18653/ 111202..mmnp-mmain..27 httpttss: htt.lattsss/aalaattollogg//
org/ 2022. emnlp-main. 756
Viola, P. and Jones, M. Rapid object detection using boosted cascade of simple features. In Proceedings of the 2001 IEEE Computer Society Conference on Computer Vision and Pattern Recognition. CVPR 2001, volume 1, pp. I-I, 2001. doi: 10. 10.110//CVR..200.990000
Wang, B., Li, G., Zhou, X., Chen, Z., Grossman, T., and Li, Screen2words: Automatic mobile ui summarization with multimodal learning. In The 34th Annual ACM Symposium User Interface Software and Technology. pp. 498-510, 2021.
Wang. C., Augenstein, S., Rush, K., Jitkrittum, W., Narasimhan, H., Rawat, A. S., Menon, A. K.,, and Go, A. Cascade-aware training of language models, 2024a. URL https::/arxiv: org/ abs/ 00060)
Wang, S., Xu, T., Li, H., Zhang, C., Liang, J.. Tang, J., Yu, P., and Wen, Q. Large language models for education: A survey and outlook. arXiv preprint arXiv:2403. 18105, 2024b.
Wang, X., Luo, Y., Crankshaw, D.,, Tumanov, A., and Gonzalez, J. E. Idk cascades: Fast deep learning by learning not to overthink. In Conference on Uncertainty in Artificial Intelligence, 2017.
Xiong. M., Hu, 7., Lu, X., III, Y., Fu, J., He, J., and Hooi, B. Can LLMs express their uncertainty? an empirical evaluation of confidence elicitation in L.MMs. In The Twelfth International Conference on Learning Representations, 2024. URL https:://ppereeaaaib http:///ppppaaeeeeeeeeeeeeeenrnp atee:/foooommdddeeeeFFFFFpg
Yaghini, M., Liu, P.. Boenisch, F., and Papernot, N. Learning to walk impartially the pareto frontier fairness, privacy, and utility. 2023.
Yang, C., Zhu, Y., Lu, W., Wang, Y., Chen, Q. Gao, C., Yan, B., and Chen, Y. Survey on knowledge distillation for large language models: Methods, evaluation, and application. ACM Trans. Intell. Syst. Technol., October 2024. ISSN 2157-6904. doi: 10.1145/3699518 URL https://///doi org/ 10. 11995/6999918 Just Accepted.
Yue, M., Zhao, J., Zhang, M., Du, L., and Yao, Large language model cascades with mixture of thought representations for cost-efficient reasoning. In The Twelfth International Conference Learning Reprexentations, 2024. URL https: http://pprreeeeeeeeeeeoetttaaaaa
15

---

A Broader Impact
This work contributes to the responsible and efficient deployment of machine learning systems by improving the decision-making capabilities of smaller, local models in model cascade architectures. By introducing loss function that calibrates model confidence with respect to correctness, our approach enhances both the performance and transparency of automated systems that must decide when to act autonomously and when to defer to a more capable model. This design can improve the accessibility and sustainability of machine learning applications by reducing reliance on large, energy-intensive models partees-patticllarly important in low-resource environments or edge computing.
At the same time, the ability to fine-tune smaller models to strategically abstain from uncertain predictions raises important considerations for fairness and accountability. In high-stakes applications such as healthcare or finance, improper tuning of the deferral threshold-or uncalibrated confidence estimates- lead the systematic denial of service or misallocation of computational resources. Care must be taken to ensure that such systems are thoroughly evaluated not only for average performance but also for differential performance across subgroups. Moreover, the use of large models as fallback decision- makers assumes their correctness, which may not always hold, especially in underrepresented domains. We therefore encourage developers and practitioners to accompany deployments of cascade- based systems with rigorous audits of fairesss reliability, and alignment with human values.
B
Additional Background
B.1
Related Work
B.1.1
LLM Routing
Ding et al. (2024) propose a hybrid LLM inference pipeline that routes each query either to a small on-device model larger high-quality model based on the query's predicted difficulty and tunable quality threshold. This cost-aware router allows dynamically trading off accuracy for efficiency, enabling up a 40% reduction in expensive model calls without degrading answer quality. Similarly, Shnitzer et al. (2023) present a method to select the best model from a pool of pre-trained LLMs for each imput by learning "router" many benchmark tasks. Without requiring labeled examples from the new target task, their approach uses existing datasets to train input- hasee model selectors, which consistently outperform always using the single best LLM for all queries.
B.1.2 Model Cascade Learning
Nie et al. (2024) introduce an online cascade-learning framework where lightweight models are incrementally trained to imitate powerful LLM's decisions data stream, deferring to the LLM only when necessary. They cast cascade construction as an imitation-learning problem with theoretical no-regret guarantees, achieving LLM-level accuracy while cutting inference cost by up to 90% and maintaining robustness to distribution shifts over time. [Chen et al.) (2023b) outline strategies for reducing LLM usage cost and present FrugalGPT, cascade approach that learns to route queries through combinations of smaller or larger LLMs to balance cost and performance. Their experiments show that an adaptive use of multiple models can match the accuracy of the strongest individual LLM (e.g.. GPT-4) with up to 98% cost savings. It can also slightly exceed GPT-4's accuracy at equal cost, highlighting the benefit of cascades that allocate queries to the most appropriate model for each input.
B.1.3 Confidence Calibration in LLMs
Jitkrittum et al. (2023) analyze the classical strategy of confidence- based deferral in model cascades, wherein model hands off a stronger model if its confidence is below threshold, to determine when this simple strategy succeeds or breaks down. They derive the optimal deferral policy in theory and show that naive confidence thresholds perform well in general but can fail when later models are specialists (only reliable on certain inputs), when there is label noise, or under distribution shift scenarios where more sophisticated deferral criteria yield better performance. Geng et al.) (2023) provide comprehensive survey of methods for confidence estimation and calibration in LLM outputs. They review recent techniques to quantify uncertainty in large language model predictions, discuss challenges unique to LLMs, and highlight advancements that improve alignment between
16

---

model's reported confidence and its actual accuracy across tasks. Azaria & Mitchell (2023) find evidence that an LLM's internal activations encode whether or not producing a truthful answer, even when the model's output is incorrect or fabricated. By training a classifier on the model's hidden state (without fine-tuning the L.IM itself), they can often detect when the model is "Iying" or unsure, suggesting that large models internally recognize their mistakes or uncertainty despite outwardly confident responses. Similarly, [Lu et al. (2024) propose a supervised approach to LLM uncertainty quantification that leverages labeled examples and the model's hidden representations to predict the correctness of its answers. They show that incorporating features from the model's ' internal layers yields significantly improved uncertainty estimates and calibration across diverse tasks, with these gains transferring robustly to new domains. Notably, their method is easy to implement and can be adapted to different levels of model access (black-box white-box), making it widely applicable.
B.1.4 Confidence Verbalization in LLMs
Lin et al. (2022) demonstrate that GPT-3 can be fine-tuned to output a calibrated verbal confidence (e.g., "I'm m 90% sure") along with each answer. This model's stated confidence levels align well with its true correctness likelihood and remain fairly well-calibrated even under distribution shift, marking the first instance of an LLM explicitly expressing useful uncertainty estimates in natural language. Xioog et al. (2024) thoroughly evaluate black-box. methods for eliciting an LLM's M's self- reported confidence through prompting and answer sampling. They find that current LLMs tend to verbalize overly high confidence (mirroring human overconfidence)) but that carefully designed prompts, consistency checks across multiple sampled answers, and improved aggregation strategies can mitigate this issue. Moreover, larger models generally show better calibration and an improved ability to predict their own failures, though room for further improvement remains in making their expressed uncertainty truly reliable. |Mielke et al. (2022) examine whether a conversational agent's s expressed certainty corresponds its actual knowledge, showing that off-the-shelf dialogue models are poorly "linguistically calibrated." They demonstrate that a model's 's likelihood of giving a correct answer can be estimated via an auxiliary model and used as control signal to adjust the agent's responses. The resulting dialogue agent exhibits far less overconfident language when it is likely to be wrong, improving transparency about uncertainty in its answers. Finally, Mahaut al. (2024) assess the reliability of various methods to estimate an LLM's ' factual confidence - the probability that its answer correct- - under both in-domain and paraphrased inputs. Through a rigorous evaluation on QA and fact-checking tasks, they conclude that the most trustworthy confidence scores come from model-ine introspeectve approaches (e.g.. trained probe on hidden states), albeit at the cost of requiring full model access and training data. They also highlight that an LLM's confidence can be unstable under pereiiiigrreervinggg input variations (paraphrases), underscoring the need for more robust and stable confidence estimation techniques for factual correctness.
B.1.5 Speculative Decoding
As noted in the introduction and our related work section, there exists a connection between model cascading and speculative decoding. Speculative decoding is technique to accelerate inference by pairing a small, fast "draft" model with a larger target model. The draft model generates multiple candidate tokens in parallel, and the large model then verifies them in single forward pass by accepting valid tokens and rejecting others. This reduces the number of expensive calls to the large model, often yielding significant speedups without sacrificing output quality. We briefly describe how this line of work relates to our goal of deferral confidence tuning.
Speculative decoding accelerates every input by asking the large model to verify draft tokens that the small model proposes. As result, the large model is still invoked on all inputs. Our cascade, by contrast, aims to avoid calling the large model on inputs the small model can already solve. Consequently, speculative decoding optimizes in- in-place token-level latency, while our method optimizes end-to-end compute and monetary cost at the request level. Because the two techniques improve fundamentally different bottlenecks we do not directly compare against this class of methods. However, we note that both cascading and speculative decoding can be composed: after our deferral gate decides to consult the large model, one could still decode that portion of the input speculatively. We direct the interested reader to Narasimhan et al. (2025) for example of such hybrids.
17

---

B.1.6
Early Exiting
Early exit networks models augmented with intermediate classifiers that allow predictions to be made before reaching the final layer. At inference time, the model can stop early on "easy" inputs while continuing deeper for "harder" ones, reducing computation without compromising much accuracy. This adaptive approach makes them well-suited for resource-constrained or latency- sensitive applications. As discuss below, this class of approaches is, while related, ultimately still distinct from our goal for the following reasons:
1. exitine aims to skip the remaining layers of single network once intermediate classifier sufficiently confident, thereby offering a layer-level latency-aceuras trade-off off within one model. Our work tackles the fundamentally different problem of model-level deferral between a small and a large model. The goal is to keep the large model entirely idle for easy requests and to invoke it only when needed. Hence our primary measure of success is joint accuracy versus cross-model compute budget, not marginal delay per layer inside a fixed backbone.
2. Early-exit usually presupposes white-box control over the entire network architecture so that branch classifiers can be inserted and jointly trained (e.g.. BranchyNet (Teerapitayanon et al., 2016), Adaptive Neural Networks (Bolukbasi et al." (017bb). Our cascade setting explicitly targets heterogeneous or API-based laperts-large language models, vision-language mod- els, and so forth-that cannot be modified. Early-exits cannot be applied in such black-box situations, whereas our GATEKEEPER fine-tuning remains feasible. Moreover, early-exiting has seen limited success for autoregressive sequence generation, where every token depends on the full hidden state; in contrast, our experiments span both classification and open-ended generation tasks.
3. Because early-exit and cascading operate different granularities, they are complementary rather than competing techniques. One could in principle insert early exits inside the small model and still rely on calibrated deferral the large model-yyelding a three-level system (exit-1 - exit-2 - large model). Evaluating that combined design is a compelling direction for future engineering work, but it lies beyond the scope of our current study, whose contribution is a ceffidenne-tuningg loss independent of architectural changes.
B.2 Model Access Levels
In Figure g
show schematic overview of different model access levels discussed in Section 2
B.3 Ideal Deferral Curve
We present the functional form of the ideal deferral curve, denoted BCCideal()) (r), for a small (student) model Ms and large (teacher) model ML. Recall that r E [0, 1] denotes the deferral ratio, i.e., the fraction of inputs that Ms "defers" to ML. Let p. aMc(Ms), and pt acc(Mz) with 0 S P. S Pr I 1. Our goal is to describe the maximum achievable joint accuracy if exactly a fraction r of the data is deferred to the large model.
Intuition and Setup Since Mg achieves accuracy p.. it misclassifies a fraction (1 p.) of the inputs. In an ideal scenario, we defer exactly those inputs that Ms is going misclassify. Because Mt is more accurate (p: 2 Pa) every example misclassified by MMg benefits from being passed to ML.
Case 1: (1 (1-p.).
We can use entire deferral "budget" to cover only those inputs Ms would get wrong. Hence, deferring a fraction r of the data (all from Ms's 's mistakes) raises the overall accuracy by substituting Ms's 's errors with M accuracy on that fraction.
Case 2: r (1 (1-px).
We have enough capacity to defer all of Ms's 's mistakes, so the joint accuracy saturates at Pr. Deferring additional examples (which Ms would have classified correctly) will not improve the overall accuracy beyond pt.
18

---

Black box
No
Local
Input
Preprocess
Ms
Outputs
Postprocess
Deferral
Yes
MB B
Outputz
Remote
Gray box
No
Local
Input
Mss
Logits
Decode
Outputs
Deferral
Yes
MB
Outputz
Remote
White box
No
Local
Input
Ms
Tuning
Logits
Decode
Outputs
Deferral
Yes
MB
Outputz
Remote
Figure 8: An overview of different uncertainty quantification strategies depending on model access level.
Piecewise Functional Form
Thus, the ideal deferral curve can be expressed as:
Pt
Ps Ps
Ps
r
S (1 ps);
(CC)deal(r)
(9)
pl
(1 - Ps) < r S
When < S (1 (s), the overall accuracy grows linearly from acCideal (0) = Pe acCideat P.) = Past r = (1 11-pp p,), it remains constant at pt.
Figure (b) in the main paper plots this ideal deferral curve (green line). It serves upper bound on how effective any real deferral strategy can be. In contrast, a purely random deferral strategy produces linear interpolation (the red line), which is strictly below the ideal curve for most Consequently, the difference acCideat aCCrand (r) represents the maximum possible gain one can achieve by carefully selecting which examples to defer rather than choosing them random.
Summary
We summarize the key take-aways below:
Ideal Deferral Routes All Mistakes:
Only the inputs misclassified by Ms get deferred, guaranteeing the highest possible joint accuracy each deferral level
Piecewise Definition: Accuracy increases linearly from ps to p over the interval r
[0, (1 Ps)], then remains pt-
19

---

Upper Bound on Realized Deferral: No actual strategy can exceed this ideal curve, as it assumes perfect knowledge of which specific inputs Mg would misclassify.
B.4 Gatekeeper in the Context of Canonical Calibration Objectives
Motivation. Section showed that the GATEKEEPER loss improves deferral performance with minimal implementation effort. Because many calibration objectives also manipulate confidence, now position GATEKEEPER relative to four widely- used losses.
Canonical calibration objectives. Let po (y x) denote the softmax (or token) distribution pre- dicted by model with parameters and let y/ be the ground-truth label. Below we recap four popular alternate calibration objectives.
(a) Temperature scaling (Guo et al. 2017) applies a single scalar T>0 at test time: pr x) exp(z./T)). It preserves the rank ordering and therefore cannot tighten the ranking- based cove-coe curve, but can improve threshold-basee acceptance.
(b) Focal loss (Lin et al. 2017) adds down-weighting factor to easy examples, (L.p, y") -(1 - Py, log Pyr with E [0, oc). It improves class imbalance calibration but does not incorrect samples.
explicitly penalize over-confidence
(c) Confidence penalty (Pereyra et al., (017) regularises high-entropy predictions through CE(p, (") XXH()) While flattens all distributions, it does not distinguish between correct and incorrect cases.
(d) Outlier exposure (OE) (Hendrycks et al. 2018 adds KL-uniform loss on auxiliary OOD data, mirroring the second term of GATEKEEPER but only outliers, not in-distribution misclassifications.
How Gatekeeper differs. GATEKEEPER combines two complementary gradients: (i) standard CE term on correct predictions with an instance level mask, thereby sharpening those logits; (ii) a KL-to-unifom term incorrect predictions, flattening their confidence. This asymmetric design forces the scalar summary g(x) max, (o(==c x) (or token-entropy) to separate correct from incorrect points without requiring additional heads, auxiliary datasets, test-time tuning.
B.5 Extension to Multi-Level Cascades
Cascading can extend beyond single deferral Ms + MM multiple deferrals M₁ - M2 + - MMg GATEKEEPER makes no explicit assumption on the number of cascade levels it is inherently modular and agmostic to model architecture. In practice, one can fine-tune each model M; in the chain independently using the GATEKEEPER loss with stage-specific hyper-parameter that governs the balance between confident acceptance and deferral. natural hierarchical training procedure is to first apply GATEKEEPER to Mi to calibrate its confidence, then apply it to M2 on the subset of inputs deferred by M₁ and so on through Mg- This approach scales linearly in the number of levels and incurs only the familiar overhead of per-model fine-tuning-te joint optimization or additional routing networks are required. We therefore expect that deeper cascades will offer flexible trade-offss between computational cost and predictive accuracy across multiple model tiers, and we leave a full empirical evaluation of > 2 cascades to future work.
B.6 Extension to Multi-Level Cascades
Additional Experimental Details
C.1 CNN Used in Image Classification Experiments
Below we include representation of the Smal1CNN model used as MM in image classification experiments discussed in Section 4.]
SmalICNN
(features): Sequential
( (3,v2d 16, kernel_size-(3. (3, 3), stride-(1, 1), padding-(C (1, 1))
20

---

Eryyimgeeee2200
6oo119
difr110
cifari 100
0.6-
02-
GABC
LABI
0.6-
Dooitive
Rode
Sode
Dositivv
Figure 9: ROC curves for image classification experiments. Each figure shows the ROC curves for each the datasets considered in Section We observe that GATEKEEPER consistently increases separation of correct and incorrect confidence scores across varying a (colored curves) compared to the baseline (denoted with black dashed line).
Baseline
12
Conrect
insormect
0.
Condiiecc
Confidence
Condalence
Confidence
Condidence
Bosties
Confilence
Confdeenc
Confidence
Confilence
Confidence
Confióence
ettr1 Dindtt
255
0.2
Condalence
Condidence
Condidence
Uondalence
Confidence
Condicience
drar100 pusstt
12-
0.5 Canfidence
Confilence
Confidence
Confidence
Confilence
Confdeenc
Figure 10: Distributional overlap for image classification experiments. Left-most column shows
the results obtained using the untuned baseline, while the remaining columns correspond the results obtained using GATEKEEPER with decreasing a values. Rows correspond to the datasets considered in Section 4.1 We see that GATEKEEPER increases separation of correct and incorrect confidence scores compared to the baseline.
(1); BatchNorm2d (16, eps=1e-05, momentum=0.1, affine=True, track Beag__runinngstatt-Thee
(2) RelU((inplace-Tre))
(3): Maxpoll2d(eerneessizze
stride-2, padding-0, dilation-1,
ceil_mode-False)
(4) Conv2d 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(5) BatchNorm2d (32, eps=1e-05, momentum=0.1, affine=True,
track_ runninggttttttttsrruue
(6): ReLU((aplacce=Tre)
10
(7): (xaPool2d(xeneel stride=2, padding=0, dilation=1, ceil_mode=Falee)
(classifier): Sequential(
13
(0) (ifeture2.. out_features=66, bias=True)
14
(1): ReLU(inplace=TTue)
(2): Linear (in features=64, out fut_features=10, bias=True)
17
C.2 Reduce Confidence and Answer "N" Baselines
(i.e., a model that was not fine-tuned with our specialized Laer loss but from which we still compute predictive entropy as deferral signal), we also examine two additional methods aimed at eliciting uncertainty from the model directly via prompt
In addition to the baseline model in Section 4.2
21

---

Baseline
0.9
0.7
0.5
0.0.33
0.1
0.4
0.4
32-
02
0.2
0.00
5.8
Pootto
0.73
0.74
6.6-
0.70
0.5
0.5
0.50
0.4
0.90
0.00
0.00
0.0-
0.00
0.5
0.85
0.85
0.6
0.55
5.8-
0.50
5.7
0.7
0.7
5.6
0.6
0.6-
6.5
0.
5.4
355
Dedenaa
Deferaal Ratlo
Defercal
Dedorall
Deferval Rotio
Detbnra
Figure 11: Deferral curves for image classification experiments. Left-most column shows the results obtained using the untuned baseline, while the remaining columns correspond to the results obtained using GATEKEEPER with decreasing a values. Rows correspond to the datasets considered in Section 4.1 The results show that GATEKEEPER brings the realized deferral (black line) closer to the ideal deferral (dashed upper line).
Deden al rato
Deferna rate 0.
Defirral mite 0.5
- Deftrral rate 0.7
Dedoral
duylinagenet200
food101
cifar10
cifar1
0.56
0.7
0.55
0.42
6.50
0.5
0.25
0.30
0.5
0.34
0.4
03
Figure 12: Joint accuracy across different levels of a. For varying fixed deferral ratios, we observe that the accuracy Ms generally decreases as - O.
modifications. Both methods are black box approaches that only rely query interface to the model via prompt injection, and we provide their implementation details below.
Reduce Confidence. In this setting, we modify the original prompt by appending an additional instruction x' that encourages the model to respond with lower confidence when it is uncertain: x f- x x'. For instance, the instruction we add is:
x "Respond with low confidence if you uncertain."
We treat this appended text hint the model self-regulate its confidence when producing an answer. This is similar in spirit to other black box approaches such as confidence quantification. rejection awareness, remote model notice, and self-critiquing. Although Xiong et al. (2024) show that large language models can express aspects of their confidence via prompting, our experiments indicate that simply prompting the model express lower confidence does reliably improve the separation of correct versus incorrect predictions, nor does it offer advantages in deferral setting. These findings are in line with those reported Kadavath et al. 20022..
Answer "N
We also consider an alternate prompt modification, in which the appended instruction
is:
x' "Respond with 'N' if you are uncertain."
22

---

This approach explicitly instructs the model to produce a special "N" token to indicate uncertainty or lack of confidence. The intuition is that by introducing designated "uncertain" response. one might isolate uncertain cases for deferral. However, our results in Section 4.2 similarly show that the model's ability to follow this instruction is inconsistent and does not substantially improve performance as deferral model. The model often remains overconfident and fails to produce "N" in cases where it is in fact incorrect.
C.3 Additional metrics
In addition to the metrics outlined in Section we also consider the Area Under the Receiver Operating Characteristic Curve (AUROC) (SAUROC). The AUROC quantifies the model's ability to discriminate between correctly and incorrectly classified data points by evaluating the trade-off between the True Positive Rate (TPR) and the False Positive Rate (FPR) across various confidence thresholds 7. Formally, given the confidence sets Coorr and Cincorr, the AUROC is defined as
SARROC S.' (PR(T) (FPR(T),
(10)
where for each threshold
[0, we compute TPR(T) 1[oECur and FPR(T)
(oECuu 122r) Note that SAUROC indicates perfect separability and Sauroc 0.5 corresponds Cincsel
random guessing baseline.
C.4 Factuality Scoring
Factuality scoring with Gemini for reference caption r and candidate caption is computed as follows:
1. Compute the log-likelhhods. Let &same r) be the log- ligelioel that the model outputs "Same" for a given candidate caption e and reference r, and let (ar r) be the log- ligeliiel that the model outputs "Different".
2. Apply softmax. To convert these log-likelihoods into probabilities, we exponentiate and normalize:
exp cep(aammc(e (c,r))
(Same (,rr exp (esamm (c, r)) + exp(fvia(c,r))
exp(fDu((c,1)
(Diff c,r)
expl (Ssam (c, (c) + exp (c,ir(,))))
3. Interpret the probability. The value p( (Same c. c. is then taken the factual alignment score, expressing how confidently the model believes the candidate caption is factually aligned with the reference.
C.5 Additional Experimental Results
In this section, we provide additional experimental results further supporting our findings reported for image classification experiments in Section 4.1] In particular, show ROC curves in Figure and distributional overlap in Figure TO both demonstrating that GATEKEEPER increases the separation of correct/incorrect confidence scores. Similarly, the deferral curves in Figure 11 clearly show that GATEKEEPER successfully pushed the realized deferral (black line) closer to the ideal one (marked with dashed upper line). Lastly, we report the joint accuracy of Ms across varying a parameter in Figure 12 As discussed in Section 4 we observe that Ms's 's accuracy generally decreases with
0.
23

---

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:
Our abstract and intro reflects the contributions accurately.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results
be expected to generalize other settings.
is fine to include aspirational goals as motivation as long is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification: See Section S
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those not discussed in the paper.
The authors are encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification. asymptotic approximations only holding locally). The authors should reflect how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low images are taken in low lighting. Or speech-to-text system might be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach address problems of privacy and faimess.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory assumptions and proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [NA]

---

Justification:
We do not provide any theoretical claims or results.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
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 affects the main claims and/or conclusions of the paper (regardless whether the code and data are provided or not)?
Answer: [Yes]
Justification: See Appendix
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution is primarily a new algorithm, the paper should make clear how to reproduce that algorithm.
(b) the contribution is primarily a new model architecture, the paper should describe the architccture clearly and fully.
(c) If the contribution is new model (e.g., a large language model), then there should either way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reprodueibility. the case of closed-source models, may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
25

---

Answer:
[No]
Justification:
We are not providing code for this work.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips.cm
for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https
/ cpippi.ccppulli//uuuia wwwtii/gguie///eeeeceoocc CoodeSubeiiieeee ssi wsionnolli onnolicy)))) for more details.
The authors should provide instructions on data access and preparation, including how access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
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:
See Appendix
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, or as supplemental material.
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:
All of our reported results are reported as mean values over random runs.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, statistical significance tests, at least for the experiments that support main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call a library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
should be clear whether the error bar is the standard deviation the standard error of the mean.
26

---

It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of enrors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out range (e.g. negative error rates).
If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments compute resources
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?
Question:
Answer: [Yes]
Justification: See Appendix
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g.. preliminary or failed experiments that didn't make into the paper).
9. Code of ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https: htttps://euur: httppss:///urps...o/m http://.......occcc.c.cccciieeea
Answer: [Yes]
Justification:
The paper conforms the NeurIPS Code of Ethics.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special cireumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., if there is a special consid- eration due to laws or regulations their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [Yes]
Justification: See Appendix
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No. they should explain why their work has no societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that improvement in the quality of generative models could be used to
27

---

generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
there are negative socictal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).
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:
We not releasing any assets requiring safeguards.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators 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?
Question:
Answer: [Yes]
Justification:
We have cited related work appropriately.
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (c.g., website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, snperswwihcoode com/ datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if has changed) should be provided.
If this information is not available online, the authors are encouraged to reach out to the asset's is creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
28

---

Answer:
[NA]
Justification:
We are not releasing any new assets.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset used.
At submission time, remember to anonymize your assets (if applicable). You can either create anonymized URL or include an anonymized zip file.
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 details about compensation (if any)?
Answer: [NA]
Justification:
We did not conduct any crowdsourcing and/or research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included the main paper.
According to the NeurIPS Code of Ethics, workers involved in data collection, curation, other labor should be paid at least the minimum wage in the country of the data collector.
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 equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification:
We did conduct any studies requiring IRB approval.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent)
may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
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, originality of the research, declaration is not required.
Answer: [Yes]
29

---

Justification: Our confidence tuning method can be applied to LLMs and we demonstrate this usecase in our experimental results (Section In terms of paper writing, LLMs were used to help with editing.
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs as any important, original, stn-standard components.
Please refer to our LLM policy for what should or should not be described.
30


Reviews:
[{"review_id": "sHPPEU5RYK", "reviewer": "Reviewer_5Nnu", "summary": "This paper suggests a loss function, \"Gatekeeper,\" for calibrating smaller models, and uses it to fine-tune small models for use in model cascades.", "questions": "- The GATEKEEPER loss looks a bit odd as it ignores the entropy of $y_i$. For example if for some example $x$ there are two equally likely options $y,y'$, then even the Bayes optimal predictor will incur $\\mathcal{L}_\\text{incorr}$. Could the authors provide some intuition for why such a loss function would still make sense in this setting?\n- I have the same question for the gating function. In settings where $M_L$ would be uncertain, the gating mechanism will always defer the input to the larger model.\n- More generally, the GATEKEEPER loss only uses the predicted and true labels and does not even require query access to $M_L$. Could the fine-tuning process be improved given access to the logits of $M_L$ by trying to predict when the output distributions of $M_S$ and $M_L$ will differ, rather than quantifying the uncertainty of $M_S$?\n- Could the authors clarify the connection between Gatekeeper and calibration? While I agree that Gatekeeper will push the model towards the uniform distribution whenever it is unsure, this is not quite the same as calibration.", "limitations": "yes", "rating": 5, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 3, "significance": 2, "originality": 2, "strengths_and_weaknesses": "Strengths:\n- The paper is well written and easy to follow\n- The proposed method is simple and only requires fine-tuning the small model $M_S$\n\nWeaknesses:\n- I don't fully understand the connection between Gatekeeper and calibration – see the questions below\n- The proposed method does not leverage the query access to the large model $M_L$ for the fine-tuning, and only uses it during inference, which suggests further improvements could be made if the logits for $M_L$ are known on the training set.", "was_revised": true, "final_justification": "I originally recommended acceptance and decided to keep my score after the discussion phase.", "rebuttal": "We thank the reviewer for their thoughtful evaluation and for **recognizing the practical relevance of our framework**. Their comments raise important points about the **assumptions behind the Gatekeeper loss, the gating behavior under uncertainty, and its connection to calibration**. We address each concern below, clarifying the loss design and discussing possible extensions. For all questions, we are happy to add extended sections to the appendix.\n\n---\n\n``1) Does Gatekeeper still make sense if we allow for examples with high inherent label noise?``\n\n``2) Same question for the gating function.`` \n\nWe acknowledge that, when the predictive distribution $p(y|x)$ for a sample $x$ is inherently ambiguous—e.g. two labels $y,y^{\\prime}$ are equally plausible—the Bayes-optimal predictor will place mass on both classes, and our KL-to-uniform term $\\mathcal{L}_{\\text{incorr}}$ will be triggered whenever the small model’s top-1 guess happens to miss the sampled ground-truth. This behaviour is in fact desirable for deferral: **the loss encourages the small model to *express* its uncertainty by flattening the output distribution whenever it is not confident**, thereby raising the probability that the subsequent gating function will defer the input. Conceptually, the system partitions samples into three mutually exclusive cases:\n\n- (i): $g(x)\\ge\\tau \\ \\land\\  \\hat{y}_{S}=y$ \n- (ii): $g(x)<\\tau \\ \\land\\  \\hat{y}_{L}=y$ \n- (iii): $g(x)<\\tau \\ \\land\\  \\hat{y}_{L}\\neq y$ \n\ncorresponding to the small model answering correctly, the large model rescuing the mistake, and an irreducible error that neither model can solve. **Optimizing $\\mathcal{L}_{\\text{incorr}}$ improves the *boundary* between (i) and (ii)**, which is exactly the region that matters for useful deferral, while leaving case (iii) unaffected. Alternative objectives that try to anticipate the large model’s internal confidence would **collapse cases (ii) and (iii) into a single target but would require the small model to emulate the large model’s latent beliefs**—an infeasible endeavour in most practical scenarios and a key difference between *cascading* and *routing* strategies. Routing methods typically decide which model to send a query to (potentially partially querying any of the candidate models for determining the final routing decision itself), whereas cascading typically only uses signals produced by the small model. Finally, note that **our three-case breakdown explicitly accounts for inputs on which the large model itself remains uncertain (case (iii))**. In such instances the gate will indeed route the sample to $\\mathcal{M}_L$, but this does not hurt joint accuracy—neither model can answer correctly—and only incurs an extra compute cost, which practitioners may cap by choosing a higher deferral threshold $\\tau$ if desired\n\n---\n\n``3) Gatekeeper only uses the predicted and true labels and does not even require query access to M_L.``\n\nWe agree that, **when query access to $\\mathcal{M}_L$ is available, its logits can provide a complementary training signal—yet a key strength of Gatekeeper is that it *does not depend* on this access**, making it broadly deployable. Importantly, our training objective is easily augmented. One candidate solution would be to simply add an alignment term $\\mathcal{L}_A = \\frac{1}{N}\\sum_i KL(p^{(S)}_i | | p^{(L)}_i)$ to our loss \n\n$$\\mathcal{L}_\\text{GK+} = \\mathcal{L} + \\beta \\mathcal{L}_A$$\n\nwhere $p^{(S)}_i$ and $p^{(L)}_i$ denote the small- and large-model posteriors for sample $i$, and $\\beta$ controls how aggressively we penalize divergences in their predictions. This term encourages $\\mathcal{M}_S$ to *anticipate* situations in which its distribution would differ from $\\mathcal{M}_L$, thereby sharpening the deferral signal when such discrepancies are likely.\n\n---\n\n``4) Connection between Gatekeeper and calibration.``\n\n\nGatekeeper is *purpose-built for deferral*: the KL term deliberately flattens the softmax on examples the student misclassifies, **inducing *structured underconfidence* on precisely those inputs that should be routed to the large model**, while the cross-entropy term leaves already-correct examples confident. This selective re-weighting has two key effects:\n- **It “unlearns” brittle patterns on hard inputs**, freeing capacity to sharpen predictions on easy ones. \n- **It creates a much clearer separation between the confidence scores of correct and incorrect predictions**. Although this yields a model that is globally underconfident, it improves the *local calibration* that matters for risk–deferral: conditional on a confidence threshold, the realized accuracy tracks the nominal threshold far more closely, enabling a near-optimal deferral curve without any additional routing head.\n\nEmpirically, we observe lower overconfidence on wrongly classified examples and substantially higher deferral performance than both the untuned baseline and prior cascading methods, confirming that this targeted form of calibration is precisely what a small local model needs to decide when to defer.\n\n---\n\n**We hope that our rebuttal has clarified the reviewer’s questions and concerns and that the reviewer considers raising their individual score(s). We are happy to further engage with them during the discussion period should further questions arise!**"}, {"review_id": "6gEY4zcAuZ", "reviewer": "Reviewer_ri54", "summary": "The paper introduces the “GATEKEEPER” loss, which fine-tunes a lightweight model to escalate hard-to-label samples to a larger, more capable model. The objective is to balance overall accuracy while minimizing queries to the computationally heavier model. Experiments validate the effectiveness of GATEKEEPER across classification, language, and vision-language tasks.", "questions": "Q1. Would GATEKEEPER scale beyond two stages? \n In other words, could you train a $k$-level cascade: $\\mathcal{M}_1 \\rightarrow \\mathcal{M}_2 \\rightarrow … \\rightarrow \\mathcal{M}_K$ with increasingly large models and apply the loss hierarchically? \n\n\nQ2. Could you also explain why you regard the assumption that “$\\mathcal{M_S}$ makes the same mistakes as $\\mathcal{M}_L$” as requisite for your method development? L139-L140", "limitations": "Oracle scenario omitted. \n \nThe current study always treats a neural model as the “large” backend. An informative upper-bound experiment would replace the backend with an Oracle (e.g., ground-truth label or a human expert) and treat today’s large model as the new $\\mathcal{M}_S$.  \n\nThis would strengthen the study by revealing how much of the joint-accuracy gain stems from better routing versus from the back-end’s inherent power and would inform human-in-the-loop deployments of autonomous systems where server cost is secondary to trusting the system’s outputs. This should be feasible with a simple simulation and no extra data collection.", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 4, "significance": 2, "originality": 3, "strengths_and_weaknesses": "Strengths \n- Well-written paper and relevant appendices. \n- The method supports heterogeneous architectures in the model pair. \n- Comprehensive related work and a well-positioned contribution. \n- Extensive experiments. \n- Flexible setup via the α-parameter: one can tune the system to query the large model for high-risk applications (e.g., medical) and lean more toward the edge model for lower-risk scenarios (e.g., traffic). \n\nWeaknesses: \n\n1. Several assumptions may not hold in real settings and are not verified: \n  - The same mistakes made by $\\mathcal{M}_L$ would be made by $\\mathcal{M}_S$), especially with heterogeneous architectures. \n  - The method assumes Acc($\\mathcal{M}_L$) > Acc($\\mathcal{M}_S$). This can fail after fine-tuning, particularly on a restricted test domain, where studies show Acc($\\mathcal{M}_S$ fine-tuned on $X$) > Acc($\\mathcal{M}_L$ on $X$). \n\n2. It is unclear whether higher joint accuracy results from a genuine deferral process or simply from fine-tuning on the test domain. \n\n3. The study does not assess the case where $\\mathcal{M}_S$ is fine-tuned on pseudo-labels from $\\mathcal{M}_L$ in a test-time adaptation setting.", "was_revised": true, "final_justification": "I increased my rating to \"accept\" as my major concerns have been adressed during the rebuttal.", "rebuttal": "We thank the reviewer for their thoughtful and positive assessment, including their **recognition of our clear writing, comprehensive experiments, and practical, flexible design**. They raise important **concerns about real-world assumptions—such as model agreement, post-fine-tuning accuracy, and pseudo-labeling**—which we address below:\n\n---\n\n``W.1 + Q.2) Why is the assumption “M_S makes the same mistakes as M_L” necessary?.``\n\nThe premise that every error made by the large expert model $\\mathcal{M}_L$ is also made by the small model $\\mathcal{M}_S$ **follows directly from well-documented scaling-law behavior**: when two models are trained on comparable data with the same objective, predictive accuracy increases monotonically with parameter count and compute budget. Empirically, larger vision transformers and language models dominate their smaller counterparts across virtually every benchmark, a pattern quantified by [1] and many subsequent works. Embracing this monotonicity gives us two critical benefits. \n1) **It lets us treat $\\mathcal{M}_L$ as an *upper envelope* of capability**: deferring from $\\mathcal{M}_S$ to $\\mathcal{M}_L$ can only maintain—or improve—performance, never degrade it. Without that guarantee, the deferral problem becomes ill-posed because sending an input to $\\mathcal{M}_L$ could occasionally lower accuracy—forcing the gate to reason about *relative* per-example competence instead of simple uncertainty in $\\mathcal{M}_S$.\n2) **The assumption is baked into our evaluation protocol and the definition of the ideal deferral curve against which we compute the deferral-performance score $s_d$**. That oracle curve assumes perfect identification of samples misclassified by $\\mathcal{M}_S$ *and* subsequent perfect correction by $\\mathcal{M}_L$; it therefore constitutes a true upper bound only when $\\mathcal{M}_L$’s error set is a subset of $\\mathcal{M}_S$’s. If that subset relation does not hold, any metric normalised by the oracle curve would be inflated in expectation, obscuring real differences between methods. \n\nIn short, the **assumption is not an arbitrary convenience**—it is a principled reflection of established scaling trends that (i) shields deferral from pathological failure cases and (ii) grounds our performance metric in a meaningful, attainable ceiling. We are happy to expand on this in our updated paper draft.\n\n[1] Kaplan, Jared, et al. \"Scaling laws for neural language models.\" arXiv preprint arXiv:2001.08361 (2020).\n\n---\n\n``W.2) It is possible that Acc(M_S fine-tuned on X) > Acc(M_L on X). Also, it is unclear whether higher joint accuracy results from a good deferral or simply from fine-tuning on the test domain.``\n\nThis is a valid concern. However, in our experimentation, **we ensure that both the small and large models are always trained or fine-tuned on the same training split of the target dataset before we apply Gatekeeper to the small model only**. Our fine-tuning stage with Gatekeeper alters only the confidence profile of the small model—its class-probability outputs—not the underlying data distribution, and the large model remains completely frozen. Consequently, whenever joint accuracy improves while the small model’s stand-alone accuracy remains constant or even drops, **the only mechanism that can explain the gain is more effective routing of hard examples to the large model** (which is precisely what we observe; see the rightmost panel in Figures 4 and 5 showing noticeably better performance for $\\mathcal{M}_L$ over $\\mathcal{M}_S$). Moreover, **all baselines are tuned under the same protocol**, so the relative improvements we report isolate the contribution of the deferral strategy itself, not differences in domain exposure or training budget. We have clarified this in our experimental setup paragraphs in Section 4.\n\n---\n\n\n``W.3) What if M_S is fine-tuned on pseudo-labels from M_L?``\n\nWe agree that one could continually tune $\\mathcal{M}_S$ on pseudo-labels from $\\mathcal{M}_L$, but **this falls outside the deployment regime we study**. In our setting **(i)** $\\mathcal{M}_L$ is an immutable resource beyond the control of the user (after initial fine-tuning); and **(ii)** the deferral gate must decide whether to call $\\mathcal{M}_L$ without first querying it, because every additional call already incurs the latency and cost we seek to avoid. Streaming logits or hard labels from $\\mathcal{M}_L$ for ongoing adaptation would therefore nullify the efficiency gains of deferral. **Gatekeeper is, however, agnostic to how the initial $\\mathcal{M}_S$ was obtained**: if practitioners start with an $\\mathcal{M}_S$ distilled from $\\mathcal{M}_L$ (like we often do in our experiments), they can still apply Gatekeeper to sharpen the confidence gap needed for reliable deferral—without any extra expert queries at inference time. Fully exploring that hybrid pipeline is an interesting avenue for future work, but it involves distinct engineering and economic trade-offs and would require a broader experimental setup than what fits in the present study.\n\n---\n\n``Q.1) Scaling Gatekeeper to multi-level cascades M_1 → M_2 → … → M_K.``\n\nThis is indeed an interesting extension. **Gatekeeper makes no explicit assumption on the number of cascade levels**—it is inherently modular and agnostic to model architecture. In practice, one can **fine‐tune each model $\\mathcal{M}_i$ in the chain independently using the Gatekeeper loss with a stage‐specific hyperparameter $\\alpha_i$** that governs the balance between confident acceptance and deferral. A natural hierarchical training procedure is to first apply Gatekeeper to $\\mathcal{M}_1$ to calibrate its confidence, then apply it to $\\mathcal{M}_2$ on the subset of inputs deferred by $\\mathcal{M}_1$, and so on through $\\mathcal{M}_K$. This approach scales linearly in the number of levels and incurs only the familiar overhead of per‐model fine‐tuning—no joint optimization or additional routing networks are required. We therefore expect that deeper cascades will offer flexible trade‐offs between computational cost and predictive accuracy across multiple model tiers, and we leave a full empirical evaluation of $K>2$ cascades to future work.\n\n---\n\n``Replace big model with a perfect label oracle.``\n\nWe appreciate the reviewer’s interest in an oracle-backed simulation; however, **replacing the back-end with ground-truth labels would shift the task from *cascading* to *selective classification***, where the objective is abstention instead of deferral. Our goal in this paper is noticeably different: we study how to *optimize the small model’s confidence* so that it routes inputs judiciously when the only fallback is a *fixed, imperfect* but highly capable model—exactly the constraint faced in real-world cascades where API calls or server latency make further human verification impractical. To keep the exposition focused, we therefore do *not* evaluate deferral quality via the risk–coverage framework but instead **resort to the normalized area metric $s_d$ (Eq. 7)**, which already benchmarks each method against its own *ideal* routing curve; thus the performance gap we report isolates routing effectiveness from raw back-end accuracy. **Performance evaluation of selective classifiers [1,2,3,4] as well as oracle-style upper bounds for rejection [5,6] on the other hand have been thoroughly explored in prior SC work**, so duplicating them here would add length without clarifying our contribution. Nevertheless, we will note this distinction in a separate appendix section in the revision and cite the relevant SC papers to guide readers interested in SC-related oracle analyses.\n\n[1] Geifman, Yonatan, and Ran El-Yaniv. \"Selectivenet: A deep neural network with an integrated reject option.\" International conference on machine learning. PMLR, 2019.\n\n[2] Huang, Lang, Chao Zhang, and Hongyang Zhang. \"Self-adaptive training: beyond empirical risk minimization.\" Advances in neural information processing systems 33 (2020): 19365-19376.\n\n[3] Liu, Ziyin, et al. \"Deep gamblers: Learning to abstain with portfolio theory.\" Advances in Neural Information Processing Systems 32 (2019).\n\n[4] Jiang, Heinrich, et al. \"To trust or not to trust a classifier.\" Advances in neural information processing systems 31 (2018).\n\n[5] Geifman, Yonatan, Guy Uziel, and Ran El-Yaniv. \"Bias-reduced uncertainty estimation for deep neural classifiers.\" arXiv preprint arXiv:1805.08206 (2018).\n\n[6] Rabanser, Stephan, et al. \"Training private models that know what they don’t know.\" Advances in Neural Information Processing Systems 36 (2023): 53711-53727.\n\n---\n\n``Re-label Stage 1 as a prerequisite``\n\nThis is a good suggestion; indeed Stage 1 is not really a part of the Gatekeeper tuning process itself! We have restructured this part of the paper to **state Stage 1 as a preliminary setup step**. With this change Stage 2 (now Stage 1) now corresponds to the Gatekeeper training stage and Stage 3 (now Stage 2) now corresponds to the deployment stage.\n\n---\n\n**We hope that our rebuttal has addressed the reviewer’s questions and concerns and that the reviewer considers raising their individual and overall score(s). Our work has already noticeably improved thanks to their feedback. We are happy to further engage with them during the discussion period!**"}, {"review_id": "nwMW6bY4iL", "reviewer": "Reviewer_Eqgu", "summary": "The authors introduce Gatekeeper, a novel loss function designed to improve the performance of small models in model cascade systems by tuning their confidence calibration. The method enables the small model to identify when to defer to a more capable large model, aiming to optimize deferral performance without modifying the architecture or requiring auxiliary deferral modules. Gatekeeper encourages high confidence on correct predictions and uniform uncertainty on incorrect ones through a linear combination of a cross-entropy and a KL-divergence loss. The approach is validated on a wide range of tasks (image classification, language modeling, and vision-language tasks ) and architectures (encoder-only, decoder-only, and encoder-decoder models). The results show notable improvements in deferral metrics over baselines.", "questions": "See weakness in loss implementation", "limitations": "No specific potential negative societal impact", "rating": 3, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 2, "significance": 2, "originality": 2, "strengths_and_weaknesses": "### **Strengths**\n\n1. **Novel adaptation of a hybrid loss for cascading systems**: The introduction of a loss function that penalizes incorrect predictions with uniform confidence is novel in the context of model cascades.\n2. **Strong empirical validation**: The method is tested across diverse tasks and model architectures, demonstrating robust improvements in deferral performance.\n\n### **Weaknesses**\n\n**1. Limited exploration of alternative confidence metrics and loss formulations:**\nWhile the authors frame the use of the Gatekeeper loss as a novel contribution to the cascading setup, the core idea—using confidence calibration to distinguish correct from incorrect predictions—has been extensively studied in the broader uncertainty estimation literature. As such, the novelty of adapting this loss to cascaded inference may be limited. To strengthen the empirical contribution, the authors should include a broader comparison against alternative loss functions and calibration techniques from prior work. I highly suggest authors to do a more through literature search on these domains (e.g., start with [A][B][C] papers and check papers that cite these works) to better contextualize the proposed approach and clarify its advantages in model cascading setup.\n\n**2. Ambiguity in loss computation details:** The paper lacks clarity regarding the implementation of the Gatekeeper loss (Eq. 1, Lines 172–193). It remains unclear how the model determines which samples are incorrect predictions, and whether the loss terms $L_{\\text{corr}}$ and $L_{\\text{incorr}}$ are computed in a single forward pass or through separate evaluation and sampling steps. For instance, do the authors generate separate batches for correct and incorrect samples after a full forward pass, or do they compute both terms within each mini-batch during training? Without a clear explanation, it is difficult to reproduce the training pipeline or assess the computational overhead. This should be clarified with an explicit algorithmic description or pseudocode.", "was_revised": true, "final_justification": "I appreciate the authors' efforts in the rebuttal. The further comparison with previous work and explanation of the implementation details were helpful. I have increased my score accordingly.", "rebuttal": "We thank the reviewer for recognizing the **novelty of our loss in the cascading setting and the strength of our empirical results**. We appreciate the constructive **suggestions regarding broader comparisons and implementation clarity**, and address both issues below:\n\n---\n\n``1)  Limited exploration of alternative confidence metrics and loss formulations.``\n\nWe appreciate the reviewer’s observation that confidence calibration has a long history and agree that situating Gatekeeper more explicitly within that literature will clarify its added value for cascaded inference. However, **the reviewer did unfortunately not specify what the listed references [A][B][C] are** (paper names and authors are not provided in the review). As a result, we were unable to accommodate any additional experimentation during the initial rebuttal. Nevertheless, **we have added an additional section in the Appendix (§B.4) that now surveys canonical calibration objectives**—including temperature scaling, focal loss, confidence penalty, and outlier-exposure style uniformity losses—and explains how Gatekeeper extends these ideas by *jointly* sharpening correct predictions *and* flattening incorrect ones so that a single scalar confidence score yields a clean separation for deferral. These additions, together with our existing comparisons to cascading methods [1,2], and prompting baselines [3], demonstrate that Gatekeeper provides a simple yet consistently stronger solution for confidence-based deferral, thereby addressing the reviewer’s concern about novelty and empirical breadth.\n\nWe include a shortened version of our new appendix section below:\n\n> ### C.4 Gatekeeper in the Context of Canonical Calibration Objectives\n>\n>**Motivation.**  Section 4 showed that the Gatekeeper loss improves *deferral performance* with minimal implementation effort. Because many calibration objectives also manipulate confidence, we now position Gatekeeper relative to four widely–used losses.\n>\n> **Canonical calibration objectives.**  Let $p_{\\theta}(y\\mid\\mathbf{x})$ denote the softmax (or token) distribution predicted by a model with parameters $\\theta$, and let $y^\\star$ be the ground-truth label.\n>\n>1. **Temperature scaling** [4]:  \n>   Applies a single scalar $T{>}0$ at test time: $p_T(y\\mid\\mathbf{x}) \\propto \\exp(z_c/T)$. It preserves the rank ordering and therefore *cannot tighten* the ranking-based risk–coverage curve, but can improve threshold-based acceptance.\n>\n>2. **Focal loss** [5]:  \n>   Adds a down-weighting factor to easy examples, $\\text{FL}(p,y^\\star)=-(1-p_{y^\\star})^\\gamma \\log p_{y^\\star}$ with $\\gamma \\in [0,\\infty)$. It improves class imbalance calibration but does not explicitly penalize over-confidence on *incorrect* samples.\n>\n>3. **Confidence penalty** [6]:  \n>   Regularises high-entropy predictions through $\\text{CE}(p,y^\\star) + \\lambda \\mathcal{H}(p)$. While it flattens *all* distributions, it does not distinguish between correct and incorrect cases.\n>\n>4. **Outlier exposure (OE)** [7]:  \n>   Adds a KL-to-uniform loss on auxiliary **OOD** data, mirroring the second term of Gatekeeper but only on outliers, not in-distribution misclassifications.\n>\n>**How Gatekeeper differs.**  **Gatekeeper** combines *two complementary gradients*:\n>- (i) a standard CE term on correct predictions with an instance-level mask, thereby **sharpening** those logits;\n>- (ii) a KL-to-uniform term on *incorrect* predictions, **flattening** their confidence.\n>\n>This asymmetric design forces the scalar summary $g(\\mathbf{x}) = \\max_c p_\\theta(y = c \\mid \\mathbf{x})$ (or token-entropy) to separate correct from incorrect points *without* requiring additional heads, auxiliary datasets, or test-time tuning.\n\n\n**References**:\n\n[1]: Narasimhan, Harikrishna, et al. \"Post-hoc estimators for learning to defer to an expert.\" Advances in Neural Information Processing Systems 35 (2022): 29292-29304.\n\n[2]: Gupta, Neha, et al. \"Language model cascades: Token-level uncertainty and beyond.\" ICLR 2024.\n\n[3]: Kadavath, Saurav, et al. \"Language models (mostly) know what they know.\" arXiv preprint arXiv:2207.05221 (2022).\n\n[4]: Guo, Chuan, et al. \"On calibration of modern neural networks.\" ICML 2017.\n\n[5]: Lin, Tsung-Yi, et al. \"Focal loss for dense object detection.\" ICCV 2017.\n\n[6]: Pereyra, Gabriel, et al. \"Regularizing neural networks by penalizing confident output distributions.\" arXiv preprint arXiv:1701.06548 (2017).\n\n[7]: Hendrycks, Dan, Mantas Mazeika, and Thomas Dietterich. \"Deep anomaly detection with outlier exposure.\" ICLR 2019.\n\n---\n\n``2) Ambiguity in loss computation details.``\n\nWe agree with the reviewer that our original paper draft did not explain the loss computation in sufficient detail. As a result, **we have added the following clarifying paragraph to our current draft**.\n\n>**Practical computation of the Gatekeeper loss**. We compute the Gatekeeper loss *once per mini-batch* within the standard training loop—no auxiliary passes or data-set re-shuffling are required. Given a mini-batch $B=\\\\{ (x_i, y_i)\\\\}_{i=1}^{N}$, we perform a single forward pass through $\\mathcal{M}_S$. This allows us to obtain \n>\n>- a logit vector $z_i=f_\\theta(\\mathbf{x}_i)$ and \n>- predicted labels $\\hat{y_{i}} = \\arg \\max_{c} z_{i,c}$.\n>\n>Two binary masks $m^{\\text{corr}}_i=I\\\\{y_i=\\hat{y}_i\\\\}$ and $m^{\\text{incorr}}_i= \\neg m^{\\text{corr}}_i$ are computed *on-the-fly*. The hybrid loss is then assembled in a fully vectorized manner:\n>\n>$$\n>\\mathcal{L}_\\text{corr} = \\frac{1}{N}\\sum_i m^{\\text{corr}}_i \\operatorname{CE} (p_i,y_i) \n>$$\n>\n>$$\n>\\mathcal{L}_\\text{incorr} = \\frac{1}{N}\\sum_i m^{\\text{incorr}}_i \\operatorname{KL} (p_i | | \\mathcal{U}) \n>$$\n>\n>where $p_i=\\operatorname{softmax}(z_i)$. Because both masks and losses are computed inside the same tensor graph, back-propagation incurs only the cost of $\\mathcal{O}(N \\times C)$ element-wise operations—identical to a vanilla cross-entropy step. This single-pass design keeps the computational overhead negligible while guaranteeing that in the full loss $\\mathcal{L} = \\alpha \\mathcal{L}\\_\\text{corr} + (1 - \\alpha) \\mathcal{L}\\_\\text{incorr}$ every sample contributes to either $\\mathcal{L}\\_\\text{corr}$ or $\\mathcal{L}\\_\\text{incorr}$ in the same optimization step.\n\n---\n\n**We hope that our rebuttal has addressed the reviewer’s questions and concerns and that the reviewer considers raising their individual and overall score(s). Our work has already noticeably improved thanks to their feedback. We are happy to further engage with them during the discussion period!**"}, {"review_id": "FaDed0ioPW", "reviewer": "Reviewer_aLX9", "summary": "This paper aims to solve the problem in model cascades, which often deploy small models along with large models and rely on routing and deferral mechanisms to offload requests to small models and only escalate the requests to large models if the small models are not confident enough. This paper tries to tame the tradeoff between model accuracy and deferral accuracy across many model architectures and workloads. The core of this paper is a general-purpose loss function, Gatekeeper, which helps to fine-tune the small model to calibrate the small model’s confidence to improve the efficacy of model cascades.", "questions": "Please also refer to my comments above for my full review.\n1. Is it correct that the tuning of $\\alpha$ and switching of the large model will require a new round of fine-tuning?\n2. Does your new loss function break the performance of small models, especially for decoder-only models? \n3. How does this paper compare with related works\n    3.1 Compare with speculative decoding in terms of latency savings for generative models.\n    3.2 Compare with other methods (e.g., early exit) in terms of latency savings for classification workloads.", "limitations": "Yes", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 3, "significance": 2, "originality": 2, "strengths_and_weaknesses": "**Strengths**  \n1. The paper is well written and easy to follow.  \n2. The evaluation covers diverse model architectures (encoder only, decoder only, and encoder-decoder models) and tasks (classification and generation).\n\n**Weaknesses**  \n1. As the paper indicates, the core contribution of this paper, the loss function, which has been studied and applied in previous works. I don’t think there are unique challenges in model cascades.  \n2. There is a scalability issue regarding this method. Though it seems that only the small model would require extra fine-tuning. However, the labels of correct and incorrect are dependent on the specific large model. Namely, it will require a new round of fine-tuning if I switch to a new large model, which can be quite expensive, especially for generative models. Also, a very important knob here is $\\alpha$, however, tuning that will require redo the finetuning phase. It sounds very suboptimal compared with the traditional methods for model cascades, which simply tune the threshold to control the deferrals.  \n3. The loss function seems too strong for fine-tuning small models, especially in generative models. As it would enforce the exact token matching, which might make it hard for the small model to converge?  \n4. Missing comparison to related works.  \n    4.1. For classification-related evaluation (Fig. 3), the baseline seems too weak. What’s the performance if I use a similar confidence calibration method in Fig. 5?  \n    4.2. I think the biggest motivation for the model cascade is the potential latency savings. How does this paper compare with speculative decoding in terms of latency savings? Note that speculative decoding can also apply fine-tuning to align the small model based on the large model.\n    \n    4.3 Other related works can be considered, for example, early exiting.", "was_revised": true, "final_justification": "Most of my conceptual questions have been resolved by the authors. I only have minor issues with the finetune related overheads and how to handle the inputs that can't be solved by both small and large models.", "rebuttal": "We thank the reviewer for their thoughtful feedback and for highlighting the **clarity of our writing and the breadth of our evaluation**. We address the reviewer’s concerns regarding **novelty, scalability, and practical comparisons** below.\n\n---\n\n\n``1) Loss function has been studied and applied in previous works.``\n\nWhile we acknowledge that a similar loss has been proposed for OOD detection (e.g., Outlier Exposure), **our work repurposes this idea in a novel context**—namely, for deferral in model cascades. Unlike OOD detection, which typically relies on auxiliary datasets and disjoint distributions, deferral requires distinguishing between easy and hard in-distribution examples without additional supervision. Applying this loss to a standalone model would reduce utility by encouraging the model to \"unlearn\" difficult examples. However, in a cascade, **$\\mathcal{M}_L$ recovers this lost performance, and $\\mathcal{M}_S$ gains sharper confidence separation**, improving deferral quality without modifying its architecture or requiring access to $\\mathcal{M}_L$. This use of the loss is specific to the cascade setting and enables practical gains in utility, latency, and deferral performance that prior work does not address. We have clarified this point in the paper.\n\n---\n\n``2) Scalability issues: labels of correct/incorrect depend on the large model; tuning \\alpha requires redoing finetuning; suboptimal setup compared to traditional model cascades reliant on threshold tuning.``\n\nWe thank the reviewer for raising the question of scalability. There are in fact two points of misunderstanding we would like to clarify. \n\n1) **The “correct” versus “incorrect” labels used during fine-tuning come solely from the small model** evaluated on a labeled calibration dataset—not from the large model’s outputs (see Eqn. 1). In practice, one simply runs the small model on the fine-tuning set, compares its predictions to the true labels, and applies our Gatekeeper objective accordingly. Because this process does not depend on any remote model, **no additional fine-tuning is needed when you replace or upgrade the large model**. \n\n2) The trade-off parameter $\\alpha$ controls only how strongly the loss encourages low confidence on misclassified examples versus high confidence on correct ones; it **does not affect the gating threshold used at inference time**. As detailed in Stage 3 (see Eqn. 3), deferral decisions remain based on a separate confidence or entropy threshold, which can be tuned cheaply at deployment without retraining. We agree, however, that if one wishes to explore a different $\\alpha$ to shift the calibration–accuracy trade-off, a new fine-tuning run of the small model would be required. As we note in lines 183–193, this flexibility enables practitioners to select the desired balance between conservatism in uncertain regions and decisiveness on easy examples, tailored to the cost and performance constraints of their specific application.\n\n---\n\n``3) Loss function seems too aggressive for fine-tuning small models in generative setups.`` \n\nWe agree with the reviewer that **Gatekeeper can be overly aggressive in generative settings where multiple token sequences may convey the same meaning**. As already noted in our limitations (lines 366–369), relying on exact token matches may over-penalize semantically valid outputs. While recent approaches such as **semantic entropy [1] aim to address this, they typically require querying a larger model to judge semantic correctness**—effectively constituting a deferral, which defeats the purpose of our cascade setup (see lines 131–137 where we already elaborate on this point). To the best of our knowledge, there is currently no lightweight alternative that enables small models to assess semantic correctness without external assistance. Nevertheless, despite this limitation, **our results in Section 4.3 (e.g., Figure 7) show consistent improvements in deferral performance for both LM and VLM tasks**, suggesting that even token-level confidence calibration can meaningfully improve reliability in generative setups.\n\n[1] Kuhn, Lorenz, Yarin Gal, and Sebastian Farquhar. \"Semantic uncertainty: Linguistic invariances for uncertainty estimation in natural language generation.\" ICLR 2023.\n\n---\n\n``4.1) For classification-related evaluation (Fig. 3), the baseline seems too weak. What’s the performance if I use a similar confidence calibration method in Fig. 5?``\n\nWe thank the reviewer for the question but believe there may be a misunderstanding: **Figure 3 is purely illustrative and does not report any empirical results**—it serves to explain the evaluation metric used throughout the paper. The actual classification and language model experiments appear in **Figures 4 and 5, which both include an untuned baseline and strong cascading methods [1,2]. Figure 5 additionally includes prompting-based baselines [3] that are only applicable to language models and cannot be transferred to image classification tasks**. We kindly ask the reviewer to follow up with us in case we misunderstood their concern.\n\n[1]: Narasimhan, Harikrishna, et al. \"Post-hoc estimators for learning to defer to an expert.\" NeurIPS 2022.\n\n[2]: Gupta, Neha, et al. \"Language model cascades: Token-level uncertainty and beyond.\" ICLR 2024.\n\n[3]: Kadavath, Saurav, et al. \"Language models (mostly) know what they know.\" arXiv preprint arXiv:2207.05221 (2022).\n\n---\n\n``4.2) How does this paper compare with speculative decoding in terms of latency savings?``\n\nWe appreciate the reviewer’s interest in speculative decoding, yet we respectfully believe that a direct latency comparison would be misleading:\n\n1) Speculative decoding accelerates every input by asking the large model to verify draft tokens that the small model proposes, so **the large model is still invoked on all inputs**. Our cascade, by contrast, aims to avoid calling the large model on inputs the small model can already solve. Consequently, speculative decoding optimizes in-place token-level latency, while our method optimizes end-to-end compute and monetary cost at the request level. Because **the two techniques improve fundamentally different bottlenecks**, a like-for-like timing chart would conflate distinct operating points.\n\n2) As noted in the introduction and our related work section, **the two approaches can be composed**: after our deferral gate decides to consult the large model, one could still decode that portion of the input speculatively (see [1] for an example of such hybrids). The proper question is therefore not *“Which one is faster?”* but “How much additional speed-up does deferral enable on top of speculative decoding?” Measuring that joint benefit would require re-engineering both systems into a single pipeline, which is orthogonal to the paper’s core contribution—namely, a loss that teaches the small model when to defer.\n\nWe have added an extended discussion on these notable differences/similarities as part of Appendix §B.1.\n\n[1] Narasimhan, Harikrishna, et al. \"Faster cascades via speculative decoding.\" ICLR 2025.\n\n---\n\n``4.3) Relationship to early exiting.``\n\nWe appreciate the suggestion to include early-exit experimentation, but we believe that—similar to the speculative decoding angle—an empirical comparison with early–exiting methods would add little insight for three reasons:\n\n1) Early–exiting aims to *skip* the remaining layers of a single network once an intermediate classifier is sufficiently confident, thereby offering a **layer–level latency–accuracy trade-off** within *one* model. **Our work tackles the fundamentally different problem of *model–level* deferral between a *small* and a *large* model**. The goal is to keep the large model entirely idle for easy requests and to invoke it only when needed. Hence our primary measure of success is *joint accuracy versus cross-model compute budget*, not marginal delay per layer inside a fixed backbone. Adding an early-exit baseline would therefore **conflate two orthogonal trade-offs** and obscure the key question we study: *How well can a calibrated small model decide when to hand off to a more capable expert?*\n\n2) Early–exit usually presupposes **white-box control over the entire network architecture** so that branch classifiers can be inserted and jointly trained (e.g., BranchyNet, Adaptive Neural Networks). Our cascade setting explicitly targets **heterogeneous or API-based experts**—large language models, vision–language models, and so forth—that cannot be modified. Early–exits cannot be applied in such black-box situations, whereas our \\loss fine-tuning remains feasible. Moreover, early–exiting has seen limited success for *autoregressive* sequence generation, where every token depends on the full hidden state; in contrast, our experiments span both classification and open-ended generation tasks.\n\n3) Because early–exit and cascading operate at different granularities, they are **complementary rather than competing** techniques. One could in principle insert early exits *inside* the small model and still rely on calibrated deferral to the large model—yielding a three-level system (exit-1 → exit-2 → large model). Evaluating that combined design is a compelling direction for future engineering work, but it lies beyond the scope of our current study, whose contribution is a general-purpose confidence-tuning loss independent of architectural changes.\n\nAgain, we have added an extended discussion to Appendix §B.1.\n\n[1] Teerapittayanon, Surat, et al. \"Branchynet: Fast inference via early exiting from deep neural networks.\" ICPR 2016.\n\n[2] Bolukbasi, Tolga, et al. \"Adaptive neural networks for efficient inference.\" ICML 2017.\n\n---\n\n**We hope that our rebuttal has clarified the reviewer’s questions and concerns and that the reviewer considers raising their individual and overall score(s). We are happy to further engage with them during the discussion period should further questions arise!**"}]

Author rebuttal:

Meta-review: The submission proposes a novel gatekeeper objective for a smaller model to decide when to route harder examples to a larger model, to achieve a balance between accuracy and computational efficiency. 5Nnu questions the intuition of the proposed gatekeeper loss and its motivation relative to calibration, both of which are addressed satisfactorily by the author response. ri54 raises the concern that there might be conditions where M_L makes errors where M_S does not, to which the authors respond by diluting the dominance assumption to be that M_L dominates M_s with high probability. Both the reviewer and AC agree with this resolution. aLX9 raises limitations which might be addressed by a three-way decision across local accept, local reject and route to remote, which the authors analyze in the feedback but consider difficult to implement, which is convincing. aLX9 also requires discussion with respect to speculative decoding and early exiting, which are adequately included in the author response. Eqgu questions the contributions given existing literature on calibration, which is answered sufficiently well in the response and acknowledged as such by the reviewer. Overall, the AC agrees with the reviewer majority that the submission proposes an interesting method, with thorough analysis and convincing experiments. Thus, the submission is recommended for acceptance. The authors are encouraged to include discussions on limitations and clarifications on intuition for the proposed objective in the final version.
Decision: Accept (poster)