Learning to Route LLMs from Bandit Feedback: One Policy, Many Trade-offs

1 INTRODUCTION

• Deployment exposes only the chosen model’s accuracy and cost, whereas most existing routers train on labels for every candidate model on each prompt.This creates a mismatch between offline training and deployment feedback.Source: §1 · p. 1
• Prior routers also lack flexible test-time preference control: some require full-information labels, while bandit-style approaches may lack controllability or rely on offline pretraining.Table 1 compares routing methods by full-information offline supervision and preference-tunable inference.Source: §1 · p. 2, Table 1 · p. 2
• BaRP treats routing as a multi-objective contextual bandit, conditioning its policy on a user preference vector to balance performance and cost under bandit feedback.Its policy-gradient training uses entropy regularization and calibrated cost scaling; the preference can be adjusted at inference without retraining.Source: §1 · p. 2
• On in-distribution tasks, BaRP exceeds full-information offline routers by 12.46% and the top-performing individual LLM by 3.81%; on out-of-distribution tasks, the gains are 25.99% and 2.45%, respectively.Figure 1 presents testing scores on both task settings.Source: Figure 1 · p. 1, §1 · p. 2

2 APPROACH

• BARP casts routing as a multi-objective contextual bandit, learning from feedback only on the selected LLM.The prompt and preference enter the router, and the chosen model's feedback drives a policy update.Source: §2 · pp. 2–3, Figure 2 · p. 3

2.1 PROBLEM SETTING

• Each round combines a prompt and performance–cost preference vector as context, selects one LLM, and receives a scalar reward combining weighted score and capped cost.The task score is scaled to [0, 1], and τ caps cost so score and normalized cost are comparable.Source: §2.1 · p. 3

2.2 POLICY ARCHITECTURE

• The policy encodes the prompt and preference separately, concatenates their embeddings, and applies a decision head and softmax to produce probabilities over candidate LLMs.The prompt encoder is a frozen sentence transformer, while the preference encoder is a small MLP.Source: §2.2 · p. 3

2.3 OBJECTIVE AND LEARNING ALGORITHM

• REINFORCE maximizes expected cumulative reward using a variance-reduction baseline and entropy regularization; the baseline is the mini-batch mean reward, and the framework also accommodates classic bandit algorithms.Entropy regularization is used for improved exploration.Source: §2.3 · p. 4

2.4 TRAINING AND INFERENCE

• Training samples preferences uniformly on the 1-simplex and restricts supervision to the sampled LLM's outcome; at inference, operators provide a preference vector and the router chooses deterministically without retraining.Training uses complete benchmark logs to simulate bandit feedback.Source: §2.4 · pp. 4–5

3.1 DATASETS AND BENCHMARKS

• RouterBench and two question-answering datasets provide prompt-level candidate-model scores and costs; training mixes five of eight tasks, with three unseen benchmarks reserved for out-of-distribution evaluation.The five training tasks use 80%/20% train/test splits; the unseen benchmarks are MBPP, Hellaswag, and HotpotQA.Source: §3.1 · p. 5

3.2 BASELINE METHODS

• Comparisons include always choosing the smallest or largest LLM and two learned routers, RouterDC and GraphRouter, which require full-information labels.RouterDC uses dual-contrastive query and model embeddings, while GraphRouter uses graph-structured query, task, and model representations.Source: §3.2 · p. 5

3.3 METRICS

• Evaluation measures normalized task-success score, derived from exact match or GPT-4 ratings, and estimated API-call cost per query in USD.Table 2 reports testing scores on in-distribution tasks.Source: §3.3 · p. 5, Table 2 · p. 6

3.4 IMPLEMENTATION DETAILS

• The policy uses a frozen all-MiniLM-L6-v2 prompt encoder and two trainable MLPs, trained for 100 epochs with Adam at learning rate 1 × 10^-4 and batch size 32.Experiments ran on NVIDIA A100 80GB GPUs; REINFORCE entropy regularization uses β = 0.05.Source: §3.4 · p. 6

4.1 PERFORMANCE ON IN-DISTRIBUTION TASKS

• 73.57% average score gives BARP a stronger performance–cost trade-off than RouterDC and GraphRouter on in-distribution tasks.BARP’s score exceeds RouterDC by a relative 15.53% and GraphRouter by 12.44%.Source: §4.1 · p. 6

4.2 GENERALIZATION ABILITY TO NEW TASKS

• 66.08% average score is the highest among methods on unseen tasks, with BARP best on HpQA and competitive with the Largest LLM on MBPP and Hellaswag.RouterDC and GraphRouter’s performance drops sharply on MBPP and HpQA.Source: Table 3 · p. 6, §4.2 · p. 6

4.3 OVERALL PERFORMANCE AND COST-EFFECTIVENESS

• 16.84% overall average score improvement over GraphRouter coincides with a 50.00% monetary cost reduction across eight tasks.The comparison averages in-distribution and out-of-distribution tasks.Source: Table 4 · p. 7, §4.3 · p. 7

4.4.1 ANALYSIS OF THE PREFERENCE TRADE-OFF

• Increasing cost weight wc shifts BARP toward cheaper models, reducing average cost from $0.074 at wc = 0.2 to $0.015 at wc = 0.8 while performance declines.Operators can adjust this preference at inference without retraining.Source: §4.4.1 · p. 7

4.4.2 IMPACT OF PROMPT ENCODER CHOICE

• 0.7432 average score at balanced preference makes all-MiniLM-L6-v2 the best-scoring of the compared prompt encoders, despite its compact size.The results average in-distribution tasks; the authors use MiniLM as the default encoder.Source: §4.4.2 · pp. 7–8, Table 5 · p. 8

4.5 IMPACT OF DECISION HEAD ARCHITECTURE

• 0.7432 average score makes the MLP the strongest decision head; the linear head is competitive, while the bilinear head underperforms.The authors attribute a benefit to the MLP’s nonlinearity for mapping prompts and preferences to LLM choices.Source: §4.5 · p. 8, Table 6 · p. 8

4.6 ANALYSIS OF LEARNING ALGORITHMS

• REINFORCE achieves a substantially higher average score than classic contextual-bandit strategies, whose linear context–reward assumptions limit performance.The comparison uses balanced preference and identical context representations.Source: §4.6 · p. 9

6 ADDITIONAL RELATED WORK

• The comparison table covers REINFORCE and classical bandit algorithms, averaged across in-distribution tasks at balanced preference (wq = wc = 0.5).The caption gives the comparison scope but no result values.Source: Table 7 · p. 9
• GraphRouter uses graph-structured representations, RouterDC uses dual-contrastive objectives, and mixture-of-experts systems dynamically allocate queries across LLMs.These are examples of learning-based routing approaches discussed alongside early ensemble and static-heuristic strategies.Source: §5 · p. 9
• Contextual bandits select an action from an observed context and receive feedback only for that action; examples include LinUCB, Thompson Sampling, and ϵ-greedy.The passage also notes neural contextual bandits and applications in recommendation, online advertising, and adaptive experiment design.Source: §5 · p. 9

6 CONCLUSION AND DISCUSSION

• BARP significantly outperforms top-performing individual LLMs and strong offline routers on both in-distribution and out-of-distribution tasks.The conclusion reports this cross-task pattern without giving numerical margins in the supplied passage.Source: §6 · p. 9
• BARP frames routing as a preference-conditioned contextual bandit and trains with policy gradients on bandit feedback.It maps a user's prompt and performance–cost preference to a suitable LLM and offers inference-time tuning without retraining.Source: §6 · p. 9
• Training uses static offline logs rather than continuous live feedback; the single-step formulation is not explicitly designed for multi-turn conversations.Experiments also focus on general-purpose LLMs, while richer preferences and constraints such as latency are left for future work.Source: §6 · p. 10

ETHICS STATEMENT

• The stated goal is improving LLM-use efficiency, including selecting smaller, less expensive models when appropriate without significant performance loss.The paper links this aim to reduced energy consumption and carbon footprint from deployment.Source: ETHICS STATEMENT · p. 10
• The work uses existing public benchmark datasets and pretrained language models, with no private or personally identifiable information or human subjects.Source: ETHICS STATEMENT · p. 10
• The authors identify possible misuse for spam or misinformation generation, but assess the risk as limited and outweighed by efficiency benefits.Source: ETHICS STATEMENT · p. 10

REPRODUCIBILITY STATEMENT

• The authors state that training, evaluation, and analysis scripts required to replicate the experiments will be publicly released upon publication.Source: REPRODUCIBILITY STATEMENT · p. 10
• Data processing scripts will be provided, and dataset splits are deterministic based on the random seed in the code.Source: REPRODUCIBILITY STATEMENT · p. 10
• The appendix lists prompt-encoder and candidate LLMs, while critical hyperparameters are reported in section 3.4; the code uses PyTorch.Source: REPRODUCIBILITY STATEMENT · p. 10
• Experiments ran on a single NVIDIA A100 GPU with 80GB memory, and main-model training takes approximately 2–3 hours.Classic bandit baseline code is also provided and runs on a standard CPU.Source: REPRODUCIBILITY STATEMENT · p. 10

• Table 8 summarizes the paper’s notations.Source: Table 8 · p. 14

A.2 ADDITIONAL RESULTS

• Tables 9 and 10 report candidate LLM testing scores (%) on in-distribution and out-of-distribution tasks, respectively.Source: Table 9 · p. 14, Table 10 · p. 15

A.3 CANDIDATE LLMS

• The candidate LLM lists vary by task, covering RouterBench and, for NQ and HpQA, a separate set of models.Source: §A.3 · p. 15

A.4 DATASET DETAILS •

• The appendix describes datasets spanning math, knowledge and reasoning, commonsense, question answering, and code-generation tasks.Source: §A.4 · pp. 15–16

A.5 USE OF LLMS

• The LLM assisted with prose refinement and consistency checks, while the authors conceived and executed the scientific contributions.Source: §A.5 · p. 16