UCCI: Calibrated Uncertainty for Cost-Optimal LLM Cascade Routing

Technical Summary: UCCI – Calibrated Uncertainty for Cost-Optimal LLM Cascade Routing

Problem Statement
Large Language Models (LLMs) offer high performance on structured prediction tasks, but inference costs scale steeply with model size. Cascade inference—routing easy queries to a small model and escalating difficult ones to a large model—promises cost reduction. However, current deployed routers often rely on uncalibrated confidence scores (e.g., raw token entropy or margins) that are sensitive to prompt wording and workload shifts. Consequently, threshold tuning is often a manual, per-workload exercise that fails to generalize, creating a gap between the theoretical promise of cascades and production performance.

Methodology
The paper proposes UCCI (Uncertainty-Calibrated Cascade Inference), a routing framework that treats cascade design as a calibrated decision problem rather than a tuning exercise. The methodology proceeds in three steps:

1. Token-Level Uncertainty Extraction: For a generated sequence, the system calculates a scalar uncertainty score $u(x)$ based on the margin between the top-1 and top-2 next-token probabilities ($p_{t,1} - p_{t,2}$) across all tokens. This margin is aggregated to represent the overall uncertainty of the small model's output.
2. Isotonic Calibration: The raw uncertainty score $u(x)$ is not a probability. UCCI maps this score to a calibrated per-query error probability $\hat{p}(x)$ using isotonic regression on a held-out calibration set. This step transforms the miscalibrated signal into a reliable estimate of the probability that the small model's output is incorrect ($P(e(x)=1 | u(x))$).
3. Constrained Cost Minimization: A threshold policy $\pi_\theta$ is selected on a validation set. The threshold $\theta^*$ is chosen to minimize expected inference cost subject to a target accuracy constraint ($\tau$). The policy routes queries with $\hat{p}(x) \leq \theta$ to the small model and escalates those with $\hat{p}(x) > \theta$ to the large model.

Theoretical Contributions
The paper provides two key theoretical results:

• Optimality of Threshold Policies: Under three explicit assumptions (large model cost > small model cost; large model accuracy is invariant to which queries are escalated; and access to the true calibrated error probability), the authors prove that a threshold policy on the calibrated error probability is cost-optimal.
• Sample Complexity: The paper establishes that isotonic regression achieves an Expected Calibration Error (ECE) of $O(n^{-1/3})$ with $n$ calibration examples, providing a theoretical bound on the data efficiency of the calibration step.

Experimental Setup and Results
The framework was evaluated on a production Named Entity Recognition (NER) workload involving 75,000 manually labeled queries from an enterprise photo management system.

• Models: A 4B-parameter small model ($f_s$) and a 12B-parameter large model ($f_\ell$), both instruction-tuned and served on H100 GPUs.
• Cost Measurement: Costs were derived from measured end-to-end latency on H100 GPUs, yielding a cost ratio of $c_\ell/c_s \approx 3.02$.
• Baselines: The study compared UCCI against single-model inference, uncalibrated entropy thresholding, split-conformal routing, and a FrugalGPT-style learned threshold.

Key Findings:

• Cost Reduction: At a target micro-F1 of 0.91, UCCI reduced inference costs by 31% (95% CI: [27%, 35%]) compared to using the large model exclusively.
• Calibration Quality: Isotonic regression reduced the Expected Calibration Error (ECE) from 0.12 (uncalibrated) to 0.03.
• Performance vs. Baselines: At the same operating point (F1 = 0.91), UCCI achieved the lowest cost (2.08 normalized units) compared to entropy thresholding (2.31), conformal prediction (2.18), and FrugalGPT-style routing (2.24).
• Assumption Validation: The assumption that the large model's accuracy remains invariant to routing was empirically tested; the accuracy on escalated queries (0.928) was nearly identical to the global accuracy (0.932), supporting the theoretical framework.

Significance and Claims
The paper argues that the primary engineering value in cascade systems lies in calibrating the routing score rather than hand-tuning the final threshold. By converting a noisy token-margin signal into a probabilistic error estimate via a small isotonic fit, UCCI achieves near-optimal routing performance without complex learning-to-defer architectures or preference data.

The authors emphasize that their results are based on end-to-end routing using actual model outputs and measured hardware latency, avoiding the selection bias inherent in simulations based on global accuracies or nominal API prices. While the framework is task-agnostic, the authors note that the specific cost savings (31%) and calibration metrics are workload-specific. They acknowledge limitations regarding the invariance assumption in domains where "hard" queries for the small model are disproportionately hard for the large model, and note that static calibration may require adaptation for streaming deployments with distribution shifts.