Contextual Bandit Model Routing With Cost

Contextual bandit model routing learns which qualified model completes a task at the best quality, safety, latency, and cost trade for this request. The learner is useful only inside a hard eligibility boundary and with evidence that corrects for the actions production happened to choose.

This guide builds mature reward, logged propensities, doubly robust offline estimates, cohort gates, and capped online exploration. The resulting policy can spend uncertainty deliberately without letting a cost objective rewrite safety rules.

Contextual bandit model routing defines one bounded action

A model router chooses an action before observing the answer's value. Context includes task family, language, prompt length, retrieved evidence, customer tier, latency budget, and consequence. Actions are eligible models or an abstain/review path.

Reward arrives later from task success, groundedness, format validity, latency, and cost. Contextual bandit model routing learns this mapping without pretending every request has the same optimum.

Start with a narrow decision such as choosing small versus large model for support classification after fixed preprocessing. Do not let the learner simultaneously change prompts, tools, retrieval, temperature, and review policy. Those factors make reward attribution ambiguous and expand the safety surface. The first policy should beat a static baseline on a logged workload before it controls broader generation.

The Vowpal Wabbit contextual bandit tutorial describes the context-action-cost interaction and common exploration approaches in an official implementation. Its framing is useful even if the production learner differs. Log the probability of the chosen action at decision time; without that propensity, reliable offline comparison becomes much harder.

Eligibility comes before optimization

Not every model can serve every request. Policy may require a model with a particular data region, modality, context window, tool support, safety qualification, or contract. Build a deterministic eligibility filter and test it independently.

The bandit ranks only the remaining actions. It must never learn that violating a hard rule produces attractive short-term reward.

A safe exploration policy also constrains where alternatives can be sampled. High-consequence actions may stay on the qualified incumbent until shadow evidence is strong; low-risk summarization can explore within a small probability budget. Contextual bandit model routing should return review when no action is eligible rather than fall back to the cheapest model. The runnable fixture encodes that terminal state explicitly.

Use LLM routing by cost, risk, and latency to define the rule-based eligibility boundary and fallback contract. The learner operates inside it. This separation makes policy review possible: security and product owners approve the candidate set, while data owners evaluate how reward and exploration choose among qualified candidates.

Runnable artifact: The fixture keeps exploration inside safety-qualified actions and returns abstain when eligibility is empty.

Save this proof as bandit-router.test.mjs and run node bandit-router.test.mjs. Expected final line: PASS: budgeted exploration.

import assert from "node:assert/strict";
const utility=(reward,cost,lambda=.002)=>reward-lambda*cost;
function route(context,models,explore){const eligible=models.filter(m=>m.safety>=context.minSafety);if(!eligible.length)return "abstain";if(explore&&context.bucket<.05)return eligible.at(-1).id;return eligible.sort((a,b)=>utility(b.reward,b.cost)-utility(a.reward,a.cost))[0].id}
const models=[{id:"small",reward:.82,cost:8,safety:.99},{id:"large",reward:.88,cost:45,safety:.995}];
assert.equal(route({minSafety:.99,bucket:.8},models,false),"small");
assert.equal(route({minSafety:.99,bucket:.01},models,true),"large");
assert.equal(route({minSafety:1,bucket:.8},models,false),"abstain"); console.log("PASS: budgeted exploration");

Build reward from completed work, not proxy applause

A useful reward reflects task completion after all downstream checks. For extraction, score exact fields and schema validity. For support drafting, include reviewer acceptance and major edits.

For code, use tests and security checks. Subtract measured inference cost and latency penalties only after defining their exchange rate. A cheap incorrect answer should never win because dollars are immediate and quality labels arrive later.

Use cost-sensitive reward with transparent units: task_value - λ_cost × cents - λ_latency × excess_ms - harm_penalty. Sweep the lambdas and publish the resulting policy, not only one chosen coefficient. Clip or bound components so one anomalous invoice cannot dominate learning. Preserve separate dashboards for quality, safety, latency, and cost even when training consumes a scalar.

The Microsoft Research contextual bandit personalization paper is a primary account of online contextual decisions and exploration. Model routing differs because actions have service cost and sometimes asymmetric harm, yet the lesson remains: the policy learns only from displayed actions and needs deliberate exploration to discover better alternatives.

Log propensities and delayed outcomes as first-class data

Each decision row needs context available at routing time, eligible action set, chosen action, probability under the behavior policy, model and prompt versions, estimated token work, actual service metrics, response ID, outcome availability, reward components, and policy version. Never add post-response facts to context; that is label leakage and makes offline performance impossible to reproduce online.

Join delayed labels through immutable request and outcome IDs. Reviewer acceptance may arrive in minutes, refund correctness in days, and customer retention much later. Define a reward maturity window and train only on mature rows or use an explicit censoring method. Contextual bandit model routing can otherwise prefer actions whose failures simply take longer to appear.

Log rejected and timeout outcomes, not just successful completions. A provider outage changes availability and reward. Distributed tracing can connect routing to downstream tool and review spans, but the learning table should remain a versioned analytical contract. Redact text features into approved derived signals where possible so training does not become an uncontrolled copy of customer prompts.

Evaluate candidate policies against logged traffic

Replay cannot reveal rewards for actions the behavior policy never chose. Off-policy estimators correct this gap using logged propensities and a reward model, with different bias and variance tradeoffs. Implement direct method, inverse propensity weighting, self-normalized variants, and doubly robust estimation. Compare them, report effective sample size, and flag contexts with poor action overlap.

The Doubly Robust Policy Evaluation and Learning paper provides the foundational estimator that combines a reward model with propensity correction. Use confidence intervals from context-preserving bootstrap or another justified method. An impressive point estimate with a tiny effective sample is not release evidence, especially when the candidate routes rare languages or long prompts differently.

An off-policy evaluation table should include incumbent replay, candidate value, uncertainty, cost, safety metrics, worst cohort, overlap, and estimator disagreement. AI evaluation measurement contracts helps freeze outcomes and cohorts before inspection. Reject a candidate whose aggregate reward improves while an important cohort loses its qualified action or crosses a safety bound.

Spend exploration like an operational budget

Exploration has a cost in money, latency, and uncertain quality. Allocate it explicitly by context and day. For low-risk traffic, reserve perhaps five percent of decisions for qualified alternatives; for medium risk, one percent under shadow review; for high risk, zero online exploration until a controlled study grants it. The budget should shrink automatically during incidents, label delays, model drift, or elevated override rates.

Use an exploration budget counter enforced outside the learner. Include per-action and per-cohort caps so one obscure model cannot consume all exploratory traffic. Randomization must be auditable and stable enough for propensity reconstruction. If the router cannot record the probability actually used after all filters and caps, the event is ineligible for learning.

Start with shadow decisions: ask the candidate policy what it would choose while the incumbent serves. Shadowing measures disagreement and service feasibility but not counterfactual quality. Then use a small randomized canary. Canary evals for AI releases supplies rollback and exposure gates; add cumulative regret, spend, and label maturity to its normal health signals.

Guard against nonstationarity and feedback loops

Models, prices, latency, traffic, prompts, and reviewers change. The policy may also alter which examples receive review, which changes the labels available for future training. Monitor context and action distributions, propensity extremes, reward delay, model-specific calibration, override reasons, and the fraction of traffic outside logged support. Retrain on a declared cadence and retain the previous policy for immediate rollback.

A router can create a feedback loop by sending easy cases to a small model, leaving only hard labels for the large model. Raw action averages then make the large model look worse. Compare on overlapping contexts and randomized traffic, not selected production means. Maintain periodic uniform exploration within safe eligibility or a fixed benchmark stream to keep action comparisons anchored.

Use LLM admission control before routing when capacity is saturated. Congestion is an operational state, not an excuse to reinterpret reward. The system can defer, shed, or narrow eligibility transparently. If the bandit learns from emergency traffic without a regime feature and policy boundary, it may encode incident behavior as the new normal.

Publish a routing policy card and rollback path

The release card lists decision scope, contexts, features, eligibility, actions, behavior-policy logs, reward equation, outcome maturity, candidate algorithm, estimator suite, overlap, effective sample size, offline value intervals, cohort gates, exploration caps, serving budget, incident behavior, retraining schedule, and rollback. Contextual bandit model routing is not ready when only a notebook can reconstruct why a request selected one provider.

Audit individual decisions with feature values, eligible set, predicted values, selection probability, exploration reason, and final reward components. Avoid exposing sensitive derived features to end users, but give operators enough evidence to diagnose systematic routing. A counterfactual explanation should say which bounded factor would have changed eligibility or rank, not invent a narrative about the model's inner reasoning.

Ship only after the incumbent reproduces under the evaluator, the candidate improves qualified value with uncertainty, critical cohorts pass, propensities are recorded, exploration cannot cross hard rules, labels mature reliably, and rollback restores static routing. The goal is a learner that spends uncertainty deliberately. Cost reduction is durable only when the evidence shows which completed tasks paid for it.

Contextual bandit model routing belongs behind qualification and budget enforcement, never in front of them. Review contextual bandit model routing as a controlled production policy whose evidence must mature before its rewards train the next policy.