Research map / Architectures and efficiency / RNNs and normalization
Gatekeeper: Improving Model Cascades Through Confidence Tuning
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes Gatekeeper, a correctness-aware loss that fine-tunes a small model to be confident on correct predictions and uniform on incorrect ones, improving deferral performance in model cascades across classification, LLM, and VLM setups.
Problem
Model cascades route inputs from a small model to a large model, but small models are often poorly calibrated, causing unnecessary deferrals or suboptimal resource use. The paper asks how to make the small model better at knowing what it does not know to maximize deferral performance.
Approach
A hybrid loss that applies cross-entropy only to correctly predicted samples/tokens and a KL-to-uniform penalty only to incorrectly predicted ones, with a trade-off weight alpha balancing small-model accuracy against deferral performance; requires no architectural changes and is applied in a single pass per mini-batch.
Key results
- Gatekeeper outperforms the untuned baseline by 0.72x on CIFAR-100 and 2x on TinyImageNet in deferral performance, and by 7x/10x on ARC-e/ARC-c.
- Lower alpha values give stronger deferral performance but reduce small-model accuracy, showing an explicit trade-off controlled by alpha.
- For captioning, Gatekeeper increases the correlation between negative predictive entropy and Gemini-judged factuality as alpha decreases.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| confidence-tune the small model with Gatekeeper | Small model in a cascade · CIFAR-100 | deferral performance: 0.72x | Helped |
| confidence-tune the small model with Gatekeeper | Small model in a cascade · TinyImageNet | deferral performance: 2x | Helped |
| confidence-tune the small model with Gatekeeper | Gemma2B cascade · ARC-e | deferral performance: 7x | Helped |
| confidence-tune the small model with Gatekeeper | Gemma2B cascade · ARC-c | deferral performance: 10x | Helped |
| decrease the Gatekeeper trade-off parameter α (ablation) | Gatekeeper-trained small models · Image classification tasks | deferral performance | Helped |
| decrease the Gatekeeper trade-off parameter α (ablation) | Gatekeeper-trained small models · Image classification tasks | small-model accuracy | Hurt |
| decrease the Gatekeeper trade-off parameter α (ablation) | Gatekeeper-trained language models · Language modeling tasks | deferral performance | Helped |
| decrease the Gatekeeper trade-off parameter α (ablation) | Gatekeeper-trained language models · Language modeling tasks | overall accuracy | Hurt |
| decrease the Gatekeeper trade-off parameter α (ablation) | Gatekeeper-trained PaliGemma1B · VLM classification tasks | deferral performance | Helped |
| decrease the Gatekeeper trade-off parameter α (ablation) | Gatekeeper-trained PaliGemma1B · Captioning tasks | deferral performance | Helped |
| use the Reduce Confidence prompt baseline | Gemma language-model cascade · Language modeling tasks | deferral performance | No clear effect |
| use the Answer “N” prompt baseline | Gemma language-model cascade · Language modeling tasks | deferral performance | No clear effect |
Surprising findings
- Prompting baselines (Reduce Confidence, Answer N) do not reliably improve deferral, and PaliGemma often failed to return responses to the modified prompts.
Limitations
- Only the small model is fine-tuned; the large model is assumed fixed.
- Token-level penalties on exact tokens may be overly aggressive for language tasks where different sequences share meaning.
- Limited evaluation across diverse model families for LLM and VLM settings; only two-stage cascades evaluated.
Experimental setup
Datasets: CIFAR-10 CIFAR-100 Food-101 TinyImageNet200 ARC-e ARC-c MMLU GSM8K VQAv2 AI2D COCO captions Screen2Words
Benchmarks: CIFAR-10 CIFAR-100 Food-101 TinyImageNet200 ARC-e ARC-c MMLU GSM8K VQAv2 AI2D COCO captions Screen2Words
Metrics: deferral performance s_d distributional overlap s_o small model accuracy AUROC Pearson correlation with factuality
Baselines: untuned baseline Reduce Confidence prompting Answer N prompting Okutomi-style learned deferral (cascade baseline for supervised learning) post-hoc confidence calibration method
Builds on: Outlier Exposure selective prediction model cascades
Closest papers by meaning
- Enhancing Large Language Model Performance with Gradient-Based Parameter Selection (AAAI 2025 · similarity 0.84)
- Online Cascade Learning for Efficient Inference over Streams (ICML 2024 · similarity 0.84)
- Small Batch Size Training for Language Models: When Vanilla SGD Works, and Why Gradient Accumulation is Wasteful (NeurIPS 2025 · similarity 0.84)
- LISA: Layerwise Importance Sampling for Memory-Efficient Large Language Model Fine-Tuning (NeurIPS 2024 · similarity 0.84)
- A Benchmark Study on Calibration (ICLR 2024 · similarity 0.83)
- Cross-model Control: Improving Multiple Large Language Models in One-time Training (NeurIPS 2024 · similarity 0.83)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.