FN et: Mixing Tokens with F ourier Transforms (2022)

Abstract: 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. 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.

## introduction the transformer architecture @xcite has achieved rapid and widespread dominance in nlp. at its heart is a attention mechanism -an inductive bias that connects each token in the input through a relevance weighted basis of every other token. many papers have prodded and probed the transformer, and in particular the attention sublayers, in an effort to better understand the architecture; see, for example, @xcite ; @xcite ; @xcite ; @xcite . although potentially limited in their effectiveness @xcite , these probes generally back the intuition that, by allowing higher order units to form out of compositions of the input, transformer models can flexibly capture diverse syntactic and semantic relationships. in this work, we investigate whether simpler token mixing mechanisms can wholly replace the relatively complex self-attention layers in transformer encoder architectures. we first replace the attention sublayer with two parameterized matrix multiplications -one mixing the sequence dimension and one mixing the hidden dimension. seeing promising results in this simple linear mixing scheme, we further investigate the efficacy of faster, structured linear transformations. surprisingly, we find that the fourier transform, despite having no parameters at all, achieves nearly the same performance as dense linear mixing and scales very efficiently to long inputs, especially on gpus (owing to the o(n log n ) fast fourier transform (fft) algorithm). we call the resulting model fnet. while fourier transforms have previously been used to approximate or speed up computations in convolutional neural networks @xcite @xcite @xcite @xcite , recurrent neural networks @xcite @xcite , transformers @xcite , and mlp layers more generally @xcite @xcite , we believe our work is the first to wholly replace particular neural network sublayers with a fourier transform. this approach of viewing the fourier transform as a first class mixing mechanism is reminiscent of the mlp-mixer @xcite for vision, which replaces attention with mlps; although in contrast to mlp-mixer, fnet has no learnable parameters that mix along the spatial dimension. given the favorable asymptotic complexity of the fft, our work also connects with the literature on "long sequence" or "efficient" transformers, which aim to make the attention mechanism scale better via sparsity patterns @xcite @xcite @xcite @xcite @xcite @xcite or via linearization of the attention matrix @xcite @xcite . as we will show in our experiments, while some of those works achieve o(n ) scaling of attention, this complexity often hides large constants, which make them less scalable in practice than fnet. the contributions of our paper are: