Adaptive LLM Routing under Budget Constraints

TL;DR

Calling GPT-4 for every query is wasteful: most queries don’t need it, and cheaper models would do fine. Existing “routers” decide which model to use, but they’re trained as supervised classifiers — which means you must first run every model on every query to build a labeled “best-model” dataset. That’s expensive and brittle when your traffic shifts. This paper reframes routing as a contextual bandit: the router only ever sees feedback (good/bad) for the one model it actually picked, exactly like a chat thumbs-up button. The trick to making bandits practical here is two-fold: (1) PILOT, a LinUCB variant that’s pre-warmed with offline human-preference data so it doesn’t start from zero, and (2) an online knapsack cost policy that spreads a fixed dollar budget across a stream of queries, spending more on the hard ones. Result on RouterBench: 93% of GPT-4’s quality at 25% of its cost, beating every bandit baseline, with routing overhead 38x smaller than a single GPT-4 call.

Problem & Motivation

You run an AI service. Every query that hits your system could go to GPT-4 (great, expensive) or a 7B open model (cheap, weaker). Sending everything to GPT-4 burns money on “what are your hours?”; sending everything to the cheap model fails on “compare these two phones and list the drawbacks of each.” You want a router that picks per-query.

The standard way to build one is supervised: collect a big dataset of (query → best LLM) labels, train a classifier. The pain is concrete:

1. Labeling cost is brutal. To know the “best” LLM for a query, you have to run all of them on that query and score the outputs. For 11 models and 36k queries that’s ~400k inference calls just to make the training set. You pay full inference cost before you’ve saved a cent.
2. It goes stale. A supervised router is frozen at training distribution. When your users start asking a new kind of question, the router has no way to learn from live traffic — it just degrades.

What you do get cheaply in production is bandit feedback: a thumbs-up/down on the single response you actually served. You never learn what the other 10 models would have said. The paper’s insight is that this is precisely the contextual bandit setting that powers news/ad recommendation (“show one article, observe one click”), and routing should be modeled the same way.

What’s New (Core Contribution)

Three genuine contributions, plus one nice theoretical result:

• Routing as a budget-constrained contextual bandit. Before: routing = supervised classification needing full query×model labels. Now: routing = online bandit learning from one-shot good/bad feedback, no exhaustive labeling, adapts as traffic drifts. This reframing is the paper’s spine.
• PILOT — preference-prior-informed LinUCB. Before: a cold-start bandit (LinUCB) explores blindly for a long time before it’s any good. Now: the bandit’s prior is initialized from offline human-preference data (ChatArena-style “which response is better”), so it starts informed and converges faster. They prove this lower bound: if the prior is within a ball of the true reward vector, regret is provably smaller than vanilla OFUL/LinUCB.
• Online cost policy as a multi-choice knapsack. Before: budgeting is usually a static threshold (“only use GPT-4 if score > x”) or an offline optimization with hindsight. Now: an online knapsack (ZCL algorithm) decides, per query, whether an expensive model is “worth it” given how much budget is left — with a provable closeness-to-optimal guarantee — and it’s decoupled from the bandit so you can change the budget dial live without retraining.
• Shared query↔LLM embedding space. Queries and LLMs live in the same vector space; affinity = cosine similarity. This makes the reward model linear (which is what lets LinUCB apply) and lets the space keep evolving from online feedback.

How It Works (Technically)

The whole system is a pipeline: embed the query → score each LLM by cosine affinity → add an exploration bonus (PILOT) → filter by what the budget allows (knapsack) → serve → collect feedback → update. Let’s demystify each piece.

1. The shared embedding space (the reward model). Every query q is first embedded by an off-the-shelf model φ (they use OpenAI text-embedding-3-small). Then a learned linear projection maps it into the shared space:

ψ(q) = W·φ(q) + b

That’s just one matrix multiply plus a bias — a single dense layer. W and b are what’s learned. Each LLM a also gets its own learned vector θ_a in that same space. The predicted quality of sending query q to LLM a is the cosine similarity of their (unit-normalized) vectors:

E[reward | a, q] = cos(ψ̂(q), θ̂_a) = ψ̂(q) · θ̂_a

Plain English: “how aligned is this query with this model?” High alignment = expect a good answer. Because both vectors are unit-normalized, the dot product is the cosine, and crucially the reward is linear in the LLM’s vector θ — that linearity is the hinge that lets a linear bandit (LinUCB) work here.

2. Pretraining the space from human preferences (warm start). Done in two phases for stability (jointly training query projection and LLM vectors via cosine is a “moving target” — both sides shift at once and training oscillates).

• Phase 1 — query projection. Using preference tuples (query, l_i, l_j, winner), learn W, b with a triplet loss: pull together queries that prefer the same LLM, push apart queries that prefer different ones. “Hard negatives” are queries where a smaller model actually won — those are the informative, surprising cases.
• Phase 2 — LLM vectors. Freeze W, b. For each preference pair, define p_i = softmax(cos(θ_i, ψ(q))) over the two competing LLMs, and train the θ_a with binary cross-entropy so the preferred model scores higher. Output: θ^pref_a per LLM — the prior.

3. PILOT: the online bandit (the heart). At runtime, treat each query’s projected embedding ψ(q_t) as the context and each LLM as an arm. PILOT is online ridge regression per arm. For arm a it maintains two running quantities:

A_a ← A_a + ψ̂(q)·ψ̂(q)ᵀ (a covariance-like matrix: how much/which directions we’ve seen) b_a ← b_a + r·ψ̂(q) (reward-weighted sum of contexts)

The point estimate of the arm’s vector is θ̃_a = A_a⁻¹·b_a (standard least-squares solution). The preference prior enters at initialization: instead of A_a = 0, set

A_a⁰ = λ_a·I, b_a⁰ = λ_a·θ^pref_a

which is exactly a Gaussian prior θ_a ~ N(θ^pref_a, (λ_a·I)⁻¹) centered on the preference-learned vector. λ_a is the prior strength: big λ_a = trust the prior, explore little; small λ_a = adapt fast. They set λ_a = 1/accuracy_a from pretraining — models that were reliable get a stronger prior.

Selection uses the UCB rule (optimism in the face of uncertainty): pick the arm maximizing predicted reward plus an uncertainty bonus.

a_t = argmax_a [ cos(ψ̂(q), θ̃_a) + α·√(ψ̂(q)ᵀ A_a⁻¹ ψ̂(q)) ]

The first term is exploitation (best guess). The second term is the exploration bonus — it’s large when A_a⁻¹ is large in the query’s direction, i.e. when we haven’t seen many queries like this one for this arm, so we should try it to learn. α tunes the explore/exploit balance. This is the classic bandit move: try uncertain options because they might be better than they currently look.

4. The online cost policy (multi-choice knapsack). PILOT ranks models by quality; it ignores money. The cost policy is a separate gate. Frame it as an online multi-choice knapsack (ON-MCKP): you have a knapsack of capacity B (your budget) and queries arrive over time; from each query’s set of LLMs you may pick at most one item, where an item’s value = predicted reward cos(ψ̂(q), θ̂_l) and weight = estimated token cost. The ZCL algorithm gives an eligibility threshold: an LLM is allowed only if its cost satisfies

C_l ≤ (cos(ψ̂(q), θ̂_l)) / ( (UB/LB)^{z_t} · (e/LB) )

where z_t ∈ [0,1] is current budget utilization and UB/LB are bounds on the reward-to-cost ratio. Read it intuitively: as you spend more (z_t → 1), the denominator grows, so the threshold tightens — expensive models must justify themselves more as the wallet empties. A binning trick (split Q queries into N bins, give each B/N, let unused budget spill forward) prevents the infinite-horizon ZCL math from leaving budget unspent at the end. Critically, this policy is model-agnostic — they apply it to every baseline too, so comparisons are fair.

Architecture & data flow

flowchart TB
  subgraph Offline["Offline pretraining (one time)"]
    P[Human preference data<br/>ChatArena pairs] --> T1[Phase 1: triplet loss<br/>learn query projection W,b]
    T1 --> T2[Phase 2: BCE loss<br/>learn LLM vectors theta_pref]
  end
  T2 -.warm-start prior.-> INIT[Init A=lambda*I, b=lambda*theta_pref]

  subgraph Online["Online serving loop (per query)"]
    Q[User query q_t] --> E[Embed phi then project psi]
    E --> SCORE[Cosine affinity to every LLM vector]
    INIT --> UCB
    SCORE --> UCB[PILOT: score + exploration bonus]
    UCB --> KNAP[Knapsack cost gate<br/>filter by remaining budget z_t]
    KNAP --> PICK[Pick best eligible LLM]
    PICK --> SERVE[Serve response]
    SERVE --> FB[User thumbs up/down = reward r_t]
    FB --> UPD[Update A_a, b_a for chosen arm]
    UPD -.next query.-> UCB
  end

Schematic of the explore/exploit tradeoff. Each bar is an LLM "arm": solid = current estimated quality, the faded cap = the UCB exploration bonus (shrinks as that arm gets pulled). PILOT picks the tallest total bar. Click an arm to "pull" it and watch its uncertainty collapse. Illustrative, not the paper's numbers.

The cost gate in action. As budget utilization (slider) rises, the eligibility threshold tightens and expensive models drop out of the eligible set — showing why the router naturally shifts toward cheaper models as the wallet empties.

The algorithm, simplified

# PILOT online routing loop — the core idea, minus I/O and tuning.
# embed(q) -> vec ; llm_call(model, q) -> response ; score(resp) -> float in [0,1]

import numpy as np

def init_arms(llms, theta_pref, lam):                 # lam[a] = 1 / pretrain_accuracy[a]
    A = {a: lam[a] * np.eye(D) for a in llms}          # prior precision (D = shared dim)
    b = {a: lam[a] * theta_pref[a] for a in llms}      # prior centered on preference vector
    return A, b

def route(q, llms, A, b, alpha, budget):
    psi = unit(project(embed(q)))                      # query -> shared space, normalized
    best, best_ucb = None, -1
    for a in llms:
        theta = unit(np.linalg.solve(A[a], b[a]))      # theta_a = A^-1 b   (ridge estimate)
        mean  = psi @ theta                            # exploitation: cosine affinity
        bonus = alpha * np.sqrt(psi @ np.linalg.solve(A[a], psi))  # exploration: uncertainty
        if cost_eligible(a, q, budget) and mean + bonus > best_ucb:
            best, best_ucb = a, mean + bonus           # knapsack gate AND highest UCB
    return best, psi

def update(a, psi, reward, A, b):                      # one arm, bandit feedback only
    A[a] += np.outer(psi, psi)                         # we saw context psi for arm a
    b[a] += reward * psi                               # weight it by observed quality

# serving loop
A, b = init_arms(llms, theta_pref, lam)
for q in query_stream:
    a, psi = route(q, llms, A, b, alpha=2.0, budget=remaining)
    resp   = llm_call(a, q)
    r      = score(resp)                               # in production: thumbs up/down
    update(a, psi, r, A, b)
    remaining -= estimated_cost(a, q)

The two things that make this this paper and not a textbook LinUCB: init_arms injecting the preference prior, and cost_eligible enforcing the knapsack gate. Everything else is vanilla contextual bandit machinery.

Built on Prior Work

Results & Evidence

Setup. RouterBench (36,497 samples, 64 tasks, 11 LLMs — Llama-70B, Mixtral, Yi-34B, GPT-4, GPT-3.5, Claude v1/v2/instant, etc.), each query pre-scored on all models so they can simulate online traffic. Preference data from ChatArena. Traffic split into a “learning bucket” (10) and “deployment bucket” (1).

Headline numbers:

• Multi-task (full RouterBench): 93% of GPT-4’s performance at 25% of its cost, beating every bandit baseline.
• Single-task (MMLU): 86% of GPT-4 at 27% of cost.
• Lowest cumulative regret and highest deployment performance across learning-bucket sizes (i.e., works with little data).
• Overhead: routing takes 0.065s (Instructor-XL) or 0.239s (OpenAI embed) vs 2.5s for one GPT-4 call — 10–38x faster than the inference it gates.
• The online cost policy matches or beats an offline P − λC policy tuned with full hindsight (Table 2) — impressive for an online method.
• Sensible qualitative routing: 90% of MMLU and 89% of ARC to GPT-4; 94% of GSM8K to Claude-v1 (good+cheap at math); Claude handles 28% of MBPP coding.

What the evidence does NOT establish — read this before you sell it:

• It’s a simulation, not live deployment. Rewards come from RouterBench’s precomputed scores, not real thumbs-up from real users. Real bandit feedback is noisier, sparser, and biased (people rate when annoyed). The clean convergence may not survive contact with production.
• Bandit baselines, weak supervised comparison. It beats LinUCB/Epoch-Greedy/random handily, but the comparison to strong supervised routers (HybridLLM) is relegated to the appendix and framed around the different supervision regime — so “we beat supervised routing” is not cleanly demonstrated.
• Budget is ignored during learning. By the authors’ own admission (Limitations), the budget constraint only applies at deployment, not during the bandit’s learning phase. Learning under budget is left to future work.
• Single-turn only. No multi-turn conversations — a big gap for real chat products.
• The regret-bound proof needs ‖θ_pref − θ*‖ ≤ ‖θ*‖ — the prior must be at least roughly right. A bad/mismatched preference dataset could hurt rather than help; this isn’t stress-tested.

How You’d Use It

This is directly buildable into an AI-services product, and it’s a genuinely sellable capability.

• Cost-control layer in front of any multi-model gateway. If you run an LLM gateway for clients (the “buy” version is LiteLLM/OpenRouter), PILOT is the brain that decides which model each request hits. The pitch to a client: “same quality your users feel, ~25% of the spend” is a CFO-friendly line, and you have a paper to point to.
• The thumbs-up button becomes training data. Most chat UIs already have like/dislike. This paper turns that latent signal into a continuously-improving router — no labeling project, no model-eval pipeline. That’s the killer operational property: it learns from what you already collect.
• A live budget dial. Because cost policy is decoupled from the learner, you can expose “monthly budget” as a literal slider per client/tenant and the router re-allocates in real time. That’s a clean SaaS feature.
• Drift resilience as a retention story. Supervised routers rot; this one adapts. For long-running deployments where query mix changes seasonally, “self-tuning” is a real differentiator.
• In a multi-agent system (your ARC MAS background): each agent role could route its sub-calls through a shared PILOT instance — cheap models for boilerplate sub-tasks, expensive for the reasoning-heavy node — with budget shared across the whole agent graph via the knapsack.

Realistic effort: a working v1 in front of an existing gateway is days, not months. The hard part is evaluation (see below), not the bandit code.

Build Your Own (Minimal Recipe)

Smallest version that captures ~80% of the value:

1. Embedder + projection. Use any embedding API for φ. Skip the fancy two-phase preference pretraining at first — initialize θ_a from a handful of labeled “this model is good at X” examples, or even random + a short warmup. Add the preference prior later when you have it; it’s an accelerant, not a prerequisite.
2. The bandit. ~40 lines (above). Per arm keep A (a D×D numpy matrix) and b (a D vector). D = embedding dim or a projected-down version (128–256 is plenty). This is the part that is the paper and it’s small.
3. Reward plumbing. Wire your UI’s thumbs-up/down to reward ∈ {0,1}. This is the highest-leverage piece — without real feedback you have nothing to learn from.
4. Cost gate. Start with the dumb version the paper itself compares against: B/Q per query with spillover, pick the highest-ranked model under that per-query budget. Upgrade to the full ZCL/knapsack threshold only if you can measure it beating the dumb one (the paper shows the gain is modest at higher budgets).

The two genuinely hard parts: (a) offline evaluation — you can’t A/B a bandit easily, so build a replay simulator from logged (query, model, reward) traffic like the paper does with RouterBench; (b) cold start — until feedback accumulates, route conservatively (the preference prior or a static heuristic) so early users aren’t your guinea pigs.

Reach for: numpy (the whole bandit), any embedding model, LiteLLM/OpenRouter for the multi-model backend, your existing analytics for the thumbs signal.

How to Improve It

1. Budget-aware learning (the paper’s own gap). Fold the cost into the reward or the UCB during the learning phase — e.g., optimize reward-per-dollar cos(ψ,θ)/cost directly — so exploration doesn’t waste money on expensive arms. Testable: compare regret-per-dollar vs the decoupled version.
2. Non-linear reward model. Cosine-in-a-shared-space is linear by design (so LinUCB applies). Swap in a neural bandit (e.g., NeuralUCB) for richer query↔model interactions; measure whether the extra capacity beats the linear model’s faster convergence.
3. Multi-turn / contextual state. Embed the conversation, not just the latest turn, and let routing depend on dialogue state (escalate to GPT-4 once a thread gets hard). Big practical win for chat products.
4. Per-task / per-tenant priors. One global prior averages over very different query mixes. Cluster queries (or use tenant ID) and maintain per-cluster θ^pref — a contextual prior. Test on RouterBench’s heterogeneous tasks.
5. Robustness to bad priors. Add a mechanism that detects when the preference prior disagrees with online feedback and decays λ_a faster for that arm — directly attacking the ‖θ_pref − θ*‖ assumption the proof leans on.
6. Reward-model uncertainty in the gate. The knapsack uses point estimates of value; feed it PILOT’s uncertainty too, so it doesn’t over-commit budget to a high-mean-but-unsure model.

Glossary

• LLM routing — deciding, per query, which model in a pool to call.
• Contextual bandit — sequential decision problem where you pick one “arm” (action) given context, observe reward only for that arm, and adapt. Here: pick one LLM per query, see only its feedback.
• Bandit feedback — you learn the outcome of the choice you made, never the choices you skipped (vs. supervised learning, which has the full answer key).
• LinUCB — a contextual bandit assuming reward is linear in the context; picks the arm with the highest upper confidence bound.
• UCB (Upper Confidence Bound) — “optimism”: rank actions by estimated value plus an uncertainty bonus, so under-explored actions get tried.
• Exploration vs. exploitation — try uncertain options to learn (explore) vs. use the current best (exploit). The bonus term tunes the balance.
• Regret — cumulative gap between your choices and the best-possible choices in hindsight; lower = better learner.
• Prior (Bayesian) — initial belief before data. Here the preference-learned LLM vector seeds the bandit so it starts informed.
• OFUL — a well-studied linear bandit (Optimism in Face of Uncertainty Linear); used for the regret-bound proof.
• Ridge regression — least squares with an L2 penalty; what each arm runs to estimate its vector (θ = A⁻¹b).
• Triplet loss — pulls an anchor toward positives and away from negatives in embedding space; used to learn the query projection.
• Hard negative — a deliberately tricky negative example (here: a query where a smaller model unexpectedly won) that sharpens learning.
• BCE (binary cross-entropy) — loss for two-class problems; used to fit LLM vectors from pairwise preferences.
• Multi-choice knapsack (MCKP) — pick at most one item from each group to maximize value within a weight budget. Online version (ON-MCKP) handles items arriving over time.
• ZCL algorithm — the online-knapsack policy (Zhou et al. 2008) that gives the cost-eligibility threshold and a near-optimal guarantee.
• Budget utilization (z_t) — fraction of budget spent so far; tightens the cost gate as it climbs.
• RouterBench — a 36k-query, 11-LLM benchmark with precomputed scores/costs, used to simulate online routing.