MoE路由机制
神经网络规模增长带来的计算瓶颈催生了条件计算的需求——并非每个输入都需要激活全部参数,路由机制正是决定"激活哪些参数"的核心调度组件。
Jacobs、Jordan、Nowlan、Hinton 在 Neural Computation 发表 "Adaptive Mixtures of Local Experts",首次提出混合专家框架。核心思想是将输入空间划分给不同的"专家"网络,由一个门控网络决定每个专家的输出权重。门控网络输出概率分布,所有专家输出按此分布加权求和。这一工作奠定了 MoE 的理论基础,但受限于当时算力,专家数量极少(通常 2-8 个),且训练极不稳定,专家容易坍缩为同质网络。
Shazeer 等人在 Google Brain 发表 "Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer",将 MoE 从密集门控推向稀疏门控。关键创新是 Top-K 路由:每个输入只激活 K 个专家而非全部,使模型参数量可扩展到数千亿而每个 token 的计算量仅与单个专家相当。他们引入了噪声 Top-K 门控(softmax 前加入可调高斯噪声增加探索性)和负载均衡辅助损失,首次系统性地解决了专家坍缩问题。该工作将 MoE 嵌入 LSTM 语言模型,在翻译任务上以约 1/10 的计算量达到当时最优性能,证明了稀疏条件计算的巨大潜力。
Fedus、Zoph、Shazeer 在 Google 提出 Switch Transformer,将路由激进简化为 Top-1——每个 token 只选一个专家。这带来两个关键好处:通信开销大幅降低(All-to-All 通信量减半),路由决策更清晰使训练更稳定。Switch Transformer 将模型扩展到万亿参数级别,首次证明稀疏 MoE 在 Transformer 架构中的可行性。同时引入容量因子(capacity factor)机制,限制每个专家处理的 token 数量上限,防止分布式训练中的负载倾斜。
Zoph 等人提出 ST-MoE,引入路由器 z-loss 惩罚路由 logits 的绝对值过大,防止 softmax 饱和导致梯度消失。这一简洁的正则化技术成为后续所有 MoE 训练的标准配置,显著提升了大规模训练的稳定性。
Mistral AI 发布 Mixtral 8x7B 采用 Top-2 路由,在开源社区引发轰动,证明 MoE 可以在消费级硬件上高效推理。DeepSeek 提出细粒度专家分割和共享专家机制——将部分专家设为所有输入都经过的"共享专家",其余为按路由选择的"路由专家",兼顾通用能力与专业化,标志着 MoE 从研究走向大规模工业部署。
本日论文提出用流形幂迭代重新设计路由矩阵,将路由行视为专家权重矩阵的最优代表向量,通过迭代优化使路由决策与专家实际能力理论对齐,代表了路由器设计从启发式走向数学最优的新方向。
设输入 token 为 $x \in \mathbb{R}^d$,共 $N$ 个专家 $\{E_1, ..., E_N\}$,路由矩阵 $W_r \in \mathbb{R}^{N \times d}$,第 $i$ 行 $w_i$ 是专家 $i$ 的代表向量。路由得分为 $h(x) = W_r \cdot x$,第 $i$ 分量 $h_i(x) = w_i^\top x$ 衡量输入与专家 $i$ 的相似度。门控概率通过 softmax 归一化:$G_i(x) = e^{h_i(x)} / \sum_{j=1}^{N} e^{h_j(x)}$。Top-K 稀疏化只保留得分最高的 $K$ 个专家:$\tilde{G}_i(x) = G_i(x)$ 若 $i \in \text{TopK}(h(x))$,否则