本文目录

1 MoE 知识蒸馏：从稀疏教师模型到密集学生模型 (来自 NUS，华为诺亚方舟实验室)
1 论文解读
1.1 背景：稀疏 MoE 模型虽强大却难以部署
1.2 MoE 基本介绍
1.3 知识整合的问题定义
1.4 第一阶段：知识整合
1.5 第二阶段：知识蒸馏
1.6 视觉任务实验结果
1.7 文本任务实验结果
1.8 消融实验结果

太长不看版

人类教育系统由多个 Expert 来训练一位 Student。专家混合 (Mixture-of-experts, MoE) 是一种强大的稀疏架构，包括多个 Expert 模型。但是，MoE 的架构容易过拟合，难以部署，对业界其实并不友好。

这个工作受到人类教育的启发，作者提出了一种新的任务：知识整合 (Knowledge Integration)，本文的方法分为2步：知识聚合 (Knowledge Gathering) 和知识蒸馏 (Knowledge Distillation)。使用稀疏的教师模型来蒸馏密集的学生模型。学生是一个密集的模型，从不同 Expert 那里获取知识。

为了从不同的预训练 Expert 那里有效地收集知识，本文研究了4种不同的可能的知识收集方法，即 Summation, Averaging, Top-K Knowledge Gathering (Top-KG), 和 Singular Value Decomposition Knowledge Gathering (SVD-KG)。然后再通过知识蒸馏精炼密集的学生模型。

1 MoE 知识蒸馏：从稀疏教师模型到密集学生模型

论文名称：One Student Knows All Experts Know: From Sparse to Dense

论文地址：

http://arxiv.org/pdf/2201.10890.pdf

1.1 背景：稀疏 MoE 模型虽强大却难以部署

回顾一个研究员，一路走来的学习历程，大多数人都是从多个教师 (即 Expert 那里) 学习到不同的知识。受人类教育的启发，这项工作侧重于通过一组 Expert 的知识来训练一个强大的学生模型。

在 MoE 中，每个 Expert 都是整个模型中的一个子神经网络。对于输入的每个子集，只有一个特定的子网络被激活，MoE 的这种稀疏计算的方式使我们能够将模型扩展到万亿参数的规模。

但是，MoE 方法也有问题，比如很容易过拟合。使用 MoE 模型时，我们通常在大型数据集上预训练 MoE，然后在各种下游任务上去做微调。在大多数情况下，这些下游任务才是我们想要解决的目标问题。MoE 的应用存在下面3大问题：

第一，容易过拟合。与 Dense 的模型相比，稀疏模型更多的可训练参数和稀疏的条件计算会引入过拟合，尤其是当数据集的规模不够大的时候。

第二，比较难部署。即使我们成功地训练了一个 MoE 模型，它其实也是很难部署的。对于参数量达到数万亿规模的 MoE 模型，我们就需要在不同设备 (GPU, TPU) 上面部署不同的专家来减小设备上的内存消耗。

第三，MoE 模型硬件不友好，专家并行性代价高昂。对于 GPU 集群，all-to-all 操作太慢，无法扩大 MoE 模型。而且，门控函数包括许多创建 token mask 的操作，选择 top-k Expert 并寻找每个 Expert 的输入 token 的 id。由于 tensor 的稀疏表示，所有这些操作都很浪费。

总之，稀疏 MoE 虽然强大，但是在实践中相对难以使用。密集的模型能力稍差，但是却容易使用，因此，本文希望结合稀疏模型的能力来训练高效又易于使用的密集模型。

1.2 MoE 基本介绍

MoE 的基本概念和定义

给定一个具有 个可训练的 Expert 和输入 的 MoE 模型, 模型的输出可以表示为:

式中, 表示第 个 Expert, 表示门控网络, 是第 个 Expert 的路由权重。通常, 和 都由神经网络参数化, 且 的输出应由 Softmax 函数激活:

式中, 是 函数, 是线性层, 是高斯噪声, 其中,选择 topK 的 Expert 来激活是 MoE 的关键模块。通常将 设置为 1 或 2 , 以便模型整体的计算成本和相应的密集模型相当。

MoE 的常见目标函数

在训练 MoE 模型时，如果我们不加任何的正则化策略，大多数 token 可能会派到一小部分 Expert 里面，其他的 Expert 只会收到很少很少的 token，这种极不平衡的分配会导致 MoE 模型的效率很低且精度较差。因此，为了实现不同 Expert 的平衡工作负载，通常使用辅助的平衡损失函数：

式中, 是 Expert 的数量, 是个向量, 其中 的定义是分配给第 个 Expert 的 token 的比例, 表达式如下:

其中, 是 token 的数量, 是分配给第 个 Expert 的 token 数量。由于这个式子设计 top-k 的计算, 因此是不可微的。

的定义是路由给第 个 Expert 的 token 的概率，表达式如下:

当我们试图优化辅助平衡损失函数时，其实是希望 的输出是均匀分布。

1.3 知识整合的问题定义

本文的方法知识整合可以分为2步：知识聚合 (Knowledge Gathering) 和知识蒸馏 (Knowledge Distillation)

对于第1步, 给定 个 Expert 模型 , 希望最大化学生模型 中的知识。给定输入的表征 , 在一个 Transformer Block 内, 每个 Expert 都是一个 FFN, 可以表示为:

式中, 和 是第 个 Expert 的线性变换, 是激活函数。对于密集的学生模型, 有相同的架构但不同的训练参数:

然后目标就是根据 和 来估计 和 的参数。

第2个阶段是对密集的学生模型进行微调，以最小化教师输出和学生输出之间的差异。

本文的目标是尽可能地保留 MoE 给密集学生模型带来的帮助。因此，本文搞了个指标来衡量密集的学生模型整合 MoE 对应的知识的能力：

其中， score 可以是衡量模型的任何指标。score 指直接训练密集的学生模型的性能。

1.4 第一阶段：知识整合

如下图2所示是知识整合的具体方法，整体的思路就是2步：知识聚合 (Knowledge Gathering) 和知识蒸馏 (Knowledge Distillation)。

在初始化学生模型的时候，大部分的参数其实是和教师模型重合的，比如 Embedding 层，Attention 层，归一化层。所以对于这些层的参数，学生模型可以直接复制教师模型的。问题的关键在于 MoE 层，MoE 层比单个 FFN 层具有更多的可学习参数，每个 Expert 都是一个有独立参数的 FFN。知识聚合的目的就是合并来自不同 FFN Expert 的知识，并将知识分配给学生模型中的单个 FFN。作者研究了4种不同的知识聚合方法，即：Summation, Averaging, Top-K Knowledge Gathering (Top-KG), 和 Singular Value Decomposition Knowledge Gathering (SVD-KG)。

假设 个 Expert 的线性层为 , 每个线性层 有权重 和偏置 。知识聚合的目标就是:

简单来讲就是把 的权重给搞成单 Expert 的权重 。式中, 是线性层, 有权重 和偏置 。

偏置 的确定方法比较简单, 这里作者直接加权平均所有 Expert 的偏置项:

采用这种简单的策略是因为可训练参数较少，存储在 bias 中的知识远小于 weight 中的知识。

偏置 的确定方法有下面4种：

累加 (Summation)：

取均值 (Averaging)：

Top-K Knowledge Gathering：

对于第 个 Expert 的权重矩阵 , 首先计算每一列的权重的范数 , 然后根据 选择 Top-K , 选择出来的矩阵是 。最后将所有专家提取的矩阵连接起来作为最终的学生初始化 。

SVD Knowledge Gathering

矩阵的低秩分解用于把一个非低秩矩阵转化成低秩矩阵，这样的低秩矩阵可以近似整个矩阵的知识，可以通过从多个低秩矩阵重构高阶矩阵。作者在本文中使用 SVD 来提取关键知识并将它们合并：

式中, 是酉矩阵, 是对角矩阵。通常选择 中的 top-K 元素并构建 和 来逼近 。

然后按照下式得到低秩分解之后的 ：

式中, 是一个 秩的矩阵, 式中 , 涵盖每个 Expert 的关键知识。

1.5 第二阶段：知识蒸馏

作者进一步对学生模型进行蒸馏，蒸馏损失可以写成：

式中, 是 Softmax 函数, 是 散度损失函数。

最后的损失函数可以写成：

式中， 取决于具体的任务。比如图像分类里面就是 CE Loss。对于 BERT 预训练，它应该是 masked language modeling loss 和 next sentence prediction loss。

1.6 视觉任务实验结果

如下图3所示是 ImageNet 图像识别任务的实验结果。预训练的数据集使用的是 ImageNet-1K，基线模型选择 WideNet。Benefits 是8式所定义的那样，表示在知识集成之后，Dense 模型保留了多少改进。根据图3可以发现，SVD Knowledge Gathering 的方法的性能最好，因此在接下来的实验中，作者将基于 SVD 的方法设置为默认选择。可以观察到 OneS-L-SVD 在 ImageNet 上实现了 78.4% 的 top-1 准确率，只有 15M 参数。与最强的 Switch-L 相比，本文的模型提高了 0.6%。OneS-L-SVD 实现了与 ViT-B 相当的性能，只有 17% 的可训练参数。这些结果显示了知识整合的有效性。

1.7 文本任务实验结果

如下图4所示是 NLU 任务实验结果。预训练的数据集使用的是 English Wikipedia 和 BOOKCORPUS，微调时候的评估基准是 General Language Understanding Evaluation (GLUE)，以及两个不同的 Stanford Question Answering (SQuAD) 数据集。基线模型选择 ALBERT。

实验结果如下图4所示，可以看到 OneS 方法在所有任务上都优于基线方法 (即 Distill 和 Switch)。在四个任务上，OneS 平均超过 Switch 0.42。在 SQuAD1.1 和 SST-2 等少数任务中，OneS 甚至可以胜过教师 MoE 模型 WideNet。MoE 模型倾向于在小数据集上过拟合。OneS 具有 MoE 的知识，但是结构依然是密集的，因此预训练的优势可以更容易地迁移到下游任务中。

与 MoE 模型相比，OneS 的另一个优点是推理速度。MoE 模型之所以慢的原因是：条件计算带来的门控函数，稀疏 einsum 算子，以上缺点会降低计算的效率。但是，OneS 模型可以实现 3.7 倍的推理速度。OneS 模型高效率的另一个重要原因是密集模型比稀疏 MoE 模型更硬件友好。

1.8 消融实验结果

为了进一步验证本文提出的知识聚合和知识蒸馏的有效性，本文进行了如下图5所示的消融实验。可以看到，在没有知识聚合的情况下，性能会显著地下降，这表明预训练的稀疏模型中包含的知识对于提高学生模型的性能至关重要。在没有知识蒸馏的情况下，性能会略微下降，说明知识蒸馏有一定帮助，教师可以帮助学生挖掘收集到的知识。对于大模型，去除知识收集和知识蒸馏也会损害性能。

公众号后台回复“数据集”获取100+深度学习各方向资源整理