F ourier Transformer: Fast Long Range Modeling by Removing Sequence Redundancy with FFT Operator (2023)

Abstract: The transformer model is known to be computationally demanding, and prohibitively costly for long sequences, as the self-attention module uses a quadratic time and space complexity with respect to sequence length. Many researchers have focused on designing new forms of self-attention or introducing new parameters to overcome this limitation, however a large portion of them prohibits the model to inherit weights from large pretrained models. In this work, the transformer’s inefficiency has been taken care of from another perspective. We propose Fourier Transformer, a simple yet effective approach by progressively removing redundancies in hidden sequence using the ready-made Fast Fourier Transform (FFT) operator to perform Discrete Cosine Transformation (DCT). Fourier Transformer is able to significantly reduce computational costs while retain the ability to inherit from various large pretrained models. Experiments show that our model achieves state-of-the-art performances among all transformer-based models on the long-range modeling benchmark LRA with significant improvement in both speed and space. For generative seq-to-seq tasks including CNN/DailyMail and ELI5, by inheriting the BART weights our model outperforms the standard BART and other efficient models. Our code will be publicly available at https://github.com/LUMIA-Group/FourierTransformer

## introduction transformers @xcite , especially when equipped with large-scale pre-training @xcite @xcite have become the core architecture in most tasks in natural language processing (nlp), including both encoder-only tasks such as sentence classification, sequence tagging @xcite , and encoder-decoder tasks such as text summarization and question answering @xcite . however, due to the quadratic complexity of its selfattention module @xcite , applying these models on long sequences can be prohibitively costly. as a result, great efforts have been put into developing various efficient transformer variants @xcite , as well as establishing standardized test-beds for long sequences such as the long range arena (lra) @xcite . most efficient transformers devise special attention variants to lower its complexity @xcite . some of them achieve this by projecting components in self-attention into its lowerrank approximations @xcite @xcite , or rely on kernelization to implicitly compute the attention matrix @xcite @xcite . due to the introduction of projection matrices or extra parameters, these models are not able to inherit pre-trained model parameters. however, since pre-trained large language models (llms) have fundamentally influenced the nlp community, deviating model architecture from llms requires pre-training from scratch on the designed model, which is prohibitively resource-demanding for most practitioners. other approaches target at computing part of the attention matrix, by following some predefined patterns @xcite @xcite . some of them allow the pattern to be learnable @xcite . most of the patterns require customized cuda kernels or special operators to achieve the claimed speedup @xcite @xcite , which casts extra challenge in deploying these models on edge devices or special hardware such as tpus. moreover, some of the approaches involve considerable additional computation steps, which in practice could counterweight the time and memory complexity they reduce, especially for short and medium-length sequences @xcite . one core factor behind various approaches is the existence of redundancy in attention matrices and hidden states. for example, @xcite provides spectrum analysis on the self-attention matrix, indicating that the attention matrix learns to be low-rank, which allows them to learn a low-rank approximation of the attention matrix. inspired by this line of research, in this work, we analyze the power spectrum of the hidden states in the time dimension through different layers in fig 1 , and show that the power spectrum increasingly concentrates on lower frequency bins as the layer gets deeper. in this work, we propose fourier transformer, which doesn't even require to learn the projection matrix in order to approximate the self-attention. fourier transformer leverages our observation on power spectra of hidden states, it progressively removes sequence redundancies through different layers by downsampling hidden states with the discrete cosine transform (dct), a variant of fourier transform that generates real values. the dct in our proposed fourier transformer can be implemented with the fast fourier transform (fft) operator. thanks to its profound application in image compression and signal processing, it is one of the most widely available and highly optimized operators in a wide variety of frameworks and even on edge devices, providing o(n log n) complexity and up to o(log n) in parallel implementations with negligible overhead. as a result, fourier transformer is easily deployable on a wide range of devices, not necessary to devise special cuda kernels. in addition, experimental results on lra tasks show that it performs significantly faster than many other efficient transformers, while achieving the state-of-the-art performance among transformer-based efficient models. on the other hand, since dct is a linear, reversible transformation, and the self-attention is not interfered in our model, the proposed fourier transformer can inherit pretrained weights from large language models without hurting performance. experimental results on @xcite and eli5 @xcite show that our model could outperform bart @xcite and other efficient transformers by inheriting and fine-tuning on bart. moreover, with tiny amount of further pretraining before fine-tuning, its performance could be further improved. ## related work downsampling hidden states there are not many work that downsample sequence length for natural language. the closest work is funnel transformer @xcite , which progressively reduces the query sequence length through strided mean pooling, while keeping key and value sequence lengths intact. fourier transformer compresses the three sequences altogether and delivers more computational speedup compared with funnel transformer. note that funnel transformer needs to re-invest the saved computations to build a larger model to achieve better performance, which disables its ability to inherit pretrained weights. for other work, charformer @xcite devises a differentiable tokenization module that also relies on strided mean pooling to downsample its byte sequence. nyströmformer @xcite approximates the attention matrix through the nyström method, which effectively downsamples query and key sequences. due to the extra depth-wise convolution, it is again not able to leverage pretrained models. in a border view, downsampling has been more favorable in computer vision. @xcite aggressively downsamples the raw input to a 1d vector. perceiver @xcite adopts an asymmetric attention mechanism to distill inputs into a tight latent bottleneck. almost all of these vision models are designed for encoder-only vision tasks rather than encoder-decoder-style nlp tasks. fourier transform for transformer there are multiple recent works that incorporate fourier transform into transformer. fnet (lee-thorp et al., 2021) takes a more radical approach by replacing the entire self-attention with 2d fft, discarding the entire imaginary part to avoid complex numbers. performer @xcite introduced orthogonal random fourier features to approximate the softmax attention. fsat @xcite uses 1d fft along the sequence dimension to learn the sparse structure of attention matrix. dct-former @xcite translates sequences into frequency domain and conducts self-attention there before projecting them back, due to the nonlinearity in the network, self-attention trained in the frequency domain significantly deviates from that in the time domain. therefore, all the models dis- figure 1 : the power spectrum of input hidden states from different layers in the pretrained roberta @xcite model. the horizontal axes stand for frequency bins, starting from low frequency components on the left. the vertical axes are the corresponding amplitudes. amplitudes are averaged over all hidden dimensions and over the entire validation set of wiki-103 @xcite . since the inputs are real numbers, the positive and negative frequency components are pairwise conjugate. thus we only plot the amplitude of the positive half of the frequencies. ## experiments in this section, we experiment with our model in both of the two encoder-only and encoder-decoder settings in various datasets that involves long sequences. ## conclusion in this work, we introduce the discrete cosine transformation to progressively downsample the hidden states in the transformer model by leveraging the local correlations between hidden states in upper layers. our approach is able to significantly reduce the computation required by the vanilla transformer, while being able to achieve even better performance in various tasks. moreover, it is able to inherit the pretrained model weights, which is an notable advantage over most efficient transformers.