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

arXiv:2601.22318 · cs.LG · Submitted 2026-01-29 · Read on arXiv

Listen

Radio episode about this paper

Transcript

Introduction to the show: ident: AI Radio. Generated commentary on the latest Artificial Intelligence papers.

Tom: Today's paper: "Federate the Router".

Jane: The gist Large language models (LLMs) are increasingly accessed as remotely hosted services by edge and enterprise clients that cannot run frontier models locally,

Tom: First, who's behind it and why it matters.

Paper summary: Tom: To recap, this paper introduces the "Federate the Router: Learning Language Model Routers with Sparse and Decentralized Evaluations" framework, which aims to let clients build a shared routing policy using only their local offline query evaluation data. The core thesis is that centralizing this data for router training is not feasible because it’s fragmented and privacy-sensitive.

Jane: They claim their method enables clients to collaborate and transfer supervision across semantically similar requests by sharing the routing policies they locally train on private datasets. This allows them to train on those diverse, privacy-sensitive query distributions, plus it supports sparse or imbalanced coverage of models.

Lu: The paper tackles two types of routers: a parametric multilayer perceptron router and a nonparametric K-means router, both designed to work under heterogeneous client query distributions and non-uniform model coverage. This flexibility is key for real-world use cases where you don't have everything in one place.

Meng: How does this translate to something practical? They are essentially letting the AI learn how good a model is for a specific type of question without ever seeing all the questions from every other client.

Lalam: It’s like if everyone in the world had a piece of a puzzle, and they could learn how their piece fits with others without showing anyone their whole picture. The goal is that this shared policy helps them find models that balance response quality against inference cost, which is essential when you can't run frontier models locally.

Conclusion: Tom: So, looking at the whole "Federate the Router: Learning Language Model Routers with Sparse and Decentralized Evaluations," what we see is a practical foundation for training LLM routers using data that is private and fragmented across many different users. They show that federated routing improves the accuracy cost frontier over baselines on both global and client-level test distributions.

Jane: The authors validate that federated learning provides a lower suboptimality against local routing, specifically showing a bound of C' (D) s two NF (one/D) delta compared to the suboptimality bound of C' (D i) s two NF (one/D i) delta for a single client.

Lu: This means federated learning offers a lower suboptimality, specifically O(one/sqrt D) versus O(one/sqrt D i), because the federated router optimizes over the dataset D = sum i D i <ref:2601.22318#pg1>. That's a significant statistical improvement when you’re dealing with many clients.

Meng: For someone just listening to this, it means that instead of having one big, centralized router that might not fit any client well, you can have a system where each client contributes to making the overall routing decision better for everyone without compromising their private data.

Lalam: It offers a practical foundation for training LLM routers from privacy-sensitive and fragmented data. It moves us toward a more decentralized way of finding the right model for any given query, which is really useful when models vary so widely in capability and price.

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

cs.LG

Submitted: 2026-01-29

Updated: 2026-10-09

Comments: Accepted to NeurIPS 2026

License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/

Importance score: 83/100

The gist: The gist Large language models (LLMs) are increasingly accessed as remotely hosted services by edge and enterprise clients that cannot run frontier models locally, making routing queries to models

Key concepts

Federated Learning (FL)
A collaborative machine learning approach where multiple clients train a shared model without sharing their raw private data. In this context, clients share learned routing policies instead of their specific query data, allowing the system to learn from diverse, decentralized evaluation sets while respecting privacy constraints.
MLP-Router
A parametric router that learns global estimators to predict expected accuracy and cost for all models based on a query's embedding. It uses a shared neural network trunk and model-specific heads to output predictions for every model, creating a policy that maximizes utility.
K-Means-Router
A nonparametric router that estimates routing utilities by exploiting the spatial locality within the query embedding space. Clients train local K-means models, share their cluster centers with a server, which then computes global centers to guide routing decisions based on proximity.

Terminology

Summary

The gist Large language models (LLMs) are increasingly accessed as remotely hosted services by edge and enterprise clients that cannot run frontier models locally, making routing queries to models that balance quality and inference cost essential <ref:2601.22318#pg6>. This paper introduces the first federated framework for LLM routing, enabling clients to learn a shared routing policy from local offline query–model evaluation data <ref:2601.22318#pg5>.

How it works

The core problem addressed is that existing router approaches assume access to centralized query–model evaluation data, which is often fragmented across clients and privacy-sensitive, making centralizing data infeasible <ref:2601.22318#pg6>. The proposed federated learning (FL) solution allows clients to collaborate and transfer supervision across semantically similar requests by sharing routing policies they locally train on private offline datasets <ref:2601.22318#pg5>. This enables training on diverse and privacy-sensitive query distributions, and additionally supports sparse and imbalanced client coverage over models <ref:2601.22318#pg6>. The framework supports both parametric multilayer perceptron router and nonparametric K-means router under heterogeneous client query distributions and non-uniform model coverage <ref:2601.22318#pg6>.

Federated MLP-Router

For the parametric MLP-Router, the goal is to learn global estimators that predict the expected accuracy and cost for all models given a prompt embedding <ref:2601.22318#pg5>. This is achieved by employing a single MLP-Router consisting of a shared MLP trunk hθ:R demb →R dh that maps a query embedding x to a hidden representation hθ(x), followed by model-specific linear heads which output Aθ(x, m) and Cθ(x, m) for each m ∈ M from hθ(x) <ref:2601.22318#pg5>. The global empirical loss function L(θ) is defined as the average of local objectives weighted with the number of local samples <ref:2601.22318#pg5>. The resulting MLP-Router induced by a parameterization θ is defined as the routing policy that maximizes the utility approximated by the estimators πθ,λ(x) = arg max m∈M Uλ(x, m; θ) <ref:2601.22318#pg5>.

Federated K-Means-Router

The nonparametric K-Means-Router provides a training-free estimator for routing utilities by exploiting locality in the query embedding space <ref:2601.22318#pg6>. The federated training involves clients running local K-means with Klocal centers on its own embeddings and transmitting the resulting centroids along with cluster sizes <ref:2601.22318#pg5>. The server performs weighted K-means over those centroids to obtain Kglobal global centers which are broadcast back to clients <ref:2601.22318#pg6>. At inference time, accuracy and cost statistics of closest cluster center to a query embedding x are used in Eq. (1) for routing <ref:2601.22318#pg5>.

Theoretical Results

The theoretical analysis provides formal statements to support the effectiveness of federated routing, particularly concerning suboptimality <ref:2601.22318#pg6>. Theorem 5.3 shows that the suboptimality of a routing policy πθbi,λ is bounded by C′ Γ(Di) s log 2 NF (1/Di) δ, whereas the federated router's suboptimality is bounded by C′ Γ(D) s log 2 NF (1/D) δ <ref:2601.22318#pg6>. This indicates that federated learning provides a lower suboptimality against local routing O˜(1/√Di) vs O˜(1/√D), as the federated router optimizes on the dataset D = SN i Di <ref:2601.22318#pg6>.

Empirical Evaluation

Empirically, federated routers consistently improve over routers trained on any single client’s data: they generalize better on aggregate query distributions and also improve routing performance on individual clients by leveraging cross-client signal without exposing private client queries <ref:2601.22318#pg6>. The results show that federated training improves the accuracy–cost frontier over client-local baselines on both global and clientlevel test distributions <ref:2601.22318#pg6>. Furthermore, adaptive personalization further improves robustness under extreme heterogeneity, matching and sometimes improving upon the best achievable tradeoff between federated and local routing <ref:2601.22318#pg6>.

Model Expansion and Adaptation

Both router families naturally accommodate new model expansions; for MLP-Router, a new model is added by appending a new head and training only that head by keeping the trunk frozen <ref:2601.22318#pg6>. For K-Means-Router, adding a new model reduces to estimating its accuracy & cost statistics over existing regions of the embedding space <ref:2601.22318#pg6>. The adaptive personalization mechanism interpolates between the federated router and the client’s locally trained router based on their empirical calibration errors on the client’s existing data points <ref:2601.22318#pg6>.

Conclusion

Federated learning offers a practical foundation for training LLM routers from privacy-sensitive, fragmented data <ref:2601.22318#pg6>. The paper validates that federated routing improves the accuracy–cost frontier over client-local baselines on both global and clientlevel test distributions <ref:2601.22318#pg6>. Overall, federated learning offers a practical foundation for training LLM routers from privacy-sensitive, fragmented data >.

Theoretical Results

The theoretical analysis provides formal statements to support the effectiveness of federated routing, particularly concerning suboptimality <ref:2601.22318#pg6>. Theorem 5.

Improvements for AI systems

  1. The system can learn a shared routing policy from local offline query–model evaluation data, enabling clients to learn a shared routing policy from local offline query–model evaluation data which supports sparse and imbalanced client coverage over models.

  2. The architecture supports two router families: parametric MLP-Router and nonparametric K-Means-Router, allowing the system to choose between learning continuous function approximators for complex behaviors or utilizing geometric summaries of query embeddings.

  3. Federated training improves generalization across diverse query distributions by leveraging cross-client signal without exposing private client queries, leading to improved routing performance on both out-of-distribution (better generalization) and in-distribution (increased effective model coverage) data.

  4. The system can adapt to evolving model pools by supporting lightweight extension: for MLP-Router, this involves appending a new head and training only that head, while for K-Means-Router, it reduces to "estimating its accuracy & cost statistics over existing regions of the embedding space."

  5. An adaptive personalization mechanism can be integrated to mitigate performance degradation under extreme client data heterogeneity by interpolating between the federated router and a client’s local router based on empirical calibration errors on the client’s existing data points.

  6. The system achieves a theoretically bounded routing suboptimality, specifically showing that for federated K-Means-Router, Subopt(πb) ≤ 2Lλ Ex∼Dtest∥x − µ(x)∥2 + 2Lλ ∆max + 2Bλs log 2KglobalMδ2 nmin, demonstrating how aggregation improves coverage compared to local learning.

Sources

Related papers