Reward-Based Online LLM Routing via NeuralUCB

以下是基于论文《Reward-Based Online LLM Routing via NeuralUCB》的详细技术总结：

1. 研究背景与问题定义 (Problem Definition)

核心问题：
大型语言模型（LLM）在不同任务上表现出显著的能力差异和推理成本差异。高端模型通常质量高但成本高，而廉价模型在处理复杂推理或代码查询时可能失败。研究旨在解决**基于奖励的在线 LLM 路由（Reward-Based Online LLM Routing）**问题。

形式化定义：

• 场景：在交互轮次 $t$，系统观察到查询 $x_t$，并从候选模型集合 $A = \{1, ..., K\}$ 中选择一个模型 $a_t$。
• 反馈机制：由于部署时通常每查询仅调用一个模型，系统只能观察到被选中模型的反馈（奖励 $r_t$），而无法观察到其他未选中模型在同一查询上的结果（即部分反馈/Partial Feedback）。
• 目标：学习一个路由策略 $\pi$，以最大化长期期望效用：
  $$\max_{\pi} E \left[ \sum_{t=1}^{T} r(x_t, a_t) \right]$$
• 挑战：反馈稀疏，且从查询到最优模型选择的映射高度非线性。

2. 方法论 (Methodology)

该研究将 LLM 路由建模为**上下文多臂老虎机（Contextual Bandit）**问题，并提出了基于 NeuralUCB 的解决方案。

2.1 效用函数设计 (Utility Function)

为了平衡质量与成本，作者定义了一个包含成本惩罚的效用奖励函数：
$$r(x, a) = q(x, a) \exp(-\lambda \tilde{c}(x, a))$$

• $q(x, a)$：模型质量。
• $\tilde{c}(x, a)$：归一化后的推理成本。
• $\lambda$：控制成本惩罚强度的超参数。
• 成本归一化：采用对数变换 $\tilde{c} = \frac{\log(1+c)}{\log(1+C_{max})}$ 将成本映射到 $[0, 1]$ 区间，以减少不同模型间成本量级差异带来的不稳定性。

2.2 网络架构：UtilityNet

UtilityNet 是路由框架中的核心效用预测器，其架构包含以下部分（如图 1 所示）：

• 输入编码：
  • 文本嵌入 ($x_{emb}$)：通过预训练文本编码器（如 Sentence Transformers）获取。
  • 辅助特征 ($x_{feat}$) 和 领域嵌入 ($d$)：分别通过 MLP 和 Embedding 层处理。
  • 动作嵌入 ($a$)：候选模型 ID 的嵌入。
• 双分支结构：
  1. 效用分支 (Utility Branch)：
     • 融合上下文和动作特征 $z_u = [h_{emb}, h_{feat}, e_a]$。
     • 通过 MLP 预测效用值 $\mu(x, a)$。
     • 使用 Huber Loss 进行回归训练，目标是拟合上述定义的效用奖励。
  2. 门控分支 (Gating Branch)：
     • 仅基于上下文特征 $z_g = [h_{emb}, h_{feat}]$。
     • 预测探索是否有益的概率 $p(x)$。
     • 使用二元交叉熵损失 (BCE) 训练，用于控制是否启用 UCB 探索机制。

2.3 路由策略：NeuralUCB

基于 UtilityNet 的预测，采用 NeuralUCB 进行动作选择：

• 特征表示：使用 UtilityNet 最后一层的隐藏表示 $h(x, a)$ 作为 UCB 的特征，并拼接常数 1 作为偏置项 $g(x, a) = [h(x, a); 1]$。
• 不确定性估计：使用共享的逆协方差矩阵 $A^{-1}$ 来估计特征空间中的不确定性，而非为每个动作维护独立的统计量。
• UCB 分数：
  $$s(x, a) = \mu(x, a) + \beta \sqrt{g(x, a)^\top A^{-1} g(x, a)}$$
  其中 $\beta$ 控制探索奖励的强度。
• 最终决策规则：
  • 若门控概率 $p(x) \ge \tau_g$（阈值），则选择 UCB 分数最高的动作（鼓励探索）。
  • 否则，回退到保守动作 $a_{safe} = \arg\max_a \mu(x, a)$（仅基于预测均值选择，确保安全）。
• 训练流程：采用模拟在线协议（Simulated Online Protocol），将数据分为 20 个切片（Slices）顺序处理。每处理一个切片，更新回放缓冲区（Replay Buffer）和 $A^{-1}$，然后训练 UtilityNet 并重建 $A^{-1}$。

3. 实验设置与结果 (Experiments & Results)

3.1 实验设置

• 数据集：RouterBench（36,497 个样本，86 个领域，11 个候选 LLM）。
• 基线模型：
  • Random（随机选择）
  • Min-cost（总是选最便宜的模型）
  • RouteLLM-BERT（基于监督学习的二分类路由）
  • Max-quality（总是选质量最高的模型，作为上限参考）
• 编码器：主要使用 all-MiniLM-L6-v2，并对比了其他编码器（如 E5, MPNet, Qwen3-Embedding）。

3.2 关键结果

1. 效用奖励表现：

   • NeuralUCB 在平均奖励和累积奖励上** consistently（一致地）优于**所有基线（包括 RouteLLM-BERT、Random 和 Min-cost）。
   • NeuralUCB 的平均奖励稳定在 0.59 - 0.61，而 Min-cost 约为 0.51-0.53，RouteLLM-BERT 约为 0.35-0.36。
   • 随着切片数量增加，NeuralUCB 与基线的累积奖励差距逐渐扩大，显示出良好的在线学习能力。

2. 成本与质量权衡：

   • 与“最大质量参考（Max-quality）”相比，NeuralUCB 仅使用了约 33% 的推理成本，同时保持了具有竞争力的选择质量。
   • 这表明该方法成功实现了显著的成本节约，同时未牺牲过多的任务表现。

3. 编码器消融实验：

   • 不同的文本编码器对性能有影响。multilingual-e5-large-instruct 表现最差，而 all-mpnet-base-v2 和 all-MiniLM-L6-v2 表现最强且最稳定。
   • 在排除 E5 后，更强/更大的编码器并未带来明显的性能提升，说明在该设置下，模型架构的优化比单纯扩大编码器规模更重要。

4. 主要贡献 (Key Contributions)

1. 提出 NeuralUCB 路由框架：首次将 NeuralUCB 应用于成本感知的 LLM 路由场景，有效解决了部分反馈下的非线性映射问题。
2. 设计效用驱动的损失函数：通过引入指数形式的成本惩罚和对数归一化，构建了一个能够灵活平衡质量与成本的效用函数。
3. 引入门控机制：设计了辅助门控分支，动态决定是否启用 UCB 探索，在“探索未知高回报”与“利用已知安全策略”之间取得了平衡。
4. 实证验证：在 RouterBench 上通过模拟在线环境验证了方法的有效性，证明了其在降低推理成本的同时能维持高奖励。

5. 意义与局限性 (Significance & Limitations)

意义：

• 为 LLM 推理成本优化提供了一种高效、自适应的解决方案，特别适用于模型池动态变化或查询分布漂移的场景。
• 证明了在无需全量标注数据（Full-information）的情况下，仅通过部分反馈即可学习出优于传统监督路由的策略。

局限性与未来工作：

• 动作区分度挑战：对于许多样本，最大质量的选择本身已经接近奖励最优决策，限制了路由带来的潜在增益。
• 数据需求：未来需要收集更多训练数据并改进效用模型，以更好地在成本约束下区分动作。
• 探索与利用的平衡：虽然门控机制有所帮助，但在极端稀疏反馈下，如何更精准地控制探索强度仍是挑战。

总结：
该论文提出了一种基于 NeuralUCB 的在线 LLM 路由方法，通过结合效用预测、不确定性估计和门控机制，成功在 RouterBench 基准上实现了比现有基线更优的成本 - 质量权衡，为构建经济高效的 LLM 应用系统提供了重要的技术参考。