Loki: Low-Rank Keys for Efficient Sparse Attention

• URL: http://arxiv.org/abs/2406.02542v1
• Date: Tue, 4 Jun 2024 17:58:03 GMT
• Title: Loki: Low-Rank Keys for Efficient Sparse Attention
• Authors: Prajwal Singhania, Siddharth Singh, Shwai He, Soheil Feizi, Abhinav Bhatele,
• Abstract summary: We propose a novel sparse attention method that ranks and selects tokens in the KV-cache based on attention scores computed in low-dimensional space. Our evaluations show that Loki is able to maintain the efficacy of the models better than other popular approximation methods.
• Score: 44.74682508879725
• License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/
• Abstract: Inference on large language models can be expensive in terms of the compute and memory costs involved, especially when long sequence lengths are used. In particular, the self-attention mechanism used in such models contributes significantly to these costs, which has resulted in several recent works that propose sparse attention approximations for inference. In this work, we propose to approximate the self-attention computation by focusing on the dimensionality of key vectors computed in the attention block. Our analysis reveals that the key vectors lie in a significantly lower-dimensional space, consistently across several datasets and models. Exploiting this observation, we propose Loki, a novel sparse attention method that ranks and selects tokens in the KV-cache based on attention scores computed in low-dimensional space. Our evaluations show that Loki is able to maintain the efficacy of the models better than other popular approximation methods, while speeding up the attention computation due to reduced data movement (load/store) and compute costs.

Related papers

• Efficiently Scanning and Resampling Spatio-Temporal Tasks with Irregular Observations [13.491183255489396]
  We propose an algorithm that alternates between cross-attention between a 2D latent state and observation, and a discounted cumulative sum over the sequence dimension. Our algorithm achieves comparable accuracy with a lower parameter count, faster training and inference compared to existing methods.
  arXiv  Detail & Related papers  (2024-10-11T10:11:31Z)
• Eigen Attention: Attention in Low-Rank Space for KV Cache Compression [9.080678336379528]
  We propose Eigen Attention, which performs the attention operation in a low-rank space, thereby reducing the KV cache memory overhead. We show that Eigen Attention results in up to 40% reduction in KV cache sizes and up to 60% reduction in attention operation latency with minimal drop in performance.
  arXiv  Detail & Related papers  (2024-08-10T22:47:12Z)
• Various Lengths, Constant Speed: Efficient Language Modeling with Lightning Attention [19.618556742380086]
  We present Lightning Attention, the first linear attention implementation that maintains a constant training speed for various sequence lengths under fixed memory consumption. To enhance accuracy while preserving efficacy, we introduce TransNormerLLM (TNL), a new architecture that is tailored to our lightning attention.
  arXiv  Detail & Related papers  (2024-05-27T17:38:13Z)
• CORM: Cache Optimization with Recent Message for Large Language Model Inference [57.109354287786154]
  We introduce an innovative method for optimizing the KV cache, which considerably minimizes its memory footprint. CORM, a KV cache eviction policy, dynamically retains essential key-value pairs for inference without the need for model fine-tuning. Our validation shows that CORM reduces the inference memory usage of KV cache by up to 70% with negligible performance degradation across six tasks in LongBench.
  arXiv  Detail & Related papers  (2024-04-24T16:11:54Z)
• Towards Model-Size Agnostic, Compute-Free, Memorization-based Inference of Deep Learning [5.41530201129053]
  This paper proposes a novel memorization-based inference (MBI) that is compute free and only requires lookups. Specifically, our work capitalizes on the inference mechanism of the recurrent attention model (RAM) By leveraging the low-dimensionality of glimpse, our inference procedure stores key value pairs comprising of glimpse location, patch vector, etc. in a table. The computations are obviated during inference by utilizing the table to read out key-value pairs and performing compute-free inference by memorization.
  arXiv  Detail & Related papers  (2023-07-14T21:01:59Z)
• CloudAttention: Efficient Multi-Scale Attention Scheme For 3D Point Cloud Learning [81.85951026033787]
  We set transformers in this work and incorporate them into a hierarchical framework for shape classification and part and scene segmentation. We also compute efficient and dynamic global cross attentions by leveraging sampling and grouping at each iteration. The proposed hierarchical model achieves state-of-the-art shape classification in mean accuracy and yields results on par with the previous segmentation methods.
  arXiv  Detail & Related papers  (2022-07-31T21:39:15Z)
• SimpleTron: Eliminating Softmax from Attention Computation [68.8204255655161]
  We propose that the dot product pairwise matching attention layer is redundant for the model performance. We present a simple and fast alternative without any approximation that, to the best of our knowledge, outperforms existing attention approximations on several tasks from the Long-Range Arena benchmark.
  arXiv  Detail & Related papers  (2021-11-23T17:06:01Z)
• Learning Optical Flow from a Few Matches [67.83633948984954]
  We show that the dense correlation volume representation is redundant and accurate flow estimation can be achieved with only a fraction of elements in it. Experiments show that our method can reduce computational cost and memory use significantly, while maintaining high accuracy.
  arXiv  Detail & Related papers  (2021-04-05T21:44:00Z)
• Kronecker Attention Networks [69.22257624495899]
  We develop Kronecker attention operators (KAOs) that operate on high-order tensor data directly. Results show that our methods reduce the amount of required computational resources by a factor of hundreds.
  arXiv  Detail & Related papers  (2020-07-16T16:26:02Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.