linformer-pytorch
My take on a practical implementation of Linformer for Pytorch.
view repo
Large transformer models have shown extraordinary success in achieving state-of-the-art results in many natural language processing applications. However, training and deploying these models can be prohibitively costly for long sequences, as the standard self-attention mechanism of the Transformer uses O(n^2) time and space with respect to sequence length. In this paper, we demonstrate that the self-attention mechanism can be approximated by a low-rank matrix. We further exploit this finding to propose a new self-attention mechanism, which reduces the overall self-attention complexity from O(n^2) to O(n) in both time and space. The resulting linear transformer, the Linformer, performs on par with standard Transformer models, while being much more memory- and time-efficient.
READ FULL TEXT VIEW PDFMy take on a practical implementation of Linformer for Pytorch.
Reproducing the Linear Multihead Attention introduced in Linformer paper (Linformer: Self-Attention with Linear Complexity)
Sparse Transformer with limited attention span in PyTorch
Transformer models (vaswani2017attention) have become ubiquitous for wide variety of problems in natural language processing (NLP), including translation (ott2018scaling), text classification, question answering, among others (raffel2019exploring; mohamed2019transformers). Over the last couple of years, the number of parameters in state-of-the-art NLP transformers has grown drastically, from the original 340 million introduced in BERT-Large to 175 billion in GPT-3 (brown2020language). Although these large-scale models yield impressive results on wide variety of tasks, training and deploying such model are slow in practice. For example, the original BERT-Large model (devlin2019bert) takes four days to train on 16 Cloud TPUs, and the recent GPT-3 (brown2020language)
consumed orders of magnitude more petaflops / day to train compared to its predecessor, GPT-2
(radford2019language). Beyond training, deploying Transformer models to real world applications is also expensive, usually requiring extensive distillation (hinton2015distilling) or compression.The main efficiency bottleneck in Transformer models is its self-attention mechanism. Here, each token’s representation is updated by attending to all other tokens in the previous layer. This operation is key for retaining long-term information, giving Transformers the edge over recurrent models on long sequences. However, attending to all tokens at each layer incurs a complexity of with respect to sequence length. Thus, in this paper, we seek to answer the question: can Transformer models be optimized to avoid this quadratic operation, or is this operation required to maintain strong performance?
Prior work has proposed several techniques for improving the efficiency of self-attention. One popular technique is introducing sparsity into attention layers (child2019generating; qiu2019blockwise; beltagy2020longformer) by having each token attend to only a subset of tokens in the whole sequence. This reduces the overall complexity of the attention mechanism to (child2019generating). However, as shown in qiu2019blockwise, this approach suffers from a large performance drop with limited efficiency gains, i.e., a 2% drop with only 20% speed up. More recently, the Reformer (kitaev2019reformer) used locally-sensitive hashing (LSH) to reduce the self-attention complexity to . However, in practice, the Reformer’s efficiency gains only appear on sequences with length (Figure 5 in kitaev2019reformer). Furthermore, the Reformer’s multi-round hashing approach actually increases the number of sequential operations, which further undermines their final efficiency gains.
In this work, we introduce a novel approach for tackling the self-attention bottleneck in Transformers. Our approach is inspired by the key observation that self-attention is low rank
. More precisely, we show both theoretically and empirically that the stochastic matrix formed by self-attention can be approximated by a low-rank matrix. Empowered by this observation, we introduce a novel mechanism that reduces self-attention to an
operation in both space- and time-complexity: we decompose the original scaled dot-product attention into multiple smaller attentions through linear projections, such that the combination of these operations forms a low-rank factorization of the original attention. A summary of runtimes for various Transformer architectures, including ours, can be found in Table 1.One predominant application of Transformers, that has seen the most gains, is using them as pretrained language models, whereby models are first pretrained with a language modeling objective on a large corpus, then finetuned on target tasks using supervised data (devlin2019bert; liu2019roberta; lewis2019bart). Following devlin2019bert, we pretrain our model on BookCorpus (zhu2015aligning) plus English Wikipedia using masked-language-modeling objective. We observe similar pretraining performance to the standard Transformer model. We then finetune our pretrained models on three tasks from GLUE (DBLP:journals/corr/abs-1804-07461)
and one sentiment analysis task, IMDB reviews
(maas2011learning). On these tasks, we find that our model performs comparably, or even slightly better, than the standard pretrained Transformer, while observing significant training and inference speedups.Model Architecture | Complexity per Layer | Sequential Operation |
---|---|---|
Recurrent | ||
Transformer, (vaswani2017attention) | ||
Sparse Tansformer, (child2019generating) | ||
Reformer, (kitaev2019reformer) | ||
Linformer |
The Transformer is built upon the idea of Multi-Head Self-Attention (MHA), which allows the model to jointly attend to information at different positions from different representation subspaces. MHA is defined as
(1) |
where are input embedding matrices, is sequence length, is the embedding dimension, and is the number of heads. Each head is defined as:
(2) |
where are learned matrices and are the hidden dimensions of the projection subspaces. For the rest of this paper, we will not differentiate between and and just use .
The self-attention defined in (2) refers to a context mapping matrix . The Transformer uses to capture the input context for a given token, based on a combination of all tokens in the sequence. However, computing is expensive. It requires multiplying two matrices, which is in time and space complexity. This quadratic dependency on the sequence length has become a bottleneck for Transformers.
There has been much prior literature on improving the efficiency of Transformers, especially the self-attention bottleneck. The most common techniques for model efficiency that can be applied to Transformers (some specific to Transformers, others more general-purpose) include:
Mixed Precision (micikevicius2017mixed)
: Using half-precision or mixed-precision representations of floating points is popular in deep learning, and is also widely used in training Transformers
(ott2019fairseq). This technique can be further improved through Quantization Aware Training (jacob2018quantization; fan2020training), where the weights are quantized during training and the gradients are approximated with the Straight-Through Estimator. This line of work is orthogonal to our approach, and we use mixed-precision training by default.
Knowledge Distillation (hinton2015distilling): Knowledge distillation aims to transfer the “knowledge" from a large teacher model to a lightweight student model. The student model is then used during inference. However this approach has drawbacks: It does not address speeding up the teacher model during training, and moreover, student models usually suffer performance degradation compared to the teacher model. For example, when distilling a 12-layer BERT to a 6-layer BERT, the student model experiences an average 2.5% performance drop on several benchmark tasks (sanh2019distilbert).
Sparse Attention (child2019generating): This technique improves the efficiency of self-attention by adding sparsity in the context mapping matrix . For example, the Sparse Transformer (child2019generating) only computes around the diagonal of matrix (instead of the all ). Meanwhile, blockwise self-attention (qiu2019blockwise) divides into multiple blocks and only computes within the selected blocks. However, these techniques also suffer a large performance degradation, while having only limited additional speed-up, i.e., 2% drop with 20% speed up.
LSH Attention (kitaev2019reformer): Locally-sensitive hashing (LSH) attention utilizes a multi-round hashing scheme when computing dot-product attention, which in theory reduces the self-attention complexity to . However, in practice, their complexity term has a large constant and it is only more efficient than the vanilla transformer when sequence length is extremely long.
Improving Optimizer Efficiency: Microbatching (huang2019gpipe) splits a batch into small microbatches (which can be fit into memory), and then separately runs forward and backward passes on them with gradient accumulation. Gradient checkpointing (chen2016training)
saves memory by only caching activations of a subset of layers. The uncached activations are recomputed during backpropagation from the latest checkpoint. Both techniques trade off time for memory, and do not speed up inference.
As we’ve noted, most common techniques have limitations in reducing both the training and inference time/memory consumption, we investigate how to optimize the self-attention layers and introduce our approach next.
In this section, we demonstrate that the self-attention mechanism, i.e., the context mapping matrix , is low-rank.
We first provide a spectrum analysis of the context mapping matrix . We use two pretrained transformer models, RoBERTa-base (12-layer stacked transformer) and RoBERTa-large (24-layer stacked transformer) (liu2019roberta) on two tasks: masked-language-modeling task on Wiki103 (merity2016pointer) and classification task on IMDB (maas2011learning). In Figure 1
(left), we apply singular value decomposition into
across different layers and different heads of the model, and plot the normalized cumulative singular value averaged over 10k sentences. The results exhibit a clear long-tail spectrum distribution across each layer, head and task. This implies that most of the information of matrix can be recovered from the first few largest singular values. In Figure 1(right), we plot a heatmap of the normalized cumulative singular value at the 128-th largest singular value (out of 512). We observe that the spectrum distribution in higher layers is more skewed than in lower layers, meaning that, in higher layers, more information is concentrated in the largest singular values and the rank of
is lower.Below, we provide a theoretical analysis of the above spectrum results.
Based on the definition of the context mapping matrix , we can write
(4) |
where is an diagonal matrix. The main idea of this proof is based on the distributional Johnson–Lindenstrauss lemma (lindenstrauss1984extensions) (JL for short). We construct the approximate low rank matrix as , where with i.i.d. entries from . We can then use the JL lemma to show that, for any column vector of matrix , when , we have
(5) |
For more details, refer to the supplementary materials. ∎
Given the low-rank property of the context mapping matrix , one straightforward idea is to use singular value decomposition (SVD) to approximate with a low-rank matrix , as follows
(6) |
where , and are the largest singular values and their corresponding singular vectors. Based on the results in Theorem 1 and the Eckart–Young–Mirsky Theorem (eckart1936approximation), one can use to approximate self-attention (2) with error and time and space complexity. However, this approach requires performing an SVD decomposition in each self-attention matrix, which adds additional complexity. Therefore, we propose another approach for low-rank approximation that avoids this added complexity.
In this section, we propose a new self-attention mechanism which allows us to compute the contextual mapping in linear time and memory complexity with respect to sequence length.
The main idea of our proposed linear self-attention (Figure 2) is to add two linear projection matrices when computing key and value. We first project the original -dimensional key and value layers and into -dimensional projected key and value layers. We then compute an -dimensional context mapping matrix using scaled dot-product attention.
(7) |
Finally, we compute context embeddings for each head using . Note the above operations only require time and space complexity. Thus, if we can choose a very small projected dimension , such that , then we can significantly reduce the memory and space consumption. The following theorem states that, when (independent of ), one can approximate using linear self-attention (7) with error.
(Linear self-attention) For any and , if , then there exists matrices such that, for any row vector of matrix , we have
(8) |
The main idea of proof is based on the distributional Johnson–Lindenstrauss lemma (lindenstrauss1984extensions). We first prove that for any row vector of matrix and column vector of matrix ,
(9) |
where and , where with i.i.d. entries from and is a small constant. Applying the result in (9) to every row vector of matrix and every column vector of matrix , one can directly prove that, for any row vector of matrix ,
(10) |
by setting . This result does not utilize the low rank property of matrix (rank()=) and the resultant has a dependency on sequence length . We will further utlize the fact that rank()= to prove the choice of can be constant and independent of sequence length . For more details, refer to the supplementary materials. ∎
In Figure 2 (top right), we plot the inference speed of Linformer and standard Transformer versus sequence length, while holding the total number of tokens fixed. We see that while standard Transformer becomes slower at longer sequence lengths, the Linformer speed remains relatively flat and is significantly faster at long sequences.
Several additional techniques can be introduced on top of Linformer to further optimize for both performance and efficiency:
Parameter sharing between projections: One can share parameters for the linear projection matrices across layers and heads. In particular, we experimented with 3 levels of sharing:
Headwise sharing: for each layer, we share two projection matrices and such that and across all heads .
Key-value sharing: we do headwise sharing, with the additional constraint of sharing the key and value projections. For each layer, we create a single projection matrix such that for each key-value projection matrix across all head .
Layerwise sharing: we use a single projection matrix across all layers, for all heads, and for both key and value.
For example, in a 12-layer, 12-head stacked Transformer model, headwise sharing, key-value sharing and layerwise sharing will introduce 24, 12, and 1 distinct linear projection matrices, respectively.
Nonuniform projected dimension: One can choose a different projected dimension for different heads and layers. As shown in Figure 1 (right), the contextual mapping matrices in different heads and layers have distinct spectrum distributions, and heads in higher layer tend towards a more skewed distributed spectrum (lower rank). This implies one can choose a smaller projected dimension for higher layers.
General projections:
One can also choose different kinds of low-dimensional projection methods instead of a simple linear projection. For example, one can choose mean/max pooling, or convolution where the kernel and stride is set to
. The convolutional functions contain parameters that require training.In this section, we present experimental results for the the techniques described above. We analyze the techniques one-by-one and explore how they impact performance.
We first compare the pretraining performance of our proposed architecture against RoBERTa (liu2019roberta), which is based on the Transformer. Following devlin2019bert, we use BookCorpus (zhu2015aligning) plus English Wikipedia as our pretraining set (3300M words). All models are pretrained with the masked-language-modeling (MLM) objective, and the training for all experiments are parallelized across 64 Tesla V100 GPUs with 250k updates.
Effect of projected dimension: We experiment with various values for the projected dimension . (We use the same across all layers and heads of Linformer.) In the Figure 3(a) and (b), we plot the validation perplexity curves for both the standard Transformer and the Linformer across different , for maximum sequence lengths and . As expected, the Linformer performs better as projected dimension increases. However, even at for and for , Linformer’s performance is already nearly on par with the original Transformer.
Effect of sharing projections: In Figure 3(c), we plot the validation perplexity curves for the three parameter sharing strategies (headwise, key-value, and layerwise) with . Note that when we use just a single projection matrix (i.e. for layerwise sharing), the resulting Linformer model’s validation perplexity almost matches that of the the non-shared model. This suggests that we can decrease the number of additional parameters in our model, and consequently, it’s memory consumption, without much detriment to performance.
Effect of longer sequences: We evaluate the effect of sequence length during Linformer pretraining. In the Figure 3(d), we plot the validation perplexity for Linformer with , holding projected dimension fixed at . Note that as sequence length increases, even though our projected dimension is fixed, the final perplexities after convergence remain about the same. This further empirically supports our assertion that the Linformer is linear-time.
Model | SST-2 | IMDB | QNLI | QQP | Average | |
---|---|---|---|---|---|---|
512 | liu2019roberta, RoBERTa-base | 93.1 | 94.1 | 90.9 | 90.9 | 92.25 |
Linformer, 128 | 92.4 | 94.0 | 90.4 | 90.2 | 91.75 | |
Linformer, 128, shared kv | 93.4 | 93.4 | 90.3 | 90.3 | 91.85 | |
Linformer, 128, shared kv, layer | 93.2 | 93.8 | 90.1 | 90.2 | 91.83 | |
Linformer, 256 | 93.2 | 94.0 | 90.6 | 90.5 | 92.08 | |
Linformer, 256, shared kv | 93.3 | 93.6 | 90.6 | 90.6 | 92.03 | |
Linformer, 256, shared kv, layer | 93.1 | 94.1 | 91.2 | 90.8 | 92.30 | |
512 | devlin2019bert, BERT-base | 92.7 | 93.5 | 91.8 | 89.6 | 91.90 |
sanh2019distilbert, Distilled BERT | 91.3 | 92.8 | 89.2 | 88.5 | 90.45 | |
1024 | Linformer, 256 | 93.0 | 93.8 | 90.4 | 90.4 | 91.90 |
Linformer, 256, shared kv | 93.0 | 93.6 | 90.3 | 90.4 | 91.83 | |
Linformer, 256, shared kv, layer | 93.2 | 94.2 | 90.8 | 90.5 | 92.18 |
Thus far, we have only examined the pretraining perplexities of our model. However, we wish to show that our conclusions hold after finetuning on downstream tasks. We finetune our Linformer on IMDB (maas2011learning) and SST-2 (socher2013recursive) (sentiment classification), as well as QNLI (natural language inference) (rajpurkar2016squad), and QQP (textual similarity) (chen2018quora) We do the same with RoBERTa, 12-layer BERT-base and 6-layer distilled BERT. All of our models, including the Transformer baselines, were pretrained with the same objective, pretraining corpus, and up to 250k updates (although our Linformer takes much less wall-clock time to get to 250k updates, and was consequently trained for less time). Results are listed in Table 2.
We observe that the Linformer model () has comparable downstream performance to the RoBERTa model, and in fact even slightly outperforms it at . Moreover, we note that although the Linformer’s layerwise sharing strategy shares a single projection matrix across the entire model, it actually exhibits the best accuracy result of all three parameter sharing strategies. Furthermore, the Linformer pretrained with longer sequence length has similar results to the one pretrained with shorter length , this empirically supports the notion that the performance of Linformer model is mainly determined by the projected dimension instead of the ratio .
In Table 3, we report the inference efficiencies of Linformer (with layerwise sharing) against a standard Transformer. We benchmark both models’ inference speed and memory on a 16GB Tesla V100 GPU card. We randomly generate data up to some sequence length and perform a full forward pass on a multiple batches. We also choose batch size based on the maximum batch size that can fit in memory, and our memory savings are computed based on this number.
From Table 3, we see that even with and , Linformer has faster inference time and allows for a larger maximum batch size than the Transformer. As sequence length increases, the inference-time speed-up and memory savings are even more dramatic. We also plot inference times of both Linformer and Transformer on the 100 data samples in the top right of Figure 2.
Transformer models are notoriously slow to train and deploy in practice since their self-attention operations have time and space complexity with respect to sequence length . In this paper, we demonstrate, both theoretically and empirically, that the stochastic matrix formed by self-attention mechanism is low-rank. We further leverage this observation to propose a new, highly efficient self-attention mechanism. Through a combination of theoretical and empirical analysis, we demonstrate that our proposed approach is with respect to sequence length.
Our work focuses on making Transformers more efficient by introducing a mechanism that reduces self-attention to linear-time complexity. Potential positive impacts of efficient transformers include increasing the accessibility of our models, both for deployment on devices, as well as during training for research purposes. It also has potential impact on training transformer on images since we can support very long sequences. Furthermore, there are positive environmental benefits associated with decreasing the power consumption of models. As such, we see no immediate negative ethical or societal impacts of our work beyond what applies to other core building blocks of deep learning.
The main proof idea is based on the distributional Johnson–Lindenstrauss lemma (lindenstrauss1984extensions) (JL, for short), the following version is from (arriaga2006algorithmic).
Let be an matrix, , with i.i.d. entries from . For any , we have
(11) | |||
(12) |
For simplicity, we will omit the subscript for matrix , , , and . We will regard as , as and as . Define
(13) |
Based on the definition of contextual mapping matrix , we have
(14) |
where is an diagonal matrix such that
(15) |
Here we provide a constructive proof. Given any approximation error , define the following matrix.
(16) |
where be an matrix, , with i.i.d. entries from . Clearly the rank of matrix satisifies
(17) |
We further show that, when , we have that, for any column vector ,
(18) |
This concludes the theorem. For any row vector of matrix and any column vector of matrix , applying the JL Lemma, we can obtain
(19) |
Therefore, we have
(20) |
The above, step (a) is based on the union bound. The step (b) is utilizing the result of JL Lemma. Let , then theorem follows. ∎
Define and , where with i.i.d. entries from , is a constant with . We will first prove that for any row vector of matrix and column vector of matrix ,
(21) |
Based on the triangle inequality, we have
(22) |
The above, step (a) is based on the Cauchy inequality and JL Lemma in (11). The step (b) utilizes the fact that exponential function is Lipchitz continuous in a compact region. Then we can choose a small enough , i.e., such that
(23) |
The step (c) is based on the JL Lemma defined in (12).
Applying the result in (21) to every row vector of matrix and every column vector of matrix , one can directly prove that, for any row vector of matrix ,
(24) |
by setting . This result does not utilize the low rank property of matrix (rank()=) and the resultant has a dependency on sequence length . We will further prove the choice of can be constant and independent of sequence length .
Based on the fact that rank()=, we can find a row submatrix of matrix such that rank()=. Applying the result in (21) to every row vector of matrix and every column vector of matrix , and , we can obtain that, for any row vector of matrix ,
(25) |
Furthermore, define the matrix as
(26) |
We have that, for any row vector of matrix , .
The above, step (a) utilizes the inequality , where ( is the largest eigenvalue) is the spectrum norm of a matrix . The step (b) is based on matrix norm inequality , where is the Frobenius norm of matrix . The step (c) is based on the results of (24). ∎