兼具Swin和ViT的优势！可用于MAE预训练的超简单层次Transformer结构

2022-06-05 16:06

©PaperWeekly 原创 · 作者 | Jason

研究方向 | 计算机视觉

摘要

最近，掩蔽图像建模（MIM）为视觉 Transformer 的自监督预训练提供了一种新的方法。高效实现的一个关键思想是在整个目标网络编码器中丢弃掩蔽图像 patch 或 token，这要求编码器是普通视觉 Transformer（例如 ViT），但是分层视觉 Transformer（例如 Swin Transformer）在形成视觉输入方面具有潜在更好的特性。

在本文中，作者提出了一种新的分层视觉 Transformer HiViT（Hierarchy ViT 的缩写），它在 MIM 中既具有高效性，又具有良好的性能。关键是删除不必要的“局部单元间操作”，产生出结构简单的分层视觉 Transformer ，其中 mask-units 可以像普通视觉 Transformer 一样序列化。实验结果研究表明，HiViT 在全监督、自监督和迁移学习方面具有优势。特别是，在 ImageNet-1K 上运行 MAE 时，HiViT-B 比 ViT-B 的准确度提高了 0.6%，比 Swin-B 的速度提高了 1.9 倍，性能提高可以推广到检测和分割的下游任务。

论文标题：

HiViT: Hierarchical Vision Transformer Meets Masked Image Modeling

论文地址：

https://arxiv.org/abs/2205.14949

Motivation

在过去十年中，深度神经网络一直是深度学习的基础，并推动了计算机视觉、自然语言处理等研究领域的发展。最近，计算机视觉界见证了从语言模型移植而来的视觉 Transformer 的出现，它取代了卷积神经网络的主导地位。Transformer 具有建立长期特征依赖关系的能力，这自然有利于视觉识别，尤其是在长距离关系很重要的情况下。

根据是否使用多分辨率特征图方面，目前主要有两类视觉 Transformer，即普通视觉 Transformer 和层次视觉 Transformer。虽然后者被认为可以捕获视觉信号的性质（大多数基于卷积的模型都使用了层次结构），但通常使用了一些空间局部操作（即，带有移动窗口的早期自注意）。当需要灵活操作 token 时，这些模型可能会遇到困难。一个典型的例子是掩蔽图像建模（MIM），这是一种最新的预训练视觉 Transformer 的方法——图像块的随机部分从输入中隐藏，分层模型很难确定每对 token 是否需要留下。

本质上，这是因为层次视觉 Transformer 在 masking units 之间使用了非全局操作（例如，窗口注意）。因此，与可以序列化所有加速 token 的普通视觉 Transformer 不同，层次视觉 Transformer 必须保持二维结构，在整个编码器中保留掩蔽 token。因此，分级 Transformer 的训练速度比普通 Transformer 慢 2 倍。

在本文中，作者首先将层次视觉 Transformer 中的操作分为“单元内操作”、“全局单元间操作”和“局部单元间操作”。作者注意到，平面视觉 Transformer 仅包含“单元内操作”（即 patch 嵌入、层归一化、MLP）和“全局单元内操作”（即全局自注意力），因此可以丢弃单元的空间坐标，并且可以将单元序列化以进行高效计算，如 MAE。这就是说，对于层次视觉 Transformer 来说，是“局部单元间操作”（即，移位窗口自注意力、patch 合并）需要基于单元的空间坐标进行额外判断，所以阻碍了序列化和移除 mask 单元。

本文的一个关键观察结果在于，“局部单元间操作”对识别性能的贡献不大——真正有意义的是层次设计（即多尺度特征图）本身。因此，为了使分层视觉 Transformer 适合 MIM，作者删除了“局部单元间操作”，从而形成了一个简单的分层视觉 Transformer，它既吸收了 ViT 的灵活性，又吸收了 Swin Transformer 的优势。在层次视觉 Transformer 中，通常有四个不同分辨率的阶段，其中第三阶段的层数最多，称之为主阶段。作者移除了 Swin 的最后一个阶段，并关闭了所有局部单元间窗口注意力，只在主阶段保留 token 之间的全局注意力。

在实验中，最后一个阶段合并到主阶段（保持模型 FLOPs 不变），早期阶段的局部窗口注意被具有相同 FLOPs 的单元内多层感知器所取代。通过这些最小的修改，作者删除了层次视觉 Transformer 中所有冗余的“局部单元间操作”，其中只采用了最简单的层次结构。与普通 ViTs 相比，本文的模型只在主阶段之前添加了几个空间合并操作和 MLP 层。由此产生的结构被命名为 HiViT，该结构能够对分层视觉信号进行建模，但所有 token 都是最大程度上独立的，并保持操作的灵活性。同时，本文的 HiViT 维护 ViT 范式，与其他层次视觉 Transformer 相比，ViT 范式的实现非常简单。

作者在 ImageNet-1K 上进行了全监督分类实验，以验证 HiViT 的优越性。HiViT 介于 ViT 和 Swin Transformer 之间，与这两个模型相比，其精度有着一致的提高，例如，HiViT-B 报告的 top-1 精度为 83.8%，比ViT-B 高出 2.0%，比 Swin-B 高出 0.3%。通过广泛的消融研究，作者发现去除“局部单元间操作”不会损害识别性能，但层次结构和相对位置编码（不是“局部单元间操作”）略微但持续地提高了分类精度，这使得 HiViT 适用于各种视觉识别场景。

基于 MIM，HiViT 的优势变得更加明显。在 800 个基于MIM的预训练阶段和 100 个微调阶段中，HiViT-B 在 ImageNet-1K 上达到了了 84.2% 的 top-1 精度，比ViT-B（使用 MAE，1600 个阶段的预训练）高出 0.6%，比 Swin-B（使用 SimMIM）高出 0.2%。更重要的是，HiViT 具有在输入阶段丢弃所有掩蔽 patch 的高效实现，因此训练速度是 SimMIM 的1.9 倍（如上图所示）。

其他视觉识别任务的优势依然存在，包括 ImageNet-1K 上的 linear probing（71.3% Top-1 精度）、ADE20K 数据集上的语义分割（48.3%mIoU），以及 COCO 数据集上的目标检测（49.5% AP）和实例分割（43.8% AP）。这些结果验证了删除“局部单元间操作”不会损害通用视觉识别。

本文的核心贡献是 HiViT，这是一种分层视觉 Transformer 结构，可用于广泛的视觉应用。特别是，由于掩蔽图像建模是一种流行的自监督学习范式，HiViT 有可能直接插入许多现有算法中，以提高其从大规模未 token 数据中学习视觉表示的有效性和效率。

方法

2.1 Preliminaries

掩蔽图像建模（MIM）是一种新兴的自监督视觉表征学习范式。流程涉及将部分 mask 的图像提供给目标模型，并训练模型进行恢复。数学上，让目标模型表示为，其中 θ 表示可学习的参数。给定训练图像 x，首先将其划分为几个 patch ，其中 M 是 patch 数。

然后，MIM 随机选择一个子集，将 ID 为的 patch（表示为）送到目标模型（也称为编码器）中，并在其上附加一个解码器，旨在恢复解码器末端的原始图像内容，无论是 tokenized 特征还是像素。如果能够解决这个问题，那么可以认为这些参数已经过良好的训练，可以提取紧凑的视觉特征。

一种适合 MIM 的高效视觉模型是 vanilla vision transformer，缩写为 ViT。在 ViT 中，每个图像块被转换成一个 token，token 通过几个 Transformer 块传播，以进行视觉特征提取。设有 L 个块，第 l 个块取 l-1 个块 token 作为输入，并输出，其中。

每个块的主要部分是自注意力，其中三个中间特征是基于，即查询、键和值，表示为，然后基于这三个中间特征进行自注意力机制。在自注意力计算之后，应用了辅助操作，包括层归一化、多层感知器、残差连接。ViT 可以适用于 MAE，其中 mask 的 token 在 encoder 中能够被丢弃，从而产生更高的建模效率。

直观地说，层次视觉 Transformer（例如，Swin Transformer）更善于捕捉多级视觉特征。它与 ViT 有三个主要区别：

1）架构被划分为几个阶段，在整个正向传播过程中，空间分辨率不是固定的，而是逐渐缩小；2）为了处理相对较大的 token 映射，将自注意计算限制在窗口网格内，并跨层移动窗口分区；3）全局位置编码被相对位置编码取代——这是为了适应窗口注意力机制。尽管层次视觉 Transformer 报告了更高的视觉识别精度，但这些模型在 MIM 方面不如 ViT 有效。因此，之前很少有作品尝试这种组合——例如，SimMIM 将整个图像（mask token 被替换为可学习的 mask token）输入编码器，从而在时间和内存上产生更大的计算成本。

2.2 HiViT: Efficient Hierarchical Transformer for MIM

在本文中，作者追求 MAE 的高效实现，即只有未掩蔽的 token 被馈入编码器——从数学上讲，模型只处理压缩的 token 列表。将其与分层视觉 Transformer （如 Swin Transformer ）集成的主要困难在于“局部单元间操作”，这使得很难序列化 token 并放弃掩蔽的 token。为了删除它们，作者首先将掩蔽单元大小设置为主阶段的 token 大小。掩蔽单元的大小为 16×16，与ViT的恒定 token 大小对齐。然后，作者调整模型如下：

1）对于主阶段之后的操作，作者不允许混合掩蔽和未掩蔽 patch 的 patch merging。为简单起见，作者直接删除了 Swin 中的最后一个阶段，并将相同数量的块附加到第三阶段。由于第三阶段具有较小的 token 维度，这样的操作保存了两个可训练参数，同时将第三阶段中的“局部单元间操作”更改为“全局单元间操作”。因此，模型的识别性能变得更好。

2）对于主阶段之前的操作，作者去除前两个阶段的窗口注意力。也就是说，作者删除了 Swin 的移位窗口，并且没有引入任何其他“局部单元间操作”，例如窗口注意力或卷积。作为替代方案，作者只在前两个阶段使用 MLP 块（用另一个 MLP 层代替自注意力）。令人惊讶的是，如上图所示，这种修改带来了 0.2% 的性能改进。与所示的普通 ViT 相比，派生的结构具有层次性，只需要两个 MLP 块，但在自监督/全监督学习方面具有更好的性能。

上述过程在 Swin Transformer（分层）和 ViT（普通）之间生成架构，上图中说明了该过程。由此产生的架构名为 HiViT（Hierarchy ViT 的缩写），该模型结构简单，为 MIM 带来了高效的实现。具体而言，HiViT 放弃了整个结构中的所有“局部单元间操作”，因此可以在输入层丢弃掩蔽图像块，并且可以在所有阶段消除其计算。因此，HiViT 既具有层次视觉 Transformer 捕获视觉表示的有效性（即识别精度远高于 ViT），又具有平面视觉 Transformer 在掩蔽图像建模任务中的效率（即 MAE 的高效实现可以直接移植，使 HiViT 比 MIM 中的 Swin Transformer 快近 2 倍）。

实验

作者提出了 HiViT-T/S/B 在内的三个模型，其配置如上表所示。

全监督的训练结果显示在上表中。与 ViT 模型相比，所有 HiViT 模型报告的结果均为更优的结果。HiViT-T/B 比 DeiT-S/B 模型分别高出 2.3% 和 2.0%，具有相似的 FLOPs 和更少的参数。与后续模型相比，本文的模型仍然显示出有竞争力的结果。特别是，HiViT-T/S/B 比 Swin-T/S/B 分别高出 0.9%、0.4% 和 0.3%，具有相似的复杂性和较少的参数。此外，本文的所有模型都是参数友好的。例如，与 Swin-T/S/B 模型相比，HiViT-T/S/B 的参数分别减少了 32.2%、24.4% 和 24.4%。

作者进行了全监督的消融研究，以显示本文方法的优势，结果如上表所示。可以发现，从 Swin-B 中删除最后一个阶段（第 4 阶段），同时对第 3 阶段使用全局注意力会带来 0.1% 的性能改进，这意味着最后一个阶段是不必要的。在前两个阶段用 MLP 块代替窗口注意力，将性能提高到 83.8%，这表明窗口注意力在早期阶段是不必要的。

RPE 很重要，去掉它会对性能造成约 0.3% 的损害。如果放弃前两个阶段，使用类似于普通 ViT 的 patch 嵌入方法向下采样 16×，但将块数增加到 24，性能将从 83.5% 下降到 82.9%。然而，这仍然高于普通 ViT 的 81.8%，这意味着层次化输入模块很重要，更深层次的结构比浅层次的结构要好得多。

作者使用 linear probing 度量对预训练模型进行评估，除了可学习的分类器层之外，编码器的所有参数都被冻结。从上表中，我们可以看到，HiViT-B 模型取得了 71.3% 的良好效果，与所有基于 MIM 的方法相比，这是最好的性能。

HiViT 只需要未掩蔽的 token 作为输入，这样本文的方法在 MIM 预训练期间就可以享受效率。如上表所示，作者报告了 MAE（ViT-B）、SimMIM（Swin-B）和本文的 HiViT-B 在不同输入大小下的预训练速度。所有结果代表 8×V100 GPU 上 1 个 epoch 的预训练时间（分钟）。由于输入图像为 192×192，HiViT-B 每 epoch 仅需 7.4 分钟，比 SimMIM 快约 1.9 倍，与MAE相当。当输入为 224×224 时，HiViT-B 大约需要 9.7 分钟，比 SimMIM 快 1.9× 与 MAE 相当。

上表展示了本文消融实验的结果，默认设置（#0）使用 2−2−20 块设置，可实现 83.8% 的性能。减少第 1、2 阶段的块数，增加第 3 阶段（#1）的块数，可以带来更多的参数和更好的性能，达到 83.9%，这与使用 Swin-B 进行 800 个 epoch 的预训练得出的 SimMIM 结果（84.0%）相当，但 71.9M 的参数仍然远低于 Swin-B 的 87.8M。

删除第 1 阶段（#3）或第 2 阶段（#2）都会对性能造成损害，即 83.6% 和 83.7%，这表明，主阶段之前的层次结构非常重要，并带来了性能改进。此外，#3 的结果低于 #2，表明第一阶段似乎比第二阶段更重要。去除前两个阶段的结果为 83.6%，这进一步验证了层次结构的重要性。

作者将相同的 CAE 设置传输到 MS-COCO 中测试本文的模型，选择第 5、9、13、19 个块作为后续 FPN 网络的输入。如上表所示，本文的方法达到了 SOTA 的性能。在 ADE20K 数据集上，MoCo-v3 报告了 300 个 epoch 的预训练 47.3% 的 mIoU 结果，低于本文的 48.3%。BEiT、CAE 和 MAE 报告的性能分别为 47.1%、48.8% 和 48.1%。通过 1600 个 epoch 的预训，MAE 达到了 48.1% 的 mIoU。与这些最先进的方法相比，HiViT-B 通过 800 个 epoch 的预训练，达到了 48.3% 的结果，这高于除 CAE 之外的所有方法，但是 CAE 使用 DALLE 的 tokenizer。

总结

本文提出了一种层次视觉 Transformer HiViT。从 Swin Transformers 开始，作者移除了在主阶段跨越 token 边界的冗余操作，并表明此类修改不会造成伤害，而会略微提高模型在全监督和自监督视觉表示学习中的性能。HiViT 在与掩蔽图像建模相结合方面显示出明显的优势，可以直接移植 ViT 上的高效实现，将训练速度提高近一倍。

尽管在实验中观察到了改进，但本文的方法仍有一些局限性。最重要的一点在于掩蔽单元的大小是固定的，这意味着需要选择一个“主阶段”。第三阶段的 Swin Transformer 提供了大多数参数和计算，因此自然会选择它，然而，在不存在主要阶段的情况下，该方法可能会遇到困难。

更多阅读