MoE架构深度解析：从KV记忆动机到负载均衡策略type
status
date
slug
summary
tags
category
icon
password
1 背景
在大模型飞速发展的几年中，从宏观来看，模型架构基本以Transformer为主流。从微观结构上看，各个Transformer的组件仍在”进化“：
- 用RMSNorm替代layer-norm（为了更好的数值稳定性和计算效率）；
- 用GQA（Group Query Attention），MLA （Multi-Head Latent Attention），Sparse Attention, Linear Attention等来代替MHA(multi-head attention)（为了降低KV-Cache的存储成本、提升计算效率）
- 用SwiGLU， MOE等来代替FFN（Feed-Forward Network）(提升FFN的表达能力)。
本文主要聚焦目MOE架构的motivation和实现原理。
2 MOE架构的Motivation和原理
根据论文 《Transformer Feed-Forward Layers Are Key-Value Memories》的insight，FFN模块可以视作Un-normalize Key-Value Memories。
- 为输入的序列的某个token，
- 为activate function。
- 为两个Linear layer的参数矩阵。
计算和模型内部memories的相关性，再根据相关性整合信息。以这个视角看，FFN的参数量越大，意味着模型的容量越大，可以存储更多的memories。
那么如何增加模型的memories呢？最简单粗暴的方法，对FFN进行横向扩展，此时有（这个形式其实等价于提升原始权重的维度.）
上述形式引出MOE的第一个核心组件——experts，它是多个FFN的集合，每一个FFN是一个expert。
虽然式2能明显提升模型的容量，但新增的计算量也等比例增加。并且式2的形式，输入对所有expert都一视同仁（权重都为1），但直觉上不同expert存储了不同的“知识”，应当赋予更相关的知识库更大的权重。如何度量这个相关性呢？
这就需要引入MOE的另一个核心组件GateNetwork , （也可以称之为Router）。这个组件的功能是预测当前token与各个expert相关性。其输出是一个维度等同于专家数量的分布，
- 表示token 和专家的亲和程度（affinity）
引入Router后，MOE的计算公式为
这个形式的MOE也称为DenseMOE。虽然DenseMOE相较Naive MOE能够更好的完成FFN层知识检索的动作，但从式3不难发现，它的计算资源需求更高了。
那么如何在尽可能不损害信息检索效果的前提下，尽可能的降低计算量呢？这就引入MOE的核心trick——稀疏化。
稀疏化的操作很直觉，前文通过Router可以得到当前token与不同expert相关度（affinity），那么我们可以只在最相关的个专家做信息检索，此时的Router的输出
此时MOE的计算形式为.放了便于描述，不妨记：
此时的MOE计算公式为：
通过稀疏化的处理，能将MOE的计算成本降低倍。
在实践中，往往会对进行归一化处理。此时
在DeepSeek-MOE中还额外引入shared expert，它不经过Router，可以理解为通用知识库。此时的MOE计算公式为：
3 MOE架构的负载均衡问题
虽然MOE的思路非常的优雅，但实际训练好还是比较有难度。根据(Shazeer et al., 2017)的论文的发现，MOE的训练中容易出现routing collapse，造成性能退化。具体来说：训练初期某个expert被激活，反向传播后，router更容易选择到最初激活的expert，强者恒强，导致其他expert没有被有效训练。因此MOE的训练需要解决expert训练过程中的负载均衡问题。
目前解决routing collapse有两种常用手段：
- auxiliary loss based
- auxiliary loss free
3.1 Auxiliary Loss Based
这个方法相比常规next-token prediction的language model loss额外引入一个辅助的balanced loss来保证expert的负载均衡。基本形式如下：
- 路由专家的总数
- Router激活的expert数量
- sequence length (每个sequence有个token)
- 路由专家对位置token的score （token-to-expert affinity）（gate network的输出）
- 相当于对进行normalize
- 专家在整个序列的平均激活score
- 在整个序列中专家被选中平均次数
- paper设置
这个公式看起来很令人疑惑，为什么加入它能够起到expert balance的效果。论文中没有给出详细解释，下面是笔者尝试解释一下。
将和代入
先看对角项
注意对于固定的，满足如下条件
显然，当为均匀分布时，上式取得最小，得到token-level的balanced loss的下界。
因此对角项对应：对于每个 token，选中的专家自身在该 token 上的归一化分量之和。
当每个 token 的 趋向均匀分布时,对角项接近它的下界。
再看交叉项
交叉项反映了：某 token 选中的专家在其它 token 上总体得到的归一化分量之和。
若某专家在很多 token 上被选中在多个为1，那么交叉项会增大，从而被惩罚。最理想情况下，如果不同token选择的专家完全正交，那么交叉项损失为0。
经过以上分析，对角项保证了token-wise激活expert的分布是均衡的。交叉项保证了不同token间激活的专家尽可能的正交。二者结合实现训练中expert的负载均衡。
虽然auxiliary loss的引入能让expert在训练过程中尽可能负载均衡，但auxiliary loss产生的梯度与语言建模的梯度不一定同向。可能会降低模型语言建模的能力。虽然可以手动调整超参数来控制梯度更新的幅值，但调大了会干扰语言建模，调小了又不能保证expert的负载均衡。
为了解决上述问题，《Auxiliary-Loss-Free Load Balancing Strategy for Mixture-of-Experts》提出一种auxiliary loss free的方法。
3.2 Auxiliary Loss Free
在Auxiliary loss base的方法中，Routed expert的选取规则为（选取激活得分topK的expert）
它完全取决于router输出的score。若不加特殊处理，极容易出现routing collapse。
auxiliary-free的方案对上式做了一个小改动
这个不是模型参数，而是一个统计量，它记录了每个专家历史的激活信息，若这个专家历史激活的比较多就会降低这个专家取得的score，否则增加（控制论中的闭环控制）。
那么如何来定义“误差”，如何确定的更新幅值呢？
论文给出的方案是：统计每个专家历史分配的token数量,从而可以计算平均数
那么可以得到每个专家分配数量的偏差，那么可以得到的更新公式
- 是符号函数，若说明这个专家被采样的频率低于平均值，需要增大
- 是更新的幅度。这个值是超参数，人为设定（论文设置为0.001）。
完整的算法流程如下：
4 小结
本文从 FFN 的 KV-Memory 视角出发，逐步推导出 MoE（Mixture of Experts）结构的动机与演化过程，并介绍了MOE架构训练过程Routing Collapse现象及解决方案。
- 作者:莫叶何竹🍀
- 链接:http://www.myhz0606.com/article/moe
- 声明:本文采用 CC BY-NC-SA 4.0 许可协议，转载请注明出处。
相关文章
