Block-Attention for Efficient Prefilling

Authors: Dongyang Ma, Yan Wang, Tian Lan
Primary area: foundation or frontier models, including LLMs
Keywords: LLM, RAG, parallel context encoding, efficient language model

Abstract:
We introduce Block-attention, an attention mechanism designed to address the increased inference latency and cost in Retrieval-Augmented Generation (RAG) scenarios. 
Traditional approaches often encode the entire context in an auto-regressive manner.
Instead, Block-attention divides retrieved documents into discrete blocks, with each block independently calculating key-value (KV) states except for the final block.
In RAG scenarios, by defining each passage as a block, Block-attention enables us to reuse the KV states of passages that have been seen before, thereby significantly reducing the latency and the computation overhead during inference.
The implementation of Block-attention involves block segmentation, position re-encoding, and fine-tuning the LLM to adapt to the Block-attention mechanism. 
Experiments on 11 diverse benchmarks, including RAG, ICL, and general domains, demonstrate that after block fine-tuning, the Block-attention model not only achieves performance comparable to that of full-attention models, but can also seamlessly switch between the block and full attention modes without any performance loss.
Notably, Block-attention significantly reduces the time to first token (TTFT) and floating point operations (FLOPs) to a very low level. It only takes 45 ms to output the first token for an input sequence with a total length of 32K. Compared to the full-attention models, the TTFT and corresponding FLOPs are reduced by 98.7\% and 99.8\%, respectively. Additionally, in Appendix A, we elaborate on how Block-attention is applied in Game AI scenario and the substantial potential benefits it entails. We strongly suggest researchers in the gaming field not to overlook this section.

Decision: accept (poster)
Reviewer scores: 6, 5, 6, 6

Reviews and discussion:
<thread>
creation date:
1730599454366

thread_type:
official_review

<thread_object>
creator: reviewer_mnr9

<content>
summary:
This paper presents the Block-Attention mechanism aimed at optimizing inference efficiency in retrieval-augmented generation (RAG) scenarios. The authors assert that by modifying the attention calculation to process tokens in independent blocks, they can significantly enhance both accuracy and computational efficiency compared to traditional self-attention. The proposed method is evaluated against well-known RAG benchmarks, and the authors claim notable improvements in inference speed and resource consumption.

While the topic is pertinent, the contributions of the paper raise several concerns. The approach, although presented as innovative, relies on established principles of attention mechanisms without sufficiently distinguishing itself from previous work. Additionally, the experimental results, while promising, lack depth in exploring the nuances of the proposed method's effectiveness across varying conditions.

strengths:
1. The paper is generally well-organized, with clear explanations of the Block-Attention mechanism and its proposed advantages.
2. The investigation into improving inference efficiency in LLMs through Block-Attention is timely.
3. The reported gains in inference speed and reduced computational load are compelling.

weaknesses:
1. The concept of segmenting attention into blocks is not entirely novel and has been explored in various forms in the literature. The paper does not convincingly articulate how Block-Attention offers distinct advantages over similar approaches, such as sparse attention methods or hierarchical attention mechanisms.
2. The paper presents empirical results without adequately establishing the theoretical underpinnings of the Block-Attention mechanism. For example, the claim that "tokens in all blocks except the last are restricted to attending only to information within their own block" (Section 3.4) needs further elaboration on how this restriction impacts the model's representational capacity compared to full self-attention.
3.  The experimental section would benefit from a broader set of benchmarks and comparison against state-of-the-art methods beyond just self-attention. 
4. The authors should discuss how the model performs under different conditions, such as varying levels of noise in the retrieved passages or the diversity of user queries.

questions:
1. How do the authors justify the choice of block size in the Block-Attention mechanism, and what implications does this have for different types of queries?
2. Can the authors elaborate on how the performance of Block-Attention varies with input length and retrieval volume? Are there any thresholds beyond which the performance might degrade?
3. Given the reliance on fine-tuning, how do the authors plan to address potential overfitting issues, particularly in diverse RAG settings?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730615481789

thread_type:
official_review

<thread_object>
creator: reviewer_1dtj

<content>
summary:
This paper tackles the efficiency challenges in RAG. The key observation is that different paragraphs retrieved in RAG are independent to each other. Therefore, it is possible to perform independent self-attention within each paragraph, and only allow the final question query to attend to all previous passages. This will dramatically improve TTFT in long-context RAG. The authors identify two important steps to recover accuracy after converting traditional self-attention to block attention (proposed in this paper): positional re-encoding and fine-tuning. Their final results show no loss of accuracy and significant speedup in long-context RAG applications.

strengths:
1. The idea is simple and reasonable: different retrieved passages are not necessarily related to each other so the attention mask can be sparsified.
2. Positional re-encoding is an elegant way to solve inconsistencies in positional encodings.
3. Finetuning results show no loss of accuracy and significant speedup over self-attention baselines.

weaknesses:
1. The method requires fine-tuning, which limits its scalability.
2. The evaluation seems to be a bit weak. 
- For example, the method trains on TQA/2Wiki's training set and evaluates on their validation sets. The results on these two benchmarks are not zero-shot and are not very representative. 
3. It's unclear what is the overhead associated with positional re-encoding.
4. Can you construct some examples where different retrieved passages are related to each other? In this case, will the proposed method fail? For example, you might retrieve several chapters in a textbook, where later chapters are dependent on earlier chapters.

questions:
Please respond to my questions in "Weaknesses".

===

Post rebuttal: my concerns have been resolved and I will keep my score.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730693281253

thread_type:
official_review

<thread_object>
creator: reviewer_wuez

<content>
summary:
The paper proposes Block Attention for retrieval augmented generation (RAG). The idea is to confine attention within each retrieved document (a.k.a. block) and only let the query attend to all documents. In this way, the KVs for each document can be cached and reused if the document re-appears in multiple queries. The method also proposes to explicitly correct RoPE positional encoding before reusing the cached KVs. After fine-tuning, this technique can retain the accuracy of the baseline while speeding up the inference due to smaller attention span and avoiding the KV recomputation thanks to the block caching.

strengths:
- The method is simple and easy to implement: It only involves limiting the attention computation within each document in RAG scenarios and caching the KVs for each block for document re-use as proposed in PromptCache [Gim et. al, 2024]. The newly introduced RoPE rotation correction is also straightforward to implement.
- Local attention and prompt caching are sound techniques and suitable for RAG applications.
- The paper is clear and easy to follow.

weaknesses:
My primary concern is regarding the paper's contribution compared to the existing works and the need for a better experimental analysis to understand where this method stands among previous methods that share very similar ideas or how it improves them. The authors can find more detailed comments and suggestions below.

- Contribution and limitations compared to existing works

To me the main contributions of the paper to improve the efficiency of RAG can be summarized as 1) limiting the attention within each document 2) caching the KVs of the documents to reuse them 3) Adjusting the rotation in the RoPE position encoding based on the position of the reappearing cached document. Points 1 and 2 have been already proposed in prior works. E.g. PromptCache [Gim et. al. 2024] proposed the same techniques to speed up the RAG applications. This makes contribution 3, the main new addition to the existing works. However, according to the experiments reported in the submission, the RoPE position correction does not lead to any improvement over PromptCache (lines 353-356). To my understanding this makes fine-tuning the only source of improvement over PromptCache, which is an expected gain which comes with known challenges of such a fine-tuning.

I want to also point the authors to [1] below that is a closely relevant prior art which is currently missing from the draft. [1] also addresses the three aforementioned points above. Namely, it proposes limiting attention to attention blocks, caching the KVs for the documents to re-use them, and addresses the position encoding problems by overlaying the documents as parallel edges without requiring any fine-tuning. I would suggest the authors add a discussion on the advantages of the proposed method compared to [1] and also add controlled experimental comparisons to better highlight the contributions.

[1] Superposition Prompting: Improving and Accelerating Retrieval-Augmented Generation, Merth et. al.

* Experimental Results

The experimental section needs to be revised to clearly highlight the main sources of improvements and contributions between the proposed method and existing prior methods (under same settings, e.g. with or without fine-tuning). Based on the submission, comparisons with similar approaches are missing (see [1] above), additional speed up on top of PromptCache (if any) is not verified, and it is not clear to me if anything other than the fine-tuning is the source of accuracy improvement over PromptCache. The authors can find more detailed comments and questions regarding the experiments in the following section.

questions:
- Additional questions and comments:

1) To my understanding, the main technical contribution of the method compared to the previous approaches is the proposal to explicitly re-adjust the positional encodings. What is the practical advantage of the proposed RoPE encoding adjustment? Does it lead to any improvement on the final RAG task quality metrics over PromptCache? From the current manuscript I did not find an ablation showing its importance to improve prior arts. 

2) What is the comparison between the proposed method and PromptCache in terms of efficiency? How much does the proposed method improve PromptCache in terms of wall-time?

3) How does the method compare with respect to the Superposition Prompting paper above [1] both in terms of accuracy and efficiency?

4) The paper mentions fine-tuning on RAG-specific datasets necessary for the proposed approach. Expectedly, fine-tuning can improve the accuracy on the specific tasks at hand, but it is also important to maintain the generalizability of the pre-trained models. This can be difficult especially when the training recipes are not available, or they are instruction fine-tuned on non-available data. How does the performance of the reported models change on common benchmarks such as MMLU and HumanEval after the suggested fine-tuning?

5) The positional encoding re-adjustment is only discussed for RoPE. As a suggestion for further improvement, it would be nice to have a more thorough analysis on different positional encoding approaches.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730766077359

thread_type:
official_review

<thread_object>
creator: reviewer_2j6u

<content>
summary:
This paper introduces the Block-Attention method, an efficient approach tailored for RAG scenarios that enhances inference efficiency while preserving performance through fine-tuning. Block-Attention operates by dividing the input sequence into multiple independent blocks, each calculating its key-value (KV) states separately via self-attention. Only the final block can attend to previous blocks, allowing the user query to reference all prior retrieved documents. Evaluations on four RAG benchmarks reveal that Block-Attention, after fine-tuning, maintains comparable accuracy to traditional self-attention models (e.g., 68.4% vs. 67.9% on Llama3) and even slightly outperforms in some cases (e.g., 62.3% vs. 59.6% on Mistral). The efficiency benefits, measured in TTFT and FLOPs, grow with input length; at 32K input length, Block-Attention’s TTFT and FLOPs reach just 1.3% and 0.2%, respectively, of those in standard self-attention models.

strengths:
- Novel and practical idea
- easy to follow
- Simplicity: The solution is relatively simple to implement and can be integrated into existing LLM architectures.
- Impressive performance: this method reduces Time to First Token (TTFT) by up to 98.7% and FLOPs by up to 99.8%.

weaknesses:
### Needs for finetuning:  
The method requires additional fine-tuning, which might be resource-intensive for larger models. It would be better to explore a training-free approach with the proposed method.

questions:
Only tested on relatively small models (7B-8B parameters) - would it work as well on larger models?
</content>
</thread_object>

</thread>
<thread>
creation date:
1732078854527

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
A general comment to all reviewer's concerns about "Needs additional fine-tuning" and "generalizability & avoid overfitting

comment:
Thank you to all the reviewers for your valuable suggestions. I'm very glad to have the opportunity to engage in such an insightful discussion with you all on OpenReview. In this general response addressed to everyone, I'd like to briefly introduce the background of Block-Attention and the key points focused on in this research paper. I believe it can effectively address your concerns regarding regarding the issues of  "Needs additional fine-tuning" and "generalizability & avoid overfitting".

Firstly, any new attention mechanism requires fine-tuning, and Block-Attention is no exception. By way of analogy, transitioning from Alibi to RoPe also requires retraining the entire model from scratch. In real-world applications, the best approach to using Block-Attention is integrating it in the post-training stage of the foundation model. Specifically, for the LLM during this stage, we simultaneously adopt the self-attention mechanism and the block-attention mechanism to represent the training data. Regarding the division of blocks, some are manually divided during the annotation process, and others are divided according to rules. For example, in the RAG domain, each passage should be divided into one block; in the code domain, each function and class should be divided into one block; for few-shot tasks, each example can be divided into an independent block. The resultant model can adapt to both attention mechanisms simultaneously, and its results on common benchmarks are almost identical to those of the model without Block-Attention, without any performance loss. More importantly, after obtaining such a foundation model, we no longer need additional fine-tuning since the model has already completely adapted to the Block-Attention manner.

However, since there are currently no open-sourced foundation models that support block-attention, if we want to use the block-attention mechanism in downstream tasks to improve inference efficiency, we must conduct an additional fine-tuning. The adaptation cost in this aspect is inevitable. The difference only lies in whether this cost is borne during the training process of the foundation model or during the fine-tuning stage of the downstream tasks. Fortunately, this cost is not high. We don't need to change the attention-mechanism in the pre-training stage. It can be achieved simply by adjusting the attention masks in the downstream fine-tuning.

In this paper, due to confidentiality reasons, we are unable to release our internal SFT and Preference Optimization data. Therefore, in this submission, for the sake of re-implementation considerations, we focus on the RAG (Retrieval-Augmented Generation) scenario and only use publicly available dataset to train the model to verify the effectiveness of block-attention. The model trained in this way will surely have a decline in generalizability. However, please kindly note that this decline is caused by our experimental settings rather than by block-attention. We will also keep on encouraging certain open-source LLM teams, such as Llama, Qwen, DeepSeek, and Hunyuan, to incorporate the block-attention mechanism into their future models. Stay tuned!
</content>
</thread_object>

</thread>


Metareview:
metareview:
a) The paper proposes Block Attention for retrieval augmented generation (RAG) with multiple documents (blocks). In this work, each document is assumed to be independent from the others, therefore, it is possible to reduce computation by performing independent self-attention within each block, and only allow the final question query to attend to all previous blocks. In this way, the KVs for each document can be cached and reused. The method introduces an explicitly correction of the positional encoding and fine-tuning, to retain the performance of the baseline model, but much faster inference.

b) The paper is well presented and easy to follow. The method is simple and easy to implement. Positional re-encoding is an elegant way to solve inconsistencies in positional encodings. The reported gains in inference speed and reduced computational load are compelling.

c) The concept of separating attention into blocks for RAG is not novel. It is already considered in works such as PromptCache and Superposition Prompting. In particular PromptCache seems very similar to the proposed idea, but missing the adjustment of the positional encoding and the fine-tuning. 

d) The final decision for this paper was not easy. The paper has many similarities with PromptCache. Therefore the proposed idea is not novel and this should be clearer in the main paper. The main contributions are re-positional encoding and fine-tuning, which bring the method to the performance of the full self-attention, but with reduced computation. However, fine-tuning requires data and has the risk of over-specializing the model, losing generalization. Authors show results for in-domain and out-domain performance and the model seems to perform comparable to full self-attention also on out-domain. In addition, authors point out that block-attention could be used on the post-training phase, which would avoid the need of fine-tuning, although they did not provide results on that probably due to the training cost and time, but I do not see any reason why it should not work. 
Considering all of that, I believe that the proposed ideas of re-positional encoding for improving PromptCache by itself can be quite important for RAG and deserves publication as it can produce similar retrieval performance as full-attention but with a fraction of the computational cost. Said that, the authors should do a good job on adapting the paper to reviewers comments and this discussion and make clearer what are the actual contributions compared with previous work.