Source-linked AI summary
Wave-ViT: Unifying Wavelet and Transformers for Visual Representation Learning
Ting Yao, Yingwei Pan, Yehao Li, Chong-Wah Ngo, Tao Mei
TL;DR
Multi-scale self-attention is computationally quadratic, and conventional key/value pooling can discard high-frequency details. Wave-ViT unifies self-attention with invertible wavelet down-sampling and IDWT-based local-context aggregation, achieving stronger results across vision tasks, including 85.5% ImageNet Top-1 accuracy versus PVT's 83.8%.
Problem
Multi-scale Vision Transformers face quadratic self-attention cost, while irreversible pooling-based key/value reduction drops high-frequency object details.
Method
Wave-ViT uses DWT for lossless invertible key/value down-sampling within self-attention and IDWT to strengthen outputs with local contexts.
Results
Wave-ViT surpasses PVT across ImageNet recognition and COCO detection and instance segmentation, reaching 85.5% ImageNet Top-1 versus 83.8%.
Takeaways & Limitations
Wave-ViT offers a multi-scale Transformer design with lossless down-sampling and a better computation-cost versus performance trade-off.
Abstract
from arXiv · showhide
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 (\textbf{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}.
1 JD Explore Academy, China
The passage lists the paper's core keywords: Vision Transformer, Wavelet Transform, self-attention learning, and image recognition.
- The paper concerns Vision Transformers for visual representation learning.
- Wavelet Transform is identified as a central technical concept.
- Self-attention learning and image recognition define additional focus areas.
1 Introduction
The introduction motivates Wave-ViT by the quadratic cost and information loss of conventional key/value down-sampling. It proposes wavelet-based invertible down-sampling and reports stronger results across recognition and dense prediction tasks.
- 1 Introduction: Self-attention cost grows quadratically with input patch number, making high-resolution multi-scale processing expensive.The stated cost is O(H^2W^2D).
- 1 Introduction: Pooling-based key/value down-sampling reduces computation but irreversibly drops high-frequency texture information.The introduction notes that this loss can especially affect dense prediction tasks.
- 1 Introduction: Wavelets block uses DWT to split keys/values into four lower-resolution frequency subbands while preserving image details.The low-frequency subband captures coarse structure, while three high-frequency subbands retain texture details.
- 1 Introduction: The block applies convolution to the stacked subbands, performs self-attention on down-sampled keys/values, and uses IDWT to reconstruct local contextual features.Attended and reconstructed feature maps are combined as the block output.
- 1 Introduction: Wave-ViT reaches 85.5% ImageNet Top-1 accuracy, improving PVT's 83.8% by 1.7% with comparable parameters.On COCO, it also exceeds PVT by 1.3% and 0.5% mAP for object detection and instance segmentation, respectively, with 25.9% fewer parameters.
2 Related Work
Related work spans CNNs, vision Transformers, multi-scale ViTs, and wavelet-based CNN methods. Wave-ViT differs by applying invertible wavelet down-sampling within Transformer self-attention.
- 2 Related Work: CNN backbones traditionally build low-resolution, high-level representations, whereas dense prediction requires high-resolution or pixel-level features.
- 2 Related Work: Vision research has combined convolution with global or local attention, while ViT and DETR use pure Transformers for image recognition and object detection.
- 2 Related Work: Multi-scale ViTs use pyramid-like representations to address the diverse scales of visual patterns in natural scenes.
- 2 Related Work: Existing multi-scale ViTs commonly reduce spatial dimensions with irreversible average pooling or pooling kernels.Wave-ViT instead uses invertible wavelet down-sampling over keys/values for a better computation–performance trade-off.
- 2 Related Work: Although wavelet transforms have been used in CNN down-sampling and up-sampling, this work applies them within Transformer blocks to down-sample keys and values.
3 Our Approach: Wavelet Vision Transformer
Wave-ViT replaces irreversible key/value pooling with a unified Wavelets block that combines wavelet-based invertible down-sampling, self-attention, and inverse reconstruction. This design preserves frequency details while reducing attention cost and adds local contextualization through reconstructed features.
- Multi-head Self-attention in ViT Backbones: Multi-head self-attention costs O(H^2W^2D), motivating spatial reduction for high-resolution feature maps.
- Wavelets Block: Wavelets blocks use DWT to decompose inputs into four lower-resolution subbands, preserving coarse structure and fine texture details without information dropping.Haar wavelets produce one low-frequency subband and three high-frequency subbands.
- Wavelets Block: The four subbands are concatenated, locally processed by a 3×3 convolution, and linearly transformed into down-sampled keys and values for attention.The resulting key/value sequence has reduced spatial resolution while retaining the wavelet components.
- Wavelets Block: Wavelet-based multi-head attention operates on input queries and down-sampled keys/values, producing long-range contextualized information for each head.
- Wavelets Block: IDWT reconstructs locally contextualized features that retain input details, while the DWT-Convolution-IDWT path enlarges receptive fields with negligible computational and memory increase.
- Wavelet Vision Transformer: Wave-ViT provides small, base, and large variants with four progressively down-sampled stages, matching comparable Swin model sizes and computational complexity.The variants are Wave-ViT-S, Wave-ViT-B, and Wave-ViT-L.
4 Experiments
Wave-ViT is evaluated across image recognition, detection, instance segmentation, and semantic segmentation, consistently outperforming comparable vision backbones. Ablations attribute these gains to invertible wavelet down-sampling and inverse wavelet reconstruction.
- Image Recognition: Image recognition uses ImageNet1K training from scratch with top-1 and top-5 validation accuracy as evaluation metrics.The benchmark contains 1.28 million training images, 50K validation images, and 1,000 classes.
- Image Recognition: Wave-ViT variants consistently outperform CNN, single-scale ViT, and multi-scale ViT backbones under similar GFLOPs for image recognition.Wave-ViT-B⋆ reaches 84.8% Top-1 accuracy, improving 0.6% over VOLO-D1⋆ at 84.2%.
- Object Detection: Wave-ViT-S achieves consistent gains over ResNet50, Swin-T, and PVTv2-B2 across GFL, Sparse RCNN, Cascade Mask R-CNN, and ATSS detectors.The evaluation tests whether pretrained multi-scale features generalize across four state-of-the-art object detectors.
- Semantic Segmentation: On ADE20K semantic segmentation, Wave-ViT-S/B achieves 49.6%/51.5% mIoU, improving 2.5%/2.0% over the best competitors.The comparisons use UPerNet and group pretrained backbones by Small and Base size.
- Ablation Study: Ablations show DWT-based invertible down-sampling raises Top-1 accuracy from 82.0% to 82.5%, while adding IDWT further boosts Top-1 and Top-5 with negligible cost increases.Average pooling and pooling-kernel variants obtain the same performance, while pooling kernels increase parameters.
- Visualization of Learnt Visual Representation: Score-CAM visualizations show Wave-ViT-S concentrates more strongly on semantically relevant objects than PVTv2-B2.The authors use six ImageNet1K images and interpret the concentration as evidence of more robust representations.
5 Conclusions
Wave-ViT unifies Transformer modules with invertible wavelet down-sampling to enable efficient, lossless multi-scale self-attention. Extensive experiments report superiority over state-of-the-art multi-scale ViT backbones across recognition and dense prediction tasks.
- Wavelets blocks use DWT for invertible down-sampling over keys and values, while IDWT reconstructs outputs to aggregate local contexts with an enlarged receptive field.The design unifies self-attention with lossless down-sampling and local-context enhancement.
- Stacked Wavelets blocks form four-stage Wave-ViT backbones with different model sizes and light computational and memory budgets.The architecture operates over multi-scale features within a multi-stage Transformer framework.
- Wave-ViT is empirically validated against state-of-the-art multi-scale ViT backbones for image recognition under comparable parameter counts.
- The models are also evaluated across vision tasks including object detection and instance segmentation.The conclusion describes validation across both recognition and dense prediction settings.