WNet: Discrete Wavelets Transform for Efficient Token Mixing

TL;DR

WNet replaces self-attention in Transformers with discrete wavelet transforms to enable efficient token mixing without quadratic cost. It uses wavelet-based mixers with linear, gated, or token-adaptive scale fusion, achieving faster training than attention at long sequence lengths while maintaining strong performance.

Why It Matters

This approach enables scalable modeling of long sequences by avoiding the computational bottleneck of self-attention, making it practical for large-scale language tasks without sacrificing contextual quality. It opens new pathways for efficient architecture design in sequence modeling.

Key Findings

• Wavelet-based mixers eliminate the need for self-attention while enabling effective token interactions through scale recombination.
• Two-tap wavelet filters like Haar limit token connectivity to fixed blocks regardless of depth, but longer filters achieve full sequence coverage in just two layers.
• The token-gated mixer trains 2.7 times faster than self-attention at 4,096 tokens and matches attention speed at 256 tokens.