WaveMix: Efficient Multi-scale Wavelet Networks

• WaveMix is a family of neural network architectures that leverages multi-scale discrete wavelet transforms for efficient token mixing, feature extraction, and scale invariance.
• It integrates convolutional projections with multi-level 2D-DWT and channel MLPs to decompose images into hierarchical feature maps while reducing memory and parameter requirements.
• Adaptable to tasks like image classification, super-resolution, and time series augmentation, WaveMix achieves state-of-the-art performance with significantly fewer parameters than conventional CNNs and ViTs.

WaveMix is a family of neural network architectures and data augmentation techniques distinguished by their use of multi-scale discrete wavelet transforms (DWT) for efficient token mixing, feature extraction, or data recombination. Initially introduced for image analysis as an alternative to vision transformers (ViTs) and convolutional neural networks (CNNs), WaveMix leverages the inherent properties of the wavelet transform to achieve competitive performance with fewer parameters, lower hardware requirements, and greater scale invariance. Extensions of the WaveMix paradigm encompass image super-resolution, magnification-invariant medical image analysis, time series augmentation, and resource-constrained scenarios.

1. Architectural Foundations of WaveMix

The canonical WaveMix image architecture is built on the following stages:

1. Initial Convolutional Projection: The input image (e.g., $X \in \mathbb{R}^{C\times H\times W}$) is first processed by a convolutional layer that increases channel dimensionality and injects image-specific inductive bias.
2. Multi-level 2D-DWT Token Mixing: Each “WaveMix block” applies the 2D discrete wavelet transform over multiple spatial scales. At each level, the DWT decomposes the input into four sub-bands (approximation and three directional details), reducing resolution by a factor of two per level and expanding the feature representation in a lossless, parameter-free manner.
3. Resolution Integration: Outputs at all wavelet levels (each with different spatial size) are upsampled—typically via transposed convolution (deconvolution) or, in WaveMixSR-V2, PixelShuffle—and concatenated along the channel axis.
4. Depth-wise Convolution and Channel MLP: The aggregated multi-scale feature map undergoes (a) a depth-wise convolution for spatial-channel interaction, and (b) an MLP sublayer (two 1×1 convolutions and GELU activation) to further mix channelwise information.
5. Residual Addition: The block output is summed with the block input, preserving gradient flow and facilitating stacking of multiple blocks.

A representative sequence of operations for a WaveMix block (as given in (Jeevan et al., 2022, Jeevan et al., 2023, Jeevan et al., 16 Sep 2024)) is:

$$\begin{align*} x_0 &= c(x_{\text{in}}, \xi) \ x &= [w_{aa}(x_0) \oplus w_{ad}(x_0) \oplus w_{da}(x_0) \oplus w_{dd}(x_0)] \ \tilde{x} &= b\big(t\big(m(x, \theta), \phi\big), \gamma\big) \ x_{\text{out}} &= \tilde{x} + x_{\text{in}} \end{align*}$$

where $c$ is convolution, $w$ denotes Haar wavelet filters, $\oplus$ is concatenation, $m$ is a channel MLP, $t$ and $b$ are upsampling (deconvolution or PixelShuffle) and batch-norm respectively.

This structure can be adapted for image classification (with global averaging and MLP head), dense prediction, super-resolution, or sequence modeling.

2. WaveMix vs. CNNs, ViTs, and Other Token Mixers

A central contribution of WaveMix is the replacement of conventional spatial mixing (CNN convolution, ViT self-attention, or MLP mixing) with linear, lossless 2D-DWT, yielding the following comparative properties (Jeevan et al., 2022, Jeevan et al., 2022, Jeevan et al., 2023):

Property	CNN	ViT	WaveMix
Token Mixing	Learnable, local	Quadratic self-attn	Multi-scale DWT (fixed)
Input Flattening	No	Yes	No
Scale Invariance	Weak (pooling)	Patch-based	Strong (wavelet-based)
Parameter Efficiency	Moderate to high	High	High
GPU RAM Usage	High (deep nets)	Very high (quadratic)	Low

WaveMix achieves competitive or state-of-the-art accuracy on multiple image recognition (CIFAR-10/100, Tiny ImageNet, EMNIST, Galaxy 10 DECals) and segmentation tasks (Cityscapes, Places-365), often surpassing model variants with 2–10× more parameters and requiring 5–21% of the GPU memory consumed by ViT- or large CNN-based baselines (Jeevan et al., 2022, Jeevan et al., 2022).

In medical image analysis with variable magnification (BreakHis dataset), WaveMix maintains accuracy above 87% for all inter-magnification splits and demonstrates greater magnification invariance than either CNNs or ViT variants (Jeevan et al., 2023).

3. Theoretical Inductive Biases and Multi-scale Properties

WaveMix incorporates two key forms of inductive bias:

• Convolutional Bias: The initial and in-block convolutions promote translational invariance and local feature extraction.
• Wavelet Multi-Scale Bias: The DWT's hierarchical decomposition captures both local detail (low-level DWT) and global structure (high-level DWT), allowing for rapid, lossless receptive field expansion (region covered increases as $2^L$ for $L$ levels) and efficient scale-invariant representation (Jeevan et al., 2022).

The use of detail sub-bands emphasizes sparse edge content, beneficial for natural images’ redundant and edge-sparse character. Shift invariance is achieved by the Haar wavelet basis, which preserves the nature of features under small spatial translations (Jeevan et al., 2022, Jeevan et al., 2022).

In data augmentation for time series, wavelet-based mixing preserves both temporal and local spectral properties, improving over Fourier-masked methods that disrupt temporal coherence (Arabi et al., 20 Aug 2024).

4. Resource Efficiency, Performance, and Implementation

WaveMix's computational and memory efficiency arises from several design features (Jeevan et al., 2022, Jeevan et al., 2022, Jeevan et al., 2023, Jeevan et al., 16 Sep 2024):

• Parameter Allocation: DWT and upsampling layers are parameter-free; parameters concentrate in lightweight convolutional and MLP layers.
• Miniaturized Models: Models such as WaveMix-128/7 achieve 2.4M parameters (vs. 11M+ for ResNet-18) and low training/inference RAM usage (0.2–2.3 GB for WaveMix vs. 13–15 GB for ViTs at matched batch size).
• Throughput and Latency: High batch sizes are possible within fixed budgets; inference throughput in WaveMixSR-V2 reaches up to 82.6 fps with an inference latency of 12.1 ms (Jeevan et al., 16 Sep 2024).
• Super-resolution Efficiency: WaveMixSR achieves state-of-the-art PSNR (e.g., 33.12 dB @ 2× on BSD100) with only 1.7M parameters, compared to 20.8M in HAT and 11.8M in SwinIR (Jeevan et al., 2023, Jeevan et al., 16 Sep 2024).
• Progressive Multistage Design: WaveMixSR-V2 employs stage-wise 2× upsampling for 4× tasks, further reducing model size and computational overhead.

In all settings, models benefit from the fixed, non-learnable DWT as opposed to learned QKV mixing (transformers) or large spatial convolutions.

5. Specialized Variants and Application Domains

WaveMixSR & WaveMixSR-V2 (Jeevan et al., 2023, Jeevan et al., 16 Sep 2024): Designed for image super-resolution, these variants retain the core WaveMix block but enhance upsampling. WaveMixSR-V2 replaces transposed convolution with PixelShuffle for computational savings and artifact mitigation and adopts a multistage approach for high-scale upsampling. The result is improved parameter efficiency, latency, and throughput, with state-of-the-art PSNR/SSIM on BSD100.

Augmentation for Time Series (WaveMix as Data Augmentation) (Arabi et al., 20 Aug 2024): Here, “WaveMix” refers to a wavelet-based augmentation strategy. For two samples $s_1$ and $s_2$, multi-level DWT is applied per channel. Learned masks exchange coefficients between the two sets. Inverse DWT reconstructs an augmented sample with mixed time-frequency profiles, preserving temporal dependencies—a capability lacking in frequency domain-only methods. The methodology yields the best or second-best results across 16 forecasting setups, with marked performance under data-scarce (cold-start) conditions.

Magnification-Invariant Image Analysis (Jeevan et al., 2023): WaveMix's capacity to robustly generalize across input scales, derived from its multi-resolution wavelet representation, outperforms CNNs, ViTs, and other token mixers in classification accuracy and stability across non-stationary (magnification-varying) domains.

Summary Table: Core WaveMix Applications

Variant	Task/domain	Distinctive technique	Resource efficiency	Notable SOTA/performance
WaveMix (Jeevan et al., 2022)	Image class/seg	Multi-scale 2D-DWT mixing	High	SOTA, ~2.4M params on EMNIST
WaveMixSR (Jeevan et al., 2023)	Super-resolution	2D-DWT + upconv	High	SOTA PSNR on BSD100
WaveMixSR-V2 (Jeevan et al., 16 Sep 2024)	Super-resolution	2D-DWT + PixelShuffle + multistage	Very high	Improved PSNR, lower latency/params
WaveMix (Arabi et al., 20 Aug 2024)	Time series aug	DWT-based coefficient mixing	N/A (augmentation)	Best/second-best (12/16 tasks)

6. Extensibility and Practical Integration

The modular, fully convolutional, and spatially aware construction of WaveMix allows for straightforward adaptation to classification, segmentation, super-resolution, and generative modeling. Key practical properties include:

• Preservation of the 2D spatial grid for images and explicit time-frequency structure for sequences.
• Plug-and-play compatibility as a backbone for larger frameworks (e.g., segmentation heads, U-Nets, detection pipelines).
• Suitability for edge, green, and embedded computing due to minimized parameterization and memory footprints (Jeevan et al., 2022, Jeevan et al., 2022, Jeevan et al., 16 Sep 2024).
• Public availability of training code and models fosters reproducibility and extension (Jeevan et al., 2022, Jeevan et al., 16 Sep 2024).

7. Conclusion

WaveMix represents a paradigm shift towards multi-scale, wavelet-based token mixing, with demonstrated advantages in efficiency and scale invariance for both discriminative and generative vision tasks, as well as data augmentation for time series. The core innovation—systematic, lossless, and multi-resolution analysis via DWT—enables robust performance in scenarios where resource constraints or input variations dominate. Successors such as WaveMixSR-V2 and augmentation adaptations affirm the extensibility and ongoing relevance of the approach in contemporary machine learning.