Source-linked AI summary
FNet: Mixing Tokens with Fourier Transforms
James Lee-Thorp, Joshua Ainslie, Ilya Eckstein, Santiago Ontanon
TL;DR
Transformer encoders rely on costly self-attention, motivating the question of whether simpler token mixers can replace it. The paper introduces FNet, which uses an unparameterized Fourier Transform in place of encoder self-attention. FNet retains 92–97% of BERT counterparts’ GLUE accuracy while training substantially faster, and matches the most accurate efficient Transformers on long-input tasks with lower speed and memory costs.
Problem
The paper asks whether simpler token-mixing mechanisms can replace self-attention despite attention’s quadratic time and memory bottleneck for long sequences.
Method
FNet replaces each Transformer encoder self-attention sublayer with a 2D Fourier Transform consisting of sequence- and hidden-dimension mixing, followed by a feed-forward sublayer.
Results
FNet achieves 92% and 97% of BERT-Base and BERT-Large accuracy on GLUE while training 70–80% faster, and matches the accuracy of the most accurate efficient Transformers on Long-Range Arena.
Takeaways & Limitations
Linear units, particularly Fourier mixing, can serve as effective and efficient replacements for attention in text classification and long-input encoder settings.
Takeaways & Limitations
The study focuses on encoder and non-generation tasks, leaving FNet decoders and encoder-decoder cross-attention for future work.
Abstract
from arXiv · showhide
We show that Transformer encoder architectures can be sped up, with limited accuracy costs, by replacing the self-attention sublayers with simple linear transformations that "mix" input tokens. These linear mixers, along with standard nonlinearities in feed-forward layers, prove competent at modeling semantic relationships in several text classification tasks. Most surprisingly, we find that replacing the self-attention sublayer in a Transformer encoder with a standard, unparameterized Fourier Transform achieves 92-97% of the accuracy of BERT counterparts on the GLUE benchmark, but trains 80% faster on GPUs and 70% faster on TPUs at standard 512 input lengths. At longer input lengths, our FNet model is significantly faster: when compared to the "efficient" Transformers on the Long Range Arena benchmark, FNet matches the accuracy of the most accurate models, while outpacing the fastest models across all sequence lengths on GPUs (and across relatively shorter lengths on TPUs). Finally, FNet has a light memory footprint and is particularly efficient at smaller model sizes; for a fixed speed and accuracy budget, small FNet models outperform Transformer counterparts.
1 Introduction
FNet investigates whether simple token-mixing transformations can replace Transformer self-attention, finding that an unparameterized Fourier Transform preserves much accuracy while improving speed and scalability.
- Core approach: FNet replaces each Transformer encoder self-attention sublayer with a Fourier Transform as its token-mixing mechanism.The model combines Fourier mixing with standard feed-forward layers.
- GLUE results: 92% and 97% of BERT-Base and BERT-Large accuracy, respectively, are achieved on GLUE while training 80% faster on GPUs and 70% faster on TPUs.These results use the FNet model with Fourier mixing.
- Hybrid models: 97−99% of BERT counterparts’ GLUE accuracy is achieved by hybrid FNet models containing only two self-attention sublayers, while running 40−70% faster.The result indicates that attention can improve accuracy without being required in every layer.
- Long inputs: On the Long-Range Arena benchmark, FNet matches the accuracy of the most accurate efficient Transformers while being significantly faster across GPU sequence lengths.On TPUs, FNet is faster at relatively shorter lengths, while faster long-sequence alternatives are less accurate.
- Implication: The paper argues that new mixing mechanisms may be more valuable than increasingly efficient approximations of attention.This conclusion follows from FNet’s speed and accuracy comparisons with efficient Transformers.
2 Related work
Related work examines Fourier methods, attention alternatives, and efficient Transformers, while motivating FNet through attention’s computational cost and the practical limits of nominally linear approaches.
- Fourier methods: Fourier methods have been used in neural networks for convolution, recurrence, dense linear-layer approximation, and several Transformer-related approaches.The paper distinguishes these prior uses from wholly replacing particular neural-network sublayers with a Fourier Transform.
- Attention alternatives: Attention models’ success is commonly attributed to the flexibility and capacity of token-dependent attention patterns, but the paper asks whether their cost is necessary.Prior work suggests token-dependent weights may be expressive without being crucial for accurate NLP models.
- Efficient Transformers: Standard attention has quadratic time and memory complexity in sequence length, limiting tasks involving long-range dependencies.This motivates efficient and long-sequence Transformer research.
- Efficient Transformers: Some efficient Transformers achieve O(N) theoretical attention complexity, but hidden constants can make them less scalable in practice.Global-token requirements can also introduce quadratic dependence on the number of global tokens.
- Efficient Transformers: The Long-Range Arena compares efficient Transformers on long-range tasks, including models reported as fastest and lowest-memory on TPUs.FNet instead completely replaces self-attention with Fourier mixing.
- Experimental framing: The paper compares vanilla BERT and vanilla FNet, treating more recent Transformer optimizations as orthogonal to its token-mixing investigation.This comparison is framed as an analysis of different mixing mechanisms.
3 Model
FNet replaces Transformer self-attention with Fourier-based token mixing, combining a 2D discrete Fourier transform with feed-forward layers. Its design uses Fourier structure to mix sequence and hidden dimensions, with implementation choices balancing asymptotic complexity and hardware-specific speed.
- Discrete Fourier Transform: The DFT represents each output Xk as a sum of all input tokens xn weighted by twiddle factors, while FFT computation reduces the cost to O(N log N).Applying the DFT directly as a matrix multiplication costs O(N^2).
- FNet architecture: FNet applies a 2D discrete Fourier transform to each embedding input, mixing tokens along both the sequence and hidden dimensions.The sequence and hidden transformations are one-dimensional DFTs, denoted Fseq and Fh; their ordering is immaterial because they commute.
- FNet architecture: Each FNet encoder layer places a Fourier mixing sublayer before a feed-forward sublayer, following Transformer-style residual connections and layer normalization.Only the real part is retained after both Fourier transformations, so the feed-forward and output layers need not handle complex numbers.
- FNet architecture: FNet can be interpreted as alternating Fourier-domain multiplications and time-domain convolutions, although residual connections and retaining only real values complicate this intuition.The real-part operation makes the transformation non-invertible, so the alternating-domain view is only an intuition.
- Implementation: FNet can perform equally well without position embeddings because positional information is encoded by the Fourier transform, although embeddings are retained for comparison with BERT.The model otherwise uses word, absolute position, and sentence-type embeddings in the reported comparison setup.
- Implementation: On GPUs, FFTs are faster than matrix multiplications across 512–8192 tokens, whereas TPUs favor cached DFT matrix multiplication for sequences up to 4096 tokens.For longer TPU sequences, the FFT is faster; the implementation therefore selects different mixing computations by hardware and sequence length.
4 Results
The results show that replacing self-attention with Fourier or other linear mixing can substantially improve efficiency while retaining competitive accuracy. FNet performs especially well in speed-accuracy trade-offs, long-sequence settings, and hybrid models that retain limited attention.
- Model configurations: FNet replaces every self-attention sublayer with a Fourier sublayer, while Linear, Random, and FF-only encoders use alternative mixing configurations.The Linear encoder uses learnable dense transformations; Random uses constant random matrices; FF-only removes token mixing entirely.
- Ablations: The FF-only model severely underperforms, while random mixing is also weak, indicating that token mixing and structured transformations are important for model expressivity.The paper reports approximately 50% accuracy on several binary classification tasks for FF-only models.
- GLUE transfer learning: 97% and 99% of the respective BERT accuracies are achieved by hybrid FNet models with only two self-attention sublayers.These hybrids incur only limited speed degradation, supporting selective rather than universal use of attention layers.
- Speed-accuracy trade-offs: FNet and Linear define the GPU pre-training efficiency frontier for smaller, faster models, while BERT and FNet-Hybrid define it for larger, slower models.The same trends are reported for GPU and TPU results.
- Long-Range Arena: FNet achieves accuracy comparable to the most accurate efficient Transformers on LRA while training and inferring faster than all evaluated architectures across GPU sequence lengths.On TPUs, FNet is faster at relatively shorter lengths; at longer lengths, faster alternatives are less accurate on LRA.
5 Conclusions
FNet replaces Transformer self-attention with unparameterized Fourier mixing, retaining much of BERT’s accuracy while improving speed and memory efficiency. The paper also identifies hybrid attention models and broader encoder-only scope as important boundaries and directions.
- FNet achieves 92% and 97% of BERT-Base and BERT-Large accuracy on GLUE while training 70–80% faster on GPUs and TPUs.The Fourier-based model offers a speed, memory, and accuracy compromise.
- FNet matches the accuracy of the most accurate efficient Transformers on Long-Range Arena while being much faster and lighter on memory.Its favorable scaling properties make it competitive for long inputs.
- The authors surveyed other linear transformations only cursory, so additional fast mixing alternatives remain worth exploring.The Fourier Transform was especially effective because of FFT speed.
- Replacing only the final two Fourier sublayers with self-attention yields 97–99% of BERT’s accuracy with limited speed penalties.Hybrid models provide a direct speed–accuracy trade-off.
- The study only evaluates encoder architectures and non-generation tasks, leaving FNet decoders and encoder–decoder cross-attention for future work.The authors note that adapting Fourier mixing to cross-attention remains open because cross-attention may be crucial to performance.
A.1 Pre-training details
The pre-training experiments follow BERT-style configurations while comparing models with different mixing mechanisms and sizes. FNet trains substantially faster than BERT, though BERT remains more accurate in the reported MLM comparison.
- The experiments use BERT’s Base and Large model and learning configurations, training on C4 with a 32,000-token SentencePiece vocabulary.GPU runs use batch size 64 on 8 V100 chips, while TPU runs use batch size 256 on 4 × 4 TPU v3 chips.
- FNet and the Linear model train nearly 80% faster than BERT on GPUs and 70% faster on TPUs, despite weaker accuracy metrics.The models without learnable mixing parameters were also the most stable during training.
- BERT-Base is more accurate than FNet-Large despite FNet-Large containing more than twice as many parameters.The comparison indicates that parameter count alone does not explain BERT’s higher MLM accuracy.
- Each FNet-Hybrid model contains two self-attention sublayers, while feed-forward size and attention-head count are fixed relative to hidden size.The model-size table ignores output projection layers and excludes hybrids with only two total layers.
A.2 TPU results
On TPUs, FNet and Linear models occupy the speed–accuracy frontier for smaller models, while FNet’s long-sequence advantage depends on sequence length and implementation. Alternative transforms generally trade accuracy for speed or match Fourier performance without improving it.
- FNet and the Linear model define the TPU Pareto frontier for smaller, faster models, while BERT defines it for larger, slower models.The frontier is shown as speed versus MLM pre-training accuracy on 4 × 4 TPU v3 chips.
- On TPU Long-Range Arena text classification, FNet trains faster than all efficient Transformers through sequence length 512, but is slower than Performer at longer lengths.The cited comparison reports that FNet is more accurate than the linear-attention alternatives discussed there.
- Alternative transforms: The DCT underperforms FNet by approximately 4% accuracy, whereas the Hadamard Transform is approximately 2% less accurate despite being slightly faster.These alternatives do not improve the overall accuracy result of Fourier mixing.
- Alternative transforms: The Hartley Transform matches the Fourier Transform on GLUE at 76.7 versus 76.7.Hartley transforms real input to real output and is defined from the real and imaginary parts of the Fourier Transform.
- Design variants: Adding learnable parameters to the Fourier sublayer was detrimental or inconsequential and generally slowed the model slightly.The experiments tried element-wise multiplication and regular matrix multiplication approaches.
- Design variants: Modified FNet block structures degraded accuracy and produced a more unstable model during training.Extra feed-forward sublayers and simpler dense replacements did not help either.
A.4 Mixing layer speeds
Fourier mixing is substantially faster than self-attention, but whole-model gains are limited because FNet and BERT share feed-forward sublayers. The reported speed measurements isolate mixing-layer computation from the rest of the architecture.
- FNet training speed is particularly fast because its mixing layer has no parameters to update.On TPUs, Linear inference is faster than FNet because it uses real rather than complex matrix multiplications.
- FNet is 70–80% faster than BERT overall because shared feed-forward sublayers bottleneck whole-model training and inference.The Fourier mixing sublayer itself is significantly faster than self-attention in forward and backward passes.
A.5 FNet-Hybrid ablations
FNet-Hybrid ablations vary how many self-attention layers are used and where they are placed, showing an accuracy–speed trade-off. The main experiments therefore use two attention layers in the top configuration.
- The ablation table reports GPU pre-training accuracy and speed for Base FNet-Hybrid models after 100k steps.
- More attention improves accuracy at the cost of speed, with diminishing returns.
- Top placement of attention layers gives the best accuracy results among the tested layouts.
- The selected main-text variant uses two attention layers in the top configuration to prioritize speed.
A.6 A note on Long-Range Arena hyperparameter settings
The Long-Range Arena comparisons use configurations adjusted where the paper and repository differed or omitted details. Accuracy comparisons with prior runs should be interpreted cautiously because some task results vary substantially across runs.
- Several Long-Range Arena hyperparameters are unspecified or mismatched between the reference paper and its code repository.
- For CIFAR10, one layer performed best for all models, unlike the three-layer setting suggested by Tay et al. (2021a).
- For Pathfinder, a base learning rate of 0.001 outperformed the 0.01 value indicated by Tay et al. (2021a) for all models.
- Comparisons with Tay et al. (2021a) are caveated because Text and Retrieval results can vary considerably between runs, especially for Transformer models.
A.7 FNet code
The FNet code is implemented in JAX/Flax with a Fourier-transform layer and a feed-forward layer. The listing omits embedding and output projection layers for simplicity.
- The FourierTransformLayer applies a real fast Fourier transform with JAX over the input.
- The FeedForwardLayer defines the feed-forward component with configurable hidden dimension and dropout rate.
- The code listing omits the embedding and output projection layers.