Skyformer: Remodel Self-Attention with Gaussian Kernel and Nyström Method

Introduction

The cost of language model training increases exponentially. Among different models, Transformer-based language models (Vaswani et al., 2017; Devlin et al., 2019; Liu et al., 2019; Lewis et al., 2020) are shown to enjoy state-of-the-art (SOTA) performances on many Natural Language Processing (NLP) tasks despite their enormous training cost. One of the computation bottlenecks lies in the self-attention mechanism, which is known to be resource-intensive with quadratic time and space complexity ($O(n)$ where $n$ is the input sequence length). Consequently, Transformers cannot support long sequence processing and large batch size with limited resources.

The challenge of improving computational efficiency of Transformers has motivated several recent studies on attention acceleration, using either sparse attention pattern (Qiu et al., 2020; Child et al., 2019; Zaheer et al., 2020; Beltagy et al., 2020; Kitaev et al., 2020) or low-rank approximation (Choromanski et al., 2020; Wang et al., 2020). However, there is usually a lack of theoretical analysis on the approximation error of these methods due to the complex softmax structure, which makes the theoretical comparison between the efficiency of each method infeasible. It is also unclear in theory how to set the hyper-parameters of those methods to attain a desired level of approximation accuracy.

Another issue of Transformers is the training instability that small perturbations in parameter updates tend to be amplified, resulting in significant disturbances in the model output (Liu et al., 2020a). Transformers on some NLP tasks have shown to be sensitive to hyper-parameters, learning schedulers, or even random seeds, which usually demands a time-costly grid search for the best configuration in real-world applications. It has also been observed in our experiments that a slight change in the learning rate may cause the failure of convergence for some models. We conjecture that the instability in Transformer training comes from the softmax structure, as the un-normalized attention score matrices before softmax tend to have extremely large condition numbers due to its fast singular value decay.

To alleviate the instability issue, an extra factor of $1/\sqrt{p}$ in the softmax kernel SM is suggested by Vaswani et al. (2017) to restrain the scale variation; Liu et al. (2020a) proposes a new scheme to control the magnitude of output change and stabilize the training in early stages. In practice, we also need to consider the lower numerical precision of GPU implementation in model training, which further deteriorates the stability.

To further improve the efficiency, we propose Skyformer (Symmetrization of Kernelized attention for NYström method) to accelerate kernelized attention. Skyformer adapts the Nyström method (Williams and Seeger, 2001; Drineas et al., 2005) to the non-PSD empirical Gaussian kernel matrix (as query matrices in general do not equal to key matrices), by instead lifting the kernelized attention score matrix into a large PSD matrix that contains the un-normalized attention score matrix as the off-diagonal block. We further conduct theoretical analysis by showing that Skyformer has a small matrix approximation error on kernelized attention in the spectral norm. Our experiments on the LRA benchmark show that Skyformer consistently uses less space and time while achieving better accuracy than other baseline methods.

(1) We revisit the intrinsic connection between self-attention and kernel methods, and explore a new kernel-based structure, kernelized attention, to stabilize the training of Transformers.

(2) We propose Skyformer, which approximates the kernelized attention via low dimensional randomized sketches by adapting the Nyström method to a non-PSD matrix. We provide the theoretical guarantee that the matrix multiplication error is small in term of spectral norm.

(3) Extensive experiments show that Skyformer achieves comparable performance to the original self-attention with fewer computational costs.Our code is released at https://github.com/pkuzengqi/Skyformer

Related Work

Among all the transformer acceleration methods, including attention layer simplification by pruning redundant attention heads (Voita et al., 2019; Michel et al., 2019) and model size reduction with knowledge distillation (Jiao et al., 2020; Tang et al., 2019; Liu et al., 2020b), we focus on attention approximation models, which are closely related to kernel methods.

To reduce the time and space complexity by avoiding exhaustive computation over the attention metric, recent studies propose to apply sparse attention patterns to limit the numbers of elements participating in matrix multiplications (Qiu et al., 2020; Child et al., 2019; Zaheer et al., 2020; Beltagy et al., 2020). Beyond limiting the attention to fixed patterns, some approaches learn the patterns by determining token assignments to relevant groups (Kitaev et al., 2020; Roy et al., 2021). Those models utilize local and global information in the attention score matrix to perform approximation, which coincides with the attempt to accelerate the computation in Gaussian processes (Snelson and Ghahramani, 2007).

The attention score matrix is known to exhibit a very fast rate of singular value decay (Bhojanapalli et al., 2020; Dong et al., 2021), similar to that of an empirical kernel matrix (Yang et al., 2017). This near singular property motivates many low-rank attention approximation methods to skillfully leverage the computation techniques in kernel methods. Among them, Linformer (Wang et al., 2020) compresses the size of the key and value matrix with random projections based on the Johnson–Lindenstrauss transform, a common randomized sketching method in Gaussian processes (Yang et al., 2017); Reformer (Kitaev et al., 2020) applies locality-sensitive hashing (LSH) (Har-Peled et al., 2012) to simplify the computation of the attention score matrix, which is widely used in kernel density estimation (Charikar and Siminelakis, 2017; Backurs et al., 2019); Performer (Choromanski et al., 2020) projects both query and key matrix through random Fourier features (Rahimi et al., 2007), heavily exploiting Bochner Theorem for stationary kernels.

The most related papers to ours are linear attention (Katharopoulos et al., 2020), Synthesizer (Tay et al., 2020a), and Nyströmformer (Xiong et al., 2021). Linear attention takes the softmax structure in self-attention as a measure of similarity and replaces it with the dot product of separately activated query and key matrices; Synthesizer aims to modify the original self-attention by replacing the dot product before softmax with Synthetic Attention, which generates the alignment matrix independent of token-token dependencies. Their attempts indicate that the softmax structure in self-attention is not the only feasible choice, and justify our usage of kernelized attention. Rather than remodeling self-attention, Nyströmformer applies the Nyström method (Williams and Seeger, 2001; Drineas et al., 2005), a powerful and effective method for large-scale kernel machines acceleration, to approximate the attention score matrix. However, Nyströmformer applies the Nyström method to a non-PSD matrix, and thus fails to utilize the full potential of the Nyström method. This issue is resolved in our proposed Skyformer by instead lifting the kernelized attention score matrix into a large PSD matrix which contains the target non-PSD matrix as its off-diagonal block. For more details on attention approximation methods, we refer readers to a survey paper on efficient transformers (Tay et al., 2020c).

Preliminaries and notations

where $\bm{Q}=\bm{X}\bm{W}_{Q}$, $\bm{K}=\bm{X}\bm{W}_{K}$, and $\bm{V}=\bm{X}\bm{W}_{V}$, and $\bm{W}_{Q}$, $\bm{W}_{K}$ and $\bm{W}_{V}$ are the query, key, and value weight metrics that linearly project the input $\bm{X}$ of $d_{0}$ dimension to an output tensor of $p$ dimensions.

To simplify the future analysis, the left softmax term can be rewritten into $\bm{D}^{-1}\bm{A}$, where $\bm{A}\vcentcolon=\exp(\bm{Q}\bm{K}^{T}/\sqrt{p})$ is the un-normalized attention score matrix; $\bm{D}$ is a diagonal matrix whose diagonal is $\exp(\bm{Q}\bm{K}^{T}/\sqrt{p})\cdot\bm{1}$ (by convention $\bm{1}$ is a size-$n$ vector with all elements being $1$). Following the notation in Performer (Choromanski et al., 2020), we define $\text{SM}(\bm{q},\bm{k})\vcentcolon=\exp(\bm{q}^{T}\bm{k}/\sqrt{p})$ as the softmax kernel function, and represent $\bm{A}$ by the notation $\text{SM}(\bm{Q},\bm{K})$, which means the element $a_{ij}$ from the $i$-th row and $j$-th column in $\bm{A}$ is equal to $\text{SM}(\bm{q}_{i},\bm{k}_{j})$. Throughout this paper $\bm{q}_{i}$ (resp. $\bm{k}_{j}$) means the $i$-th (resp. $j$-th) row in $\bm{Q}$ (resp. $\bm{K}$).

We close this subsection with a short lemma to show $\text{SM}(\cdot,\cdot)$ is a positive semidefinite (PSD) kernel function (Wainwright, 2019, Definition 12.6) by relating it to Gaussian kernels.

We first state an important equation to connect the softmax kernel and Gaussian kernels as follows:

The middle part $\exp\left(\frac{\|\bm{q}_{i}-\bm{q}_{j}\|^{2}}{2\sqrt{p}}\right)$ is exactly a Gaussian kernel with bandwidth $p^{\frac{1}{4}}$. (Choromanski et al. (2020) have more discussion on the findings.)

Through this equation, we can rewrite $\bm{C}$ as

where $\bm{D}_{Q}$ is a diagonal matrix with elements $(\bm{D}_{Q})_{ii}=\exp\left(\frac{\|\bm{q}_{i}\|^{2}}{\sqrt{p}}\right),\forall i\in[n]$, and $\kappa(\bm{q}_{i},\bm{q}_{j})\vcentcolon=\exp\left(-{\|\bm{q}_{i}-\bm{q}_{j}\|^{2}}/{2}\right)$ is the standard Gaussian kernel function.

We prove the lemma by using the fact that $\kappa$ is a PSD kernel and $\kappa\left(\frac{\bm{Q}}{p^{1/4}},\frac{\bm{Q}}{p^{1/4}}\right)$ is a PSD matrix. $\diamondsuit$

2 Nyström method

We close this subsection with a remark that it is not appropriate to directly extend the Nyström method from kernel method to self-attention due to a core requirement that $\bm{B}$ should be PSD with consideration of approximation performance improvement. We will show in the next section how to address this challenge and properly adapt Nyström method to non-PSD matrices.

3 Approximation evaluation

Beyond the time and space complexity, attention acceleration methods have been mostly evaluated with empirical experiment results, such as the perplexity of pretrained language models and the fine-tuned performance on downstream natural language understanding tasks. Specifically, Long Range Arena benchmark (Tay et al., 2020b) has been proposed to systematically evaluate the performance of efficient transformers with ten NLP tasks in long-context scenarios. However, such empirical results are indirect for theoretical analysis. Therefore, we introduce a common criterion used in matrix approximation, spectral norm, to ease the future discussion on performance.

In previous works, the direct analysis of the approximation error to the entire output $\bm{D}^{-1}\bm{A}\bm{V}$ in the $(\varepsilon,\delta)$-MA manner is usually spared due to the difficulty caused by the complex softmax structure. In this paper, with the new kernelized attention, we are allowed to perform the analysis through the existing theoretical results in kernel methods. Consequently, in Section 4.5 we are able to give a relatively precise error analysis on the approximation of Skyformer to the entire kernelized attention, which eases the future comparison with other methods approximating kernelized attention.

Method

Kernelized Attention replaces the softmax structure in vanilla self-attention with a Gaussian kernel, and the new attention model is stated as:

where we define the $n$-by-$n$ matrix $\bm{C}$ as the kernelized attention score matrix $\kappa(\bm{Q}/p^{1/4},\bm{K}/p^{1/4})$.

The justification for using the kernelized attention model is as follows. A significant advantage of softmax attention is that tokens are allowed to attend to a limited number of other important tokens in the sequence. We observe that Gaussian kernel function can play a similar role. The expression of a Gaussian kernel is $\kappa(\bm{q}_{i},\bm{k}_{j})\vcentcolon=\exp\left(-\|\bm{q}_{i}-\bm{k}_{j}\|^{2}/2\right)$. Via this expression, for token $i$ in the query, Gaussian kernel assigns a large attention score to the token $j$ when $\bm{k}_{j}$ is close to $\bm{q}_{i}$. The distance-based weight assignment is indeed considered as a major reason why kernel methods are powerful. The form of kernelized attention also leads to an automatic normalization. Based on Equation (1), the new attention model can be rewritten in terms of the un-normalized attention score matrix $\bm{A}$ as

where $\bm{D}_{Q}$ (resp. $\bm{D}_{K}$) is a diagonal matrix with elements $(\bm{D}_{Q})_{ii}=\exp\left(\frac{\|\bm{q}_{i}\|^{2}}{\sqrt{p}}\right)$ (resp. $(\bm{D}_{K})_{ii}=\exp\left(\frac{\|\bm{k}_{i}\|^{2}}{\sqrt{p}}\right)$), $\forall i\in[n]$. We remark the kernelized attention model can thus be formally taken as a variant of the original self-attention, which instead normalizes the matrix $\bm{A}$ in a form of $\bm{D}^{-1}\bm{A}$. The intrinsic normalization allows kernelized attention to have a more reasonable condition number than self-attention, which benefits the stability of model training. To demonstrate the improvement in stability, we additionally provide a toy experiment in Appendix F, which shows the “condition number" of kernelized attention is smaller than self-attention. Moreover, empirical evaluation in Section 5 supports our claim that the new attention model can attain a comparable performance to the original attention model.

2 Skyformer: a modified Nyström method

Before jumping into details of Skyformer, we first propose a method to apply Nyström method to approximate an asymmetric (and thus non-PSD) empirical kernel matrix $\bm{B}$ constructed with any PSD kernel $\phi(\cdot,\cdot)$. Specifically, with two different $n$-by-$p$ design matrices $\bm{Q}$ and $\bm{K}$, its element $b_{ij}$ from the $i$-th row and $j$-th column in $\bm{B}$ is equal to $\phi(\bm{q}_{i},\bm{k}_{j})$, where $\bm{q}_{i}$ (resp. $\bm{k}_{j}$) is the $i$-th (resp. $j$-th) row in $\bm{Q}$ (resp. $\bm{K}$). We remark this type of empirical kernel matrices involves the un-normalized attention score matrix $\bm{A}\vcentcolon=\text{SM}(\bm{Q},\bm{K})$, and the empirical Gaussian kernel matrix $\bm{C}\vcentcolon=\kappa(\bm{Q}/p^{1/4},\bm{K}/p^{1/4})$. Therefore this method leads to a low-rank approximation to the output of either self-attention $\bm{D}^{-1}\bm{A}\bm{V}$ or Kernelized Attention $\bm{C}\bm{V}$. ($\bm{D}$ in self-attention can be obtained by computing $\bm{A}\cdot\bm{1}$, and thus a low-rank approximation to $\bm{A}$ also implies an approximation to $\bm{D}$.)

Computational details are stated as follows. To tackle the challenge of approximating a non-PSD matrix $\bm{B}$, our first step is to complete the matrix into a PSD matrix $\bar{\bm{B}}$:

where $\bm{S}$ is a $2n$-by-$d$ uniform sub-sampling matrix as defined in Definition 1. The final approximation will be given as

and thus we show our task of approximating the non-PSD matrix $\bm{B}$ boils down to well approximating the PSD matrix $\bar{\bm{B}}$.

With the derivation above, we officially introduce our proposed Skyformer as an approximation to Kernelized Attention, which applies the modified Nyström method to the kernelized attention score matrix $\bm{C}$. The next two subsections will continue our discussion on it, and respectively state the theoretical analysis of its approximation error and some details of its implementation in practice.

3 Error analysis of Skyformer

As mentioned, an implicit advantage of using Kernelized Attention is that we can leverage the existing conclusions for kernel methods to analyze the theoretical properties of the model. In this subsection, we aim to provide some theoretical analysis of its approximation error.

We state a high probability bound on the size $d$ of the sub-sampling matrix used in Skyformer to attain $(\varepsilon,\delta)$-MA property for the kernelized attention score matrix $\bm{C}$ by the following theorem. We refer the readers to the proof in Appendix D to take a closer look at our claim that the matrix to be approximated should be PSD is a key to the theoretical guarantee of Nyström method.

This theorem implies the time and space complexity of our proposed approximation depends on the statistical dimension $d_{stat}$. If we directly use the conclusion from Gaussian kernels, $d_{stat}$ should be $\widetilde{\mathcal{O}}(1)$ (complexity modulo poly-log term) (Yang et al., 2017) due to the exponential eigenvalue decay rate of Gaussian kernels, which is comparable to the complexity of most other efficient transformers. However, different than the case in the classical kernel methods, the distribution of the query and key matrix $\bm{Q}$ and $\bm{K}$ changes during the training procedure, which may invalidate the conclusion about $d_{stat}$. We leave the exact non-asymptotic analysis of the computational complexity for future work.

4 Workaround in implementation

A potential limitation with the implementation of the proposed method lies in the tricky fact that the matrix inversion on GPU is much slower and numerically less stable than the same operation on CPU due to the different back-end libraries in the two platforms. We attempt to circumvent the problem by adapting the strategy in Nyströmformer (Xiong et al., 2021) to our setting. Specifically, we use the matrix-product-based iterative method (Razavi et al., ) for finding approximate inverses, instead of some division-based methods (such as the conjugate gradient method) which induces some instability in model training.

To apply the iterative method and inverse matrix $\bm{M}=\bm{S}^{T}\bar{\bm{C}}\textbf{S}$, we need to satisfy its assumption (Razavi et al., , Theorem 2) that $\|\bm{I}-\bm{M}\|<1$. In practice, we instead pass the matrix $\bm{D}_{M}^{-1/2}(\bm{M}+\gamma\bm{I})\bm{D}_{M}^{-1/2}$ as an input to the iterative method, where $\gamma>0$ is a small constant and the diagonal matrix $\bm{D}_{M}$ is defined as $\text{diag}\left((\bm{M}+\gamma\bm{I})\bm{1}\right)$. We give the following lemma to justify our practical usage of the method. The proof is deferred to Appendix E.

Given a constant $\gamma>0$, if matrices $\bm{M}$ is constructed as $\bm{S}^{T}\bar{\bm{C}}\textbf{S}$, and $\bm{D}_{M}$ are defined as above, then all the singular values of $\bm{D}_{M}^{-1/2}(\bm{M}+\gamma\bm{I})\bm{D}_{M}^{-1/2}$ are within $(0,1)$, which implies that $\|\bm{I}-\bm{D}_{M}^{-1/2}(\bm{M}+\gamma\bm{I})\bm{D}_{M}^{-1/2}\|<1$.

We further comment that numerically an implicit risk of the Schulz-type iterative method we use is the unintended consequence of “zero fill-in". If we use some sparse kernels (e.g. test functions with bounded support) other than Gaussian kernels, the empirical kernel matrices are sparse while the approximate inverse will converge to a dense matrix, which increases the computational cost.

5 Empirical approximation evaluation

Spectral norm, the maximum singular value of a matrix, is a computation-light indicator of matrix approximation performance. In this work, we compare the spectral norm of the difference between the outputs from attention functions and the output from vanilla self-attention with the same input.

We use the initialized and pretrained bert-base-cased models from Huggingface’s implementation (Wolf et al., 2019) . The input vector $X$ is embedded from the tokenized raw text in Wikitext-2 dataset (Merity et al., 2017). The query, key and value weight matrices in initialized or pretrained models transform input $X$ into $Q,K,V$ of different distributions. We compare the results with different sequence lengths and different numbers of features used in attention approximation methods. We set the number of features in the range of $2^{4}$ to $2^{8}$. More features usually require more computation resources.

Figure 1 shows the performance of the modified Nyström method on approximation error with regards to the number of features. We conclude that for Skyformer the approximation is significantly better with the increased number of features, while for other methods the gain is not obvious. The good performance of the modified Nyström method also validates our previous claim that the Nyström method is currently one of the most powerful methods in large-scale kernel machines acceleration.

Remark. Although in a single step the modified Nyström method in Section 4.2 can give low approximation error, we do not recommend directly applying it to the original self-attention. With some exploratory experiments on classification tasks, we find the variant suffers a more severe gradient explosion issue than usual transformers. We speculate that it is because the matrix $\bm{S}^{T}\bar{\bm{A}}\textbf{S}$ (in the middle of Equation (6)) inherits the high condition number of the original attention score matrices $\bm{A}$, while the derivative of matrix inverse ($\left(\bm{A}^{-1}\right)^{\prime}=-\bm{A}^{-1}\bm{A}^{\prime}\bm{A}^{-1}$) further amplifies the condition number during backpropagation.

Experimental Results

Tasks and Datasets. We evaluate the proposed methods on five classification tasks on LRA benchmark (Tay et al., 2020b), which focuses on model quality under long-context scenarios: ListOps (Nangia and Bowman, 2018), Text Classification on IMDb review dataset (Maas et al., 2011), Document Retrieval on AAN dataset (Radev et al., 2013), Pathfinder (Linsley et al., 2018), and Image Classification on CIFAR-10 (Krizhevsky et al., 2009). The LRA benchmark covers diverse long-sequence tasks in sequence length, task difficulty, and inspected model abilities. For example, ListOps and Pathfinder evaluate the abilities to capture the long-range hierarchical dependency and spatial dependency, respectively, which poses challenges for sparse attention pattern based methods. We report the classification accuracy on the test set, training time, and peak memory usage during training for each task.

Baselines. Aside from the vanilla quadratic self-attention, we compare with Big Bird (Zaheer et al., 2020), Performer (Choromanski et al., 2020), Linformer (Wang et al., 2020), Nyströmformer (Xiong et al., 2021), Informer (Zhou et al., 2020), and Reformer (Kitaev et al., 2020). Most methods are approximating the vanilla full attention for efficiency and thus are not expected to have better performance. As it is not realistic to exhaustively fine-tune all models and search for the best performance under limited computation resources, we instead only replace the self-attention module with the various attention methods and keep other experimental settings the same for fair comparisons.

Implementation Details. We conduct each experiment on one Tesla V100 SXM2 16GB. We use the LRA evaluation benchmark reimplemented in PyTorch by Xiong et al. (2021). We use a 2-layer transformer model with $64$ embedding dimension, $128$ hidden dimension, $2$ attention heads, and mean pooling for classification. Batch size is selected conditioned on the memory requirements of the standard self-attention method, which leads to $16$ for Text Classification, $32$ for ListOps, $16$ for Document Retrieval, $128$ for Pathfinder, and $256$ for Image Classification. Learning rate is set to $1e-4$ for Text Classification, ListOps, and Image Classification, and $2e-4$ for Retrieval and Pathfinder. Each model on each task is trained for $50k$ steps, during which the best checkpoint with the highest accuracy on the development set will be saved for evaluation. For comparable computation complexity, we control the number of features to be $128$ used in all methods (except Big Bird), under which setting the models will visit $128\cdot n$ elements in the attention matrix. For numerical consistency, all experiment results are averaged across three runs with different random seeds.

We do not follow all settings in (Xiong et al., 2021) due to the hardware limitation. The compromises, such as approximation dimension and gradient accumulations steps, might bring performance differences comparing to results reported in (Xiong et al., 2021). The training instability problem also helps explain the performance gap.

Results. The training process of the standard softmax-based method is unstable as observed in Figure 2: it takes more steps to reach the stationary distribution of its long-time limit, and it is more easily getting stuck in a local minimum. Runs with different random seeds may bring divergent performances, and probably leads to lower averaged scores. We have also tried directly approximating the self-attention method with the Nyström method and observed numerical instability during training.

Replacing the softmax structure with Gaussian kernel somehow alleviates this instability problem with boosted performance as shown in Table 1. However, the time and space requirement of Kernelized Attention is not significantly improved compared to the original version, which serves as the motivation to approximate Kernelized Self-Attention with Nyström method.

Though not necessarily the fastest, our proposed Skyformer can efficiently converge to the long-time limit with comparable general performance in classification accuracy (Table 1) and resource consumption (Table 2). The advantages over the standard self-attention are significant with consistently less training time and generally better performance. For example, Skyformer brings nearly 4 times speed-up on text classification and document retrieval while with 2.75% and 1.37% accuracy improvement over the standard self-attention.

Limitations. The applications of Skyformer might be limited to long sequence tasks because for small sequence length $n$ the statistical dimension $d_{s}tat$ might be close to $n$.

To make the claim above clear, we first reiterate that the efficiency of Skyformer is related to $d_{s}tat$. As implied by Theorem 2, the intrinsic difficulty of approximating a raw attention score matrix is concluded as $d_{stat}$, which corresponds to the effective rank of matrix $\bar{C}$. The complexity of Skyformer depends on the sub-sample size $d$ (the size of the sub-sampling matrix $S$). A large $d_{stat}$ leads to a large $d$ , and an inefficient application of the Nyström method.

The classical theory for statistical dimension only guarantees that $d_{stat}$ is small (compared to $n$) when $n$ is large enough, and it is possible the statistical dimension associated with a short sequence might be even close to the sequence $n$. Therefore a large $n$ serves as a condition to make the method work. Figure 1 empirically shows that our method performs better with larger $n$’s.

Conclusions and future work

Motivated by the connection between kernel methods and self-attention, we introduce Kernelized Attention, which replaces the softmax structure in self-attention with a Gaussian kernel. We also propose Skyformer, which adapts the Nyström method to Kernelized Attention to improve its efficiency. We expect the new model can enjoy more stable training while inheriting the strong performance from self-attention. Extensive experiments verify our intuitions and show that both Kernelized Attention and its Nyström approximation variant have comparable accuracy to the original Transformer on the LRA benchmark.

Direct development of this work is the incorporation of further computation tricks in kernel methods, such as the local and global approximation for gram matrix (Snelson and Ghahramani, 2007) and the importance sampling in Nyström methods (Musco and Musco, 2017; Chen and Yang, 2021b, a). Other related questions include the choice of the kernel other than the Gaussian kernel in our kernelized attention model. It is expected that for different tasks there will be specific kernels more proper than the original self-attention. The results in this work also shed new light on the design of the attention mechanism, which may benefit board downstream NLP tasks.

Acknowledgments and Disclosure of Funding

This research is based upon work in part supported by the Office of the Director of National Intelligence (ODNI), Intelligence Advanced Research Projects Activity (IARPA), via contract No. FA8650-17-C-9116, and U.S. DARPA KAIROS Program No. FA8750-19-2-1004. This work is also in part supported by NSF grant DMS-1810831. 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, ODNI, IARPA, or the U.S. Government. The U.S. Government is authorized to reproduce and distribute reprints for governmental purposes notwithstanding any copyright annotation therein.

References

Appendix A Validation loss

Figure 3 shows the validation loss changes with respect to training time for $50$k steps as supplementary results for the experiments in Section 5. In general, Skyformer converges faster and finishes $50$k steps earlier than vanilla Attention and Kernelized Attention over all tasks. We further remark that on Text Classification, all models quickly fall into over-fitting, and thus the validation losses rise quickly. On Pathfinder, due to the difficulty of training, in the trial shown in the figure vanilla Attention fails to reach the best long-time limit under a certain setting.

Appendix B Singular value decay rate

Figure 4 shows the singular value distribution of attention output from the second layer of a trained vanilla transformer. Results are averaged across one random batch from the test set in each LRA task.

The singular values decay fast and thus justify the low-rank approximation, as analyzed by Wang et al. , Dong et al. . We propose to measure the task difficulty with the singular value decay rate in attention output, as higher intrinsic task difficulty forces the model to output a matrix with more large singular values. Such matrices are considered more informative since they are harder to approximate, requiring more ranks even in the truncated SVD approximation. With the observation in Figure 4, we conclude that the singular values in Document Retrieval and Pathfinder tasks decay slower, and those two tasks are more difficult than Text Classification and ListOps.

Appendix C Useful facts

This section introduces some useful facts, which are key in the proof in the next section. To start with, we provide a matrix concentration inequality as follows.

Consider a finite sequence $\{\bm{X}_{k}\}$ of independent, random, self-adjoint matrices with dimension $n$. Assume that each random matrix satisfies

For a certain $n$-by-$n$ orthogonal matrix $\bm{H}$ ($\bm{H}\bm{H}^{T}$ is a diagonal matrix) and an $n$-by-$d$ uniform sub-sampling matrix $\bm{S}$ (as defined in Definition 1 in the main paper), we denote the sketching matrix $\bm{\Pi}\vcentcolon=\sqrt{n}\bm{S}$. We aim to show $\bm{H}\bm{\Pi}\bm{\Pi}^{T}\bm{H}^{T}$ can satisfy $(\frac{1}{2},\delta)$-MA property for $\bm{H}\bm{H}^{T}$ by the following lemma.

Denote the stable rank $s\vcentcolon=\frac{\|\bm{H}\|_{F}^{2}}{\|\bm{H}\|^{2}}\geq 1$, and a constant $\delta<1/2$. Suppose there exists a constant $\beta\in(0,1]$ such that $\beta\leq\frac{\|\bm{H}\|_{F}^{2}}{n\|\bm{H}^{(i)}\|^{2}},\forall i=1,\dots,n$, where $\bm{H}^{(i)}$ is the $i$-th column of $\bm{H}$. There exists a constant $C_{0}$ that if

then $\bm{H}\bm{\Pi}\bm{\Pi}^{T}\bm{H}^{T}$ satisfies $(\frac{1}{2},\delta)$-MA property for $\bm{H}\bm{H}^{T}$.

The main idea is to utilize Lemma C.1 by setting $t=\frac{1}{2}\|\bm{H}\bm{H}^{T}\|=\frac{1}{2}\|\bm{H}\|^{2}$. Specifically, we denote the matrices

We still need two steps to give control of $R$ and $\sigma^{2}$. For $R$, we have

where $\{z_{ki}\}_{i=1}^{n}$ are the indicators of whether the $i$-th column is chosen. The first inequality of the preceding display holds due to the fact that $\bm{H}$ is an orthogonal matrix. Using the condition $n\leq\frac{\|\bm{H}\|_{F}^{2}}{\beta\|\bm{H}^{(i)}\|^{2}},\forall i=1,\dots,n$, we further have

and we thus set $R\vcentcolon=\frac{\|\bm{H}\|_{F}^{2}}{d\beta}$. On the other hand,

Again using the condition that $n\|\bm{H}^{(i)}\|^{2}\leq\frac{\|\bm{H}\|_{F}^{2}}{\beta},\forall i=1,\dots,n$, we reach

and set $\sigma^{2}\vcentcolon=\frac{\|\bm{H}\|_{F}^{2}}{d\beta}\left\|\bm{H}\right\|^{2}$.

Finally we plug $R$ and $\sigma^{2}$ into Lemma C.1 and obtain:

To ensure the right-hand-side is smaller than $\delta$, we just need

which validates the lemma. $\diamondsuit$

Appendix D Proof of Theorem 2 in the main paper

For the second conclusion, we utilize the following important identity:

Taking $\varepsilon=\frac{1}{2}$ and applying Lemma C.2, we can conclude that with the conditions on $d$ in the theorem, $\bm{\Psi}^{T}\bm{\Pi}\bm{\Pi}^{T}\bm{\Psi}$ satisfies $(\frac{1}{2},\delta)$-MA property for $\bm{\Psi}^{T}\bm{\Psi}$. Therefore it holds with probability $1-\delta$ that,

From identity $\bm{V}\bar{\bm{\Sigma}}^{\frac{1}{2}}\bar{\bm{\Sigma}}^{\frac{1}{2}}\bm{V}^{T}=\bm{B}^{T}\bm{B}+\frac{\lambda}{2}\bm{I}$, we obtain

Finally, we multiply two sides of Eq. (8) by $(\bm{I}-\bm{P_{\Pi}})$ to obtain

Based on the conclusion above, the last implication is direct:

which completes the proof. $\diamondsuit$

Appendix E Proof of Lemma 3 in the main paper

As $\bar{\bm{C}}$ is constructed based on a PSD kernel, $\bar{\bm{C}}$ is also PSD. Consequently $\bm{M}=\bm{S}^{T}\bar{\bm{C}}\textbf{S}$ is PSD, and $\bm{D}_{M}^{-1/2}(\bm{M}+\gamma\bm{I})\bm{D}_{M}^{-1/2}$ is positive definite, with all eigenvalues positive. To prove the claim in the lemma we only need to show the eigenvalues of $\bm{D}_{M}^{-1/2}(\bm{M}+\gamma\bm{I})\bm{D}_{M}^{-1/2}$ are bounded from above by $1$. It is equivalent to prove that $\bm{I}-\bm{D}_{M}^{-1/2}(\bm{M}+\gamma\bm{I})\bm{D}_{M}^{-1/2}$ is PSD, which can be induced by another statement that $\bm{L}\vcentcolon=\bm{D}_{M}-(\bm{M}+\gamma\bm{I})$ is PSD.

where the last equation holds due to the fact that $(\bm{D}_{M})_{ii}=\sum_{j=1}^{d}\bm{W}_{ij}$.

Combining the pieces above we can conclude that $\|\bm{I}-\bm{D}_{M}^{-1/2}(\bm{M}+\gamma\bm{I})\bm{D}_{M}^{-1/2}\|<1$. $\diamondsuit$

Appendix F Additional discussions about the stability in model training

For our argument about stability, we mainly refer to the paper [Liu et al., 2020a], which identifies that the amplification of small parameter perturbations in the self-attention module is the root cause of training instability. We take kernelized attention as mitigation since it contains an automatic normalization. We have empirically used Figure 2 and Figure 3 in Appendix A to support our claim.

For further analysis we conduct a toy experiment adapting from Figure 4 in the aforementioned paper [Liu et al., 2020a]. We aim to show that in kernelized attention (and Skyformer) the output changes $f(x,W^{*})-f(x,W)$ for parameter changes $W^{*}-W$ is smaller than in self-attention (and its approximation Nyströmformer). This concept involved is somewhat similar to condition number and below we will formalize it as “instability score".

We show a table of the averaged ratios between the instability scores of kernelized attention (we also add Skyformer and Nyströmformer for reference) and self-attention to conclude our statement about stability. A ratio smaller than $1$ means higher stability compared to self-attention. We follow all the settings in Table 1 in the main paper except here we only update the model for 20 steps (we limit the number of steps as suggested by Liu et al. [2020a] to make the results of the same step comparable among different models). In step $i$ for each model we compute the instability score $\tau_{i}=\frac{\|f(x_{i},W_{i})-f(x_{i},W_{i-1})\|^{2}_{F}}{\|W_{i}-W_{i-1}\|^{2}_{F}},i=1,\cdots,20$, where $f()$ gives the embedding after two layers, $x_{i}$ is the $i$-th input sequence batch, $W_{0}$ represents the initial parameters, and $W_{i}$ represents the parameters after step $i$. In each step we compute the ratio of a certain method’s $\tau_{i}$ to the $\tau_{i}$ of self-attention, and finally average the $20$ ratios in Table 3 in the appendix.

As we can observe, both kernelized attention and Skyformer consistently have a lower instability score than self-attention, while the instability score of Nyströmformer, an approximation to self-attention, fluctuates around $1$ in all the tasks. The results support our claim that the proposed kernelized attention can improve stability.