深入研究大型语言模型中的专家混合机制

作者

引言

随着大型语言模型（LLMs）的出现，自然语言处理领域发生了革命性的变化。研究人员通过增加模型规模和训练数据量，不断推升这些模型的能力。然而，这也带来了训练成本和复杂性的急剧增加，迫切需要新的解决方案来更好地平衡预训练效率和模型性能。专家混合（Mixture-of-Experts, MoE）架构通过动态路由输入到一部分专家，从而在不牺牲计算效率的情况下，实现模型的大规模扩展。

友情链接：ACEJoy

 

Mixture-of-Experts 概述

MoE模型通过用N个并行的前馈网络（FFNs）替换原有的FFNs，并结合一个路由器来增强Transformer。这些并行的FFNs被称为专家，每个专家通过路由器的评分分布来接收输入。公式如下：

[ \mathbf{y} = \sum_{n \in N} g_{n}(\mathbf{x}; \mathbf{G}, k) E_{n}(\mathbf{x}) ]

当Top-k小于N时，仅有部分专家参与计算，这被称为稀疏专家混合（SMoE）。

方法概述

本文选取了Mixtral 8x7B、DeepSeekMoE和Grok-1三个模型进行分析，并附加了标准的Transformer模型Mistral 7B作为对比。研究分为静态参数分析和动态行为分析两部分。

静态参数分析

专家权重矩阵

在矩阵级别，我们计算了所有专家在每一层的投影矩阵Wup、Wgate和Wdown之间的相似性，并用PCA将这些矩阵扁平化为二维空间。结果显示，不同模型的矩阵相似性模式存在显著差异。

门嵌入

门嵌入的相似性分析显示，Mixtral和DeepSeek的门嵌入与专家选择的权重矩阵Wgate之间存在较强的相关性，这表明门嵌入和Wgate可能学习了相似的知识来执行选择操作。

动态行为分析

专家输出

通过输入短序列和长序列，我们测量了专家输出特征向量之间的相似性。结果显示，被选择的专家输出更相似，尤其是在Mixtral模型中。

专家输出的L2范数和门评分

在Mixtral模型中，我们发现被选择的专家通常输出范数最大，这表明范数可能是门决策的关键因素之一。这一发现与CompeteSMoE的设计一致，该设计基于输出范数选择专家。

讨论

基于我们的分析，我们提出了一些关于MoE模型的建议：

  • 神经元级别的专家：门嵌入决定了专家选择，而Wgate负责选择激活的神经元，两者之间的相似性表明神经元可能是更细粒度的专家。
  • 模型架构：可以考虑在深层增加专家数量，而在最后一层减少专家数量。此外，基于输出范数作为路由机制是合理的。
  • 相关性测量：权重矩阵之间的相似性测量在某种程度上等同于输出特征向量的相似性测量。

相关工作

现有的大部分工作从路由器的角度分析MoE，通过观察专家选择来研究其行为。近年来的研究发现标准路由器并未在领域级别表现出明显的专业化，主要基于token ID进行路由。

结论

本文首次尝试通过研究参数和输出来揭示MoE模型的内部工作机制。我们总结了实证观察，并提出了一些实用建议。虽然目前还无法全面回答MoE是否学习了异质的专家，但我们的一些实验表明，特定的架构设计和训练框架可能更有利于专家的专业化。

限制

我们的分析主要基于观察，缺乏进一步验证的干预措施。

GitHub – kamanphoebe/Look-into-MoEs: A Closer Look into Mixture-of-Experts in Large Language Models

这个GitHub仓库包含了论文《A Closer Look into Mixture-of-Experts in Large Language Models》的代码。以下是仓库的主要内容和使用指南：

概述

我们首次尝试理解基于MoE的大型语言模型的内部工作机制。具体来说，我们全面研究了三个基于MoE的模型（Mixtral 8x7B, DeepSeekMoE, Grok-1）的参数和行为特征，揭示了一些有趣的观察结果，包括：

  • 神经元像细粒度专家一样行为：门嵌入决定了专家选择，而专家的投影矩阵则负责选择激活的神经元。它们的相似性值显示出关联性，这表明专家神经元可能是细粒度的专家。
  • MoE的路由器通常选择输出范数较大的专家：以Mixtral为例，我们发现输出特征向量范数最大的专家最有可能被门赋予最高分数。
  • 专家多样性随着层数增加而增加，但最后一层是一个例外：在多个实验中，我们观察到专家之间的相似性在深层中通常较低，而在最后一层相似性增加。

基于这些观察结果，我们还为MoE实践者提供了一些建议，如路由器设计和专家分配。

使用指南

  1. 下载模型检查点：
    默认情况下，我们的代码从ckpt目录加载预下载的模型。你也可以修改代码直接从HuggingFace下载。我们使用的模型下载链接如下：
  • Mixtral 8x7B Base
  • Mixtral 8x7B Instruct
  • Mistral
  • DeepSeekMoE
  • Grok-1
  1. 创建conda环境：
   git clone https://github.com/kamanphoebe/Look-into-MoEs.git
   cd Look-into-MoEs
   conda create -n analyze --file env.txt

创建conda环境后，需要将其选为Jupyter内核。

  1. 运行实验：
    两个Jupyter notebooks（static_analysis.ipynb和dynamic_analysis.ipynb）分别包含了关于静态参数和动态行为的实验代码。你可以简单地运行每个实验对应的代码块，这些代码块的标题与论文中的实验名相同。

引用

如果你觉得我们的工作有用，请引用我们的论文：

@misc{lo2024closerlookmixtureofexpertslarge,
      title={A Closer Look into Mixture-of-Experts in Large Language Models}, 
      author={Ka Man Lo and Zeyu Huang and Zihan Qiu and Zili Wang and Jie Fu},
      year={2024},
      eprint={2406.18219},
      archivePrefix={arXiv},
      primaryClass={cs.CL},
      url={https://arxiv.org/abs/2406.18219}, 
}

致谢

我们的模型配置和建模文件借用了HuggingFace仓库中的相应文件，感谢这些作者的出色工作。

通过这个GitHub仓库，研究者可以复现论文中的实验，并进一步探索MoE模型的内部机制。

发表评论