UCB Drift Router for Cheap-or-Expert Inference
Implementation & benchmark of arXiv:2607.06017 — Learning When to Automate: Queue Control in Human-AI Service Systems
Source paper: Learning When to Automate: Queue Control in Human-AI Service Systems arXiv:2607.06017 ⓘ · analyzed Aug 30, 2026
AI-generated research hypothesis, automatically tested. Not peer-reviewed.
Idea description
Replace a fixed confidence threshold in cascaded inference or mixture-of-experts routing with a queue-aware UCB-DPP controller. The controller sends an input to a cheap model when its optimistic estimated success is sufficiently high and the expert backlog is large, while escalating uncertain or high-value inputs when the penalty for an error dominates congestion. This should reduce expensive-model utilization without allowing latency or escalation queues to diverge.
Formulas
Mathematical statement
For task type k in [K], arrivals have probability lambda_k, the cheap model has unknown success probability p_k, Q_{t,k} is the fallback queue length, A_{t,k} is the number of new fallback jobs, and S_{t,k} is completed expert service. Queue stability is represented by Q_{t+1,k}=[Q_{t,k}+A_{t,k}-S_{t,k}]^+, where [z]^+=max(z,0). After n_k observations with s_k successes, estimate cheap-model success using hat{p}_k(t)=s_k(t)/max(1,n_k(t)) and the optimistic upper confidence bound U_k(t)=min(1,hat{p}_k(t)+sqrt(2 log(t)/max(1,n_k(t)))). For candidate routing action a, drift-plus-penalty minimizes V times immediate compute cost minus quality reward plus the backlog-weighted expected new fallback work, sum_k Q_{t,k} E[A_{t,k}|a]. Here V>0 controls the quality/compute tradeoff. This comes from the quadratic Lyapunov function L(Q)=1/2 sum_k Q_k^2: large queues increase the score of actions that create additional fallback work, providing negative feedback. UCB optimism prevents under-explored request types from being permanently assigned to the wrong route.
Implementation notes
Integrate the controller at the inference gateway immediately after request classification and before model execution. Partition requests into K buckets using an existing difficulty classifier, prompt-length bin, predicted value, or domain label. For every bucket maintain n[k], s[k], and a virtual backlog q[k]. Increment n[k] whenever the cheap model is evaluated; increment s[k] when its answer passes a task-specific verifier, delayed expert label, or large-model judge. At request t compute ucb[k]=clip(s[k]/max(1,n[k]) + sqrt(2*log(t)/max(1,n[k])), 0, 1). Enumerate candidate actions such as cheap-only, cheap-then-expert, and direct-expert. For cheap-then-expert, estimate fallback probability as 1-ucb[k]; for cheap-only, use the estimated probability of an uncorrected failure as its quality loss; for direct-expert, use zero fallback arrivals but higher compute cost. Score each action as V*(compute_cost - quality_reward) + sum_k q[k]*expected_fallback_arrivals[k], and select the minimum. After execution update q[k]=max(0, q[k]+observed_fallback[k]-completed_expert_service[k]); service can be measured per request or per fixed time window. The formulas determine the confidence radius, queue recursion, and objective; empirical calibration supplies costs, rewards, verifier labels, and service capacity. First test on a GSM8K or CIFAR-100 cascade with a small model and larger verifier, comparing fixed confidence thresholds, greedy cheap-only routing, and UCB-DPP at matched accuracy. Record expert-call rate, p95 latency, queue length, accuracy, and cumulative compute. Success means at least 20% fewer expert calls at equal accuracy while queue length remains bounded under sustained load.
Verification
Stage 1 · Toy mechanism gate: Failed ✗
Stage 2 · Mechanism transferred to benchmark: Not tested
Stage 2 · Practical benchmark result: Not run
Methodology: Toy-system gate first; the benchmark stage runs only after a pass. How verification works
Stage 1 — Mechanism check agent confidence 9/10
Built a toy UCB-DPP cascade simulator with queue recursion, DPP routing scores, UCB updates, stability sweeps, and baseline comparisons. The math checks succeeded: queue drift changed sign at lambda*fallback=service, the observed DPP switching backlog was 121 versus predicted 120, and UCB radii followed the predicted n^-1/2 scaling exactly. However, the router never escalated in the mini-experiment because optimistic UCB made fallback risk zero; it matched cheap-only at 65.0% accuracy rather than preserving expert-level quality, so the promised practical win was not demonstrated.
- Agent confidence
- 9/10
- Baseline
- fixed threshold: expert-call rate 99.996%, accuracy 99.996%, mean queue 0.00003, average compute cost 4.9999
- Idea
- UCB-DPP: expert-call rate 0%, accuracy 65.04%, mean queue 5.66, p95 queue 17, average compute cost 1.0; cheap-only was identical
Limitations: Synthetic Bernoulli toy rather than GSM8K/CIFAR or a production cascade. It does not test multiple task buckets, delayed verification, batching, real latency, calibrated confidence, or a corrected quality-aware UCB policy.
How to run: python3 router_experiment.py
Verdict computed by deterministic test code from paired-seed statistics — not by the language model.
Artifacts
- report.md 1.4 KB View
- results.json 2.9 KB View Raw
- router_experiment.py 5.4 KB View Raw