[论文解读] HiViT: Hierarchical Vision Transformer Meets Masked Image Modeling

Xiaosong Zhang, Yunjie Tian+5 more|arXiv (Cornell University)|May 30, 2022

Advanced Neural Network Applications被引用 12

一句话总结

HiViT 提出了一种分层视觉变换器架构，通过移除局部单元间操作（如移位窗口自注意力机制）来实现高效掩码图像建模（MIM），同时保留多尺度特征图，使令牌序列化类似于普通 ViT。该设计在 MAE 预训练下，相比 ViT-B 提升了 0.6% 的 ImageNet-1K 图像分类 top-1 准确率，且训练速度比 Swin-B 快 1.9 倍。

ABSTRACT

Recently, masked image modeling (MIM) has offered a new methodology of self-supervised pre-training of vision transformers. A key idea of efficient implementation is to discard the masked image patches (or tokens) throughout the target network (encoder), which requires the encoder to be a plain vision transformer (e.g., ViT), albeit hierarchical vision transformers (e.g., Swin Transformer) have potentially better properties in formulating vision inputs. In this paper, we offer a new design of hierarchical vision transformers named HiViT (short for Hierarchical ViT) that enjoys both high efficiency and good performance in MIM. The key is to remove the unnecessary "local inter-unit operations", deriving structurally simple hierarchical vision transformers in which mask-units can be serialized like plain vision transformers. For this purpose, we start with Swin Transformer and (i) set the masking unit size to be the token size in the main stage of Swin Transformer, (ii) switch off inter-unit self-attentions before the main stage, and (iii) eliminate all operations after the main stage. Empirical studies demonstrate the advantageous performance of HiViT in terms of fully-supervised, self-supervised, and transfer learning. In particular, in running MAE on ImageNet-1K, HiViT-B reports a +0.6% accuracy gain over ViT-B and a 1.9$ imes$ speed-up over Swin-B, and the performance gain generalizes to downstream tasks of detection and segmentation. Code will be made publicly available.

研究动机与目标

• 解决由于空间依赖的局部操作导致的分层视觉变换器在掩码图像建模（MIM）中效率低下的问题。
• 通过消除阻碍 MIM 效率的‘局部单元间操作’，实现在分层变换器中的令牌序列化。
• 在保留分层特征表示性能优势的同时，实现类似 ViT 的计算灵活性。
• 在自监督、全监督以及下游目标检测/分割任务中展示最先进性能。
• 为普通 ViT 和复杂分层模型（如 Swin Transformer）提供一种简单、高效且适合 MIM 的替代方案。

提出的方法

• 将分层视觉变换器中的操作分类为单元内操作、全局单元间操作和局部单元间操作。
• 识别出‘局部单元间操作’（如移位窗口自注意力机制和补丁合并）是阻碍令牌序列化的主因。
• 在 Swin Transformer 的主阶段中完全移除所有局部单元间操作，仅保留全局自注意力机制和分层特征图。
• 将早期阶段的局部操作替换为等效的单元内 MLP，以维持 FLOPs 和模型容量。
• 将最后一级与主阶段合并，以在简化架构的同时保持 FLOPs。
• 通过丢弃被掩码的单元并仅处理可见单元，实现在 MIM 预训练过程中对可见令牌的完全序列化。

实验结果

研究问题

• RQ1 能否在不损失性能的前提下，使分层视觉变换器在掩码图像建模中的效率与普通 ViT 相当？
• RQ2 分层变换器中的‘局部单元间操作’是否对识别准确率有显著贡献，还是分层结构本身才是关键？
• RQ3 对 Swin Transformer 进行最小修改——即移除局部注意力——是否能实现比 ViT 和 Swin 更优的 MIM 效率与准确率？
• RQ4 所提出的架构 HiViT 是否在自监督、全监督以及下游目标检测/分割任务中均保持强大性能？
• RQ5 能否在不依赖空间感知注意力机制的前提下，有效利用分层结构进行 MIM？

主要发现

• HiViT-B 在 ImageNet-1K 上达到 83.8% 的 top-1 准确率，相比 ViT-B 提升 0.6%，相比 Swin-B 提升 0.3%。
• 在 MAE 预训练下，HiViT-B 的训练速度比 Swin-B 快 1.9 倍，同时准确率超过 ViT-B。
• HiViT-B 的自监督预训练在 ImageNet-1K 上达到 84.2% 的 top-1 准确率，与监督基线具有竞争力。
• 通过改进特征金字塔的使用，HiViT-B 在 COCO 目标检测和实例分割任务上分别达到 51.2% AP^box 和 44.2% AP^mask。
• 在 ADE20K 语义分割任务上，HiViT-B 达到 51.2% 的 mIoU，优于相同预训练协议下的 ViT-B。
• 消融实验表明，移除局部单元间操作对性能影响极小，验证了核心设计选择的合理性。