Table of Contents
Fetching ...

Cost-Aware Contrastive Routing for LLMs

Reza Shirkavand, Shangqian Gao, Peiran Yu, Heng Huang

TL;DR

CSCR tackles the problem of cost-aware routing among diverse LLM pools by learning a unified embedding space that jointly encodes prompt semantics, model capability, and per-model cost. It introduces two lightweight descriptors, logit footprints for open-weight LLMs and perplexity fingerprints for API-based models, and trains a contrastive InfoNCE objective augmented with cost bands to prioritize cheaper yet accurate experts. Inference reduces to a fast k-NN lookup over FAISS, enabling microsecond routing while avoiding retraining when the pool changes. Across three benchmarks, CSCR improves the accuracy–cost Pareto frontier and generalizes robustly to unseen models and out-of-distribution prompts, illustrating practical gains for scalable, economical deployment of multi-LLM systems.

Abstract

We study cost-aware routing for large language models across diverse and dynamic pools of models. Existing approaches often overlook prompt-specific context, rely on expensive model profiling, assume a fixed set of experts, or use inefficient trial-and-error strategies. We introduce Cost-Spectrum Contrastive Routing (CSCR), a lightweight framework that maps both prompts and models into a shared embedding space to enable fast, cost-sensitive selection. CSCR uses compact, fast-to-compute logit footprints for open-source models and perplexity fingerprints for black-box APIs. A contrastive encoder is trained to favor the cheapest accurate expert within adaptive cost bands. At inference time, routing reduces to a single k-NN lookup via a FAISS index, requiring no retraining when the expert pool changes and enabling microsecond latency. Across multiple benchmarks, CSCR consistently outperforms baselines, improving the accuracy-cost tradeoff by up to 25%, while generalizing robustly to unseen LLMs and out-of-distribution prompts.

Cost-Aware Contrastive Routing for LLMs

TL;DR

CSCR tackles the problem of cost-aware routing among diverse LLM pools by learning a unified embedding space that jointly encodes prompt semantics, model capability, and per-model cost. It introduces two lightweight descriptors, logit footprints for open-weight LLMs and perplexity fingerprints for API-based models, and trains a contrastive InfoNCE objective augmented with cost bands to prioritize cheaper yet accurate experts. Inference reduces to a fast k-NN lookup over FAISS, enabling microsecond routing while avoiding retraining when the pool changes. Across three benchmarks, CSCR improves the accuracy–cost Pareto frontier and generalizes robustly to unseen models and out-of-distribution prompts, illustrating practical gains for scalable, economical deployment of multi-LLM systems.

Abstract

We study cost-aware routing for large language models across diverse and dynamic pools of models. Existing approaches often overlook prompt-specific context, rely on expensive model profiling, assume a fixed set of experts, or use inefficient trial-and-error strategies. We introduce Cost-Spectrum Contrastive Routing (CSCR), a lightweight framework that maps both prompts and models into a shared embedding space to enable fast, cost-sensitive selection. CSCR uses compact, fast-to-compute logit footprints for open-source models and perplexity fingerprints for black-box APIs. A contrastive encoder is trained to favor the cheapest accurate expert within adaptive cost bands. At inference time, routing reduces to a single k-NN lookup via a FAISS index, requiring no retraining when the expert pool changes and enabling microsecond latency. Across multiple benchmarks, CSCR consistently outperforms baselines, improving the accuracy-cost tradeoff by up to 25%, while generalizing robustly to unseen LLMs and out-of-distribution prompts.

Paper Structure

This paper contains 87 sections, 2 theorems, 25 equations, 7 figures, 16 tables.

Key Result

Theorem 5.2

Let $\hat{r}_k$ be trained on $n$ i.i.d. prompt embeddings. Under Assumption ass:lipschitz, for any $\lambda\!\ge\!0$ and any $k\!\le\!n$, where $C$ depends only on $L$ and the diameter of $\Phi_q(\mathcal{X})$, and $r^\star$ is the Bayes-optimal rule $r^\star(x)=\mathop{\rm arg\,min}_m\bigl[\gamma(x,h_m)+\lambda c(h_m)\bigr]$.

Figures (7)

  • Figure 1: Accuracy–cost/size deferral curves on three expert pools. Across all benchmarks, our Cost-Spectrum Contrastive Router (blue) consistently dominates the Pareto frontier, achieving higher accuracy at lower model size and latency (left, middle) and reduced cost (right).
  • Figure 2: Cosine similarity of perplexity descriptors for RouterBench LLMs. Despite using a shared scorer, the descriptors distinctly separate the experts.
  • Figure 3: Deferral curves of test on new LLMs.
  • Figure 4: Effect of Descriptor Type. Perplexity descriptors slightly improve AUDC but require an extra pass compared to the faster logit descriptors. Mixing descriptors has no impact on results.
  • Figure 5: Real sample of routing decisions made by UMR, vanilla contrastive router, and CSCR. CSCR chooses a more expensive but accurate expert than the vanilla router, while also selecting a cheaper option than UMR, achieving better accuracy–cost trade-offs on both ends.
  • ...and 2 more figures

Theorems & Definitions (2)

  • Theorem 5.2: Excess risk
  • Lemma 5.3: Directional alignment with cost bands