Loki: Low-rank Keys for Efficient Sparse Attention

Summary

This method propose the PCA based attention score approximation for top-k attention. Perform PCA on offline dataset, and store the PCA vectors for inference.

Strengths

This work easily makes a QK approximator without gradient-based training but with only simple PCA for top-k attention selection.

Weaknesses

You have to store the PCA vectors and perform the projection of KV during attention. This projection matrix should be holded on threadblock in GPU, to minimize the GPU memory access. But this matrix might be huge if the hidden size of QKV is large. But untill now, most of LLM using less than $256$, therefore, should be fine. (NOTE: I need to checkout the implementation whether this is correct or not)

Questions

1. The .gz file seems broken. Can you upload the file in other format again? I want to look into source codes, for details. 2. Fig. 1. might leads mis understand that this attention mechnism using low rank PCA attention only. Can you add that this method perform top-k operation on top of approximated scores in the figure using diagram? (I think text only is not sufficient) 3. Top-k operation is usually very slow in GPU, due to synchronization and information exchange of top-k values via global-memory. I think per-kernel latency breakdown should be presented on this paper. (Fig. 6 only shows the single sequnce length in left.) 4. Some figures are not sized properly. (e.g. Fig. 6.) Can you adjust the size of figures to avoid stretching the texts? 5. I think there should be some plot that show latency-performance trade-off (ms - accuracy) Can you add this plot using some downstream tasks? 6. Fig.4, they evaluate downstream tasks, but every task is quite short sequence length. Can you try LongBench (https://github.com/THUDM/LongBench)? Question 5 should be solved using LongBench rather than short sequence tasks. I have concern about performance evalution. The only metic used is Perplexity, HellaSwag, TQA, Winogrande, ARC, MMLU, which is relatively short sequence compared to long context LLM such as Qwen2 and Phi3. I suggest to add evaluation on long context.

Rating

8

Confidence

5

Soundness

4

Presentation

4

Contribution

4

Limitations

PCA should be performed during train stage, and PCA projection required for top-k attention selection. This may leads additional effort to optimize the GPU kernels in many scale, and sometimes impossible if the device cannot hold PCA projection matrix in shared memory.

Summary

This paper reveals that key vectors lie in a significantly lower-dimensional space. Inspired by this finding, the author approximates the computation of the original attention score using PCA, then selects the top-k keys based on the approximate attention scores. Experiments across different models and datasets show that PCA-TopK can achieve speedups of up to 40% with minor reductions in generation quality.

Strengths

1. The author discovers that the key vectors in multi-head attention lie in a lower-dimensional space, which may inspire future work on Sparse Attention. 2. TThe experiments across various models and datasets indicate that PCA-TopK can achieve speed improvements with only minor reductions in generation quality.

Weaknesses

1. There are some typos (lines 44, 150) in this article, and some figures are unclear with text overlaps (Figures 2, 6). 2. There are few baselines about Sparse Attention in the experiments. How does the PCA-TopK method compare to SPAR-Q Attention? What is the trade-off curve between their acceleration ratio and effect?

Questions

1. The analysis that key vectors lie in a lower-dimensional space is interesting. Do value vectors and query vectors share similar characteristics? What about the vector after merging multi-head attention? 2. Given the differences in Rankl@90 across layers, what is the impact of varying the policy depending on the layer? 3. If post-training with PCA-TopK, will it yield better performance?

Rating

6

Confidence

3

Soundness

3

Presentation

2

Contribution

3

Limitations

N/A

Summary

The paper introduces a method for approximating attention in LLMs, with the benefit of improved inference efficiency. The insight is to focus on the dimensionality of the key vectors computed in the attention block. Principal component analysis reveals that the keys lie in a low dimensional space. This gives rise to the proposed method PCA-TopK, which uses PCA to compute approximate attention scores in a reduced dimension, and then selects the top-k tokens based on these scores, and compute the equivalent of full attention only for the selected tokens. The method is evaluated on multiple LLMs from the Llama family, Pythia, Mistral, Mixtral, Phi, and various datasets. The results show comparable performance to full attention while having significant speedups.

Strengths

The study of low intrinsic dimensionality of attention keys across multiple models and datasets is insightful. Theoretical support is provided through lemmas and proofs. The evaluation includes several models, tasks and datasets, showing it could be generalized. The authors developed optimized Triton kernels for practical implementation.

Weaknesses

The method does not seem to reduce the memory usage.

Questions

For some models the pre-rotary PCA transforms outperform the post-rotary ones, which is intriguing. The authors acknowledge that they do not have a clear explanation, however if more understanding was developed it would be great to include it in the paper.

Rating

6

Confidence

3

Soundness

3

Presentation

3

Contribution

2

Limitations

yes

Summary

This paper proposes a sparse attention mechanism for large language models by leveraging the low-dimensionality of key vectors in the attention block. The approach ranks and selects tokens in the KV-cache based on attention scores computed in the reduced dimensional space, leading to significant speedups in attention computation without major sacrifices in model quality.

Strengths

Empirical Validation: Extensive evaluations demonstrate significant speedups with minimal accuracy degradation. Theoretical Soundness: The use of low-rank keys is backed by robust theoretical analysis, enhancing the credibility of the approach.

Weaknesses

Implementation Complexity: The method's complexity might pose challenges for achieving the theoretical speedups without specialized knowledge. Memory Footprint Consideration: While the approach excels in computation, it does not address memory footprint reduction, a limitation when compared to other sparse attention techniques.

Questions

1. Model Class and Size Performance Variance: How does Loki's performance vary with different model classes and sizes not included in the paper? 2. Integration with High-Efficiency Attention Mechanisms: How does the PCA-TopK method align with or complement high-efficiency attention mechanisms such as Flash Attention or Ring Attention? Can it be orthogonally integrated with these techniques, or are there specific challenges that need to be addressed? 3. Comparison with Other Sparsity Strategies: The paper does not compare with attention sparse strategies beyond H2O, such as SnapKV[1]. Could the authors discuss the positioning of PCA-TopK relative to these methods and possibly include comparative analysis in future work? 4. Figure Clarity: Regarding Figure 6 (left), it appears there might be a compression issue. Could the authors verify the image quality and ensure that it is legible in subsequent versions of the paper? 5. Information Loss Management: How does PCA-TopK handle the potential loss of information when reducing the dimensionality of key vectors, especially in the context of Rotary Positional Embeddings (RoPE) which increase dimensionality? 6. Combination with Compression Techniques: How does PCA-TopK interact with other model compression techniques, and is there a combined approach that could yield better results? 7. Handling of Longer Sequences: How does Loki handle sequence lengths longer than those in the evaluation datasets? 8. Quantization Integration: Can the PCA-TopK mechanism be integrated with other optimization techniques like quantization?

Rating

6

Confidence

4

Soundness

2

Presentation

2

Contribution

2

Limitations

N/A