MoA: Mixture of Sparse Attention for Automatic Large Language Model Compression

[W1] Compression overhead seems high

The process of identifying the optimal sparse attention configuration appears to be complex, involving profiling the model, evaluating configurations, and pinpointing the best plan. This complexity may pose a barrier to practical implementation.

We appreciate the reviewer's concern regarding the overhead of our compression process. We quantify the overhead of each compression step in Table 13. For the Vicuna-13B model, the entire compression pipeline can be completed in 105 minutes. Even for the 70B model, the compression can be done within 9 hours of latency and 35 GPU hours. Importantly, once compression is achieved, the model can be utilized across various tasks without any additional overhead.

We also quantify the progressive overhead of our compression method for different model sizes in Table 14. The calibration dataset generation, profiling, and validation stages exhibit linear complexities with respect to both parameter size and dataset size. During the optimization stage, the complexity is determined by the number of attention heads, which typically scales with parameter size. Although this leads to an exponentially growing search space, advancements in MIP solvers allow for polynomial time solutions under many conditions [1, 2], as confirmed by our empirical results.

We would also like to emphasize our efficiency design considerations during compression, particularly in handling models with extensive input lengths and numerous parameters:

1. Elastic Attention Rules for Large Input Length: Our method searches for length-adaptive elastic rules instead of length-dependent sparse attention masks. As demonstrated in Figure 5, by compressing within 12k input length, the same compression plan effectively generalizes up to 256k, ensuring broad applicability without re-compression.
2. Efficient Optimization for Large Parameter Sizes: We use efficient white-box optimization for the compression pipeline. During optimization, each sparsification rule for an attention head is abstracted with only two float numbers: density and accuracy loss. This abstraction greatly simplifies the problem and formalizes it as Mixed-Integer Planning (MIP), whose solving is highly optimized with existing solvers [3, 4].

In conclusion, the empirical results, asymptotic analysis, strong generalization ability, and effective design choices ensure the efficiency of our automatic compression pipeline.

Lastly, the compression overhead can be further reduced with system-level optimizations. Currently, we use HuggingFace model parallelism during profiling, which does not support pipelining and results in low GPU utilization. By integrating pipelining and tensor parallelism, the total GPU time can be significantly reduced in future implementations.

[W2] Robustness for calibration dataset

The effectiveness of MoA seems to depend heavily on the availability of appropriate long-range contextual data, such as the MultiNews dataset. The performance gains might diminish without such data. How robust the method is against different long-range contextual dataset?

We appreciate the reviewer’s concern regarding the robustness of MoA with respect to the calibration dataset. MoA demonstrates strong robustness as long as the calibration dataset satisfies the long-dependency and model alignment properties discussed in Section 5. We conduct the robustness study detailed in Appendix B.3.1. The experiments show that MoA maintains strong performance across varied calibration datasets, regardless of whether the same or different datasets are used during the compression and testing phases. Additionally, we explore the possible reasons for this robustness in Appendix D.1 (Figure 14). The oracle experiment on attention patterns across diverse tasks (e.g., few-shot learning, chatting, and coding) reveals that, while task-specific details may vary, the overall attention span of each head remain consistent. This consistency likely explains the strong robustness of MoA across different calibration datasets and tasks.

[W3] Add more baselines

Overall I believe this paper is good at its idea and implementation. While the topic of efficient decoding of LLMs is very hot, it is expecte to compare with more competitive related methods in the experiments to demonstrate the effectiveness of the proposed method among its competitors.

We thank the reviewer and add the latest SnapKV [5] and PyramidKV [6] baselines. As shown in the updated Figure 4, MoA consistently achieves a superior throughput-accuracy trade-off across varied densities, outperforming all six baselines. Additionally, we have included the retrieval performance of SnapKV [5] and PyramidKV [6] on longer contexts in Table 9. MoA demonstrates 4% to 21% higher retrieval accuracy compared to these baselines for 64K to 256K lengths.

Dear Reviewer pLr2,

We sincerely hope our revisions address your concerns regarding compression overhead and calibration dataset robustness. We also add two more baselines to further demonstrate the effectiveness of MoA. We look forward to your feedback on these updates. If there are any remaining questions, please let us know at your earliest convenience. We are happy to engage in further discussions. Once again, thank you for your valuable time and contributions to refining our work.

For the complexity I mentioned in weakness 1 mainly refers to the cumbersome process involved in this approach, i.e., profiling the model, evaluating configurations, and pinpointing the best plan, but not the computational overhead only. Introducing so many processes just in the sparsification of the attention modules would introduce more uncertainty in training large models, let alone mixture-structured models. And the reported compression time in your response will vary with different experimental setups. To summarize, 5 is the highest score I can give to this work.

[W1] SSM+Uniform attention may replace heterogeneous attention

While the proposed method targets trained full-attention models, the need for heterogeneous window sizes might diminish with the emergence of uniform sliding window methods augmented with Mamba-like state space layers.

We thank the reviewer for raising such a new and interesting perspective and we wish to add a few personal options to it.

We believe some insights that MoA provides may assist the design of future hybrid-attention designs. For example, MoA discovers that the first few layers require denser attention than its subsequent layers, as shown in Figure 11. Based on such observation, it may be sensible to use more transformer blocks in initial layers, and gradually add state-space blocks in subsequent layers. In addition, in the final layers, MoA finds that while most attention heads require low density, some high-density heads is still needed. Such result hints that it may be important to keep few transformer blocks at last layers, rather than replace them all with state-space blocks.

[W2] Writing clarity for MIP section

The methodology description, especially the MIP section, could be clearer.

We apologize for the confusion and revise Appendix D.3.2 with an example and figure illustration for better clarity.

[Q1] Clarify: loss function at profiling

During profiling, what is the loss function? Is it the sum of entropy losses at all output positions or just the largest position? You mention that perplexity is evaluated only on the answer part - does this apply during profiling as well? Please clearly define the loss function.

The loss function used during profiling is the sum of cross-entropy losses computed over the answer part generated by the original model. Using the answer part as supervision enhances model alignment and long-range dependency properties.

We thank the reviewer for pointing out this potential confusion. To better connect the profiling stage discussed in Section 4.1 with the supervision details in Section 5, we have revised the paper to explicitly link the profiling loss definition to Section 5. Additionally, we have clarified the loss definition within Section 5 to improve overall clarity and coherence.

[Q2] Clarify: search space for hyperpareter alpha and beta

Are the choices of ($\alpha$, $\beta$) pairs finite? If yes, what are they?

Yes, the choices of ($\alpha$, $\beta$) pairs are finite. In our experiments, we use 6 values for $\alpha$ and 9 values for $\beta$, creating a search space of 54 pairs for each attention head. $\alpha$ is uniformly sampled from the range $[-2048, 8192]$, and $\beta$ is uniformly sampled from $[0,1]$. The resulting attention span lengths are clipped to the range between 0 and the current input length.

We revise Section 3.2 and add hyperparameter details in Appendix A.1 for further clarity.

Dear Reviewer UNvV,

We sincerely hope our revisions clarify the presentation of the method details and setups. If you have any remaining questions or require further clarification, please feel free to let us know at your earliest convenience. Once again, thank you for your insightful perspectives, valuable feedback, and contributions to improving our work.

[W1] Add popular short-context benchmark: AlpacaEval

As this study targets general sparse attention, it would be better to also conduct evaluation on popular LLM evaluations (e.g., AlpacaEval). Since the prevailing application for LLM is generation, it is necessary to inspect the potential performance degeneration on those tasks.

We thank the reviewer for the suggestion and have included the AlpacaEval 2.0 benchmark in Appendix B.1.5. As shown in Table 10, MoA achieves the highest length-controlled win rate, outperforming other sparse attention baselines and even the original dense model.

We would also like to note that our work mainly targets the long-context inference scenario, where the attention computation and kv cache storage is heavy. Therefore, we mainly adopt long-context benchmarks, instead of benchmarks with shorter sequences (e.g., approximately 50 tokens of input and 450 tokens of output in AlpacaEval).

Nevertheless, we agree that the addition of AlpacaEval can strengthen the evidence for MoA’s generalizability and robust performance across a wide range of input and output lengths.

[W2] Efficiency results at 128K length

Also, I suggest the author to also evaluate the wallclock time and memory consumption for a longer context (e.g., 128k), since the proposed method should have the largest generation acceleration potential at those cases.

We thank the reviewer for the valuable suggestion and have added the 128K efficiency test results in Appendix B.2.2. Thanks to the reduced KV-Cache, MoA efficiently processes 128k input using only one A100 GPU, whereas FlashAttention2 and vLLM baselines require at least two GPUs to handle a single request. As shown in the Table, MoA achieves a $4.7$-$4.9\times$ decode speedup compared to FlashAttention2, while using half the number of GPUs. Additionally, it demonstrates a $1.9$-$2.1\times$ reduction in GPU memory usage. Compared to vLLM, which utilizes tensor parallelism, MoA delivers $1.3$-$1.4\times$ higher throughput per GPU, alongside significant memory savings.

[Q1] The reason for the name: MoA

Why the proposed method is named "Mixture of Attention"? I feel the name is a bit confusing and would suggest the author to change it to a more straightforward name.

We thank the reviewer for raising this discussion. The name "Mixture of Sparse Attention" (MoA) was chosen to highlight the importance of heterogeneity in sparse pattern designs. The core idea of our method is to mix different attention spans and elastic rules across attention heads, aligning with their distinct roles within the model. Inspired by "Mixture of Experts" (MoE), the name also reflects the diversity of attention heads' functions in LLMs and aims to encourage further exploration in this area.

We apologize for any confusion the name may have caused. If this explanation does not fully address your concern, we are open to discussing alternative names that better capture the method’s contributions.

Dear Reviewer 2JPU,

We sincerely hope we address your concerns by adding the AlpacaEval benchmark and efficiency tests at 128k. We look forward to your feedback on these updates. If there are any remaining questions, please let us know at your earliest convenience. We are also happy to engage in further discussions regarding the name of the paper. Once again, thank you for your valuable time and contributions to refining our work.

[W1.1] Difference between prior sparse attention works [1-3]

Sparse attention has been extensively explored in prior research, notably in [1] and [2], which diminishes the novelty of this study.

We acknowledge the extensive prior works on uniform sparse attention. They primarily address the coherence problem in long contexts. Instead, our approach explores the heterogeneous elastic rules for attention patterns. We further address the distinct challenge of extending the effective context length of LLMs. Specifically:

• Prior work [1] analyzes massive activations in LLMs, revealing the sparse nature of attention.

• StreamingLLM [2] proposes uniform sparse pattern and attention sink, improving response coherence, especially for streaming scenarios and out-of-distribution lengths.

• H2O [3] proposes uniform dynamic sparse pattern, reducing KV-Cache during the decoding stage.

• Our work proposes heterogeneous sparse patterns and their elastic rules, uniquely expanding the effective context length of LLMs.

[W1.2] Sliding window approach has been explored by H2O [3]

Additionally, H2O [3] has already thoroughly examined the feedback from using a sliding window approach.

We wish to clarify that the sliding window pattern is not our contribution. Instead, we construct and optimize heterogeneous elastic rules for this pattern across heads and lengths. MoA shows significant improvements in both accuracy and efficiency (see Figures 1, 4, 5 and Tables 4, 5 in the main paper). Our results demonstrate clear advantages over [2] and [3], which serve as our primary baselines.

[W2] Add comparative baselines H2O [3], SnapKV [4], and PyramidKV [5]

There is a lack of comparative baselines, specifically H2O [3], SnapKV [4], and PyramidKV [5].

We thank the reviewer and add all the queried baselines. As shown in the updated Figure 4, MoA consistently achieves a superior throughput-accuracy trade-off across varied densities, outperforming all six baselines. Additionally, we include the long context retrieval performance of SnapKV [4] and PyramidKV [5] in the following Table. MoA shows 4% to 21% higher retrieval accuracy compared to [4,5] for 64K to 256K lengths. H2O [3] encounters OOM errors at these lengths.

We also wish to note that H2O [3] is one of the primary baselines in our work, with comprehensive comparisons provided in Figures 4 and 5, as well as Tables 4 and 5 of the main paper.

[W3] Add Needle-In-A-Haystack (NIAH) evaluation

An evaluation of Needle is necessary to demonstrate the motivation behind preserving long-range dependencies effectively.

We add the Needle-In-A-Haystack (NIAH) retrieval evaluation [7] in Appendix B.1.2, Figure 8. As shown in the Figure, MoA shows perfect retrieval accuracy across 8k to 256k lengths.

Note that our primary benchmark LongEval[8] also validates the retrieval ability of models, which are shown in Figure 1,4,5 and Table 1,3,4 of the main paper. The result of NIAH evaluation aligns with the results observed in the LongEval benchmark.

[W4] Reference concurrent work: RazorAttention [6]

The proposed method is really similar to a previous method RazorAttention [6]. I recommend adding this to citations.

We thank the reviewer and discuss the concurrent work RazorAttention[6] in Appendix E. Both MoA and RazorAttention use different strategies for different attention heads. They mainly distinguish from each other in the following aspects:

1. Strategies for attention heads: RazorAttention categorizes attention heads into two types: retrieval and non-retrieval. It adopts bipolar strategies, applying either full attention or fixed-sized local attention. In contrast, MoA recognizes the diverse attention spans of different heads and employs a broader range of strategies, covering the entire spectrum from very limited local attention to full attention.

2. Adaptation to different input lengths: RazorAttention uses a fixed density for non-retrieval heads across all input lengths. MoA, on the other hand, applies heterogeneous elastic rules for each head, dynamically adjusting densities based on input length while maintaining overall density constraints.

3. Determination of strategies: RazorAttention relies on a heuristic approach to assign strategies. It uses attention scores between the current token and specific tokens (e.g., echo and induction tokens) to identify retrieval heads and assign full attention. MoA employs a loss-based method. It uses a first-order Taylor expansion to estimate the impact of a strategy on end-to-end prediction loss and optimizes strategies to minimize this loss under density constraints.

By discussing the concurrent work RazorAttention [6], we provide a more comprehensive context for the paper.

Dear Reviewer 9X5F,

We sincerely hope our revisions address your major concerns regarding prior works and baselines. We look forward to your feedback on these updates. If there are any remaining questions, please let us know at your earliest convenience. We would be happy to engage in further discussions. Once again, thank you for your valuable time and contributions to refining our work.

Thanks for the authors' efforts. I have several questions and concerns to the authors:

1. What is the main difference from MoA's heterogeneous sparse patterns to other paper that also discusses about the attention sparsity? For example, [1]'s method is head-level adaptive compression, [2]'s method also identify the unique patterns of diverse heads. I think [1] and [2] both apply the uniform sparse attention across heads, and then build methods based on these patterns. I am not pretty sure about the unique novelty of heterogeneous elastic rules for attention patterns from MoA.

2. It may be also helpful to discuss MoA with [3], [3] also identify diverse unique patterns in long-context attention, which is a close setting compared with MoA.

Since the rebuttal period is close to the end, authors could just briefly answer these questions without empirical results. Thanks!

[1] SnapKV: LLM Knows What You are Looking for Before Generation [2] MODEL TELLS YOU WHAT TO DISCARD: ADAPTIVE KV CACHE COMPRESSION FOR LLMS [3] MInference 1.0: Accelerating Pre-filling for Long-Context LLMs via Dynamic Sparse Attention