Title: Gatekeeper: Improving Model Cascades Through Confidence Tuning
Venue: NeurIPS 2025
Accepted: yes
Keywords: cascading

deferral

selective prediction

confidence tuning
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.


Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS AND WEAKNESSES:\n**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.\n\nQUESTIONS:\nPlease 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.\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "CLARITY": 3, "IMPACT": 2, "ORIGINALITY": 2, "DATE": "2025-06-15"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS AND WEAKNESSES:\n### **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.\n\nQUESTIONS:\nSee weakness in loss implementation\n\nLIMITATIONS:\nNo specific potential negative societal impact", "TITLE": "Official Review", "RECOMMENDATION": 3, "REVIEWER_CONFIDENCE": 3, "CLARITY": 2, "IMPACT": 2, "ORIGINALITY": 2, "DATE": "2025-07-03"}

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

comments: 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.

STRENGTHS AND WEAKNESSES:
Strengths:
- The paper is well written and easy to follow
- The proposed method is simple and only requires fine-tuning the small model $M_S$

Weaknesses:
- I don't fully understand the connection between Gatekeeper and calibration – see the questions below
- 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.

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?
- 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.
- 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$?
- 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
TITLE: Official Review
CLARITY: 3
IMPACT: 2
ORIGINALITY: 2
DATE: 2025-07-05