树路由取代扁平的专家列表
传统的混合专家路由把专家视为一个无结构的集合:门控矩阵为每个专家产生一个分数,softmax 把分数转成概率,再分派 top-K 专家。专家索引不携带任何拓扑含义,路由器本身也无法阻止某些专家被选中的频率远高于其他专家。BRANCH-MoE 改变的是分配规则,而不是专家网络。它把 E 个专家放在深度为 log2 E 的平衡二叉树的叶节点上,于是每个 token 依次经过一串二元决策,叶节点分数是各条根到叶路径上分支概率的乘积。
论文的动机是:扁平路由器可能造成专家利用率不均衡,而平衡通常要靠辅助负载均衡损失或批次耦合的分配来恢复。BRANCH-MoE 则把平衡嵌入决策的几何结构中。在每个内部节点,分支概率以到达该节点流量的到达加权平均分数为中心,因此决策边界落在实际到达该节点的分布中部。作者称这能在没有辅助负载均衡损失的情况下促进两个子树的利用率。
到达加权锚点决定每次分裂
该机制是一个锚点,而不是负载误差反馈变量。在内部节点 j,偏置为 bj = -νj,即该节点到达加权平均分数的指数移动平均的相反数。训练时锚点是当前批次在该节点的到达加权平均分数;推理时它是一个不可训练的指数移动平均,衰减率为 0.99,每个训练步更新一次。论文将其与 DeepSeek-V3 的动态逐专家偏置作对比:后者对欠载专家增大偏置、对过载专家减小偏置。BRANCH-MoE 的偏置把节点决策边界锚定在自身输入分数分布的中心,而不是对负载误差作出反应。
指数移动平均引入了一个作者明确指出的权衡。定理 4.2 表明,在平稳锚点估计下,初始化误差按几何速率衰减,而渐近均方波动至多为 (1-alpha)s^2/(1+alpha),因此 alpha 越大,采样噪声下限越低,但 EMA 响应越慢。命题 4.3 将其推广到漂移锚点:大 alpha 抑制采样噪声,小 alpha 减小滞后。论文还指出,在软路由下每个节点在精确算术中都会收到正质量,但深层节点的到达权重可能高度集中,从而增大方差并带来数值下溢风险。
线性路由在对数凹到达下可防塌缩
论文证明,在特定条件下,到达加权锚定能防止路由质量塌缩。引理 4.4 指出,如果输入分布在其仿射支撑上具有对数凹密度,且某节点的每个祖先都使用仿射分数,那么该节点的到达加权分布也是对数凹的。推论 4.5 随后表明,对于权重向量非零的线性节点映射,在精确总体锚点处,两个子节点概率中较小的一个至少为 1/(2e)。在任意活动锚点下,该下界变为 (1/2)(1/e - A|delta|)+,其中 A 是中心化投影分数密度的上界,delta 是活动锚点误差。
沿通往叶节点的路径反复条件化,可得到期望叶质量的下界。在线性对数凹情形且总体锚点精确时,每个叶节点的期望路由质量至少为 (2e)^-D。作者谨慎说明了这能推出什么、不能推出什么:这些下界本身并不蕴含每个叶节点都会被确定性 top-K 截断保留。论文指出,软质量下界消除了专家饥饿的一个重要机制,应使实际执行频率比无约束学习路由下更大、更稳定,但它并不在形式上蕴含正的 top-K 执行概率。
执行频率决定优化瓶颈
另一项结果把利用率与训练动态联系起来。在冻结路由器下,论文把 pe 定义为专家 e 在 top-K 选择后的实际执行概率,这与期望叶质量 qe 不同。定理 4.7 表明,在冻结路由器专家 SGD 且条件风险强凸且光滑时,到专家最优点的期望平方距离以由 pe 乘以学习率和强凸常数所决定的速率收缩。如果每个专家的执行概率至少为 pmin,那么当全局训练步数大致达到 (1/(pmin gamma mu)) log(初始距离平方除以 epsilon) 时,暂态会降到 epsilon 以下。
作者由此得出的含义是:实际执行频率是稀疏激活专家的优化瓶颈,最小频率越大,最差专家在全局训练步数下的收敛越快。BRANCH-MoE 的设计目标就是改善这一量:把每个内部决策锚定到到达该节点的总体,从而抑制会同时压制整棵子树中所有专家的持续高层不平衡。论文明确指出,均匀随机路由仍是边际平衡的自然基准,而 BRANCH-MoE 追求的是在保留任务相关特化和树对齐通信局部性的同时,达到可比的利用率。
根部的高置信决策限制跨设备通信
当专家按树前缀分配到设备时,树结构还会带来通信上的后果。论文分析了一个硬件边界,其域与深度 r 的路由子树对齐,token 向每个被接触的域传输一次,之后可在本地复制。定理 4.8 表明,首选深度 r 子树之外的叶质量至多为前 r 层 exp(-gamma_i) 之和,其中 gamma_i 是第 i 层学到的分数间隔。包含至少一个被选专家的深度 r 通信域数量,受 K 与 1 加上该外部质量与 top-K 专家中最小保留概率之比这两者中的较小值约束。
实际阈值被直接给出:如果前 r 层 exp(-gamma_i) 之和低于 top-K 专家中最小的保留概率,那么所有 top-K 专家都位于首选子树内,token 只接触一个域,通信成本仅为 token 表示大小乘以每字节传输成本。如果前 r 个分数间隔每个都不小于 gamma,则外部质量至多为 r exp(-gamma),因此 gamma 大于 log(r/tau_K) 就足以在深度 r 边界实现单域分派。作者指出,这一逐 token 结果与总体层面的防塌缩结果相容:一个 token 可以把大部分概率分配给一个分支,而其他 token 偏好相反分支,因此跨 token 的广泛专家利用率可以与每个 token 集中且有结构的分配共存。
四个基准上任务质量保持竞争力
评估将 BRANCH-MoE 与四种路由范式比较:带辅助均衡损失的 Switch softmax、DeepSeek-V3 动态偏置、Skywork logit 归一化门控,以及确定性哈希路由。所有条件都使用 E = 16、top-4 路由和容量因子 1.25,仅门控不同,因此比较是在匹配容量下进行的。基准包括作为主要大规模任务的 Criteo 点击率预测,以及作为可低成本多随机种子复现的表格基准 Forest Covertype、HIGGS 和 YearPredictionMSD。每个条件都重复五个随机种子。
在 Criteo 上,logloss 和 AUC 是 15 次评估中的最佳留出值,论文报告称在 5 对 5 个种子下,logloss 差异低于约 3.7e-4 或 AUC 差异低于 2.7e-4 时,双侧 Welch t 检验在 p < 0.05 下无法分辨。Criteo 比较是非配对的,因为 TPU 执行不是逐位可复现的;用相同种子重跑相同配置对留出 logloss 的改变,与更换种子带来的改变大致相当。在 UCI 基准上,种子决定初始化但不决定数据顺序,同一数据集内所有实验组共享相同的种子集合,因此比较按种子配对。摘要的总结是:分层路由可以保持任务质量和均衡利用率,同时形成支持局部化专家共激活并减少通信的拓扑。
一个预处理细节与单层研究范围
Forest Covertype 上的一个具体结果表明树路由器对预处理有多敏感。54 个特征中有 40 个是二元土壤类型指示变量,其中最稀有的一个只在 464,809 个训练行中的 2 行上触发。使用 epsilon 为 1e-7 的分母保护进行标准 z 分数标准化,会把一个在拟合归一化器所用样本中恒定的列,在它确实触发的行上映射为约 1e7 量级的尖峰,从而破坏运行统计量,并在树路由器中诱发双峰专家塌缩。把任何自适应方差低于 1e-4 的列固定为单位方差,可消除这种塌缩,并把 BRANCH-MoE 的种子间准确率波动从 6.2 个百分点降到 0.46 个百分点。
该研究的范围是单个 BRANCH-MoE 层,作者指出这无法捕捉多层之间路由决策的相互作用。理论也是有条件的:防塌缩保证是针对线性节点映射和对数凹到达分布证明的,收敛结果假设冻结路由器、独立同分布样本以及强凸且光滑的条件风险。通信界假设前缀对齐放置,并依赖学到的分数间隔,而不是自动成立。论文自身的重点是总体层面平衡与逐 token 集中之间的权衡:到达加权锚定防止 token 总体在任一子节点上分配任意小的路由质量,但并不要求每个 token 均匀分配其质量。证据支持的是这一区分,而不是逐 token 均匀路由的主张。
本文引用的文献
BRANCH-MoE: Balance-Aware Tree Routing for Large Embedding Models
BRANCH-MoE 将 16 个专家放在二叉树的叶节点上,并把每个内部节点的分裂中心对准到达流量的到达加权平均分数,该均值由指数移动平均估计。在 Criteo、Forest Covertype、HIGGS 和 YearPredictionMSD 上,使用 top-4 路由和五个随机种子,它在没有辅助负载均衡损失的情况下保持了任务质量和均衡利用率,其树拓扑支持局部化的专家共激活并减少通信。
