FNet: Mixing Tokens with Fourier Transforms

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.

PDF Abstract NAACL 2022 PDF NAACL 2022 Abstract
Task Dataset Model Metric Name Metric Value Global Rank Result Benchmark
Linguistic Acceptability CoLA FNet-Large Accuracy 78% # 9
Semantic Textual Similarity MRPC FNet-Large Accuracy 88% # 23
Natural Language Inference MultiNLI FNet-Large Matched 78 # 43
Mismatched 76 # 34
Natural Language Inference MultiNLI BERT-Large Matched 88 # 15
Mismatched 88 # 10
Natural Language Inference QNLI FNet-Large Accuracy 85% # 40
Paraphrase Identification Quora Question Pairs FNet-Large F1 85 # 5
Natural Language Inference RTE FNet-Large Accuracy 69% # 57
Sentiment Analysis SST-2 Binary classification FNet-Large Accuracy 94 # 37
Semantic Textual Similarity STS Benchmark FNet-Large Spearman Correlation 0.84 # 28