Wave-ViT: Unifying wavelet and transformers for visual representation learning

• Main Authors: YAO, Ting, PAN, Yingwei, LI, Yehao, NGO, Chong-wah, MEI, Tao
• Format: text
• Language: English
• Published: Institutional Knowledge at Singapore Management University 2022
• Subjects: Vision transformer; Wavelet transform; Self-attention; learning; Image recognition; Artificial Intelligence and Robotics; Graphics and Human Computer Interfaces
• Online Access: https://ink.library.smu.edu.sg/sis_research/7508; https://ink.library.smu.edu.sg/context/sis_research/article/8511/viewcontent/2207.04978.pdf
• Institution: Singapore Management University

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 (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 https://github.com/YehLi/ImageNetModel.