SlowGuess/NAACL_2022_OCR huggingface dataset:
Yifan Chen $^{1*}$ , Qi Zeng $^{1*}$ , Dilek Hakkani-Tur $^{2}$ , Di Jin $^{2}$ , Heng Ji $^{1}$ , Yun Yang $^{1}$
$^{1}\{yifanc10, qizeng2, hengji, yy84\} @illinois.edu$
$^{2}\{hakkanit, djinamzn\} @amazon.com$
Transformer-based models are not efficient in processing long sequences due to the quadratic space and time complexity of the self-attention modules. To address this limitation, Linformer and Informer reduce the quadratic complexity to linear (modulo logarithmic factors) via low-dimensional projection and row selection, respectively. These two models are intrinsically connected, and to understand their connection we introduce a theoretical framework of matrix sketching. Based on the theoretical analysis, we propose Skeinformer to accelerate self-attention and further improve the accuracy of matrix approximation to self-attention with column sampling, adaptive row normalization and pilot sampling reutilization. Experiments on the Long Range Arena benchmark demonstrate that our methods outperform alternatives with a consistently smaller time/space footprint<sup>1</sup>.
Transformer (Vaswani et al., 2017) utilizes softmax self-attention modules to capture the dependency between tokens in a sequence and has been widely used in various Natural Language Processing tasks. The time and space complexity of the dot-product self-attention is quadratic in the input sequence length, which restricts the largest sequence length and batch size. To adapt transformers to long sequences, documents have to be truncated, chunked using a sliding window, or processed in parallel on multiple GPUs. These additional operations usually cause the loss of long-range dependency and introduce additional computational costs.
In this paper, we focus on efficient self-attention methods (Xiong et al., 2021; Qiu et al., 2020; Zaheer et al., 2020; Beltagy et al., 2020; Kitaev
et al., 2020a; Roy et al., 2021), among which Linformer (Wang et al., 2020b) and Informer (Zhou et al., 2020) are two representative approaches to reducing the  $\mathcal{O}(n^2)$  self-attention to an  $\widetilde{\mathcal{O}}(n)$  operation ( $\widetilde{\mathcal{O}}(\cdot)$  means  $\mathcal{O}(\cdot)$  modulo poly-log terms and  $n$  is the sequence length) in both space and time complexity. Linformer forms a low-rank factorization of the original attention by decomposing it into smaller attentions, while Informer allows each key to only attend to a certain number of queries.
To better understand self-attention, we introduce a theoretical framework, sketching (Woodruff, 2014), to help explain the key ideas in Informer and Linformer from the perspective of matrix approximation. Specifically, sketching methods replace the original matrix  $\mathbf{B}$  with its random sketch  $\mathbf{BS}$  to reduce computations. In Section 3.3 we introduce some concrete instances of commonly used distributions for constructing the random sketching matrix  $\mathbf{S}$ . Furthermore, taking matrix approximation as a general guideline, we recognize the deficiency in Informer and Linformer, that they either do not fully utilize the information in the value matrix  $\mathbf{V}$ , or deviate from the original self-attention output. This guideline also motivates us to propose Skeinformer through the theoretical analysis under the sketching framework.
To improve the approximation accuracy in terms of the original attention output, Skeinformer applies sub-sampling sketching to reduce time complexity and exploits the information from the value matrix  $V$  with column sampling. Skeinformer also incorporates an adaptive row normalization step, which approximates the un-selected rows by a vector with all elements  $\frac{1}{n}$  and has significantly boosted the performance of Informer. In addition, we introduce a simple yet effective step, pilot sampling reutilization, which reuses the computation from pilot sampling to improve both approximation accuracy and training efficiency. Our experiments on the LRA benchmark show that Skeinformer con
*Equal contribution. The majority of this work was done prior to the first author's internship at Amazon Alexa AI.  
<sup>1</sup>Our code is released at https://github.com/ pkuzengqi/Skeinformer
Human Language Technologies, pages 5187 - 5199
sistently uses less space and time while achieving better accuracy than most baseline methods.
In summary, our contributions are twofold:
- We introduce sketching as a theoretical framework for analyzing and developing efficient transformers. Specifically, the randomized sketching theory covers these two methods from the perspective of approximate matrix multiplication. This framework connects the studies on efficient transformers and randomized sketching theory, so that future development in efficient transformers and sketching can benefit each other.
- We propose Skeinformer as a straightforward product of the sketching framework to accelerate the training and inference of transformers. Skeinformer consists of three components: the initial column sampling that incorporates the information from the value matrix  $\mathbf{V}$  into the sampling probabilities, the adaptive row normalization that fills un-selected columns with the averaged selected columns, and the pilot sampling re-utilization.
The ability to process long sequences is critical for many Natural Language Processing tasks, including Document Summarization (Xiao and Carenini, 2019; Huang et al., 2021), Question Answering (Wang et al., 2020a), Information Extraction (Li et al., 2021; Du and Cardie, 2020; Ebner et al., 2020; Du et al., 2022), and Machine Translation (Bao et al., 2021). However, the quadratic computational cost of self-attention in transformer-based models limits their application in long-sequence tasks. Recent methods have been proposed to accelerate attention computation by selectively attending to a subset of the tokens or with low-rank matrix approximation (Tay et al., 2020b).
Selective attention methods limit the scope of matrix operation with sparse attention patterns or column/row sampling methods. BlockBERT (Qiu et al., 2020) introduces sparse block structures into the attention matrix. Sparse Transformer (Child et al., 2019) introduces dilated patterns. Big Bird (Zaheer et al., 2020) proposes a combination of random, window, and global attention. Longformer (Beltagy et al., 2020) combines local windowed attention with task-motivated global attention. The most related work to ours is In
former (Zhou et al., 2020), which allows each key to only attend to the top queries under the Kullback-Leibler divergence based sparsity measurement.
Low-rank attention matrix approximation methods are based on the assumption of low-rank structure in the full self-attention matrix. Linformer (Wang et al., 2020b) compresses the size of the key and value matrices by the Johnson-Lindenstrauss transform (Johnson and Lindenstrauss, 1984). Performer (Choromanski et al., 2020) recognizes the attention score matrix as an empirical Gaussian kernel matrix and constructs a low-rank projection for both the query and key matrices through random Fourier features (Rahimi and Recht, 2007). Nyströmformer (Xiong et al., 2021) instead utilizes Nyström method (Williams and Seeger, 2000; Drineas and Mahoney, 2005) to approximate the attention score matrix. Skyformer (Chen et al., 2021) replaces the softmax structure with a Gaussian kernel and adapts the Nyström method to accelerate the computation.
Some other methods follow a similar principle to decompose the attention score matrix, although they are not necessarily aiming to approximate the original self-attention itself. The representative methods include Linear Transformer (Katharopoulos et al., 2020), which claims that the exponential transform of the dot-product in the softmax operation can be replaced by the direct matrix multiplication of the projected query and key matrices, and Reformer (Kitaev et al., 2020b), which forces the query and key matrices to be identical and applies locality-sensitive hashing (LSH) (Har-Peled et al., 2012) to simplify the computation of the attention score matrix. Those methods are effective alternatives of the original self-attention, while they do not fall into the scope of matrix approximation. We spare the discussion of those methods in this paper.
3 Sketching Framework
Given an input sequence  $X\in \mathbb{R}^{n\times d_{input}}$  , where  $n$  is the sequence length and  $d_{input}$  is the embedding dimension, the dot-product attention for a single attention head in transformer (Vaswani et al., 2017) is defined as
$$
\operatorname {A t t e n t i o n} (Q, K, V) = \operatorname {s o f t m a x} \left(\frac {Q K ^ {T}}{\sqrt {p}}\right) V
$$
where  $Q = XW_{Q}$ ,  $\pmb{K} = X\pmb{W}_{K}$ , and  $\pmb{V} = X\pmb{W}_{V}$ .  $W_{Q}, W_{K}, W_{V} \in \mathbb{R}^{d_{input} \times p}$  are the query,
Algorithm 1: Skeinformer.
key, and value weight metrics that linearly project the input  $X$  of dimension  $d_{input}$  to an output tensor of dimension  $p$ .
To ease the future analysis, the softmax term can be rewritten into  $D^{-1}A$ , where  $A := \exp(\mathbf{Q}\mathbf{K}^T / \sqrt{p})$ , and  $D$  is a diagonal matrix whose diagonal is  $\exp(\mathbf{Q}\mathbf{K}^T / \sqrt{p}) \cdot \mathbf{1}$  (1 is a size- $n$  vector with all elements being 1).
3.2 Sketching Methods
Beyond current attempts to accelerate self-attention, research in the random matrix approximation community can be potentially applied to fast attention. Among the theoretical frameworks, the sketching method (Woodruff, 2014) is the most comprehensive one as its general concept can incorporate many different approaches.
The core idea of the sketching method is to replace an original matrix  $B \in \mathbb{R}^{n_B \times n}$  with its random sketch  $BS$ , where  $S \in \mathbb{R}^{n \times d}$  is a random sketching matrix. In practice, to apply the sketching method we plug an identity matrix into the original expression, and then formally replace the identity matrix with the product  $SS^T$ , as the distribution of  $S$  is usually designed to satisfy the constraint that
$$
\mathbb {E} \left(\boldsymbol {S} \boldsymbol {S} ^ {T}\right) = \boldsymbol {I}. \tag {1}
$$
Common methods to construct a sketching matrix include sub-Gaussian maps (Vershynin, 2010; Halko et al., 2011), subsampled randomized Hadamard transform (SRHT) (Ailon and Chazelle, 2006; Lu et al., 2013; Yang et al., 2017), sparse oblivious subspace embeddings (Cohen et al.,
2016), very sparse random projection (Li et al., 2006), accumulative sketching (Chen and Yang, 2021), and sub-sampling sketching (Monte Carlo algorithms) (Drineas et al., 2006). Specifically, Informer and Linformer, two efficient transformer-based methods mentioned above, can be understood as applications of sub-sampling sketching and sub-Gaussian maps, respectively. We further elaborate the connections in the next subsection.
3.3 Sketching in Self-attention Approximation
A naive step in applying sketching method to approximate the self-attention output  $D^{-1}AV$  is to construct a random sketch of the un-normalized attention score matrix  $\mathbf{A}$ , the bottleneck in computation. Informer and Linformer construct two types of sketches,  $\mathbf{A}^T\mathbf{S}$  and  $\mathbf{A}\mathbf{S}$  respectively.
3.3.1 Informer
Informer selects  $d$  important rows of  $D^{-1}A$ , though deterministically, to represent  $D^{-1}A$ . This process can be related to a sketched approximation  $D^{-1}SS^T A$ , where  $S$  is a sub-sampling matrix defined as follows:
Definition 3.1 (Sub-sampling matrix). Consider a discrete distribution which draws  $i$  with probability  $p_i > 0, \forall i \in [n]$ . For a random matrix  $S \in \mathbb{R}^{n \times d}$ , if  $S$  has independent and identically distributed (i.i.d.) columns and each column  $S^{(j)}$  is  $\frac{1}{\sqrt{dp_i}} e_i$  with probability  $p_i$ , where  $e_i$  is the  $i$ -th column of the  $n$ -by- $n$  identity matrix  $I_n$ , then  $S$  is called a sub-sampling matrix with sub-sampling probabilities  $\{p_i\}_{i=1}^n$ .
Some researchers in the field of approximate ma
trix multiplication have provided a practical guideline for the choice of the sub-sampling probabilities  $\{p_i\}_{i = 1}^n$  in  $S$ . Specifically for the matrix multiplication  $BC$  of two arbitrary matrices  $B$  and  $C$ , Drineas et al. (2006) approximate it with  $BSS^{T}C$  and set the sampling probability  $p_i$  in  $S$  proportional to the product  $\| B^{(i)}\| _2\| C_{(i)}\| _2$ , where  $B^{(i)}$  is the  $i$ -th column in matrix  $B$  and  $C_{(i)}$  is the  $i$ -th row in matrix  $C$ . For the product  $D^{-1}A$ , the probability in sketching will be  $p_i = \frac{\sqrt{\sum_{j = 1}^n a_{ij}^2}}{\sum_{j = 1}^n a_{ij}}$ , where  $a_{ij}$  is the  $j$ -th element of the  $i$ -th row in matrix  $A$ .
The above sampling probability choice  $\{p_i\}_{i = 1}^n$  is highly related to the sparsity measurement used in Informer, which is  $M_{i} = \ln \frac{\sum_{j = 1}^{n}a_{ij}}{(\prod_{j = 1}^{n}a_{ij})^{1 / n}}$ . Here  $p_i$  is the ratio between the quadratic mean and the arithmetic mean of  $\{a_{ij}\}_{j = 1}^n$ ;  $M_{i}$  is the logarithm of the ratio between the arithmetic mean and the geometric mean. It is clear that  $M_{i}$  will increase with  $p_i$  as these two ratios will both be large when  $\{a_{ij}\}_{j = 1}^n$  are highly non-uniform. We conclude that in Informer, the main idea to select the rows with high sparsity measurement can be taken as a special variant of the sub-sampling method above with probabilities  $\{p_i\}$ .
3.3.2 Linformer
Another type of sketch  $\mathbf{A}\mathbf{S}$  is mentioned (but not finally used) in Linformer. The sketching matrix  $\mathbf{S}$  takes a form different from sub-sampling. The construction of  $\mathbf{S}$  in Linformer is motivated by Johnson-Lindenstrauss (JL) transform, which applies the sketching matrix  $\mathbf{S}$  satisfying the  $(\varepsilon, \delta)$ -JL guarantee:
Definition 3.2 (Oblivious Johnson-Lindenstrauss guarantee (Johnson and Lindenstrauss, 1984)). A distribution  $\mathcal{D}$  over  $\mathbb{R}^{n\times d}$  satisfies "oblivious Johnson-Lindenstrauss guarantee" if for some  $\varepsilon, \delta \in (0,1/2)$ :
$$
\forall \boldsymbol {b} \in \mathbb {R} ^ {n}, \underset {\boldsymbol {S} \sim \mathcal {D}} {\mathbb {P}} \left(\left| \| \boldsymbol {S b} \| _ {2} ^ {2} - \| \boldsymbol {b} \| _ {2} ^ {2} \right| > \varepsilon \| \boldsymbol {b} \| _ {2} ^ {2}\right) <   \delta . \tag {2}
$$
Specifically, a matrix with i.i.d. Gaussian elements can meet the above requirement. It has been proven (Johnson and Lindenstrauss, 1984) that with  $d = \mathcal{O}(\varepsilon^{-2}\log (1 / \delta))$ , a Gaussian sketching matrix  $S$  can satisfy the oblivious  $(\varepsilon ,\delta)$ -JL guarantee. To extend the conclusion from a single vector  $\pmb {b}\in \mathbb{R}^n$  to a matrix  $B\in \mathbb{R}^{n_B\times n}$ , the size  $d$  still needs to suffer from an additional  $\log n_B$  term (Vershynin, 2010), which matches the bound in sub-sampling sketching (Drineas et al., 2006, Theorem 1).
However, the direct use of Gaussian sketching matrix, i.e. the approximation  $D^{-1}A S S^T V$  (Wang et al., 2020b, Eqn. (5)), requires the computation of the whole matrix  $A$ . To avoid this computational burden, Linformer replaces the form of sketching method with  $\text{softmax}\left((Q K^T / \sqrt{p}) S\right) S^T V$ , which sacrifices the accuracy for efficiency in some tasks as shown in later experimental results.
4 Methodology: Skeinformer
Based on the previous discussion, we observe that Informer omits the information from the value matrix  $V$ , and Linformer deviates from the usual sketching form for efficiency. To address these issues and fully exploit the power of sketching, we strengthen the attention approximation with the following components.
In Section 4.1, we introduce column importance sampling, which allows the information incorporation from  $V$  to accelerate the matrix multiplication (compared to JL transform); in Section 4.2 and Section 4.3, we leverage the sampled columns to perform the row normalization and reuse the pilot row sampling, which further improves the approximation and makes the training more stable.
We describe the proposed method Skeinformer in Algorithm 1 and verify its performance on matrix approximation in Section 5. We also provide complexity analysis in Section 4.5 to show that our method enjoys the same  $\mathcal{O}(n\log n)$  complexity as other methods.
4.1 Column Sampling
The row selection in Informer has been derived as a special variant of the sketching method and can be further improved by utilizing the information from  $V$ , in a form similar to Linformer:
$$
\boldsymbol {D} ^ {- 1} \boldsymbol {A} \boldsymbol {S} \boldsymbol {S} ^ {T} \boldsymbol {V},
$$
where  $S$  above is a sub-sampling matrix defined in Definition 3.1 with sampling probabilities
$$
p _ {i} \propto \| (\boldsymbol {D} ^ {- 1} \boldsymbol {A}) ^ {(i)} \| _ {2} \| \boldsymbol {V} _ {(i)} \| _ {2}, \quad i = 1, 2, \dots , n.
$$
We remark that using the sub-sampling sketching in this way can both circumvent the computation burden of Gaussian sketching, and also allow the incorporation of the information from  $V$ .
As  $S$  formally samples some columns from  $D^{-1}A$ , we name the procedure as column sampling in our method. The performance regarding
Figure 1: Spectral norm results with the sequence length of 1024 and 4096. Y axis: Lower percentage score means better approximation. X axis: Higher number of feature means larger computation cost. "V-Mean" is an artificial baseline against sampling that always uses a rank-one matrix  $\frac{1}{n}\mathbf{1}^T V$  to approximate the original self-attention.
the Frobenius norm loss of the approximate matrix multiplication can be guaranteed by the following proposition:
Proposition 1 (Adapted from Theorem 1 (Drineas et al., 2006)). Suppose the attention score matrix  $\mathbf{B} := \mathbf{D}^{-1}\mathbf{A} \in \mathbb{R}^{n \times n}$ , the value matrix  $\mathbf{V} \in \mathbb{R}^{n \times p}$ , the number of sampled columns  $d \in \mathbb{Z}^+$  such that  $1 \leq d \leq n$ , and the sampling probabilities  $\{p_i\}_{i=1}^n$  are such that  $\sum_{i=1}^n p_i = 1$  and such that for a quality coefficient  $\beta \in (0,1]$
$$
p _ {i} \geq \beta \frac {\| B ^ {(i)} \| \| V _ {(i)} \|}{\sum_ {i ^ {\prime} = 1} ^ {n} \| B ^ {(i ^ {\prime})} \| \| V _ {(i ^ {\prime})} \|}, \forall i \in [ n ]. \tag {3}
$$
Construct a sub-sampling matrix  $S \in \mathbb{R}^{n \times d}$  with sub-sampling probabilities  $\{p_i\}_{i=1}^n$  as in Definition 3.1, and let  $\mathbf{BSS}^T \mathbf{V}$  be an approximation to  $\mathbf{BV}$ . Let  $\delta \in (0,1)$  and  $\eta = 1 + \sqrt{(8 / \beta) \log(1 / \delta)}$ . Then with probability at least  $1 - \delta$ ,
$$
\left\| \boldsymbol {B} \boldsymbol {V} - \boldsymbol {B} \boldsymbol {S} \boldsymbol {S} ^ {T} \boldsymbol {V} \right\| _ {F} ^ {2} \leq \frac {\eta^ {2}}{\beta d} \| \boldsymbol {B} \| _ {F} ^ {2} \| \boldsymbol {V} \| _ {F} ^ {2}. \tag {4}
$$
Remark. Proposition 1 guides Informer and our method to pick the important rows and columns of the attention score matrix  $B$ . In self-attention, it is feasible to compute the norm  $\| V_{(i')}\|$  of each row in  $V$  with  $\mathcal{O}(n)$  time, assuming the dimension  $p$  in each head is fixed and independent of  $n$ . However, similar to Informer, it is inefficient to exactly compute the  $\ell -2$  norm of each column in the  $n$ -by-  $n$  matrix  $B$ , and we need pilot sampling as well to estimate the norm of the columns in  $B$ . We show that  $\mathcal{O}(\log n)$  samples in the pilot sampling are sufficient to guarantee the quality coefficient  $\beta \geq \sqrt{\frac{1}{3}}$  with high probability by the following lemma. (See proof in Appendix.)
Lemma 1. Assume for any  $i\in [n]$ $\| B^{(i)}\| ^2 /n$  is uniformly lower bounded by a constant  $C$  , where
$\pmb{B} \coloneqq \pmb{D}^{-1}\pmb{A}$ . For another constant  $\delta \in (0, 1/2)$ , we uniformly sample  $d$  indices  $\{j_k\}_{k=1}^d$  from  $[n]$  with replacement, and let  $d$  be a constant multiple of  $\log(n/\delta)$ . Then with probability at least  $1 - \delta$ , the estimated sub-sampling probabilities
$$
\hat {p} _ {i} := \frac {\left(\sum_ {k = 1} ^ {d} b _ {j _ {k} i} ^ {2}\right) ^ {\frac {1}{2}} \| V _ {(i)} \|}{\sum_ {i ^ {\prime} = 1} ^ {n} \left(\sum_ {k = 1} ^ {d} b _ {j _ {k} i ^ {\prime}} ^ {2}\right) ^ {\frac {1}{2}} \| V _ {(i ^ {\prime})} \|}, \forall i \in [ n ], \tag {5}
$$
satisfy the constraints (3) with  $\beta = \sqrt{\frac{1}{3}}$ , where  $b_{ji}$  is the element of  $\mathbf{B}$  from the  $j$ -th row and  $i$ -th column.
This lemma states the sub-sampling weights used in our proposed method. Its computation only requires accesses to  $\{B^{(j_k)}\}_{k=1}^d$  obtained from the pilot sampling, and thus has greatly reduced the time cost. Combining the preceding lemma and Proposition 1, we conclude that with the sampling probabilities  $\{\hat{p}_i\}_{i=1}^n$  estimated by  $\mathcal{O}(\log n)$  pilot samples, the sampled  $d$  important columns suffice to capture the essence of the original output  $BV$ . We conclude this subsection with a remark that the theoretical result that sub-sampling sketching can well approximate the original self-attention, indeed matches the rank collapse phenomenon observed by Dong et al. (2021) that self-attention can be well approximated by a low-rank matrix.
4.2 Adaptive Row Normalization
In addition to the theoretical guarantee of the subsampling sketching method, we identify an important component behind Informer, row normalization, which implicitly fills the un-selected rows with  $\frac{1}{n}$ . The experiments in Section 5 reveal that even the rank-one pure row normalization benchmark  $\frac{1}{n} \mathbf{1} \mathbf{1}^T V$ , as an ablation, will have acceptable spectral norm loss  $\| D^{-1} A V - \frac{1}{n} \mathbf{1} \mathbf{1}^T V \|_2$ . Therefore, we incorporate adaptive row normalization
to provide an even better attention score distribution in each row. It fills un-selected columns with the averaged selected columns. Moreover, from the model training perspective, it allows the whole value matrix  $\mathbf{V}$  in Skeinformer to participate in the computation (compared to only using the subsampling sketch  $S^T V$ ), and thus can improve the efficiency of updating  $W_V$  during the training.
Specifically, in adaptive row normalization any row in the matrix  $A$  can be divided into two parts, the exactly computed elements in the selected columns with indices  $\{j_k'\}_{k=1}^d \subset [n]$  and the other elements in the un-selected columns. For the latter, in each row, we set all the un-selected elements as the geometric mean of the selected ones, considering the exponentiation in softmax. We then perform row normalization based on the above construction, in which the  $i$ -th diagonal element in  $D$  is estimated as
$$
\hat {d} _ {i i} = \sum_ {k = 1} ^ {d} a _ {i j _ {k}} + (n - d) \left(\prod_ {k = 1} ^ {d} a _ {i j _ {k}}\right) ^ {\frac {1}{d}}, \tag {6}
$$
where each  $a_{ij}$  is the corresponding element in matrix  $\mathbf{A}$ . Next we normalize rows composed of exact elements in the selected columns, and the other elements estimated with the mean value above. We comment that though the component of adaptive row normalization makes the proposed method inapplicable to Proposition 1, it benefits the performance on matrix approximation and avoid the cost to compute the diagonal normalization matrix  $\pmb{D}$ . (c.f. Section 5)
4.3 Pilot Sampling Reutilization
Since we have already computed  $B_J$  in pilot sampling step (defined in Ln. 3 of Algorithm 1), we can exactly reproduce the  $d$  rows in the original self-attention output with an additional product  $B_J V$  in  $\mathcal{O}(n \log n)$  time. This allows for more precise approximation with little cost. In addition, the computation of those rows involves the whole key matrix  $K$ , which benefits the training of the parameters  $W_K$ .