[Paper Review] 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 Applications8 citations

TL;DR

Wave-ViT proposes a novel Wavelets block that unifies wavelet transforms with Vision Transformers to enable invertible, information-preserving down-sampling in self-attention mechanisms. By replacing irreversible pooling with Discrete Wavelet Transform (DWT) for down-sampling keys/values and using Inverse DWT (IDWT) to reconstruct and enhance feature maps, Wave-ViT achieves superior performance on image recognition, object detection, and segmentation with comparable FLOPs to SOTA models.

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}.

Motivation & Objective

• To address the information loss caused by irreversible down-sampling operations (e.g., average pooling) in multi-scale Vision Transformers.
• To preserve high-frequency details (e.g., texture) during self-attention computation by replacing pooling with invertible wavelet-based down-sampling.
• To enhance local context modeling in self-attention outputs by reconstructing down-sampled features via inverse wavelet transforms.
• To design a scalable, efficient Transformer block that maintains high accuracy with minimal computational overhead.
• To validate the effectiveness of the proposed Wavelets block across multiple vision tasks, including image recognition and dense prediction.

Proposed method

• Introduces the Wavelets block, which applies Discrete Wavelet Transform (DWT) to decompose keys and values into four subbands (LL, LH, HL, HH) for invertible down-sampling.
• Stacks the four DWT subbands into a single down-sampled feature map and applies a 3×3 convolution to enforce local invariance across frequency subbands.
• Performs multi-head self-attention on the down-sampled keys/values and the original query to maintain computational efficiency.
• Applies Inverse DWT (IDWT) to the down-sampled keys/values to reconstruct a high-resolution feature map with enlarged effective receptive field.
• Fuses the attended feature map and the IDWT-reconstructed map to produce the final output, enhancing local context aggregation.
• Integrates the Wavelets block into a multi-stage, multi-scale Transformer architecture to build Wave-ViT, a new backbone for visual representation learning.

Experimental results

Research questions

• RQ1 Can invertible wavelet transforms effectively replace irreversible pooling operations in Vision Transformers to preserve high-frequency image details?
• RQ2 Does the integration of inverse wavelet transforms improve local context modeling in self-attention outputs without significant computational cost?
• RQ3 How does the proposed Wavelets block compare to standard pooling-based self-attention blocks in terms of accuracy and efficiency across vision tasks?
• RQ4 Can Wave-ViT achieve state-of-the-art performance on image recognition, object detection, and instance segmentation with comparable FLOPs to existing models?
• RQ5 To what extent does the use of wavelet-based down-sampling improve feature representation robustness, as shown by saliency map visualization?

Key findings

• Wave-ViT-S achieves 51.5% mIoU on ADE20K validation set, surpassing the best competitor (Swin-S) by 2.0 percentage points.
• Wave-ViT-B achieves 82.7% Top-1 accuracy on ImageNet-1K, outperforming the baseline with irreversible pooling (82.0%) by 0.7 percentage points.
• The ablation study shows that replacing average pooling with DWT alone improves Top-1 accuracy from 82.0% to 82.5%, demonstrating the benefit of invertible down-sampling.
• Adding IDWT to the DWT-based block further boosts performance to 82.7% Top-1 with negligible computational increase, confirming the effectiveness of enhanced receptive field.
• Score-CAM visualization confirms that Wave-ViT-S learns more robust and semantically focused representations than PVTv2-B2, with higher attention concentration on relevant objects.
• Wave-ViT generalizes well across tasks, achieving state-of-the-art performance on image recognition, object detection, and instance segmentation under comparable FLOPs.