【V-MoE】论文学习笔记
文章目录
- 1 Abstract
- 2 Introduction
- 3 Methods
-
- 3.1 The Vision Mixture of Experts (V-MoE)
-
- 3.1.1 Conditional Computation with MoEs
- 3.1.2 MoEs for Vision
- 3.1.3 Routing
- 3.1.4 Expert's Buffer Capacity
- 3.2 Skipping Tokens with Batch Prioritized Routing（BPR）
-
- 3.2.1 Vanilla Routing
- 3.2.2 Batch Prioritized Routing(BPR)
- 4 Transfer Learning
-
- 4.1 Upstream results
- 4.2 Linear few-shot results
- References
arxiv论文
GitHub代码
1 Abstract
背景：
- MoE在NLP领域表现出出色的可拓展性；
- CV的几乎所有SOTA网络都是 “dense”，i.e.对于每个输入，每个参数都要参与计算 。
- MoE是近年来快速发展的一种模型加速方法。
MoE的思想：
研究者认为在面对不同类型的数据时，model也可以选择更“擅长”的Expert（特定的网络结构，MoE中是MLP）来处理。
在训练过程中，与其对每个输入都要训练整个模型，不如根据输入的不同，仅触发并更新该模型中的某一部分，以降低模型总体的训练代价，学习出更大参数量的模型；
在推理过程时，可以针对不同的下游任务，将大模型拆解得到一个体积更小，但是效果更好的子模型，以实现更加高效的推理。
本质：V-MoE是一个Vision Transformer（ViT） 的稀疏版本。
优点：
- V-MoE有出色的可扩展性；
- V-MoE与最大的dense网络相当的性能；
- 图像识别领域，V-MoE与最新的SOTA网络性能相近，但推理时只需要一半的计算开销。
- Google Brain提出的路由算法可以对对整个batch中的每个输入的子集进行优先级排序，从而实现每幅图像的自适应计算；这使得V-MoE可以更好地权衡性能和计算量；
- V-MoE具有vision model缩放的潜力：在ImageNet上训练一个15B parameter model，精度达到90.35%。
2 Introduction
conditional computaion：通过仅将参数的子集(a subset of paranmeters)应用于每个样本，增加model的容量，与保持训练和推理成本的大致恒定。
目的：Google Brain探索了大规模vision model的条件计算，提出V-MoE。用于图像分类任务。
创新点：
- V-MoE是ViT的一种稀疏变体。V-MoE用稀疏的MoE层代替ViT中密集的前馈层的子集，其中每个图像patch被 “路由” 到 “Expert” 的子集。“Expert”本质上是多层感知机（MLPs）。
- 由于路由的独特失效模式和不可微性，在深度稀疏模型中应用路由充满挑战。Google Brain提出了一种有效的路由算法——批量优先路由算法（Batch Prioritized Routing Algorithm，BPR）：该算法重新利用模型稀疏性来跳过某些patch的计算，从而减少了对不提供信息的图像区域的计算。这使得V-MoE可以很好的权衡性能与计算量，其性能显著优于目前的密集网络。
主要贡献：
- 增大vision model的规模：V-MoE是一个分布式、稀疏激活的ViT。Google Brain训练的model有多达24个MoE层，每层32个Expert（MLPs），几乎达到15B（15-bilion）参数。models可以被稳定地训练，无缝地用于迁移，并只需1000个数据点上成功地fine-tune。这个最大的model在经过fine-tuned后，在ImageNet上实现了90.35%的测试精度。
- 更强的性能：V-MoE在上游任务（upstream：预训练）、小样本学习（Few-shot）和充分finr-tune指标上，显著的优于密集的counterpart。在推断阶段，V-MoE可以调整为：(i)达到最大的密集model的性能，而只使用一半的计算量；(ii)在相同的成本下显著优于最大的密集model。
- 批量优先路由：我们提出了一种新的BPR，允许V-MoE丢弃最没用的patches。这样，我们对每个图像的计算量就减少了。而且，V-MoEs可以在节省20% FLOPs的同时，达到和密集models一样的性能。
- 理论上的进步：提供了一些路由决策的可视化，揭示了有助于激励设计决策的模式和结论，并可能进一步提高对该领域的理解。
3 Methods
3.1 The Vision Mixture of Experts (V-MoE)
3.1.1 Conditional Computation with MoEs
condtional computation旨在针对不同的输入，激活网络的不同子集。
MoE是一种特定的实例化，其中不同模型“Expert”负责输入空间不同的区域。
MoE ( x ) = ∑ i = 1 E g ( x ) i e i ( x ) \operatorname{MoE}(\mathbf{x})=\sum_{i=1}^{E} g(\mathbf{x})_{i} e_{i}(\mathbf{x}) MoE(x)=i=1∑Eg(x)iei(x)
其中， x \mathbf{x} x表示这一层的输入， e i e_i ei表示Expert i i i计算的函数， g i g_i gi表示用来为Expert i i i决定输入条件权重的路由函数。
但是按照这个定义，这仍然是一个密集的网络。然而，如果 e i e_i ei和 g i g_i gi是稀疏的，i.e. model被限制为仅分配 k < < E k<
3.1.2 MoEs for Vision
ViT：将图像处理为patch序列，输入图像首先被分为大小相等的patch，然后被线性投影为token。添加positional embeddings后，tokens由Transformer处理，该Transformer主要由交替的self-attention层和MLP层组成。MLP由2个FC层(fully-connected layers)和GeLU非线性函数组成：
MLP ( x ) = W 2 σ gelu ( W 1 x ) \operatorname{MLP}(\mathbf{x})=\mathbf{W}_{2} \sigma_{\text {gelu }}\left(\mathbf{W}_{1} \mathbf{x}\right) MLP(x)=W2σgelu (W1x)
V-MoE的模型结构和MoE Transformer的结构类似，如下图所示。研究者将ViT中的feedforward层（FFN，即图中的 M L P MLP MLP）替换为多个FFN（即 M L P i , i = 1 , . . . , E MLP_i,i=1,...,E MLPi,i=1,...,E）的混合集成（其中每一个FFN为一个Expert）。之后，对于每一个模块的输入，ViT通过学习路由单元（Router）来选择性地激活若干个Expert，并为每一个Expert赋予其权重。
对于V-MoE，这里用MoE层替换其中的一个子集，其中的每一个Expert都是一个MLP。如上图所示，Expert具有相同的结构： e i ( x ) = M L P θ i ( x ) e_i(\mathbf{x})=MLP_{\theta_i}(\mathbf{x}) ei(x)=MLPθi(x)，但有不同的权重 θ i = ( W 1 ı , W 2 ı ) \theta_{i}=\left(\mathbf{W}_{1}^{\imath}, \mathbf{W}_{2}^{\imath}\right) θi=(W1ı,W2ı)。
3.1.3 Routing
在V-MoE中，每张图片中不同的片段（即不同的Token）将被Router分配到不同的几个FFN中。一般而言，Router会为每个Token选择Top-K个FFN（K一般取1或2），而在每一个ViT模块中大概会有E个FFN（在本论文中E取值为32）。通过这种方式，V-MoE模型的大小可以得到提升，而且在模型规模提升的情况下，每一个Token对应的计算量都不会发生太大变动。
对于V-MOE中的每个MoE层，上述路由函数（Router）定义为：
g ( x ) = TOP k ( softmax ( W x + ϵ ) ) g(\mathbf{x})=\operatorname{TOP}_{k}(\operatorname{softmax}(\mathbf{W} \mathbf{x}+\epsilon)) g(x)=TOPk(softmax(Wx+ϵ))
其中， TOP k \operatorname{TOP}_{k} TOPk是一种操作：将输入向量中除了最大的k个元素外元素都设置为0。实践中，一般取 k = 1 k=1 k=1或 k = 2 k=2 k=2。 x \mathbf{x} x表示网络某一层上的图像token。因此，V-MoE只路由patch representations，而非整幅图像。
不同于MoE中先 TOP k \operatorname{TOP}_{k} TOPk再 softmax \operatorname{softmax} softmax，V-MoE中先 softmax \operatorname{softmax} softmax再 TOP k \operatorname{TOP}_{k} TOPk。这样才使得 k = 1 k=1 k=1s时也可以训练，否则路由梯度会处处为0。
最后，通过添加少量符合正态分布标准差的噪声，进一步提高性能。
3.1.4 Expert’s Buffer Capacity
Routing带来的问题：
在训练期间，稀疏模型可能只训练一小部分Expert 。这种模式会导致两个问题。
- 统计效率低：model专注于单一Expert的学习，该model不会比密集的model更强大。
- 计算效率低：向Expert分配数据的不平衡可能会导致硬件利用率低。
解决方案：
将每个Expert的缓冲容量（Buffer Capacity，即每个Expert处理的tokens数量）固定，并使用辅助损失函数训练model，以保持负载均衡。
Expert的缓冲能力 B e B_e Be：
B e = r o u n d ( k N P C E ) B_e=round(\frac{kNPC}{E}) Be=round(EkNPC)
其中， N N N是batch重要处理的图像数量， P P P是每个图像的token数量， k k k是每个token的选定Expert数量， E E E是总的Expert数量， C C C是容量比率。
如果路由器向Expert分配了多于 B e B_e Be个tokens，则仅处理其中的 B e B_e Be个tokens。剩余的tokens并不完全 “丢掉”，因为它们的信息由剩余的连接保留。另外，如果 k > 1 k>1 k>1，则有会多位Expert处理每个token，tokens永远不会完全丢弃。
如果路由器为Expert分配少于 B e B_e Be个tokens，则其缓冲区的其余部分就用零填充。
Google Brain使用容量比率 C C C来调整Expert的能力。当 C > 1 C>1 C>1时，将添加一个空闲容量，以考虑潜在的路由不平衡。当新数据可能来自与上游训练非常不同的分布时，这种设置非常有助于fine-tuning。在 C < 1 C <1 C<1的情况下，路由器会被迫忽略一些分配。在 3.2 3.2 3.2中提供了一个新算法，利用 C < < 1 C<<1 C<<1来丢弃最无用的tokens，并在推理阶段节省算力。
3.2 Skipping Tokens with Batch Prioritized Routing（BPR）
我们前面无数次提到过：为了提高路由效率，Google Brain提出了一种新的路由算法（BPR）。
基本思想： 允许模型对重要token进行优先排序。通过同时减少每个专家的容量，丢弃最没用的token。直观地，并非每个patch对于给定图像进行分类都同样重要的，例如，可以删除大多数背景patch，使模型仅关注具有相关实体的patch。
3.2.1 Vanilla Routing
路由函数逐行应用于一个batch的输入 x \mathbf{x} x。一个batch处理包含 N N N个图像，每个图像由 P P P个token组成； x \mathbf{x} x的每一行对应于图像的特定token的 D D D维表示。相应地， g ( x ) t , i g(\mathbf{x})_{t,i} g(x)t,i表示第 t t t个token和第 i i i个Expert的路由权重。在所有路由算法中，对于 i < j i
3.2.2 Batch Prioritized Routing(BPR)
为了关注“最重要”的token，作者提出了计算每个token的优先级分数 s ( x ) s(\mathbf{x}) s(x)，并在分配之前对 g ( x ) g(\mathbf{x}) g(x)进行相应排序。根据token的最大路由权重对token进行排序，即 s ( x ) t = m a x i g ( x ) t , i s(\mathbf{x})_t=\mathbf{max}_{i}g(\mathbf{x})_{t,i} s(x)t=maxig(x)t,i。前k个权重的总和用同样的方式计算，即 s ( x ) t = ∑ i g ( x ) t , i s(\mathbf{x})_t=\sum_{i}g(\mathbf{x})_{t,i} s(x)t=∑ig(x)t,i。
作者将路由器输出重用为分配优先级的代理。上图可视化了容量越来越小的Batch Prioritized Routing (BPR)算法的token优先级结果。由于batch中所有图像上的所有token彼此竞争，因此不同的图像可以接收不同的计算量。BPR的算法如下所示：
BPR通过有效地选择合适的tokens，提供了一种减少缓冲大小的方法。这会对整个稀疏model的计算成本产生巨大影响。根据Model Analysis可以得到这样的结论：越深路由决策与图像类型的相关性越强。
4 Transfer Learning
4.1 Upstream results
从上表可以看出，本文的模型在各个数据集上都表现不错，并且在ImageNet上基于fine-tuning设置达到了90.35%的准确率。
上图显示：在JFT-300M和ImageNet数据集上，不同V-MoE和ViT变体的总训练计算量和时间。
4.2 Linear few-shot results
上表还显示了：在1/5/10 shot中，V-MoE性能优于ViT。
上图显示，在5-shot ImageNet下的计算量和训练时间的结果。
References
- Riquelme C, Puigcerver J, Mustafa B, et al. Scaling vision with sparse mixture of experts[J]. Advances in Neural Information Processing Systems, 2021, 34.
- 150亿参数，谷歌开源了史上最大视觉模型V-MoE的全部代码
- ImageNet准确率超过90%！谷歌大脑开源V-MoE，用稀疏条件计算来训练目前最大的视觉模型！（NeurIPS 2021）…
- 迈向更高效通用的加速之路：谷歌提出视觉和多任务MoE方法