Single-gated MoE : CV领域MoE模型应用

文章链接：Revisiting Single-gated Mixtures of Experts

会议：BMVC 2022 （CV领域C类会议）

1.背景介绍

在计算机视觉（computer vision，CV）领域，神经网络的设计目的是为任何给定图像提取一组固定的细节特征。但不同图像呈现出的内容和层次是相异的，彼此的复杂度也不相同。如果对于简单的图像（比如空白背景）进行复杂图像一样的处理操作，从效率上讲，可能是计算资源的浪费。
这一现象在conditional computing(条件计算)和early-exiting(早退机制) 领域已经有了很详尽的研究。在条件计算中，神经网络会根据图像动态地打开或关闭网络的子部分。在早退机制中，对于简单的图像，网络可以提前完成预测，避免不必要的进一步计算。

近年来，MoE在条件计算方面取得了很大进展，比如在NLP领域，得益于MoE的易拓展性和次线性增加的推理成本，基于MoE的大参数transformer模型正成为新常态。类似的模型也开始出现在CV领域，训练模型使用了非常大的数据集和大量的Routing decisions。

这些大规模MoE条件模型的成功引发了这样一个问题:对于规模更小的数据集和架构(如ImageNet上的ResNet-18)，是否也可以得到类似的结果?

---

2.内容摘要

文章提出了一种简单的single-gate MoE模型，用于训练大规模模型，在推理时计算成本合理。引入了一个 基网络 分支，它的特征有助于初始专家选择，该 基网络 在与专业专家进行集成时可以作为一个很好的正则化器，搭建了一个简单高效的异步训练管道来避免训练不稳定性，提出了一个简单的阈值规则来动态地适应推理时的计算预算。实验结果表明，该模型实现了与更复杂的MoE模型相当的效率与精度权衡，并且优于非混合基线。

模型介绍

本文设计的MoE模型主要用于图像分类。

MoE模型由K个$e_{k,k\in [1,\mathrm{2...}K]}$组成，每个专家输出图片在目标类的分布，专家的执行取决于Gate (or router) $g$，最终结果是所有专家的概率分布和。对于图像分类任务，本文的目标是最大化训练数据集$D$上的模型总似然值：
$$\mathcal{L}(D)={\mathbb{E}}_{(x,y)\in D}\left[\sum _{k=1}^{K}g(k|x)\cdot e_k(y|x)\right]$$
成功的MoE模型依赖于Gate学习跨K个集群的输入空间解构，这样专家就可以专门研究结果子集。这种假设是有依据的，比起在整个数据集上训练单个模型，组合式的方法更具优势。在推理时，对门进行阈值设置，以便只执行一个或几个专家，以控制准确性和效率的平衡。

从上述公式可以看出MoE模型可以被看作是由 Gate 权重的 Experts 组成的集合，而非仅为了提高效率而只使用排名第1位的 Expert 。对 Gate 输出设置阈值，以确定在推理时动态包含哪些专家，这是MoE模型常常用到的一个技巧。

挑战

• 受Gate控制，experts 对于训练集只有一个局部视野，这意味着它们更容易过拟合；
• Gate 和 experts 的联合训练提出了一个 “chicken-and-egg” 问题：Gate 必须将样本发送给最有可能成功分类的 experts，但较弱的 experts 需要数据来改进；
• Gate 定义的初始数据子集对 experts 训练有很大的影响，例如在初始化时有的 expert 比其他人更受青睐。

措施

1. 泛型知识分支（基础模型$\Phi$）：
   • 处理潜在的过拟合：它作为专家的一种正则化形式，将专家们的输出与$\Phi$初始基础预测相结合；
   • 初始化 Gate ，利用$\Phi$诱导的特征空间对训练样本进行聚类；
   • 作为一个早退机制，避免在不必要的时候执行任何专家，并根据输入图像有条件地激活。
2. 简单轻量的训练方案：
   该方案首先通过对基本模型的嵌入进行聚类来初始化专家子集，然后在训练过程中保持 Gate 和experts 的独立性，以避免门崩溃的问题。

本文提出的架构维护一个全深度的基本模型$\Phi$，在训练阶段处于冻结状态。$\Phi$主要有三个作用：
1.作为Gate和Experts的输入（$\Phi$早期层）；2.与专家输出集成；3. 作为推断的早退出口。
模型定义一个Gate $g$，它根据输入 $x$ 的当前表示选择执行哪个专家。在推断时，选择唯一的最可能执行的专家。$$g_{test}(k|x)=\mathbb{I}(k=argma{x}_{k^{{}^{\prime }}}g(k^{{}^{\prime }}|x))$$，其中$\mathbb{I}(\cdot )$为指示函数。
Experts的输入是$\Phi$的中间特征映射。这种设计选择有两个好处:(1)共享和冻结了专家的早期特征，减少了可训练参数的数量，降低了专家过拟合的风险（特别是在小数据集上训练时);(2)这允许模型在推断时重复使用来自$\Phi$的计算，从而进一步提高效率。

---

模型实现

训练

算法1展示了 Gate 和 Experts 并行独立训练的训练过程。该训练逻辑的设计有三个出发点：

1. 为了避免 Gate 倒塌，在训练 Experts 时保持 Gate 的权重固定。这会令模型严重依赖于 Gate 初始化。因此，为了定义一个有意义的初始 Gate $g_0$，使用K-means对预训练的基本模型$\Phi$的嵌入进行了聚类;
2. 由于未校准的输出，与experts $e_k$联合训练一个集成开发器$d_k$，常常导致$d_k$严重偏向于基本模型$\Phi$，使专家无法进行专门化。为了避免这个问题，在对相应的专家$e_k$进行充分的训练后才开始训练$d_k$。
3. 由于 experts 是用基本模型$\Phi$的预先训练的权重初始化的，但随后会在 Gate 分配的一个特定子集上进行训练，因此如果在推理阶段 Gate 分给它们一个从未见过的类（比如这个类是交由隔壁Expert处理的)，会导致该专家输出的结果乱套。尽管提出的集合器通过提供额外的正则化部分缓解了这个问题，将非指定样本分发给 experts 通常是有益的。这在上述算法的步骤4和步骤5中有体现， Gate $g_0$通过 $\Gamma :\cdot \mapsto clip(\cdot ,\gamma ,1.)$ ，其中$\gamma$是超参数，的变换进行样本分配的平滑。本文提出了两种方法来达到这个目的：
   • 用平滑Gate weights对所有样本的损失重新加权，包括负样本(如算法1所示)
   • 在组建训练批时，用这些Gate weights作为抽样概率
   • 前者实现起来更简单，后者对于大型数据集更实用，可以帮助模型更快的收敛，从而减少训练时间

推理

通过早退机制和动态集合进行推理。
本文设计的模型框架集成了一个简单的提前退出选项，在简单的情况下直接输出基本模型的预测，从而进一步提高计算效率。这一成果借鉴了前人[1，2，3]的研究成果，在推理时对基础模型的置信度进行阈值化，来决定是提前退出还是执行门选专家，类似于二分类。我们可以引入下方公式：

$${\alpha }_k(x)=g(k|x)\cdot (1-ma{x}_y\varphi (y|x))$$
其中，${\alpha }_k(x)$可以解释为样本$x$不早退，且 Gate 将$x$ routing 到第$e_k$的联合概率。从直观上看，如果${\alpha }_k(x)$的结果大小对于所有 Experts 来说都低于某一阈值，则意味着基本模型$\Phi$具有较高的置信度， Gate 并没有把握地将$x$ routing 到任何 expert，因此该样本应该尽早退出。
给定一个训练有素的Gate $g$和 experts (及其集成电路)e，将早退机制和 experts 集合行为结合起来，可以定义任意时间模型状态$p^{at-\tau }$，$\tau \in [0,1]$是个超参数，变化$\tau$能够使模型在推断时间快速地实现大范围的计算预算，而不需要任何再训练:
$$\begin{gathered} ee(x)=1⟺\forall k\in [1,K],{\alpha }_k(x)<\tau \\ {p}^{at-\tau }(y|x)=ee(x)\cdot \varphi (y|x)+(1-ee(x))\sum _K^{k=1}\mathbb{I}({\alpha }_k(x)\ge \tau )\cdot g(k|x)\cdot e_k^{{}^{\prime }}(y|x) \end{gathered}$$

实验结果

作者在不同尺度的数据集上进行了实验:CIFAR-100， tiny-ImageNet (imagenet的一个缩小的子集，有200个类和110k张图片)，以及ILSVRC2012。

模型使用不同深度的ResNets作为主要骨干架构。对于CIFAR-100和tiny-ImageNet，采用ResNets的一种改进版本 tiny-ResNet(tr)，该版本消除了文献中常用的前两种降尺度操作(大步卷积和max-pooling)。

表一报告了不同深度的tiny-ResNets在CIFAR-100和tiny-ImageNet上的结果。所有结果都是由20名专家报告的，结果表明，即使是简单的Single-gated MoE模型也可以显著提高与标准CNNs相比的效率/精度权衡。而且异步训练方案允许跨多个设备独立地训练专家。

在准确性方面，使用更深入的基础模型通常比使用更深入的专家更有益。

表2结果显示:所有的MoE模型都优于总体基线。此外对于专家数量对模型的影响的分析，增加更多的专家会导致更高的专门化，但也有更多潜在的Routing fault。因此，增加专家的数量对模型是否有利并不明显。

表3中报告了ResNet和MobileNet骨干的ImageNet实验结果。之前的观察结果仍然成立:MoE模型模型优于基线和主干，基于基础模型置信度的早退机制有助于进一步减少计算量，从而带来有限的精度下降。

3.文章总结

本文重新讨论了 Single-gated MoE 模型。具体来说，文章使用了一个新颖的集成方案和一个利用基于集群的初始化的简单异步和稳定的训练管道来增强MoE模型。模型的推理精度高于分层分类器和1专家集成基线，揭示了每个样本 routing 训练专业专家的好处。此外，将基本模型作为一个独立的分支，这允许系统嵌入一个简单的基于阈值的条件规则。