跳转至

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

会议: NeurIPS 2026
arXiv: 2601.22318
OpenReview: https://openreview.net/forum?id=MABQINnqfs
领域: 联邦学习
关键词: 大语言模型路由, 联邦学习, 稀疏评估, 参数化路由器, 非参数化路由器

一句话总结

本文提出了首个用于大语言模型（LLM）查询路由的联邦学习框架，使各客户端能够在私有查询和稀疏、偏置的模型评估数据不出本地的前提下，协同训练参数化 MLP 路由器与非参数化 K-Means 路由器，显著提升了准确率与推理成本的帕累托前沿。

研究背景与动机

随着大语言模型（LLM）在编程辅助、智能客服、复杂推理及科学写作等语言密集型任务中的广泛应用，端侧设备与企业端点因本地显存、算力及功耗限制，通常无法直接运行顶尖前沿模型，因而普遍转向通过 API 访问远程托管的云端模型池。然而，实际应用中的用户查询在难度与领域上高度异构，不同模型在能力上限与调用成本（输入与输出 token 计费）上差异悬殊。如果所有查询均调用昂贵的旗舰模型会造成巨大的经济浪费，而单一轻量模型又无法胜任复杂逻辑任务。因此，如何为每个输入查询自适应选择性价比最高的目标模型——即语言查询路由（LLM Query Routing），成为降低大模型推理开销的核心方向。

现有的路由器训练范式普遍依赖于集中式的查询-模型全量评估数据集，即假设所有训练查询都在模型池中的全部候选模型上进行了基准测试与成本记录。然而，在现实商业与个人场景中，这一假设完全不成立：一方面，真实用户查询往往包含高度隐私敏感的商业机密、专有代码或个人对话，受合规性与隐私法规约束根本无法集中上传；另一方面，每个客户端通常只会在其特定的业务领域内尝试极少数模型，形成极度稀疏（每条查询往往仅由一个模型处理）且严重倾斜（某些模型使用率极高、其余模型几乎未被采样）的本地评估日志。若客户端孤立地训练本地路由器，会因查询分布狭窄和模型覆盖不足而导致严重过拟合与高决策方差。

面对分散、私有且稀疏的评估数据，本文的核心切入角度在于引入联邦学习（Federated Learning）机制，在严格保留原始查询文本和模型反馈于本地的前提下，通过跨客户端聚合路由策略来转移语义相似请求的监督信号。核心 idea：将 LLM 路由建模为去中心化离线策略学习问题，通过联邦协作联合训练参数化 MLP 路由器与非参数化 K-Means 路由器，借助跨客户端聚合提升有效模型覆盖与全局查询泛化能力，并结合自适应个性化插值保障高异质性下的局部鲁棒性。

方法详解

整体框架

在包含 \(M\) 个语言模型的候选池 \(\mathcal{M}\) 中，每个模型 \(m\) 对给定查询嵌入 \(\mathbf{x} = \mathrm{Enc}(s) \in \mathbb{R}^{d_{\text{emb}}}\) 具有未知的期望准确率 \(\text{acc}(\mathbf{x}, m) \in [0, 1]\) 和期望推理成本 \(\text{cost}(\mathbf{x}, m) \in [0, c_{\max}]\)。系统定义带有权衡参数 \(\lambda \ge 0\) 的效用函数： $\(U_{\lambda}(\mathbf{x}, m) := \text{acc}(\mathbf{x}, m) - \lambda \cdot \text{cost}(\mathbf{x}, m)\)$ 在联邦路由架构中，全网存在 \(N\) 个客户端，每个客户端 \(i\) 仅持有本地离线数据集 \(\mathcal{D}_i = \{(\mathbf{x}_i^{(j)}, m_i^{(j)}, \widehat{\text{acc}}_i^{(j)}, \widehat{\text{cost}}_i^{(j)})\}_{j=1}^{D_i}\)。这里的关键约束是每条历史查询仅记录了单一模型 \(m_i^{(j)}\) 的采样反馈，且各客户端之间的任务类型分布和模型选择偏置高度非独立同分布（Non-IID）。

整个联邦路由架构分为两种互补范式：参数化 MLP 路由器通过多轮端云通信联合拟合准确率与成本预估网络；非参数化 K-Means 路由器则通过轻量级的聚类中心与单元均值聚合建立分块常数效用场。推理时，客户端无需重新训练网络，只需根据自身实时预算偏好选取 \(\lambda\) 即可完成最优模型的确定性派发。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["本地私有评估数据<br/>每条查询仅含单模型反馈"] --> B["联邦参数化优化<br/>共享主干表征与模型头"]
    A --> C["联邦分层聚类<br/>两阶段几何划分与统计聚合"]
    B --> D["自适应个性化混合<br/>按本地残差动态加权"]
    C --> D
    D --> E["效用最大化路由决策<br/>依据偏好参数选择最优模型"]

关键设计

1. 联邦参数化优化：解耦特征提取与多模型头预估

针对不同模型在稀疏日志下监督信号缺失的问题，参数化路由器设计了一个两层隐藏层（宽度 512，辅以 LayerNorm、GELU 与 Dropout）的共享主干网络 \(h_{\mathbf{\theta}}(\mathbf{x}) \in \mathbb{R}^{d_h}\)，并在其上挂载 \(M\) 个独立的线性预测头，分别预估模型 \(m\) 的准确率 \(\mathsf{A}_{\mathbf{\theta}}(\mathbf{x}, m)\) 与标准化成本 \(\mathsf{C}_{\mathbf{\theta}}(\mathbf{x}, m)\)。在联邦训练过程中，客户端采用 FedAvg 算法，在本地利用私有离线数据对所采样的模型头及共享主干计算均方误差（MSE）损失： $\(\mathcal{L}_i(\mathbf{\theta}) = \frac{1}{|\mathcal{D}_i|} \sum_{(\mathbf{x}, m) \in \mathcal{D}_i} \left[ (\widehat{\text{acc}}(\mathbf{x}, m) - \mathsf{A}_{\mathbf{\theta}}(\mathbf{x}, m))^2 + (\widehat{\text{cost}}(\mathbf{x}, m) - \mathsf{C}_{\mathbf{\theta}}(\mathbf{x}, m))^2 \right]\)$ 通过服务端的多轮加权聚合，共享主干学习到了跨任务的通用查询语义表征，而未在客户端 \(i\) 本地出现过的模型则间接获得了其他客户端的更新监督。理论证明，在标准平滑性与有界异质性假设下，该优化算法在 \(\eta = \Theta(\sqrt{N/(\tau T)})\) 时满足： $\(\min_{t \in \{1,\dots,T\}} \mathbb{E}\left[\|\nabla \mathcal{L}(\mathbf{\theta}^{(t)})\|^2\right] \le \tilde{\mathcal{O}}\left(\frac{1 + A\sigma^2}{\sqrt{N\tau T}}\right)\)$ 其收敛速度享受随参与客户端数量 \(N\) 与本地步数 \(\tau\) 的线性加速，有效弥补了单客户端样本量不足与模型监督缺失的问题。

2. 联邦分层聚类：无须梯度的非参数化流形聚合

针对模型池高频动态更替导致深度网络反复微调开销过大的问题，非参数化 K-Means 路由器将查询嵌入空间划分为局部 Voronoi 胞元，并在胞元内假设各模型的效用近似恒定。在联邦设定下，客户端无需上传原始嵌入，而是先在本地执行 Lloyd 聚类提取 \(K_{\text{local}}\) 个局部中心及其样本权重；服务端接收后运行加权 K-Means 将其合并压缩为 \(K_{\text{global}}\) 个全局中心并广播回传。各客户端再将本地样本映射至最近的全局中心 \(k\)，统计每个 (中心 \(k\), 模型 \(m\)) 对的样本数 \(n_{i,k}^{(m)}\) 及其平均准确率与成本并回传服务端进行无偏加权合并： $\(\mathsf{A}_k^{(m)} = \frac{\sum_i n_{i,k}^{(m)} \bar{a}_{i,k}^{(m)}}{\sum_i n_{i,k}^{(m)}}, \quad \mathsf{C}_k^{(m)} = \frac{\sum_i n_{i,k}^{(m)} \bar{c}_{i,k}^{(m)}}{\sum_i n_{i,k}^{(m)}}\)$ 该方案单轮通信开销仅约 0.5 MB，完全免去了超参数调优与梯度反向传播。理论分析表明其路由次优度上界受限于 \(2 B_{\lambda} \sqrt{\frac{\log(2 K_{\text{global}} |\mathcal{M}| / \delta)}{2 n_{\min}}}\)，联邦聚合使每个胞元中各模型的最小观测样本数从 \(\min_{i, k, m} n_{i,k}^{(m)}\) 剧增至 \(n_{\min} = \min_{k, m} \sum_i n_{i,k}^{(m)}\)，彻底消除了局部冷门模型无统计支撑的硬伤。

3. 自适应个性化混合：校准极端数据异质性

当客户端之间的任务分布呈现极端病态偏置（如狄利克雷参数 \(\alpha \le 0.03\) 时，某客户端仅有生物专业文本而全网其他客户端几乎从未涉足）时，全局聚合的模型策略可能被多数客户端的通用分布主导，反而在该特殊客户端的局部测试集上发生精度倒挂。为此，本文设计了基于本地经验残差的免重训自适应插值机制。客户端利用本地现有数据点直接计算全局估计器与本地隔离估计器的平均绝对误差（MAE），并在推理时根据误差大小反向分配权重： $\(w_a^{(i,m)} = \frac{e(\mathsf{A}^{(m)})}{e(\mathsf{A}^{(m)}) + e(\mathsf{A}^{(i,m)})}, \quad w_c^{(i,m)} = \frac{e(\mathsf{C}^{(m)})}{e(\mathsf{C}^{(m)}) + e(\mathsf{C}^{(i,m)})}\)$ 随后在推理时代入插值预估值 \(w_a^{(i,m)} \mathsf{A}^{(i,m)} + (1 - w_a^{(i,m)}) \mathsf{A}^{(m)}\) 与对应成本计算效用。该机制在客户端特异性领域保留了本地局部几何的高灵敏度，而在本地缺乏采样的稀有模型上则平滑借用全局联邦估计，自动取得了全局泛化与局部保真的最佳折中。

损失函数 / 训练策略

对于 MLP-Router，客户端使用 AdamW 优化器（学习率 \(\eta = 10^{-3}\)，权重衰减 \(3 \times 10^{-4}\)，梯度裁剪阈值 1.0），通信轮数固定为本地 1 个 epoch，批大小为 128。当有新模型加入候选池时，网络主干与原有模型头全部冻结，仅初始化新增模型的线性头，客户端只需在 10% 的校准查询上进行极少步微调；当有新客户端接入系统时，采用带有输出蒸馏正则项的目标函数进行增量微调，防止对已有任务模式发生灾难性遗忘。

实验关键数据

主实验

在 RouterBench-Data 基准（包含 11 个涵盖开源与闭源的 LLM，跨越 8 个代表性任务数据集）上，模拟 \(N=10\) 个客户端、参与率 0.6、任务分布服从狄利克雷分布（\(\alpha=0.6\)）且每条查询仅有一个非均匀采样模型被记录的严格稀疏评估设定。通过在对数网格 \(\lambda \in [10^{-2}, 10^7]\) 内扫描 100 个点绘制准确率-成本曲线，并计算归一化曲线下面积（Normalized AUC，数值越高代表在相同成本预算下准确率越高）。

路由模型范式 训练范式 全局测试集 AUC (mean±std) 集中式全量上限 (Centralized AUC) 相对本地独立基线提升
MLP-Router (参数化) 本地独立训练 (Client-local) \(0.683 \pm 0.003\) \(0.759\) 基准
MLP-Router (参数化) 联邦协同训练 (Federated Ours) \(0.750 \pm 0.003\) \(0.759\) +0.067 (+9.8%)
K-Means-Router (非参数化) 本地独立训练 (Client-local) \(0.655 \pm 0.011\) \(0.747\) 基准
K-Means-Router (非参数化) 联邦协同训练 (Federated Ours) \(0.747 \pm 0.003\) \(0.747\) +0.092 (+14.0%)

消融实验

在客户端数量扩展性（固定总样本量切分至不同客户端规模）以及查询异质性强度（调节狄利克雷浓度参数 \(\alpha\)）下的归一化 AUC 对比如下表所示：

消融维度 具体设置参数 MLP-Router (联邦) MLP-Router (本地) K-Means-Router (联邦) K-Means-Router (本地)
客户端规模扩展性 \(N = 5\) \(0.747 \pm 0.004\) \(0.679 \pm 0.010\) \(0.733 \pm 0.007\) \(0.656 \pm 0.023\)
\(N = 10\) \(0.750 \pm 0.005\) \(0.679 \pm 0.008\) \(0.742 \pm 0.007\) \(0.646 \pm 0.018\)
\(N = 20\) \(0.750 \pm 0.003\) \(0.661 \pm 0.007\) \(0.745 \pm 0.003\) \(0.601 \pm 0.011\)
\(N = 40\) \(0.749 \pm 0.009\) \(0.652 \pm 0.003\) \(0.749 \pm 0.005\) \(0.584 \pm 0.008\)
任务异质性鲁棒性 \(\alpha = 0.03\) (极度异质) \(0.742 \pm 0.007\) \(0.640 \pm 0.015\) \(0.731 \pm 0.007\) \(0.601 \pm 0.033\)
\(\alpha = 0.10\) \(0.743 \pm 0.012\) \(0.666 \pm 0.010\) \(0.742 \pm 0.005\) \(0.614 \pm 0.014\)
\(\alpha = 0.60\) (基准) \(0.746 \pm 0.015\) \(0.672 \pm 0.011\) \(0.741 \pm 0.008\) \(0.650 \pm 0.018\)
\(\alpha = 5.00\) (接近 IID) \(0.748 \pm 0.012\) \(0.669 \pm 0.013\) \(0.749 \pm 0.004\) \(0.625 \pm 0.015\)

关键发现

  • 非参数化路由器的增益更为显著：在全局测试集上，K-Means 路由器的性能提升幅度（+0.092 AUC）超过了 MLP 路由器（+0.067 AUC）。这是因为在本地样本极其有限且偏置时，孤立的聚类中心极其脆弱，且许多流形胞元内完全缺乏某些候选模型的采样反馈；联邦加权中心与统计量累加直接补全了流形死角。
  • 客户端切分越细，联邦优势越明显：当客户端数量从 5 激增至 40 时，每个客户端分配到的本地样本量缩水至原本的 1/8，本地独立路由器的性能急剧崩溃（K-Means AUC 从 0.656 暴跌至 0.584），而联邦路由策略保持了近乎恒定的高水准（0.749），充分展示了联邦池化效应。
  • 联邦性能几乎无损逼近集中式上限：在 RouterBench-Data 与 ProxRouter-Data 上，联邦 MLP 路由器和 K-Means 路由器的 AUC 表现几乎与汇聚所有私有数据的集中式训练曲线完全重合（0.750 vs 0.759；0.747 vs 0.747），实现了隐私零泄露下的无损协同。

亮点与洞察

  • 问题建模视角新颖：敏锐指出了传统大模型路由器在走向真实生产落地时的致命瓶颈——评估数据极度分散且隐私敏感，首次将联邦学习与模型路由结合，开辟了去中心化模型调度的新方向。
  • 兼顾理论保真与轻量部署：同时给出了参数化与非参数化两种成熟架构的联邦协议，不仅为参数化路由器提供了严格的优化收敛性与样本复杂度上界，更使 K-Means 方案实现了 0.5 MB 通信量、零超参数微调的即插即用体验。
  • 模型动态扩容无需重训：提出新模型引入时仅微调线性输出头或局部追加流形统计量的轻量校准协议，成功解决了开源模型与 API 生态高速迭代导致的路由器全量重训成本危机。

局限与展望

  • 依赖离线静态评估日志：当前工作主要针对固定离线数据集设计，未考虑在线推理过程中不断涌现的新增调用反馈对路由器的流式连续更新（Online Continual Routing）。
  • 句子嵌入模型的统一假设：框架假设所有客户端使用统一且静态的句向量编码器（如 all-mpnet-base-v2），若未来客户端需要使用各自私有的轻量编码器，如何对齐多源异构特征空间尚待探索。
  • 对抗鲁棒性未经验证：若联邦参与方中存在恶意客户端故意上传虚假的准确率或成本统计信息（中毒攻击），可能会导致高性价比模型被恶意抑制定向分流，未来需要引入拜占庭容错机制。

相关工作与启发

  • vs. 集中式路由器 (RouterBench / EmbedLLM)：传统方案必须将所有提示词、生成响应与打分上传至中央服务器集中训练网络，且通常假设每条查询在所有模型上都有全量测试数据。本文完全打破这一严苛假设，支持每条查询仅有单模型采样记录且数据不出本地。
  • vs. 联邦微调大模型 (FLoRA / Federated Instruction Tuning)：传统联邦大模型研究聚焦于耗费海量算力与通信带宽去协同微调千亿大模型参数本身。本文另辟蹊径，将联邦目标聚焦于轻量级路由器策略（参数量仅 670K，通信开销低至数兆字节），以极低代价实现了全网大模型推理经济性的系统级跃升。

评分

  • 新颖性: ⭐⭐⭐⭐⭐ 首次提出去中心化与稀疏评估下的大模型路由范式，兼具参数化与非参数化双方案。
  • 实验充分度: ⭐⭐⭐⭐⭐ 覆盖三大主流基准，详细考察了客户端扩展性、任务极端异质性、模型扩容与新客户端冷启动。
  • 写作质量: ⭐⭐⭐⭐⭐ 数学表达严密，理论收敛性证明与直观的流形统计解释高度自洽。
  • 价值: ⭐⭐⭐⭐⭐ 对隐私敏感型企业与边缘计算环境下调度云端混合大模型具有极高的工程借鉴与落地价值。

相关论文