SlowGuess/NAACL_2022_OCR huggingface dataset:
Applying the sub-sampling-based methods requires the support for padding masks commonly used in Natural Language Processing tasks. However, a naive implementation of Algorithm 1 will result in the unnecessary sampling of the padding tokens. Therefore, we count the number of the unpadded
tokens  $m$ , and only perform the pilot sampling within the certain range  $[m]$ . After the matrix  $B_J$  is computed, we set its columns belonging to the padded part to be all zero, so that the probability  $\hat{p}_i$  of choosing column  $i$  from the padded part will be zero and the column will not be sampled in the later importance sampling. Similar modifications can also be applied to Informer to enable its applications in NLP tasks in Section 6.
4.5 Complexity Analysis
With the mild assumption in Lemma 1, we claim that our method can have an  $\mathcal{O}(n\log n)$  time and space complexity. The claim is shown by the following analysis of the complexity, which heavily relies on the notations in Algorithm 1.
First, we point out that the row/columnretrieving operation after index sampling can be implemented by only forming a view and thus the cost is negligible. For Line  $1\sim 4$  in Algorithm 1, the time complexity of the uniform pilot sampling is  $\mathcal{O}(d) = \mathcal{O}(\log n)$ , while the computation of the matrix  $B_{J}$  and the corresponding probabilities  $\{\hat{p}_i\}_{i = 1}^n$  costs  $\mathcal{O}(nd) = \mathcal{O}(n\log n)$  time and space. For Lines  $5\sim 7$ , with probabilities  $\{\hat{p}_i\}_{i = 1}^n$ , the importance sampling takes  $\mathcal{O}(n + d\log n) = \mathcal{O}(n)$  time, and similar to the computation above it takes  $\mathcal{O}(n\log n)$  time and space to obtain  $A_{J^{\prime}}$  and  $R_{J^{\prime}}$ . For Lines  $8\sim 10$ , it is clear that the three vectors  $\pmb {g},\pmb {d}$ , and  $\pmb{v}$  can be computed in  $\mathcal{O}(n\log n)$  time. As for the last step in Line 11, since it just requires the matrix product involving a diagonal matrix, we can finish the computation also in  $\mathcal{O}(n\log n)$  time and space. In summary, the total time and space complexity for Algorithm 1 is at most  $\mathcal{O}(n\log n)$ , much lower than the  $\mathcal{O}(n^{2})$  complexity for the original softmax self-attention.
Remark. The complexity above is derived based on the high probability bound in Proposition 1, which is different than the derivation by some previous methods to claim the linear  $\mathcal{O}(n)$  complexity.
5 Approximation Evaluation
As a preliminary justification of our proposed methods, we compute the spectral norm loss, a common metric for approximate matrix multiplication, to evaluate the effect of different models to approximate the original self-attention. We compare the spectral norms of the differences between the outputs from vanilla self-attention and other fast attention methods given the same input  $\mathbf{Q},\mathbf{K},\mathbf{V}$
Specifically we compute  $\| BV - R\| _2$  , where  $B\coloneqq D^{-1}A$  is the attention score matrix in the original method, and  $R$  is the output of each approximation method.
To construct the inputs  $\mathbf{Q},\mathbf{K},\mathbf{V}$ , we first truncate the raw text from Wikitext-2 dataset (Merit et al., 2017) into sequences of length 512. Then we transform the input  $X$  into  $\mathbf{Q},\mathbf{K},\mathbf{V}$  with the query, key, and value weight metrics from a pretrained model or a randomly initiated model.
We report the spectral norm loss of different sketching-based methods in Figure 1. The results are averaged over 768 trials, and the error bars in the figure represent the standard error of the reported values. For size  $d$  ( $x$ -axis), either the number of columns/rows selected or the projection dimension, it is set in the range from  $2^{3}$  to  $2^{8}$ .
V-Mean uses a rank-one matrix  $\frac{1}{n}\mathbf{11}^T V$  to approximate the original self-attention, and thus its approximation error does not change with the size  $d$ . V-Mean can also be seen as an ablation for the row normalization step (equivalent to adaptive row normalization without sub-samples). We observe the row normalization step greatly contributes to the approximation of self-attention that involves a softmax structure. Among the candidates, Skeinformer tends to have the smallest spectral norm loss, especially when  $d$  is large. It attains a higher accuracy than Informer and Linformer regarding the matrix approximation performance.
6.1 Benchmark
We experiment on Long Range Arena (LRA) benchmark (Tay et al., 2020a), including ListOps (Nangia and Bowman, 2018), Text Classification (Maas et al., 2011), Document Retrieval (Radev et al., 2013), Pathfinder (Linsley et al., 2018), and Image Classification (Krizhevsky et al., 2009). LRA is designed for long-context scenarios and more appropriate for evaluating efficient transformers comparing to GLUE (Mutton et al., 2007) with shorter input context. Following (Xiong et al., 2021) we use a 2-layer transformer model with 64 embedding dimensions, 128 hidden dimensions, and 2 attention heads for all experiments. More details can be found in Appendix.
6.2 Baseline Methods
We compare our method with the standard quadratic self-attention (Vaswani et al., 2017), Big
Bird (Zaheer et al., 2020), Linformer (Wang et al., 2020b), Informer (Zhou et al., 2020), Performer (Choromanski et al., 2020), and Nyströmformer (Xiong et al., 2021). In addition to their vanilla implementations, we compare with standard self-attention without dropout (since most fast attention methods do not have this component), Linformer with unreduced Johnson-Lindenstrauss Transform (the original form that Linformer deviates from), and Informer with padding masks.
Ablation studies include replacing the column sampling with uniform sampling, disabling the adaptive row normalization or replacing it with the simple row normalization implemented in Informer, and disabling the pilot sampling reutilization.
For clarification, deep transformers or pretrained language models are not appropriate baselines. Training a deep transformer from scratch requires large computational resources and much more data to converge, and therefore is not adopted by previous work. A shallow transformer structure, on the other hand, has been justified by previous work to be enough for fair comparison in attention acceleration performance. Pretrained models are trained for token-level text-based tasks, and are not suitable for image pixel sequences (as in Pathfinder and Image Classification), character sequences (as in Text Classification) and math operation sequences (as in ListOps).
We conclude the results in Table 1 and Table 2 with the following observations:
Most  $\widetilde{\mathcal{O}}(n)$  attention acceleration methods have comparable or better performance with vanilla attention. After all models converge to their long-time limits, Linformer tends to have worse performance possibly due to the violation of the sketching form, while Skeinformer has the best overall performance.
While surprising, those approximation methods tend to outperform the original transformer in most tasks. We speculate the reason behind this phenomenon is that a good approximation can recover the main signals in the original self-attention matrix, and also restrain the noise via the sparse / low-rank structure. A similar phenomenon can be found in CNN (Sanyal et al., 2018), that a low-rank regularizer, such as SVD, applied to the representation of the intermediate layers can allow the model to have lower prediction errors. This speculation
Models	Text	ListOps	Retrieval	Pathfinder	Image	Average
Standard (Vaswani et al., 2017)	57.69	38.15	80.10	73.59	37.97	57.50
· w/o dropout	59.44	38.17	79.35	72.35	37.58	57.38
V-Mean	65.29	28.78	80.49	61.01	34.33	53.98
BigBird (Zaheer et al., 2020)	61.91	38.86	79.73	71.75	35.00	57.45
Performer (Choromanski et al., 2020)	57.67	37.70	75.69	56.50	37.40	52.99
Nystromformer (Xiong et al., 2021)	60.91	37.76	79.87	72.53	31.93	56.60
Reformer (Kitaev et al., 2020a)	62.69	37.94	78.85	69.21	36.42	57.02
Linformer (Wang et al., 2020b)	58.52	37.97	77.40	55.57	37.48	53.39
· w/ unreduced JLT	59.12	37.48	79.39	68.45	35.96	56.08
Informer (Zhou et al., 2020)	61.55	38.43	80.88	59.34	36.55	55.35
· w/ padding mask	60.98	37.26	79.92	62.51	37.19	55.57
Skeinformer	62.47	38.73	80.42	71.51	37.27	58.08
· w/o column sampling	64.48	30.02	80.57	64.35	36.97	55.28
· w/o row normalization	60.67	37.69	78.67	66.35	37.06	56.09
· w/ simple row normalization	60.26	38.35	78.97	65.41	39.72	56.54
· w/o pilot sampling reutilization	62.39	38.12	79.88	71.53	37.20	57.83
Table 1: Classification accuracy (%) on the test sets of LRA benchmark. Skeinformer does not always outperform other baseline methods but has consistently comparable general performance. The approximation methods are not expected to outperform the original methods (standard self-attention) though they surprisingly do.
Models	Text	ListOps	Retrieval	Pathfinder	Image
time ↓	bz ↑	time ↓	bz ↑	time ↓	bz ↑	time ↓	bz ↑	time ↓	bz ↑
Standard	50.63	16	22.30	64	53.27	16	13.91	128	21.40	64
·w/o dropout	39.49	8	19.50	32	41.88	8	11.79	64	14.88	32
V-Mean	3.62	128	4.14	256	3.90	64	3.67	512	4.44	256
BigBird	20.59	64	17.28	128	21.73	32	17.83	256	18.84	256
Performer	2.63	64	9.31	128	12.50	32	10.40	256	8.94	256
Nyströmformer	12.18	64	12.28	128	13.35	32	19.58	128	10.30	256
Reformer	10.53	64	8.28	128	11.27	64	9.25	256	11.88	256
Linformer	7.91	64	6.25	128	8.08	64	6.90	256	6.65	256
·w/ unreduced JLT	36.87	8	21.49	32	35.93	4	15.17	128	22.03	128
Informer	33.13	16	21.89	32	36.52	16	26.14	64	24.92	128
·w/ mask	25.94	32	21.50	64	35.95	32	15.79	128	22.58	128
Skeinformer	9.60	64	9.66	128	10.61	64	9.25	256	11.86	256
·w/o column sampling	7.60	128	6.66	256	6.70	64	7.27	512	7.76	256
·w/o row normalization	25.02	16	16.02	64	55.72	4	11.12	256	15.52	128
·w/ simple row normalization	6.80	128	8.16	256	8.03	64	6.84	512	11.27	256
·w/o pilot sampling reutilization	7.15	128	7.31	256	8.68	64	7.09	512	10.19	256
Table 2: Training time (minute per thousand steps) and actual batch size (in batch accumulation) on LRA benchmark. Less training time per thousand steps indicates higher time efficiency. Higher batch size indicates higher space efficiency, and within a certain range means more accurate gradient estimations. we simulate the case of real-world applications of efficient transformers that models are trained with their maximum batch size conditioned on memory.
motivates us to turn to some theoretical framework for matrix approximation to better analyze the fast attention methods, which will potentially benefit transformer pruning, compression and distillation.
Skeinformer has the comparable general performance in terms of time/space complexity and classification accuracy. A long transformer is considered efficient when it (1) reduces space complexity and supports larger sequence length and larger batch size, (2) reduces time complexity with less training time per step and less total time to converge, and (3) shows comparable performance with vanilla softmax without much loss from approximation.
For convergence efficiency, Skeinformer effi
ciently converges to the long-time limit. Regarding the training efficiency, we focus on how soon the model can attain the stationary distribution of its long-time limit (He et al., 2019). The loss decay plot on ListOps in Appendix shows significant differences in the convergence rate of each method in addition to classification accuracy.
Though our method does not always outperform others (with the fastest convergence or the highest accuracy), but we remark that Skeinformer attains the best accuracy-efficiency trade-off based on experimental results. On the opposite, some model converges fast but gets stuck in a local optimum, like Linformer in some cases.
We conclude in this paper that sketching can be applied as a theoretical framework for analyzing fast attention models, through which we are able to recognize the potential improvements upon previous work. Theoretical results are provided to guarantee the high accuracy of the approximation to the original self-attention by our proposed Skeinformer. We empirically validate the contributions of the components in Skeinformer, including column sampling, adaptive row normalization and pilot sampling reutilization, with extensive comparisons with various baseline and ablation methods.
We thank the anonymous reviewers for their helpful suggestions. This research is based upon work supported by U.S. NSF grant DMS-1810831, DARPA AIDA Program No. FA8750-18-2-0014 and U.S. DARPA KAIROS Program No. FA8750-19-2-1004. The views and conclusions contained herein are those of the authors and should not be interpreted as necessarily representing the official policies, either expressed or implied, of DARPA, or the U.S. Government. The U.S. Government is authorized to reproduce and distribute reprints for governmental purposes notwithstanding any copyright annotation therein.
Nir Ailon and Bernard Chazelle. 2006. Approximate nearest neighbors and the fast johnson-lindenstrauss transform. In Proceedings of the 38th Annual ACM Symposium on Theory of Computing, Seattle, WA, USA, May 21-23, 2006, pages 557-563. ACM.
Guangsheng Bao, Yue Zhang, Zhiyang Teng, Boxing Chen, and Weihua Luo. 2021. G-transformer for document-level machine translation. In Proceedings of the 59th Annual Meeting of the Association for Computational Linguistics and the 11th International Joint Conference on Natural Language Processing, ACL/IJCNLP 2021, (Volume 1: Long Papers), Virtual Event, August 1-6, 2021, pages 3442-3455. Association for Computational Linguistics.
Iz Beltagy, Matthew E. Peters, and Arman Cohan. 2020. Longformer: The long-document transformer. CoRR, abs/2004.05150.
Stéphane Boucheron, Gábor Lugosi, and Pascal Massart. 2013. Concentration Inequalities - A Nonasymptotic Theory of Independence. Oxford University Press.
Yifan Chen and Yun Yang. 2021. Accumulations of projections - A unified framework for random sketches in kernel ridge regression. 130:2953-2961.
Yifan Chen, Qi Zeng, Heng Ji, and Yun Yang. 2021. Skyformer: Remodel self-attention with gaussian kernel and nystrom method. In Advances in Neural Information Processing Systems 35: Annual Conference on Neural Information Processing Systems 2021, NeurIPS 2021, December 6-14, 2021, virtual.
Rewon Child, Scott Gray, Alec Radford, and Ilya Sutskever. 2019. Generating long sequences with sparse transformers. CoRR, abs/1904.10509.
Krzysztof Choromanski, Valerii Likhosherstov, David Dohan, Xingyou Song, Andreea Gane, Tamás Sarlós, Peter Hawkins, Jared Davis, Afroz Mohiuddin, Lukasz Kaiser, David Belanger, Lucy Colwell, and Adrian Weller. 2020. Rethinking attention with performers. CoRR, abs/2009.14794.
Michael B. Cohen, Jelani Nelson, and David P. Woodruff. 2016. Optimal approximate matrix product in terms of stable rank. In 43rd International Colloquium on Automata, Languages, and Programming, ICALP 2016, July 11-15, 2016, Rome, Italy, volume 55 of LIPics, pages 11:1-11:14. Schloss Dagstuhl - Leibniz-Zentrum für Informatik.
Yihe Dong, Jean-Baptiste Cordonnier, and Andreas Loukas. 2021. Attention is not all you need: pure attention loses rank doubly exponentially with depth. 139:2793-2803.
Petros Drineas, Ravi Kannan, and Michael W. Mahoney. 2006. Fast monte carlo algorithms for matrices I: approximating matrix multiplication. SIAM J. Comput., 36(1):132-157.
Petros Drineas and Michael W. Mahoney. 2005. On the nystrom method for approximating a gram matrix for improved kernel-based learning. J. Mach. Learn. Res., 6:2153-2175.
Xinya Du and Claire Cardie. 2020. Document-level event role filler extraction using multi-granularity contextualized encoding. In Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics, ACL 2020, Online, July 5-10, 2020, pages 8010-8020. Association for Computational Linguistics.
Xinya Du, Sha Li, and Heng Ji. 2022. Dynamic global memory for document-level argument extraction. In Proceedings of the 60th Annual Meeting of the Association for Computational Linguistics.
Seth Ebner, Patrick Xia, Ryan Culkin, Kyle Rawlins, and Benjamin Van Durme. 2020. Multi-sentence argument linking. In Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics, ACL 2020, Online, July 5-10, 2020, pages 8057-8077. Association for Computational Linguistics.
Nathan Halko, Per-Gunnar Martinsson, and Joel A. Tropp. 2011. Finding structure with randomness: Probabilistic algorithms for constructing approximate matrix decompositions. SIAM Rev., 53(2):217-288.
Sariel Har-Peled, Piotr Indyk, and Rajeev Motwani. 2012. Approximate nearest neighbor: Towards removing the curse of dimensionality. Theory Comput., 8(1):321-350.
Fengxiang He, Tongliang Liu, and Dacheng Tao. 2019. Control batch size and learning rate to generalize well: Theoretical and empirical evidence. In Advances in Neural Information Processing Systems 32: Annual Conference on Neural Information Processing Systems 2019, NeurIPS 2019, December 8-14, 2019, Vancouver, BC, Canada, pages 1141-1150.
Luyang Huang, Shuyang Cao, Nikolaus Nova Parulian, Heng Ji, and Lu Wang. 2021. Efficient attentions for long document summarization. In Proceedings of the 2021 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, NAACL-HLT 2021, Online, June 6-11, 2021, pages 1419-1436. Association for Computational Linguistics.
William B Johnson and Joram Lindenstrauss. 1984. Extensions of lipschitz mappings into a hilbert space. Contemporary mathematics, 26(189-206):1.
Angelos Katharopoulos, Apoorv Vyas, Nikolaos Pappas, and François Fleuret. 2020. Transformers are rnns: Fast autoregressive transformers with linear attention. In Proceedings of the 37th International Conference on Machine Learning, ICML 2020, 13-18 July 2020, Virtual Event, volume 119 of Proceedings of Machine Learning Research, pages 5156-5165. PMLR.
Diederik P. Kingma and Jimmy Ba. 2015. Adam: A method for stochastic optimization. In 3rd International Conference on Learning Representations, ICLR 2015, San Diego, CA, USA, May 7-9, 2015, Conference Track Proceedings.
Nikita Kitaev, Lukasz Kaiser, and Anselm Levskaya. 2020a. Reformer: The efficient transformer. In 8th International Conference on Learning Representations, ICLR 2020, Addis Ababa, Ethiopia, April 26-30, 2020. OpenReview.net.
Nikita Kitaev, Lukasz Kaiser, and Anselm Levskaya. 2020b. Reformer: The efficient transformer. In 8th International Conference on Learning Representations, ICLR 2020, Addis Ababa, Ethiopia, April 26-30, 2020. OpenReview.net.
Alex Krizhevsky, Geoffrey Hinton, et al. 2009. Learning multiple layers of features from tiny images.
Ping Li, Trevor Hastie, and Kenneth Ward Church. 2006. Very sparse random projections. In Proceedings of the Twelfth ACM SIGKDD International Conference on Knowledge Discovery and Data Mining, Philadelphia, PA, USA, August 20-23, 2006, pages 287-296. ACM.
Sha Li, Heng Ji, and Jiawei Han. 2021. Document-level event argument extraction by conditional generation. pages 894-908.
Drew Linsley, Junkyung Kim, Vijay Veerabadran, Charles Windolf, and Thomas Serre. 2018. Learning long-range spatial dependencies with horizontal gated recurrent units. In Advances in Neural Information Processing Systems 31: Annual Conference on Neural Information Processing Systems 2018, NeurIPS 2018, December 3-8, 2018, Montreal, Canada, pages 152-164.
Yichao Lu, Paramveer S. Dhillon, Dean P. Foster, and Lyle H. Ungar. 2013. Faster ridge regression via the subsampled randomized hadamard transform. In Advances in Neural Information Processing Systems 26: 27th Annual Conference on Neural Information Processing Systems 2013. Proceedings of a meeting held December 5-8, 2013, Lake Tahoe, Nevada, United States, pages 369-377.
Andrew L. Maas, Raymond E. Daly, Peter T. Pham, Dan Huang, Andrew Y. Ng, and Christopher Potts. 2011. Learning word vectors for sentiment analysis. In The 49th Annual Meeting of the Association for Computational Linguistics: Human Language Technologies, Proceedings of the Conference, 19-24 June, 2011, Portland, Oregon, USA, pages 142-150. The Association for Computer Linguistics.
Stephen Merity, Caiming Xiong, James Bradbury, and Richard Socher. 2017. Pointer sentinel mixture models. In 5th International Conference on Learning Representations, ICLR 2017, Toulon, France, April 24-26, 2017, Conference Track Proceedings. Open-Review.net.
Andrew Mutton, Mark Dras, Stephen Wan, and Robert Dale. 2007. GLEU: automatic evaluation of sentence-level fluency. In ACL 2007, Proceedings of the 45th Annual Meeting of the Association for Computational Linguistics, June 23-30, 2007, Prague, Czech Republic. The Association for Computational Linguistics.
Nikita Nangia and Samuel R. Bowman. 2018. Listops: A diagnostic dataset for latent tree learning. In Proceedings of the 2018 Conference of the North American Chapter of the Association for Computational Linguistics, NAACL-HLT 2018, New Orleans, Louisiana, USA, June 2-4, 2018, Student Research Workshop, pages 92-99. Association for Computational Linguistics.
Jiezhong Qiu, Hao Ma, Omer Levy, Wen-tau Yih, Sinong Wang, and Jie Tang. 2020. Blockwise self-attention for long document understanding. In Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing: Findings, EMNLP 2020, Online Event, 16-20 November 2020, pages 2555-2565. Association for Computational Linguistics.
Dragomir R. Radev, Pradeep Muthukrishnan, Vahed Qazvinian, and Amjad Abu-Jbara. 2013. The ACL
anthology network corpus. Lang. Resour. Evaluation, 47:919-944.
Ali Rahimi and Benjamin Recht. 2007. Random features for large-scale kernel machines. In Advances in Neural Information Processing Systems 20, Proceedings of the Twenty-First Annual Conference on Neural Information Processing Systems, Vancouver, British Columbia, Canada, December 3-6, 2007, pages 1177-1184. Curran Associates, Inc.
Aurko Roy, Mohammad Saffar, Ashish Vaswani, and David Grangier. 2021. Efficient content-based sparse attention with routing transformers. Trans. Assoc. Comput. Linguistics, 9:53-68.
Amartya Sanyal, Varun Kanade, Philip HS Torr, and Puneet K Dokania. 2018. Robustness via deep low-rank representations. arXiv preprint arXiv:1804.07090.
Yi Tay, Mostafa Dehghani, Samira Abnar, Yikang Shen, Dara Bahri, Philip Pham, Jinfeng Rao, Liu Yang, Sebastian Ruder, and Donald Metzler. 2020a. Long range arena: A benchmark for efficient transformers. CoRR, abs/2011.04006.
Yi Tay, Mostafa Dehghani, Dara Bahri, and Donald Metzler. 2020b. Efficient transformers: A survey. CoRR, abs/2009.06732.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. 2017. Attention is all you need. In Advances in Neural Information Processing Systems 30: Annual Conference on Neural Information Processing Systems 2017, December 4-9, 2017, Long Beach, CA, USA, pages 5998-6008.
Roman Vershynin. 2010. Introduction to the non-asymptotic analysis of random matrices. CoRR, abs/1011.3027.
Martin J Wainwright. 2019. High-dimensional statistics: A non-asymptotic viewpoint, volume 48. Cambridge University Press.
Shuohang Wang, Luowei Zhou, Zhe Gan, Yen-Chun Chen, Yuwei Fang, Siqi Sun, Yu Cheng, and Jingjing Liu. 2020a. Cluster-former: Clustering-based sparse transformer for long-range dependency encoding. CoRR, abs/2009.06097.
Sinong Wang, Belinda Z. Li, Madian Khabsa, Han Fang, and Hao Ma. 2020b. Linformer: Self-attention with linear complexity. CoRR, abs/2006.04768.
Christopher K. I. Williams and Matthias W. Seeger. 2000. Using the nyström method to speed up kernel machines. In Advances in Neural Information Processing Systems 13, Papers from Neural Information Processing Systems (NIPS) 2000, Denver, CO, USA, pages 682-688. MIT Press.
David P. Woodruff. 2014. Sketching as a tool for numerical linear algebra. Found. Trends Theor. Comput. Sci., 10(1-2):1-157.