202501 arxiv

1 intro

• MoE中常被忽视的一个关键问题是路由器的决策过程与专家执行之间的分离
  • 路由器无法直接评估专家的能力，因此它对专家的选择基本上是没有标签的预测
  • 如果路由器做出了错误的预测，选择的专家可能会试图处理这些令牌，但未能有效地处理，导致训练损失增加
    • 为减少损失，专家可能需要调整其参数以处理这些令牌，可能与原本的专业能力产生冲突
    • 另一种可能性是路由器必须通过试错来学会做出更好的决策，因为它没有意识到哪些专家更适合特定任务，这样会浪费大量的训练步骤。
• ——>论文提出了一种新的MoE范式——专家自治（AoE）
  • 这一设计基于专家能够意识到其处理输入的能力，并且这种意识反映在其内部激活的规模上
  • AoE允许专家决定是否自主处理输入
• 通过专家的自治，路由器的决策过程得以消除，从而降低了训练损失。
  • 预训练了拥有最多4百亿参数的AoE语言模型，并且在下游任务中优于传统的MoE模型，表现出更高的计算效率

2 方法

3 实验

3.1 一般设置

• 训练了由12层组成的小型语言模型，每层包含12个注意力头。
• 每层包含8个专家，选取Top-K = 2个专家
• 为MOE设置dmodel=768，dffn=3072
• AOE的dlow和dwide值是可变的
  • 为了确保AoE模型的总参数量与MoE模型相当，当调整dlow时
    • 

3.2 AOE下游表现 VS MOE下游表现

• AoE配置在平均准确率方面超越了最佳表现的MoE配置
  • AoE在没有任何辅助损失的情况下，超越了传统的MoE模型
• 调整dlow之后，性能变化不大
• 性能的提升是否源自Wg矩阵的分解？
  • 通过比较配置3和配置2，检查了专家权重矩阵分解对性能的影响
  • ——>分解对性能的影响并不显著，这与LLM权重本身就是低秩的发现一致

3.3  不同专家选择策略的影响

3.4 负载均衡

AoE在负载平衡方面优于传统的MoE模型，无论是否使用负载均衡loss