Wavelet-based token mixing speeds up transformer models for long sequences
WNet: Discrete Wavelets Transform for Efficient Token Mixing
Computation and Language
Summary
Transformers process language by mixing information between tokens, but this can be slow for long texts because every token compares itself to all others. The authors introduce WNet, which uses a discrete wavelet transform to mix tokens efficiently without relying on self-attention at every step. Their method speeds up training, especially on very long sequences, while still capturing important relationships between words. They show that using longer filters in wavelets helps tokens connect across the entire sequence more quickly.
What this means in practice
- •For natural language processing engineers: Train Transformer-based language models faster on very long text sequences using wavelet-based token mixing instead of conventional self-attention.
- •For speech recognition developers: Incorporate wavelet token mixing to efficiently handle long audio sequences in models requiring context spread across many frames.
Authors
Rana Aref Salama, Abdou Youssef, Mona Diab
Abstract
In a Transformer, token mixing is the step that lets each token draw information from other tokens, and it dominates the cost of encoding long sequences. Self-attention does this mixing very well: every token weighs every other token by content, which gives strong contextual modeling. That all-pairs comparison is also why its cost grows quadratically with sequence length. We introduce WNet, a Transformer encoder that replaces self-attention with token mixing based on the discrete wavelet transform (DWT). Three attention-free mixers recombine the scales: by linear fusion, by learned gating, or by letting each token choose its scales. A hybrid adds self-attention in the last layer only. A receptive-field analysis shows that wavelet mixers built from two-tap filters, such as Haar, never relate tokens outside fixed blocks, however deep the network, even when the filters are learned. Longer filters reach the whole sequence within two layers. We pre-train every model with masked language modeling on a fixed-token subset of C4 and fine-tune on GLUE, using one controlled setup with size-matched BERT and FNet baselines and a control that cannot mix tokens. The token-gated mixer trains as fast as attention at 256 tokens and 2.7 times faster at 4,096.