MoEfication

Moefication: Transformer feed-forward layers are mixtures of experts

链接：https://arxiv.org/abs/2110.01786

源代码：https://github.com/thunlp/MoEfication

MoEfication可以使用模型的预训练权重，将非MoE的LLM转换为MoE版本。这种方法消除了对MoE模型昂贵预训练的需求。为了实现这一点，MoEfication首先将预训练的LLM中的FFN神经元划分为多个组。在每个组内，神经元通常由激活函数同时激活。然后，它将每个神经元组重构为一个专家。

通过综述的描述，这篇文章最吸引我的点是它不需要预训练就能把一个训练好的模型转成MoE。

本文思路

首先是发现llm的FFN在推理时只有一小部分会被激活，具有稀疏性；于是就觉得把FFN拆成多个专家来组成一个MoE。带来的好处就是加速推理，以及为研究 FFN 的内部机制提供了细粒度的视角。

FFN激活的稀疏性：对经过微调的 T5-Large 模型进行推理时，90%的输入只激活小于5%的神经元

相关工作：

1. 研究发现，Transformers 的 FFN 可以被视为存储网络，存储着从语言建模中学到的大量知识

2. MoE被提出的较早，并且在深度学习中得到了应用，第一个大规模 MoE 语言模型是在2017年提出的，它在两个 LSTM 层之间添加了一个 MoE 层，并独立地将token分配给专家组合。

3. 在推理加速方面也有很多工作，比如知识蒸馏、模型剪枝、注意力近似等。在这些技术中，动态推理探索有选择地省略不必要的加速计算，这与 MoEfication 的目标类似。

如何将FFN拆成多个部分：

核心思想是将经常同时激活的神经元（每一列视为一个神经元）分组在一起。如果专家数量为k，那么FFN的中间维度是。

本文提出了拆分的两种方法：

• 第一种是参数聚类分割，为了考虑参数信息，将FFN中第一个线性层的列视为具有维度的向量集合，也就是将每一列视为一个神经元。基于这样的直觉：具有相似向量的神经元将会同时被激活，我们对向量集合应用balanced K-Means聚类算法来获得k个簇，以构建映射函数。

• 第二种是Co-Activation图分割(Co-Activation Graph Split),为了直接使用共激活的信息，通过计算LLM对训练集样本的共激活次数来构建一个共激活图。图中的每个神经元将由一个节点表示，两个节点之间的边权重是它们的共激活值。共激活值是通过下面公式计算得出的：

  其中表示第n个神经元获得输入x之后得到的输出，这个公式同时要求这两个激活值都是正的。如果两个神经元的co-activation大，说明这两个神经元经常被同时激活，那么它俩就应该在MoE的同一个专家中，对应在co-activation图中边的权重就高。最后通过图分割算法，就把众多神经元划分成了不同的类。

专家选择：

专家选择核心思想是：选到一组分区，使其输出和全部神经元的输出尽可能接近，给每个分区打分，固定选n个分最高的分区。所以关键在于选择一个打分的方式：

• ground truth选择：ground truth实际就是没有被拆分的FFN的输出，以专家输出和ground truth之间的差作为损失。在这个过程中，假设要选择激活k个专家中的n个，那么用贪心的方式，从空集合开始，每一步选择一个专家使得该专家加入后损失函数降低的最多。这只是在训练门控模块时的打分方式，在推理时并不会被使用，因为它没有节约推理

• 相似性选择：将一个专家中的所有列，加和取平均，将这个结果与输入计算相似度

• MLP选择：训练一个MLP，接受x作为输入并预测每个专家中正数值的总和

所以，可以使用第一种打分方式来训练一个MLP，让MLP作为门控模块，在推理时选择激活哪些专家。

论文指出，专家构建和专家选择的最好方式是Co-activation Graph Split和MLP Selection