WaveMix: Resource-efficient Token Mixing for Images
This addresses the need for practical, efficient neural architectures for green, edge, or desktop computing, offering an incremental improvement over existing methods.
The paper tackles the problem of high computational requirements in vision models like ViTs and CNNs by introducing WaveMix, a resource-efficient architecture using multi-scale 2D discrete wavelet transform for token mixing, achieving competitive generalization on several datasets with lower GPU RAM, computations, and storage, including SOTA results on EMNIST Byclass and Balanced datasets.
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.