Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Loki: Low-rank Keys for Efficient Sparse Attention

Prajwal Singhania, Siddharth Singh, Shwai He, Soheil Feizi, Abhinav Bhatele

TL;DR

This work tackles the computational bottlenecks of self-attention in large language models by uncovering a low-dimensional structure in attention keys and exploiting it with PCA-based Top-K selection (Loki). Loki performs approximate scoring in a reduced space to identify a small set of top tokens and computes final attention with full dimensionality only for those tokens, achieving substantial speedups (up to 45%) with limited accuracy degradation. The method is training-free, non-deleting, and implemented with optimized Triton kernels, and it generalizes across multiple models, datasets, and tasks. The results suggest a practical path to accelerate autoregressive inference while preserving model quality, with potential gains from further integration with advanced inference stacks.

Abstract

Inference on large language models (LLMs) 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 LLM inference contributes significantly to these costs, which has sparked an interest in approximating the self-attention computation to reduce such costs. In this work, we propose to approximate self-attention by focusing on the dimensionality of key vectors computed in the attention block. Our analysis reveals that 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 speed up the attention computation due to reduced data movement (load/store) and compute costs while maintaining the efficacy of the models better than other popular approximation methods.

Loki: Low-rank Keys for Efficient Sparse Attention

TL;DR

This work tackles the computational bottlenecks of self-attention in large language models by uncovering a low-dimensional structure in attention keys and exploiting it with PCA-based Top-K selection (Loki). Loki performs approximate scoring in a reduced space to identify a small set of top tokens and computes final attention with full dimensionality only for those tokens, achieving substantial speedups (up to 45%) with limited accuracy degradation. The method is training-free, non-deleting, and implemented with optimized Triton kernels, and it generalizes across multiple models, datasets, and tasks. The results suggest a practical path to accelerate autoregressive inference while preserving model quality, with potential gains from further integration with advanced inference stacks.

Abstract

Inference on large language models (LLMs) 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 LLM inference contributes significantly to these costs, which has sparked an interest in approximating the self-attention computation to reduce such costs. In this work, we propose to approximate self-attention by focusing on the dimensionality of key vectors computed in the attention block. Our analysis reveals that 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 speed up the attention computation due to reduced data movement (load/store) and compute costs while maintaining the efficacy of the models better than other popular approximation methods.
Paper Structure (28 sections, 2 theorems, 7 equations, 16 figures, 5 tables, 2 algorithms)

This paper contains 28 sections, 2 theorems, 7 equations, 16 figures, 5 tables, 2 algorithms.

Table of Contents

  1. Introduction
  2. Background and Related Work
  3. Related Work
  4. Dimensionality Analysis of Attention Keys
  5. Models and Datasets Used
  6. Findings and Discussion
  7. Loki: Low-Dimensional Key Representations
  8. Properties of Attention in the PCA-transformed Space
  9. PCA-based Top-K Algorithm
  10. Implementation in Triton
  11. Experimental Setup
  12. Results
  13. Comparison with Full Attention
  14. Comparison with Other Sparse Attention Methods
  15. Generalizability
  16. ...and 13 more sections

Key Result

Lemma 4.1

Let $D$ be the dimension of an attention head and $\mathbf{P} \in \mathbb{R}^{D \times D}$ be the PCA projection matrix of key tensors calibrated offline on a dataset. Assuming we are generating the $S^{\mathit{th}}$ token in the sequence, let $\mathbf{q}_S \in \mathbb{R}^{1 \times D}$ be the query

Figures (16)

  • Figure 1: Rank at which 90% of the variance is explained, averaged across all layers and heads for different models. Full rank is represented by the black dashed line (left). Overview of Loki (right).
  • Figure 2: Rank at which 90% of the variance is explained for pre-rotary and post-rotary keys produced by each layer averaged across all heads ($Rank_{l}@90$) for different models. We observe that all models exhibit significantly low rank (full dimensionality is 128 or 96 represented by the black dashed line) consistently across all datasets.
  • Figure 3: Evaluation of Loki on perplexity (left plot) and short-context tasks (right plot) for different models. Task accuracy is an average across all short-context tasks mentioned in \ref{['sec:setup']}.
  • Figure 4: Evaluation of Loki on LongBench tasks for the Llama2-7B-Chat model.
  • Figure 5: Downstream task performance for Loki and other approaches for different models (higher is better). GSM8K is excluded, as we were unable to run H$_2$O for this task.
  • ...and 11 more figures

Theorems & Definitions (4)

  • Lemma 4.1
  • proof
  • Lemma 4.2
  • proof