Swin Transformer V2！MSRA原班人马提出了30亿参数版本的Swin Transformer！

【写在前面】

在本文中，作者提出了将Swin Transformer缩放到30亿个参数的技术 ，并使其能够使用高达1536×1536分辨率的图像进行训练。通过扩大容量和分辨率，Swin Transformer在四个具有代表性的视觉基准上创造了新的记录：ImageNet-V2图像分类的84.0%top-1 准确度，COCO目标检测上的63.1/54.4box / mask mAP ，ADE20K语义分割的59.9 mIoU ，以及Kinetics-400视频动作分类的86.8%top-1 准确度。目前，视觉模型尚未像NLP语言模型那样被广泛探索，部分原因是训练和应用中的以下差异:

1） 视觉模型通常在规模上面临不稳定性问题；

2） 许多下游视觉任务需要高分辨率图像，如何有效地将低分辨率预训练的模型转换为高分辨率模型尚未被有效探索。当图像分辨率较高时，GPU显存消耗也是一个问题。

为了解决这些问题，作者提出了几种技术，并在本文中以Swin Transformer进行了说明：

1）提高大视觉模型稳定性的后归一化（post normalization） 技术和缩放余弦注意力（scaled cosine attention） 方法；

2）一种对数间隔连续位置偏差技术（log-spaced continuous position bias technique） ，用于有效地将在低分辨率图像中预训练的模型转换为其高分辨率对应模型。

此外，作者还分享了重要的实现细节，这些细节可以显著节省GPU显存消耗，从而使使用常规GPU训练大型视觉模型成为可能。利用这些技术和自监督的预训练，作者成功地训练了一个强大的30亿参数的Swin Transformer模型，并将其有效地迁移到涉及高分辨率图像的各种视觉任务中，在各种基准上实现了SOTA的精度。

1. 论文和代码地址

Swin Transformer V2: Scaling Up Capacity and Resolution

论文地址：https://arxiv.org/abs/2111.09883

代码地址：尚未开源

2. Motivation

扩大语言模型的规模已经取得了成功。它显著地提高了模型在语言任务上的表现，并且该模型展示了与人类相似的Zero-shot能力。BERT大型模型具有3.4亿个参数，语言模型在几年内迅速扩大了1000倍以上，达到5300亿个密集参数和1.6万亿个稀疏参数。

另一方面，视觉模型的扩展一直滞后。虽然人们早就认识到，较大的视觉模型通常在视觉任务上表现得更好，但最近，模型大小刚刚能够达到约10-20亿个参数。更重要的是，与大型语言模型不同，现有的大型视觉模型仅适用于图像分类任务。

为了成功地训练大型和通用的视觉模型，需要解决几个关键问题。首先，对大型视觉模型的实验揭示了训练中的不稳定性问题。作者发现，在大型模型中，各层之间的激活幅度差异显著增大。仔细观察结构可以发现，这是由直接添加回主分支的残差单元的输出引起的。结果是激活值逐层累积，因此深层的振幅明显大于早期层的振幅。为了解决这个问题，作者提出了一种新的归一化配置，称为post norm，它将LN层从每个残差单元的开始移动到后端 ，如上图所示。作者发现，这种新的配置在网络层上产生了更温和的激活值。作者还提出了一种缩放余弦注意（scaled cosine attention）来取代以前的点积注意（dot product attention） 。缩放余弦注意使得计算与块输入的振幅无关，并且注意值不太可能陷入极端。在本文的实验中，提出的两种技术不仅使训练过程更加稳定，而且提高了精度，特别是对于较大的模型。

其次，许多下游视觉任务，如目标检测和语义分割，需要高分辨率的输入图像或大的注意力窗口。低分辨率预训练和高分辨率微调之间的窗口大小变化可能相当大。当前的常见做法是对位置偏移map执行双三次插值（bi-cubic interpolation）。这个简单方法的结果通常是次优的。作者引入了一种对数间隔连续位置偏差（Log-CPB），它通过在对数间隔坐标输入上应用一个小型元网络，为任意坐标范围生成偏差值 。

模型容量和分辨率的放大也会导致现有视觉模型的GPU显存存消耗过高。为了解决显存问题，作者结合了一些重要技术，包括zero optimizer 、activation check pointing 和顺序自注意计算的新实现 。通过这些技术，大大降低了大型模型和分辨率的GPU显存消耗，对训练速度的影响微乎其微。

通过上述技术，作者成功地训练了一个30亿的Swin Transformer模型，并使用Nvidia A100-40G GPU将其有效地迁移到图像分辨率高达1536×1536的各种视觉任务中。在本文的模型预训练中，作者还采用了自监督预训练来减少对超大标注数据的依赖。与JFT3B数据集相比，标注数据减少了40倍，30亿模型在广泛的视觉基准上达到了SOTA的精度。

3. 方法

3.1. A Brief Review of Swin Transformer

Swin Transformer是一个通用的计算机视觉主干网络，它在各种粒度的识别任务上都有很强的性能，包括区域级的目标检测、像素级的语义分割和图像级的图像分类。Swin Transformer的主要思想是将几个重要的视觉信号先验引入到vanilla Transformer编码器结构中，包括层次性 、局部性 和平移不变性 ，这结合了两者的优点：基本Transformer单元具有强大的建模能力 ，视觉信号先验使得它对各种视觉任务都很友好 。

Normalization configuration

众所周知，归一化技术对于训练更深层次的网络结构以及稳定训练过程至关重要。原始的Swin Transformer继承了language Transformers和vanilla ViT中的设置，利用了预归一化设置。

Relative position bias

相对位置偏差是原始Swin Transformer中的一个关键设计，该Transformer引入了一个额外的参数偏差项，用于解释自注意计算中的几何关系：

\operatorname{Attention}(Q, K, V)=\operatorname{SoftMax}\left(Q K^{T} / \sqrt{d}+B\right) V

B \in \mathbb{R}^{M^{2} \times M^{2}} 是每个head的相对位置编码；Q, K, V \in \mathbb{R}^{M^{2} \times d} 是query, key 和 value矩阵。d是query, key 的维度，M^{2} 是窗口中的patch数。相对位置偏差计算视觉元素的相对空间配置，并在各种视觉任务中显示出其关键性，特别是对于密集识别任务，如目标检测。

在Swin Transformer中，每个轴的相对位置范围位于[−M+1，M− 1]，相对位置偏差由偏差矩阵 \hat{B} \in \mathbb{R}^{(2 M-1) \times(2 M-1)} 进行参数化。当在不同窗口大小之间转换时，预训练中学习的相对位置偏差矩阵用于通过双三次插值方法在微调中初始化不同大小的偏差矩阵。

Issues in scaling up model capacity and window resolution

作者观察到在缩放Swin Transformer的容量和窗口分辨率方面存在两个问题。

1）放大模型容量时出现的不稳定性问题。 如上图所示，当将原始的Swin Transformer模型从小尺寸放大到大尺寸时，深层的激活值会急剧增长。振幅最高和最低的层之间的差异达到了10 ^4 。当进一步将它扩展到一个巨大的规模（6.58亿个参数）时，它无法完成训练，如下图所示。

2）跨窗口分辨率迁移模型时性能降低。 如上表第一行所示，当通过双三次插值方法在更大的图像分辨率和窗口大小上直接测试预训练的ImageNet-1K模型（256×256图像，窗口大小为8×8）的精度时，精度显著下降。

3.2. Scaling Up Model Capacity

作者观察到，当放大模型容量时，深层的激活值显著增加。事实上，在预归一化（pre-normalization）配置中，每个残差块的输出激活值直接合并回主分支，并且主分支的振幅在更深层会越来越大。不同层中的振幅差异过大可能会导致训练不稳定问题。

Post normalization

为了缓解这个问题，作者提出使用后归一化（post normalization） 方法。在这种方法中，每个残差块的输出在合并回主分支之前被归一化，并且当层加深时，主分支的振幅不会累积。如上图所示，这种方法的激活幅度比原始预归一化配置中的激活幅度要温和得多。

Scaled cosine attention

在原始的自注意计算中，像素对的相似性项被计算为query和key向量的点积。作者发现，在使用这种方法进行大型视觉模型时，某些block和head的学习注意力通常由几个像素对控制，特别是在后归一化配置中。为了缓解这个问题，作者提出了一种缩放余弦注意（scaled cosine attention）方法，该方法通过缩放余弦函数计算像素对i和j的注意力：

\operatorname{Sim}\left(\mathbf{q}_{i}, \mathbf{k}_{j}\right)=\cos \left(\mathbf{q}_{i}, \mathbf{k}_{j}\right) / \tau+B_{i j}

其中B_{i j} 是像素i和j之间的相对位置偏移；τ是一个可学习的标量，在head和层之间不共享。余弦函数是自然归一化的，因此可以具有较低的注意值。

3.3. Scaling Up Window Resolution

在本小节中，作者提出了一种对数间隔连续位置偏差（log-spaced continuous position bias） 方法。

Continuous relative position bias

连续位置偏差法不是直接优化参数化偏差，而是在相对坐标上采用小型网络：

B(\Delta x, \Delta y)=\mathcal{G}(\Delta x, \Delta y)

其中，\mathcal{G} 是一个小型网络，默认情况下是中间有ReLU激活的2层MLP。

网络\mathcal{G} 生成任意相对坐标的偏差值，因此可以自然地转换为具有任意不同窗口大小的微调任务。在推理中，每个相对位置处的偏差值可以预先计算并存储为模型参数，因此，与原始参数化偏差方法相比，它在推理时同样方便。

Log-spaced coordinates

当在很大程度上改变窗口大小时，将有很大一部分相对坐标范围需要外推。为了缓解此问题，作者提出使用对数间隔坐标，而不是原始线性间隔坐标：

其中\Delta x, \Delta y,\widehat{\Delta x}, \widehat{\Delta y} 分别为线性比例坐标和对数间隔坐标。

3.4. Other Implementation

Implementation to save GPU memory

另一个问题是在容量和分辨率都很大的情况下，常规实现无法负担GPU显存消耗。为了解决显存问题，作者采用以下实现：

Zero-Redundancy Optimizer (ZeRO)

数据并行模式的常规优化器实现向每个GPU或主节点广播模型参数和优化状态。这对于大型模型非常不友好，例如，当使用AdamW优化器和fp32权重时，一个30亿参数的模型将消耗48G GPU显存。通过ZeRO优化器，模型参数和相应的优化状态将被划分并分配到多个GPU，从而显著降低显存消耗。作者使用DeepSpeed框架，并在实验中使用了ZeRO 优化器。这种优化对训练速度几乎没有影响。

Activation check-pointing

Transformer层中的特征映射也会消耗大量GPU内存，这在图像和窗口分辨率较高时会构成瓶颈。Activation check-pointing优化将使训练速度最多降低30%。

Sequential self-attention computation

以非常大的分辨率训练大型模型，在采用上述两种优化策略后，常规GPU（40GB显存）仍然无法承受。作者发现，在这种情况下，自注意力模块构成了一个瓶颈。为了缓解这个问题，作者实现了顺序的自注意计算，而不是使用以前的批处理计算方法。该优化在前两个阶段应用于各层，对整体训练速度几乎没有影响。

通过这些方法，作者成功地使用Nvidia A100-40G GPU训练了一个3B模型，用于输入图像分辨率为1536×1536的COCO目标检测和输入分辨率为320×320×8的Kinetics-400动作分类。

Joining with a self-supervised approach

更大的模型更需要数据。为了解决数据匮乏的问题，以前的大型视觉模型通常要么利用巨大的标注数据，如JFT-3B，要么利用自监督的预训练。在这项工作中，作者结合了两种策略：一方面，作者将ImageNet-22K数据集适度放大5倍，以达到7000万张带噪标签的图像 ；虽然该数据量表仍远远落后于JFT3B，但作者还采用了一种自监督学习方法 来更好地利用该数据。通过结合这两种策略，作者训练了一个30亿参数的强Swin-Transformer模型，并在几个具有代表性的视觉基准上实现了SOTA的精度。

3.5. Model configurations

作者为4种配置的Swin Transformer V2保持原始Swin Transformer的阶段、块和通道设置：

C表示第一阶段中隐藏层的通道数。

作者进一步将Swin Transformer V2扩展到huge尺寸和giant尺寸，分别具有6.58亿个参数和30亿个参数：

对于SwinV2-H和SwinV2-G，作者进一步在主分支上每隔6层引入一个层归一化单元。

4.实验

4.1. Scaling Up Experiments

ImageNet-1K image classification results

上表展示了一系列大模型在ImageNet-1K V1和V2上的实验结果。

COCO object detection results

上表展示了SOTA方法和SwinV2在COCO目标检测上的实验结果。

ADE20K semantic segmentation results

上表展示了SOTA方法和SwinV2在ADE20K语义分割上的实验结果。

Kinetics-400 video action classification results

上表展示了SOTA方法和SwinV2在Kinetics-400视频动作分类上的实验结果。

4.2. Ablation Study

Ablation on post-norm and scaled cosine attention

上表展示了post-norm和cosine attention的消融实验结果。可以看出这两个模块对于模型性能都有促进作用。

上表展示了不同模型大小下，L-CPB的消融实验结果。

Scaling up window resolution by different approaches

上表通过分别在ImageNet-1K图像分类、COCO目标检测和ADE20K语义分割的3个下游视觉任务中，将预训练中的分辨率从256×256缩放到更大的尺寸，消融了3种方法的性能。

5. 总结

在本文中，作者提出了将Swin Transformer缩放至30亿个参数的技术，使其能够使用高达1536×1536分辨率的图像进行训练，包括后归一化 和缩放余弦注意 ，以使模型更易于在容量上进行缩放；以及对数间隔连续相对位置偏差 方法，该方法可使模型更有效地跨窗口分辨率传输。经过调整的结构名为Swin Transformer V2，通过扩展容量和分辨率，它可以在多个数据集上达到SOTA的结果。