SlowGuess/NAACL_2022_OCR huggingface dataset:
James Lee-Thorp and Joshua Ainslie and Ilya Eckstein and Santiago Ontañón  
Google Research
{jamesleethorp, jainslie, ilyaeck, santiontanon}@google.com
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 TPU's 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 TPU's). 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.
The Transformer architecture (Vaswani et al., 2017) 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, Tenney et al. (2019); Vig and Belinkov (2019); Clark et al. (2019); Voita et al. (2019). Although potentially limited in their effectiveness (Hewitt and Liang, 2019), 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  $\mathcal{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 (El-Bakry and Zhao, 2004; Mathieu et al., 2014; Highlander and Rodriguez, 2015; Pratt et al., 2017; Lin et al., 2018; Chitsaz et al., 2020; Goldberg et al., 2020), Recurrent Neural Networks (Koplon and Sontag, 1997; Zhang and Chan, 2000; Zhang et al., 2018), Transformers (Choromanski et al., 2020; Tamkin et al., 2020), and MLP layers more generally (Cheng et al., 2015; Moczulski et al., 2016; Sindhwani et al., 2015), 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 (Tolstikhin et al., 2021) 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,
1Code is available at https://github.com/google-research/google-research/tree/master/f_net.
Human Language Technologies, pages 4296 - 4313
which aim to make the attention mechanism scale better via sparsity patterns (Child et al., 2019; Qiu et al., 2020; Parmar et al., 2018; Beltagy et al., 2020; Ainslie et al., 2020; Zaheer et al., 2020; Wang et al., 2020; Tay et al., 2020b,a; Kitaev et al., 2020; Roy et al., 2021; Vyas et al., 2020; Liu et al., 2018) or via linearization of the attention matrix (Katharopoulos et al., 2020; Choromanski et al., 2021; Peng et al., 2021). 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:
- We show that simple linear transformations, including even (parameter-free) Fourier Transforms, along with standard MLPs in feedforward layers, are competent at modeling diverse relationships in text. That such a simple linear transformation works at all is surprising, and suggests that, for at least some NLP problems, attention may not be the principal component driving the performance of Transformers.
- We introduce a new model, FNet, that uses the Fourier Transform as a mixing mechanism. FNet offers an excellent compromise between speed, memory footprint, and accuracy, achieving  $92\%$  and  $97\%$ , respectively, of the accuracy of BERT-Base and BERT-Large (Devlin et al., 2019) on the GLUE benchmark (Wang et al., 2018), while training  $80\%$  faster on GPUs and  $70\%$  faster on TPU's.
- We find that FNet hybrid models containing only two self-attention sublayers achieve  $97 - 99\%$  of their BERT counterparts' accuracy on GLUE, while still running  $40 - 70\%$  faster. This indicates that, while attention can improve accuracy, it may not be necessary to use in every layer.
- We demonstrate FNet scales very well to long inputs and offers a better compromise between speed and accuracy than the efficient Transformers evaluated on the Long-Range Arena (LRA) benchmark (Tay et al., 2021a). Specifically, FNet achieves accuracy comparable to the most accurate efficient Transformer architectures but is significantly faster at both
training and inference than all of the evaluated Transformer architectures across all sequence lengths on GPUs. On TPUs, FNet is faster for relatively shorter sequence lengths; for longer sequences, the only efficient Transformers that are faster than FNet on TPUs are less accurate on the LRA benchmark. Based on this, we argue that rather than seeking more efficient approximations of the attention, there may be more value in seeking out completely new mixing mechanisms.
2.1 Fourier Transforms in neural networks
Fourier analysis features heavily in studies of the universal approximation properties of neural networks; see, for example, (Cybenko, 1989; Barron, 1993). In terms of practical applications, discrete Fourier Transforms (DFT), and in particular the Fast Fourier Transform (FFT), have been used to tackle signal processing problems such as fitting neural networks to FFTs of electrocardiogram signals (Minami et al., 1999; Gothwal et al., 2011; Mironovova and Bila, 2015) and vibration signals (Zhang et al., 2013), or to evolve solutions of Partial Differential Equations (Li et al., 2021).
Because ordinary multiplication in the frequency domain corresponds to a convolution in the time domain, FFTs have been deployed in Convolutional Neural Networks to speed up computations, in Recurrent Neural Networks to speed up training and reduce exploding and vanishing gradients, and generally to approximate dense, linear layers to reduce computational complexity; see references cited in Section 1. DFTs have also been used indirectly in several Transformer works. The Performer (Choromanski et al., 2020) linearizes the Transformer self-attention mechanism by leveraging random Fourier features to approximate a Gaussian representation of the softmax kernel. In our work, rather than approximating attention, we replace attention with the Fourier Transform, which acts as an alternate hidden representation mixing mechanism. Tamkin et al. (2020) use spectral filters to generate hierarchical features, showing that the filtered embeddings perform well in different tasks (word-level, sentence-level or document-level), depending on which frequency scales are filtered. In contrast to FNet, they separate Fourier frequencies, rather than using the transform to combine features.
2.2 Modeling semantic relations via attention
Attention models have achieved state of the art results across virtually all NLP tasks and even some image tasks (Dosovitskiy et al., 2021). This success is generally attributed to the flexibility and capacity of attention. Although some works (Ramsauer et al., 2021) have endeavoured to gain a deeper understanding of attention, the pervasive intuition is that the success of attention models derives from the token-dependent attention patterns in different layers; see, for example, (Tenney et al., 2019). However, it is natural to ask: Do we really need the flexibility, and associated cost, of attention?
Tay et al. (2020a) empirically investigated the importance of the dot product operation in the attention mechanism in their Synthesizer model (related to our "Linear" baseline below). They find that learnt token-dependent attention weights are highly expressive, but not necessarily crucial for realizing accurate NLP models. You et al. (2020) replace attention weights in the Transformer encoder and decoder with unparameterized Gaussian distributions, showing minimal performance degradation provided they retain learnable cross-attention weights. Similarly, Raganato et al. (2020) find little to no accuracy degradation when replacing all but one of the attention heads of each attention layer in the encoder with fixed, non-learnable positional patterns. Finally, Tolstikhin et al. (2021) present MLP-Mixer, where attention is replaced by MLPs, with limited performance degradation in image classification tasks.
2.3 Efficient and long sequence models
The standard attention mechanism (Vaswani et al., 2017) has a quadratic time and memory bottleneck with respect to sequence length. This limits its applicability in tasks involving long range dependencies. Most efforts to improve attention efficiency are based on sparsifying the attention matrix. Tay et al. (2020c) survey many of the recent efficient attention works; see also citations in Section 1. Several "efficient Transformers" achieve  $\mathcal{O}(N\sqrt{N})$  or even  $\mathcal{O}(N)$  theoretical complexity. However, the constants hidden by this notation can be large. For example, in models such as Longformer (Beltagy et al., 2020), ETC (Ainslie et al., 2020), and BigBird (Zaheer et al., 2020), attention is  $\mathcal{O}(N)$  as a function of the input length, but quadratic in the number of "global tokens"; the latter must be sufficiently large to ensure good performance.
The Long-Range Arena benchmark (Tay et al., 2021a) attempts to compare many of the efficient Transformers in a series of tasks requiring long range dependencies, finding that the Performer (Choromanski et al., 2021), Linear Transformer (Katharopoulos et al., 2020), Linformer (Wang et al., 2020), and Image Transformer (Local Attention) (Parmar et al., 2018) were the fastest on TPUs and had the lowest peak memory usages per device. Instead, in this paper we completely replace self-attention with a different mixing, namely the Fourier Transform, which offers: (1) performance, (2) reduced model size (no learnable parameters), and (3) simplicity.
Finally, we note that, in an effort to investigate different token mixing mechanisms, we compare a vanilla BERT model (Devlin et al., 2019) with a vanilla FNet, ignoring more recent Transformer optimizations, which we consider orthogonal to this work; see, for example, (Narang et al., 2021; Kim and Hassan, 2020; Shleifer and Rush, 2020).
3.1 Discrete Fourier Transform
The Fourier Transform decomposes a function into its constituent frequencies. Given a sequence  $\{x_{n}\}$  with  $n\in [0,N - 1]$ , the discrete Fourier Transform (DFT) is defined by the formula:
$$
X _ {k} = \sum_ {n = 0} ^ {N - 1} x _ {n} e ^ {- \frac {2 \pi i}{N} n k}, \quad 0 \leq k \leq N - 1. \tag {1}
$$
For each  $k$ , the DFT generates a new representation  $X_{k}$  as a sum of all of the original input tokens  $x_{n}$  with so-called "twiddle factors". There are two primary approaches to computing the DFT: the Fast Fourier Transform (FFT) and matrix multiplication. The standard FFT algorithm is the Cooley-Tukey algorithm (Cooley and Tukey, 1965; Frigo and Johnson, 2005), which recursively re-expresses the DFT of a sequence of length  $N = N_{1}N_{2}$  in terms of  $N_{1}$  smaller DFTs of sizes  $N_{2}$  to reduce the computation time to  $\mathcal{O}(N\log N)$ .
An alternative approach is to simply apply the DFT matrix to the input sequence. The DFT matrix,  $W$ , is a Vandermonde matrix for the roots of unity up to a normalization factor:
$$
W _ {n k} = \left(e ^ {- \frac {2 \pi i}{N} n k} / \sqrt {N}\right), \tag {2}
$$
$^{2}$ Memory usage is often overlooked, but empirical studies have shown that Transformer architectures are often memory-bound (Ivanov et al., 2020; Shazeer, 2019).
Figure 1: FNet architecture with  $N$  encoder blocks.
where  $n, k = 0, \dots, N - 1$ . This matrix multiplication is an  $\mathcal{O}(N^2)$  operation, which has higher asymptotic complexity than the FFT, but turns out to be faster for relatively shorter sequences on TPUs.
3.2 FNet architecture
FNet is an attention-free Transformer architecture, wherein each layer consists of a Fourier mixing sublayer followed by a feed-forward sublayer. The architecture is shown in Figure 1. Essentially, we replace the self-attention sublayer of each Transformer encoder layer with a Fourier sublayer, which applies a 2D DFT to its (sequence length, hidden dimension) embedding input - one 1D DFT along the sequence dimension,  $\mathcal{F}_{\mathrm{seq}}$ , and one 1D DFT along the hidden dimension,  $\mathcal{F}_{\mathrm{h}}$ :<sup>3</sup>
$$
y = \Re \left(\mathcal {F} _ {\operatorname {s e q}} \left(\mathcal {F} _ {\mathrm {h}} (x)\right)\right). \tag {3}
$$
As indicated by Equation (3), we only keep the real part of the result; hence, we do not need to modify the (nonlinear) feed-forward sublayers or output layers to handle complex numbers. We found that FNet obtained the best results when the real part of the total transformation was only extracted at the end of the Fourier sublayer; that is, after applying both  $\mathcal{F}_{\mathrm{seq}}$  and  $\mathcal{F}_{\mathrm{h}}$ . We also experimented with the Hadamard, Hartley and Discrete Cosine Transforms. Of these three, the Hartley Transform was the strongest alternative, obtaining comparable
accuracy to Equation (3); see Appendix A.3 for details.
The simplest interpretation for the Fourier Transform is as a particularly effective mechanism for mixing tokens, which provides the feed-forward sublayers sufficient access to all tokens. Because of the duality of the Fourier Transform, we can also view each alternating encoder block as applying alternating Fourier and inverse Fourier Transforms, transforming the input back and forth between the "time" and frequency domain. Because multiplying by the feed-forward sublayer coefficients in the frequency domain is equivalent to convolving (with a related set of coefficients) in the time domain, FNet can be thought of as alternating between multiplications and convolutions.
We use the same embedding layers as in Devlin et al. (2019); namely, we combine the word embeddings, absolute position embeddings of the tokens and type embeddings of the sentences. Because of the positional information encoded by the Fourier Transform in Equation (1) (see  $n$ ,  $k$  indices), FNet performs just as well without position embeddings. Nevertheless, we include the position embeddings to allow for a cleaner comparison with BERT.
Empirically, we found that on GPUs: the FFT is faster than matrix multiplications for all sequence lengths we consider (512 - 8192 tokens), whereas on TPUs: for relatively shorter sequences ( $\leq$  4096 tokens), it is faster to cache the DFT matrix and then compute the DFT through matrix multiplications than using the FFT; for longer sequences, the FFT is faster. As a result, our GPU FNet implementation always uses the FFT, while our TPU implementation computes the 2D DFT using matrix multiplications for sequences up to lengths of 4096 and the FFT for longer lengths. Presumably the GPU vs TPU difference is primarily a result of two factors: (1) TPUs are even more highly optimized for matrix multiplications than GPUs, and (2) GPUs offer a more efficient FFT implementation than TPUs. We suspect that FNet will only become more performant on TPUs as the TPU implementation of the FFT improves. Our model uses JAX and, in particular, the Flax framework<sup>5</sup>. Core
4This is merely an intuition; the reality is more complicated due to the presence of residual connections and since the transformation in Equation (3) is no longer invertible if we only use the real component.
<sup>5</sup>https://github.com/google/flax
3The relative ordering of  $\mathcal{F}_{\mathrm{seq}}$  and  $\mathcal{F}_{\mathrm{h}}$  in Equation (3) is immaterial because the two 1D DFTs commute.
Table 1: Number of mixing layer operations (forward pass) and learnable parameters, excluding any task specific output projection layers.  $n$  is the sequence length and  ${d}_{h}$  is the model hidden dimension. The mixing layer operations are given on a per layer basis.
Model	Mixing layer ops (per layer)	Model	params
Base	Large
BERT	2n2dh + 4nd2h	112M	339M
Linear	n2dh + nd2h	94M	269M
FNet (mat)	n2dh + nd2h	83M	238M
FNet (FFT)	ndh log(n)+ ndh log(dh)	83M	238M
Random	n2dh + nd2h	83M	238M
FF-only	0	83M	238M
model code is given in Appendix A.7 and the full source core is available online. $^6$
4.1 Transfer learning
We compare FNet and Transformer architectures in a common transfer learning setting. For a fuller picture, we compare multiple models (see Table 1 for parameter counts in "Base" configuration):
- BERT-Base: a Transformer encoder model.
- FNet encoder: we replace every self-attention sublayer with a Fourier sublayer.
- Linear encoder: we replace each self-attention sublayer with a two learnable, dense, linear sublayers, one applied to the hidden dimension and one to the sequence dimension.
- Random encoder: we replace each self-attention sublayer with a two constant random matrices, one applied to the hidden dimension and one applied to the sequence dimension.
- Feed Forward-only (FF-only) encoder: we completely remove the self-attention sublayer, so that this model has no token mixing.
Despite its simplicity, the Linear baseline turns out to be surprisingly accurate and fast. Our Linear model is similar to the MLP-Mixer (Tolstikhin et al., 2021) (for vision) and also the Random Synthesizer (Tay et al., 2020a), but simplifies the latter model further by removing the multiple heads and
softmax projections, resulting in just two matrix multiplications in the mixing sublayer.
It is reasonable to expect that the Linear encoder, which uses densely parameterized mixing layers, will learn more flexibly than FNet, which uses parameter-free mixing layers. As we will show, although the Linear-Base model outperforms FNet-Base slightly on GLUE (0.3 points), it has several efficiency drawbacks relative to FNet: it has a much larger memory footprint (see Table 4b), it is slower to train on regular 512 sequence lengths (see Table 3), and scales significantly worse on long sequence lengths (see Tables 4b-4c). We also found that Linear-Large was more difficult to train due to gradient blow up (see "Large" scores in Table 2).
We adopt the same fixed "Base" and "Large" model and training configurations as for the original BERT (Devlin et al., 2019), except that we pretrain on the much larger C4 dataset (Raffel et al., 2020) and use a 32000 SentencePiece vocabulary model (Kudo and Richardson, 2018) (see Appendix A.1 for full pre-training details). For fine-tuning on the GLUE benchmark (Wang et al., 2018), we found that different BERT runs with the same base learning rate could yield slightly different results. Consequently, for the Base (Large) models, we performed 3 (6) trials, respectively, for each base learning rate and reported the best result across all experiments. This reflects our observation that BERT-Large was less stable than BERT-Base, as noted in Devlin et al. (2019).
We report the results for the best base learning rate (no early stopping) on the GLUE Validation split in Table 2. For Base models, results mirror the pre-training metrics (see Appendix A.1): BERT performs best. FNet and the Linear model both underperform BERT by  $7.5 - 8\%$ . Referring to Table 3, we see that although less accurate, FNet trains significantly faster than BERT  $-80\%$  faster on GPUs and  $70\%$  faster on TPUs - and performs  $63\%$  of BERT's FLOPS. Measured in isolation, the Fourier sublayers perform forward and backward passes an order of magnitude faster than the self-attention sublayers (see Appendix A.4), but FNet's overall training speed is impeded by the feed-forward sublayers that all models share.
On the other hand, the smaller sized Linear models do generally perform well on 512 sequence lengths; see Figure 2.
${}^{8}$  WNLI is excluded in Devlin et al. (2019). BERT's accuracy on WNLI is below baseline,unless a special training recipe is used. See also (12) in https://gluebenchmark.com/faq.
$^{6}$ https://github.com/google-research/google-research/tree/master/f_net
Table 2: GLUE Validation results on TPUs, after finetuning on respective tasks. We report the mean of accuracy and F1 scores for QQP and MRPC, Spearman correlations for STS-B and accuracy scores for all other tasks. The MNLI metrics are reported by the match/mismatch splits. Average scores exclude any failure cases. After controlling for batch size and training steps, the GPU metrics (not shown) are similar.
Model	MNLI	QQP	QNLI	SST-2	CoLA	STS-B	MRPC	RTE	Avg.
BERT-Base	84/81	87	91	93	73	89	83	69	83.3
Linear-Base	74/75	84	80	94	67	67	83	69	77.0
FNet-Base	72/73	83	80	95	69	79	76	63	76.7
Random-Base	51/50	70	61	76	67	4	73	57	56.6
FF-only-Base	34/35	31	52	48	67	FAIL	73	54	49.3
FNet-Hybrid-Base	78/79	85	88	94	76	86	79	60	80.6
BERT-Large	88/88	88	92	95	71	88	86	66	84.7
Linear-Large	35/36	84	80	79	67	24	73	60	59.8
FNet-Large	78/76	85	85	94	78	84	88	69	81.9
FNet-Hybrid-Large	79/80	87	89	92	81	88	86	70	83.6
Table 3: Pre-training and inference speeds in milliseconds per batch of 64 examples on GPU (8 V100 chips) and 256 examples on TPU  $(4\times 4\mathrm{v}3$  chips), alongside GFLOPS for a forward pass of a single example. Speed-up multipliers relative to BERT are given in parentheses.
Model	Pre-training	Inference	GFLOPS /example
GPU	TPU	GPU	TPU
BERT-Base	305	213	82	32	98
Linear-Base	199 (1.5x)	149 (1.4x)	52 (1.6x)	20 (1.6x)	71 (73%)
FNet-Base	169 (1.8x)	128 (1.7x)	46 (1.8x)	23 (1.4x)	62 (63%)
Random-Base	182 (1.7x)	130 (1.6x)	52 (1.6x)	22 (1.4x)	71 (73%)
FF-only-Base	162 (1.9x)	118 (1.8x)	43 (1.9x)	16 (2.0x)	59 (60%)
FNet-Hybrid-Base	198 (1.5x)	149 (1.4x)	51 (1.6x)	24 (1.3x)	68 (69%)
BERT-Large	OOM	503	263	111	337
Linear-Large	592	397 (1.3x)	170 (1.5x)	108 (1.0x)	247 (73%)
FNet-Large	511	275 (1.8x)	149 (1.8x)	82 (1.4x)	217 (64%)
FNet-Hybrid-Large	541	294 (1.7x)	157 (1.7x)	84 (1.3x)	227 (67%)
Returning to Table 2: the FF-only model severely underperforms all other models: as expected, token mixing is critical to the expressivity of the model. For example,  $50\%$  accuracy scores on the binary classification tasks (QNLI, SST-2, RTE), indicate that the model fails to learn the tasks. The weak accuracy of the Random model suggests that not just any mixing will do; rather, a structured mixing is required. We also include metrics from a hybrid FNet attention model. In the hybrid model, we replace the final two Fourier sublayers of FNet with self-attention sublayers - other configurations are possible, but we generally found that replacing the final layers worked best; see Appendix A.5. With the addition of just two self-attention sublayers, the hybrid FNet models achieve  $97\%$  and  $99\%$  of their respective BERT counterpart's accuracies with only limited speed degradations (see Table 3).