Skip to content

Gatekeeper: Improving Model Cascades Through Confidence Tuning

🕒 Published (v1): 2025-01-01 · Source: NeurIPS · Venue: NeurIPS 2025 · link

Ask a follow-up

Open an assistant pre-loaded with this paper's context.

💬 Ask ChatGPT✦ Ask Claude

TL;DR

Gatekeeper is a fine-tuning loss function that calibrates a small model's confidence in model cascade setups: correct predictions are trained toward high confidence, incorrect ones toward uniform (low) confidence. By directly optimizing the correct/incorrect confidence separation, the small model learns when to handle a request itself versus defer to a larger model, improving deferral performance across image classification, LLMs, and VLMs without architectural changes.

Problem

Existing model cascade systems rely on off-the-shelf confidence signals from small models that are poorly calibrated for the deferral decision—small models are often overconfident on errors, causing excessive or incorrect deferrals and suboptimal compute/accuracy tradeoffs. Prior white-box methods either require auxiliary architectures or static data partitioning; no principled loss existed to train a small model to know what it does not know.

Method

Gatekeeper introduces a hybrid fine-tuning loss \(\mathcal{L} = \alpha \mathcal{L}_\text{corr} + (1-\alpha)\mathcal{L}_\text{incorr}\): - \(\mathcal{L}_\text{corr}\): standard cross-entropy applied only to correctly predicted samples (reinforces confident correct predictions). - \(\mathcal{L}_\text{incorr}\): KL divergence between the model's predicted distribution and the uniform distribution \(U\), applied only to incorrectly predicted samples (pushes incorrect predictions toward maximum entropy).

Binary masks (correct/incorrect) are computed on-the-fly in a single forward pass via \(\hat{y} = \arg\max_c z_c\), adding negligible overhead (\(O(N \times C)\)). The hyperparameter \(\alpha \in (0,1)\) controls the accuracy–deferral tradeoff. For token-based models (LLMs, VLMs), the loss extends per-token. At inference, classification models use max softmax; sequence models use negative predictive entropy averaged over tokens as the gating signal \(g(x)\); inputs below threshold \(\tau\) are deferred to the large model.

Key Contributions

  • Novel Gatekeeper loss that jointly optimizes correct-prediction confidence and incorrect-prediction uncertainty within a standard training loop, requiring no architectural changes.
  • Extension of the loss to token-level sequence models (decoder-only LLMs, encoder-decoder VLMs).
  • Explicit hyperparameter \(\alpha\) that provides a controllable knob trading small-model accuracy against deferral performance.
  • Evaluation across three architecture families (encoder-only CNNs/ResNets, decoder-only Gemma 2B/7B, encoder-decoder PaliGemma 1B/7B) and multiple task types including VLM captioning with LLM-based factuality scoring.
  • Two new deferral evaluation metrics: distributional overlap \(s_o\) (overlap of confidence PDFs for correct vs. incorrect predictions; lower is better) and deferral performance \(s_d\) (fraction of potential improvement over random deferral achieved; higher is better).

Results

  • Image classification (CIFAR-100, TinyImageNet200): Gatekeeper improves deferral performance \(s_d\) by 0.72× on CIFAR-100 and 2× on TinyImageNet200 over the untuned baseline.
  • Language modeling (ARC-e/c): 7× improvement on ARC-e and 10× on ARC-c over the baseline; outperforms Gupta et al. (2024) and uncertainty-prompting baselines ("Reduce Confidence", "Answer N") on correct/incorrect separation and deferral.
  • VLM classification (VQAv2, AI2D): Consistent improvement in distributional overlap and deferral performance at lower \(\alpha\).
  • VLM captioning (COCOcap, Screen2Words): Increased Pearson correlation between negative predictive entropy and factuality score; improved factuality-adapted deferral metric.
  • Baseline from Narasimhan et al. (2022) preserves accuracy better but requires an explicit estimate of the large model's correctness and architectural modifications; Gatekeeper achieves stronger deferral with simpler deployment.

Limitations

  • Only the small model is fine-tuned; co-adapting the large model is not explored.
  • Token-level incorrectness is judged by exact token match, not semantic equivalence—penalizing valid paraphrases and introducing suboptimal deferral signals for generative tasks.
  • Evaluation of LLM/VLM settings uses only one model family (Gemma/PaliGemma); generalization across diverse families is not established.
  • Only two-stage cascades are evaluated; multi-stage extension is left to an appendix without empirical results.
  • Factuality evaluation for captioning uses Gemini as an LLM judge, which is itself imperfect.
  • Lower \(\alpha\) values that maximize deferral performance cause the small model to "unlearn" harder examples, degrading raw accuracy on the full distribution.

Relevance to Vision-Language Models

Gatekeeper directly targets encoder-decoder VLMs (PaliGemma 1B/7B) for cascade-based efficient inference, demonstrating that confidence tuning transfers to multimodal sequence generation beyond classification. For researchers tracking VLMs, this addresses the practical deployment problem of when a cheap local VLM should invoke an expensive API model, offering a principled alternative to post-hoc confidence thresholding. The token-level extension and factuality-based evaluation pipeline (using an LLM judge for caption quality) are reusable patterns for VLM deferral evaluation. This connects to the broader line of work on efficient VLM inference alongside model compression and speculative decoding.