Causal LLM Routing: End-to-End Regret Minimization from Observational Data
Authors: Asterios Tsiourvas
Wei Sun
Georgia Perakis
Venue: NeurIPS 2025 poster ()
Keywords: LLMs
Routing
Causal ML
Primary area: deep_learning

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
Asterios Tsiourvas
Wei Sun IBM Research
Georgia Perakis MIT
MIT
Abstract
LLM routing aims to select the most appropriate model for each query, balancing competing performance metries such as accuracy and cost across pool language models. Prior approaches typically adopt decoupled strategy. where 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-feedbuck data, where each query is evaluated by all candidate models, which 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: dlassification-base weighted regret approximation shown to recover the optimal policy at convergence. We further extend our framework to handle heterogeneous cost preferences via nneevva-conditional architecture. Experiments on public benchmarks show that method outperforms existing baselines, achieving state-of-the-art performance across different embedding models.
upper bound, and softmax-
Introduction
LLM routing is emerging research area focused optimizing model selection for each input query, balancing performance and cost across pool of available LLMs. Because L.M performance varies significantly by task and input (Hu al. 2024). as well computational cost [Ong al. 2024], dynamic routing strategies have been proposed to select the most suitable model per query [Shnitzer al. 2023. This challenge becomes even more critical in agentic applications, where multiple LLM calls may be made within a single workflow, making efficient model selection essential for user experience and resource allocation. As LLM deployment scales, routing also contributes to environmental sustainability by reducing unnecessary computation (Singh et al. 2025].
Routing methods can be broadly classified based whether they invoke query. Multi-model approaches include non-predictive routing, which cascades models sequentially from light heavy [Wang al., 2023 Chen et al. 2023), and predictive ensembles, which learn to combine outputs or scores from multiple LLMs (Shekhar et al., [0224] Huang et al. (2025.. While these methods can improve accuracy and robustness, they incur significant computational cost and latency due to repeated inference. By contrast, predictive routing methods, including our approach, select a single LLM for each query by training a router that maps input queries the most appropriate model (Shnitzer et al. 2023 Ong et al. 2024 Somerstep et al. 2025) This framework offers more scalable and cost-efficient solution, particularly in settings where minimizing latency and compute resources is critical.
multiple models per
common formulation in predictive routing is to maximize utility function of the form Yu (t) a. (t) À (t), where (t) and (t) denote the quality (e.g.. accuracy) and cost for model t on given query x, and X > captures user cost sensitivity, or willingness-to-pay Existing methods
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

typically adopt decoupled approach: separate predictors are trained for each metric, and routing performed by selecting the model with the highest estimated utility. However, decision quality is highly sensitive to these predictors, and errors can compound, especially when incorporating additional metrics (e.g., latency, faithfulness, alignment), increasing complexity and uncertainty.
Another fundamental limitation in the existing literature on predictive routing is its reliance on full-feedback datasets. Prior work [Ong et al. 2024) [Somerstep al. 2025 assumes access to data where each query has been evaluated by all available LLMs. This assumption is impractical: (i) the computational and monetary cost of exhaustively querying all models is prohibitive; and (ii) the rapid pace of LLM development makes it challenging to maintain comprehensive, up-to-date evaluation datasets. In contrast, observational data, where each query is evaluated by only one model, is readily available from real-world LLM deployments. making it far more scalable than full feedback datasets. However, introduces treatment bias from historical routing policies, which can lead to suboptimal decisions if not properly addressed (Swaminathan and Joachims 2015 Künzel et al. 2019].
To the best of our knowledge, this is the first work to (i) learn LLM routing from observational data and (ii) introduce an integrated learning framework for routing. Our main contributions are:
We propose causal end-lo-end framework that learns routing policies by directly minimiz- ing decision-making regret using observational data. Unlike the predominant decoupled paradigm, where various performance metrics (e.g.. accuracy, cost) are first predicted and then used to inform routing decisions, our method integrates prediction and prescription into a unified objective. By optimizing for regret directly, the framework explicitly aligned with the final routing decision quality. Furthermore, it is designed to scale efficiently and leverage readily available observational data, while accounting for treatment bias without requiring costly full-feedback datasets.
As the original regret minimization objective is not directly differentiable, we derive two surrogate objectives to enable end-to-end policy learning. The first is a classification-bass upper bound that reframes regret minimization as a multiclass prediction problem under mild Lipschitz assumptions, allowing efficient training with standard methods. The second is a softmax-weighht regret surrogate that smoothly approximates regret using a softmax distribution and provably recovers optimal decisions at convergence.
We extend our framework to support heterogeneous cost preferences by introducing a unified model that conditions on both the query and the user's s cost sensitivity. Leveraging the affine structure of the utility function, we design an efficient interpolation scheme using only two endpoint models per interval. We theoretically show that the optimal treatment is piecewise constant in the cost parameter and that our architecture can exactly represent the optimal policy, enabling flexible and scalable routing across diverse preferences.
We conduct comprehensive experiments on two public benchmarks, demonstrating that our regret-minimizing and heterogeneous cost-aware approaches consistently outperform exist- ing baselines. Our methods achieve sate-off-hhe-rr performance across both BERT-based and LLaMA-based embeddings, highlighting their robustness and practical effectiveness.
2 Methodology
2.1 Problem Formulation
We consider dataset of observational samples, denoted by D ((z,,t, Ce c:)) where each sample is independently drawn from the joint distribution p(z, t, a, c). Here, I, E X C Rd denotes feature vector, typically an embedding, that characterizes the query; 44 E [T] := 10 T 1} specifies the LLM assigned to the query; Og E R2o denotes a numeric quality score of the LLM's response, such as accuracy, or a preference rating; and c, E Rzu represents the cost incurred by model t; when processing query E4.
Given a query x E X, the objective of an LLM router is to select a model 1 E T that maximizes the cost-aware performance, or utility, defined as v=(()) == (!) X cz().. Here, X 2 0 is user-specified parameter, modeling the trade-off between accuracy and cost. Higher values of yz
corresponding to greater performance, are preferred.

---

2.2 End-to-End Regret Minimization
Our goal to route each prompt x to the LLM such that the decision leads to the highest possible utility (t). To this end, we aim to learn an end-to-eue policy f X' + T that minimizes the
decision- mariinn regret [Feráández-Lone Loria and Provost 2022 Zou 2022
More formally.
f == arg min Regret(f)
where regret is defined as Regret(f) = Bx(xx((X) - Xx(x(x))) =x(Xx(Ex) - Ex(xx(((X)))
(I)
with f == arg max (t()).
We want to point out that unlike full-feedback datasets used in prior routing work, which record outcomes for all models t E T, observational datasets contain only partial feedback, logging the outcome of single model f, selected by historical policies. As result, counterfactual outcomes for unobserved LLMs are missing, making it necessary to estimate them while correcting for treatment bias. We address this challenge of estimating Yx (-) using causal inference techniques, as detailed in Section 2.3
With an accurate approximation (x(-), the empirical regret can be approximated as
1 Σ (ti) -Ya (F(r.))
Regret(f) ~
where to == arg maXTeT Y,, (t) is the estimated optimal decision for query Z;.
its dependence on the discrete routing decision f(;;), which makes the regret non-differentiable. To address this, we introduce two surrogate loss functions that serve as differentiable approximations.
key challenge with the objective Equation 2
Surrogate Loss I: Classification- Based Upper Bound
Our first approach is to derive a tractable upper bound on the regret and directly minimize To do so, we define the following notion of Lipschitz continuity for utility functions over the probability simplex.
Definition 1. Let Y, T + R be estimated utility function assigning scalar utility to each model t E T for given input x. We extend Y, to the probability simplex AITI by defining
Yz(v) = Ep(D),,))
LET
for any p E AT) We say that Y, L-Liscchitz over the simplex with respect the & norm for all P.q E TI,
(²(()) - 82(9 S L- 1-1pp-lli:
Here, the constant L can be taken as L = maxiET (1)-
This condition ensures that small changes in the model distribution lead to bounded changes in expected utility. Since T is finite set and Y, (-) is typically learned via bounded smooth function approximators (e.g., neural networks), it is natural to expect bounded variation in utility values across nearby treatments. Note that when p is one-hot vector e, and q f(x) is probabilistic policy, this setting corresponds to our model selection problem, where we seek to minimize the regret between the optimal choice and stochastic routing decision.
Proposition 1. Suppose the estimated utility function Y, T + R is L-Lipschitz continuous over the probability simplex with respect to the & norm, as in Definition Then, for a policy f + AT that outputs a distribution (x) over T, the regret can be upper bounded by:
Regret (f) 5 L. Σ V2 /2-CE(1,, CE(,,, ((x)),
(5)
where 17 = arg max Yn (1) is the optimal treatment for input x, and CE(C, =
- -log f (24)); denotes the cross-entropy loss.
3

---

This motivates classification- based surrogate objective: rather than modeling the full utility surface, we directly learn a policy f X'' 7 by solving a supervised learning problem, where the target label for each input is the estimated optimal decision 1 Optimizing classification loss d(C), fxx)) serves as a tractable surrogate for minimizing the regret in Equation (2), as it upper bounds the regret under mild assumptions. This formulation reduces policy learning to a multiclass classification task, enabling efficient training using standard techniques.
2.2.1 Surrogate Loss 2: Softmax- Softmax-Weighted Regret
The second proxy directly minimizes the regret using differentiable softmax approximation. Specif- ically, we model the policy function f as a neural network with [7] outputs passed through softmax layer with temperature parameter > O, which makes the regret surrogate in Equation 2 differentiable. The first term of the regret, Ex ((")), approximated as:
Ex(Yx((")] ~
Yz (17), where 17 == arg max Ya, (1).
(6)
IET
The second term, Ex Bx(Xx(f(X))] is estimated by treating the softmax output as distribution over treatments:
T Y2, (t) softmax(f(z)))
Ex(Xx
Combining the two, we minimize the following differentiable surrogate objective:
T Ya, (000)ssffmxx(((x))
min
Σ 1=1
(8)
After training, the learned policy prescribes for each E X' the treatment arg max.eT f(x) We now show that this objective recovers pointwise optimal treatment assignment, thus providing a consistent and differentiable approximation the original regret minimization objective.
X + T be a neural network whose output is passed through a softmax layer with fixed temperature - > O, and define t == arg maxte7 Yz, (t). Then, optimizing the softmax- softmax--weighted surrogate regret objective via gradient descent
Proposition 2. Let f :
ITI Σ
(2,(1)
Yz, (t) softmax(f(z:)))
(9)
leads the model f concentrate all probability mass on the optimal treatment t. That is, at convergence,
softmax(((xx)) + {o i
(10)
otherwise.
2.3 Estimating Counterfactual Utility via Causal Inference
In the previous sections, assumed that we have access to Y-(-). Given the observational nature of the data, the potential utility function Yz(.) (.) is not directly observable. We follow the potential [Ranewor/RRseebbuum outcomes framework Rosenbaum and Rubin [9833, , Rubin [1984] and assume the existence of potential utility function (t). We adopt the following assumptions, which are standard in the causal inference literature.
Assumption (Stable Unit Treatment Value), The potential outcome of one sample independent of the treatment assignments on the other samples. Assumption 2 (Ignorability). The assigned treatments and potential outcomes are independent condi- tional on observed covariates, i.e. t. (V"(("))(' Hx (Hirano and Imbens) 2004).
Assumption 3 (Support). For x E X' such that p(x) > 0, we have that p(t[z) > for each t E T.
In the causal inference literature, counterfactual outcomes can be estimated using various methods. Such examples include the "meta-learner" [Künzel et al., 2019), or the Inverse Propensity Weighting

---

(IPW) estimator [Hrvvizz and Thompson 1952 In this work, we utilize the doubly robust estimator
introduced by Dudik et [ defined as:
(1) (y fr ( + (z)), VEE
(11)
(1(2)
where Å. : + y denotes the direct outcome regression model for treatment 4, p(((2) is the estimated propensity score, and X T is the logging policy observed in the dataset D.
The doubly robust estimator combines an outcome model (x(x)) and a propensity model p(t[z))
yielding consistent estimates if either is correctly specified (Dudik et al., 201T]. It offers a favorable bias-variance trade-off: the propensity model corrects for selection bias, while the outcome model reduces variance by leveraging structure in the data.
Remark While we use the doubly robust (DR) estimator in our experiments due to its favorable bias- tradeoff and strong practical performance, our framework is estimator-agnoste any valid counterfactual estimator can be used to compute utility estimates. This includes more ad advanced approaches that relax or mitigate these assumptions
In the experimental section, we show that ignoring the treatment bias leads to inaccurate counterfactual estimates and causes substantial degradation in routing quality, highlighting the limitations of standard supervised learning approaches that assume full feedback.
Routing under Heterogeneous Cost Preferences
In the previous section, introduced causal end-to-end framework for learning optimal routing policies from observational data, where the objective is to maximize a utility function of the form y Ac, with fixed À > 0 representing the trade-off between accuracy and cost.
In practice, however, user preferences vary, i.c.. different queries may be associated with different values of X. From system design perspective, training and maintaining a separate router for each possible À impractical. In this section, we propose unified approach that supports routing under heterogeneous cost sensitivities. We first present joint model architecture that conditions on both the query and the cost parameter, and then provide a theoretical analysis to justify the proposed design.
3.1 Intevval-ooniitionel Joint Router
We design a neural network f X' x Rz0 + R! that jointly takes a query x E X and a cost sensitivity parameter À E Rgo as input, outputs score vector over available LLMs and the routing decision is then made via:
i= arg max / (x,, A)z- TET
(12)
assume access to a finite, representative set of cost preferences A = (A1.... Am C R2o. For ease of notation we assume that Å₁ < X2 <... < Am In practice, these values may correspond to discrete service quality tiers (e.g., basic, standard, premium) that reflect users' varying willingness to trade off cost for performance. For each À E A, partition the training data by cost preference and estimate X-specific utility Y (t), which forms the basis of our joint interval-conditions architecture.
Training Procedure.
For each À E A, we first train an individual router fx X + R!7l using the methods introduced in Section 2.2
For each interval [Ay, Ay+1lā E 14...,m-1k. we initialize a joint network f(x, A) X'x Tl that uses as input both æ, À E (Ay,))+1).
3. The shared model is fine-tuned to minimize regret over the interval:
min 2n
Σ
Regret (f( A)),
(13)
XE(A,Ag11)
where regret is computed using the doubly robust estimator as described earlier, under the corresponding X-specific utility (t).
5

---

Deployment Strategy. At inference time, given a user-specified cost sensitivity parameter À E Rzo:
I E A, we use the individual model (x)
If A & A, we identify the closest neighbors A, A E A such that À < < < X, and we use the corresponding joint network fx, X) trained to generalize across the preference in (A, X).
3.2 Model Architecture
key component of the heterogeneous cost preference routing setup introduced in Section 3.1 the dilerva-cooddttiona joint model f far A), that for each interval [Ay, Ay+al, interpolates between the two corresponding individual models fa, and fAS+3- Specifically, the architecture is designed to exploit the affine structure of the utility function with respect to X, namely
Ac. This
motivates lightweight parameterization that uses only the two endpoints of the interval [Ay, A1+1l rather than all pre-trained models. Concretely, the joint model that we propose is defined as:
f(z, A) = Linear (Ifx, (x), (x)]+ g())),
(14)
where .] denotes concatenation, and g(A) := Activation[Linear A)) is learnable representa- tion of the cost sensitivity parameter.
This architecture enables smooth interpolation between fx, and fxy+1 within [ A4+1). allowing the router to adapt intermediate values of without requiring an individual model for each one. By conditioning only on the two bounding models, this design achieves computational efficiency and strong generalization across cost preferences. The proposed architecture is illustrated in Figure []
linacur layer
λ
add
yes
30
Enear layer
activa tinn
Interpolate router [à, A]
fA
concall
layer
router fa
router fx
Jearnable
frowen
embeddings
Figure 1: Overview of the proposed interval-conditionel joint router framework. Left: Decision logic for handling a given cost sensitivity parameter Right: Joint router architecture.
3.3 Theoretical Guarantees
We now present theoretical guarantees that justify the structure and training strategy of joint cost- preference router. These guarantees leverage the affine nature of the utility function, which is linear in the cost parameter À for fixed accuracy and cost enabling exact interpolation for specific cost sensitivity across fixed intervals.
Proposition 3 (Piecewise Constant Optimal Policy). Fix query X and assume the estimated utility is affine in X, i.e., ()) a. (t) X ex (1) for all E T. Then the optimal treatment
t" (A) = arg max Y²)(t) 1ET
is piecewise constant in X. That is, the cost parameter R2 can be partitioned into intervals over which the optimal treatment remains fixed. Proposition 4 (Affine Closure of Utility Function). Let >3 < Aj+1 be two adjacent cost values and let [) A111) Suppose the utility function is affine in X, i... (t) (t. A Ct)).. Then

---

for all t E T, the utility at À is a convex combination of utilities at the endpoints:
(1) = a. y (t) + (1 (1- . (t), where a = As+1
A++1 - X3
Corollary (Sufficiency Two Models per Interval). Under the affine assumption, the utility (2) >y+1) can be exactly reconstructed using only the endpoints Y: (t) and & (t). Thus, it is sufficient to use only the two corresponding models fA, and fAz41 for interpolation within the interval.
for any À E D
Proposition (Expressivity of Additive Two-Model joint Architecture). Let À [λ, A5+1]. and suppose that for each E T the utility satisfies (1) = a=(()) X (4). Then the optimal treatment (A) == arg max, (t) can be exactly represented by a softmax policy over a function of the form:
f(z, A) = Linear (Ifx, fay) (x)) +
where g()) is any differentiable embedding of x, and fxy: faymI accurate predictors trained at endpoints X and A1+1-
Implications for Architecture and Training. The theoretical results above provide strong justifi- cation for both the proposed model architecture and the associated training procedure. Proposition shows that the optimal treatment changes only across small number of cost sensitivities, supporting our interval-conditionee strategy for routing. Proposition f guarantees that the utility for any inter- mediate cost preference can be exactly recovered through convex interpolation of endpoint models. estaablishes that our joint model is expressive enough to capture the optimal policy within each interval. Our method provides a principled approach for learning a joint routing model that accommodates heterogeneous cost preferences from observational data.
Finally, Proposition [
4 Experiments
4.1
Datasets
We evaluate methods two publicly available benchmarks for LLM routing: RouterBench Huu et al. 2024 and SPROUT (Somerstep et al., 2025
ReuterBench is standardized benchmark comprising 35,712 prompt-repoons pairs from 11 language models. The prompts are drawn from eight evaluation suites spanning reasoning, factual recall, dialogue, mathematics, and code generation. Each prompt is annotated with model accuracy and execution cost, enabling supervised training and evaluation of routing policies.
SPROUT is larger and more diverse benchmark focused on cost-aware routing, consisting of 44,241 prompts and responses from 13 state-of-the-arr LLMs. The prompts cover six challenging
tasks: GPQA [Rein et al., 2.24). MuSR [Sprague et al. 2023) MMI.U-Pr [Wang et al. 2024
MATH (Hendrycks et al. 2021J. OpenHermes [Teknium] 2023 and RAGBench [Friel et al. 2024
SPROUT includes a predefined split: 80% for training, with the remaining 20% evenly divided
between validation and test sets.
4.2 Embeddings and Model Architecture
To encode input queries into vector representations we generate embeddings using two compact, publicly available language models: BERT-base-unde (2048 dimensions). Each input is passed once through the model, and the final hidden states are mean-pooled to obtain a fixed-length embeddings. These models were selected for their efficiency and suitability for real-time routing.
(768 dimensions) and Llama-3. 2-22-B
The embeddings are processed by a two-layer fully connected network with GELU activations and 200 hidden units per layer. The model is trained with the Adam optimizer (learning rate 1 x 10-4) for up to 10,000 epochs, using early stopping with a patience
100. A softmax output with temperature = 100 is used to control the sharpness of the output probabilities. This architecture is used consistently across all benchmarked methods for fair comparison. For the doubly robust estimator, the same network models the direct outcomes f, (z),, while the propensity scores p(t z) are estimated
7

---

using XGBoost. To reduce variance from extreme inverse propensity weights, we apply clipping the 5th and 95th percentiles. The only architectural modification is for the interval-basee based model, where the softmax temperature is increased to 1000 to enable smoother interpolation across X.
Hyperparameters are summarized in Appendix D
4.3 Methods
We evaluate our proposed routing strategies against range of baselines from the causal machine learning and LLM routing literature. Since both SPROUT and RouterBench provide full-feedback datasets (i.e., responses from all models), simulate observational data by sampling single model per prompt. Specifically, for each prompt, we sample model E T with probability proportional to its accuracy P(r- = 7] c" where o, is the accuracy of model T on that prompt.
As an optimistic oracle, we include a Full-Feedback model that learns a model f X + R using the complete outcome vector for each query and optimizes standard multi-class classification loss. We benchmark against a common decoupled routing strategy denoted Baseline, which independently estimates model accuracy (t) and cost C, (f), without accounting for selection bias. This reflects the approach taken in prior predictive routing methods such as CARROT [Somerstep et al. 2025),
which has demonstrated superior performance over alternatives like RouteLLM [Ong et al. 2024
and RoRF [Jain et al. 2023). J. To adjust for treatment assignment bias, we consider a Regress- and
Compare (R&C) method, which fits outcome models (t) for each treatment f, and selects the action f = arg max, Yg (t). Building on this. we implement a Gausal-CARROV variant by adapting both the parametric and KNN instantiations of CARROT to the R&C framework. We additionally include CF-
Regression, which models f X - RITl and is trained to minimize MSE against the counterfactual utility function from the doubly robust estimator: miny Σi= Σ (Yz, (1) - f((x)))2. Decisions are made by selecting the treatment with the highest predicted value, i.e., { = arg max, f(z)).
Finally, we evaluate our regret- mineemiaiiiiiiioonn methods. CM-Classiffiattinn formulates the task as multi-class prediction over optimal treatments, serving as classífication-base upper bound. RM- Softmax directly minimizes softmax- regret surrogate. In the heterogeneous preference setting, we also assess RM-Interval (Section D. which generalizes across a continuum of cost sensitivities by interpolating between models trained at discrete À values.
Table Comparison of routing methods by causal reasoning and end-to-end training.
Method
Causal End-to-End
Full-Feedback
Baseline
R&C
CARROI- (kNN & Embed)
CF-Regression
RM-Classification (Ōurs)
RM-Sofimax (Ours)
RM-Interval (Ours)
4.4 Evaluations
We evaluate methods in two settings. In the X-specific setting, each cost sensitivity value À {0, 100, 1000} defines a separate routing task, with models trained and evaluated independently. In the heterogeneous preference setting, RM-Interval is trained on a subset {0, 200, ... 1000} and tested on held-out values {100, 300, 500, 700,7009900 900) to assess generalization across cost sensitivities.
We report the average utility across 10 independent trials for each routing method for SPROUT dataset and BERT embeddings in Table where each trial involves randomly sampling observational data and retraining all models. In Figure 2 we also visualize the corresponding accuracy-com curve associated with each method. Additional plots for the rest of the datasets and embeddings, along with detailed router performance are provided in Appendix B
8

---

Accuracy
SPROUT Embedding:
Zoom-in: E [100, 900]
Full range: [0, 1000]
Full-Foedback
82.5
interval
80.0
R&C
77.5
75
SAuuaa-CAAROOY (8mbed)
CF-Repression
Accuracy
75.0
70
GPT-4o
63
GPT-Ao-mini
72.5
Granite 3-20
60
Granke-3-8B
70.0
LLaM6-313-7788
55
LLaMA 2-8
67.5
LLaMA: 3
LLaMA: 3-2
LLaMA-3-3-100
65.0
45
LLAMA-3-4088
a0c0:
0.0022 0.003 0.0004 0.0005
6.001
0.004
0.005
Nixtral-ax78
Cost ($)
(S)
Figure 2: Accuracy-cost trade-off curve for SPROUT with BERT embeddings.
Table 2: Utility on SPROUT with BERT embeddings. Full Feedback serves as an oracle upper bound. The best-performing method for each column is highlighted in bold.
Method
À 0
100
À 200
À 300
À = 400
À = 500
Feedbackk
85.99 83.99+0.17
79.34 19.34
75,55+
72. 72.16-024
69.47 -044
66.97 4:L4T
Baseline
66.88 1.55
664.6210.66
39.15_o.8. 15.85500.7
36.7310.67
34.08 10.57
R&C
83.34 32
76.45
72.40+0.56
68.59=0.69
65.34
63.1940.s0
CaRRAl-CARROT (kNN)
84.52 36
76.55
71.34
67.82_0.43
65.3840.40
63.3840.43
Causal-CARROT (EmbedNet)
83.46
68.73
62.4442.40
56.72-1.68
$4.37
48.91
CF-Regression
83.54
76.65
72.42+0.58
68.95
65.81
63.5840.07
RM-Classification
84.20
77.58 L0.34
73.3610.49
69.84 10.48
66.49 66.4910063
64.0311.02
RM--Sofmax
84,97
77.53 10.81
73.89
10.00
70.47 _0,00
67.38 1..47
65.51
RM-Interval
84.97 1.0.99
777.6010.2
73.8910.00
69.922...7
67.3810.47
65.20 10.20
Method
= 600
700
800
X 900
1000
Full-Feedback
64.79
63. 18.221
61.43+0.27
59.98
58.83
Buseline
51.17
48.79-..67
44..74+1.21
42.58_1.99
38.97+2.0
R&C
60.98
59.01
57.21
55.97
54.33+1.37
Causal-CARROS (KNN)
61.77
60.39
59.07
57.99
56,99
Causal-CARROT (EmbedNet)
46.35+1.58
43.66-3.044
41.81+2.34
37.42-6.86
34,69+4.44
CF-Regression
61.37
59.52-0.54
57.7240072
56.31
54.56+0.71
RM-Classification
61.8411.12
59.68
58.0941.22
56.52
54.85
RM-Softmax
64.03
62.04
60.32
$8.85
56.95
RM- Intervll
64,03
61.54
60.32
58.58
56.95
Our RM-based approaches consistently deliver the strongest performance overall, with RM-Softmax and R.M-Iterrvl standing out for both high utility and low variance. Notably, RM-Interval general- izes remarkably well to unseen budget levels (i.e., odd values), many times even outperforming models trained specifically those points. These results underscore the effectiveness of regret- minimization framework both fixed and variable cost settings.
The standard Baseline method, which reflects the common decoupled approach used in prior work and ignores treatment selection bias, performs the worst across most values of X, underscoring the importance of accounting for treatment bias in observational data. Notably, the simple R& C method and the Causal-CARROT variants (which incorporate causal corrections) achieve substantial improvements over it, validating our claim that bias-award routing significantly improves performance.
The performance gap between CF-Regression and our RM-based methods demonstrates the benefit of an integrated and-to-end approach. Whereas CF-Regression focuses on approximating the counterfactual utility function and then selecting the best model based on predicted outcomes, our methods directly minimize regret, leading to superior and more stable results. Comparing our two surrogate formulations, RM-Softmax generally outperforms IM-Classification variance, indicating the advantage of optimizing differentiable surrogate objective. Finally, among the Causal-CARRes CRROT variants, the kNN version consistently outperforms the EmbedNet variant, suggesting that non-parametric estimators may offer greater robustness in this setting.
both utility and

---

Conclusion
We propose causal end-to-end framework for routing queries LLMs under observational data. Our approach introduces a regret-minimizing objective grounded in counterfactual estimation, enabling principled policy learning that accounts for treatment selection bias without requiring full-feedback data. Unlike prior approaches that rely on decoupled prediction of accuracy and cost, where errors can compound, our method directly optimizes the decision objective. To support heterogeneous user preferences, we develop interval-conditionel routing architecture that generalizes across continuum of cost-sensitivity parameters. Theoretical analysis provides guarantees on interpolation sufficiency and regret bounds, while empirical evaluations on public routing benchmarks demonstrate that our methods consistently outperform strong baselines, including recent routing algorithms, across multiple embedding models. Future work includes extending the framework to accommodate additional user-defined metrics or hard constraints that cannot be readily incorporated as soft penalties in the objective. Another promising direction is to explore online or adaptive routing in dynamic environments, as well as extending causal regret minimization to multi-turn settings.
References
Lars Buitinck, Gilles Louppe, Mathieu Blondel, Fabian Pedregosa, Andreas Mueller, Olivier Grisel, Vlad Niculae, Peter Prettenhofer, Alexandre Gramfort, Jaques Grobler, Robert Layton, Jake VanderPlas, Arnaud Joly, Brian Holt, and Gaël Varoquaux. API design for machine learning software: experiences from the scikit-learn project. CoRR, abs/1309.0233 2013. URL http //arxiv.org/ ha////bs///110 0238 0238
Lingjiao Chen, Matei Zaharia, and James Zou. Frugalgpt: How to use large language models while reducing cost and improving performance. arXiv preprint arXiv:2305.05 2023.
Dheeru Dua and Casey Graff. UCI Machine Learning Repository, 2017. URL http:://acccive ics. uci. edu/ml
Miroslav Dudík, John Langford, and Lihong Li. Doubly robust policy evaluation and learning. arXiv preprint 103 4601, 2011.
Carlos Fernández-Loria and Foster Provost. Causal classification: Treatment effect estimation vs. outcome prediction. Journal of Machine Learning Research, 23(59):1-355 2022.
Robert Friel, Masha Belyi, and Atindriyo Sanyal. Ragbench: Explainable benchmark for retrieval- augmented generation systems, 2025. URL https://arxiv. org/abs/2407. 2024.
Gurobi Optimization, LLC. Gurobi Optimizer Reference Manual, 2023. URL https://ww/
gurobi. com.
Dan Hendrycks, Collin Burns, Saurav Kadavath, Akul Arora, Steven Basart, Eric Tang, Dawn Song, and Jacob Steinhardt. Measuring mathematical problem solving with the math dataset. arXiv preprint arXiv:2103.0388 2021.
Keisuke Hirano and Guido W Imbens.
The propensity score with continuous treatments. Applied Bayesian modeling and causal inference from incomplete-dati perspectives, 226164:73-84, 2004.
Daniel G Horvitz and Donovan Thompson. A generalization of sampling without replacement from a finite universe. Journal of the American statistical Association, 47(260):663-685, 1952.
Qitian Jason Hu, Jacob Bieker, Xiuyu Li, Nan Jiang, Benjamin Keigwin, Gaurav Ranganath, Kurt Keutzer, and Shriyash Kaustubh Upadhyay. Routerbench: A benchmark for multi-llm routing system. arXiv preprint arXiv:2403. 12031,, 2024.
Keke Huang, Yimin Shi, Dujian Ding, Yifei Li, Yang Fei, Laks Lakshmanan, and Xiaokui Xiao. Thrifillm: On cost-effective selection of large language models for classification queries. arXiv preprint arXiv:2501. 25t.:22000..0 2025.
D. Jain, T.-Y. Tung, and T. H. Kofman. RoRF: Routing on random forests. notdiamond. ai/blog/ /orff 2023. Accessed: 2025-01-02.
https: : www
10

---

Sören R Künzel, Jasjeet S Sekhon, Peter J Bickel, and Bin Yu. Metaleamers for estimating heteroge- neous treatment effects using machine learning. Proceedings of the national academy of sciences, 116110)4156-4165 2019.
Isaac Ong, Amjad Almahairi, Vincent Wu, Wei-Lin Chiang, Tianhao Wu, Joseph E Gonzalez, M Waleed Kadous, and lon Stoica. Routellm: Learning to route Ilms from preference data. In The Thirteenth International Conference Learning Representations, 2024.
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Kopf, Edward Yang. Zachary Detito, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang, Junjie Bai, and Soumith Chintala. Pytorch: An imperative style, high-performance deep learning library. In Advances in Neural Information Processing Systems 32, pages 8024-8035. Curran Associates, Inc., 2019. URL http:///: .ww.rpppereeeuuupppppppp.p..p.pp cc/paper/ 9015 2015-pytorch- an an- in-imperative style- high- wwwprerrrinneeeeeeeeeeeeereaeeae learning -iibarr...... pdf
David Rein, Betty Li Hou, Asa Cooper Stickland, Jackson Petty, Richard Yuanzhe Pang, Julien Dirani, Julian Michael, and Samuel R Bowman. Gpqa: A graduate-level google-proof q&a benchmark. In First Conference on Language Modeling. 2024.
Paul R Rosenbaum and Donald B Rubin. The central role of the propensity score in observational studies for causal effects. Biometrika, 70(1):41-55, 1983.
Donald Rubin. Bayesianly justifiable and relevant frequency calculations for the applied statistician. The Annals of Statistics, pages 1151-1172, 1984.
Shivanshu Shekhar, Tanishq Dubey, Koyel Mukherjee, Apoorv Saxena, Atharv Tyagi, and Nishanth Kotla. Towards optimizing the costs of Ilm usage. arXiv preprint arXiv:2402.01742, 2024.
Tal Shnitzer, Anthony Ou, Mírian Silva, Kate Soule, Yuekai Sun, Justin Solomon, Neil Thompson, and Mikhail Yurochkin. Large language model routing with benchmark datasets. arXiv preprint arXiv:2309.15789 2023.
Aditi Singh, Nirmal Prakashbhai Patel, Abul Ehtesham, Saket Kumar, and Tala Talaei Khoei. A survey of sustainability in large language models: Applications, economics. and challenges. In 2025 IEEE 15th Annual Computing and Communication Workshop and Conference (CCWC), pages 00008-00014. IEEE, 2025.
Seamus Somerstep, Felipe Maia Polo, Allysson Flavio Melo de Oliveira, Prattyush Mangal, Mírian Silva, Onkar Bhardwaj, Mikhail Yurochkin, and Subha Maity. Carrot: cost aware rate optimal router. arXiv preprini urXiv:2502.0326) 2025.
Zayne Sprague, Xi Ye, Kaj Bostrom, Swarat Chaudhuri, and Greg Durrett. Musr: Testing the limits of thain-of-thought with multistep soft reasoning. arXiv preprint arXiv:2310. 16049, 2023.
Adith Swaminathan and Thorsten Joachims. Batch learning from logged bandit feedback through counterfactual risk minimization. The Journal of Machine Learning Research, 16(1):1731-1755, 2015.
Teknium. Openhermes 2.5. 2023. Accessed: 2025-01-30.
Guido Van Rossum and Fred I.. Drake. Python 3 Reference Manual. CreateSpace, Scotts Valley, CA, 2009. ISBN 1441412697.
Hongyi Wang, Felipe Maia Polo, Yuekai Sun, Souvik Kundu, Eric Xing, and Mikhail Yurochkin. Fusing models with complementary expertise. arXiv preprint 10.00:2310.015422 2023.
Yubo Wang, Xueguang Ma, Ge Zhang, Yuansheng Ni, Abhranil Chandra, Shiguang Guo, Weiming Ren, Aaran Aruiraj, Xuan He, Ziyan Jiang, et al. Mmlu-pro: A more robust and challenging multi- task language understanding benchmark. In The Thirty-eight Conference on Neural Information Processing Systems Datasets and Benchmarks Track, 2024.
11

---

Hao Zou, Bo Li, Jiangang Han, Shuiping Chen, Xuetao Ding, and Peng Cui. Counterfactual prediction for outcome-oriented treatments. In Kamalika Chaudhuri, Stefanie Jegelka, Le Song, Csaba Szepesvari, Gang Niu, and Sivan Sabato, editors, Proceedings of the 39th International Conference on Machine Learning, volume 162 of Proceedings of Machine Learning Research, pages 27693-27706. PMLR, 17-23 Jul 2022. URL https: httpss:/pproceedingg mlr press/v162. zou22a. 1m
URLpttpss:
12

---

A Proofs of Propositions
Suppose the estimated utility function Yg T-Ri L-Lipschitz continuous over the probability simplex with respect to the & norm, in Definition Then, for policy f :.X A!7I that outputs a distribution f(x over T, the regret can be upper bounded by:
Proposition 1.
Regret (f) S.
/2 CE(
(15)
where to arg m Ye, (t) is the optimal treatment for input Hy, and CE(E),
- logg (z))ee denotes the cross-entropy loss.
Proof. Let E A77 denote the one-hot distribution over the optimal treatment G. By the definition of regret:
Regret( = Σ (
(16)
Using the L-Lipschitz continuity of Yz, (-) under the l norm:
(e, &z, (f(x))) LL lle, f(x))111
(17)
Applying Pinsker's inequality:
- f(=))la S /2 KL(e;; flx)))) /2 CE( f(x)),
(18)
where the equality follows because KI. divergence from one-hot distribution probability vector reduces to cross-entropy. Combining the above:
Σ /2 5/(t
Regret (1) SL n
(19)
which completes the proof.
Proposition 2. Let f X Rl7 be neural network whose output is passed through softmax layer with fixed temperature 0, and define t == arg maxieT Yz, (t). Then, optimizing the following objective using gradient descent
min ni
12,
Σ &z (t) softmax(f(x:)).
(20)
leads the model f place all prohability mass on the optimal treatment 47 That is, at convergence,
softmax(f(zi))e {
if 17, otherwise.
(21)
Proof. Let Yo, RITI denote the vector of estimated potential outcomes for input z;, and let f(z;) RI7 be the output of the neural network before the softmax layer. The objective for single instance v; can be written as minimizing the regret surrogate:
Yn. (t) softmax(f
(22)
This is equivalent to maximizing the inner product:
æイ softmax
(23)
13

---

Let us denote == softmax(f(zi)) E AITI-I the probability simplex. We now show that the inner product P increases each gradient step. Since p softmax(f we can compute the gradient of the objective with respect to f(x;) as:
V f(x,)) (Yz, softmax (f (x+))) Jaoitmax (f(x))) Ye..
(24)
where Jsoftmax (1 (x:)) is the Jacobian of the softmax function, given by:
(fottmax(f(xx)) softmaa (f (f(z)) softmax(f(zi)) (ô.,, softmax(f((;)))) (25)
Of(x))
This gradient direction corresponds to increasing the logit value of actions with higher Yz, (f) and decreasing those with lower values, pushing the softmax distribution toward the mode of Y In other words, the gradient ascent step increases the inner product at each iteration k:
(xz)) (k+1)
(fftmax((((((;))))
(26)
Since Y,,, is fixed and the softmax is smooth and bounded, this sequence is monotonically increasing and converges to the maximum possible value:
(Yzz softtox((((/))/// (xxx))) + max (I,(1) (t)),
(27)
which implies:
(oftmax(((zi)ee + 1 ift=
(28)
0 otherwise.
Thus, the regret surrogate converges to zero:
&g,41) ((4) - (V softmax((((z)))) + 0,
(29)
and the learned policy selects the treatment maximizing the estimated outcome.
Proposition 3 (Piecewise Constant Optimal Policy). Fix query E X and assume the estimated
utility function is affine in X, i.e., (²) dx-(t) A (t) for all E T. Then the optimal treatment
(A) arg max (t)) TET
piecewise constant in That is, the budget space Ry0 can be partitioned into intervals over which the optimal treatment remains fixed.
Proof. For fixed each (t))) is an affine function of X. The pointwise maximum of finite collection of affine functions is piecewise affine, and the argmax corresponds to the highest line at each A. Since each pair of lines can intersect at most once, the number of intervals over which single treatment optimal is bounded by [7] 1. Therefore, t" (A) changes only at these intersection points and remains constant within each interval.
Proposition 4 (Affine Closure of Utility Function). Let Ay < Ay+1 be two adjacent budget values and let X [Ay, Aj+1]. Suppose the utility function is affine in A:
(1) = (1) - A (x))
Then for all f T, the utility at À is convex combination of utilities at the endpoints:
22 = a (1) + a) a) ()n where a = Aj+1
Aj+1 Xj
Proof. We expand each term:
(2)(1) = az()) - ->, (t), G (t) = (t) - AU+1 Co (t).
14

---

Then:
a (1- ( = (ax(t) - (ycx(t)) + (1- -1- (ax(t) (a)1 (x+1cx(t)
= dz)(t) - (adj + (1 (1- (11-a)4 a)1
= az((1) - x.)(()) = ((),
since:
ady + (1-a)))+1==
Corollary (Sufficiency of Two Models per Interval). Under the affine assumption, the utility (t)
for any Aj:A can be exactly reconstructed using only the endpoints Y (1) and A (t). Thus, it is sufficient use only the two corresponding models Sa, and says2 for interpolation within the interval.
Proof. This follows immediately from the statement of Proposition 4.
Proposition (Expressivity of Additive Two-Model joint Architecture). Let E [A5, A 1]. and suppose that for each E T the utility function satisfies (() az (t)(t) - X- Oz (t). Then the optimal treatment 8" (A) = arg max, (2(()) be exactly represented by softmax policy over function of the form:
far X) Linear (1))(1)) fAz+1 (x)) g())),
where g(A) is any differentiable embedding of x, and fx,, fayma accurate predictors trained at endpoints X and A++1-
Proof. From Proposition 4 the utility (t) is convex combination of Y (t) and Y (t). If the network f(x X) linearly combines the outputs of fxy (x) and fxy+s (z), then its scores can match (t) up scalar transformation. Applying softmax preserves the argmax.
Including (A) allows the architecture learn any additional monotonic reweighting of the interpola- tion, ensuring the output scores can be shaped to approximate the true utility surface exactly. Thus, the architecture can represent the optimal policy within each interval.
B Additional Results
In this section, we present the additional plots for the rest of the datasets as well as the exact values of utility for value A. We begin by presenting the rest of the figures.
B.1
Additional Figures
Accuracy vs. Cost Dataset: ROUTERBENCH Embedding: BERT
Zoom-in:
range: [0. 1000]
Full-Feedback
Clsssfication
70
intervarval
Baseline
R&C
65
Cousal-CARROT (dNN)
Accuracy (86)
Causal CARROT (Embed)
CF-Regrecsion
60
Accuracy
WizardLM-11 138
Clbude-nntant
Claude-1
55
Claude-2
GPT-3.5
GPT-4
50
30
45
Mixtral-8x77
20
Y--4pp
0.0006
00000000050
00100.001.00
Cost ($)
Cost (5)
Figure Accuracy- cost trade-off curve for RouterBench with BERT embeddings.
15

---

Accuracy vs. Cost
Dataset: ROUTERBENCH Embedding: LLAMA-3.2-18
Zoom-in: À E [100, 900]
Full range: [0. 1000]
Fullf oul-bec
80
RM Cassificatio
Softmax
interval
70
Baseline
RSC
65
Accuracy (%)
60
(Embed)
Regression
WizardLN-138
60
50-
Claude-Instant
Claude 2
Cloude-2
55
GPT-3.5
GPI--
CodeLLaMA- 34B
LLaMA-2-708 708
50
Mistral
Mixtral --x
11-34B
0.0001 0.0002 0.0003 0.0004 0.0005 0.0006
000000,0001000000000001100000000 0.00000130000000000 50250.0030
Cost (S)
Figure 4: Accuracy-cost trade-off curve for RouterBench with LLaMa-3.2-1B embeddings.
Accuracy vs. Cost - Dataset: SPROUT | Embedding: BERT
Zoom-in: A E (100, 900]
Full range: [o. 1000]
<8- Full-Feedback
Gassification
82.5
Softmax
80.0
(INN)
77.5
75
Causal-CARREI (Embed)
Accuracy (96)
Accuracy (%)
CF Regression
Titan
GPT-4o
4o-mini
72.5
Granite: 3:28
80
Granite:3 3-90
10.0
LL3MA-3-1-706
LLaMA-3-
67.5
LLaMA
50
1LaMA-3-2-38
LLaMA-3-3-708
0.0000
0.0001 0.0002 0.000s
0.0004
0.0005
0.000
0.001
0.002
0.004
Cost (S)
Cost ($)
Figure 5: Accuracy-com cost trade-off curve for SPROUT with BERT embeddings.
Accuracy Cost Dataset: SPROUT Embedding: LLAMA-3.2-1B
Zoom-in: 900]
Full range: 1000]
Full-Foedback
Ia-CCassifitttonn
82.5
RM-interval
80.0
Baseline
R&C
Causal CARROT (dNN)
75
CauSSl-CARROT (Embed)
Accuracy (6))
75.0
Accuracy (%)
CF-Regression
litan
GPT-40 40
GPT-4e-mini
Granite--77e
80
Granite 3-88
67.5
LLaMA-3 3-14--3-1--
LLaMA-3-1-46
LLaMA-3-22
50
L.377-3-3.708
62.5
ILaMA-3-4058
0.003
Mixtral 8x78
Cost ($)
Cost ($)
Figure 6: Accuracy- cost cost trade-off curve for SPROUT with LLaMa-3.2 3.2-1B embeddings.
16

---

B.2 Additional Tables
Table Utility for RouterBench with BERT embeddings. Full Feedback serves as an oracle upper bound. The best-performing method for each column is highlighted in bold.
FullFeelluck
56.26
55.0
66:66
68.00
33,8511.
31.8711
30.5111
36.73
2f.
RAC
15.08.14
65.88,13
01.73ess
M..33a
56.
55.47aa
54 M4
51.08
51.88
50.18113
Cassal CARROT 8000)
75.12.0.1
65.12.a5
62.21,85
60.56.1.5
58,94114
57.39
55.561835
54.1..40
52.95144
51.47.a.8
1008
Cauall-AARO CARROT (EmbeilNec)
36.87a:31
47.7lamm
44.06.:10
42.08330
38.85
55.224330
11:25
26.56.m00
28.89
34.5
Regresson
24.1
05.56:317
61,95:14
59:71
57.91
50.561mm
54.572055
53 53.44:144
52.19:2
45.43:133
AM-Chhoifiaaion
76.30-44:
66.43.1.x
62.6
60333
.9.6...
5800 2..0
46
55.39....
34.07..0
52.45
51.14.14
58.50:11
56.18
53.79gm74
51.0021
47.09::4
77 51
60.5ggm
63.01.mm
58.58
54.53g17
57.11am
51.00ge
39.00gm
47.04:
Table Utility for RouterBench with Llama-2 1B embeddings. The best-performing method for each column is highlighted in bold.
A-300
A-
A-
Bisenn
Festbock
2235
59.39111
50,5
M.
36
2124
RAC
60.30.ss:
61.62.131
58.04.s
36.30a11
S4.33a14
$3.02.00
SL.3
50.30
48.44.
48.63.1
46.73,113
Chrsat CARROT HNN)
47
64.72:14
61:50
0833
56.71:1
S188:0
m1.50
51,00.0
89.54:311
Casal-CARRON (EabedNed)
69.30
44.52.m:
38.36a24
33.522
23.50.11
2406
2047-4
15:1
12.32.4
9.64.1
13,48
Ci-Rogmason
07.38
5900
540
518
51.77m
50.80ams
#:10gm
Clasification
72.76-1m
63.22.11
60:58...
36.6.119
35.44
$3.95.
51.44am
50.06.m1s
48/63-131
72.8.0.. o.
61.99
4096,03
"8.36ar
50.5%.n
Table Utility for SPOUT with BERT embeddings. The best best-performing method for each column highlighted in bold.
FullFeesluck
23 15.02:
3939
66366.15
697
66135...0
56.15
50
58.75A
151 1i64
165 Mi1.
45 45,60
RAC
63.34
76.45
12.40..
68.59..m
65.36gm
63.19...
60.9
59.01...
57.21
55.97
54.33.
CHARR-CAROOT (NN)
64.52.035
76.55
11 11.34.03
6738..0
61.88an
61.8a.
61.771a
60 60M
59.00
$1.9
51
56391131
CARROC (Embellie))
63.48.1.
68.75amm
62.44-21
56.72.1.4
S4.3azu
48.91am
46.54x1
43.66.114
41.81a3m
31.42
34.59..4
Regresson
72442.05
6885
60.3
63:51
61.37294
59.52.14
57.72:am
50.1190
54.561m
64.20-..:
73.36...
66:0-.s:
61.00..0
61.8..m
30:8....
38.09
36.52.143
54.5
7339-0.0
67.8..
6551
6410:a00
6104::17
08.32:1
5695:111
Table 6: Utility for SPOUT with LLaMa-3. 2-1B embeddings. The best-performing method for each column is highlighted in bold.
000
A=
1000
Feedback
57.25
36.22
Bradine
61 61122a0
"S.35.11.
-4786...1
-1250
4605.2
144 34.25...
30.30a332
27.99
BBC
$1.34anaa
74.75
64.77331
67.09.ss
64.02cm77
62.33sm
56.8611.
55.6..117
54. 11.10234
Casal-CARROT (NN))
8450
76.30
7:24
18.22ms:
65.7
63.0.e
61.33amm
60.33mm
59.6
$736
56 56.00ass
CARRIUUUR (EmbedNee))
2929
62
53.0
31
31.29amm
15 2147-0
CF-Rogession
62.32
15.55.4
70.24.177
68.25us:
65.5...
63.51es4
61,9.a3
59.90..e
38.19
84.0
71.45 A3)
57.12
71138.11
67.16...1
-6.71mm
RUR
31
36.04
$7.020
St.
We also report the AUC (Area Under the Curve) of the accuracy-cor trade-ofi curve for each method.
This is computed using the sklearn. metrica.... sklearn.metrics. auc function Buitinck al. 2013). AUC provides a single scalar summary of performance that captures how well model balances accuracy and computational cost. higher AUC indicates a more favorable overall trade-off across budgets, making it robust evaluation metric for comparing routing strategies.
17

---

Table 7: Average AUC over 10 trials across datasets and embedding models. Higher is better. Abbreviations: RouterBench, SPROUT.
Method
RB-BERT
RB-LLaMa
SP-BERT
SP-LLaMa
Full-Feedback
0.1839
0.1719
0.1727
0.1655
Baseline
0.0890
0.0134
0.0255
0.0148
R&C
0.1539
0.1108
0.1316
0.1105
Causal-CARROT (KNN)
0.1441
0.1433
0.1642
0.1618
Causal-CARROT (EmbedNet)
0.1376
0.0842
0.1148
0.0768
CF-Regression
0.1412
0.1119
0.1352
0.1233
RM- CMMCClassfficttion
0.1665
0.1389
0.1477
0.1436
RM-
EM-Softmax
0.2286
0.2213
0.3320
0.2444
RMMinteervll
0.2285
0.2196
0.3320
0.2464
C
Contribution of Each Component
The results highlight the impact of causal bias correction, end-to-end training, and regret-focused objectives follows:
Causal Inference (Bias Correction): Our Baseline corresponds to CARROT without any causal correction that is, decoupled predictor trained directly observational data without accounting for treatment bias. consistently performs the worst across all settings. contrast, incorporating causal techniques for bias correction yields substantial gains: both R&C and Causal-CARROYS which integrate causal adjustments into routing, achieve +10-15% routing accuracy at comparable cost, demonstrating that accounting for selection bias is critical for effective routing. These results validate that bias-aware routing significantly improves utility over naive predictors trained on biased data.
End-to-End Learning vs. Two-Stage: Among bias-corrected approaches, end-to-end regret- minimizing methods (RM) consistently outperform the two-stage methods CCF-Regression, Causal- CARROT, R&C). The performance gap (+1-3% routing accuracy comparable cost) demonstrates the benefit of an integrated end-to-end approach: by directly optimizing the decision -uaalityy objective (regret) rather than optimizing intermediate predictions, method achieves superior and more stable results.
Regret Minimization Objective: Even compared to other causal learners, our specific training objective provides an edge. For instance, RM-Softmax (diffcrentiable surrogate) slightly outperforms RM-Classification (upper-bound surrogate) in most cases, with lower variance (+0.5-1% routing accuracy at comparable cost). This highlights the advantage of our softmax-weighted regret objective and its alignment with the true decision loss. Moreover, both RM methods outperform Causal- CARROT and Ceggersssi,, underscoring that minimizing expected regret more effective than surrogate approaches that focus only on accuracy/cost prediction.
D Experimental Details
All experiments were implemented Python 3.8. 12 Van Rossum and Drake 2009]. using PyTorch 2.4. 1ctl211 Paszke al. [2019] and Scikit-learn Buuitinck et al. [2013]. Experiments were conducted on an internal compute cluster equipped with an Intel(R) (eon(R) Platinum 8260 CPU @ 2.40GHz, 512 GB of RAM, and two NVIDIA V100 GPUs with 16 GB memory each.
Prompt Encoding Augmentation To encode input queries into vector representations x, employ two-stage embedding process, First, we enrich each prompt with contextual metadata by prepending a natural language prefix that identifies the source dataset. Specifically, for prompt originating from dataset D (e.g., openhermes/ (eermmess/eeiium)) we construct the following context- augmented input: "The following prompt comes from the datasel D. The prompt is: p". This step provides the embedding model with useful dataset-level context, which is particularly beneficial in multi-domain routing scenarios. The template is flexible and can be extended to include additional metadata if desired.
18

---

Datasets. RouterBench [Hu et al. 2024 is standardized benchmark for LLM routing, comprising 35,712 prompt-response pairs collected from 11 LLMs. The prompts span eight different evaluation benchmarks covering reasoning. factual knowledge, dialogue, mathematics, and code generation. Each prompt annotated with model accuracy and execution cost, enabling response-basee -based decision- making. To maintain consistency in evaluation, we adopt the same split strategy for RouterBench, applied deterministically at the prompt level to ensure reproducibility. SPROUT [Somerstep et al. 2025) is more recent and larger benchmark for cost-aware routing, consisting of 44,241 prompts and responses from 13 state-of-the-ar language models. The prompts are drawn from six diverse benchmarks, including GPQA [Rein et al. 2024], MuSR (Sprague et al. 2023]. ,, MMLU-Pro (Wang 2024 MATH Henddryck et al., 2021J, OpenHermes [Teknium 2023]. and RAGBench [Friel
et al.
et al.
2024 SPROUT includes predefined vrain/aaliattion/ess split, using 80% of the data for training and splitting the remaining 20% equally between validation and test sets.
Neural Router Models. All neural models used in experiments share the same architecture for fairnsss and comparability. We use a 2-layer feedforward neural network with GELU activation and 200 hidden units per layer. Models are trained using the Adam optimizer with learning rate of 10-4 batch size of 128, and a maximum of 10,000 epochs. Early stopping is applied with patience of 100 epochs based on validation regret. The temperature parameter for the bofsex-bbssed regret objective is set to 100, and to 1000 for the interval model to allow smoother gradients across budget intervals.
the same neural architecture described above, trained separately for each treatment 6. For the propensity model p(t(r), we use an XGBoost classifier with the following hyperparameters: maximum depth [1,2,3,5, number of trees = [10,20,50, 100J. The estimated DR scores are clipped to the [5th 95th percentile to reduce the impact of extreme propensity weights and improve training stability.
Doubly Robust Estimation. For the outcome model Åz (x)).
Embedding Generation. We generate sentence-level embeddings using the serr-bbas-uncarn (768-dim) and meta-llama/Il ama-3. 3.2-1BB (2048-dim) models. Embeddings are extracted via mean pooling over the final hidden states and are precomputed in batches using GPU acceleration. These embeddings are fixed during training of all downstream routing models.
RM-Interval Network. The joint model used for budget interpolation is implemented using a small feedforward network that takes as input the concatenation of outputs from fx,, fx.1. and linear embedding of A. The architecture mirrors the router described above and is fine-tuned using the regret objective over -peerfcc training data. The proposed architecture is presented in Figure
Inference latency and computational efficiency Latency is critical in real-time applications, Our method is designed minimize this by using lightweight routing networks (2-layer MLP with 200 neurons per hidden layer) and precomputed light-weight embeddings (Llama-3.2-) and BERT). For instance, with Llama-3.2-1 1B embeddings and an M-.P-based router on the RoutherBench dataset, the end-to-end routing latency is under 2.5 ms on a single A100 GPU for a batch size of almost 25,000. To contextualize this, recent benchmarks show that on 8x A100s, LLaMA-2 or LLaMA-3 70B can generate 100-token paragraph in 1.5 to 2.5 seconds under realistic conditions using optimized inference stacks. Even when using fewer GPUs or smaller models, generation latency typically remains one to two orders of magnitude higher than our routing time. In many cases (e.g.. longer outputs or cold starts), the difference can be significantly larger. Thus, the routing overhead is negligible relative to the cost of LLM generation and does not meaningfully impact user experience. Additionally, the
architecture reduces deployment complexity by avoiding the need train store separate models for different preferences.
19

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
The abstract and introduction accurately reflect the paper's contributions and scope. They clearly articulate the core innovations: (i) learning LLM routing policies from observational data, (ii) introducing an end-to-end regret minimization framework that avoids the limitations of decoupled approaches, (iii) developing differentiable surrogate objectives for training, and (iv) extending the method to support heterogeneous cost preferences via interval-conditionel model. These contributions theoretically and empirically, and are consistent with the stated goals of enabling scalable, bias-aware LLM routing in realistic deployment settings.
Justification:
substantiated in the main text, both
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.
2. Limitations
Question:
Does the paper discuss the limitations the work performed by the authors?
Answer: [Yes]
Justification:
The paper discusses its limitations explicitly. We note that the current frame- work is restricted to affine utility functions, which limits flexibility modeling more complex user preferences (e.g., hard constraints that cannot be regularized in the objective). In the conclusion, we also outline directions for future work, including extending the method to dynamic or online learning settings and adapting it to multi-turn interactions. Additionally, the paper makes certain assumptions to support theoretical guarantees, including a mild Lipschitz continuity condition, and others that standard in the causal inference literature (e.g., unconfoundednees). These are discussed in the relevant sections and are typical for counterfactual utility estimation from observational data.
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
20

---

If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't It acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community, Reviewers will be specifically instructed
penalize honesty concerning limitations.
3. Theory assumptions and proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?
Answer: [Yes]
Justification: All theorems, assumptions, propositions, and corollaries are included in the paper. Their complete proofs can be found in the Appendix.
Guidelines:
The answer NA means that the paper does include theoretical results. the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but they appear in the suppiemental material, the authors are encouraged to provide short proof sketch provide intuition.
Inversely, any informal proof provided the core of the paper should complemented by formal proofs provided in appendix supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Answer: [Yes]
Justification: You can find all the details to reproduce the experiments in Section and in Section of the Appendix. We will also release code in the supplementary material.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution a novel architecture, descríbing the architecture fully might suffice, or if the contribution is specific model and empirical evaluation, it may be necessary to either make possible for others to replicate the model with the same dataset, provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution is primarily a new algorithm, the paper should make clear how reproduce that algorithm.
21

---

(b) If the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution is new model (e.g.. a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset instructions for how to construct the dataset).
(d)
We recognize that reproducibility may be tricky in some cases, in which case authors are welcome describe the particular way they provide for reproducibility, In the case of closed-source models, it may be that access to the model is limited in some way (e.g., registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
Justification: Our paper provides all the details to reproduce the experimental results in Section 4 and Section of the Appendix. Furthermore, the code used is provided in the supplemental materials.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips.cc
sos.uumiiissiooooooocoooo for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https
/ //nips. httpppscccppbllc/uuudddoec htt.bl/cccuuiiesooeeeuuiiiiiieie
for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification:
We report training and evaluation details in Section 4 and Appendix D
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core the paper to a level of detail that is necessary appreciate the results and make sense of them.
The full details can provided either with the code, appendix, or as supplemental material.
7. Experiment statistical significance
22

---

Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [Yes]
Justification: We each experiment 10 times and report mean utility along with standard deviation (I-sigma). The variability reflects random initialization and train/validation splits.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
It should be clear whether the error bar is the standard deviation the standard error of the mean.
It is OK to report I-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables figures symmetric error bars that would yield results that are out range (e.g. negative error rates).
If error bars are reported in tables or plots. The authors should explain the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Question:
Answer: [Yes]
Justification:
We state the computer resources used in Section Diof the Appendix.
Guidelines:
The answer NA means that the paper does include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported the paper (e.g., preliminary or failed experiments that didn't 't make it into the paper).
9. Code of ethics
Question: Does the research conducted the paper conform, in every respect, with the NeurlPS Code of Ethics http://pww..aaaro..ccc///a//aaaa
Answer: [Yes]
Justification: We have reviewed the NeurlPS Code of Ethics, and our work complies with its guidelines in all respects.
Guidelines:
The answer NA means that the authors have reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special cireumstances that require deviation from the Code of Ethics.
23

---

The authors should make sure to preserve anonymity (e.g., if there is special consid-
eration due to laws or regulations their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [Yes]
Justification: The paper highlights potential positive societal impacts of this work. By enabling efficient LLM routing using observational data, our framework reduces the need for full-feedbade supervision, significantly lowering computational and monetary costs. LLM routing itself contributes to more sustainable Al deployment by reducing energy consumption and resource usage. We do not anticipate negative societal impacts, as the method focuses on infrastructure-levell efficiency improvements and does not involve sensitive user data or content generation.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g.. disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms monitor how system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data or models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [Yes]
Justification: Yes. The paper uses publicly available benchmark datasets, and we plan to release our code to support transparency and reproducibility. As our work focuses on routing policies rather than generating model outputs, and does not involve training or releasing new pretrained language models or sensitive datasets, we do not identify any specific risks or misuse concerns requiring additional safeguards.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided relcasing unsafe images.
24

---

We recognize that providing effective safeguards is challenging, and many papers do require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
Justification: All external assets used in the paper, including benchmark datasets, embedding models, are properly cited in the main text. We have respected the terms of use and licenses associated with these resources, and references provided to acknowledge the original creators.
Guidelines:
The answer NA means that the paper does existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version asset used and, possible, include URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g.. website). the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, pappsswiteeeateeee thcode com/ datasets has curated licenses for some datasets. Their licensing guide can help determine the license dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if has changed) should be provided. this information is available online, the authors are encouraged to reach out to the asset's s creators.
13. New assets
Question: Are new assets introduced in the paper well documented and the documentation provided alongside the assets?
Answer: [Yes]
Justification: We introduce new code assets to implement our causal end-to-end LLM routing framework. These assets well documented, with usage instructions and reproducibility details provided to facilitate adoption and verification. Documentation will be released alongside the code repository.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/moddl as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset used.
At submission time, remember to anonymize your assets (if applicable). You can either create anonymized URL or include an anonymized zip file.
14. Crowdsourcing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification: We did not conduct crowdsourcing experiments and research with human subjects.
25

---

Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included the main paper.
According to the NeurIPS Code of Ethics, workers involved in data collection, curation, other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals or equivalent for research with human subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification:
The paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent)
may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
16. Declaration of LLM usage
Question: Does the paper describe the usage of LLMs if it is an important, original, or non-standard component of the core methods this research? Note that if the L.M used only for writing, editing, or formatting purposes and does not impact the core methodology. scientific rigorousness, or originality the research, declaration is not required.
Answer: [NA]
Justification: We used LLMs for grammar, formatting, and rephrasing assistance (e.g.. improving proof presentation). They did not influence the methodological originality the work.
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs as any important, original, non-standard components.
Please refer to our LLM policy for what should or should not be described.
26


Reviews:
[{"review_id": "o7TNb2YJMo", "reviewer": "Reviewer_J5tc", "summary": "The paper introduces a novel approach by combining causal inference with end-to-end regret minimization for LLM routing. It addresses a gap in the literature where existing methods (RouteLLM, CARROT) rely on full-feedback data that is impractical to obtain in real deployments", "questions": "How sensitive is the method to violations of key assumptions (Lipschitz continuity, unconfoundedness, affine utility)?\nHow does the method compare to more sophisticated baselines and recent routing advances?\nHow realistic is the accuracy-proportional sampling for simulating observational data?", "limitations": "Authors mention restrictions to affine utility functions and suggesting extensions to hard constraints and online settings. However, there are a few more missing limitations. There is no discussion of robustness to causal inference assumption violations. There is limited analysis of scenarios where the method would perform poorly. This approach \n\nThe approach relies on strong assumptions including Lipschitz continuity of utility functions (which may not hold for discrete model capabilities), standard causal inference assumptions (unconfoundedness, overlap, SUTVA) that are difficult to verify in practice, and affine utility structures that may not capture complex user preferences. The convergence guarantees for surrogate objectives lack finite-sample bounds. The experimental evaluation relies on simulated observational data that may not reflect the complexity of real-world routing biases, limiting the generalizability of the empirical results.", "rating": 5, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 2, "originality": 4, "strengths_and_weaknesses": "Strengths\nThe theoretical framework is well-constructed and mathematically principled. The two surrogate objectives (classification-based upper bound and softmax-weighted regret) are motivated with formal propositions. Proposition 2's proof that the softmax surrogate recovers optimal treatment assignment provides crucial theoretical grounding. \n\nThe paper also addresses the reliance on expensive full-feedback datasets for current LLM routing. The treatment of observational data with historical routing policies is practically important and represents a significant advancement over existing approaches that assume unrealistic data availability.\n\nEvaluation on RouterBench and SPROUT benchmarks with multiple embedding models (BERT, LLaMA) is comprehensive and good. It includes reasonable baselines and the heterogeneous preference evaluation shows practical applicability.\n\nWeaknesses\nThere is no comparison with recent state-of-the-art methods like the latest RouteLLM variants or more sophisticated ensemble approaches. There is also limited analysis of individual component contributions (causal inference vs end-to-end learning vs regret minimization)\n\nThere are some assumptions made that may not be true for all models. The L-Lipschitz assumption may not hold for discrete model capabilities across diverse tasks. The restriction to linear cost functions limits applicability to more complex preference structures", "was_revised": false, "final_justification": "", "rebuttal": "We want to thank the reviewer for their thoughtful and encouraging feedback. We're grateful for your recognition of our theoretical contributions, the practical importance of using observational data, and the strong empirical results. Below we address your concerns.\n\n**Comparison with Sota Methods and Baselines**\n\nWe agree that evaluating against the latest baselines is important. In our experiments, we included CARROT (Somerstep et al., 2025) as a strong recent baseline, since it has been shown to outperform earlier routing methods like RouteLLM (Ong et al., 2024) and RoRF (Jain et al., 2023). By adapting CARROT to our causal setting (Causal-CARROT), we effectively compared to the sota RouteLLM-class methods. Our results show that our regret-minimization framework consistently outperforms all these strong baselines, achieving sota utility across embeddings and datasets. This comparison demonstrates a performance gain over prior approaches. \n\nRegarding *sophisticated ensemble approaches*: we focused on single-LLM routing rather than multi-model ensembles or cascades. As noted in the paper, multi-model pipelines boost accuracy but incur significant extra cost and latency due to repeated model calls. By contrast, our approach (like other predictive routing methods in the literature) selects a single best model per query for efficiency. We will clarify this distinction. \n\n**Contribution of Each Component**\n\nWe appreciate the request for a deeper analysis of the contributions of each component. While our evaluation does include ablations that isolate these elements, we agree this could be made more explicit in the presentation, and we will revise accordingly. The results indeed highlight the impact of causal bias correction, end-to-end training, and regret-focused objectives:\n\n- **Causal Inference (Bias Correction):** Our Baseline corresponds to CARROT without any causal correction - that is, a decoupled predictor trained directly on observational data without accounting for treatment bias. It consistently performs the worst across all settings. In contrast, incorporating causal techniques for bias correction yields substantial gains: both R&C and Causal-CARROT which integrate causal adjustments into routing, achieve +10–15% routing accuracy at comparable cost, demonstrating that accounting for selection bias is critical for effective routing. These results validate that bias-aware routing significantly improves utility over naive predictors trained on biased data.\n\n- **End-to-End Learning vs. Two-Stage:** Among bias-corrected approaches, our end-to-end regret-minimizing methods (RM) consistently outperform the two-stage methods (CF-Regression, Causal-CARROT, R&C). The performance gap (+1–3% routing accuracy at comparable cost) demonstrates the benefit of an integrated end-to-end approach: by directly optimizing the decision-quality objective (regret) rather than optimizing intermediate predictions, our method achieves superior and more stable results. \n\n- **Regret Minimization Objective:** Even compared to other causal learners, our specific training objective provides an edge. For instance, RM-Softmax (differentiable surrogate) slightly outperforms RM-Classification (upper-bound surrogate) in most cases, with lower variance (+0.5–1% routing accuracy at comparable cost). This highlights the advantage of our softmax-weighted regret objective and its alignment with the true decision loss. Moreover, both RM methods outperform Causal-CARROT and CF-Regression, underscoring that minimizing expected regret is more effective than surrogate approaches that focus only on accuracy/cost prediction.\n\nWe will add a more explicit discussion to quantify each component’s contribution.\n\n---\n\n**Lipschitz Continuity**\n\nThank you for raising this point. We agree that the Lipschitz assumption is a simplification that may not always hold for discrete models. In practice, we did not observe optimization issues that would suggest discontinuities, and training was always stable. We will clarify in the paper that this assumption is primarily used to derive a regret bound, and that our second, more efficient, surrogate, does not rely on it. \n\nAfter your comment, we observed that the assumption applies to the estimated utility function over the softmax output space (i.e., the predicted distribution over models). If utility is computed as an expectation over estimates (e.g., $\\hat{Y}\\_{x}(p) = \\sum_{t=1}^T p_t \\hat{Y}_x(t)$), then Lipschitz continuity holds exactly with $L = \\max_t |\\hat{Y}_x(t)|$. Note that when $p$ is a one-hot vector, as is our case, this corresponds to pure model selection. \n\nTherefore, based on your comment, we revisited the assumption and provide a more appropriate **refined surrogate regret bound**. This establishes a relationship between regret and cross-entropy, offering a more faithful interpretation of Lipschitz continuity. Let:\n\n- $f(x) \\in \\Delta^{T-1}$ the predicted softmax distribution,  \n- $e_{t^{\\*}} \\in \\Delta^{T-1}$ the one-hot vector indicating the optimal model,  \n- $CE(t^\\*, f(x)) = -\\log f_{t^*}(x)$ the cross-entropy loss.\n\n**Step 1: Pinsker’s Inequality**\n\nFor any distributions $p, q \\in \\Delta^{T-1}$, Pinsker’s inequality states: $\\|p - q\\|_1 \\leq \\sqrt{2 \\cdot KL(p || q)}$\n\nSetting $p = e_{t^\\*}$ and $q = f(x)$, we have: $ \\|e_{t^{*}} - f(x)\\|_1 \\leq  \\sqrt{2 \\cdot KL(e\\_{t^\\*} || f(x)) }$\n\n**Step 2: Lipschitz Continuity**\n\nIf the utility function $Y(x, \\cdot)$ is $L$-Lipschitz with respect to $\\ell_1$ (it is when $Y(x, p) = \\sum_t p_t \\hat{Y}\\_x(t)$), then: $ Regret(f(x)) = Y(x, e_\\{t^\\*}) - Y(x, f(x)) \\leq L \\cdot \\|e_{t^*} - f(x)\\|_1.$\n\nCombining: $Regret(f(x)) \\leq L \\cdot \\sqrt{2 \\cdot CE(t^*, f(x))}. $ We will include this in the updated manuscript.\n\n---\n\n**Affine utility (linear cost trade-off)**\n\nIndeed, we assume the utility takes the affine form utility = accuracy – λ · cost. This formulation is the most widely used in prior routing literature and effectively captures a broad range of practical scenarios where users weigh performance against cost. It enables theoretical tractability and yields interpretable decisions. While real-world preferences may sometimes involve more complex or non-linear interactions, we view the affine model as a reasonable default. We will note in the paper that this is a limitation, and that our framework could be extended if the utility function is known or learnable, e.g., by incorporating it directly into the regret formulation. \n\n---\n\n**Causal Assumptions**  \n\nThank you for raising this point. Our framework currently relies on standard assumptions for counterfactual identification, unconfoundedness, overlap, and SUTVA, which are foundational in causal inference and offline policy learning. We acknowledge that these assumptions may not always be fully testable in practice, and their violation may affect utility estimation and downstream routing decisions. While we use the doubly robust (DR) estimator in our experiments due to its favorable bias–variance tradeoff and strong practical performance, our framework is estimator-agnostic: any valid counterfactual estimator can be used to compute utility estimates. This includes more advanced approaches that relax or mitigate these assumptions (e.g., balancing methods [1], sensitivity-aware estimators [2], or methods for limited overlap [3]). We will clarify this flexibility more explicitly in the revised version.\n\n---\n\n**Sample Complexity**\n\nWhile we do not provide formal finite-sample guarantees, we expect the sample complexity of our method to be comparable to standard offline policy learning approaches that rely on utility estimation and classification losses. Specifically, when using the classification-based surrogate, the sample complexity is similar to that of multiclass classification, scaling roughly as $\\mathcal{O}(T/\\epsilon^2)$ for $T$ actions and target regret $\\epsilon$, assuming bounded Lipschitz utilities and proper regularization. For the RM-Softmax, the optimization resembles smooth regret minimization and inherits the sample efficiency of learning with strongly convex, smooth surrogate losses. We will clarify this in the paper.\n\n---\n\n**Accuracy-Proportional Sampling**  \n\nWe agree that real-world logging policies may not follow accuracy-proportional sampling. We use it as a controlled and interpretable proxy that introduces bias while maintaining support over all models. Importantly, our method does not assume this form—it only requires logged model choices and observed features. As an additional ablation, we also experimented with cost-proportional sampling, and observed similar trends in performance. We will include these results in the appendix of the revised version, and we would also be glad to re-run experiments with any alternative sampling policy the reviewer suggests.\n\n---\n\n**Robustness to violations** \n\nWhile we did not analyze robustness explicitly, we expect our method to be resilient to mild violations (e.g., approximate Lipschitzness or minor confounding). We use a doubly robust estimator, which may mitigate some bias. Severe violations, such as extreme selection bias or unobserved confounders, would challenge *any offline learner*. We will expand our discussion. Please also see our response to Reviewer jrAR for an experimental stress-test of our approach.\n\n---\n\n**References**\n\n[1] Stable Weights that Balance Covariates for Estimation With Incomplete Outcome Data, José R. Zubizarreta. JASA, 2015.\n\n[2] Interval Estimation of Individual-Level Causal Effects Under Unobserved Confounding, Kallus et al. NeurIPS 2019.\n\n[3] Overlap in Observational Studies With High-Dimensional Covariates, D’Amour et al. Journal of Econometrics, 2021.\n\n\n---\n\nAgain, thank you for supporting our work with your positive evaluation! We hope our work and answers are to your satisfaction. Please let us know if you have additional questions or comments that we can address."}, {"review_id": "bAuP2LZl2x", "reviewer": "Reviewer_HNML", "summary": "This paper proposes a causal framework to route queries to LLMs using observational (bandit) data. They authors formulate the problem as regret minimization via counterfactual estimation, addressing selection bias without full-feedback data. They also propose methods that can also support diverse cost preferences through an interval-conditioned architecture. The authors provide theoretical guarantees for the proposed algorithms, and also conduct experiments on public benchmarks to show the superior performance of their proposed methods over baselines.", "questions": "Could you please discuss the relationship between your work and bandit-based approaches? Your work is studying the offline setting, but is it related to offline bandits?", "limitations": "Yes", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strengths:\n1. The problem is well-motivated and important. \n2. The paper is well-written and easy to follow.\n3. The authors provide theoretical guarantees of the proposed algorithms.\n4. The authors conduct extensive experiments to show the outperformance of their proposed algorithms over existing methods.\n\nWeaknesses:\n1. The main concern is that the related works are not discussed. I think there are many works on LLM routing, but they are not discussed. The authors mention that previous works rely on full-feedback data, but there are some existing works on multi-armed bandits for LLM routing, which depend on bandit feedback instead of full-information feedback. I think discussions and comparisons with related works are needed to place the position of this work appropriately. For example, some works of the bandit approaches:\n[1] MixLLM: Dynamic Routing in Mixed Large Language Models\n[2] LLM Bandit: Cost-Efficient LLM Generation via Preference-Conditioned Dynamic Routing\n[3] A Multi-Armed Bandit Approach to Online Selection and Evaluation of Generative Models\n2. In terms of the terminology, the \"observational data\" is just bandit feedback in the literature?\n3. There are some typos. For example, in the definition of regret (Eq.(1)), I think $t^{\\star}$ should be $t^{\\star}_X$, which depends on the prompt $X$ ($t^{\\star}_X=\\arg\\max_t Y_X(t)$)?", "was_revised": true, "final_justification": "The authors have addressed my concerns, and I believe this work makes some valuable contributions to the field.", "rebuttal": "We thank the reviewer for their thoughtful and encouraging feedback. We’re glad you found the paper clear, technically sound, and addressing an important problem. We address your questions and concerns below.\n\n---\n\n**Related work** \n\nThank you very much for giving us the opportunity to discuss existing works on multi-armed bandits for LLM routing. As you mentioned, several recent works have explored routing in LLM systems using bandit-based approaches, notably including MixLLM [1], LLM Bandit [2], and Hu et al. [3]. These methods typically adopt an online learning framework with partial feedback and assume either the ability to actively explore routing decisions or access to full-feedback datasets for supervised pretraining. For instance, MixLLM trains model-specific predictors and a meta-controller using query embeddings and leverages online feedback to adapt routing policies over time. Similarly, LLM Bandit learns preference-conditioned policies via contextual bandits, requiring known cost-performance trade-offs for all candidate models and further assuming availability of model identity vectors or quizzing mechanisms to support new model integration. In contrast, our work addresses the offline routing setting where only observational data, i.e., logs of past routing decisions with outcomes from a single deployed model, is available. This setup is more restrictive, as it precludes either online exploration or full-feedback supervision. Furthermore, unlike these prior methods, we directly address treatment selection bias induced by historical policies through the use of doubly robust estimators from causal inference, enabling unbiased counterfactual utility estimation. Our framework also departs from the conventional decoupled paradigm of predicting utility metrics followed by greedy selection: instead, we optimize a fully integrated, regret-minimizing objective that directly aligns learning with the quality of the final decision. Finally, while [3] propose a multi-armed bandit approach for evaluating generative models in an online setting, their work focuses on model selection without context, rather than per-query routing. To our knowledge, our work is the first to formulate LLM routing as a causal decision-making problem in the observational setting, providing both theoretical guarantees and empirical validation under realistic deployment constraints. We will move this discussion in the main text.\n\n---\n\n**Terminology Clarification on \"Observational Data\" and \"Bandit Feedback\"**\n\nThank you for the request to clarify the distinction. While “observational data” and “bandit feedback” are related concepts, they are not strictly interchangeable. In the bandit literature, bandit feedback typically refers to partial feedback settings, where only the outcome of the selected action is observed - this aligns with our setting in that we do not observe counterfactual outcomes. However, most bandit-based methods assume the logged data was collected under a *stochastic* or *exploratory* policy, enabling more reliable off-policy evaluation and coverage across actions.\n\nIn contrast, our work operates in the *fully offline observational setting*, where data is logged under a fixed, possibly biased routing policy, often reflecting real-world deployments. This observational data lacks exploration and often exhibits selection bias due to confounding with input features. As such, our framework draws from causal inference, emphasizing counterfactual reasoning and the use of techniques like inverse propensity weighting and doubly robust estimation to correct for treatment assignment bias.\n\nWhile observational data can be viewed as a form of logged bandit feedback, our formulation reflects stronger real-world constraints and requires tools that go beyond standard bandit assumptions. We will incorporate this clarification in the revised manuscript to better delineate the differences in data assumptions and methodological implications.\n\n---\n\n**Typo**\n\nWe thank the reviewer for catching this typo. You are correct: in equation (1), the optimal treatment should indeed be written as $t_X^\\star := \\arg\\max_t Y_X(t)$, to indicate that it depends on the query $X$. We will correct this in the final version of the paper.\n\n---\n\n**Discuss the relationship between your work and bandit-based approaches**\n\nWe appreciate the reviewer’s insightful question. While our work shares many conceptual similarities with the literature on offline (or batch) contextual bandits, particularly in using partial-feedback data to learn decision policies, it also departs in several ways. Traditional offline bandit approaches often assume that the logging policy is stochastic with known or estimable propensities, and typically optimize policy value via importance weighting or its variants. In contrast, our framework is grounded in the causal inference literature, where we explicitly model and correct for treatment selection bias using a doubly robust estimator that combines outcome modeling and propensity scoring. Moreover, we focus on minimizing regret rather than maximizing expected reward, which aligns the learning objective more directly with decision quality. From a modeling perspective, our approach integrates counterfactual estimation with end-to-end policy learning, enabling direct optimization of a surrogate regret objective under observational data. Thus, while methodologically related to offline contextual bandits, our work contributes a causal treatment of LLM routing with stronger robustness to bias and a novel application in large-scale decision-making under cost–performance trade-offs.\n\n---\n\nWe hope our work and answers are to your satisfaction and would appreciate it if you would consider raising your score in light of our response. Please let us know if you have additional questions or comments that we can address. Thank you!"}, {"review_id": "bIy1xdynkS", "reviewer": "Reviewer_jrAR", "summary": "This paper addresses the critical problem of LLM routing, which aims to select the most appropriate model from a pool of candidates for a given query to balance performance and cost. In detail, the paper introduces a novel causal end-to-end framework for learning routing policies directly from more practical observational data. Furthermore, they extend their framework to handle heterogeneous user cost preferences via an elegant interval-conditioned architecture. Experiments on several benchmarks demonstrate that the proposed methods significantly outperform existing methods.", "questions": "1. As mentioned in the weaknesses, could you provide an ablation study or sensitivity analysis on the performance of your `RM-*` methods under different levels of misspecification for the outcome model ($\\hat{r}_t$) and the propensity score model ($\\hat{p}$)? For instance, one could use simpler models (e.g., Logistic Regression) instead of XGBoost to estimate propensity score model, or intentionally add noise to the training data for these models.\n\n2. Equation *11* suggests that the joint network for an interval is fine-tuned by minimizing the average regret at only the two endpoints ($\\lambda_j, \\lambda_{j+1}$). Is there a risk of \"overfitting\" to these two points and performing poorly for $\\lambda$ values in the middle of the interval? While Proposition *5* suggests the architecture is expressive enough, the training objective seems sparse. Could you provide some intuition or empirical evidence (e.g., plotting performance for many $\\lambda$ values within a single interval) to show that this training scheme indeed leads to good generalization across the entire interval?", "limitations": "yes", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 2, "significance": 3, "originality": 2, "strengths_and_weaknesses": "**Strengths:**\n\n1. The paper addresses a crucial and practical problem. Developing an efficient, cost-aware LLM routing strategy is important for the real-world deployment of large language models. This paper tackles a high-impact problem that is of great interest to the communities.\n\n2. The authors provide a novel framework, Causal LLM Routing, which uses causal inference to handle treatment bias and the end-to-end formulation to avoid error compounding. This moves the field beyond traditional supervised methods or decoupled pipelines.\n\n3. A variety of experiments have been conducted to show the effectiveness of the proposed approach across different baselines.\n\n\n**Weaknesses:**\n\n1. The causal LLM routing framework's performance relies on the counterfactual estimation of the doubly robust estimator, which means it depends on the quality of the outcome model ($\\hat{r}_t(x)$) and the propensity score model ($\\hat{p}(t|x)$). While doubly robust estimator is robust, if both models are poorly specified, the final counterfactual estimates can be inaccurate. It’s necessary to discuss the sensitivity to the quality of these nuisance models.\n\n2. A potential weakness of the proposed framework lies in its scalability as the number of candidate models, $|\\mathcal{T}|$, increases significantly. The training process for the final routing policy (e.g., `RM-Classification` or `RM-Softmax`) requires computing the full vector of counterfactual utilities $\\hat{Y}_{x}(t)$, $t \\in \\mathcal{T}$ for each training sample $x$. This is necessary to determine the optimal action $t^*$ for the regret target and to compute the surrogate loss itself (as seen in Eq. 6 and the objective for Eq. 3). Consequently, the computational complexity of the policy training stage appears to scale at least linearly with $| \\mathcal{T}|$, potentially making it prohibitively expensive for scenarios with hundreds or thousands of models. Could the author provide a more formal analysis or empirical evidence to characterize the computational complexity of the approach? This would help clarify the practical limits of the current framework.", "was_revised": true, "final_justification": "The authors have addressed my concerns, and I will maintain my positive score.", "rebuttal": "We thank the reviewer for their thoughtful and constructive comments. We are pleased that you found our paper to address a high-impact and practical problem, and that you appreciated the novelty of the causal framework, the integration of counterfactual estimation with regret minimization, and the strength of our experimental evaluation. Below, we address each of your concerns in turn and outline additional steps we are taking to strengthen the final version of the paper.\n\n---\n\n**Sensitivity and Model Quality**\n\nWe fully agree that the performance of our framework depends on the accuracy of the components in the doubly robust estimator. As noted, the robustness stems from the fact that consistency is guaranteed if either the outcome model or the propensity model is correctly specified.\n\nTo assess sensitivity, we conducted additional ablations:\n- *Propensity model:* We replaced the XGBoost classifier used to estimate $ \\hat{p}(t \\mid x) $ with logistic regression (scikit-learn, with regularization). The performance of RM-Classification and RM-Softmax remained nearly unchanged across all $ \\lambda $ values, with deviations typically below 1 utility point. This is consistent with the fact that we clip propensity scores between the 5th and 95th percentiles, which improves robustness to model misspecification.\n- *Outcome model:* Replacing also the two-layer MLP outcome model with a simpler linear regression model (scikit-learn, with ridge regularization) led to performance degradation, typically a $20-30\\\\%$ drop in average utility, depending on the benchmark and $\\lambda$. This highlights the importance of flexible outcome modeling in counterfactual estimation.\n\nWe will report the exact utility values for both experiments in the updated manuscript. These findings confirm that our framework, when using DR, is robust to moderate misspecification in the propensity model and highlight the importance of appropriate outcome modeling.\n\nImportantly, we would like to emphasize that **our framework is estimator-agnostic**. While we use DR due to its simplicity, widespread adoption in the literature, and favorable robustness–bias tradeoff, the proposed framework is designed to operate with any estimated utility function $Y_X(t)$. For instance, more advanced counterfactual estimators, such as balancing methods [1], sensitivity-aware techniques [2], or approaches designed to address limited overlap [3], can be seamlessly integrated into our regret formulation. We will clarify this flexibility more explicitly in the revised paper.\n\n---\n\n**Generalization Within Intervals in RM-Interval**\n\nWe appreciate the reviewer’s concern about potential overfitting to the endpoints of each interval in RM-Interval. While Proposition 5 guarantees that the architecture is expressive enough to interpolate between endpoints, empirical validation is indeed important.\n\nTo test this, we evaluated RM-Interval at three unseen intermediate $\\lambda$ values within each interval. For example, between $\\lambda = 100$ and $\\lambda = 200$, we tested the model at $\\lambda \\in \\\\{125, 150, 175\\\\}$, between $\\lambda = 200$ and $300$, we tested the model at $\\lambda \\in \\\\{225, 250, 275\\\\}$ and so on. The resulting estimated utility values looked accurate and followed a smooth and monotonic trend: e.g., utility(100) $\\geq$ utility(125) $\\geq$ utility(150) $\\geq$ utility(175) $\\geq$ utility(200) and so on, as expected from the structure of the utility function.\n\nThis pattern was consistent across all embedding types and benchmark datasets. We will include full numeric results in the updated paper to illustrate the generalization behavior. These findings suggest that our interpolation scheme does not overfit to the endpoints and generalizes effectively across the entire cost-sensitivity interval. We believe this is facilitated by the relatively fine-grained partitioning of the $\\lambda$ space (in steps of 100), which provides sufficient resolution for training. We will add this discussion to the main text, along with a note that substantially coarser intervals (e.g., in larger increments of 1,000) may warrant further investigation.\n\n---\n\n**Scalability with Number of Models $T$**\n\nWe thank the reviewer for raising this important point. We agree that as the number of candidate models $T$ increases, the computational cost of computing the full counterfactual utility vector $\\hat{Y}_x(t)$ grows linearly. However, our framework can incorporate several practical design elements that mitigate this cost and support scalability:\n\n\n- **Shared architectures and fine-tuning:** Both the outcome and propensity models can be implemented using parameter-sharing across treatments, for example, via multi-output neural networks or treatment-conditioned heads. This architecture allows for efficient batching across treatments and makes it possible to fine-tune these components starting from pretrained direct models already used for utility estimation. For instance, if a new model (e.g., LLaMA-8B) is added from a family already represented in the pool (e.g., LLaMA-2-70B), we can initialize its outcome head using the existing model’s representation and fine-tune accordingly. This can reduce significantly training time and make updating the model pool feasible even as $T$ grows.\n\n- **Model pruning during training:** In practice, some models become obsolete or are consistently dominated in utility across queries and cost preferences. These models can be excluded from the candidate pool entirely, reducing the number of direct outcome models that must be maintained and evaluated during training.\n\n- **Model pruning at inference time:** In practical deployments, users often operate under specific constraints, such as latency requirements, minimum accuracy thresholds, or cost budgets, that effectively eliminate certain models from consideration. As a result, it is uncommon to evaluate the full pool of LLMs for every query. Instead, a pruned subset of candidate models is typically selected at inference time based on these constraints. Our method is naturally compatible with such runtime pruning: once trained, the router can be applied to any valid subset of models without retraining. We will incorporate this clarification in the revised manuscript.\n\n---\n\n**References**\n\n[1] Stable Weights that Balance Covariates for Estimation With Incomplete Outcome Data, José R. Zubizarreta. JASA, 2015.\n\n[2] Interval Estimation of Individual-Level Causal Effects Under Unobserved Confounding, Kallus et al. NeurIPS 2019.\n\n[3] Overlap in Observational Studies With High-Dimensional Covariates, D’Amour et al. Journal of Econometrics, 2021.\n\n\n---\n\n\nWe hope our work and answers are to your satisfaction and would appreciate it if you would consider raising your score in light of our response. Please let us know if you have additional questions or comments that we can address. Thank you!"}, {"review_id": "jO6PuNFaYE", "reviewer": "Reviewer_9Lmz", "summary": "This paper proposes a causal end-to-end LLM routing framework based on observational data, aiming to optimize model selection by minimizing decision regret while balancing accuracy and cost. Due to the limitations of observational data, it employs a **Doubly Robust Estimator** from causal inference to estimate counterfactual utility $\\hat{Y}_x(t) $. This estimator integrates a direct outcome regression model with a propensity score model. Since directly optimizing the regret $\\operatorname{Regret}(f) $is non-differentiable due to the discrete nature of routing decisions $f(x) $, the paper introduces two differentiable surrogate loss functions to address this issue. In practice, users’cost sensitivity $\\lambda $may vary per query. To avoid training separate models for each $\\lambda $, the paper proposes an interval-conditioned joint router.", "questions": "1. The paper uses XGBoost to estimate the propensity score $\\hat{p}(t \\mid x) $and applies clipping to mitigate the impact of extreme weights. However, the stability of the propensity score model is not thoroughly analyzed. How robust is this model to variations in data distribution or feature quality, and how do errors in propensity estimation affect routing performance?\n\n2. The method assumes query independence (SUTVA assumption), but in multi-turn dialogues or agentic applications, queries may be interdependent (e.g., one turn's response influences the next). How can the method adapt to such scenarios where query dependencies violate the SUTVA assumption?\n\n3. The paper claims high computational efficiency (e.g., interval interpolation using two endpoint models) and direct optimization of decision regret in an end-to-end manner. However, it does not provide specific data on inference times. In real-time applications, the latency of the router (embedding generation + routing decision) could impact user experience. How does the method ensure low latency to meet real-time requirements?", "limitations": "The paper has several key limitations, including dependence on embedding quality, constraints of observational data quality, and consequences of causal assumption violations. Additional potential issues, such as unquantified computational complexity, further restrict its applicability.", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 2, "strengths_and_weaknesses": "### Advantages\n\n1. Unlike traditional methods requiring \"full feedback data\" (testing every query with all models), which is costly and impractical, this paper uses only \"observational data\" (one model per query). This aligns with real-world scenarios, saving computational and economic costs.\n\n2. Traditional routing predicts each model's accuracy and cost separately, risking errors that accumulate in the final selection. The paper's end-to-end approach directly optimizes decision regret, avoiding error accumulation and improving model selection accuracy.\n\n3. Observational data often has selection bias (e.g., favoring cheaper models). The paper uses causal inference (via a doubly robust estimator) to estimate untested models' performance, correcting bias and enhancing routing reliability.\n\n4. Users have varying cost sensitivities (some prioritize savings, others quality). The paper's interval-conditioned joint router adapts strategies to different budgets without training separate models for each preference, making it efficient and practical.\n\n### Disadvantages\n\n1. The method assumes sufficient data support (each model has some query assignments), but real-world data may heavily favor certain models (e.g., only the cheapest). This can lead to inaccurate counterfactual utility estimates. The paper lacks detailed analysis of robustness under imbalanced data distributions.\n\n2. Although the method is claimed to be efficient (e.g., using two endpoint models for interpolation), the paper does not provide detailed computational cost analysis (e.g., training time or inference latency). In high-throughput scenarios, these costs could be critical.\n\n3. Experiments are conducted on two public datasets (RouterBench and SPROUT) with diverse tasks, but these may not fully represent complex real-world LLM applications (e.g., real-time dialogue or multimodal queries). The simulated observational data (sampled by accuracy) may not reflect actual historical biases.\n\n4. The method relies on causal assumptions (e.g., ignorability, support, and SUTVA). These may not hold in practice, such as when query features \\( x \\) fail to explain model selection (due to unobserved confounders) or when queries are interdependent (e.g., in multi-turn dialogues, where prior model choices affect subsequent outcomes, violating SUTVA). The paper does not thoroughly address the impact of assumption violations.", "was_revised": true, "final_justification": "The response from the authors resolves my concerns. I will maintain my positive rating.", "rebuttal": "We thank the reviewer for the detailed and thoughtful assessment. We are encouraged that you found our contributions well-aligned with real-world deployment constraints, and we appreciate your recognition of the value in leveraging observational data, causal inference, and cost-aware regret minimization. Below, we address each of your concerns and outline the revisions we are making to strengthen the final version of the paper.\n\n---\n\n**Robustness under imbalanced data / limited support**\n\nWe appreciate the reviewer’s insightful concern regarding the reliability of counterfactual estimation in the presence of imbalanced logged data. We would like to emphasize that **our framework is counterfactual estimator-agnostic** as it is designed to work with any estimated utility function $\\hat{Y}_X(t)$. This means that more advanced counterfactual estimators, such as balancing methods [1], sensitivity-aware techniques [2], or approaches specifically designed to handle limited overlap [3], can be seamlessly incorporated into our regret formulation. We will make this flexibility more explicit in the revised manuscript.\n\nThat said, we chose the doubly robust (DR) estimator to demonstrate our framework due to its strong theoretical guarantees and solid empirical performance. DR estimation provides consistent utility estimates as long as either the outcome model or the propensity model is well-specified, making it particularly suitable for biased and observational data where certain treatments (models) may be underrepresented.\n\nTo further validate robustness in such scenarios, we conducted a stress test where we artificially downsampled logged data from higher-utility models (usually GPT-4o in our datasets, retaining only 20-30% of queries routed to them), simulating extreme skew toward cheaper or default models. Despite this imbalance, our routing performance remained stable (didn't observe extremely low utilies, still our models were the best-performing), especially when using the DR estimator with clipped propensity scores. This suggests that our approach is reasonably robust to the types of bias encountered in practice. We will include these new experimental results in the updated manuscript.\n\n---\n\n**Stability and role of the propensity score model**\n\nWe appreciate this important point. As noted in our response to Reviewer jrAR, we tested logistic regression for the propensity model, and observed minimal impact on routing utility, typically <1 point deviation, due to the robustness of the doubly robust estimator and our percentile-based clipping. These results will be included in the revised paper. Additionally, our framework allows for fine-tuning or regularizing the propensity estimator, and in practice, moderate performance is sufficient so long as it preserves relative likelihoods across treatments. \n\n---\n\n**SUTVA and sequential (multi-turn) settings**\n\nThank you for highlighting this important point. Our framework assumes query-level independence, which is standard in most current LLM routing setups in the literature, where each input is routed independently, e.g., in tasks such as question answering, summarization, or retrieval-augmented generation. This aligns with the Stable Unit Treatment Value Assumption (SUTVA), as the outcome for a given query depends only on the selected model and not on other queries.\n\nWe agree that multi-turn dialogue introduces dependencies that could violate this assumption. We will clarify this in the paper and include a discussion in the limitations section. That said, our method could be extended to such sequential settings by redefining the “query” to include recent interaction context (e.g., a dialogue window) and learning routing policies over stateful representations.\n\nHowever, we also want to emphasize that multi-turn routing introduces qualitatively different challenges, particularly around cache efficiency. In many LLM systems, keeping the same model across turns enables KV (key-value) caching, which drastically reduces computation by reusing past activations. Switching models mid-dialogue invalidates this cache, requiring recomputation of the full conversation history, potentially negating any cost savings from routing to a cheaper model. This creates a complex trade-off between model flexibility and inference efficiency, framing multi-turn routing as a sequential decision-making problem with long-term cost implications. While this is a highly compelling direction for future work, we believe it falls outside the scope of the current paper, which focuses on single-turn routing.\n\n---\n\n**Inference latency and computational efficiency**\n\nWe agree that latency is important in real-time applications. Our method is designed to minimize this by using lightweight routing networks (only a 2 layer MLP with 200 neurons per hidden layer) and precomputed light-weight embeddings (Llama-3.2-1B and BERT). For instance, we measured that with Llama-3.2-1B embeddings, RoutherBench dataset, and an MLP-based router, the end-to-end routing latency is under 2.5 ms on a single A100 GPU for a batch size of almost 25,000. \n\nTo contextualize this, recent benchmarks show that on 8× A100s, LLaMA‑2 or LLaMA‑3 70B can generate a ~100-token paragraph in ~1.5 to 2.5 seconds under realistic conditions using optimized inference stacks. Even when using fewer GPUs or smaller models, generation latency typically remains one to two orders of magnitude higher than our routing time. In many cases (e.g., longer outputs or cold starts), the difference can be significantly larger. Thus, our routing overhead is negligible relative to the cost of LLM generation and does not meaningfully impact user experience. Additionally, our interval-conditioned architecture reduces deployment complexity by avoiding the need to train or store separate models for different user preferences. We will add these comparisons and latency numbers to improve clarity and practical relevance.\n\n---\n\n**Benchmark limitations and generalization**\n\nWe appreciate the reviewer’s concern. We chose RouterBench and SPROUT because they cover diverse NLP tasks, including question answering, generation, reasoning, and summarization, and offer a structured and reproducible evaluation framework for LLM routing under varying cost-accuracy trade-offs. While we acknowledge that these datasets may not fully capture all real-world complexities (e.g., multimodal inputs, interactive or agentic settings), they currently are the only publicly available benchmarks, to the best of our knowledge, that include data suitable for LLM routing evaluation. We believe they provide a valuable testbed for principled progress in this space. We would be more than happy to re-run our experiments on additional datasets if the reviewer is aware of others that better reflect practical deployment scenarios.\n\n---\n\n**References**\n\n[1] Stable Weights that Balance Covariates for Estimation With Incomplete Outcome Data, José R. Zubizarreta. JASA, 2015.\n\n[2] Interval Estimation of Individual-Level Causal Effects Under Unobserved Confounding, Kallus et al. NeurIPS 2019.\n\n[3] Overlap in Observational Studies With High-Dimensional Covariates, D’Amour et al. Journal of Econometrics, 2021.\n\n---\n\nWe hope our work and answers are to your satisfaction and would appreciate it if you would consider raising your score in light of our response. Please let us know if you have additional questions or comments that we can address. Thank you!"}]

Author rebuttal:

Meta-review: The paper introduces a novel framework for LLM routing with observational data. The proposed method is theoretically sound, and its efficacy is well supported by extensive experimental validation. As all reviewers are positive about the work, I recommend acceptance.
Decision: Accept (poster)