PaperlayerView exact revision on arXiv

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

arXiv:2207.04978v1cs.CVcs.LG

TL;DR

Multi-scale ViT self-attention faces quadratic computational costs, while pooling keys and values can discard information. Wave-ViT uses wavelet transforms for invertible down-sampling in self-attention, and experiments report that it surpasses state-of-the-art ViT backbones across multiple vision tasks.

  • Problem

    Quadratic self-attention costs motivate down-sampling keys and values, but pooling-based operations are irreversible and drop information, including high-frequency texture details.

    Source: §1 · p. 2

  • Method

    Wave-ViT integrates wavelet-based invertible down-sampling of keys and values into self-attention and uses inverse wavelet transforms to aggregate local contexts.

    Source: §3.2 · p. 6, §5 · p. 14

  • Results

    Across image recognition, object detection, and instance segmentation, Wave-ViT surpasses state-of-the-art ViT backbones; it reaches 85.5% top-1 accuracy on ImageNet.

    Source: §1 · p. 3

  • Takeaways & Limitations

    For multi-scale ViTs, wavelet transforms offer a better trade-off between computation cost and performance than irreversible pooling-based down-sampling.

    Source: §2.1 · p. 4

Abstract

from arXiv · show

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 Introduction

Multi-scale ViTs reduce the quadratic cost of self-attention by down-sampling keys and values, but pooling can discard high-frequency detail. Wave-ViT uses wavelet transforms for invertible down-sampling and reports gains over PVT on ImageNet and COCO.

1 Introduction

  • Because ViT’s single-scale, low-resolution features are not straightforward to use for dense prediction, prior work introduced multi-scale designs such as PVT and MViT.Source: §1 · p. 2
  • Pooling keys and values reduces quadratic self-attention cost but drops information, particularly high-frequency object texture details, which can hurt dense prediction performance.The introduction also notes that pooling in CNNs can hurt shift-equivariance.Source: §1 · p. 2
  • DWT downsamples inputs invertibly into four lower-resolution frequency subbands, retaining coarse object structure and fine texture details for lossless self-attention learning.IDWT can reconstruct the original image from the subbands.Source: §1 · pp. 2–3
  • Wavelets blocks apply DWT and a 3×3 convolution to keys and values before attention, then combine attended features with IDWT-reconstructed features for stronger local context.The DWT–convolution–IDWT process enlarges the receptive field with negligible increases in computation and memory compared with a single 3×3 convolution.Source: §1 · p. 3
  • On ImageNet, Wave-ViT reaches 85.5% top-1 accuracy versus PVT’s 83.8% (+1.7%); on COCO, it surpasses PVT by 1.3% detection and 0.5% instance-segmentation mAP with 25.9% fewer parameters.Source: §1 · p. 3

2 Related Work

Prior work developed multi-scale CNN and Transformer backbones and used wavelets in CNNs; Wave-ViT brings invertible wavelet down-sampling to Transformer keys and values. This contrasts with multi-scale ViTs that reduce spatial dimensions through pooling.

2.1 Visual Representation Learning

  • CNN backbones evolved from deep, low-resolution representations toward multi-scale designs such as Res2Net and HRNet that preserve high-resolution features.These designs address dense prediction tasks that need high-resolution or pixel-level representations.Source: §2.1 · p. 3
  • Vision research progressed from hybrid and pure Transformers to multi-scale backbones including Swin, PVT, PVTv2 [54], and MViT.These architectures introduce hierarchical or multi-scale features, with pooling used for spatial reduction in PVTv2 and MViT.Source: §2.1 · pp. 3–4
  • Wave-ViT replaces irreversible pooling in multi-scale ViTs with invertible wavelet down-sampling of keys and values for self-attention over multi-scale features.The paper presents this as a way to trade off computation cost and performance.Source: §2.1 · p. 4

2.2 Wavelet Transform in Computer Vision

  • Wavelets have supported CNN vision tasks including restoration, super-resolution, recognition, texture classification and annotation [15], and parameter-efficient recognition.Prior work used wavelet transforms for information-preserving receptive-field expansion as well as feature decomposition and dimensionality reduction.Source: §2.2 · p. 4
  • Unlike prior CNN uses of wavelets, Wave-ViT applies the transform within Transformer blocks to down-sample keys and values, extending its role in feature learning.The paper states that wavelet down-sampling had not previously been explored for Transformer architectures.Source: §2.2 · p. 5

3 Our Approach: Wavelet Vision Transformer

Wave-ViT combines wavelet-based invertible down-sampling of attention keys and values with inverse-wavelet local-context enhancement. Its Wavelets blocks form a four-stage multi-scale vision Transformer backbone.

Multi-head Self-attention in ViT Backbones.

  • Standard multi-head self-attention projects inputs into queries, keys, and values, then computes attention per head; its O(H^2W^2D) cost grows quadratically with patch count.The outputs of the heads are concatenated and linearly transformed.Source: Multi-head Self-attention in ViT… · p. 6

Self-attention with Down-sampling in Multi-scale ViT Backbones.

3.2 Wavelets Block

  • The Wavelets block uses classical Haar DWT to produce four down-sampled subbands that retain all input details—one encoding coarse object structure and three preserving fine texture—for lossless self-attention down-sampling.Source: §3.2 · pp. 6–7
  • An inverse wavelet transform reconstructs locally contextualized features that retain input details and provide a larger receptive field than a single 3×3 convolution at negligible extra cost.The block concatenates this reconstructed local information with long-range self-attention outputs.Source: §3.2 · p. 7

3.3 Wavelet Vision Transformer

  • Wave-ViT-S, -B, and -L use four stages of patch embedding, Wavelets blocks, and feed-forward layers, with increasing channel dimensions and shrinking spatial resolutions.Their model sizes and computational complexities are similar to Swin-T/S/B, respectively.Source: §3.3 · p. 8

4 Experiments

Experiments evaluate Wave-ViT on image recognition, object detection, instance and semantic segmentation, and block ablations. Across these tasks, the results support wavelet-based invertible down-sampling and the full Wavelets block.

4.1 Image Recognition on ImageNet1K

  • ImageNet1K recognition is evaluated from scratch using top-1 and top-5 validation accuracy after 300 training epochs following 10 warm-up epochs.Training uses 1.28 million training images, 50K validation images, AdamW, batch size 1,024 across 8 V100 GPUs, and RandAug, CutOut, and Token Labeling with MixToken; upgraded Wave-ViT variants use Token Labeling and conv-stem for comparison with upgraded VOLO variants.Source: §4.1 · p. 9

Performance Comparison.

  • Wave-ViT-B⋆ reaches 84.8% Top-1 accuracy versus 84.2% for VOLO-D1⋆, a 0.6% absolute improvement under comparable Base-size compute.Without Token Labeling and conv-stem, Wave-ViT-S also outperforms the best Small-size multi-scale ViT, PVTv2-B2. Wave-ViT-L⋆ retains 56×56 inputs while VOLO-D3⋆ reduces them to 28×28, and the paper reports a better efficiency-versus-accuracy trade-off for Wave-ViT-L⋆.Source: Table 2 · p. 9, §4.1 · p. 9, Performance Comparison · p. 10

4.2 Object Detection and Instance Segmentation on COCO

  • Wave-ViT-S gains consistently over ResNet50, Swin-T, and PVTv2-B2 across GFL, Sparse RCNN, Cascade Mask R-CNN, and ATSS object detectors.The COCO evaluation also uses RetinaNet for detection and Mask R-CNN for instance segmentation, but the supplied passages do not state their comparative results.Source: §4.2 · p. 10, Dataset and Optimization Setups · pp. 11–12, Table 4 · p. 11

4.3 Semantic Segmentation on ADE20K

  • 49.6%/51.5% mIoU for Wave-ViT-S/B on ADE20K exceeds the respective best competitors, Twins-SVT-S/Swin-S, at 47.1%/49.5%.The evaluation uses UPerNet and follows Swin’s hyperparameters and detection heads; the reported gains apply to comparable model sizes.Source: §4.3 · pp. 12–13, Table 5 · p. 12

4.4 Ablation Study

  • The Wavelets-block ablation raises Top-1 from 82.0% with average-pooling down-sampling to 82.5% with DWT-based invertible down-sampling.The full block adds IDWT to strengthen outputs with an enlarged receptive field, boosting Top-1 and Top-5 with negligible computational-cost and memory increases; pooling-kernel down-sampling matches average pooling but increases parameters.Source: §4.4 · pp. 13–14, Figure 3 · p. 13

4.5 Visualization of Learnt Visual Representation

  • Score-CAM visualizations show Wave-ViT-S concentrating saliency more on semantically relevant objects than similarly sized PVTv2-B2, which the authors interpret as more robust representations.The comparison visualizes six ImageNet1K images.Source: §4.5 · p. 14, Figure 4 · p. 14

5 Conclusions

Wave-ViT unifies Transformer modules with invertible wavelet down-sampling for efficient multi-scale self-attention, and the paper reports strong results on image recognition and downstream vision tasks.

5 Conclusions

  • Wave-ViT outperforms state-of-the-art multi-scale ViT backbones for image recognition at comparable parameter counts and generalizes to detection, instance and semantic segmentation.The conclusion reports these task-level findings without supplying metric values.Source: §5 · p. 14
  • Wavelets blocks use DWT for invertible, lossless down-sampling of self-attention keys and values, then IDWT to reconstruct outputs and aggregate local context over an enlarged receptive field.Source: §5 · p. 14
  • Stacked Wavelets blocks on multi-scale features form lightweight four-stage Wave-ViT backbones in different model sizes, with low computational and memory costs.Source: §5 · p. 14

Related papers

Loading 2207.04978v1…