Paper: SEA: Sparse Linear Attention with Estimated Attention Mask (ICLR-2024)

Reviewer comment (Experiments):
One of the motivations of the paper is that other linear attentions cannot distill the learned attention patterns, and hence need to train from scratch. However, the authors in the paper still need to train their Performer and Decoder from scratch. I haven't seen any discussion about the inherent cost of doing that. Intuitively, it should be cheaper than training from scratch, but can you point me to the text (or elaborate in a new discussion) about how expensive it is to do this training?

Author response:
Thank you for pointing this out. We included Figure 4(b) which shows the validation curves for distillation using Performer, Reformer, and SEA. We can see that **SEA converges extremely fast** due to the direct distillation of the attention matrix. We have added the discussion on this point to section 4.1 in the revised text. 

Additionally, we added an experiment in Table A.6 which evaluates SEA on the OpenWebText dataset. **In this experiment, the dataset size is only $0.064$\% (less than one tenth of one percent) the size of the pretraining dataset**, and **SEA is still able to maintain competitive performance with the quadratic baseline** while Performer delivers a much worse perplexity.

|  | PPL $\downarrow$ | Memory $\downarrow$ | Latency $\downarrow$ |
|-|-|-|-|
|Vanilla|$19.82$|$408$|$4.88$|
|Performer|$61.20$($+41.38$)|$51$|$1.21$|
|SEA|$22.64$ ($+2.82$)|$187$| $6.76$|

Reviewer comment (Writing):
This is my subjective view, but the paper is extremely dense and hard to follow. I'd recommend reducing the notation significantly and moving most of the details to the appendix. I appreciate the figures, but overwhelming them with notation does aid my understanding of your method in the current version of the paper. I recommend: 1) simpler figures with less notation that can clearly explain your method conceptually; 2) moving a lot of the notation in the text for the appendix.

Author response:
Thank you for this suggestion, we have **removed some of the dense notation from the method section (section 3) and moved some minor details to the appendix**. 

- All updated text is displayed in a $\textbf{\textcolor{blue}{blue color}}$.
- We have **added higher level explanations and motivation** to the paragraphs describing the attention matrix estimation, CNN decoder, and top-k selection to make it easier for readers to follow from a high level. 
- We have also added a **visualization of the Performer output and that of the CNN decoder (figure 3)** in order illustrate why the **CNN decoder is necessary for forming the compressed attention matrix.** 
- We have moved Table 4 to the appendix (Table A.7) to allow space for the explanations previously mentioned.
- We have also moved the discussion related to our novel sparse format FlatCSR to the appendix A.12, and briefly introduce it and its benefit in the main text.
- We have **updated Figure 2 to be more intuitive, and moved the previous version of Figure 2 to the appendix A.10 as a more detailed illustration of SEA.**
- We have **moved some notation related to lower level details of the implementation to section A.5.**

Reviewer comment (Evaluation):
It seems to me that some of the results are a bit underwhelming. For example, in Figure 4, panel (a), right figure, what is the motivation to use your method, since I could use a Vanilla transformer and achieve an on-par accuracy with much less latency?

Author response:
*This is an update to the original experiment which was posted in our previous response, as we were able to improve the result substantially.*

We would like to add that the **experiment conducted in Figure 5a held the context length at 2k.** However, our model shows **more savings as the context length increases, as depicted in Figure 8**. Therefore, we performed an additional experiment with a **longer context length ($T = 4096$)** on WikiText2. 

For this experiment, we **did not use a finetuned quadratic teacher, and instead resorted to self-distillation** from the full attention matrix within our SEA model. We utilized the **dynamic k** described in Section 4.3, and **query skipping** described in Section A.8. The results can be seen in the text as Table A.6 and also in the below table.

**Table A.6 Perplexity (latency (ms)/ memory (MB)**

**X-axis:** dynamic value of k, **Y-axis:** number of query-skip rows 
**Color:** Redder --> worse | Greener --> better

|   |96|104|112|120|128|
|---|---|---|---|---|---|
|$\textbf{16}$|$\textbf{\textcolor[RGB]{189, 67, 59}{23.43}}$($\textbf{\textcolor[RGB]{71, 203, 21}{14.31}}$/$\textbf{\textcolor[RGB]{64, 212, 19}{448.37}}$)|$\textbf{\textcolor[RGB]{161, 99, 50}{23.00}}$($\textbf{\textcolor[RGB]{80, 193, 24}{14.70}}$/$\textbf{\textcolor[RGB]{94, 176, 28}{455.65}}$)|$\textbf{\textcolor[RGB]{137, 127, 43}{22.63}}$($\textbf{\textcolor[RGB]{96, 174, 29}{15.40}}$/$\textbf{\textcolor[RGB]{134, 130, 42}{463.06}}$)|$\textbf{\textcolor[RGB]{118, 149, 36}{22.34}}$($\textbf{\textcolor[RGB]{111, 156, 34}{16.06}}$/$\textbf{\textcolor[RGB]{174, 85, 54}{470.28}}$)|$\textbf{\textcolor[RGB]{103, 166, 31}{22.11}}$($\textbf{\textcolor[RGB]{125, 141, 39}{16.65}}$/$\textbf{\textcolor[RGB]{213, 40, 67}{477.59}}$)|
|$\textbf{8}$|$\textbf{\textcolor[RGB]{178, 80, 56}{23.26}}$($\textbf{\textcolor[RGB]{78, 195, 23}{14.61}}$/$\textbf{\textcolor[RGB]{64, 212, 19}{448.37}}$)|$\textbf{\textcolor[RGB]{150, 113, 47}{22.82}}$($\textbf{\textcolor[RGB]{93, 177, 28}{15.28}}$/$\textbf{\textcolor[RGB]{94, 176, 28}{455.65}}$)|$\textbf{\textcolor[RGB]{127, 138, 39}{22.48}}$($\textbf{\textcolor[RGB]{104, 165, 32}{15.74}}$/$\textbf{\textcolor[RGB]{134, 130, 42}{463.06}}$)|$\textbf{\textcolor[RGB]{109, 159, 33}{22.20}}$($\textbf{\textcolor[RGB]{118, 149, 36}{16.35}}$/$\textbf{\textcolor[RGB]{174, 85, 54}{470.28}}$)|$\textbf{\textcolor[RGB]{95, 176, 29}{21.98}}$($\textbf{\textcolor[RGB]{133, 132, 41}{16.97}}$/$\textbf{\textcolor[RGB]{213, 40, 67}{477.59}}$)|
|$\textbf{4}$|$\textbf{\textcolor[RGB]{163, 98, 51}{23.02}}$($\textbf{\textcolor[RGB]{94, 176, 29}{15.33}}$/$\textbf{\textcolor[RGB]{64, 212, 19}{448.37}}$)|$\textbf{\textcolor[RGB]{135, 129, 42}{22.60}}$($\textbf{\textcolor[RGB]{107, 162, 32}{15.85}}$/$\textbf{\textcolor[RGB]{94, 176, 28}{455.65}}$)|$\textbf{\textcolor[RGB]{115, 153, 35}{22.29}}$($\textbf{\textcolor[RGB]{121, 146, 37}{16.46}}$/$\textbf{\textcolor[RGB]{134, 130, 42}{463.06}}$)|$\textbf{\textcolor[RGB]{97, 173, 29}{22.02}}$($\textbf{\textcolor[RGB]{134, 131, 41}{17.03}}$/$\textbf{\textcolor[RGB]{174, 85, 54}{470.28}}$)|$\textbf{\textcolor[RGB]{85, 187, 25}{21.83}}$($\textbf{\textcolor[RGB]{147, 116, 46}{17.59}}$/$\textbf{\textcolor[RGB]{213, 40, 67}{477.59}}$)|
|$\textbf{2}$|$\textbf{\textcolor[RGB]{144, 119, 45}{22.73}}$($\textbf{\textcolor[RGB]{118, 149, 36}{16.35}}$/$\textbf{\textcolor[RGB]{64, 212, 19}{448.37}}$)|$\textbf{\textcolor[RGB]{117, 150, 36}{22.33}}$($\textbf{\textcolor[RGB]{130, 135, 40}{16.86}}$/$\textbf{\textcolor[RGB]{94, 176, 28}{455.65}}$)|$\textbf{\textcolor[RGB]{98, 172, 30}{22.02}}$($\textbf{\textcolor[RGB]{145, 118, 45}{17.50}}$/$\textbf{\textcolor[RGB]{134, 130, 42}{463.06}}$)|$\textbf{\textcolor[RGB]{82, 191, 24}{21.78}}$($\textbf{\textcolor[RGB]{159, 102, 49}{18.08}}$/$\textbf{\textcolor[RGB]{174, 85, 54}{470.28}}$)|$\textbf{\textcolor[RGB]{69, 206, 20}{21.58}}$($\textbf{\textcolor[RGB]{172, 87, 54}{18.64}}$/$\textbf{\textcolor[RGB]{213, 40, 67}{477.58}}$)|
|$\textbf{1}$|$\textbf{\textcolor[RGB]{121, 146, 37}{22.38}}$($\textbf{\textcolor[RGB]{159, 102, 50}{18.12}}$/$\textbf{\textcolor[RGB]{64, 212, 19}{448.35}}$)|$\textbf{\textcolor[RGB]{99, 171, 30}{22.04}}$($\textbf{\textcolor[RGB]{173, 86, 54}{18.68}}$/$\textbf{\textcolor[RGB]{94, 176, 28}{455.65}}$)|$\textbf{\textcolor[RGB]{81, 192, 24}{21.76}}$($\textbf{\textcolor[RGB]{187, 70, 59}{19.31}}$/$\textbf{\textcolor[RGB]{134, 131, 41}{462.98}}$)|$\textbf{\textcolor[RGB]{67, 208, 20}{21.55}}$($\textbf{\textcolor[RGB]{201, 53, 63}{19.92}}$/$\textbf{\textcolor[RGB]{173, 86, 54}{470.19}}$)|$\textbf{\textcolor[RGB]{64, 212, 19}{21.38}}$($\textbf{\textcolor[RGB]{214, 39, 67}{20.47}}$/$\textbf{\textcolor[RGB]{213, 40, 67}{477.54}}$)|

In light of this result, we think the motivation to use our model is clear. **SEA can produce a perplexity on longer context which is competitive (+4.47) relative to quadratic attention with a 3.55ms reduction in latency and a 1135.63MB reduction in memory usage.**

As noted in the general response above, we think **these savings are of great interest to anyone who wishes to deploy an efficient transformer in a production setting.**

Reviewer comment (Writing):
The paper's clarity and explanation of the algorithm's functionality are lacking, making it challenging to determine its applicability, particularly with regard to pre-trained models.

Author response:
We clearly mentioned that our method aims to distill the knowledge of the quadratic attention from a pretrained transformer into linear attention in the abstract and the introduction. In the introduction section, we clearly stated this in the bullet points that summarize our contributions.

 - (bullet point #1) We propose a novel, test-time linear attention mechanism (SEA) **which distills knowledge
from a pretrained quadratic transformer into a compressed estimated attention matrix**...

We are clearly stating that we provide a way to distill knowledge from pretrained quadratic transformers into a model with linear time and space complexity at inference time. From a high level point of view, the basic outline of our method is as follows:

1. Obtain a pretrained quadratic attention model.
2. Initialize SEA attention with $Q, K, V$ and MLP weights from the pretrained teacher model.
3. Perform distillation + finetuning on a target task.
4. Discard quadratic model and use SEA for linear inference 

As noted in the general response, we have added the following to the revised text:

 - Added **higher level explanations and motivation** to the paragraphs describing the attention matrix estimation, CNN decoder, and top-k selection to make it easier for readers to follow from a high level. 
 - **Updated figure 2 to show a higher level concept** of our model.
 - Added a **visualization of the Performer output and that of the CNN decoder (figure 3)** in order illustrate **why the CNN decoder is necessary for forming the compressed attention matrix.**
 - Moved the lower level detailed explanations of the implementation to appendix (A.5) 

Regarding the applicability of our method:

We have added another experiment on longer context on WikiText2 where we do not use a finetuned teacher model, and **only apply self-distillation on our model for finetuning.** 

From the initial result in our general response, we **improved PPL from 28.9 to 23.43, while achieving 14.31 ms (80.12 %) latency and 448 MB (28.28 %) memory compared to vanilla quadratic attention.** This is presented in Table A.6 of the revised text and in our general responses above.

Reviewer comment (Experiments):
Can your algorithm be applied to pre-trained models for faster inference without requiring fine-tuning? Can one apply your algorithm for faster inference with no fine-tuning whatsoever?

Author response:
Our method requires a distillation step from the quadratic teacher in order to get the linear attention benefits at test time, but we would like to note that **it does not require the huge dataset originally used for pretraining, but only a small surrogate dataset**. Additionally, as shown in Figure 4b, **the convergence rate is extremely fast** due to the direct distillation of the full attention matrix.  

To illustrate the scale of savings, we conducted an experiment on the OpenWebText dataset which has been included in our revision in Table A.6. In this experiment, we were able to distill the knowledge from the teacher into a linear model with minimal performance degradation with a dataset that is **0.064\% of the size of the pretraining dataset**. That is less than one tenth of one percent of the pretraining dataset. 

- Our model **only requires low-cost knowledge distillation for training SEA, which converges extremely fast compared to the baselines (please see figures 4b and A.5 for fast training curves)**
- **Finetuning requires only tiny subset of dataset compared to the pre-training from scratch (0.064% of the pretraining dataset size)**

Reviewer comment (Presentation):
Could you provide a detailed explanation of Figure 2? I'm having difficulty understanding this diagram. Is the decoder (MLP CNN) trained from scratch, or can it be extracted from pre-trained models?

Author response:
Regarding Figure 2, we have updated it to illustrate a higher level concept of our attention distillation framework, and we moved the original Figure 2 to the appendix. In Figure 2, we are attempting to illustrate the following steps. 

1. Performer + Decoder receives $Q, K, V$ and outputs a compressed attention matrix
2. The top values from the compressed attention matrix are selected and then decompressed via interpolation to construct a sparse attention mask.
3. With the attention mask, we utilize the original $Q, K, V$ given to performer to calculate the sparse attention matrix.

The decoder MLP and CNN are newly introduced components of SEA, therefore must be initialized from scratch. However the $Q, K, V$ matrices and the MLP's within the original quadratic teacher are used in the initialization of SEA. Therefore, **we re-use the parameters from the pretrained teacher model** wherever possible, while initializing the parameters for new components introduced by SEA, such as the decoder.

To further explain the role of the decoder CNN in addition to the writing in section 3.1, we used a CNN to estimate the compressed attention matrix since the compressed attention matrix can be thought of as an image with the width dimension compressed, which makes a CNN a natural choice over using the raw Performer output. Empirically, we also found the CNN to obtain better performance than the raw Performer when comparing the attention matrix distillation loss during training.

Reviewer comment (Experiments):
Comparison with FlashAttention [1]: It would be fair to add FlashAttention among the baselines. Especially, FlashAttention would also be competitive in terms of memory.

Author response:
We added a comparison to FlashAttention to the appendix in Figure 7 and A.7. **Flash attention's implementation allows for linearly scaling memory consumption, but the latency still scales quadratically**. Therefore, linear attention is crucial in cases where latency is a major concern, as it would be in the deployment of any large LLM. FlashAttention also cannot store the attention matrix, so if the attention matrix is needed for any kind of token importance analysis, such as token pruning, FlashAttention cannot be used.

Reviewer comment (Theory):
The justification of the method for autoregressive language modeling is unclear. As most causal models are used for sequence generation, sampling one token at a time, the need to sparsify the attention matrix at inference time is reduced.

Author response:
While in the single sequence case, it is true that the bottleneck tends to be in parameter loading and thread scheduling rather than attention matrix calculation. However, the current trend of production LLMs is to provide centralized processing where batches of inputs are masked for length and passed through the model. In this setting with long contexts and where parameter loading and thread scheduling is no longer the bottleneck, the sparse attention operation of SEA will start to show huge savings.

Reviewer comment (Experiments):
How would SEA compare to FlashAttention?

Author response:
We added a comparison to FlashAttention to the appendix in Figure 7 and A.7. **Flash attention's implementation allows for linearly scaling memory consumption, but the latency still scales quadratically**. Therefore, linear attention is crucial in cases where latency is a major concern, as it would be in the deployment of any large LLM. FlashAttention also cannot store the attention matrix, so if the attention matrix is needed for any kind of token importance analysis, such as token pruning, FlashAttention cannot be used.