Vision MLP 之 Sparse-MLP A Fully-MLP Architecture with Conditional Computation

原始文档：https://www.yuque.com/lart/pa…

从摘要读文章

Mixture-of-Experts (MoE) with sparse conditional computation has been proved an effective architecture for scaling attention-based models to more parameters with comparable computation cost. In this paper, we propose Sparse-MLP, scaling the recent MLP-Mixer model with sparse MoE layers, to achieve a more computation-efficient architecture.

本文引入 MoE 的思维，应用条件计算来实现再裁减参数的同时，还能够保障计算成本不会太高。

We replace a subset of dense MLP blocks in the MLP-Mixer model with Sparse blocks. In each Sparse block, we apply two stages of MoE layers:

1. one with MLP experts mixing information within channels along image patch dimension,
2. one with MLP experts mixing information within patches along the channel dimension.

对于空间 MLP 和通道 MLP 都会应用 MoE 进行革新。

Besides, to reduce computational cost in routing and improve expert capacity, we design Re-represent layers in each Sparse block. These layers are to re-scale image representations by two simple but effective linear transformations.

这里引入了放缩特色维度的构造升高计算量。

When pre-training on ImageNet-1k with MoCo v3 algorithm, our models can outperform dense MLP models by 2.5% on ImageNet Top-1 accuracy with fewer parameters and computational cost. On small-scale downstream image classification tasks, i.e., Cifar10 and Cifar100, our Sparse-MLP can still achieve better performance than baselines.

次要内容

• 将 MoE 技术引入 MLP 模型中，用于替换原始的 MLP 层。其中，通过构建多个具备不同权重的专家，从而扩充了模型的容量和表达能力，并且通过门控路由机制从而束缚了理论应用的专家数量，即所谓的 conditional computation。从而不至于带来过多的计算成本和工夫损耗。
• 因为原始空间 token 数量和通道数量差别较大，这会导致在路由和专家前向计算时不均衡的计算成本（computational cost），所以作者们在空间 MoE 层的前后通过对空间 token 数量和通道数量进行均衡（应用线形层从新投影），从而保障了更加均衡和无效的计算过程。

专家混合结构（Mixture of Experts)

外围操作

$\text{MoE}(x) = \sum_{i=1}^{N}G(x)_iE_i(x), \, G(x) = \text{TopK}(\text{softmax}(W_g(x) + \epsilon)):\mathbb{R}^{D} \rightarrow \mathbb{R}^{N}, \, E_i(x):\mathbb{R}^D \rightarrow \mathbb{R}^D.$

这里从左到右别离是蕴含 $N$ 个专家的 MoE 层的聚合操作，用于计算以输出为条件的路由权重的门控网络（应用 softmax 生成归一化权重，这里引入了噪声 $\epsilon \sim \mathcal{N}(0, \frac{1}{N^2})$ 从而摸索更好的调配策略），以及第 $i$ 个专家层。

相近工作

• MoE 的思维次要来自于 ICLR 2017 的一篇文章：OUTRAGEOUSLY LARGE NEURAL NETWORKS: THE SPARSELY-GATED MIXTURE-OF-EXPERTS LAYER (_Noam Shazeer, Azalia Mirhoseini, Krzysztof Maziarz, Andy Davis, Quoc Le, Geoffrey Hinton and Jeff Dean_)，该文提到“_The capacity of a neural network to absorb information is limited by its number of parameters. Conditional computation, where parts of the network are active on a per-example basis, has been proposed in theory as a way of dramatically increasing model capacity without a proportional increase in computation._”。这篇文章通过在 LSTM 两头插入多个专家构造，通过可训练的门控网络来对专家们进行 稠密组合，从而将模型的能力（capacity）晋升了超过 1000 倍并且只有大量的效率损失。除了次要的想法与这篇文章相似外，针对 MoE 构造的优化策略也根本一脉相承。
• 另外，这篇 SparseMLP 的工作与作者们的另一篇针对 ViT 改良的工作 Go Wider Instead of Deeper 有着潜在的共通点。

Go Wider Instead of Deeper

• 当然，从损失函数上来讲，其实更像 Scaling Vision with Sparse Mixture of Experts 这篇文章。损失函数模式 十分类似。

本文细节

本文次要将 Mixer-MLP 中的最初几层空间和通道 MLP 进行了替换，替换成了 MoE 构造（蕴含空间和通道两种构造）。这样的设定有助于引入更多的参数，晋升模型的能力。

多专家模型的训练是不容易的。次要是因为稠密的门控路由机制导致并不是所有的专家必然能够被失去充沛的训练，也就是所谓的负载不平衡问题。所以应用多专家设定的办法大多数都须要非凡的损失来进行针对性的解决。

对于损失函数，本文连续了之前工作的设定，利用了负载平衡损失（Load Balance Loss）。该损失激励横跨专家对输出进行平衡调配。

该损失蕴含两局部设定：

• 重要性损失（Importance Loss）：目标是让各个专家信息流传过程中的重要性尽量相近，这样能够保障各个专家能够被尽量抉择到并被充沛的训练。

  • 为此，引入了对于专家重要性的定义：$\text{Imp} (X) = \{\sum_{x \in X} \text{softmax} (W_gx)_i\}_{i=1}^{N}$。
  • 权重 $W_g \in \mathbb{R}^{D \times N}$ 示意 MoE 层的门控权重矩阵（gating weight matrix），其将输出的 $D$ 维的 $x$ 的映射到专家数量 $N$，由 softmax 解决后即取得每个样本 $x$ 被调配到各个专家的权重。这里将与第 $i$ 个专家无关的各个输出对应的权重加和后取得其对于 batch 输出 $X$ 的重要性度量。这能够反映出各个专家相较于其余专家，在整体输出被调配解决的过程中所起到的绝对作用。
  • 所以为了尽可能平衡各个专家的重要性，使大家都能更好的“体现”，各个专家的重要性应该尽量平衡。于是应用重要性的平方变异系数（_the squared coefficient of variation of the importance distribution over experts_）作为重要性损失：$L_{imp}(X) = (\frac{\text{std}(\text{Imp}(X))}{\text{mean}(\text{Imp}(X))})^2$。

论文 Scaling Vision with Sparse Mixture of Experts 的表 4

• 负载损失（Load Loss）：重要性损失旨在保障所有专家均匀而言具备类似的路由权重。然而可怜的是，不难想到这些看上去有着总体趋于均衡的权重的路由配置，依然有一小部分专家取得了所有调配（可见上表，尽管输出 1~4 对专家的权重之和均为 2，然而却在最终额抉择中，仅仅只会抉择到专家 1 和 3，而 2 则无奈失去适合的学习）。

  • 为此这里引入了对于专家负载的定义：$\text{Load}(X) = \{\sum_{x \in X} p_i(x)\}_{i=1}^{N}$。
  • $p_i(x) := Pr(G(x)_i >= threshold_k(G(x)))$，示意专家 $i$ 在输出 batch 数据时，对各个样本而言被选中（门控路由大于阈值，即位于前 k 个最大权重的专家中）的概率和。
  • 这个概率看上去不太好搞定，然而作者们这里引入了一个正态分布的噪声，使得所有都能够计算了，大抵如下式，最初是一个正态分布变量的概率的计算：
    $$p_i(x) := Pr(G(x)_i >= threshold_k(G(x))) = Pr(W_g(x)_i + \epsilon >= threshold_k(W_g(x) + \epsilon)) = Pr(\epsilon >= threshold_k(W_g(x) + \epsilon) – W_g(x)_i)$$
  • 负载损失则示意为负载散布的平方变异系数：$L_{Load}(X)=(\frac{\text{std}(\text{Load}(X))}{\text{mean}(\text{Load}(X))})^2$。

所以 MoE 层损失为：$L_{aux} = \lambda(\frac{1}{2}L_{imp} + \frac{1}{2}L_{load})$。这里的超参数 $\lambda$ 用来管制辅助损失在激励跨专家路由的均衡，也保障不会压抑（overwhelm）原始的模型损失。理论与之前的工作设置一样，都设为 0.01。依照之前的工作，这个参数对于性能的影响不太显著。

除了多专家层自身的设定，思考到在原始的 MLP-Mixer 中，基于 patch 的 token 解决形式导致空间 token 数量小于通道数量的 1/3。这对于 MOEs，也就是空间层面上的 MoE 层而言，会导致路由局部与专家局部计算成本的不平衡。所以作者们引入了一个重表征层（re-present layer)来从新调整 MOEs 的输出和输入的空间和通道尺寸。其中次要是通过专门的线性层进行解决，伪代码如下：

理论中设置 $S_1 = 2S, C_1 = C/2$。

这里蕴含两层，一个用于输入，一个用于输出。二者搭配，用于均衡二者两头包裹的 MOEs 的运算（升高 MOEs 运算时的通道数量并减少空间 patch 数量。

试验后果

We find that scaling MLP models in parameters and training them from scratch with limited training data will lead to an overfitting problem. Such finding is consistent with previous work on MLP models (Tolstikhin et al. 2021) and attention based models (Chen, Xie, and He 2021; Dosovitskiy et al. 2020; Xue et al. 2021).

In order to better obtain model capacity, we adopt MoCo v3 algorithm as our default self-supervised training algorithm (He et al. 2019; Chen et al. 2020b; Chen, Xie, and He 2021), and fine-tune models on downstream image classification tasks.

所有的模型都是应用 MoCo V3 自监督训练得来的

论文的融化试验次要探讨了以下四点：

• 专家数量的影响

这里是别离固定 MoEs 和 MoEc 来做试验的，能够看到，MoEs 的减少能够带来性能的晋升。然而 MoEc 却会导致降落，作者们认为是造成了过拟合（对于减少针对通道特色的专家数量会导致过拟合的景象在作者们之前的工作 Go Wider Instead of Deeper 中也有体现）。

• 路由抉择的专家的数量 K

这里针对不同的地位别离尝试了不同的 K 值，这里都是基于 B 构造进行的试验。能够看到，对于通道专家须要同时利用更多，而空间单个即可。

• Sparse Blocks 的地位，即 MoE 构造的地位

对于构造 B，本来 Sparse Blocks 被搁置在最初两个阶段，这里进行了比照，能够看到，这样的构造搁置到模型最初能够起到更好的成果。

• 重表征层的作用

能够看到，应用重表征层后尽管速度晋升了，然而性能却没有降落，反而晋升了。这是个乏味的景象。然而作者没有给出正当的解释和剖析。仅仅是提了下其对于均衡路由和专家的计算成本的作用。那这样的构造间接用于 MLP-Mixer 中是否也会有晋升？

参考链接

• Sparse-MLP: A Fully-MLP Architecture with Conditional Computation: [https://arxiv.org/pdf/2109.02…] (https://arxiv.org/pdf/2109.02…)
• OUTRAGEOUSLY LARGE NEURAL NETWORKS: THE SPARSELY-GATED MIXTURE-OF-EXPERTS LAYER: https://arxiv.org/pdf/1701.06538.pdf
• Go Wider Instead of Deeper: https://arxiv.org/pdf/2107.11817.pdf
• Scaling Vision with Sparse Mixture of Experts: https://arxiv.org/pdf/2106.05974.pdf