💬 NLP

Expected Reward Prediction, with Applications to Model Routing

This paper introduces a method to predict the expected reward an LLM will achieve on a given prompt before generating a response, demonstrating that this simple prediction enables an efficient and extensible model routing protocol that outperforms category-based baselines by optimizing the trade-off between reward and computational cost.

Original authors: Kenan Hasanaliyev, Silas Alberti, Jenny Hamer, Dheeraj Rajagopal, Kevin Robinson, Jasper Snoek, Victor Veitch, Alexander Nicholas D'Amour

Published 2026-03-24

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Kenan Hasanaliyev, Silas Alberti, Jenny Hamer, Dheeraj Rajagopal, Kevin Robinson, Jasper Snoek, Victor Veitch, Alexander Nicholas D'Amour

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are the manager of a massive, high-tech call center. You have a team of AI assistants (the Large Language Models) with different skill levels and different "hourly rates." Some are junior assistants who are cheap and fast but might make mistakes. Others are senior experts who are slow and expensive but usually get things right.

Every time a customer calls with a question (a prompt), you have to decide: Which assistant should I send this call to?

The Old Way: Guessing or Checking Everything

Traditionally, to make this decision, you might have to:

Guess based on the category of the question (e.g., "Math questions go to the math expert").
Or, ask every single assistant to write an answer, read all of them, pick the best one, and then pay for all the work you didn't use. This is incredibly wasteful and slow.

The New Way: The "Crystal Ball" (Expected Reward Prediction)

This paper introduces a clever trick called Expected Reward Prediction (ERP). Think of it as a crystal ball that can look at a customer's question before any assistant answers it and predict exactly how well a specific assistant will do.

Here's how it works, broken down into simple steps:

1. The "Scorecard" (Reward Models)

First, imagine you have a strict judge (a Reward Model) who listens to conversations and gives a score from 0 to 100 based on how helpful the answer is.

Old problem: The judge only scores the answer after it's written.
New insight: The authors realized that the judge's scoring system is so consistent that you can predict the average score an assistant will get just by looking at the question.

2. The "Simple Math" (Linear Probes)

You might think predicting this requires a super-complex AI brain. Surprisingly, the authors found that a very simple math formula (a linear model) works perfectly.

The Analogy: Imagine you have a map of the question (an "embedding"). You just draw a straight line on that map. Where the line hits tells you, "If Assistant A answers this, they will likely get an 85/100. If Assistant B answers it, they will likely get a 40/100."
It's like looking at a weather forecast. You don't need to simulate every raindrop to know if you need an umbrella; you just look at the barometer.

3. The "Smart Dispatcher" (Model Routing)

Now, you have your crystal ball. When a call comes in:

The crystal ball predicts: "The Junior Assistant will score 60. The Senior Expert will score 95."
The dispatcher checks the cost: "The Junior costs $1. The Senior costs $10."
The Decision: If the question is simple, the crystal ball says the Junior will do fine (score 80). The dispatcher sends it to the Junior to save money. If the question is hard, the crystal ball says the Junior will fail (score 20), so it sends it to the Senior.

Why is this a Big Deal?

It's Fast: You don't need to generate answers from everyone to know who is best. You just do a quick calculation.
It's Flexible: If you hire a new assistant tomorrow, you just teach the crystal ball how they usually perform. You don't have to retrain the whole system or compare them against everyone else pairwise.
It Beats the "Category" Rule: Sometimes a "Math Expert" is bad at a specific type of math question, but a "Generalist" is great at it. The crystal ball sees the nuance; a simple rule like "Math goes to the Math Expert" does not.

The "Zooter" Connection

The paper mentions a complex system called "Zooter" that does something similar but is much harder to build. The authors show that Zooter works because it is secretly trying to do exactly what this simple crystal ball does: predict the expected score. By using the simple method, you get 90% of the benefit with 10% of the effort.

The Bottom Line

This paper proves that we can predict how good an AI will be at a specific task just by reading the question, without waiting for the AI to actually do the work. This allows us to route tasks to the right AI instantly, saving huge amounts of money and time while still getting high-quality answers.

In short: Instead of asking five chefs to cook a meal to see who is best, you just look at the recipe and know instantly which chef is the right one for that specific dish.

1. Problem Statement

Large Language Models (LLMs) are typically evaluated using Reward Models (RMs), which score individual responses ( $y$ ) to a prompt ( $x$ ). However, standard RMs are designed to rank responses from a single model or compare two specific responses. They do not inherently provide a metric for a model's suitability for a specific prompt before the model generates a response.

The paper addresses two core challenges:

Theoretical Feasibility: Can we predict the Expected Reward ($ER$) of a specific LLM for a given prompt a priori (without sampling responses)?
- Challenge 1: LLMs are non-deterministic; their response distribution varies.
- Challenge 2: Standard RMs are often trained on pairwise preference data (Bradley-Terry model), which only guarantees relative ordering ( $r(x, y_1) - r(x, y_0)$ ) and allows for arbitrary, prompt-dependent offsets. This theoretically makes the absolute expected value unidentifiable.
Application: How can this prediction be used to optimize Model Routing?
- In a pool of models with varying costs (e.g., 8B vs. 70B parameters), how do we route a prompt to the model that maximizes reward while minimizing computational cost?

2. Methodology

A. Defining Expected Reward

The authors formalize the expected reward of a model $\pi$ for a prompt $x$ under a reward function $r$ as:
$ER_\pi(x) := \mathbb{E}_{Y \sim \pi(Y|x)} [r(x, Y)]$
The goal is to learn a function $f(x)$ that predicts this quantity directly from the prompt embedding.

B. Data Generation & Training

Dataset: The open-perfectblend dataset, containing 4,000 prompts across four categories: General Chat, Instruction Following, Math, and Coding.
Ground Truth Construction: For each prompt $x_i$ and model $\pi$ , the authors generate $K=32$ responses. They compute the reward for each response using a pre-trained RM and calculate the empirical mean as the target label: $\hat{ER}_\pi(x_i) = \frac{1}{K} \sum r(x_i, y^{(k)})$ .
Predictor Architecture: A surprisingly simple linear model (ridge-regularized linear regression) is trained.
- Input: Fixed-length vector embeddings of the prompt $x$ (using gte-large-en-v1.5).
- Output: The predicted expected reward.
- Training Objective: Minimize squared error between the predicted value and the empirical mean reward.

C. Model Routing Protocol

The paper proposes an Expected Reward Prediction (ERP) routing strategy:

Train a separate linear predictor for each model in the pool to predict its $ER$ for any given prompt.
At inference time, for a prompt $x$ , calculate the predicted $ER$ for all models.
Select the model $i$ that maximizes a cost-adjusted score:
$\rho_\lambda(x) = \arg\max_i \left( \hat{ER}_{\pi_i}(x) - \lambda \cdot c(i) \right)$
Where $c(i)$ is the computational cost (proportional to model size) and $\lambda$ is a tunable parameter balancing reward vs. cost.

3. Key Contributions

Empirical Discovery of Predictability: The paper demonstrates that despite theoretical ambiguities in reward modeling (specifically regarding the Bradley-Terry offset), the expected reward of an LLM is highly predictable from the prompt alone using a simple linear probe on prompt embeddings.
Lifting Response-Level to Model-Level Rewards: The authors successfully "lift" response-level scoring to model-level suitability scoring. This allows for the discrimination of models based on specific prompt characteristics without generating actual responses.
Scalable Model Routing: The ERP approach offers a significant scalability advantage over preference-based routing (like Zooter or pairwise classifiers).
- Preference-based: Requires training a model on pairwise comparisons between all model pairs ( $O(N^2)$ complexity).
- ERP-based: Requires training one predictor per model ( $O(N)$ complexity). Adding a new model only requires training one new predictor, not retraining the entire system.
Theoretical Justification: The paper provides a proposition showing that if reward distributions are sub-Gaussian, the difference in expected rewards directly correlates with the win-rate probability between two models.

4. Experimental Results

Predictability (R² Scores):
- Linear models achieved high $R^2$ scores (ranging from 0.19 to 0.59 depending on the category and model) in predicting expected rewards on held-out test data.
- This indicates that prompt embeddings contain sufficient information to estimate how well a specific model will perform on a specific task.
- Performance was consistent across different model families (Llama 3.1, Gemma 2, Gemma 1) and sizes (8B to 70B).
Routing Performance:
- Pareto Frontier: The ERP-based routing policy established a Pareto frontier (trade-off between cost and regret) that dominated all non-oracle baselines.
- Comparison to Baselines:
  - Outperformed "Same Model" (fixed model) and "Random" routing.
  - Outperformed "Per-Category Best" (an oracle that knows the category and picks the best average model for that category).
  - vs. Zooter: ERP performed comparably to the more complex "Zooter" method (which predicts a softmax distribution over winning models) but with significantly less training complexity.
- Pairwise Prediction: When used to predict which of two models would win on a prompt, ERP achieved AUROC scores nearly identical to logistic regression trained directly on pairwise preference data.

5. Significance and Implications

Efficiency in Inference: This method enables dynamic, cost-aware model routing at inference time. Systems can route easy prompts to small, cheap models and hard prompts to large, expensive models without the latency of generating responses from multiple models to compare them.
Simplification of RLHF/Alignment: The findings suggest that the complex task of learning preference distributions can be simplified to learning expected reward distributions, provided the reward model is well-behaved.
Extensibility: The approach is trivially extensible. As new models are added to a production pool, they can be integrated into the routing system simply by training a new linear predictor, avoiding the need for massive pairwise retraining.
Insight into Reward Models: The success of this method implies that modern reward models, trained on diverse datasets, have learned a consistent "scale" of quality that generalizes across prompts, rather than just relative rankings.

In conclusion, the paper establishes that Expected Reward Prediction is a robust, simple, and scalable technique for optimizing LLM deployment, effectively bridging the gap between prompt analysis and model selection.