[论文解读] VSA: Learning Varied-Size Window Attention in Vision Transformers

Qiming Zhang, Yufei Xu+2 more|arXiv (Cornell University)|Apr 18, 2022

Advanced Neural Network Applications被引用 4

一句话总结

本文提出可变尺寸窗口注意力（VSA），一种可学习机制，用自适应、数据驱动的窗口配置替代视觉变换器中固定的窗口大小。通过使用窗口回归模块为每个注意力头预测最优窗口大小和位置，VSA增强了长距离建模能力，提升了在不同物体尺度下的特征表示，并以可忽略的计算开销提升了性能——在 Swin-T 上实现 ImageNet 上 1.1% 的准确率提升，且在更大输入图像下最高可达 1.9% 的准确率增益。

ABSTRACT

Attention within windows has been widely explored in vision transformers to balance the performance, computation complexity, and memory footprint. However, current models adopt a hand-crafted fixed-size window design, which restricts their capacity of modeling long-term dependencies and adapting to objects of different sizes. To address this drawback, we propose extbf{V}aried- extbf{S}ize Window extbf{A}ttention (VSA) to learn adaptive window configurations from data. Specifically, based on the tokens within each default window, VSA employs a window regression module to predict the size and location of the target window, i.e., the attention area where the key and value tokens are sampled. By adopting VSA independently for each attention head, it can model long-term dependencies, capture rich context from diverse windows, and promote information exchange among overlapped windows. VSA is an easy-to-implement module that can replace the window attention in state-of-the-art representative models with minor modifications and negligible extra computational cost while improving their performance by a large margin, e.g., 1.1\% for Swin-T on ImageNet classification. In addition, the performance gain increases when using larger images for training and test. Experimental results on more downstream tasks, including object detection, instance segmentation, and semantic segmentation, further demonstrate the superiority of VSA over the vanilla window attention in dealing with objects of different sizes. The code will be released https://github.com/ViTAE-Transformer/ViTAE-VSA.

研究动机与目标

• 为解决视觉变换器中固定大小窗口注意力的局限性，该局限性限制了长距离依赖关系的建模并降低了对不同尺寸物体的适应能力。
• 使视觉变换器能够基于输入数据动态学习每个注意力头的最优窗口大小和位置。
• 在不显著增加计算成本或显存使用的情况下提升视觉任务的性能。
• 通过实现重叠自适应窗口之间的信息交换，减少对复杂机制（如移位窗口）的依赖。
• 在多种任务上展示泛化能力，包括分类、检测和分割，尤其在高分辨率输入下表现更优。

提出的方法

• VSA 引入一个窗口回归模块，基于默认固定大小窗口内的令牌预测目标注意力窗口的大小和空间位置。
• 随后，键和值令牌从预测的目标窗口中采样，而查询仍来自默认窗口，从而实现动态感受野。
• 每个注意力头独立应用 VSA，使不同头可采用多样的窗口配置以捕捉多尺度上下文。
• 该模块为即插即用设计，仅需极少架构改动即可替换 Swin-T 等模型中的标准窗口注意力，且计算量增加可忽略。
• 通过重叠窗口促进跨窗口特征交换，减少了对移位窗口机制的依赖。
• 窗口回归与模型其余部分端到端联合训练，使模型能够基于数据驱动方式适应物体尺度和空间布局。

实验结果

研究问题

• RQ1 可学习的窗口机制是否能超越固定大小窗口的限制，在视觉变换器中实现更优的长距离建模？
• RQ2 自适应窗口大小和位置预测是否能带来对不同尺度物体更优的特征表示？
• RQ3 VSA 是否能在不显著增加计算成本的前提下提升图像分类和密集预测任务的性能？
• RQ4 在更大输入分辨率下，VSA 表现如何，尤其是在固定窗口难以捕捉大物体时？
• RQ5 通过实现有效的跨窗口信息交换，VSA 是否能减少甚至消除对移位窗口机制的需求？

主要发现

• 在 Swin-T 上，VSA 在 ImageNet 上实现了 1.1% 的 top-1 准确率增益，当使用更大图像（如 480×480）进行训练和测试时，性能提升进一步增至 1.9%。
• 在下游任务（如目标检测、实例分割和语义分割）中，VSA 始终优于标准窗口注意力，尤其在大物体上表现更优。
• 采用 VSA 的模型在 t-SNE 可视化中表现出更优的特征聚类，表明其具有更具判别性和语义意义的表示。
• VSA 通过重叠窗口实现了有效的跨窗口信息交换，使得 Swin 模型中的移位窗口机制在 VSA 变体中变得冗余。
• 尽管由于 PyTorch 采样操作不够高效，导致速度下降 12–17%，但 VSA 仅使 GPU 显存使用增加 2%，表明其计算开销极低。
• 可视化结果证实，VSA 生成了多样化、上下文感知的窗口，其对物体的覆盖效果优于固定大小窗口，且在显著区域表现出更高的窗口大小和位置方差。