Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Federate the Router: Learning Language Model Routers with Sparse and Decentralized Evaluations

Baris Askin, Shivam Patel, Anupam Nayak, Andrea Vigano, Jiin Woo, Gauri Joshi, Carlee Joe-Wong

TL;DR

This work tackles the problem of routing queries among multiple remotely hosted LLMs when client data and evaluations are privacy-sensitive and sparsely distributed. It introduces a federated learning framework for LLM routing and develops two router families: a parametric Federated MLP-Router and a nonparametric Federated K-Means-Router, both trained with FedAvg to predict per-model accuracy $\text{acc}(\mathbf{x},m)$ and cost $\text{cost}(\mathbf{x},m)$ and to maximize the utility $U_{\lambda}(\mathbf{x},m)=\text{acc}(\mathbf{x},m)-\lambda\text{cost}(\mathbf{x},m)$. Theoretical results establish convergence and suboptimality bounds for both routers, showing that federated training improves data coverage and reduces routing errors compared to client-local training. Empirically, across RouterBench-Data and ProxRouter-Data, federated routers achieve better accuracy-cost frontiers, with adaptive personalization enhancing robustness under extreme heterogeneity and new-model adaptation enabling lightweight, scalable extension. The findings highlight the practical impact of privacy-preserving, cross-client routing signals for efficient, high-quality LLM deployment in edge and enterprise settings, even as model pools evolve.

Abstract

Large language models (LLMs) are increasingly accessed as remotely hosted services by edge and enterprise clients that cannot run frontier models locally. Since models vary widely in capability and price, routing queries to models that balance quality and inference cost is essential. Existing router approaches assume access to centralized query-model evaluation data. However, these data are often fragmented across clients, such as end users and organizations, and are privacy-sensitive, which makes centralizing data infeasible. Additionally, per-client router training is ineffective since local evaluation data is limited and covers only a restricted query distribution and a biased subset of model evaluations. We introduce the first federated framework for LLM routing, enabling clients to learn a shared routing policy from local offline query-model evaluation data. Our framework supports both parametric multilayer perceptron router and nonparametric K-means router under heterogeneous client query distributions and non-uniform model coverage. Across two benchmarks, federated collaboration improves the accuracy-cost frontier over client-local routers, both via increased effective model coverage and better query generalization. Our theoretical results also validate that federated training reduces routing suboptimality.

Federate the Router: Learning Language Model Routers with Sparse and Decentralized Evaluations

TL;DR

This work tackles the problem of routing queries among multiple remotely hosted LLMs when client data and evaluations are privacy-sensitive and sparsely distributed. It introduces a federated learning framework for LLM routing and develops two router families: a parametric Federated MLP-Router and a nonparametric Federated K-Means-Router, both trained with FedAvg to predict per-model accuracy and cost and to maximize the utility . Theoretical results establish convergence and suboptimality bounds for both routers, showing that federated training improves data coverage and reduces routing errors compared to client-local training. Empirically, across RouterBench-Data and ProxRouter-Data, federated routers achieve better accuracy-cost frontiers, with adaptive personalization enhancing robustness under extreme heterogeneity and new-model adaptation enabling lightweight, scalable extension. The findings highlight the practical impact of privacy-preserving, cross-client routing signals for efficient, high-quality LLM deployment in edge and enterprise settings, even as model pools evolve.

Abstract

Large language models (LLMs) are increasingly accessed as remotely hosted services by edge and enterprise clients that cannot run frontier models locally. Since models vary widely in capability and price, routing queries to models that balance quality and inference cost is essential. Existing router approaches assume access to centralized query-model evaluation data. However, these data are often fragmented across clients, such as end users and organizations, and are privacy-sensitive, which makes centralizing data infeasible. Additionally, per-client router training is ineffective since local evaluation data is limited and covers only a restricted query distribution and a biased subset of model evaluations. We introduce the first federated framework for LLM routing, enabling clients to learn a shared routing policy from local offline query-model evaluation data. Our framework supports both parametric multilayer perceptron router and nonparametric K-means router under heterogeneous client query distributions and non-uniform model coverage. Across two benchmarks, federated collaboration improves the accuracy-cost frontier over client-local routers, both via increased effective model coverage and better query generalization. Our theoretical results also validate that federated training reduces routing suboptimality.
Paper Structure (63 sections, 15 theorems, 73 equations, 24 figures, 1 table, 2 algorithms)

This paper contains 63 sections, 15 theorems, 73 equations, 24 figures, 1 table, 2 algorithms.

Table of Contents

  1. Introduction
  2. Related Work
  3. Federated Learning and LLMs
  4. LLM Query Routing
  5. Problem Formulation
  6. Proposed Federated Routing Methods
  7. Federated MLP-Router
  8. Federated K-Means-Router.
  9. Theoretical Results
  10. Federated MLP-Router
  11. Federated K-Means-Router
  12. Experiments
  13. Federated Learning Improves Generalization with Better Query Distribution Coverage
  14. Federated Learning Even Improves In-distribution Local Performance via Better Model Coverage
  15. When New Models are Added to the Pool
  16. ...and 48 more sections

Key Result

Theorem 5.1

Under standard bounded heterogeneity, unbiased gradient, and smoothness assumptions wang2020tacklingkoloskova2020unified, with step size $\eta = \Theta\!\bigl(\sqrt{N/(\tau T)}\bigr)$, the empirical risk objective in eq:fed_objective converges at the rate Here, $\tilde{\mathcal{O}}$ hides the faster decaying $1/T$ terms and $\sigma^2$ denotes the upper bound on heterogeneity in mini-batch gradien

Figures (24)

  • Figure 1: Federated LLM Router System Description. (i) Clients possess queries from different tasks, and each query is evaluated only on some LMs (in our experiments, we consider the extreme case of one LM per query, where the choice of LM is client-specific and can be highly non-uniform). (ii) Federated Learning setup where clients train their local router on local data, and central server aggregates these local models and returns global model to clients. (iii) Functioning of router, where modelwise accuracy and cost is calculated for the given query, which is routed to the most suitable model that maximizes the objective value.
  • Figure 2: Federated vs. client-local (no-FL) routers on the global test distribution. Accuracy--cost curves obtained by sweeping the trade-off parameter $\lambda$ for MLP-Router (top) and K-Means-Router (bottom). AUC scores are shown in parentheses in the legend. Federated training improves generalization to the global distribution, with larger gains for K-Means-Router.
  • Figure 3: Federated vs. client-local (no-FL) routers evaluated on local test sets. We show representative clients for MLP-Router (top row) and K-Means-Router (bottom row). Numbers in parentheses show AUC scores in the legend. Federated training improves the accuracy--cost frontier even in-distribution, primarily due to increased effective model coverage under sparse and imbalanced per-client query--model evaluations.
  • Figure 4: Adapting to new models. Accuracy--cost frontiers on the global test set for routers trained with three models withheld and after those models are introduced and incorporated via a lightweight calibration step.
  • Figure 5: Adaptive personalization under extreme heterogeneity ($\alpha=0.03$). We plot accuracy--cost frontiers evaluated on local test sets for representative clients, comparing federated training, isolated local training, and the proposed adaptive mixture of the two. Top: MLP-Router. Bottom: K-Means-Router.
  • ...and 19 more figures

Theorems & Definitions (29)

  • Theorem 5.1: Convergence: \ref{['alg:mlp_fl']} - Informal
  • Definition 5.2: Suboptimality
  • Theorem 5.3: Suboptimality - Informal
  • Remark 5.4
  • Theorem 5.5: Suboptimality-Informal
  • Proposition 7.4: Convergence of federated optimization error in \ref{['alg:mlp_fl']}
  • proof
  • Remark 7.5: Implication
  • Definition 7.6: Router suboptimality
  • Remark 7.7
  • ...and 19 more