SeerAttention: Learning Intrinsic Sparse Attention in Your LLMs

Weakness 1 (Require of Training)

SeerAttention is the first to use learning representation method that resolve the challenging Long-context Attention Sparsity problem in a simple but effected manner. SeerAttention can be applied in the post-training stage with minimal cost or directly used in long-context fine-tuning with more training resources. However, we cannot agree that require of training/learning is a weakness of our work especially in this (ICLR) International Conference of Learning Representation conference. We understand the concern of cost but we want to emphasize that (1) a major contribution of this work is actually to make the training process very efficient (2) the capability of training makes our method a general and efficient methods, which have more potential to perform well in the long run.

1.We make the training of SeerAttention very efficient

• Our customized kernel make training cost as low as inference in post-training. In Post-training usage, only the AttnGates are trained, and the original model’s weights do not require gradient. The training cost is similar to inference thanks to our customized kernel to generate ground truth (explained in fig3). Without this kernel, no matter using torch naive implementation or old Flash-Attention, you can not avoid the quadratic memory consumption to save the attention map to train the AttnGates.
• The additional RoPE in AttnGate enable its context-length extrapolation ability. First, to clarify, a long-context LLM does not require training with data that is as long as its context window size. In most of the cases, a model is trained with both long and short context data and scaling RoPE can also help to increase the context length. In SeerAttention, another key contribution we made is to add the additional Block RoPE in AttnGate. This RoPE helps the AttnGate to learn the block-level positional information and enable its context-length extrapolation ability similar to original LLM. In other word, we do not necessarily need to train a 128k Model using 128k length data. Here is a detailed experiments on Post-training cases with 500 steps different training sequence length on Llama-3.1-8B with global batchsize 16. The GPU hours is calculated as num_GPUs x num_hours.

With a single 4x or 8xA100 machine, it only takes several hours in Post-training . In addition, the Fine-tuning experiments with YaRN (32k batch size 8) in our original paper only takes 32 GPU hours.

Here we also show the results of different training and evaluation context length to demonstrate that the block RoPE in AttnGate actually learns the position information. The table shows the PPL results on Llama-3.1-8B-Instruct with 70% sparsity. As shown in the table, you still can get very good performance with training on shorter data like 16k and 32k.

In summary, despite our method requires training, with our customized kernel and additional RoPE design in AttnGate, it’s very efficient to train AttnGate in post-training setting with your original model untouched. This is different from Reformer, Linear attention or Mamba-like approaches that requires you to retrain the whole network from scratch.

2. Concern with LoRA

To clarify, MInference and MoA do not need training, but they require offline calibration as shown in previous table. Thus, the situation of LoRA is the same. For different LoRA models, they still need to redo the calibration process. However, as a training-based method, SeerAttention can be co-trained with any LoRA or Fine-tuning scheme just like our long-context extension experiment. This means that your model can natively obtain the sparse capability during your LoRA or Fine-tuning, which is bound to deliver the better accuracy performance.

Weakness 1 (Require of Training )

3. Try Quest like reducing in Prefill

Quest is mainly designed for per-Q decode. To fit in our prefill setting, we use a Quest-like design by replacing per-Q prediction with avgpooling of Q similar to SeerAttention. A major problem with Quest is that it uses element-wise max operation between QavgKmax and QavgKmin before reduction (summation). When hidden dimension being 128, this element-wise operation costs 128x more GPU memory compared to SeerAttention as we directly do matmul. Thus, Quest-like method is super inefficient in prefill and it gets OOM on 64k and 128k test on a single A100. We also tried adding a similar block RoPE and it improves the performance a little bit. However, it still fails to deliver reasonable accuracy. Here is the results on Llama-3.1-8B-Instruct and PG-19 with 50% sparsity.

4. Generalizability

First, we think the Table 1’s 128k results is not evidence that SeerAttention fails at 128k. SeerAttention still get very good results on sparsity ≤ 80%. The degradation of 90% sparsity is mainly due to the fact that we apply 90% sparsity on every single head. However, MInference allocates different sparsity ratios on different head with an averaged of 90%. We acknowledge that different head should be treated differently to get better performances. Nevertheless, a more sophisticated sparsity allocation method can be explored in the future and it is orthogonal to our current story.

On the other hand, we cannot agree that training-based method does not generalize well. In our experiment, we use a single pre-trained dataset (Redpajama) and test all the downstream benchmark without seeing any generalizability issue. On the contrary, previous sparse attention methods mostly rely on human observation, which is hard to serve as a general and scalable solution.

Weakness 2 (Memory complexity)

The output of AttnGate is quadratic. However, when block size is 64, it’s only 1/4096 of the original size. Also to clarify, our TopK is applied per-row bases instead of global TopK on the 2D map. Since we only need [Seq/Block, TopK] shape tensor to store indices, the memory related to TopK is much less than your calculation. (We will add more clarification in the paper.)

As the computation of AttnGate is similar to self-attention (without the multiply with V part), if memory is really an issue, we can apply exactly the same tiling strategy as FlashAttention in AttnGate. In such case, for each local block, you can first do local TopK and then merge as global TopK in the end.

Weakness 3 (Limited methodological contribution)

Thanks for the suggestion. We have discussed many meaningful insights in ablation study. We will consider moving these earlier so that readers can better understand the contribution.

In terms of TDSA, it looks like a similar approach at first glance. However, it did not solve any challenges related to long-context LLM because it reduces on hidden dimension instead of sequence dimension. For a model like Llama-3-8B, the hidden dimension is only 128 and there is not much space to squeeze. A small hidden dimension may also hurt the utilization of TensorCore in modern GPUs. However, the sequence dimension can be as large as 128K or longer. Actually, we have done an experiment using a modified version of TDSA for long context scenario, which first reduce hidden dimension and reshape the tensor in sequence dimension.

q = linear(q) # [b,h,s,d] to [b, h, s, 4] reducing hidden dim to 4 
q = q.reshape(b, h, s/Block, 4 * Block) # change the reduction to seq dim similar to us
#… linear + rope similar to us

Here are the results on this modified TDSA on PG-19 with identical post-training setting as SeerAttention. It fails to deliver reasonable performance similar to Quest.

In general, we are the first to systematically formulate and solve the Long-context sparse attention problem using a learning-based method. Similar to MoE that helps scaling the model size using learnable gates, we aim to scale the context length by using learnable AttnGate. We provide a comprehensive methodology study on:

1. How to formulate the AttnGate to address sparse long-context attention?
2. What’s the ground truth? How to obtain it efficiently by avoiding O(N^2) memory explosion?
3. How to choose Pooling methods?
4. How to retain the block-wise positional information after pooling?
5. How to train in both Post-training (only train AttnGate) and Fine-tuning settings (train AttnGate and original weights)?

In summary, we have tackle many methodological or practical challenges to really make the design works.

Weakness 4 (Limited evaluated task)

Thanks for the suggestion of adding more downstream test. We have added experiments on ruler benchmark and compare with MInference under the same Post-training setting in the paper (Llama-3.1-8B-Instruct). In this experiment, we demonstrate that SeerAttention performs better than MInference in both accuracy and speedup. We found that MInference spends a huge amount of time doing online search of per-head sparse index on complex tasks. While our method is more efficient and effective than MInference.

Question 1 & 2 are answered in previous section

Question 3 (Block Size Experiments)

Here we add a simple experiment on PPL (PG-19) with block sizes 32, 64, and 128. The results come from post-training setting Llama-3.1-8B-Instruct with sparsity 70%.

The results show that block 64 and 128 are similar and block 32 are slightly worse. We think this might be the fact when block size is larger, the sentence level information can be better preserved. However, it might be strongly related to what tasks is evaluated. We will try more experiments in our revision should you find this experiment critical.

Weakness 1 (Improve figures of Section 3.1):

Thanks for the suggestion, we have added a dataflow figure similar to Figure 2 that explains the AttnGate operations (see Figure 11 in appendix in our rebuttal revision pdf).

Weakness 2 (FlashAttention symbol unexplained):

Thanks for the suggestion, we will add more explanation of related symbol in sec 4.2. We will keep faithfully the symbols used in FlashAttention such that readers can always refer to the original paper.

Weakness 3 (The meaning of "max and min pooling on K matrix" is not clear):

Thanks for pointing out this issue. We will improve the wording and figure caption. The pooling operation happens in the Seq dimension. For example, original Q tensor shape is [Batch, Head, Seq, Dim], the pooled result will be [Batch, Head, Seq / Block, Dim]. When more than one pooling method are used like Max+Min pooling in K, the tensor will first be concated in hidden dim before the linear layer ( [Batch, Head, Seq / Block, Dim * 2] in this case).

Weakness 4 (Rounding Issue when Seq not evenly divided by B):

When seq not evenly divided by B, the standard way is to use padding. In fact, the pooling function in torch can automatically support “padding” by using ceil_mode = True.

In current work, we focus on accelerating Prefill stage and leave decode for future exploration as mentioned in discussion section. Since Prefill and Decode are distinct two phases, many previous works also only focus on one side. For example, MInference merely focuses on Prefill, while Quest, H2O only focus on Decode.

Our work has potential to also be applied to Decode phase (doing per-Q token AttnGate without pooling). In general, the distribution of average pooling on Q tensor is similar to each token’s Q vector. As a result, we can enable similar training strategy to train AttnGate with generated tokens on the same linear layer’s weights. This is left for future work.

Weakness 5 (Pseudo code of Inference Kernel):

Thanks for the suggestion. We will add pseudo code of our inference kernel as well. A very simple version will be like this:

Input: Q, K, V; Output: O, A
for i from 1 to Tr
	Qi = load(Q, i) # Load Q block to HBM
	acc = zeros() # Output buffer
	for j from 1 to TopK # Only Topk block to compute
		Kj_index, Vj_index = k_index[j], j_index[j] # Spare BLock Index
		Kj, Vj = load(K, Kj_index), load(V, Vj_index) # Spare load KV
		sij = dot(Qi,Kj), rij = rowmax(sij)
		mi = max(mi, rij)
		p = exp(sij - mi)
		lij = sum(p, dim=1)
		alpha = exp(mi - rij)
		li = li * alpha + lij
		acc = acc * alpha + dot(p, Vj)
	acc /= li
	store(acc, O, i) # Store Acc to Output
	Return O

Weakness 6 (Inconsistencies in algorithm Symbol):

Thanks for pointing out this detail mistake, we will make it consistent.

Weakness 7 (Evaluation in other task):

To clarify, we do not discard tokens. We skip unimportant blocks in Attention computation, which is different than completely drop this token in the input. We also add new experiment on ruler, which contains various types of needle-in-a-haystack test. Here is the average score under llama-3.1-8B-instruct under post-training setting:

The results shows that our method is better than MInference and StreamingLLM (mentioned by other reviewers).

Weakness 8 (Figure 8 meaning):

Figure 8 shows the performances of our kernel for training (the kernel that generate the max-pooled attention map as ground truth). This kernel is not used in inference, so no concept “prefill” or “decode”. We use figure 8 to show that our training process is very efficient so that our model can be trained end-to-end by using this kernel. This makes the overall post-training cost as low as inference.

Besides, the inference kernel results lie in fig5 fig6 and table4. We will highlight “training/inference” our figure title and caption for clarity.

Weakness 9 (What training data is used):

We have mentioned in Sec5 line 280 that we use Redpajama dataset, which is the training dataset we use in all experiments (post-training and fine-tuning). Redpajama dataset is widely used in other long-context training works like LongLora and Longrope. We did not use any training data related to specific benchmark or test. For future work, we believe by doing SFT-like training, SeerAttention can achieve better performance in instruct-following tasks and also enable the sparse decoding ability.

Weakness 10 (Merge Table 1 and Table 4):

Thanks for the suggestion. We will consider merging the two tables. However, the resulting table might be huge and unclear to reader. Separating the table into accuracy and efficiency evaluation might better fit our current flow.

If you combine the orginal results, you can find that we already outperforms SOTA methods (MoA and MInference).

• MoA: MoA needs to predefine sparsity ratio. Our sparsity=0.5 results are all higher than MoA's sparsity=0.35 results.
• MInference: It does not support user-defined sparsity. It directly searches a sparsity for you. As we explain before, short context data have less sparsity. In order to preserve accuracy, MInference only use very small sparity. For 0-4k, it uses an avg of 0.06 sparsity and our sparsity=0.1 is higher than MInferece (55.91 > 55.23). For 4-8k, its avg sparsity is 0.25, our sparity=0.25 is better than MInference (54.09 > 53.87). For 8k+, it uses an avg of 0.45 sparsity, our score is also higher when sparsity=0.5 (52.43 > 52.18).

In summary, under higher or similar sparsity ratio. We achieve better accuracy results. In terms of speedup, please refere to Fig6, where we show that our kernel speedup is better than those works under different sparsity and context length.

Weakness 1 (Limitations in Related Work Discussion):

Thanks for pointing out these related works. We will add more discussion in our revision.

Weakness 2 (Limited Context Length Evaluated):

Thanks for the suggestion of increasing the evaluation length. Since you mentioned StreamingLLM can extend the context length of a model without training, we perform an additional needle-in-a-haystack experiments in ruler benchmark and found that StreamLLM can not perform well in this challenging task. Here is the average score on Llama-3.1-8B-Instruct under post-training setting.

This is mainly because that not all attention heads demonstrate this attention sink pattern in StreamingLLM, which highlight the advantage of our method that is agnostic to any pre-defined sparse pattern. We did not include [1,2] in our results comparison because they are relatively earlier works and our baseline MInference and MoA already surpass them.

In terms of 1M context length, it requires additional system optimization like memory offloading, tiling, and etc to avoid OOM on single A100, which is orthogonal to our current study. However, we did a simple test with up to 0.5 million context length using gradientai/Llama-3-8B-Instruct-Gradient-1024k model. Here is the PPL results on PG19 (filtered with length ≥ 0.5M).

Due to the time limitation of rebuttal, we can not fully explored the optimization of training and system implementation of 1M model. However, we think there are great opportunities for SeerAttention to apply in longer context models.

Weakness 3 (Deployment Cost of the Proposed Approach):

To clarify, methods like MInference and MoA do not need training, but they also require offline calibration. In SeerAttention, we have made two contributions to make our deployment cost low. On the one hand, we customize a training kernel to generate ground truth, which avoid the quadratic memory issue of storing attention map (shown in Fig3). On the other hand, we add a block-level RoPE in AttnGate to enable context-length extrapolation ability of AttnGate.

Here is the GPU hours (num_GPU * num_hours) statistics. For SeerAttention, we use a batch size of 16 with 500 training steps.

As for fine-tuning experiment with YaRN in the paper, it takes 32 GPU hours to train with 32k data with batch size of 8 for 400 steps.

We show that our block-level RoPE design in AttnGate enables the ability of to learn the block postition information. This means that you do not have to use data as long as your context window to train your AttnGate. Here is the PPL results with different training context length with and without block RoPE in AttnGate (a more detailed version of Fig9).

In summary, you can achieve reasonable performance by only spending a small amount of training resources.

Weakness 4 (Fig8 clarification):

Fig8 reflects the performance kernel that we use to generate the 2D maxpool attention map kernel for training AttnGates as illustrated in Fig2(b). For clarification, this is not the kernel to run AttnGate, but the kernel to generate the ground truth to train AttnGate. We want to use Fig8 to show that with the help of this kernel, our training process can be very efficient compared to manual torch implementation. This is because we modify flash-attention kernel and do not need to explicated generate the O(N^2) full attention map.

Memory overhead in inference:

The major overhead of our methods lies in the AttnGate outputs because it’s the only quadratic part. For a layer with 32 heads like llama-8B, the extra tensor shape is [32, seq/B, seq/B]. For a model with 128K seq length. The extra cost is only about 0.25GB. For very extremely case like 1M model, you can tile the AttnGate computation similar to FlashAttention to save memory if necessary.

Memory overhead in training:

The memory cost of post-training method is quite small, since we do not need to compute the gradient of original weights and LM loss (logits). You can easily train with 128k sequence data without any fancy trick to save memory like tensor parallelism or pipeline parallelism on a single A100.

In terms of fine-tuning, we also customized a block sparse flash-attn kernel for backward (not discussed in the paper). As a result, the major memory consumption part is similar to dense fine-tuning baseline. The overhead of AttnGate is similar to inference case mentioned above.

Question 1 (Theoretical or Intuitive Explanation for Pooling Selection):

We have provided a short discussion sec 6 (Pooling Ablation). We think that K prefers max min while Q prefers avg can be related to a well-known phenomenon in Quantization field that K has more outliers than Q and V. As a results, we think the outliers in K might carry more information and thus are better preserved by max-min pooling.

Question 2: (Row-Normalization of Max-Pooled Attention Map):

The Row-Normalization of max-pooled Attention map is mainly used to adhere to the output distribution of AttnGate in training where softmax is used (row sum = 1). Currently we found this ground truth works well and have not tried other methods. We tested by using this generated ground truth to run different benchmarks under high sparsity ratios and it gives nearly no accuracy degradation. This means that as long as the trained AttnGate can faithfully learn from the ground truth, it will perform well.

On the other hand, you can still combine AttnGate with other method like StreamingLLM that always requires it to attend to first token. However, we did try this before but did not see much accuracy gain. In fact, those head with distinct streaming pattern can be learned easily by our AttnGate as shown in Fig7.

Question 3: (Compatible with MoE)

Since MoE works in MLP layers and SeerAttention works in Attention layers, we can definitely combine with each other. Here are the PPL results on Mixtral-8x7b MoE model, which validate that SeerAttention can be effectively applied to MoE.

Weakness 1 (How to choose sparsity?):

Thanks for the suggestion. Currently, we expose the value of K (TopK) to users, allowing them to adjust it freely based on their targeted accuracy-efficiency tradeoff. We can incorporate offline calibration methods that search for a “recommended” sparsity ratio allocation strategy based on input sequence length or complexity of input question. For example, in the newly added MMLU test below, we use a simple linear strategy based on input sequence length. However, more sophisticated strategies with per-layer/head adjustments are left for future exploration.

It should be noted that previous works like MoA and MInference cannot adjust the sparsity ratio as flexibly as SeerAttention. For example, MoA and MInference need to search for a pre-defined sparsity pattern offline, which do not support changing the sparsity ratio by user during runtime. In SeerAttention, since the output of AttnGate provides a “soft” score for each block, it has the ability to apply different TopK ratios.

Question 1 (Why RoPE in AttnGate?):

Please refer to section 6 line 482 and figure 9, where we did a RoPE ablation study that address your question. If we directly use RoPE’d QK, the AttnGate loss the ability of context length extrapolation.

Here is a more detailed version of the study on Llama-3.1-8B-Instruct:

Results show that our AttnGate can learn block-wise postition information effectively with the block-wise RoPE.

Question 2 (Currently use MSE loss, how about e2e LM loss?):

Yes, the MSE loss can help the AttnGate converge fast and we also believe using end-to-end LM loss have better potential than MSE. However, this requires additional effort such as making operation like TopK and attention mask differentiable, which is left for future exploration.

Question 3 (What’s the sparsity pattern):

Please refer to Sec6 Figure 7, where we visualized some AttnGate’s outputs. In fact, without any human prior knowledge, it can automatically infer common sparsity pattern you mentioned: sliding window, attn sink, vertical & slash, and etc. This highlight SeerAttention’s ability as a more general and scalable method than manual sparse feature engineering as the sparsity pattern may change over different models architecture, training data, context and etc.

Question 4 (MMLU experiment):

We have done additional experiment on MMLU, which is a short context-length benchmark (input sequences range from 200 to 3000 tokens). In short input like this, attention layer is not the bottleneck in inference and enforcing large sparsity may hurt the accuracy. As a results, we use a very simple linear sparsity allocation method: input with seqlen 3000 uses 80% sparsity and no sparsity when input seqlen is 0 (the equation will be sparsity = 0.8/3000 * seqlen). Under this setting, the MMLU (5 shots) test results got similar performance as dense baseline.

Q1 More about in AttnGate's RoPE

Here is a more detailed analysis similar to fig9 in the paper. We train and evaluate different context length. It turns out that without the additional RoPE in AttnGate (Directly use the RoPE'd Q K), the evaluation length are limited (overfitting to training length). We think this is because the RoPE properties are damaged after pooling in sequence dimention. And using the non-RoPE'd QK for pooling with the additional block-RoPE in AttnGate can solve this problem.

Q2 Visualization of Sparse Pattern

Thanks for the suggestion, we will add more visualization in the appendix should it get accepted. However, a 8B model has 32x32 =1024 heads, we can not draw all of them. We will open-source the visualization code to help the community to better study the attention mechanism.

Here are some empirical findings that we have observed:

• A pattern of a head might change wrt different input. This is the reason that we can not simply categorize whether a head is pure streaming or vertical head at this stage. It also highlights our method that might be more adaptive to different pattern on a single head.
• For each layer, there is a portion of heads that demonstrates more random/dynamic patterns than the others.

In summary, we can add a small section in appendix to discuss some empirical findings. And we will open-source the visualiation code to allow people studying the attention patterns.

Thank you for the reply. It answers some of my questions. I have raised my score.