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