# WaveMix: Resource-efficient Token Mixing for Images

arXiv: https://arxiv.org/abs/2203.03689 · HTML: https://mlsanity.com/paper/2203.03689/

## Key facts

- Cited by 9 papers in the map (data through 14 Oct 2026).
- Builds on 23 of its 47 references that are in the map.
- First posted on arXiv 7 Mar 2022.

## Builds on (most cited first, 23 of 23)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [U-Net: Convolutional Networks for Biomedical Image Segmentation](https://mlsanity.com/paper/1505.04597/) (2015)
- [mixup: Beyond Empirical Risk Minimization](https://mlsanity.com/paper/1710.09412/) (2017)
- [You Only Look Once: Unified, Real-Time Object Detection](https://mlsanity.com/paper/1506.02640/) (2015)
- [Fashion-MNIST: a Novel Image Dataset for Benchmarking Machine Learning Algorithms](https://mlsanity.com/paper/1708.07747/) (2017)
- [CutMix: Regularization Strategy to Train Strong Classifiers with Localizable Features](https://mlsanity.com/paper/1905.04899/) (2019)
- [RandAugment: Practical automated data augmentation with a reduced search space](https://mlsanity.com/paper/1909.13719/) (2019)
- [Rethinking Attention with Performers](https://mlsanity.com/paper/2009.14794/) (2020)
- [MLP-Mixer: An all-MLP Architecture for Vision](https://mlsanity.com/paper/2105.01601/) (2021)
- [Random Erasing Data Augmentation](https://mlsanity.com/paper/1708.04896/) (2017)
- [CvT: Introducing Convolutions to Vision Transformers](https://mlsanity.com/paper/2103.15808/) (2021)
- [Efficient Transformers: A Survey](https://mlsanity.com/paper/2009.06732/) (2020)
- [Nystr"omformer: A Nystr"om-Based Algorithm for Approximating Self-Attention](https://mlsanity.com/paper/2102.03902/) (2021)
- [FNet: Mixing Tokens with Fourier Transforms](https://mlsanity.com/paper/2105.03824/) (2021)
- [Why gradient clipping accelerates training: A theoretical justification for adaptivity](https://mlsanity.com/paper/1905.11881/) (2019)
- [Escaping the Big Data Paradigm with Compact Transformers](https://mlsanity.com/paper/2104.05704/) (2021)
- [Patches Are All You Need?](https://mlsanity.com/paper/2201.09792/) (2022)
- [Multi-level Wavelet-CNN for Image Restoration](https://mlsanity.com/paper/1805.07071/) (2018)
- [Benchmarking Representation Learning for Natural World Image Collections](https://mlsanity.com/paper/2103.16483/) (2021)
- [A Closer Look at Deep Learning Heuristics: Learning rate restarts, Warmup and Distillation](https://mlsanity.com/paper/1810.13243/) (2018)
- [Wavelet Integrated CNNs for Noise-Robust Image Classification](https://mlsanity.com/paper/2005.03337/) (2020)
- [SpinalNet: Deep Neural Network with Gradual Input](https://mlsanity.com/paper/2007.03347/) (2020)

## Built on by (most cited first, 9 of 9)

- [Scattering Vision Transformer: Spectral Mixing Matters](https://mlsanity.com/paper/2311.01310/) (2023)
- [An Evolutionary Approach to Dynamic Introduction of Tasks in Large-scale Multitask Learning Systems](https://mlsanity.com/paper/2205.12755/) (2022)
- [WavSpA: Wavelet Space Attention for Boosting Transformers' Long Sequence Learning Ability](https://mlsanity.com/paper/2210.01989/) (2022)
- [Cross-Domain Knowledge Distillation for Low-Resolution Human Pose Estimation](https://mlsanity.com/paper/2405.11448/) (2024)
- [Focus on Low-Resolution Information: Multi-Granular Information-Lossless Model for Low-Resolution Human Pose Estimation](https://mlsanity.com/paper/2405.12247/) (2024)
- [Multi-Task Neural Architecture Search Using Architecture Embedding and Transfer Rank](https://mlsanity.com/paper/2504.00772/) (2025)
- [Edge-Native Digitization of Handwritten Marksheets: A Hybrid Heuristic-Deep Learning Framework](https://mlsanity.com/paper/2508.16295/) (2025)
- [Is Haar Enough? Exploring Symlets and Coiflets for Wavelet Convolution Layers](https://mlsanity.com/paper/2608.17662/) (2026)
- [WNet: Discrete Wavelets Transform for Efficient Token Mixing](https://mlsanity.com/paper/2610.04720/) (2026)

## Abstract (arXiv)

Although certain vision transformer (ViT) and CNN architectures generalize well on vision tasks, it is often impractical to use them on green, edge, or desktop computing due to their computational requirements for training and even testing. We present WaveMix as an alternative neural architecture that uses a multi-scale 2D discrete wavelet transform (DWT) for spatial token mixing. Unlike ViTs, WaveMix neither unrolls the image nor requires self-attention of quadratic complexity. Additionally, DWT introduces another inductive bias -- besides convolutional filtering -- to utilize the 2D structure of an image to improve generalization. The multi-scale nature of the DWT also reduces the requirement for a deeper architecture compared to the CNNs, as the latter relies on pooling for partial spatial mixing. WaveMix models show generalization that is competitive with ViTs, CNNs, and token mixers on several datasets while requiring lower GPU RAM (training and testing), number of computations, and storage. WaveMix have achieved State-of-the-art (SOTA) results in EMNIST Byclass and EMNIST Balanced datasets.