GreenMiM：接地气的MiM，强强联合Swin Transformer和何凯明的MAE，没有大显卡也能跑起来

ArXiv: https://arxiv.org/abs/2205.13515
Open Code: https://github.com/layneh/greenmim

摘要

这篇文章提出了一种针对MiM（Masked image modeling，遮挡图像建模）采用分层ViT如SwinT的有效方法，允许分层ViT丢弃掩码块只在可见图像块上操作。这个方法包含两个关键组成：第一，对于窗口注意力，设计了组窗口注意力的方案，按照分而治之的策略。为了减轻自注意力相对图像块数量的平方复杂度，组注意力采用了同一分区，即每一个任意尺度的局部窗口内的可见图像块能够被分成相同大小的组别，其中掩码自注意在每一个分组内执行。第二，进一步提升分组策略，通过动态规划算法去最小化分组图像块上注意力的总体运算成本。因此，MiM可以以一种绿色的高效的方式在分层ViT上工作。例如，我们可以训练分层ViT，快2.7倍，同时减少GPU内存70%，然而在ImageNet分类上保持有竞争力的性能，以及在下游的COCO目标检测上体现优势。

---

---

介绍

在NLP中的MLM（遮挡语言建模）的巨大成功和ViT的先进性的推动下，MiM以一种有前途的自监督预训练模式而出现，对于计算机视觉而言。MiM从无标注的数据中学习表征，以掩码预测的方式，例如预测随机掩码的图像输入块的离散令牌，隐特征，或原始像素等。其中，有代表性的工作即MAE，展示了有竞争力的性能和令人印象深刻的效率。本质上，MAE提出了针对MiM的一个非对称的编码-解码结构，其中，编码器例如一个标准的ViT模型只在可见的图像块上操作，和一个轻量级的解码器为了掩码预测而复原所有图像块。

一方面，这个非对称的编码解码架构极大地减小了预训练的运算负担；另一方面，MAE仅支持相同尺度的ViT架构作为编码器，但是大多数现代视觉模型采用分层的架构，部分原因是需要处理视觉元素的多尺度变化。实际上，这个分层的架构和局部的传导偏置在不同的CV任务上很关键，它需要不同级别或尺度的表征去做预测，包括图像识别，目标检测。然而，如何将分层的ViT如SwinT直接集成到MAE中是不明显的。并且，尽管SimMiM已经探索Swin用于MiM，但它在可见的和掩码的图像块上操作导致了严重的计算开销。举一个具体的例子，即使基本大小的SimMiM模型也不能在8个32G的GPU的单体机上训练，更不用说更大的模型。这个计算开销让广大研究者难以深入该领域的研究，更不用提模型开发过程中的碳排放量了。

因此，我们本着绿色AI的精神，努力为具有层次模型的MiM设计新的绿色方法。我们的工作专注MAE这种非对称的编码解码架构扩展到分层的视觉Transformer中，特别是具有代表性的SwinT模型中，以便仅对可见块进行有效的预训练。我们确定主要的障碍在于局部窗口注意力的限制。尽管广泛用于分层视觉Transformer中，局部窗口注意力不能很好配合随机掩码，因为它创建了各种大小的局部窗口，而不适用于并行计算。

方法概述。输入图片首先被随机掩码，然后输入到一个四阶段的分层ViT架构中，最后一个轻量级的decoder获取可见块的表征和掩码令牌去重构遮挡部分。

这篇文章尝试解决这个问题。我们的方法概念简单，包含两个组成。第一，借鉴分而治之的原则，我们提出了一个组窗口注意力方法，第一次将不均匀数量的可见补丁块的局部窗口划分为几个相等大小的分组，然后在组内使用掩码注意。第二，我们将上述分组形式化成一个带有约束的优化问题，其目标是找到一组划分使得分组后的令牌上计算注意力的开销最小化。启发于动态规划和贪心算法，我们提出了一个优化分组算法，能够动态选择优化的分组大小和将局部窗口分区成最小数量的组别。我们的方法是通用的，不需要针对骨架作出任何的修改，因此我们能够对基于可见和掩码的补丁块的基线方法作出平等的对比。在我们的实验评估中，发现我们的发现需要的训练时间更短，消耗的GPU内存更少，但是与基线方法取得相当的性能。更准确的是，使用Swin-B作为骨架，我们的方法只需要SimMiM一半的训练时间和约40%的GPU内存消耗，然后再ImageNet-1K上取得了top1的83.7%的精度。

方法

基础

Notations. 符号

假设$X\in R^{C\times H\times W}$表示输入特征，其中$C,H,W$表示$X$的通道数，高宽； M ∈ { 0 , 1 } H × W M \in \{0,1\}^{H \times W} M∈{0,1}H×W表示在训练过程中随机产生的掩码，其中$0$表示补丁块对于编码器是不可见的，反之。

Masked Image Modeling

MIM通过从输入X的部分观察$\hat{X}\leftarrow Mask(X,M)$预测遮挡部分，以此来学习表征。当前的MIM方法根据$Mask(\cdot ,\cdot )$操作，可以分为两个分类。多数方法使用哈达玛乘积（即对应位置的元素相乘）来掩码，并保留补丁块，通过在通道维度上进行如下计算$C$次，即$\hat{X}\leftarrow Mask(X,M)=X\odot M$。不同以上方法，MAE提出在掩码部分抛弃遮挡部分的补丁块，即
X ^ ← M a s k ( X , M ) = { X i , j : M i , j = 1 } . ( 1 ) \hat X \leftarrow Mask(X,M) =\{X_{i,j}:M_{i,j}=1\}. (1) X^←Mask(X,M)={Xi,j​:Mi,j​=1}.(1)
MAE设计了一个非对称的和各层同分辨率大小的编码解码器去利用部分输入：其中编码器仅在可见的补丁块$\hat{X}$上操作，而不包含掩码令牌；解码器在可见块的表征和遮挡令牌上重构原始的图像。这个设计使得MAE获得可竞争的性能和令人印象深刻的效率，例如比在所有的图像块上操作，训练速度提升了三倍。然而，MAE仅仅在各层同分辨率的ViT上有效，并且如何将MAE的高效转移到分层的ViT上不清晰，而在多数任务上它展示了比原始ViT的一致优越性。这篇文章尝试去回答这个问题，并且提出一个更加绿色的方法对于带有分层ViT的遮挡图像建模。

GreenMiM

Base architecture.

我们采用具有代表性的ViT，即SwinT作为我们的基准，它主要包含FFN和滑动窗口注意力。但FFN是在point-wise即逐点操作的，并只在可见的补丁块上操作，窗口注意力无法做到这一点。例如，SwinT的窗口尺度设为$p=7$，窗口注意力首先将特征图$X$分区成 $n_w=\frac{H}{p}\times \frac{W}{p}$个无重叠的局部窗口，即 X → { X i } i = 0 n w X \rightarrow \{X_i\}^{n_w}_{i=0} X→{Xi​}i=0nw​​，其中每个$X_i$包含$p\times p$个图像块。然后多头注意力MSA在每个窗口内部独立地并行执行，每个窗口内的图像块是相等的。然后，当布局窗口内的图像块的数量是不均匀的时候，例如由于发生随机掩码，如何有效并行计算注意力是不清晰的。因此我们提出了一个有效的组窗口注意力策略，直接替换所有的移位窗口注意力，以一种绿色的方法让它只在可见的图像块上执行。

Group Window Attention

继续讨论上述问题，我们提出了一个Group窗口注意力策略，能够在掩码特征上有效提升窗口注意力的计算有效性。参考凯明的MAE做法，即公式(1)能够给出掩码特征$\hat{X}=Mask(X,M)$，我们获取了一组不均匀的布局窗口 X ^ → { X ^ i } i = 0 n w \hat X \rightarrow \{\hat X_{i} \}^{n_w}_{i=0} X^→{X^i​}i=0nw​​，其中每一个窗口只包含可见的图像块，对应的数量为 { w i } i = 0 n w \{w_i \}^{n_w}_{i=0} {wi​}i=0nw​​。如下图所示，文章的组窗口注意力首先使用一个优化的分组算法去讲不均匀的窗口划分成几个相等大小的组别，然后在每个组内执行掩码注意力Masked Attention以免信息泄漏。

带有滑动窗口的Group Window Attention策略。表明，该方法与窗口分区无关。

提示：这里省略如何进行有效地划分，一句话，作者将之概括成背包问题并通过动态规划解决，感兴趣去看文章。

Masked Attention

由于不相邻的局部窗口被划分在同一组内，注意力权值也需要对应掩码以避免这些局部窗口之间的信息交换。如下图所示，在获取注意力图后，我们只保留窗口内注意力权重（即对角线上的元素）和去掉窗口间的注意力权重。一个类似的掩码方法也应用于相对位置偏置的检索，其中我们存储每个令牌的原始绝对位置，并即时计算相对位置以检索对应的偏置。

Illustration of the Masked Attention scheme. Given a group of tokens, we first compute their pairwise attention weights and then set the attention weights between tokens from different local windows to −∞ (indicated by the gray cells). The final attention output is then computed with the masked attention map.
掩码注意力示例。给定一组token，首先计算它们成对的注意力权重，然后将来自不同局部窗口的token之间的注意力权重设置为-∞(即灰色部分的格子)，然后使用掩码注意力图去计算最终的注意力输出。

结论

这篇文章提出了一个GreenMiM框架，使用SwinT作为基线，结合有效的组窗口注意力和动态规划分组策略，该方法能够比SimMiM训练分层ViT快2.7倍，减少70%的GPU用量，同时保持有竞争力的性能。

限制。这个算法的限制是它需要batch-wise掩码方案。尽管这个限制对MiM预训练影响很小，但值得关注。

---