SEA: Sparse Linear Attention with Estimated Attention Mask

Thank you for taking the time to review our work, we will address each of your comments below

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.

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.

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

Thank you for your recommendation. We have replaced figure 2 with a higher level representation of our method and moved the original figure 2 to the appendix. We are currently working on revising the notation in the method section during the discussion period.

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?

Indeed, in this experiment, the sequence length is short enough that the benefits of linear attention are not fully realized. Our motivation in showing this figure was to show that SEA preserves the performance of the quadratic teacher better than the other linear baselines.

To show the positive effects of our model when running on longer sequence lengths, we conducted an experiment where we took the OPT-125M trained models from figure 4 and interpolated the positional encodings to expand the number of encodings for the attention operations. This table has been added to the appendix A.8 as Table A.5. This result shows that SEA outperforms quadratic attention in both latency and memory cost. Additionally, the experiment result shows our model is much stronger than the other linear baselines after positional embedding interpolation. This result is interesting, as we think this shows that sparse attention masking helps to preserve the important attention relations while masking out unimportant relations.

Thank you again for taking the time to review our work. If you have any remaining concerns after our responses, we will do our best to answer them.

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

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.

Thank you for your continued discussion regarding our work. If you have any remaining concerns, please let us know before the end of the rebuttal period. Thank you.

Thank you again for taking the time to review our work. Regarding your comments below, we have posted an updated version of the text with the following modifications:

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.

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 **\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.

Update To Concern 3: in Figure 5 (updated figure number), 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?

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.

(continued in next post...)

Thank you for taking the time to review our work. We address your concerns below:

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.

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

Can one apply your algorithm for faster inference with no fine-tuning whatsoever?

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.

Could you provide a detailed explanation of Figure 2? Is the decoder (MLP CNN) trained from scratch, or can it be extracted from pre-trained models?

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.

In general, do you train all components depicted in Figure 2, or do you incorporate some parts from pre-trained models while keeping them fixed?

As SEA makes a big change to the attention mechanism, we require training all SEA weights during distillation while the teacher is fixed. However, we initialize the SEA weights from the pretrained quadratic teacher wherever possible in order to avoid training everything from scratch.

In addition to the two concerns raised above about the parameters of the model, we would like to note that as depicted in Figure 4(b) SEA converges much faster than the linear attention baseline models. We attribute this faster convergence due to the fact that Performer cannot directly distill the attention matrix, and Reformer can only distill the sparse range of values which are actually used in the Reformer attention calculation, which provides a much weaker training signal during distillation.

Thank you again for your time and efforts in reviewing our work. Please do let us know if there are any remaining concerns which have not been fully addressed in our responses.

We would like to thank you again for reviewing our work. We would like to draw your attention to a few updates in our current revision of the paper.

Updates to Concern 1: 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.

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.

Updates to Concern 2: Can one apply your algorithm for faster inference with no fine-tuning whatsoever?

• 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)

Thank you again for reviewing our work, if you have any remaining concerns please let us know in the remaining time of the rebuttal period. Thank you.

The latency results do not show a clear advantage [...] over baselines.

While our method does not have the lowest latency of all linear attention models, a clear advantage can be seen in Table 3 (revised to include OPT 1.3B) where the linear attention baselines deliver a roughly 2x increase in error, while SEA delivers performances within $0.4$% of quadratic attention in the worst case. Additionally, our method incurs a huge cost savings over quadratic attention as can be seen in Figure 7 which is logarithmically scaled.

The justification [...] 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.

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.

How difficult is it to tune the weights given to the different loss terms?

We did not perform extensive hyperparameter tuning on the loss terms, and in our experience SEA was robust to the intuitive values we chose. The hyperparameters on page 17 were chosen by:

1. Following common convention where KL divergence receives a lower weight as a regularization term.
2. Emphasizing the distillation loss terms over the downstream task loss terms to put more emphasis on distillation and avoid overfitting to the downstream task.

Wikitext is a relatively narrow dataset, how would the approximation handle more diverse datasets such as openwebtext?

Thank you for this suggestion. To conduct this experiment, we ran OPT-125M on OpenWebText. 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.

Thank you again for taking time and care in reviewing our work. Please do let us know if there are any remaining concerns which have not been addressed.

Thank you for your time in reviewing our work, we will respond to each of your comments below.

How would SEA compare to FlashAttention?

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.

The method is still quite complex, making it hard to deploy.

While there are many components to our final model, we implement them using open source HuggingFace implementations where possible. We believe the benefit of deploying SEA is huge due to the fact that our results show that SEA can maintain the performance of the quadratic teacher while incurring linear inference cost which can lead to enormous savings when deploying LLMs in a production environment.

The current trend of LLM's is centralized processing, (like ChatGPT). Therefore, inference cost saving is extremely important because it will reduce variables costs of serving the model. For example, currently ChatGPT costs 2$per 1M tokens to user [1]. Assuming this price is linearly dependent with the true cost, *OpenAI is spending approximately$21M per month approximately [2] on inference*. If the cost can be reduced by 1%, it translates to roughly $200K in savings per month, which is enough to pay several developers to improve the product.

Furthermore, dialog models are requiring longer context windows to improve quality (ChatGPT has 16k variants). Increasing the context window is very expensive if you use quadratic attention, even if using a memory efficient implementation (flash-attention). In figure 7 of the updated text we compare to FlashAttention, and SEA is significantly better than quadratic and FlashAttention in terms of latency.

In that context, the cost saving of our method is significant, because SEA reduces memory and latency while preserving performance and interpretability better than other linear methods.

For development cost in deployment, our method is quite cheap, because all developers have to do is replace the attention layer, and finetune the model with knowledge distillation. The training cost is also significantly lower than Performer and Reformer, as shown in Figure 4b. Moreover, when training our model, we would like to note that we do not require the huge pretraining dataset, and only require a much smaller finetuning dataset.

[1] https://topapps.ai/chat-gpt-pricing/

[2] https://www.digitaltrends.com/computing/chatgpt-cost-to-operate/

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

In summary, 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.

We would like to stress that this was done without a finetuned quadratic teacher and only self-distillation of our own model, and we were able to achieve a perplexity which is very close to that of the baseline vanilla model with better latency and memory consumption.

Thank you again. We hope this new experiment has addressed your further comments. If you have any remaining concerns, please do let us know in the remaining rebuttal time.

Dear Reviewer qYba,

Thank you for your reply to our comments, we will answer your further questions below:

Would SEA still approximate vanilla attention in terms of PPL?

The first experiment we posted on the longer context WikiText2 used a teacher which had only been pretrained and not finetuned for the task. Therefore the distillation target was not ideal for our model which led to the larger gap between the quadratic teacher and SEA.

It was difficult to decide on how to go about this during the rebuttal period due to the constrained time, we did not have resources to run much larger experiments on 8k or 16k sequences. We expect to need approximately 40GB VRAM GPUs (such as RTX 6000 Ada). But we only have A5000, so we can not do this at this point.

However, we were able to improve the result to show the strength of SEA in the 4k setting.

• We first trained longer than the previous response (previously 750 steps, now trained for 4k steps. For reference, Table 3 in our paper was 10k steps).
• We also use self distillation during training, because we do not have proper teacher for this setting (we think this technique should be researched further in the future).
• We use post training compression techniques introduced in this paper (dynamic-k and query skipping introduced in A.8) from the trained checkpoint for further compression.
• Below are the results (also shown in Table A.6 in the revised text)

(continued in next post...)