Sketching as a Tool for Understanding and Accelerating Self-attention for Long Sequences (2022)

Abstract: 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.

## introduction transformer @xcite 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 @xcite @xcite @xcite , among which linformer @xcite and informer @xcite are two representative approaches to reducing the o(n 2 ) self-attention to an o(n) operation ( o(•) means o(•) 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 @xcite , to help explain the key ideas in informer and linformer from the perspective of matrix approximation. specifically, sketching methods replace the original matrix b with its random sketch bs to reduce computations. in section 3.3 we introduce some concrete instances of commonly used distributions for constructing the random sketching matrix 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 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 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-sistently uses less space and time while achieving better accuracy than most baseline methods. in summary, our contributions are twofold: ## related work the ability to process long sequences is critical for many natural language processing tasks, including document summarization @xcite , question answering @xcite , information extraction @xcite @xcite , and machine translation @xcite . 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 @xcite . selective attention methods limit the scope of matrix operation with sparse attention patterns or column/row sampling methods. blockbert @xcite introduces sparse block structures into the attention matrix. sparse transformer @xcite introduces dilated patterns. big bird @xcite proposes a combination of random, window, and global attention. longformer @xcite combines local windowed attention with task-motivated global attention. the most related work to ours is in-former @xcite , 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 @xcite compresses the size of the key and value matrices by the johnson-lindenstrauss transform @xcite . performer @xcite 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 @xcite . nyströmformer @xcite instead utilizes nyström method @xcite to approximate the attention score matrix. skyformer @xcite 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 @xcite , 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 @xcite , which forces the query and key matrices to be identical and applies locality-sensitive hashing (lsh) @xcite 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. ## 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 o(n log n) complexity as other methods. ## 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 q, k, v. specifically we compute ∥bv -r∥ 2 , where b := 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 q, k, v, we first truncate the raw text from wikitext-2 dataset @xcite into sequences of length 512. then we transform the input x into q, k, v with the query, key, and value weight metrices 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 1 n 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. ## experiment 6.1 benchmark we experiment on long range arena (lra) benchmark @xcite , including listops @xcite ), text classification @xcite , pathfinder @xcite , and image classification @xcite . lra is designed for long-context scenarios and more appropriate for evaluating efficient transformers comparing to glue @xcite with shorter input context . following @xcite 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. ## conclusion 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.