[论文解读] Rethinking Spatial Dimensions of Vision Transformers

Byeongho Heo, Sangdoo Yun+4 more|arXiv (Cornell University)|Mar 30, 2021

Advanced Neural Network Applications参考文献 39被引用 67

一句话总结

PiT 引入基于池化的池化层到 ViT，为视觉变换器创建类似 ResNet 的空间维度缩减，在 ImageNet、检测和鲁棒性基准测试中，相较于 ViT 提供了更好的准确性和泛化能力。

ABSTRACT

Vision Transformer (ViT) extends the application range of transformers from language processing to computer vision tasks as being an alternative architecture against the existing convolutional neural networks (CNN). Since the transformer-based architecture has been innovative for computer vision modeling, the design convention towards an effective architecture has been less studied yet. From the successful design principles of CNN, we investigate the role of spatial dimension conversion and its effectiveness on transformer-based architecture. We particularly attend to the dimension reduction principle of CNNs; as the depth increases, a conventional CNN increases channel dimension and decreases spatial dimensions. We empirically show that such a spatial dimension reduction is beneficial to a transformer architecture as well, and propose a novel Pooling-based Vision Transformer (PiT) upon the original ViT model. We show that PiT achieves the improved model capability and generalization performance against ViT. Throughout the extensive experiments, we further show PiT outperforms the baseline on several tasks such as image classification, object detection, and robustness evaluation. Source codes and ImageNet models are available at https://github.com/naver-ai/pit

研究动机与目标

• 研究 CNN 风格的空间维度缩减（下采样）是否有益于基于 Transformer 的视觉模型。
• 用基于池化的机制扩展 ViT，以在不牺牲基于 token 的处理的前提下实现空间下采样。
• 在图像分类、目标检测和鲁棒性基准测试中评估 PiT，以评估泛化性和效率提升。

提出的方法

• 为 ViT 设计一个池化层，将 2D token 映射重塑为 3D 张量，应用基于深度卷积的空间降维和通道扩展，再重塑回 2D tokens。
• 在 PiT 中嵌入两个池化层，以创建三个空间尺度，通道扩展通过增加多头注意力头实现。
• 构建与 ViT（tiny、small、base）相对应尺度的 PiT 变体，并调整 FLOPs/参数以在降低延迟的同时与 ViT 具有竞争力。
• 在 ImageNet-1k 上，在 vanilla、CutMix、DeiT 和 Distillation 训练方案下对 PiT 与 ViT 进行比较。
• 通过熵和空间距离分析注意力模式，以理解 PiT 如何在各层重塑 token 之间的交互。
• 评估 PiT 作为 Deformable DETR 的骨干网络在 COCO 上的表现，并评估在 ImageNet 基准测试中的鲁棒性。

Figure 1 : Schematic illustration of dimension configurations of networks. We visualize ResNet50 [ 13 ] , Vision Transformer (ViT) [ 9 ] , and our Pooling-based Vision Transformer (PiT); (a) ResNet50 gradually downsamples the features from the input to the output; (b) ViT does not change the spatial

实验结果

研究问题

• RQ1 将 CNN 风格的空间维度缩减应用于 Vision Transformers 能否提升模型能力和泛化能力？
• RQ2 基于池化的 ViT（PiT）是否在图像分类、目标检测和鲁棒性基准测试中优于 ViT？
• RQ3 PiT 的池化对注意力模式相对于 ViT 的影响是什么，这与性能和鲁棒性有何关系？
• RQ4 在准确性、吞吐量和延迟方面，PiT 在不同尺度上是否具备与卷积骨干网络的竞争力？
• RQ5 在扩展训练方案（长时间训练、高分辨率微调）下，PiT 的表现如何？

主要发现

• PiT 在多个 ImageNet 训练方案中普遍优于 ViT，在更低 FLOPs 下获得更高的准确性，且吞吐量往往更好。
• 在 ImageNet 上，PiT-B 在相似设置下达到 84.0%，而 ViT-B 为 83.4%；PiT-S 在若干配置中也显示出对 ViT-S 的提升。
• 随着模型规模增大，PiT 的性能提升，缓解了 ViT 在更大尺度和常规训练技巧下出现的泛化差距。
• 作为目标检测（COCO 上的 Deformable DETR）的骨干网络，PiT 的 AP 高于 ViT-S（39.4 对 36.9），且延迟具有竞争力。
• 在鲁棒性基准测试中，PiT-S 相对于 ViT-S 在标准、Occ、IN-A、BGC 和 FGSM 测试中显示出持续改进。
• 总体而言，PiT 证明了在 ViT 中引入基于池化的空间降维可以提升性能和鲁棒性，同时实现更有利的效率权衡。

(a) Model capability