Organizations: George Washington University, Washington DC, USA · Faculty of Computers and Artificial Intelligence, Cairo University, Egypt · Carnegie Mellon University, Pittsburgh, Pennsylvania, USA
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.
Figures & tables
Figure 1: WNet. (a) Each attention sublayer of a pre-LN encoder is replaced by a wavelet mixer. (b) The mixer decomposes its input U along the sequence with a smoothing filter s that is dilated at every level; each detail is what one more smoothing step removes. (c) A learned layer recombines the bands.
Figure 2: Receptive field vs. depth at S=256 (exact). Two-tap designs with downsampling stay flat, whatever the depth. Full-resolution bands overlap and grow, and db10 reaches the whole sequence in two layers. Dotted line: the depth we use.
Mixer
Wavelet, J
Mixing params
RF 1
RF 4
ms @256
ms @4096
baselines
BERT (attention)
–
1053k
256
256
28.3
87.6
FNet
–
0
256
256
19.1
20.6
FNet + attn (last layer)
–
263k
256
256
19.4
36.6
LMWT-style
Haar, 5
1594k
32
32
20.0
20.5
No mixing (control)
–
798k
1
1
18.6
19.0
WNet
WNet-MS
db10, 3
1050k
127
256
27.9
28.1
Table 1: Token mixers in a 4-layer encoder with d=256 . Mixing parameters are summed over the four layers. RF n is the number of tokens that can influence one output position after n layers at S=256 (exact). Time is one encoder training step (forward, backward and update) on one T4 at a fixed 8,192 tokens per batch, median of 20 steps; repeated runs agree within ± 3%. FNet is timed with the rescaled transform (§ 6 ). ‡ Wavelet layers only; the attention layer is global.
Figure 3: Training-step time at a fixed 8,192 tokens per batch on one T4. Attention slows down with length; the attention-free WNet mixers do not.
Model
Loss ↓
SST-2
QNLI
MNLI m/mm
Core
Majority
–
50.9
50.5
35.4/35.2
45.6
BERT
7.946
62.3 ± 0.7
56.4 ± 0.3
34.1/34.2
50.9
FNet
8.048
50.9 ± 0.5
54.9 ± 0.6
34.0/34.3
46.6
FNet + attn
8.606
54.8 ± 0.4
54.8 ± 0.1
33.6/33.8
47.7
LMWT-style
7.905
51.1 ± 0.2
55.2 ± 0.4
33.4/33.5
46.6
No mixing
7.921
61.0 ± 0.3
57.3 ± 0.2
34.9/34.7
51.1
Table 2: Validation MLM loss and GLUE development accuracy on the core tasks (subscript: s.d.); Core averages SST-2, QNLI and MNLI-m. These numbers come from a quick run that validates the pipeline: 500 pre-training steps (one seed), with GLUE fine-tuned on 5,000 examples per task, mean of two fine-tuning seeds.
Variant
Loss
MNLI-m
ms@4k
WNet-Dyn (db10, J=3 )
7.919
34.5
33.0
downsampled DWT + inverse
7.920
34.5
42.1
Haar instead of db10
7.920
34.4
30.8
learned filter (init. db10)
7.916
34.6
35.1
no wavelet bands ( J=0 )
7.921
34.9
19.0
WNet-MS (db10, J=3 )
7.986
34.9
28.1
Table 3: Ablations, changing one component at a time. “Downsampled DWT + inverse” computes Eq. 1 and returns each band to full length with the inverse DWT; “+ interpolation” stretches the bands by linear interpolation instead ( ∗ coif4). ms@4k: step time at 4,096 tokens. Quick run as in Table 2 .
Appendix figures & tables4 assets
Supplementary material from the paper’s appendix.
Appendix
Figure 4: Coefficients affected by changing one token ( S=64 , position 21). The Fourier transform spreads the change over all coefficients with equal magnitude; a three-level Haar DWT confines it to one coefficient per band.
Layers / hidden / FFN
4 / 256 / 1024, pre-LN
Attention heads
8 (BERT, hybrids)
Tokenizer
FNet SentencePiece (32,000)
Sequence length
256, packed (no padding)
Pre-training data
C4, 85M tokens (one pass)
Objective
MLM 15%, 80/10/10
Batch / steps
32 / 10,000
Appendix
Table 4: Hyperparameters shared by all models.
Model
Loss
SST-2
QNLI
MNLI-m
MNLI-mm
Majority class
–
50.9
50.5
35.4
35.2
BERT
7.946
62.3 ± 0.7
56.4 ± 0.3
34.1 ± 0.6
34.2 ± 0.2
FNet (unscaled)
8.698
53.5 ± 1.1
54.5 ± 0.1
33.8 ± 0.3
34.0 ± 0.1
FNet (scaled)
8.048
50.9 ± 0.5
54.9 ± 0.6
34.0 ± 0.8
34.3 ± 1.4
FNet + attn
8.606
54.8 ± 0.4
54.8 ± 0.1
33.6 ± 0.6
33.8 ± 0.1
LMWT-style
7.905
51.1 ± 0.2
55.2 ± 0.4
33.4 ± 0.2
33.5 ± 0.1
Appendix
Table 5: All models in the quick run (500 pre-training steps, one seed; GLUE fine-tuned on 5,000 examples per task, mean of two fine-tuning seeds). “DWT + inverse”: downsampled bands returned to full length with the inverse DWT; ∗ coif4. The quick run evaluates only the core tasks.
Models, 500 steps
WNet-MS, 300 steps
LMWT-style
7.905
Haar
8.614
WNet-Dyn (db10)
7.919
sym10
8.628
WNet-Dyn (Haar)
7.920
db4
8.660
No mixing
7.921
coif4
8.663
BERT
7.946
db20
8.667
FNet (scaled)
8.048
db10
8.668
Appendix
Table 6: Validation loss early in training (one seed). Left: a model that cannot mix tokens matches the wavelet mixers and beats attention. Right: short runs favour short filters. Neither ordering reflects the ability to relate distant tokens.
LIGM, CNRS, Univ Gustave Eiffel, ENPC, Institut Polytechnique de Paris, France · LIX, École Polytechnique, CNRS, IP Paris, France · AMIAD, Pole recherche +4