Mixture of experts(MoE)是一种机器学习技术,其中多个专家网络(或称为学习者)将问题空间划分为同质区域,代表了一种集成学习形式。在此类系统中,路由机制(通常称为门控函数或加权函数)决定输入如何分配给专家,以及它们的输出如何组合。因此,Mixture of experts 路由指的是控制这一选择和加权过程的具体算法与设计选择,这些选择对于平衡模型容量、计算成本和输出质量至关重要。该概念已从1990年代的早期统计模型演变为当代大型语言模型中使用的稀疏路由方案,后者以相对较低的每令牌计算量实现了海量参数。
任何 mixture of experts 系统的核心架构都包括一组专家函数 \(f_1, ..., f_n\),每个函数接收相同的输入 \(x\) 并产生输出,以及一个加权函数 \(w\),它将 \(x\) 映射到非负权重向量 \((w(x)_1, ..., w(x)_n)\)。最终输出通常计算为加权和:\(f(x) = \sum_i w(x)_i f_i(x)\)。专家和加权函数通过最小化损失函数(通常使用梯度下降)联合训练。路由机制决定了这些权重如何计算,是密集的(所有专家都参与)还是稀疏的(仅激活少数专家),以及系统如何处理负载均衡和训练稳定性。
早期路由设计
最早的路由公式之一是 meta-pi 网络,由 Hampshire 和 Waibel 在1990年代初提出。在该设计中,输出是专家输出的加权和,训练通过梯度下降最小化均方误差损失。专家可以是任意函数,门控网络学习根据输入分配权重。在他们的原始出版物中,研究人员将此方法应用于对六位日语说话者(两女四男)语音信号中的音素进行分类。他们训练了六个专家,每个专家是一个在梅尔频谱图上运行的时间延迟神经网络。值得注意的是,学习到的路由将五个专家专门分配给五位个体说话者,而第六位男性说话者的语音则由其他三位男性说话者的专家的线性组合进行分类,这表明路由可以发现共享子空间,而非一对一的映射。
另一种早期方法是自适应局部专家混合,它使用高斯混合模型作为门控函数。在此方法中,每个专家预测一个高斯分布,通常忽略输入而仅学习均值向量。加权函数是线性 softmax 函数,其中专家 \(i\) 的权重计算为 \(w(x)_i = \exp(k_i^T x + b_i) / \sum_j \exp(k_j^T x + b_j)\)。这种 softmax 路由产生了一个归一化的专家概率分布,整体模型输出是高斯预测的混合。该公式允许概率解释,并通过最大似然进行训练,为后来的概率路由方法奠定了基础。
稀疏路由与 Transformer 时代
现代 mixture of experts 路由随着 Transformer (architecture) 架构在 Deep learning 中的兴起而获得 prominence。在大规模模型中,密集路由(每个专家处理每个输入)随着专家数量的增长在计算上变得不可行。稀疏路由在2017年由 Noam Shazeer 及其在 Google DeepMind 的同事发表的论文“Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer”中引入,通过仅对每个输入令牌激活一小部分专家来解决此问题。门控函数对专家分数计算 softmax,但随后应用 top-k 选择,仅保留得分最高的专家(通常 k = 1 或 2),并将其余部分置零。这使得模型可以拥有数十亿参数,而每次前向传播仅计算其中一小部分。
稀疏门控函数通常使用一个可训练的权重矩阵,将输入令牌表示映射到每个专家的 logit 向量。这些 logit 通过 softmax 转换为概率,然后选择 top-k 专家。所选专家的输出按其归一化概率加权并求和。该路由机制相对于门控参数是可微的,从而允许通过反向传播进行端到端训练。然而,稀疏路由引入了负载不均等挑战,即少数专家可能主导训练,而其他专家接收到的训练信号很少,这导致了辅助负载均衡损失的发展。
负载均衡与辅助损失
mixture of experts 路由中的一个持续问题是专家崩溃,即门控网络学会将大多数输入路由到一小部分专家,导致其他专家利用不足。为了应对这一问题,现代实现添加了辅助损失,以鼓励均匀路由。一种常见方法,如 Google DeepMind 在2021年推出的 Switch Transformer 中所用,是添加一个负载均衡损失,当分配给专家的令牌分布偏离均匀分布时进行惩罚。该损失通常计算为平均路由概率与每个专家路由令牌比例之间的缩放点积,并以小系数添加到主训练损失中。另一种技术是使用专家容量限制,其中每个专家每批次只能处理固定数量的令牌;超出此容量的令牌会被丢弃或路由到残差连接,从而防止任何单个专家成为瓶颈。
大型语言模型中的路由
Mixture of experts 路由已成为扩展 Large language model 的基石。例如,由 Mistral AI 开发的 Mixtral 8x7B 模型使用稀疏 MoE 层,其中每个令牌被路由到八个专家中的两个,每个专家是一个前馈网络。这使得模型拥有470亿总参数,但每个令牌仅使用约130亿参数,从而在推理成本上与更小的密集模型相当。类似地,Switch Transformer 证明了在自然语言任务中,通过稀疏路由实现数万亿参数的扩展是可行的,并且相比密集基线实现了加速。在这些模型中,路由决策是针对每个令牌做出的,而非每个序列,从而允许模型将输入的不同部分分配给不同的专家。
大型语言模型中的路由机制通常作用于 Multi-Head Attention 层产生的隐藏状态。门控函数是一个线性投影,后接 softmax 或 sigmoid 函数,并与网络其余部分联合训练。一个关键的设计选择是是否使用带噪声的 top-k 门控,其中在 logits 选择之前添加可训练的高斯噪声,以在训练期间鼓励探索,并防止门控过早变得过于确定性。另一个选择是使用专家并行,其中专家分布在多个设备上,并且必须协调路由以最小化通信开销。这导致了在 Amazon Web Services SageMaker 和 Google Cloud TPU 环境等框架中的专门实现。
挑战与最新进展
尽管 mixture of experts 路由有效,但它面临若干开放性挑战。一个是路由粒度与计算效率之间的权衡:专家过多可能导致内存开销和通信成本增加,而专家过少则限制了容量。另一个是训练的不稳定性,因为离散的 top-k 选择可能导致梯度问题,尽管已经探索了直通估计器和 softmax 松弛等方法。最近的研究还调查了学习路由策略,这些策略可以根据输入复杂度动态调整活跃专家数量,对简单输入减少计算,对复杂输入增加计算。此外,关于专家合并与剪枝的研究(即在训练后合并或移除冗余专家)可以在不牺牲质量的情况下减少内存占用。随着模型持续扩展,路由机制将在决定其性能和实用性方面发挥越来越核心的作用,使其成为理论和应用研究的丰富领域。
与密集模型的比较
Mixture of experts 路由为密集模型提供了一种根本性的替代方案,后者对每个输入使用所有参数。密集模型(如原始 Transformer (architecture) 架构)对每个令牌具有固定的计算成本,而 MoE 模型根据路由决策具有可变的成本。这使得 MoE 模型能够在不按比例增加推理成本的情况下实现更高的容量,使其在资源受限环境中的部署具有吸引力。然而,密集模型通常更容易训练和微调,因为它们不需要负载均衡损失或仔细的容量调整。密集与 MoE 架构之间的选择取决于具体应用,MoE 特别适用于大规模训练,其目标是在固定计算预算下最大化质量。
在实践中,包括 OpenAI、Anthropic 和 Google DeepMind 在内的许多组织已在其生产模型中采用 MoE 层,尽管它们通常不披露确切的路由细节。该技术也已应用于语言之外的领域,包括计算机视觉和语音识别,并显示出类似的优势。随着硬件持续发展,诸如 AWS Trainium 和 Groq 等专门加速器针对稀疏计算进行了优化,MoE 路由的效率可能会进一步提高,使其成为未来 Artificial intelligence 系统的标准组件。
未来方向
Mixture of experts 路由的未来在于使其更具适应性和效率。一个方向是开发完全可微的路由,其中离散的 top-k 选择被连续近似替代,从而允许更平滑的梯度并可能实现更好的优化。另一个是使用学习路由策略,该策略可以根据输入复杂度动态调整活跃专家数量,对简单输入减少计算,对复杂输入增加计算。此外,关于专家合并与剪枝的研究(即在训练后合并或移除冗余专家)可以在不牺牲质量的情况下减少内存占用。随着模型持续扩展,路由机制将在决定其性能和实用性方面发挥越来越核心的作用,使其成为理论和应用研究的丰富领域。