SlowGuess/NAACL_2022_OCR huggingface dataset:
Long range arena: A benchmark for efficient transformers. In 9th International Conference on Learning Representations, ICLR 2021, Virtual Event, Austria, May 3-7, 2021.
Yi Tay, Mostafa Dehghani, Dara Bahri, and Donald Metzler. 2020c. Efficient transformers: A survey. arXiv preprint arXiv:2009.06732.
Yi Tay, Mostafa Dehghani, Jai Prakash Gupta, Vamsi Aribandi, Dara Bahri, Zhen Qin, and Donald Metzler. 2021b. Are pretrained convolutions better than pretrained transformers? In Proceedings of the 59th Annual Meeting of the Association for Computational Linguistics and the 11th International Joint Conference on Natural Language Processing (Volume 1: Long Papers), pages 4349-4359. Association for Computational Linguistics.
Ian Tenney, Dipanjan Das, and Ellie Pavlick. 2019. BERT rediscovers the classical NLP pipeline. In Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics, pages 4593-4601, Florence, Italy. Association for Computational Linguistics.
Ilya O Tolstikhin, Neil Houlsby, Alexander Kolesnikov, Lucas Beyer, Xiaohua Zhai, Thomas Unterthiner, Jessica Yung, Andreas Steiner, Daniel Keysers, Jakob Uszkoreit, et al. 2021. Mlp-mixer: An all-mlp architecture for vision. Advances in Neural Information Processing Systems, 34.
Iulia Turc, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. 2019. Well-read students learn better: On the importance of pre-training compact models. arXiv preprint arXiv:1908.08962.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. 2017. Attention is all you need. Advances in neural information processing systems, 30.
Jesse Vig and Yonatan Belinkov. 2019. Analyzing the structure of attention in a transformer language model. In Proceedings of the 2019 ACL Workshop BlackboxNLP: Analyzing and Interpreting Neural Networks for NLP, pages 63-76, Florence, Italy. Association for Computational Linguistics.
Elena Voita, David Talbot, Fedor Moiseev, Rico Senrich, and Ivan Titov. 2019. Analyzing multi-head self-attention: Specialized heads do the heavy lifting, the rest can be pruned. In Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics, pages 5797-5808, Florence, Italy. Association for Computational Linguistics.
Apoory Vyas, Angelos Katharopoulos, and François Fleuret. 2020. Fast transformers with clustered attention. Advances in Neural Information Processing Systems, 33:21665-21674.
Alex Wang, Amanpreet Singh, Julian Michael, Felix Hill, Omer Levy, and Samuel Bowman. 2018. GLUE: A multi-task benchmark and analysis platform for natural language understanding. In Proceedings of the 2018 EMNLP Workshop BlackboxNLP: Analyzing and Interpreting Neural Networks for NLP, pages 353-355, Brussels, Belgium. Association for Computational Linguistics.
Sinong Wang, Belinda Li, Madian Khabsa, Han Fang, and Hao Ma. 2020. Linformer: Self-attention with linear complexity. arXiv preprint arXiv:2006.04768.
Weiqiu You, Simeng Sun, and Mohit Iyyer. 2020. Hard-coded Gaussian attention for neural machine translation. In Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics, pages 7689-7700. Association for Computational Linguistics.
Manzil Zaheer, Guru Guruganesh, Kumar Avinava Dubey, Joshua Ainslie, Chris Alberti, Santiago Otonan, Philip Pham, Anirudh Ravula, Qifan Wang, Li Yang, et al. 2020. Big bird: Transformers for longer sequences. Advances in Neural Information Processing Systems, 33:17283-17297.
Jiong Zhang, Yibo Lin, Zhao Song, and Inderjit Dhillon. 2018. Learning long term dependencies via fourier recurrent units. In International Conference on Machine Learning, pages 5815-5823. PMLR.
Y Zhang and Lai-Wan Chan. 2000. Forenet: fourier recurrent networks for time series prediction. In 7th International Conference on Neural Information Processing (ICONIP 2000), pages 576-582.
Zhenyou Zhang, Yi Wang, and Kesheng Wang. 2013. Fault diagnosis and prognosis using wavelet packet decomposition, fourier transform and artificial neural network. Journal of Intelligent Manufacturing, 24(6):1213-1227.
Table 5: Loss and accuracy pre-training metrics on TPUs. The GPU metrics are very similar. "B" denotes Base, "L" is Large and "H" is Hybrid.
Model	Total	Loss	Accuracy
MLM	NSP	MLM	NSP
BERT-B	1.76	1.48	0.28	0.68	0.86
Linear-B	2.12	1.78	0.35	0.62	0.83
FNet-B	2.45	2.06	0.40	0.58	0.80
Random-B	5.02	4.48	0.55	0.26	0.70
FF-only-B	7.54	6.85	0.69	0.13	0.50
FNet-H-B	2.13	1.79	0.34	0.63	0.84
BERT-L	1.49	1.23	0.25	0.72	0.88
Linear-L	1.91	1.60	0.31	0.65	0.85
FNet-L	2.11	1.75	0.36	0.63	0.82
FNet-H-L	1.89	1.58	0.31	0.67	0.85
A.1 Pre-training details
We adopt the same fixed "Base" and "Large" model and learning configurations as for the original BERT (Devlin et al., 2019). We train on the much larger C4 dataset (Raffel et al., 2020) and use a 32000 SentencePiece vocabulary model (Kudo and Richardson, 2018) trained on a 100 million sentence subset of C4. Our TPU experiments use a batch size of 256 as in Devlin et al. (2019) and are each run on  $4 \times 4$  TPU v3 chips. Our GPU experiments use a smaller batch size of 64 and are run on 8 V100 chips. Because the training configuration is lifted from Devlin et al. (2019), it may be slightly biased towards the BERT attention model.
Table 5 summarizes the pre-training metrics for the different models; the pre-training speeds are shown in Table 3 in the main text. Although they have weaker accuracy metrics, the Linear model and FNet train nearly  $80\%$  faster than BERT on GPUs, and  $70\%$  faster on TPU's (see Table 3). We also find that the three models with no learnable parameters in their mixing layer, namely FNet, the Random model and the FF-only model, are the most stable during training.
BERT's higher accuracy on the MLM pretraining task is not simply a result of having more parameters than the other models. Indeed, Table 5 shows that BERT-Base is actually more accurate than FNet-Large, which contains more than twice as many parameters. BERT is presumably more expressive because the mixing (attention) weights are both task specific and token dependent, determined
Table 6: Pre-training model sizes (ignoring output projection layers). As in Turc et al. (2019), for all models, we fix the feed-forward size to  $4d_{h}$  and the number of self-attention heads to  $d_{h} / 64$ . Smaller architectures have a similar number of parameters across all models because the majority of parameters are in the embedding layers. Each FNet-Hybrid ("FNet-H") model contains 2 self-attention sublayers. We exclude FNet-Hybrid models with only 2 total layers.
Dimensions dhLayers	Parameters (millions)
BERT	Linear	FNet	FNet-H
768	12	111	93	83	88
512	12	55	49	42	44
512	8	42	38	34	36
256	8	15	15	13	13
512	4	30	28	26	28
256	4	12	12	11	11
256	2	10	10	10	-
128	2	5	5	4	-
by token-token (query-key) dot products; see also Tay et al. (2020a). FNet's mixing weights, on the other hand, are neither task specific nor token dependent.
Finally, Table 6 shows the model sizes that were used to construct Figure 2 (main text) and Figure 3 (Appendix A.2).
A.2 TPU results
In this section, we report FNet efficiency results for TPUs; the main text focuses on GPUs. Figure 3 shows the speed vs MLM pre-training accuracy curve when training on TPU (4 × 4 v3 chips). As on GPUs, FNet and the Linear model define the Pareto efficiency frontier for smaller, faster models, while BERT defines the frontier for larger, slower models.
Table 7 shows Long Range Arena Text classification efficiency results on TPUs (4 × 4 v3 chips). The Linear model and FNet train faster than all the efficient Transformers for sequence lengths ≤ 2048 and 512, respectively. For longer sequences, FNet is slower than the Performer and, based on results in Tay et al. (2021a), likely also slower than the other efficient Transformers that linearize attention, namely Local Attention (Parmar et al., 2018), Linformer (Wang et al., 2020) and Linear Transformer (Katharopoulos et al., 2020). However, it is worth noting that Table 4a suggests that FNet is more accurate than all of the aforementioned models. Moreover, we expect that the GPU speed gains will
Figure 3: Speed-accuracy trade-offs for TPU pre-training. The dashed line shows the Pareto efficiency frontier.
Table 7: TPU training speeds (in steps per second; larger is better), inference speeds (in milliseconds per batch; smaller is better) and peak memory usage during training (in GB; smaller is better) on the Long-Range Arena Text classification task. Speed up multipliers relative to the Transformer are given in parentheses.
Seq. length	512	1024	2048	4096	8192	16386
Training Speed (steps/s)
Transformer	8.0	5.6	1.7	OOM	OOM	OOM
Linear	9.4 (1.2x)	9.1 (1.6x)	7.6 (4.5x)	3.9	1.4	OOM
FNet (mat)	9.5 (1.2x)	9.1 (1.6x)	6.1 (3.6x)	3.0	0.8	0.2
FNet (FFT)	8.6 (1.1x)	6.0 (1.1x)	3.2 (1.9x)	1.6	0.8	0.3
Performer	9.2 (1.2x)	8.4 (1.5x)	6.9 (4.1x)	4.2	2.2	1.1
Inference Speed (ms/batch)
Transformer	7.0	13.2	39.4	129.9	490.2	OOM
Linear	5.6 (1.2x)	6.5 (2.0x)	9.6 (4.1x)	20.4 (6.4x)	54.6 (9.0x)	OOM
FNet (mat)	6.0 (1.2x)	7.7 (1.7x)	15.4 (2.6x)	40.7 (3.2x)	137.0 (3.6x)	454.5
FNet (FFT)	10.8 (0.7x)	16.8 (0.8x)	29.9 (1.3x)	58.8 (2.2x)	113.6 (4.3x)	263.2
Performer	6.1 (1.2x)	7.2 (1.8x)	10.1 (3.9x)	17.5 (7.4x)	31.8 (15.4x)	61.0
Peak Memory Usage (GB)
Transformer	1.1	2.1	5.8	9.1	OOM	OOM
Linear	0.9	1.1	1.9	4.9	14.8	OOM
FNet (mat)	0.8	0.9	1.3	2.2	4.8	11.9
FNet (FFT)	0.8	0.9	1.3	2.0	3.5	6.3
Performer	1.0	1.3	1.8	3.0	5.1	9.6
transfer to TPUs as the TPU FFT implementation improves.
A.3 Additional configurations that we experimented with
We experimented with a number of additional ideas to improve FNet.
Fourier Transform algorithm. On GPUs, the FFT was the fastest algorithm for computing the DFT across all sequence lengths that we experimented with (512 - 8192). On TPUs, it is faster to compute the DFT directly using matrix multiplications for relatively shorter sequence lengths (up to lengths of 4096; see Table 7). This efficiency boundary between matrix multiplication and FFT on TPUs will change depending on the XLA precision for the matrix multiplications. We found that, although (slower) HIGHEST XLA precision was required to very accurately reproduce FFT in computing the DFT, (faster) DEFAULT XLA precision was sufficient to facilitate accurate model convergence.
Modifying the Fourier Transform computation. To keep the entire FNet architecture simple, the Fourier sublayer accepts real input and returns real output. The standard Fourier sublayer in FNet simply extracts the real part after computing the 2D DFT. We found that FNet was less accurate and less stable during training if only the real part of the DFT was used throughout the computation. Simply extracting the absolute value (instead of the real part) also led to a significantly less accurate model. Because the feed-forward sublayer mixes the hidden dimension, we experimented with applying a 1D DFT along the token dimension only in the Fourier sublayer (i.e. no hidden dimension mixing in the Fourier sublayer). This yielded some training speed gains but hurt accuracy. The 1D (token mixing only) DFT model still significantly outperformed the (no token mixing) FF-only model, indicating that token mixing is most important mechanism in the Fourier sublayer.
Other transforms. We experimented with three natural alternatives to the Fourier Transform:
- Discrete Cosine Transform (DCT). The DCT is closely related to the DFT but transforms real input to real output. However, we found that the DCT model underperformed FNet ( $\sim$  4% accuracy degradation).
- Hadamard Transform<sup>12</sup>. Although the Hadamard Transform was slightly faster than the DFT, it yielded less accurate results ( $\sim 2\%$  accuracy degradation).
- Hartley Transform. The Hartley Transform, which transforms real input to real output, can be described in terms of the Fourier Transform:  $\mathcal{H} = \Re \left\{\mathcal{F}\right\} -\Im \left\{\mathcal{F}\right\}$ . We found that the Hartley Transform matched the Fourier Transform on GLUE (76.7 vs. 76.7).
Introducing learnable parameters to the Fourier sublayer. Our attempts to introduce learnable parameters into the Fourier sublayer were either detrimental or inconsequential, and generally slightly slowed the model. For the (sequence length, hidden dimension) input in each Fourier sublayer, we tried two approaches to introduce learnable parameters: (1) element wise multiplication with a (sequence length, hidden dimension) matrix, and (2) regular matrix multiplication with (sequence length, sequence length) and (hidden dimension, hidden dimension) matrices. We experimented with these approaches in various configurations: preceding and/or following the DFT, and also in combination with inverse DFT (e.g. transform to frequency domain, apply element wise multiplication, transform back to time domain), but most setups degraded accuracy and reduced training stability, while a few did not change accuracy but lead to small speed decreases. In a slightly different set of experiments and in an effort to provide more flexibility to the model, we added (complex) learnable weights to the 2D DFT matrix. This model was stable but did not yield any accuracy gains, suggesting that the DFT is locally optimal in some sense.
FNet block modifications. The standard FNet encoder block structure follows that of the Transformer: a Fourier sublayer followed by a feedforward sublayer, with residual connections and layer norms after each sublayer; see Figure 1. We tried several modifications to this structure, based on the intuition of moving in and out of the frequency domain between multiplications. For example, the sandwiching of Fourier, feed-forward, Fourier (or inverse Fourier) sublayers and only applying the residual connections and layer norms to the final result, yields a structure that more closely
<sup>12</sup>Whereas the DFT matrix in Equation (2) contains  $N$  roots of unity, the Hadamard Transform simply contains two roots of unity:  $\{\pm 1\}$ ; see also Kunz (1979).
Table 8: Training (forward and backward passes; left) and inference (forward pass; left) speeds for only the mixing sublayers – all other model sublayers are removed. Both speeds are measured in milliseconds per batch (smaller is better), with batch sizes of 64 (GPU) and 256 (TPU). All batch examples have the sequence length fixed at 512. FNet uses the FFT for GPUs and matrix multiplications for TPUs. Speed up multipliers relative to self-attention are given in parentheses.
Training speed (ms/batch)	Inference speed (ms/batch)
GPU	TPU	GPU	TPU
Self-attention (Base)	136	76	43	16
Linear (Base)	36 (3.7x)	12 (6.1x)	15 (2.8x)	4 (3.9x)
FNet (Base)	11 (12.2x)	8 (9.9x)	11 (4.0x)	8 (2.1x)
Self-attention (Large)	404	212	128	43
Linear (Large)	103 (3.9x)	35 (6.1x)	36 (3.6x)	10 (4.5x)
FNet (Large)	18 (22.2x)	22 (9.7x)	18 (7.3x)	22 (2.0x)
mimics convolutions. However, these setups degraded accuracy and lead to a more unstable model during training. Adding extra feed-forward sublayers to this layering, or swapping out the feedforward sublayers for simpler dense sublayers, did not help either.
A.4 Mixing layer speeds
Table 8 summarizes the inference and training speeds for the different mixing layers. For each of the Base and Large configurations, we have removed all other sublayers and transformations and then calculated the speed per batch of input examples. The FNet training speeds are particularly fast because no parameters are updated. The Linear model has faster inference than FNet on TPUs because it is performing real matrix multiplications, whereas FNet performs complex matrix multiplications; see Equation (2).
Although the Fourier mixing sublayer itself performs forward and backward passes significantly faster than the self-attention sublayer, FNet is overall  $70 - 80\%$  faster than BERT because the overall training and inference speeds are bottle-necked by the feed-forward sublayers that all models share.
A.5 FNet-Hybrid ablations
Table 9 shows the effects of varying the number of attention sublayers and the attention layout in the FNet-Hybrid model. For the "BOTTOM" layout, all attention sublayers are placed in the first few encoder layers, where they replace the Fourier mixing sublayers. For the "TOP" layout, attention sublayers are placed in the final encoder layers; for the "MIDDLE" layout they are placed in the middle layers; and for the "MIXED" layout, they are distributed through the model.
Table 9: GPU pre-training accuracy and speed ablations for FNet-Hybrid models in the Base configuration. Batch size is 64. Metrics are recorded after 100k steps, which we have generally found to be a good indicator of final relative performance. See text for a description of the layouts.
Attention	Accuracy	Speed (ms/batch)
Layers	Layout	MLM	NSP
2	BOTTOM	0.497	0.733	193
2	MIDDLE	0.499	0.686	196
2	MIXED	0.509	0.727	194
2	TOP	0.526	0.738	193
0	TOP	0.486	0.679	173
2	TOP	0.526	0.738	193
4	TOP	0.539	0.740	214
6	TOP	0.546	0.746	235
From the Table 9, we can make two observations: (1) more attention improves accuracy at the cost of speed, and ultimately with diminishing returns; (2) placing attention layers at the top of the model gives the best accuracy results. Given our focus on speed, we chose to focus FNet-Hybrid experiments in the main text of the paper on the 2 attention layer, "TOP" configuration variant.
A.6 A note on Long-Range Arena hyperparameter settings
Concerning the Long-Range Arena setup, several hyperparameters are not described in Tay et al. (2021a) and there a few mismatches between the configurations described in the paper and the code repository. Where possible, we prioritize configurations described in the paper with only two exceptions. Firstly, for the CIFAR10 (Image) task, we perform a sweep of the number of layers in the
range [1, 2, 3, 4]. We found that 1 layer worked best for all models; Tay et al. (2021a) suggest 3 layers yielded the best results. Secondly, for the Pathfinder task, we found that a base learning rate of 0.001 (as given in the code repository) yielded better results for all models than the 0.01 value indicated in Tay et al. (2021a). We also perform a very small sweep over the embedding dimension and batch size, which are not listed in Tay et al. (2021a).
We also remark that the accuracy comparisons between our runs and those from Tay et al. (2021a) should be performed with the caveat that we found that results for certain tasks – Text and Retrieval in particular – can vary quite a bit between runs, especially for the Transformer; we report the best results.
A.7 FNet code
Listing 1: FNet code written in JAX/Flax. Embedding and output projection layers are omitted for simplicity.
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66