专栏名称: 极市平台

极市平台是由深圳极视角推出的专业的视觉算法开发与分发平台，为视觉开发者提供多领域实景训练数据库等开发工具和规模化销售渠道。本公众号将会分享视觉相关的技术资讯，行业动态，在线分享信息，线下活动等。网站: http://cvmart.net/

MoE 系列超详细解读 (四)：MoE 知识蒸馏：从稀疏教师模型到密集学生模型

极市平台 · 公众号 · · 2024-03-09 22:00

正文

↑ 点击蓝字关注极市平台

作者丨科技猛兽

编辑丨极市平台

极市导读

该模型保留了61.7%的MoE模型的效益，并在15M参数的情况下达到了78.4%的top-1准确率。 >> 加入极市CV技术交流群，走在计算机视觉的最前沿

本文目录

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)。