Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Causal LLM Routing: End-to-End Regret Minimization from Observational Data

Asterios Tsiourvas, Wei Sun, Georgia Perakis

TL;DR

The paper tackles LLM routing from observational data by formulating end-to-end regret minimization that directly optimizes decision quality while correcting for selection bias with doubly robust counterfactual estimates. It introduces two differentiable surrogate objectives and extends the routing framework to heterogeneous cost preferences via an interval-conditioned joint router that interpolates between endpoint models. Theoretical guarantees show piecewise constant optimal policies in the budget parameter and exact interpolation under affine utility, while experiments on RouterBench and SPROUT demonstrate state-of-the-art performance across embedding types. The work enables scalable, bias-aware, cost-aware routing without requiring full feedback, with practical implications for efficient and economical LLM deployment.

Abstract

LLM routing aims to select the most appropriate model for each query, balancing competing performance metrics such as accuracy and cost across a pool of language models. Prior approaches typically adopt a decoupled strategy, where the metrics are first predicted and the model is then selected based on these estimates. This setup is prone to compounding errors and often relies on full-feedback data, where each query is evaluated by all candidate models, which is costly to obtain and maintain in practice. In contrast, we learn from observational data, which records only the outcome of the model actually deployed. We propose a causal end-to-end framework that learns routing policies by minimizing decision-making regret from observational data. To enable efficient optimization, we introduce two theoretically grounded surrogate objectives: a classification-based upper bound, and a softmax-weighted regret approximation shown to recover the optimal policy at convergence. We further extend our framework to handle heterogeneous cost preferences via an interval-conditioned architecture. Experiments on public benchmarks show that our method outperforms existing baselines, achieving state-of-the-art performance across different embedding models.

Causal LLM Routing: End-to-End Regret Minimization from Observational Data

TL;DR

The paper tackles LLM routing from observational data by formulating end-to-end regret minimization that directly optimizes decision quality while correcting for selection bias with doubly robust counterfactual estimates. It introduces two differentiable surrogate objectives and extends the routing framework to heterogeneous cost preferences via an interval-conditioned joint router that interpolates between endpoint models. Theoretical guarantees show piecewise constant optimal policies in the budget parameter and exact interpolation under affine utility, while experiments on RouterBench and SPROUT demonstrate state-of-the-art performance across embedding types. The work enables scalable, bias-aware, cost-aware routing without requiring full feedback, with practical implications for efficient and economical LLM deployment.

Abstract

LLM routing aims to select the most appropriate model for each query, balancing competing performance metrics such as accuracy and cost across a pool of language models. Prior approaches typically adopt a decoupled strategy, where the metrics are first predicted and the model is then selected based on these estimates. This setup is prone to compounding errors and often relies on full-feedback data, where each query is evaluated by all candidate models, which is costly to obtain and maintain in practice. In contrast, we learn from observational data, which records only the outcome of the model actually deployed. We propose a causal end-to-end framework that learns routing policies by minimizing decision-making regret from observational data. To enable efficient optimization, we introduce two theoretically grounded surrogate objectives: a classification-based upper bound, and a softmax-weighted regret approximation shown to recover the optimal policy at convergence. We further extend our framework to handle heterogeneous cost preferences via an interval-conditioned architecture. Experiments on public benchmarks show that our method outperforms existing baselines, achieving state-of-the-art performance across different embedding models.

Paper Structure

This paper contains 36 sections, 12 theorems, 39 equations, 6 figures, 7 tables.

Table of Contents

  1. Introduction
  2. Methodology
  3. Problem Formulation
  4. End-to-End Regret Minimization
  5. Surrogate Loss 1: Classification-Based Upper Bound
  6. Surrogate Loss 2: Softmax-Weighted Regret
  7. Estimating Counterfactual Utility via Causal Inference
  8. Routing under Heterogeneous Cost Preferences
  9. Interval-Conditioned Joint Router
  10. Training Procedure.
  11. Deployment Strategy.
  12. Model Architecture
  13. Theoretical Guarantees
  14. Implications for Architecture and Training.
  15. Experiments
  16. ...and 21 more sections

Key Result

Proposition 1

Suppose the estimated utility function $\hat{Y}_x : \mathcal{T} \to \mathbb{R}$ is $L$-Lipschitz continuous over the probability simplex with respect to the $\ell_1$ norm, as in Definition def:lip. Then, for a policy $f : \mathcal{X} \to \Delta^{|\mathcal{T}|}$ that outputs a distribution $f(x)$ ove where $t_i^* := \arg\max_{t \in \mathcal{T}} \hat{Y}_{x_i}(t)$ is the optimal treatment for input $

Figures (6)

  • Figure 1: Overview of the proposed interval-conditioned joint router framework. Left: Decision logic for handling a given cost sensitivity parameter $\lambda$. Right: Joint router architecture.
  • Figure 2: Accuracy–cost trade-off curve for SPROUT with BERT embeddings.
  • Figure 3: Accuracy–cost trade-off curve for RouterBench with BERT embeddings.
  • Figure 4: Accuracy–cost trade-off curve for RouterBench with LLaMa-3.2-1B embeddings.
  • Figure 5: Accuracy–cost trade-off curve for SPROUT with BERT embeddings.
  • ...and 1 more figures

Theorems & Definitions (20)

  • Definition 1
  • Proposition 1
  • Proposition 2
  • Remark 1
  • Proposition 3: Piecewise Constant Optimal Policy
  • Proposition 4: Affine Closure of Utility Function
  • Corollary 1: Sufficiency of Two Models per Interval
  • Proposition 5: Expressivity of Additive Two-Model joint Architecture
  • Proposition 1
  • proof
  • ...and 10 more