Source-linked AI summary
WNet: Discrete Wavelets Transform for Efficient Token Mixing
Rana Aref Salama, Abdou Youssef, Mona Diab
TL;DR
Because language has short-span morphology and syntax but long-span discourse, DWT matches this organization by splitting signals into coarse summaries and local details in O(S) time. WNet replaces attention with a full-resolution wavelet decomposition and recombines its scales through linear fusion, learned gating, or token-chosen scales. WNet-Dyn trains as fast as attention at 256 tokens and 2.7× faster at 4,096.
Problem
Because language has short-span morphology and syntax but long-span discourse, DWT matches this organization by splitting signals into coarse summaries and local details in O(S) time.
Source: §1 · p. 1Method
Results
WNet-Dyn trains as fast as attention at 256 tokens and 2.7× faster at 4,096.
Source: §9 · p. 7Takeaways & Limitations
Two-tap wavelet designs remain confined to fixed blocks at any depth, whereas longer filters reach the whole sequence within two layers.
Source: §1 · p. 1, §9 · p. 7Quality is evaluated only at 256 tokens, so the paper makes no claim about long-context accuracy.
Source: Limitations · p. 7
Abstract
from arXiv · showhide
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.
1 Introduction
WNet replaces attention with linear-time, full-resolution wavelet token mixing, motivated by the multiscale structure of language and the quadratic cost of attention. Its mixers recombine wavelet scales in three ways, while receptive-field analysis distinguishes fixed-block two-tap designs from longer filters that reach the sequence.
1 Introduction
- Language has short-span morphology and syntax as well as long-span discourse, while DWT represents coarse summaries and local details in O(S) time.By contrast, self-attention has content-dependent all-pairs comparisons and quadratic sequence-length cost; Fourier mixing lacks locality.Source: §1 · p. 1
- WNet replaces each attention sublayer with a linear-time wavelet mixer built on a full-resolution multiscale decomposition.The mixer keeps one value per token at every scale, and the scales add up to the input and shift with it.Source: §1 · p. 1
- Linear fusion, learned gating, and token-chosen scales are the three attention-free ways WNet recombines its wavelet bands.WNet-Hybrid adds self-attention in the last layer only.Source: §1 · p. 1
- Two-tap wavelet mixers, including learned Haar designs, remain confined to fixed blocks at any depth, whereas longer filters reach the whole sequence within two layers.This is a receptive-field result, not an empirical benchmark comparison.Source: §1 · p. 1
- The standard downsampled DWT represents a phrase differently depending on its starting position.The full-resolution construction instead preserves shifts, as described in the method section.Source: §1 · p. 2
- The comparison uses shared data, depth, width, optimiser, and budget, with size-matched BERT and FNet, learnable Haar mixing, and a no-mixing control.Speed is reported up to 4,096 tokens.Source: §1 · p. 2
2 Related Work
Prior attention-free mixers include learned token projections, Fourier-domain methods, and other long-sequence approaches; WNet replaces attention rather than approximating it. The paper builds on the DWT’s recursive decomposition into coarse approximations and fine details.
2 Related Work
- Efficient-attention methods approximate attention, whereas WNet replaces it with a token mixer.Prior approaches include sparse patterns, low-rank projections, kernel approximations, pooling, and cosine transforms.Source: §2 · p. 2
- Attention-free alternatives include learned token projections, fixed or learned Fourier mixers, long convolutions, structured matrices, and state-space models.WNet-Dyn’s gating follows gated linear units, but its mixing remains a fixed wavelet transform.Source: §2 · p. 2
- For sequences, WavSpA and Wavelet GPT retain attention, while LMWT replaces it with five levels of learned two-tap Haar decomposition.LMWT reports 27.2 BLEU on WMT16 En–De against 27.8 for attention.Source: §2 · p. 2
3 Discrete Wavelet Transform
- The DWT recursively filters and downsamples a signal into approximation coefficients for low-varying components and detail coefficients for fine-scale variation.The two coefficient sets recover the input exactly, and repeating the process on the approximation yields coarser summaries.Source: §3 · p. 2
- Haar filters have two taps, while db10 has 20; longer filters summarize wider token stretches, unlike Fourier coefficients that depend on the whole sequence.After J levels, the standard DWT represents the input with one approximation and details at progressively finer resolutions.Source: §3 · p. 2
4 WNet
WNet retains a pre-LN Transformer encoder but replaces each attention sublayer with a mixer that decomposes each channel along the sequence and recombines full-resolution bands. Its recombination variants use concatenation and projection, learned band gating, token-selected scales, or a final attention layer.
4 WNet
- Each WNet mixer decomposes hidden states separately by channel along the sequence, then recombines the resulting bands in the encoder’s attention-sublayer position.The encoder retains pre-LN structure.Source: §4 · p. 2
4.1 Token Mixing
- The mixer forms a smooth band summarizing a wide neighborhood and detail bands recording progressively finer changes, with one value per token in every band.A learned layer recombines the bands, and stacking layers widens context.Source: §4.1 · p. 3, Figure 1 · p. 3
- The undecimated à trous DWT keeps every band at full sequence length, makes the bands sum exactly to the input, preserves shifts, and costs one convolution per level.Skipping downsampling also avoids interpolation.Source: §4.1 · p. 3
4.2 Recombining the scales
- WNet-MS concatenates the bands and projects them, with its projection parameter count matching attention’s four projections when J=3.The learned fusion combines the full-resolution scales.Source: §4.2 · p. 3
- WNet-G uses learned per-channel weights to combine bands into a gate and signal, while WNet-Dyn lets each token choose its bands by channel group.In WNet-Dyn, the mixed context multiplies a per-token value; initialization favors the coarse band.Source: §4.2 · p. 3
- WNet-Hybrid uses WNet-Dyn with self-attention in the last layer only.Source: §4.2 · p. 4
5 How Far Can Information Travel?
The receptive field measures how many input tokens can influence an output position after a stack of mixers. Downsampled two-tap designs stay within fixed blocks, while full-resolution overlapping bands grow with depth and longer filters can reach the whole sequence quickly.
5 How Far Can Information Travel?
- The receptive field is computed exactly by propagating each mixer’s dependency pattern through the stack, including the residual path.Source: §5 · p. 4
- Two-tap downsampled filters keep each output within an aligned block of 2^J tokens at every depth, even when their weights are learned.In a five-level Haar design, each token sees 32 tokens after one or six layers; a 128-token sentence remains split into four noncommunicating blocks.Source: §5 · p. 4
- With overlapping full-resolution bands, Haar’s receptive field grows to 8, 15, 22, and 29 tokens over four layers; db10 WNet-MS and WNet-Dyn reach all 256 tokens in two layers.WNet-G with four levels sees 254 tokens after one layer.Source: §5 · p. 4, Figure 2 · p. 4
- The standard downsampled DWT is position-sensitive: shifting a random 256-token input by one token changes its three-level bands by 110% of the norm for Haar and 67% for db10.Full-resolution bands shift with the input, with measured error below 10^-6.Source: §5 · p. 4
- Wavelet mixers cost O(JmSd), while attention costs O(S^2d) plus O(Sd^2) for projections; with d=256, attention’s quadratic term overtakes projections around S=1024.The paper therefore measures cost at longer lengths.Source: §5 · p. 4
6 Experimental Setup
The experiments hold the encoder and training pipeline fixed while comparing token mixers against BERT, FNet, and other controls. Evaluation combines masked-language-model loss, selected GLUE tasks, and training-step time.
6 Experimental Setup
- All models use a four-block, 256-hidden-unit encoder and the same fixed C4 subset, tokenization, optimizer, batch size, and 10,000-step pre-training budget.Pre-training uses masked language modeling with 15% of tokens selected for replacement.Source: §6 · p. 4
- The baselines include size-matched BERT, FNet with a rescaled Fourier transform, FNet + attn, and a five-level learnable-Haar LMWT-style reimplementation.Scaling FNet's transform by 1/Sd lowers its 500-step validation loss from 8.70 to 8.05.Source: §6 · pp. 4–5
- The no-mixing control retains WNet-Dyn's gates and projections but gives each token access only to itself, isolating gains attributable to token mixing.It uses no wavelet bands (J=0).Source: §6 · p. 5
- Evaluation reports held-out MLM loss, GLUE development scores, and median encoder-step time at a fixed 8,192 tokens per batch.CoLA, MRPC, RTE, and STS-B remain near majority-class performance for every model, so SST-2, QNLI, and MNLI are the core tasks.Source: §6 · p. 5, Table 1 · p. 5
7 Results
WNet's token-gated and scale-fusion mixers reduce long-sequence step time relative to BERT, and every WNet mixer exceeds FNet and LMWT-style on the quick-run core average. Results also show that reach and task requirements matter: pair-local Haar falls short on MNLI, while longer filters expand receptive fields.
7.1 Efficiency
- At fixed tokens per batch, BERT's step time rises from 28 to 88 ms between 256 and 4,096 tokens, while attention-free WNet mixers stay flat.WNet-Hybrid inherits some attention-related growth but is still 1.9× faster than BERT at 4,096 tokens.Source: §7.1 · p. 5
- At 4,096 tokens, WNet-Dyn is 2.7× faster than BERT and WNet-MS is 3.1× faster; FNet remains fastest overall at about 19 ms.WNet-Dyn's step time at 256 tokens is 29.8 ms versus BERT's 28.3 ms.Source: §7.1 · p. 5
7.2 Pre-training and GLUE
- Every WNet mixer scores 48.8–50.1 on the core average, ahead of FNet and LMWT-style at 46.6.The core average covers SST-2, QNLI, and MNLI-m.Source: Table 2 · p. 6, §7.2 · p. 6
- On quick-run GLUE tasks, MNLI is at chance for every model; BERT leads SST-2 at 62.3, while WNet-Dyn exceeds BERT on QNLI (57.9 vs. 56.4).WNet-Hybrid scores 57.8 on QNLI and 58.1 on SST-2.Source: Table 2 · p. 6, §7.2 · p. 6
- In the 500-step run, WNet-Hybrid has the lowest MLM loss (7.902 vs. BERT's 7.946), while the no-mixing control reaches 7.921.WNet-Dyn and WNet-G also score below BERT on loss; WNet-MS scores above it.Source: §7.2 · p. 6
- In a preliminary full-protocol run, single-level Haar trails BERT by 4.8 points on MNLI-m (49.6 vs. 54.4), despite nearly matching it on SST-2 and QNLI.The authors identify MNLI as requiring premise–hypothesis relations that a pair-local mixer lacks.Source: §7.2 · p. 6
7.3 Ablations
- Full-resolution decomposition is 22%, 27%, and 38% faster than downsampled DWT plus inverse for WNet-Dyn, WNet-MS, and WNet-G, respectively.The authors attribute this to convolving each band directly without reconstructing it; transform, wavelet, and filter choices barely change loss, and removing mixing costs 0.002.Source: Table 3 · p. 6, §7.3 · p. 6
8 Discussion
The discussion cautions that early training loss can reward locality rather than useful token mixing, and argues that receptive fields can be checked before training. It connects the limits of fixed mixing to sentence-alignment tasks.
8 Discussion
- Early loss cannot reliably evaluate token mixing: the no-mixing control matches wavelet mixers after 500 steps, and short runs favor short filters.The authors therefore chose db10 in advance for its receptive field.Source: §8 · p. 6
- A receptive-field check takes no training and a few lines of code, and identifies whether distant tokens can interact.A pair- or block-local mixer cannot relate distant tokens, however long it is trained.Source: §8 · p. 6
- Two-tap filters keep tokens within fixed blocks at any depth, whereas longer filters reach the whole sequence within two layers.This analysis explains the MNLI gap of the single-level Haar mixer and applies to published designs such as learnable Haar mixing.Source: §8 · p. 7
- Fixed mixing combines tokens the same way regardless of their words, while WNet-Dyn lets each token choose scales and WNet-Hybrid adds one full-attention layer.The passage contrasts fixed mixing for local-cue tasks with sentence alignment, which fixed mixing does not suit.Source: §8 · p. 7
9 Conclusion
WNet replaces self-attention with wavelet-based token mixing. Its full-resolution scales are recombined by a learned layer, while filter length determines how far information can travel.
9 Conclusion
- WNet replaces self-attention with wavelet token mixing.It is a family of linear-time mixers built on a full-resolution multiscale decomposition.Source: §9 · p. 7
- The decomposition's scales retain sequence length, sum to the input, and shift with it before a light learned layer recombines them.Source: §9 · p. 7
- Two-tap filters confine tokens to fixed blocks, whereas longer filters reach the whole sequence within two layers; WNet-Dyn is 2.7× faster than attention at 4,096 tokens.WNet-Dyn trains as fast as attention at 256 tokens.Source: §9 · p. 7
Limitations
The experiments use small models and report quality only at 256 tokens, limiting the scope of the conclusions. The study covers English bidirectional encoders, not causal decoding.
Limitations
- All models have four layers and hidden size 256; absolute GLUE scores are far below published BERT and FNet results, and rankings may change at larger scale.Four GLUE tasks stay near the majority class.Source: Limitations · p. 7
- Cost is measured up to 4,096 tokens, but quality is measured only at 256, so the paper makes no claim about long-context accuracy.Source: Limitations · p. 7
- The study covers English bidirectional encoders; causal decoding would require one-sided filters because centred filters let tokens see the next one.Source: Limitations · p. 7
- Three runs are used for comparisons carrying claims, so only large, consistent differences reach significance; some GPU kernels are nondeterministic.Reruns of those models agree to about three decimals.Source: Limitations · p. 7
Ethical Considerations
The pretraining corpus is a filtered web corpus that may contain social biases and personal information. The authors describe the models as small research artefacts not intended for deployment.
Ethical Considerations
- Pretraining uses C4, a filtered web corpus that contains social biases and may contain personal information.Source: Ethical Considerations · p. 7
- The authors describe their models as small research artefacts not intended for deployment.Source: Ethical Considerations · p. 7
- All experiments ran on a single GPU.Source: Ethical Considerations · p. 7
A The DWT in One Picture
The section contrasts Fourier’s global coefficient response to a token change with the DWT’s region-local response, then describes full-resolution multiscale decomposition built from smoothing and detail bands.
A The DWT in One Picture
- Changing one token affects every Fourier coefficient equally, while a three-level Haar DWT confines the change to one coefficient per band.The figure uses sequence length 64 and token position 21.Source: A The DWT in One Picture · p. 8, Figure 4 · p. 9
B The À Trous Decomposition
- À trous levels smooth the prior approximation and retain removed detail; the bands telescope to the input for any filter, including learned filters, and shift with the input.Each level is a dilated depthwise convolution, and no inverse transform is needed.Source: B The À Trous Decomposition · p. 9
C Implementation Details
- Table 4 specifies hyperparameters shared across models; implementation checks compare transforms against reference operations and test full-resolution reconstruction and shift behavior.Reported errors are 0.0 for the downsampled DWT, 2.4×10−7 for interpolation, and <10−6 and 0.0 for full-resolution reconstruction and shift behavior, respectively.Source: Table 4 · p. 9, C Implementation Details · p. 9
D Full Results
- Table 5 covers 500 pre-training steps with one seed, GLUE fine-tuning on 5,000 examples per task with two fine-tuning seeds averaged, and core tasks only.The table includes all models in the quick run.Source: Table 5 · p. 9
E Early Training and the Wavelet Sweep
- In one-seed early validation, a model that cannot mix tokens matches wavelet mixers and beats attention; short runs favor short filters, but neither ordering reflects distant-token ability.Table 6 reports early-training validation loss.Source: Table 6 · p. 9