[论文解读] Wave-ViT: Unifying Wavelet and Transformers for Visual Representation Learning

Ting Yao, Yingwei Pan+3 more|arXiv (Cornell University)|Jul 11, 2022

Advanced Neural Network Applications被引用 8

一句话总结

Wave-ViT 提出了一种新颖的小波块（Wavelets block），将小波变换与视觉 Transformer 相结合，实现在自注意力机制中的可逆、信息保真的下采样。通过用离散小波变换（DWT）替代不可逆的池化操作来实现键/值的下采样，并利用逆向离散小波变换（IDWT）进行特征图的重建与增强，Wave-ViT 在图像识别、目标检测和分割任务中实现了优于 SOTA 模型的性能，且计算量（FLOPs）相当。

ABSTRACT

Multi-scale Vision Transformer (ViT) has emerged as a powerful backbone for computer vision tasks, while the self-attention computation in Transformer scales quadratically w.r.t. the input patch number. Thus, existing solutions commonly employ down-sampling operations (e.g., average pooling) over keys/values to dramatically reduce the computational cost. In this work, we argue that such over-aggressive down-sampling design is not invertible and inevitably causes information dropping especially for high-frequency components in objects (e.g., texture details). Motivated by the wavelet theory, we construct a new Wavelet Vision Transformer ( extbf{Wave-ViT}) that formulates the invertible down-sampling with wavelet transforms and self-attention learning in a unified way. This proposal enables self-attention learning with lossless down-sampling over keys/values, facilitating the pursuing of a better efficiency-vs-accuracy trade-off. Furthermore, inverse wavelet transforms are leveraged to strengthen self-attention outputs by aggregating local contexts with enlarged receptive field. We validate the superiority of Wave-ViT through extensive experiments over multiple vision tasks (e.g., image recognition, object detection and instance segmentation). Its performances surpass state-of-the-art ViT backbones with comparable FLOPs. Source code is available at \url{https://github.com/YehLi/ImageNetModel}.

研究动机与目标

• 为解决多尺度视觉 Transformer 中由不可逆下采样操作（如平均池化）引起的信息损失问题。
• 通过用可逆的小波下采样替代池化操作，在自注意力计算过程中保留高频细节（如纹理）。
• 通过逆向小波变换重建下采样特征，增强自注意力输出中的局部上下文建模能力。
• 设计一种可扩展、高效的 Transformer 块，在计算开销极小的前提下保持高精度。
• 在多个视觉任务（包括图像识别与密集预测）中验证所提出小波块的有效性。

提出的方法

• 引入小波块，对键和值应用离散小波变换（DWT），将其分解为四个子带（LL、LH、HL、HH），实现可逆下采样。
• 将 DWT 的四个子带拼接为单个下采样特征图，并应用 3×3 卷积以在频带间强制实现局部不变性。
• 在下采样的键/值与原始查询上执行多头自注意力，以保持计算效率。
• 对下采样的键/值应用逆向离散小波变换（IDWT），重建出高分辨率特征图，扩大有效感受野。
• 将注意力输出的特征图与 IDWT 重建的特征图进行融合，生成最终输出，增强局部上下文聚合能力。
• 将小波块集成到多阶段、多尺度的 Transformer 架构中，构建 Wave-ViT，作为视觉表征学习的新骨干网络。

实验结果

研究问题

• RQ1 可逆小波变换能否有效替代视觉 Transformer 中的不可逆池化操作，以保留高频图像细节？
• RQ2 在计算成本无显著增加的前提下，逆向小波变换的集成是否能改善自注意力输出中的局部上下文建模？
• RQ3 与基于标准池化的自注意力块相比，所提出的波小波块在视觉任务中的准确率与效率表现如何？
• RQ4 Wave-ViT 是否能在图像识别、目标检测与实例分割任务中实现 SOTA 性能，且计算量（FLOPs）与现有模型相当？
• RQ5 通过显著性图可视化，小波基下采样在多大程度上提升了特征表示的鲁棒性？

主要发现

• Wave-ViT-S 在 ADE20K 验证集上达到 51.5% 的 mIoU，较最佳对比模型（Swin-S）高出 2.0 个百分点。
• Wave-ViT-B 在 ImageNet-1K 上达到 82.7% 的 Top-1 准确率，优于使用不可逆池化的基线模型（82.0%），提升 0.7 个百分点。
• 消融实验表明，仅将平均池化替换为 DWT 即可使 Top-1 准确率从 82.0% 提升至 82.5%，证明了可逆下采样的优势。
• 在 DWT 基础块中加入 IDWT 后，性能进一步提升至 82.7% 的 Top-1 准确率，且计算量增加可忽略，证实了增强感受野的有效性。
• Score-CAM 可视化结果表明，Wave-ViT-S 学习到的表征比 PVTv2-B2 更具鲁棒性与语义聚焦性，注意力更集中于相关物体。
• Wave-ViT 在各类任务中泛化能力出色，在图像识别、目标检测与实例分割任务中均实现了 SOTA 性能，且计算量与现有模型相当。