PaperlayerView exact revision on arXiv

Source-linked AI summary

UCCI: Calibrated Uncertainty for Cost-Optimal LLM Cascade Routing

Varun Kotte

arXiv:2605.18796v1cs.LGcs.CL

TL;DR

LLM cascades can lower inference cost, but miscalibrated confidence scores and workload-specific threshold tuning limit routing. UCCI calibrates token-margin uncertainty into error probabilities before choosing an escalation threshold; on production NER, it cuts cost by 31% at micro-F1 = 0.91 versus large-model-only inference.

  • Problem

    Cascades promise lower inference cost, but routers often rely on uncalibrated confidence and thresholds that must be retuned across workloads.

    Source: §1 · p. 1

  • Method

    UCCI maps token-margin uncertainty to calibrated error probabilities and selects an escalation threshold through constrained cost minimization.

    Source: §4.1 · p. 2, §8 · p. 6

  • Results

    31% lower inference cost (95% CI: [27%, 35%]) at micro-F1 = 0.91 versus large-model-only inference on a 75,000-query production NER workload.

    Source: §1 · p. 1

  • Takeaways & Limitations

    A small isotonic fit converts a noisy token-margin signal into a routing score that is close to optimal in this setting.

    Source: §8 · p. 6

    Evaluation covers one photography NER domain, so the absolute results are workload-specific and broader validation across tasks, languages, and model families would strengthen the generalization claim.

    Source: §7 · p. 5

Abstract

from arXiv · show

LLM cascades and model routing promise lower inference cost by sending easy queries to a small model and escalating hard ones to a large model, but most deployed routers use uncalibrated confidence scores and require per-workload threshold tuning. We present UCCI, a calibration-first router that maps token-level margin uncertainty to a per-query error probability via isotonic regression and selects the escalation threshold by constrained cost minimization. Under three explicit assumptions, threshold policies on the calibrated score are cost-optimal, and isotonic calibration achieves O(n^{-1/3}) sample complexity for expected calibration error (ECE). On a production named entity recognition workload of 75,000 queries served by 4B and 12B instruction-tuned LLMs on H100 GPUs, UCCI cuts inference cost by 31% (95% CI: [27%, 35%]) at micro-F1 = 0.91 while reducing ECE from 0.12 to 0.03. At the same operating point, UCCI beats entropy thresholding, split-conformal routing, and a FrugalGPT-style learned threshold. All cascade results use end-to-end routing on actual model outputs and measured H100 latency, not simulated routing from global accuracies or nominal API prices.

1. Introduction

Cascades can reduce inference cost, but miscalibrated, prompt-sensitive confidence scores make routing thresholds difficult to transfer across workloads. UCCI calibrates uncertainty before thresholding and reports lower cost at a matched accuracy target on a production NER workload.

  • Threshold policies on calibrated error probabilities are cost-optimal under three explicit assumptions.Source: §1 · p. 1
  • Isotonic regression maps token-margin uncertainty to error probability, reducing expected calibration error from 0.12 to 0.03 on the workload.Source: §1 · p. 1
  • 31% lower cost (95% CI: [27%, 35%]) at micro-F1 = 0.91 versus large-model-only inference on a 75,000-query production NER workload.Routing uses actual model outputs at test time, and costs use measured H100 latency.Source: §1 · p. 1
  • At the same F1 = 0.91 operating point, UCCI beats entropy thresholding (cost 2.31), conformal-prediction routing (2.18), and a FrugalGPT-style learned threshold (2.24).Source: §1 · p. 1
  • UCCI’s operational claim is that calibrating a raw token-margin signal before thresholding produces a near-optimal routing score in this setting.The introduction contrasts calibration with hand-tuning the final threshold.Source: §1 · p. 1

2. Related Work

UCCI connects selective prediction and calibration to LLM structured outputs, where token-level uncertainty must be corrected and routing evaluated end to end. It differs from related routing approaches by using calibrated error probabilities and analyzing threshold optimality.

  • Selective classification and learning-to-defer study abstention or expert deferral under cost or risk constraints, typically assuming scalar confidence scores on classification heads.Source: §2 · p. 2
  • Calibration maps confidence to correctness probabilities; UCCI combines token-margin aggregation with isotonic regression to produce calibrated error probabilities for cost-constrained routing.The related-work passage reports ECE = 0.03 on the workload.Source: §2 · p. 2
  • Conformal selective-classification methods provide coverage guarantees; UCCI’s conformal-routing baseline uses raw token-margin uncertainty as its nonconformity score.Source: §2 · p. 2
  • Unlike FrugalGPT, Hybrid LLM, RouteLLM, and RouterBench, UCCI makes the routing score a calibrated error probability and proves when thresholding it is cost-optimal.The related work describes these systems as API routing, quality-aware routing, preference-data routing, and standardized heterogeneous-routing evaluation, respectively.Source: §2 · p. 2

3. Problem Formulation

The formulation considers a structured prediction task with small and large models, each with an accuracy and per-query cost. It seeks a routing policy that meets a target accuracy while minimizing expected cost.

  • For each input, the small and large models produce structured predictions whose accuracy is evaluated against ground truth, with per-query costs c_s and c_ℓ and c_ℓ > c_s.Source: §3 · p. 2
  • A routing policy π selects whether each input receives the small or large model’s prediction, determining the resulting cost and accuracy.Source: §3 · p. 2
  • For a target accuracy τ between the small- and large-model average accuracies, the objective is to find a cost-optimal cascade.Source: §3 · p. 2

4. Uncertainty-Calibrated Cascaded Inference

UCCI extracts uncertainty from the small model’s generated tokens, calibrates it as an error probability, and chooses an escalation threshold through constrained cost minimization. Its procedure uses a held-out calibration set and validation-set model outputs.

4. Uncertainty-Calibrated Cascaded Inference

  • UCCI’s three stages are extracting a scalar uncertainty score, calibrating it to an error probability, and selecting a threshold by constrained cost minimization.Source: §4 · p. 2

4.1. Token-level Uncertainty

  • For each generated token, UCCI defines the margin as the difference between the top-1 and top-2 next-token probabilities, then aggregates token margins into uncertainty u(x).Larger u(x) indicates greater uncertainty, and the method uses greedy decoding.Source: §4.1 · p. 2

4.2. Calibrating Uncertainty to Error Probability

  • The calibration target is an exact-match JSON correctness event, and isotonic regression fits a monotonic map from raw uncertainty to error probability using a calibration set.Source: §4.2 · pp. 2–3

4.3. Threshold Policy Selection

  • UCCI applies a threshold policy to the calibrated error probability and selects it on a validation set where both models have been run, using actual model outputs.Source: §4.3 · p. 3

5. Theoretical Analysis

The analysis establishes cost-optimality of threshold routing on calibrated error probabilities under explicit assumptions, and bounds isotonic calibration’s ECE rate. The routing guarantee depends most restrictively on large-model accuracy remaining invariant to which queries are escalated.

5. Theoretical Analysis

  • Under cℓ > cs, routing-invariant large-model accuracy, and calibrated error probabilities, a threshold policy meets the accuracy constraint at minimum cost among policies depending only on u(x).The characterization holds up to tie-breaking among queries with equal calibrated error probability.Source: §5 · p. 3
  • Escalating queries in decreasing calibrated error probability minimizes cost because their marginal accuracy gain per unit cost is non-decreasing.Since isotonic regression is piecewise constant, ties can occur without changing expected cost or accuracy.Source: §5 · p. 3
  • Under standard regularity conditions and bounded u(x), isotonic regression on n calibration examples has an ECE bound of O(n^-1/3).The proof sketch derives the L1 bound from an O(n^-2/3) MSE rate and uses concentration for the high-probability statement.Source: §5 · p. 3
  • Routing-invariant large-model accuracy is the most restrictive assumption; systematically harder escalated queries can lower subset accuracy below marginal accuracy.The paper tests this assumption in Section 6.3 and discusses when it can fail in Section 7.Source: §5 · p. 3

6. Experiments

End-to-end production NER tests show UCCI lowering cost at a matched F1 while isotonic calibration reduces ECE. Ablations and diagnostics examine routing signals, entity-level variation, and a theoretical assumption.

6.1. Setup

  • The production NER dataset has 75,000 labeled queries split into disjoint calibration, validation, and test sets; cascade evaluation routes actual test outputs end-to-end.The 4B and 12B instruction-tuned models have measured mean latencies of 47.2 ms and 142.3 ms per query, respectively.Source: §6.1 · p. 3

6.2. Results; 6.3. Ablations and Diagnostics

  • ECE fell from 0.12 to 0.03 with isotonic regression, whose calibrated curve tracks the diagonal across [0, 1].Isotonic regression also reached ECE = 0.03 versus 0.08 for temperature scaling; token margin gave the best cost-accuracy curve among three tested signals.Source: Figure 1 · p. 4, §6.2 · p. 4, §6.3 · p. 5

6.2. Results

  • At F1 = 0.91, UCCI has the lowest cost (2.08), beating entropy (2.31), conformal prediction (2.18), and FrugalGPT-style routing (2.24).At matched cost budget 2.00, UCCI exceeds the F1 target by 0.01 while entropy falls 0.02 below it.Source: Table 2 · p. 4, §6.2 · p. 4
  • At micro-F1 = 0.91, UCCI costs 2.08 versus 3.02 for large-model-only inference, a 31% reduction (95% CI: [27%, 35%]).Figure 2 reports the test-set cost-accuracy frontier under end-to-end cascade evaluation with measured costs.Source: §6.2 · p. 4, Figure 2 · p. 5
  • Across hypothetical large/small cost ratios, the accuracy-feasible threshold is unchanged; only the reported cost scale changes.Table 3 reports the same UCCI routing decisions under those hypothetical ratios.Source: §6.2 · p. 4, Table 3 · p. 5

6.3. Ablations and Diagnostics

  • The small-large F1 gap ranges from 1.8 points for focal length to 14.3 points for lens.The reported variation is across entity types on the test set.Source: §6.2 · p. 4, Table 4 · p. 5
  • At F1 = 0.91, the large model achieves micro-F1 = 0.928 on escalated queries versus 0.932 on the full test set, a difference of 0.004.The paper reports that the assumption of routing-invariant large-model accuracy holds approximately in this setting.Source: §6.3 · p. 5
  • Heavy-tailed margins could leave high-error queries in the score distribution’s body, where threshold routing escalates fewer than an oracle; margin-tail behavior was not measured.The paper identifies the largest residual gap from the large model on camera and lens entity types, where the small-large F1 gap is largest.Source: §6.3 · p. 5

7. Discussion and Limitations

The evaluation covers one domain, and its reported absolute gains are workload-specific. The discussion also identifies assumptions and deployment choices that bound how the method’s guarantees and measurements transfer.

Scope

  • The evaluation is limited to one photography NER domain; the 31% cost reduction and ECE = 0.03 are workload-specific, and broader validation would strengthen the generalization claim.The authors propose replication on a public NER benchmark and a different small/large model family.Source: §7 · p. 5

Theoretical assumption

  • The optimality guarantee assumes the large model’s accuracy is invariant to routing, which may fail when queries hard for the small model are also disproportionately hard for the large model.In that case, the threshold policy is no longer guaranteed globally optimal, though the empirical baseline comparison is unaffected.Source: §7 · pp. 5–6

Static calibration

  • Static held-out-batch calibration may need online or continual refitting under distribution shift in streaming deployments.The calibrated-threshold framework still applies, but the fitting procedure must be adapted.Source: §7 · p. 6

Cost model

  • Costs are measured in H100 latency; optimizing dollars, throughput, energy, or batched end-user latency requires re-deriving the model costs.The authors state this does not change the calibration step or optimality argument.Source: §7 · p. 6

8. Conclusion

UCCI calibrates token-level uncertainty into error probabilities and chooses an escalation threshold by constrained cost minimization. On the evaluated NER workload, it reduces inference cost while maintaining the reported accuracy and outperforming several routing baselines at the same operating point.

Conclusion

  • 31% lower inference cost at micro-F1 = 0.91 versus large-model-only inference; UCCI also beats entropy, conformal, and FrugalGPT-style baselines at that operating point.On 75,000 production NER queries, a small isotonic fit turns token-margin uncertainty into a routing score close to optimal in this setting.Source: §8 · p. 6

Ethics Statement

Lower inference cost may broaden access to LLM-based systems, while also lowering the barrier to deploying systems with potential for harm. The authors advise considering the use case and escalated-query population, especially in high-stakes settings.

Ethics Statement

  • Cost reductions may broaden access to LLM systems but also lower deployment barriers for systems with potential for harm, warranting attention to use case and escalated-query populations.The authors particularly flag high-stakes settings.Source: Ethics Statement · p. 6

A. Full Proofs

The appendix proves the threshold characterization and calibration-rate result, then records the prompt, implementation settings, and measured model latencies used in the evaluation.

A.1. Proof of Theorem 1

  • Escalating queries in decreasing calibrated error probability order minimizes cost under the accuracy constraint; this is implemented by thresholding the calibrated score, with ties on isotonic level sets immaterial to expected cost and accuracy.The proof’s allocation compares each query’s marginal accuracy gain with the common positive marginal cost of escalation.Source: §A.1 · p. 8

A.2. Proof of Proposition 2

  • Under mild regularity conditions and bounded uncertainty scores, the isotonic fit has a stated mean-squared-error rate, which yields an expected calibration error rate through an L1 bound and Cauchy-Schwarz.The proof also invokes a concentration argument to obtain a high-probability statement.Source: §A.2 · p. 8

B.1. Prompt Format

  • The prompt requests JSON entities for camera, lens, aperture, shutter speed, ISO, and focal length from a photo search query.The template places the user input after “Query:” and expects output afterward.Source: §B.1 · p. 8

B.2. Hyperparameters

  • The implementation uses greedy decoding, a 256-token maximum, vLLM 0.4.2, an NVIDIA H100 80GB GPU, and default isotonic-regression settings.Source: §B.2 · pp. 8–9

B.3. Cost Measurement Details

  • 142.3 ± 8.7 ms per query for the 12B model versus 47.2 ± 3.1 ms for the 4B model; the measured latency ratio cℓ/cs is 3.02.End-to-end latency was averaged over 100 cold-cache queries.Source: §B.3 · p. 9

B.4. Calibration Method Ablation

Isotonic regression better calibrates token margins than temperature scaling or uncalibrated routing on this workload. Heavy-tailed margins may still leave some high-error queries under-escalated relative to an oracle.

B.4. Calibration Method Ablation

  • ECE was 0.03 with isotonic regression, versus 0.08 with temperature scaling and 0.12 with uncalibrated routing on the calibration set.The authors attribute the difference to isotonic regression's ability to correct non-monotone miscalibration, unlike temperature scaling's single-parameter rescaling.Source: §B.4 · p. 9

B.5. Falsification Regime: Heavy-Tailed Margins

  • With heavy-tailed margins, finite-sample calibration may miss some high-error queries, so threshold routing escalates fewer than an oracle; the workload's tail statistics were not measured.The residual cost gap is concentrated on camera and lens entity types, consistent with—but not proof of—the heavy-tail hypothesis.Source: §B.5 · p. 9
Loading 2605.18796v1…