Loki: Low-rank Keys for Efficient Sparse Attention
Authors: Prajwal Singhania
Siddharth Singh
Shwai He
Soheil Feizi
Abhinav Bhatele
Venue: NeurIPS 2024 poster (2024)
Keywords: Approximate Attention
Low-Dimensional Attention Keys
PCA
Top-K Attention
Primary area: infrastructure

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
Prajwal Singhania, Siddharth Singh, Shwai He, Soheil Feizi, Abhinav Bhatele
Department of Computer Science, University of Maryland College Park, MD 20742 prajwal hamdd edu, bhatele&cs. umd edu
Abstract
Inference 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 reduce such costs. In this work, we propose approximate self-attention by focusing on the dimensionality of key vectors computed in the attention block. Our analysis reveals that key vectors lie 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- www-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.
Introduction
As large language models (LLMs) grow in size, deploying them for efficient inference presents substantial challenges, largely due to computation and memory access bottlenecks in the self-attention block (32), especially when handling long sequences. These challenges stem from the autoregressive nature of attention, which generates the output one token at a time. At each step, the entire preceding state, stored in the key-value (KV) cache, must be fetched from memory, which can sometimes exceed the size of the model parameters itself (18). This frequent K KV-cache access from GPU DRAM to registers becomes costly, as it scales quadratically with the output sequence length. In addition, matrix multiplications in the attention layers also have a quadratic scaling cost with sequence length, compounding the overall computational burden.
Several strategies (39 26 20 have been proposed to address this challenge by reducing the com- putational complexity and/or memory demands associated with the self-attention mechanism. One promising category of approaches focuses on approximating attention, employing techniques such as quantization using a subset of the tokens in the KV-cache [I] (sparse attention).
In contrast to other sparse attention approaches that either permanently prune tokens from the key- value cache [3 or impose fixed sparsity pattern (35), our proposed method dynamically selects key tokens at each generation step based on approximate attention scores and avoids deletions. This approach is inspired by critical observation: across range of LLMs and datasets, key tensors consistently occupy significantly lower-dimensional space than the full attention head dimension.
For instance, in Figure [ (left), we show that across various LLMs [ [6] 90% of the variance explained by PCA is captured at an effective key vector rank of around 80, despite the key tensor dimension being much larger (128).
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

Dimensionality of Attention Keys
140
Rank 90% Explaincd Variacce
128
PCA
Sire
100
Caloration
80
60
40
INTERINNC
Query
Approc Arrnntttttonn
Scormi
20
Select
Ulam2277
Liamaz11
Llaa L78
Llama
1-ma3.88
Lima3. 708
Misiral7B
Mixoral-hh]
Mixtral-Bi B
Bytha 6,5
Tips TOURS
AAsaatton on
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).
Based on this observation, we introduce Loki, a sparse attention method that leverages the low- dimensional structure of key vectors reduce data movement and computation costs without significantly impacting model quality. First, we apply PCA to keys generated from a calibration dataset, storing all principal components but using only the top d (25-50%) compute approximate attention scores during inference. This dimensionality reduction, informed by our previous observa- tion that key vectors have low effective rank, allows to efficiently identify the top-k (12.5-25%) most relevant tokens using the approximate scores. For these selected keys, we then revert to the full dimensionality compute the final attention scores, ensuring both efficiency and accuracy. Figure (right) illustrates our approach.
Our theoretical complexity analysis demonstrates that Loki can provide significant speedups in the attention step. However, actually realizing these gains requires an efficient implementation of our method to minimize data movement in the additional operations introduced top of the original self attention algorithm. Thus, we implement optimized sparse matrix multiplication kernels for Loki in Triton, leading speedup of up to 45% over the standard HuggingFace Transformer's 's (34 attention implementation (vanilla attention) for Llama2-13B. For this setting, the average degradation in model accuracy (measured across 6 different benchmarks and 8 different models) is only 6.8%.
Our contributions can be summarized as follows:
Detailed analysis showing the intrinsic low-dimensionality of keys in self-attention, its variation across layers for different models, and consistency across different datasets.
Loki: sparse attention method that exploits the aforementioned low dimensionality of keys to make the attention computation faster without sacrificing model quality.
Optimized kernels for efficient implementation of Loki in Py PyTccch.
Evaluation of Loke] on multiple LLMs and downstream tasks, showing that it can achieve significant speedups with minimal degradation in model quality.
2 Background and Related Work
The attention mechanism 32 is at the core of the transformer architecture. Consider single attention query head with head dimension D, processing an input token sequence of length S. During auto-regressive generation, the output of the attention head is calculated as:
qk
y softmax
V
where E RIxD is the query, and K E SXD and V E RSXD are the key and value caches respectively. Additionally, newer transformer models add Rotary Position Embeddings (RoPE) [2] to the keys and query, before computing the attention scores. Since every query attends to all past keys, the mechanism has a quadratic complexity (SS2) ) in number of input + generated tokens.

---

2.1 Related Work
Numerous studies have explored the low-rank structures in transformers for various purposes. Lin- former [33] demonstrated that the attention score matrix is low-rank and proposed alternative low-rank attention formulations during training for linear computational complexity. LoRA [33 showed that parameter updates to a transformer model during fine-tuning reside in low-dimensional subspace. To the best of our knowledge, our work is the first to study the intrinsic low dimensionality of the attention keys themselves and demonstrate the generalizability of this low-dimensional structure across different models (for natural language data).
Sparse-transformers S was one of the first works to introduce a sparse-attention method employing strided sparsity patterns in the attention mechanism. Reformer [17] used locally-sensitive hashing to compute attention scores in sparse manner. Performer [6] used positive orthogonal random features to approximate the attention mechanism. Unlike these methods, which require training or fine-tuning, our approach operates entirely post-training without any fine-tuning.
Another category of sparse attention methods employ token eviction policies permanently delete tokens from the KV-ccche based on some heuristic. StreamingLLM [35] uses initial tokens and a rolling K V-cache for processing infinite-lnngth sequences. Zhang et al. 39 retain only "Heavy Hitters" tokens in the KV-cache based on accumulated attention scores. Scissorhands (20 prioritizes important tokens based on the "Persistence of Importance Hypothesis". .. Ge al. [0 propose an adaptive eviction policy for each transformer layer. These methods are effective in reducing the memory and compute footprint of the attention but suffer from permanent loss of information leading to non-trivial degradation in model quality, Our method does not involve any permanent loss of information with the trade-off of not reducing the memory footprint. Quantization- are complementary our work and can be applied in tandem.
approximate
approaches [4 23
SparQ Attention 26 is a recent work that inspires our approach. They use high-magntudde query dimensions and corresponding key dimensions for approximate attention scoring, followed by computing the full attention scores for the top-k keys. However, their method requires costly non- contiguous column indexing of the key vectors. Further, they store two copies of the past keys for efficiency, increasing memory use by 50%. In contrast, Loki avoids the extra memory and leverages the natural ordering of principal components, allowing for more efficient slicing operation.
concurrent work, InfiniGen 19), accelerates attention by pre-fetching top-k keys from CPU to GPU memory, using SVD-based low-rank approximation of the attention scores. While their low- rank approximation is similar to Loki, our work provides deeper analysis of the intrinsic low-rank structure of attention keys and focuses on speeding up attention computation without CPU offloading. Importantly, their results affirm the benefits of the -ineennnsio nature of attention keys applied in other contexts.
3 Dimensionality Analysis of Attention Keys
As noted in Section D Loki, our proposed method for sparse self-attention, is based on the observation that key tensors consistently reside in a lower-dimensional space than the full attention head dimension suggests. Here, we present empirical evidence supporting this claim by performing PCA on the keys generated in several language models and datasets.
3.1 Models and Datasets Used
To investigate the dimensionality of attention keys, we run 11 transformer-based models: Llama-2
7B/13B/70B BT), Llama-3 8B/70B [8) TinyLlama-1. 1. (38. Pythia-6.9B [4], Mistral-7B [15]. Mixtral-8x7B/8x222 [6]. and Phi3-Mini-4K WikiText-2 (21) (Validation Split), C4 25] (Custom Split), and BookCorpus (40] (Custom Split).
22 on three popular English language datasets:
Custom splits are used for datasets where the validation split not available. We run perplexity evaluation on these datasets and save the generated attention keys, before and after the application of rotary embeddings [29, referred to as pre-rotary and post-rotary keys, respectively throughout the paper. We then perform PCA all the keys generated for each layer and head individually.
3

---

The metric we use in our analysis is the rank at which 2% of the variance is explained by the principal components. We calculate this metric for each layer and head of the models as follows:
d
R Au min d E Z+
Σ Nr 2 v/100
where, Mrn is the 5th normalized eigenvalue of the covariance matrix of the keys for layer, and head, h. We average this metric ranks across all heads of layer, I and refer to it as Ranky @v.
3.2 Findings and Discussion
Figure (left) shows the average Rank; @90 across all layers for models with full key dimensionality of 128. We can see that the average rank is significantly lower than the full dimensionality of the keys for all models. Diving deeper, we present layer-wise wise analysis for few models: Llama2-7B. Llama3-70B, Mixtral-8x4 -xx, and Phi3-Mini-4K in Figure 2 The results for the other models are similar and can be found in Appendix A.T
BookCCopus
Lirrod Arrarrdo Keps Roney
Limea---00 Actintion Kept (Pore Rostry)
Soprss-BcTSAAmrree Amantinn Koys (TeetRoony
PND (1a 1-kk--kkkk Amsedon Keps Rlaseoorry
Ubred 7B Actention Kegs (Pre Rosury)
Liorah 700 Ansention Koys (Pro-kotor)
Maryel Au78 Arteiion Kays Rataay)
PhD Mas Accint on Keys Rocary
1.0
140.
188
109
25
16
25
13
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ı@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.
We observe that the dimensionality of the keys (both pre-rotary and post-rotary) is significantly lower than the full dimensionality of the keys across all calibration datasets. Furthermore, the Ranky @90 for particular layer is consistent across datasets, for all combinations of models and datasets. This indicates that the lower-dimensional structure of the keys is consistent when calculated using different calibration datasets. Another trend we observe is that the initial layers of most models have very low rank, as compared to the later layers, and this trend is particularly prominent for the pre-rotary keys. Lastly, we also observe that for most models, the average of Rank; (90 across all layers is lower for pre-rotary keys as compared to post-rotary keys, indicating that the rotary embeddings increase the dimensionality of the keys. Further analysis on the variation of the rank across different heads within a layer and across different layers within a model can be found in Appendix A.
These results indicate the existence of the following properties: (1) The keys produced by the attention layers of transformer models lie in significantly lower-dimensional space. (2) The dimer-dimen structure of the keys is consistent across different calibration datasets. (3) Rotary embeddings increase the dimensionality of the keys for most models. We now use the first two properties to propose Loki, an efficient sparse-attention method.
4 Loki: Low-Dimensional Key Representations
We now describe our proposed algorithm for sparse attention - Loki. Loki leverages low dimensional projections of the keys in the KV-cache to efficiently and accurately select the top-k (12.5-25%) most relevant tokens for self attention. Before discussing our approach, let us first look at some theoretical properties of attention in the PCa-trranformed space of the key tensors.

---

4.1 Properties of Attention in the PCA-transformed Space
We begin by proving two lemmas that provide the rationale for our approach to compute attention in the ICA-transformed space.
Lemma 4.1. Let D be the dimension of an attention head and P E R DxD be the PCA projection matrix of key tensors calibrated offline sequence. let qs E RIXD be the query vector for the 5th token, K. E R be the key vectors, including the past (S 15-1] 1) keys and the current key. Then, the attention scores computed using the (CA-transformed query and keys are equivalent to the atlention scores computed using the original query and keys.
dataset. Assuming
generating the 5th token in the
Proof. Let Âs- = qsP and Ks = S P be the PCA transformed query and key vectors. Focusing on the dot product term in the attention computation (Equation Dp. we have:
as qs (Rs pT)T [inverting the PCA transform]
= =qs (Ps)("") = (qsP) = ask's
It is important note here that Lemma 4.1 holds for any orthogonal P.
Lemma 4.2. Let Kis.d Sxd (d < D) be the reduced dimension key vectors obtained by projecting the key vectors onto the first d principal components of Then, the atlention scores computed using Kid Proof. Let a E RdxD be an orthogonal transformation that transforms the keys into the reduced dimension space as L,s.d R. Our objective is to minimize the following expression:
good approximation of the the actual attention scores.
min lasK's qs RR Ra
Using |lssK's qs (Ls... RT) |
inequality, we have:
K
(L.s.aRR)))
We change our objective to minimize the upper bound in the RHS instead of the original objective. We know that PCA minimizes the reconstruction error (2nd term in the RHS) among all the orthogonal transformations. Thus, follows that the optimal value R = P.d, and Lisdd K.s.dd
Since we minimize an upper bound when proving Lemma 4.2 is possible that some other transfor- mation might give better approximation to the dot product. Thus, in our experiments, we use PCA transforms computed on both the pre-rotary and post-rotary keys as candidate transformations.
Based on these lemmas and the inherent low-dimensional nature of key tensors in attention, we now introduce the workings of the Loki algorithm.
4.2 pase Top-K Algorithm
Loki implements a PCA-based Top-K Attention approach. Previous works have shown that attention scores for query are highly concentrated motivated several methods that compute attention using only the top-k most relevant keys. However, these previous works either compute the exact attention scores and then select the top-k keys [12 compute non-exact scores but have significantly higher memory requirements [26]. Loki alleviates these issues by computing approximate attention scores (for ranking the keys) in the reduced lower- dimensional space, without any significant increase memory requirements. Algorithm [ shows our Loki method. Line 5 of the algorithm computes the approximate attention scores using d principal dimensions of the query and key vectors. Lines 6-7 select the top-k keys based on the approximate attention scores. Line computes the exact attention scores using the selected top-k keys, directly in the transformed space (Lemma 4.1
small subset of keys [36] 30). This observation has
For vanilla attention, the complexity of computing 9sK7 is O(DS) and the complexity of multiplying the values with the attention scores is O(DS). For Loki, the complexity of calculating the approximate attention scores (Line 5) is (dS)) The complexity of selecting the top-k keys (Lines 6-7) is approximately O (Slog(S) (Sog(S)= + + (sorting followed by selection). The complexity of calculating the exact attention scores and multiplying with the values (Line 8-9)
Compute and Memory Analysis:
5

---

Algorithm Loki: PCA-based Top-K Attention
Require: At the 5th step Input: X5 E RIxD, KV-cache: K,5- sss 1 E ((S-1)xD, Projection
Matrix: P E IR DxD, Configuration parameters (reduced dimensionality, top-k): d, k
1: function LOKI-
K.s-1, V.s-i,P.d,K
2: 4s,ks, Vg + computeQKV(xs))
3: gs t qsp, ks + ksP
4: R.s t (m ks), + coneell
5:
Aappror t (s s
6:
indices topk (aappror: k)
7:
K's c [
Vislindices
8: apzant + softmax
VD
9: return acract Vis
10: end function
is O(2Dk). Additionally, the complexity of projections into the PCA space (Line 3) is O(2D²). Assuming the complexity of selecting the top-k keys small compared to the other operations, the overall complexity of the algorithm is OdSS + 2D + 2D²). Then, we have:
2DS
speedup ds + 2Dk + 2D2 4/2D + K/S + D/S d,/2 ky (given << (5)
where, dj = d/ and ky- = k/S. The memory requirement of the KV-cache is the same as the original attention, with small overhead of storing the PCA transformation matrix.
4.3 Implementation in Triton
Performing Loki efficiently involves complex indexing operations within the KV-cache (lines 5 and 7 of Algorithm [ Standard PyTorch operations create temporary, dense copies of the KV-cache data in memory, leading to slowdowns due expensive memory access. To alleviate this issue, develop optimized kernels in Triton [ for the three matrix multiplication operations in Loki. Our kernels directly access relevant subsets of the K KV-cache (both feature and sequence dimensions) and perform computations within GPU registers. This eliminates the need for creating dense copies, significantly improving performance. Our approach builds on SparQ [26). which introduced similar kernels for top-k attention calculations. However, we identified and addressed inefficiencies in the SparQ kernels, which resulted in speedups of nearly 3x in certain scenarios. (see Appendix
5 Experimental Setup
We evaluate Loki on the basis of perplexity using the WikiText-2 (1] dataset (test split), and on the basis of downstream task performance for short contexts using the LM-harness benchmark [9] and long contexts using LongBench [2]. For the short-context evaluation, we choose the same tasks and associated metries the HuggingFace OpenLLM leaderboard 3 For the LongBench tasks, evaluate on all the English language tasks.
We compare our method against three methods full attention without any approximations, the exact TopK approach which computes the exact attention scores and then uses the top-k tokens to compute the final output, and H₂O [39], popular token-eviction method. For these comparisons, we show the results with a budget size of k 0.25 and 0.125, For our method, we additionally use dy 0.25 and 0. This configuration of our represents 2.6x theoretical speedup. Table provides overview of the methods compared and the associated budget terms. H2O's 's budget was split equally between the heavy hitter and recent tokens, as per the author's recommendations. For H2O, we were unable to run the GSM8K task as the the author's 's ML benchmarking code was too memory intensive to run for that task. For the aforementioned experiments, dataset. For the LongBench tasks, we compare our method with the full attention baseline as we were unable to run H₂O due memory constraints.
generate PCA transforms using the WikiText- 103
For the generalizability study, we compare the results of our method with PCA transforms from different calibration datasets: WikiText-103 (11) C4 (25), and BookCorpus 40). Additionally,

---

Table 1: Explanation of key-budget and dimensionality (Dim.) for different approaches, along with the expected speedup and memory savings.
Method
Budget
Dim.
Description
Speedup
Memory Savings
Exact Top-K
ky
Full
ky fraction of keys selected using exact attention scores
No
No
H₂O
ky
Full
by fraction of keys values selected using H₂O policy
F
Loki
ky
dy
by fraction of keys &values selected using atlention scores compured with dy fraction of full dimensionality
No
(2y BRy
also benchmark our triton based implementation of Loki by running an attention microbenchmark on a Baama2-13Bllike setup (same hidden size and number of heads) for various prompt and generation lengths, and demonstrate speedups over vanilla attention.
All experiments are run on NVIDIA A100 GPUs with 40 and 80 GB of memory on the Perlmutter [24 supercomputer. For larger models, we use AxoNN 27 28 to shard the model across multiple GPUs.
6
Results
We now present the comparisons of Loki with full attention and other sparse attention methods, including a comparison of the computation times.
6.1 Comparison with Full Attention
Let us begin our discussion with Figure showing the perplexity (left) and short-context downstream task evaluation (right) results for Loki on different models. We focus on the Llama2-7B model, comparing pre-rotary (light green/purple) and post-rotary (dark green/purple) PCA transforms for different k and y values. For Llama2-7B.
that the performance of both candidate transforms similar. This trend is consistent across all the models except for Llama3-8B/70) and Mistral-7B, where the post-rotary PCA transform performs significantly worse than the pre-rotary one. For Llama3-8B, perplexity jumps from about 5 for the full attention to over 10, significant decline not with the pre-rotary transform. Mistral-7B shows similar pattern. This is surprising observation since attention scores are calculated from post-rotary keys in the original attention mechanism. A possible explanation is that post-rotary PCA captures token distributions tied to specific positions in the calibration dataset, while pre-rotary PCA may generalize better by using less positional information. Nevertheless, at least one of the PCA transformations performs well for every model. For subsequent results, we only show the better-performing transformation for each model.
Loki vs. Full Attention: Perplexity Evaluation
Loki vs. Full Attention: Short-context Tasks Evaluation
100
Full Attention
Loki (pre, Ki = d, 025)
Lold (post, ky= 025, d, 025)
Avuraga Task Accuracy
10-
Loki (pre, ky 0.125 d 0.)
Parplosity
Loki (post, ky= o22.dd 0.5)
23395
20
2
Liama2-78
Liam 2318
Limma 888
708
Masmal78 78
Maaral-8d88
Llemaa-78
Lama2-138 138
Limad 1.8
Liama 7788
Mistral-B 78
Mixtral-877
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 5
Figure 4 shows the performance of Loki on the LongBench tasks for the Llama2-7B 7B-Chat model. We see that for all tasks, either one of the two candidate transforms performs similarly to full attention. For Summarization, Few Shot Learning. Synthetic, and Code Completion task categories, the best performing Loki configuration at par or better than the full attention model. For the Single-Doo QA and Multi-Dos Doc QA task categories, Loki performs slightly worse than the full attention model,
7

---

with the biggest drop in performance observed for HotpotQA of around 3%. Comparing different (ky. d,) settings, we see that using K, = 0.25 and d 0.25 (green), is better than using k 0.125 and d = 0.5 (purple) for all models and tasks (short-context and long-context). These two settings balance speed and performance well, with the first being superior for accuracy.
Lold vs. Full Atteentionnoorgeeeeea. English Tasks Evaluation
100
FdlAtteriin
80-
5000
Lok: fpon. Ro- 025. 0251
jpons iy di o.g)
Task Apcurnty
60-
grn 6:155,
20-
Namneekk
Maradok
HagorCa
Modpe
God Aport
Cop
Mare-Dr QA
Samma
Lerring
Completen
Figure 4: Evaluation of Loki on LongBench tasks for the Llama2-70 10a--B-Char B--hat model.
6.2 Comparison with Other Sparse Attention Methods
Next, we compare the performance of Loki with other methods, using ky 0.25 for all methods and d, = 0.25 for ours. Table the perplexity results for Llama2-7B/13BB Llama3-8B, and Mistral-7B. Loki's perplexity drop is within 0.1 of full attention across all models, threshold considered acceptable for attention mechanism approximations 37 In contrast, H₂O's perplexity drop nears 0.2 for all models. similar to full attention for all models, except Llama3-88 8B, where the performance is notably worse, but still better than H2O. Importantly, on the challenging MML.U task, Loki degrades less than H2O.
this trend on short-context evaluation. Loki performs
Table 2: Perplexity evaluation of Loki and other approaches for different models (lower is better).
Method
ky
dj
Speedup
Llama2-7B
Llama2-13B
Llama3-8B
Mistral-7B
Full Attention
No 5.1101
4.5680
5.5696
4.9140
Exact-TopK
0.25
No 5.1809
4.5926
5.5716
4.9171
H₂O
0.25
Yes 5.2810
4.7009
5.7056
5.0805
Loki
0.25
0.25
Yes 5.2017
4.6102
5.6648
4.9233
Full Attention
Exact-Too Saact-TopKaa (Ky = 025)
Lold (ky = 025, dy= 025)
H;O (ky = 025)
Llama2-76
Lhma2-
100
100
Task Accuracy
75
Hellowsg
TQA
Whngrande
ARC
Art two GSNBK)
MMLU
Helawat
TQA
Whngrande
ARC
Auz (wio GSMeK)
NALU
Halawwa
AQA
Winggrndo
ARC
Avt felo GSMMK)
MMLU
Helawar
TQA
Wengrande
ARC
Art (wo CSMPO
NeLL
Figure 5: Downstream task performance for Loki and other approaches for different models (higher better). GSM8K is excluded, as we were unable to run H₂O for this task.
is important note here that Loki is designed to be compatible with other sparse attention methods. For instance, token-eviction methods like H₂O delete tokens to save KV-cache memory, whereas Loki reduces memory bandwidth by selecting the top-k tokens without deletion, making them orthogonal. A combined approach could involve using H₂O to delete tokens, then applying Loki to select top-k tokens from the remaining cache. Similarly, Loki is theoretically orthogonal to quantization methods.
8

---

Comparing Loki with Exact- TopK, we find similar performance for Llama2-7B/13B and Mistral-7B. Exact-TopK represents the upper performance bound for Loki if it could perfectly select the top-k tokens. To understand why Loki works well, we examined the top-k agreement between Loki's reduced dimensional attention scores and exact attention scores. [igure[Jshows the Jaceard similarity between the top-k tokens selected by both methods across all layers and heads for Llama2-7B, For the settings: ( = 0.25, dj = 0.25) and (kj = 0.125, dj = 0.5), evaluated in Figure 3 the Jaccard similarity is around 0.9, validating that the Loki is able to select the top-k tokens with high accuracy.
Agraemert s/w ha and Raduced Din Astendion Scores
Laki Generaltation Calbbration Datazets
Uarra-1130 Attention
KN acte Uydons
10-
BeokCorpa
898
A0o Tull Accuray
30-
Tim par lape M
0+
Lann
00
Lodi
Figure 6: Top-k agreement between Loki and Exact- TooK methods for Llama2-7B (left plot). Performance of Loki using transformations derived from different calibration datasets (middle plots). Benchmarking vanilla attention and Loki for Llama2- using huggingface transformers with cache append times (right plot, prompt length = 3072, generation length = 512).
6.3 Generalizability
We now turn attention to the generalizability of the PCA transformations used in our method. Figure 6 (middle) shows the performance of Loki using PCA transformations derived from different calibration datasets (ky- = 0.25, df 0.25). We see that the performance of Loki is consistent across different calibration datasets, indicating that the PCA transformations used in our method are generalizable. This is an important observation as shows that the PCA keys can be generated using a variety of calibration datasets and still achieve good performance.
6.4 Computational Efficiency
We now turn our attention to the computational efficiency of Loki. Analyzing Llama2-13B with Hugging Face Transformers exposed an interesting bottleneck (Figure rightmost). Regardless of the attention type (vanilla or Loki), more than 80% of the time consumed within the Hugging Face framework for appending key-value pairs of the latest token the KV-cache. This shared bottleneck minimizes the overall performance improvement of our optimizations. We hypothesize that using a more advanced inference system like vLLM [8] could significantly reduce this append time, but leave that exploration for future work. To isolate the impact of our optimizations, the plots in Figure] focus solely on the attention computation time, excluding the KV-cache append time.
In the left plot of Figure 7 we see that Loki speeds up the total attention compute time (excluding K KV-cache appends) compared to vanilla attention across various prompt and generation lengths. For a prompt length of 3072 and generation length of 512, Loki achieves nearly a 45% speedup, despite the fact that it incurs an extra matrix multiplication operation. The breakdowns also show that the top-k operation is nearly as expensive as the smaller matrix multiplications. which is a significant bottleneck. Replacing PoTorhh's 's top-k with custom kernel could improve this. For the shorter prompt length of 2048 we observe a speedup of around 40% (generation length = 512), slightly lower than the speedup at 3072. This trend is expected as larger prompts result in a bigger KV-cache, amplifying the impact of our optimizations.
Figure 7 (Right) shows the accuracy vs. attention time trade-off across various ky, dy settings of Loki, with accuracy measured on LongBench and attention times from our microbenchmark. The previously evaluated settings, ky 0.25, dy = 0.25 and ky = 0.125, d, = 0.5, provide a good balance between performance and accuracy, with ky = 0.25, dy = 0.25 favoring accuracy slightly and ky 0.125, d = 0.5 favoring performance.

---

Llama2-1 138: Attention Breakdown w/o KV<cache Updates
Liamll778 LongSench Accuracy rcterriion Times
100
Pronp:Lonyyy 2016
3072
q [boy-d)
St
Avarage Accuracy
80
Accuracy
Time 12
0.5
(po--kc)
0.5
Time par Layer (x)
0:+
S. (oo- 40)
60
1000
C28
03
0.3
0.2
hat harkooni
de es do
Be o2t di 031)
Si Brwezt eg.
Pro ar 009
(e drs (10
Read dadddaal
0.0
V DA.4
LA.4
14
0.0
LAU4
LA40
128
256
312
128
256
512
Number Generazed Tokers
Configrazion
Time per layer for vanilla attention (V) and Loki (L-A: ky 0.25, d 0.25; L-B: K, = 0.125, d = 0.25) for Llama2-13B using huggingface transformers (left two plots). Long- Bench average accuracy for different Loki configurations, alongside micro-benchmen
Figure 7:
attention times (right plot, all layers, prompt length = 3500 & generation length = 512). We choose the prompt length to match LongBench's 's configuration for this model and generation length to match the maximum in any LongBench task. For both figures, we use a batch size of 16 and report the average time over 10 trials (std. dev. in measured times was less than 0.05 percent of the mean).
7 Conclusion
In conclusion, we introduced Loki, an algorithm for efficient sparse attention that does not compromise the model quality while reducing the computational complexity of self attention. We made a crucial observation that key vectors in attention lie low-dimensional space, across different models and datasets. Leveraging this insight, Loki uses attention scores computed in lower-dimensional space to rank and select the top-A most relevant tokens from the KV-cache. dimensionality only for the selected tokens to compute the final attention. Our theoretical analysis shows that Loki can provide significant speedups in the attention step. To implement this efficiently, we develop optimized kernels for the various sparse matrix multiplications in our approach. Our empirical evaluation shows that Loki performs better than popular approximation methods on a variety of models and tasks, with respect to preserving model quality. Finally, show that Loki can provide speedups of up to 45% over the vanilla attention empirically, making it a promising approach to address the computational challenges in transformer inference.
then uses the full
Limitations and Future Work: Loki does not focus on reducing memory usage of the KV-cache currently. As mentioned previously in 6.2 it can potentially be combined with other sparse attention method for improved memory-perrie perooryyppoooom---acceee the KV-cache in CPU memory and transferring only the top-k keys and values to the GPU 19
trade-offs. Another direction involves storing
While Loki outperforms vanilla attention in our benchmarks. practical deployment would require inte- gration with efficient attention kernels like FlashAttention [7]. As seen in our compute benchmarking. the top-k selection operation could introduce a bottleneck towards achieving this. Investigating this bottleneck and integrating Loki with optimized attention kernels is left for future work.
Our finding of the keys' low intrinsic dimensionality suggests promising research directions. The variation of this dimensionality across heads and layers could further be explored. We briefly experimented with variable d/ policy per layer (see Appendix B.2), but did not observe significant significant improvements. A more sophisticated policy could be explored in future work.
Acknowledgments and Disclosure of Funding
This research used resources of the National Energy Research Scientific Computing Center (NERSC), a Department of Energy Office of Science User Facility using NERSC award ADR-ERCAP002989. Soheil Feizi was supported in part by the following grants: NSF CAREER AWARD 1942230, ONR YIP award 1-00114-2-1-22711 ARO's Early Career Program Award 310902-00001, Army Grant No. 1991NNF120076 NSF award CCF2212458, NSF award 2229885, an Amazon Research Award and award from Capital One.
10

---

References
[1] Introducing triton: Open-source gpu programming for neural networks. com htt/////xx//tt/o///////// 2021.
https://openal
[2] Yushi Bai, Xin Lv, Jiajie Zhang, Hongchang Lyu, Jiankai Tang, Zhidian Huang, Zhengxiao Du, Xiao Liu, Aohan Zeng, Lei Hou, Yuxiao Dong, Jie Tang, and Juanzi Li. Longbench: A bilingual, multitask benchmark for long context understanding. arXiv preprint arXiv:2308. /4508, 2023.
[3] Edward Beeching, Clémentine Fourrier, Nathan Habib, Sheon Han, Nathan Lambert, Nazneen Rajani, Omar Sanseviero, Lewis Tunstall, and Thomas Wolf. Open Ilm leader- board. https:/ / http:///hgggnnaceeeee////o/ ctppsaccccccepcccee/opennnn -emeeen--lmmmme Lin--eederbbard wwwpo/aarrolld//eeeeeeeeeeee old/open_nnie 11m leaderboard 2023.
[4] Stella Biderman, Hailey Schoelkopf, Quentin Gregory Anthony, Herbie Bradley, Kyle O' O'Brien, Eric Hallahan, Mohammad Aflah Khan, Shivanshu Purohit, USVSN Sai Prashanth, Edward Raff, et al. Pythia: A suite for analyzing large language models across training and scaling, In International Conference on Machine Learning, pages 2397-2430. PMLR, 2023.
[5] Rewon Child, Scott Gray, Alec Radford, and Ilya Sutskever. Generating long sequences with sparse transformers. arXiv preprint arXiv: 105009 2019.
[6] Krzysztof Choromanski, Valerii Likhosherstov, David Dohan, Xingyou Song, Andreea Gane, Tamas Sarlos, Peter Hawkins, Jared Davis, Afroz Mohiuddin, Lukasz Kaiser, David Belanger, Lucy Colwell, and Adrian Weller. Rethinking attention with performers, 2022.
[7] Tri Dao, Dan Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. Flashattention: Fast and memory-efficient exact attention with io-awareness. Advances in Neural Information Processing Systems, 35:16344-16355 2022.
(8] Abhimanyu Dubey, Abhinav Jauhri, Abhinav Pandey, Abhishek Kadian, Ahmad Al-Dahle, et al. The llama 3 herd of models, 2024.
[9] Leo Gao, Jonathan Tow, Baber Abbasi, Stella Biderman, Sid Black, Anthony DiPofi, Charles Foster, Laurence Golding, Jeffrey Hsu, Alain Le Noac'h, Haonan Li, Kyle McDonell, Niklas Muennighoff, Chris Ociepa, Jason Phang, Laria Reynolds, Hailey Schoelkopf. Aviya Skowron, Lintang Sutawika, Eric Tang, Anish Thite, Ben Wang, Kevin Wang, and Andy Zou. A framework for few-shot language model evaluation, 12 2023.
[10] Suyu Ge, Yunan Zhang, Liyuan Liu, Minjia Zhang, Jiawei Han, and Jianfeng Gao. Model tells you what to discard: Adaptive cache compression for 1lms. arXiv preprint arXiv:2310.01800 0000 2023.
[11] Suyu Ge, Yunan Zhang, Liyuan Liu, Minjia Zhang, Jiawei Han, and Jianfeng Gao. Model tells you what discard: Adaptive cache compression for Iims, 2024.
[12] Ankit Gupta, Guy Dar, Shaya Goodman, David Ciprut, and Jonathan Berant. Memory-efficient transformers via top-k attention. CoRR, abs/2106. 0bss2190600889, 2021.
[13] Edward J. Hu, Yelong Shen, Phillip Wallis, Zeyuan Allen-Zine Yuanzhi Li, Shean Wang, and Weizhu Chen. Lora: Low-rank adaptation large language models. CoRR, abs/2106.0968.. 2021.
[14] Benoit Jacob, Skirmantas Kligys, Bo Chen, Menglong Zhu, Matthew Tang, Andrew Howard, Hartwig Adam, and Dmitry Kalenichenko. Quantization and training of neural networks for efficient integer-arithneticonnll inference, 2017.
[15] Albert Q Jiang, Alexandre Sablayrolles, Arthur Mensch, Chris Bamford, Devendra Singh Chaplot, Diego de las Casas, Florian Bressand, Gianna Lengyel, Guillaume Lample, Lucile Saulnier, et al. Mistral 7b, arXiv preprint arXiv:2310.06825, 2023.

---

[16] Albert Q. Jiang, Alexandre Sablayrolles, Antoine Roux, Arthur Mensch, Blanche Savary, Chris Bamford. Devendra Singh Chaplot, Diego de las Casas, Emma Bou Hanna, Florian Bressand, Gianna Lengyel, Guillaume Bour, Guillaume Lample, Lélio Renard Lavaud, Lucile Saulnier, Marie-Anne Lachaux, Pierre Stock, Sandeep Subramanian, Sophia Yang. Szymon Antoniak, Teven Le Scao, Théophile Gervet, Thibaut Lavril, Thomas Wang, Timothée Lacroix, and William EI Sayed. Mixtral of experts, 2024.
[17]
Nikita Kitaev, Lukasz Kaiser, and Anselm Levskaya. Reformer: The efficient transformer. arXiv preprini Sitiv:2001.0445555 2020.
[18] Woosuk Kwon, Zhuohan Li, Siyuan Zhuang, Ying Sheng, Lianmin Zheng, Cody Hao Yu, Joseph E. Gonzalez, Hao Zhang, and Ion Stoica. Efficient memory management for large lan- guage model serving with pagedattention. In Proceedings of the ACM SIGOPS 29th Symposium on Operating Systems Principles, 2023.
[19] Wonbeom Lee, Jungi Lee, Junghwan Seo, and Jaewoong Sim. InfiniGen: Efficient generative inference of large language models with dynamic KV cache management. In /8th USENIX Symposium on Operating Systems Design and Implementation (OSDI 24), pages 155-172, Santa Clara, CA. July 2024. USENIX Association.
[20] Zichang Liu, Aditya Desai, Fangshuo Liao, Weitao Wang, Victor Xie, Zhaozhuo Xu, Anastasios Kyrillidis, and Anshumali Shrivastava. Seissorhands: Exploiting the persistence of importance hypothesis for Ilm kv cache compression at test time. arXiv preprint arXiv: 17/18, 2023.
[21] Stephen Merity, Caiming Xiong, James Bradbury, and Richard Socher. Pointer sentinel mixture models. CoRR, abs/1609.07843, 2016.
[22] Microsoft.
Introducing
phi-3:
Redefining what's www/coom///c/ee /e//e/nnuuuubogg 2024.
possi-
ble with slms.
introducing htt.----rddddiinn----aaa-aa-----
[23] Markus Nagel, Marios Fournarakis, Rana Ali Amjad, Yelysei Bondarenko, Mart van Baalen, and Tijmen Blankevoort. A white paper on neural network quantization, 2021.
[24] NERSC.
Perlmutter system architecture.
https://docs nersc. .erecc.oov/////// /yor/sytems/////
perlmutterraacciieeuus
[25] Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J. Liu. Exploring the limits of transfer learning with a unified text-to-text transformer, 2023.
[26] Luka Ribar, Ivan Chelombiev, Luke Hudlass-Galley Charlie Blake, Carlo Luschi, and Douglas Orr. Sparq attention: Bandwidth-efffiient 11m inference, 2023.
[27] Siddharth Singh and Abhinav Bhatele. AxoNN:
An asynchronous, message-driven parallel framework for extreme-scale deep learning. In Proceedings of the IEEE International Parallel Distributed Processing Symposium, IPDPS 222. IEEE Computer Society, May 2022.
[28] Siddharth Singh, Prajwal Singhania, Aditya K. Ranjan, Zack Sating, and Abhinav Bhatele, A 4d hybrid algorithm to scale parallel training to thousands of gpus. 2024.
[29] Jianlin Su, Yu Lu, Shengfeng Pan, Ahmed Murtadha, Bo Wen, and Yunfeng Liu. Roformer: Enhanced transformer with rotary position embedding, 2023.
[30] Mingjie Sun, Xinlei Chen, J. Zico Kolter, and Zhuang Liu. Massive activations in large language models, 2024.
[31] Hugo Touvron et al. Llama 2: Open foundation and fine-tuned chat models. Technical report, 2023.
[32] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. CoRR, abs/1706.03766. 03762, 2017.
[33] Sinong Wang. Belinda Z. Li, Madian Khabsa, Han Fang, and Hao Ma. Linformer: Self-attention with linear complexity, 2020.
12

---

[34] Thomas Wolf. Lysandre Debut, Victor Sanh, Julien Chaumond, Clement Delangue, Anthony Moi, Pierric Cistac, Tim Rault, Rémi Louf, Morgan Funtowicz, Joe Davison, Sam Shleifer. Patrick von Platen, Clara Ma, Yacine Jernite, Julien Plu, Canwen Xu, Teven Le Scao, Sylvain Gugger, Mariama Drame, Quentin Lhoest, and Alexander M. Rush. Transformers: State-of-the- art natural language processing. In Proceedings of the 2020 Conference Empirical Methods in Natural Language Processing: System Demonstrations, pages 38-45, Online, October 2020.
Association for Computational Linguistics.
[35] Guangxuan Xiao, Yuandong Tian, Beidi Chen, Song Han, and Mike Lewis. Efficient streaming language models with attention sinks. arXiv preprint arXiv:2309. / /7453, 2023.
[36] Guangxuan Xiao, Yuandong Tian, Beidi Chen, Song Han, and Mike Lewis. Efficient streaming language models with attention sinks, 2024.
[37] Zhewei Yao, Xiaoxia Wu, Cheng Li, Stephen Youn, and Yuxiong He. Zeroquant-v2: Exploring post-training quantization in Ilms from comprehensive study to low rank compensation. 2023.
[38] Peiyuan Zhang, Guangtao Zeng, Tianduo Wang, and Wei Lu. Tinyllama: An open-source small language model. Technical report, 2024.
[39] Zhenyu Zhang, Ying Sheng, Tianyi Zhou, Tianlong Chen, Lianmin Zheng, Ruisi Cai, Zhao Song, Yuandong Tian, Christopher Ré, Clark Barrett, et al. H _2 Heavy-hitter oracle for efficient generative inference of large language models. arXiv preprint arXiv:2306. 14048, 2023.
[40] Yukun Zhu, Ryan Kiros, Richard Zemel, Ruslan Salakhutdinov, Raquel Urtasun, Antonio Torralba, and Sanja Fidler. Aligning books and movies: Towards story-like visual explanations by watching movies and reading books, 2015.
13

---

A Comprehensive Dimensionality Analysis
A.1
Complete Sey-Dimeesiinnltty Analysis for All Models
Limmz Asseronn Keys (Pre-eory)
Laam2 70022-7556 Ammmmiio Koys Pars-
Laree Aperiion Koy
H
hakk a
Asyosa
Rusk tES
Liper
Liper
Liarb 705 Autoriinn Koys (iteelettrr
Liame2 708 Atsention Koy. Roost
Liama
(Pre-lenry)
8 8 130
Wilss
www.
BookCoouu
Cork
Arpention Keys Arrer-on
Takii 33% Syleett 2es
( Ameedinn Kays Pooteeorr
TTrylane-1 ABAeman RotoRotor)
Peet-Rotry)
ArteetinAA Keys (FrerRocary)
Mis Arsentor Keys Ronur)
Mocraa (Prokotard)
_.Mtrrd-BBBB Aceeson Kon Part-eaart
S 3 122
Wiker
BeckCrrpus
Liper
Uperr
M.bcrra 8c228 Amaraton Kaga (ProRotary)
8x225
190 andtaaae
Roter!
Pythib Amerr Keys (Post Romay!
Tna Ornua
BookCorpue
20
30
P. 4K. Attoviin Key. Prrrrtarr))
2... MM...MMm444K.. Attentin
025
Figure 8: Rank at which 90% of the variance is explained for pre-rotary and post-rotary keys produced by each layer averaged across all heads (Ranky (900) for different models. We observe that all models exhibit significantly low rank consistently across all datasets.
In this section, present extended dimensionality analysis results (from 3) for all the models we experimented with. Figure X displays the Rank;090 values for all models referenced in Section 3 Our analysis reveals that the low dimensionality of the keys is consistently observed across all models and datasets. Results for all models resemble those shown in Figure 2 of the main text. The models we tested cover wide range of sizes, architecture types (dense vs. MoE), well older and newer architectures trained on various datasets. Despite these differences, our main observation remains robust.
An intriguing trend is the variation in Ranky @90 across layers for different models, indicating that the intrinsic dimensionality of the keys is not uniform across model layers. A potential future direction could be to investigate the reasons for this variation from a semantic perspective.
For more fine-grained analysis, we plot the normalized eigenvalues of the covariance matrix of
the keys for few layers and heads of Lama2-7B, Mistral-7B, and Pythia-6.98 on the WikiText-2 datasct example in Figure 9 Here again, we observe that the explained variance significantly
14

---

decreases after the initial principal dimensions. The results for the other models are similar to the ones shown here.
Llama2-78 78, Layer I, Head 0
Mistral-7B, Layer Head 0
Pytha.6.99 Layer Head 0
0.30
0.30
0.30
Normalsed Eign-rabbes
0.25
Preiloarry Kopry
0.25
Roary
Normilied Egun-valars
0.25
Tre Pra-e
0.20
0.30
0.20
0.15
0.15
0:15
Normelied
0.10
0.10
0.10
0.05
0.05
0.00
20
40
50
80
100
120
0.00
20
40
60
50
100
120
0.00
20
40
60
60
100
120
Feature Dimension: 1 Rank
Feature Dimansion/ Rank
Feature Dimansion Rank
Liama2-78, Layar Huad
Mis Mistrrll78, Layur Huad
Pythisu6 Lapar Hexd
Prs-lccary
Pre-Roory
Pre-lopy
0.25
Porr-Rome Romry
025
Pare-toary
6.5
Reary
0.20
C20
oor
620
6.15
0.15
0.15
Nomalaod
8.10
2.10
645
0.05
200
600
28
20
1000
0.00
28
60
100
120
28
80
120
Faalard Dinorsion Rank
Fastur: Denenion: Ravic
Faxture Dineraion: / Ravk
Figure 9: Normalized eigenvalues of the covariance matrix of the keys produced by Layer Head 1 (top row), and Layer 28, Head 6 (bottom row) of Llama2-7B (left), Mistral-7B (middle), and Pythia-6.9B (right) on the WikiText-2 dataset. We observe that the explained variance significantly decreases after the initial principal dimensions. The dashed lines represent the rank at which 90% of the variance is explained (Rank,,
A.2 Variation of Rank across Attention Heads
In this section, we discuss the variation of the rank at which 90% of the variance is explained (Rank)090) across different heads within a layer for two models: Llama2-7B and Mistral-7B. Figure 10 shows the heatmap of the Ronky090 for the pre-rotary (top) and post-rotary (bottom) keys across all layers and heads for Mistral-7B. We observe that the Ranky090 is considerably lower for pre- rotary keys vs post-rotary keys. Focusing on the pre-rotary keys, we see that the initial layers have a lower rank compared to the later layers. In each layer, there are some heads heads with high-rank values even though the median rank is low. This might indicate that some head in that layer is more important and uses more complex information about the keys. Interestingly for post-rotary keys, we see a pattern where 4 out of the 8 heads in each layer have the same rank. This might have to do with how the rotary embeddings are applied Mistral-7B as we do not see this pattern in Llama2-7B.
Figure shows the heatmap of the Rank, @90 for the pre-rotary (left) and post-rotary (right) keys across all layers and heads for Llama2-7B. We observe similar trend as Mistral-7B where the initial layers have a lower rank compared to the later layers. However, we do not see the same pattern in the post-rotary keys as we saw in Mistral-7B. This might indicate that the rotary embeddings are applied differently in Llama2-77 7B compared to Mistral-7B.
In this section, we examine the variation in the rank at which 90% of the variance is explained (Rank,090) across different heads within a layer for two models: Llama2-7B and Mistral-7B. Figure 10 shows the heatmap of Ronk, @90 for the pre-rotary (top) and post-rotary (bottom) keys across all layers and heads for Mistral-7B. We observe that the Ranky @90 is significantly lower for pre-rotary keys compared to post-rotary keys. Focusing the pre-rotary keys, it is evident that the initial layers exhibit lower rank values than the later layers. In each layer, some heads demonstrate high-rank values, even though the median rank remains low. This suggests that certain heads in those layers are more important and leverage more complex information about the keys. Interestingly. for the post-rotary keys, notice pattern in which 4 out of the 8 heads in each layer share the same rank. This phenomenon may be related to how rotary embeddings are applied in Mistral-7B, as we do not observe this pattern in Llama2-7B. Further investigation needed understand this trend.
Figure I illustrates the heatmap of Rank, @90 for the pre-rotary (left) and post-rotary (right) keys across all layers and heads for Llama2-7B. A similar trend emerges as seen in Mistral-7B, where
15

---

60
Hod
330
Heel
50
Layers
Figure 10: Heatmap showing the rank at 90% explained variance for the pre-rotary(top) and post- rotary(bottom) key vectors across all layers and heads for Mistral-7B.
Figure 11: Heatmap showing the rank at 90% explained variance for the pre-rotary(top) and post- rotary(bottom) key vectors across all layers and heads for Llama2-7B.
the initial layers have lower rank values compared to the later layers. However, the same pattern in post-rotary keys that observed in Mistral-7B is absent here, suggesting that rotary embeddings may be applied differently Llama2-7B compared to Mistral-7B.
For both the models, we can see that in each layer, there are some heads with very high-rank values, even when the median rank low. This might indicate that some heads that layer are more important and use more complex information about the keys. Analysis into chosing the best reduced dimensionality based upon the the distribution of ranks across heads could be a potential future direction.
A.3 Dimensionality Analysis for Queries and Values
While our main focus has been on the dimensionality of the keys, we also performed exploratory analysis on the dimensionality of the queries and values. Figures 12 and 13 show the Ronk @90 for queries and values, respectively, for Llama2-7B and L-amaa-70B. We observe that the queries and values also exhibit low dimensionality across all layers and heads, similar to the keys, while values tend to have a considerably higher dimensionality and close the full dimensionality vectors. This observation intuitively be explained by the fact that both keys and queries used
the value
16

---

Rank 90% Explained Variance
Llama2-7B Attention Queries (Pre-Rotary)
Rank at 90% Explained ExplainedVariance Variance
Llama2-7B Attention Queries (Post-Rotary)
140
140
138
128
100
100
80
60
60
Wikitext
Wikitext
40
20
BookCarpus
BookCorpus
5
10
15
20
25
30
10
15
20
25
30
Layer ID
Layer ID
Rank 90% Explained Variance
Llama3-70B Attention Queries (Pre-Rotary)
Llama Attention Atneetiio Queries (Post-Rotary)
Rank at 90% Expla asplanne
140
128
Wikitext
128
Wikitext
C4
100
100
BaokCorpus
BookCorpus
80
60
60
40
40
20
20
60
20
60
Layer
Layer ID
Figure 12: Rank at which 90% of variance is explained (Ranky @90) for the query vectors across various models and datasets.
Rank at 90% Explained Variance
Llama2-7B Attention Values
Rank at 90% Explained Explaind Variance
Llama3-70B Attention Values
140
140
128
22
100
100
80
60
Wikitext
Wikitext
40
C4
40
C4
20
BookCorpus
BookCorpus
10
20
30
20
60
80
Layer ID
Layer ID
Figure 13: Rank which 90% of variance is explained (Rank, (9an))900 for the value vectors across various models and datasets.
compute the scalar attention scores and thus do not need to be high high-dimensional, while values weighted by these scores and used to compute the final output, and thus need to be high-dimensional capture the complexity of the data.
17

---

B Comprehensive Evaluation Results
B.1 Performance of Loki on Perplexity and Short-Context Downstream Tasks
Loki vs. Full Actention (Perplexit))
0.125
0.55
0.
Llama2-78
Liuma2-138 138
Pro Rotar
Rout-ootry
Pro- Roury
Posi-ouary
0.125
0.5
0.125
0.25
0.25
05
0.125
0:25
0.5
Pro- Rotary
Liama2-708
Past-Rotury
Pre-Rotrry
Liama3-88
Ross-Rotary
Purkoky
10.0
4.0
3.5
7.5
10
0.125
0.125
0.125
Liama3-70B
TinyLama-I.IB 1.11
Pro-Rotary
zxrn Post- Rotrr
Pro- Rouy
Rosta
50
10
25
10
0.125
0.125
0.
3.004
Pre- Pra-Rotrry
Mistral-7B
20
sajar
Pamt-Roury
Pre-
Mixtral-xx7B
Post-Ratary
Perplesty
5.5
6
5.0
1-221
0.125
0.25
0.5
0.125
0.25
0.125
0.25
03
0.125
0.25 kr
0.5
Kr
Kr
d
okl Full Attention (Average Task Accuracy)
0.5
0.25
0.5
Pre-Roory ry
Llama2-78
Pas-- haa-Rocary
Pre-kacary
LJama2-138 33
Retary CO0
45
Arange Tank
30
6.12
0.25
0.5
d.125
0.25
0.5
0.125
0.25
0.5
6.125
0.25
0,5
Llama2-70B
L.ma3-88
Aveug Aakkkuurey
Pre-Rosry
Post. Rotary
Pre-Rosary
Poss Rosary
60
40-
50
witfis
0.125
0.25
0.5
25
0,25
0.5
0.1 1125
0.25
0.5
0.125
0.25
0.5
Llama3-70B
TimyLlamaa I.IB
Avarage TakkAcuus Accuasy
Pre Rosuy
Post: Rotary
Rotary
Poss Rotary
35.0
35.0 15.0
32.5
32.5
40
30.0
0.5
0.125
0.5
0.125
0.35
0.125
0.25
0.5
Mistral-76
Mixtral-8x78
Amerage Taak kyuuak
Pre-Rooryy
Pou- Roory
Pre- Rotary
Poss Romay woga
30
50
10
40
40
40
0.5
0.25
ky
Figure 14: Performance of Loki on Perplexity (top) and Short-Context Downstream Task evaluation for different models using pre-rotary and post-rotary PCA transformation. For each model and each transform type, we run Loki with different values of k and d.
In this section, we provide a detailed evaluation of our method across a wide range of models and tasks. Figure [4illustrates the performance of Loki on perplexity and downstream tasks compared to
18

---

Table 3: Performance of different models compared hugging face baseline with different configura- tions of k and d using pre-rotam PCA transformation.
Model
Method
PPL.
Hellowwagi
TQA1
Winogrande
ARC1
GSMRK1
MMI U:
ArgT
lama2-78
Full Atteniion
5.1101
75.99
38.95
60.06
46.33
13.87
41 14
47.6
Laki
0.5
0.5
5.1195
1.
Loki
75.84
Laki
0.125
75.00
39.07
Loki
0:25
0.5
5.1881
75.73
38.04
67.25
44.20
11.30
39.
46.
I 1amo2-78
Laki
0.25
0.25
5.2185
73.43
38.35
61.61
41.21
7.96
36
43.50
Loki
0.25
0.125
5.3054
53.23
40.08
59.25
36.09
2,81
30.99
37.09
Loki
0,125
0.5
5.4980
70.42
39
52.49)
35.92
7.13
33.22
39.76
Loki
0.125
0.25
6,0729
56.04
42.76
49.57
31,91
2.27
27.15
34,95
Loki
0.125
0.125
23.64
28.98
Temma-1158 138
Full Attantion
4.5680
79.58
36,90
72.22
49.15
22.97
$2.06
$2.11
Loki
0.5
0.5
4.5701
79.34
37.05
72.09
48.81
23.20
52.19
52,28
Loki
0.5
0.25
4.5708
79.27
37.14
72.14
49,40
22.44
$2.00
$2.07
Loki
015
0.125
78.45
50.58
Laki
0.25
4.5979
51.
Loki
Laki
0.25
0.125
71-17
41.33
Loki
0.125
0.5
4.8153
77.38
38.45
56.27
41,64
14.94
48.63
46.22
Laki
0.125
0.25
5.3912
61.85
36.70
13.09
32.08
2.96
36 36.40
$7.03
Loki
0,125
0:125
7.6575
38,67
43,00
50.20
24.32
0.88
23.63
30.08
Full
Loki
0.5
0.5
3.1319
Loki
0.5
0.25
3.1295
83.65
39.75
76.95
56.91
41.93
63.33
60.42
Loki
0.5
0.125
3.1316
82.58
39.33
72.85
54.61
37.45
60.85
57.91
Loki
0.25
0.5
3.2986
M0.54
42.46
75.85
$7.05
22.21
57.14
55.88
I 11m
Laki
0:25
0.25
3,7830
76.05
44.88
63.54
50.26
15.92
51.79
50.41
Loki
0.25
0.125
29.37
Loki
0.125
0.125
46.59
Loki
0.125
0.125
6,4963
30.07
49.19
$1.30
22.78
24.75
29.87
L. Llama1-88
Full Attantion
5.5696
79.17
43.80
72.93
53.24
50.11
62.19
60.26
Laki
0.5
0.5
Loki
57.71
Loki
0.5
0.125
5.5876
74.83
40.41
$6.07
Loki
0:25
0.5
5,5986
76.54
48.32
60.93
43.43
48.56
58.33
54.70
Liama3-8
Loki
0.25
0.25
5.6648
69.42
41.50
50.36
34.64
33.06
44.
45.58
Loki
0.25
0.125
6.0558
56.11
42.14
50.36
27.13
9.17
30.46
35.90
Loki
0,125
0.5
5.7356
66.13
44.00
50.04
28.33
31.77
80.61
43.48
Loki
0.125
0.25
6,5780
45.14
23.89
3.18
26.05
31.43
Loki
0.125
0.125
11.1097
28.64
Llamm3-708
Full Anaation
2.5653
$4.89
45.57
90.43
64.53
80.67
75.03
71.82
Loki
0.5
0.5
2.5656
#5.17
45.66
79.95
63.99
79.91
74.50
71.60
Loki
0.5
0.25
2.5665
84.22
45.78
75.06
39.81
78.77
73.68
68. 29.55
Loki
0.5
0.125
2.5712
82:21
54.78
74.54
70.28
66.25
Laki
0:25
2.5727
$4.09
75.12
68
Llama3
Loki
0.25
Laki
0.25
0.125
2.7577
50.67
31.48
42:31
47.16
Loli
0.125
0.5
2.6285
78.96
46.43
51.14
40.70
74.53
62:19
59,00
Laki
0.125
0.25
2,8796
63.03
41.69
46.33
27.65
$1.19
36.08
44.31
Loki
0.125
0.125
4.1495
39.07
41.09
49.88
23.58
3.03
25.73
30.36
TinyLlamm-
Foll Attention
24.56
36.
Loli
0.5
0.5
8.0040
60.39
38. 38.19
59.98
32.08
1.90
24.62
36. 16.19
Loki
0.5
0.25
S.U342
59.96
38.30
59:27
32.85
2.20
24.33
36.25
Taki
0.5
0.125
1,1057
57.93
39.10
57.14
31.91
1.52
24.08
35.43
Laki
0.25
0.5
$.3475
58.05
411.05
53.17
31.05
1.52
24.83
35.62
TinyL LBB
Loki
0.25
0.25
86652
52.69
42.96
$2.01
29.18
1.29
24.76
33.82
Loki
0.25
0.125
45.43
24.34
31.51
Eoki
0.125
33.38
Loki
0.125
11.5887
Loki
0.125
0.125
19,9200
30.13
48.50
$1.30
24.66
24:11
20.00
Mistral-7B
Full Attention
4.9140
$1.07
42.62
73.95
53.92
38.59
90.65
58. 38.30
Loki
0.5
0.5
4.9147
90.34
59.83
57.60
Laki
0.125
Loki
0.25
0.5
4.9185
79.00
43.41
70.17
49.23
36 16.16
58.25
56.04
Mistral-7B
Loki
0.25
0.25
4.9255
77.63
42.15
62.99
46.59
32.08
53.70
$2.63
Laki
0.25
0.125
4.9986
66.95
39.58
53.64
36.35
14.96
38,
41.43
Loki
0.125
0.5
4.9911
72.66
43.80
$2.25
35.58
33.36
S0.01
47.96
Loki
0,125
0.25
1.9636
65.93
41.12
51.78
29.18
18.42
38.14
40.76
Loki
0.125
0.125
24.00
Mixrrl-Sx72
Pull Attention
3,5967
84,01
48.53
76.32
39.73
58.38
67.90
65.81
Loki
0.5
0.5
3.5979
83.36
46.36
75.53
60.15
$7.32
67.83
66.25
Loki
0.5
0.25
3,6047
83.70
46.70
76.24
59.73
57.01
67.21
65. 10
Loki
0.5
0.125
3,6201
82.91
12.27
73.18
57.42
43.44
65.71
60,87
Loki
11.25
3.0076
82.58
66.72
63. 63.89
Loki
57.17
Laki
0.25
0.125
73.16
43.98
Loki
0.125
0.5
3.6117
76.93
18.21
41.72
50.87
58.
$1.49
Loki
0.125
0.25
335467
70000
37.35
49.17
32.68
11.52
39.25
40.09
Taki
0.125
0.125
6.9799
42.34
43.50
$4.18
24.65
0.45
24-09
31.77
19

---

Table 4: Performance of different models compared hugging face baseline with different configura- tions of k and d using
PCA transformation.
Model
Method
PPL.
Hellawwagi
TQA1
Winogrande
ARC1
GSMRK1
MMI U:
ArgT
lama2-78
Full Atteniion
5.1101
75.99
38.95
60.06
46.33
13.87
41 1/.54
47.6
Laki
0.5
0.5
5.1195
68.59
47.55
Loki
Laki
0.125
Loki
0:25
0.5
5.1838
75.19
38. 16.16
62.12
41.21
10.69
60.42
44.63
I 1lmo2-78
Laki
0.25
0.25
5.2017
72.59
39.16
56.50 59
37:57
10.34
37:
42.
Loki
0.25
0.125
5.4425
08.49
38.83
56.51
32.17
10.92
32.68
39,93
Loki
0,125
0.5
5.3601
70.42
39.40
52.49)
35.92
7:13
33. 22
39.
Loki
0.125
0.25
5,5606
59.98
41.72
48. 15.86
26.96
6.22
28.38
35.35
Loki
0.125
0.125
7.4062
24.13
31,53
Temma-1158 138
Full Antantion
4.5680
79.38
36,90
72.22
49.15
22.97
$2.06
$2.78
Loki
0.5
0.5
4.5731
79.34
37.05
72.09
48.81
23.20
52.19
52,28
Loki
0.5
0.25
4.57.7
79.05
37.46
72.69
48.22
23.58
$1.94
52.17
Loki
015
0.125
78.45
19.85
50.58
Laki
0.25
4.5937
00
Loki
Laki
0.25
0.125
4.8082
Loki
0.125
0.5
4.7029
77.38
38.45
56.27
41.64
14.94
48.63
46.22
Laki
0.125
0.25
4.0668
72.71
40.00
51.14
33.28
9.10
39.3
40.92
Loki
0,125
0:125
6.1436
34,81
46.07
$2,09
24.15
8,79
26.50
32,07
Full
Loki
0.5
0.5
3.1411
83.89
41.32
73.06
$7.68
50.42
64.35
62.69
Loki
0.5
0.25
3.1455
83.00
43.42
76.30
56.51
52.99
64.73
62.99
Loki
0.5
0:125
3,5457
83:41
43.51
75.45
55.89
52.54
64.12
62.49
Loki
0.25
0.5
3.4619
82.36
41.91
76.87
56.45
42.61
60.11
00:06
I 11m
Laki
0:25
0.25
3.5701
81.42
45.25
71.11
49.74
44.28
59.56
58.
Loki
0.25
0.125
Loki
0.125
50.
Loki
0.125
67.03
46,19
Loki
0.125
0.125
4,6898
64.84
50.51
20.95
38.74
30.08
44.67
L. Llama1--8
Full Attantion
5.5696
79.17
43.80
72.93
53.24
50.11
62.19
60.26
Laki
0.5
0.5
Loki
Loki
0.5
0.125
5.7479
71.38
58.64
51.72
Loki
0:25
0.5
5,6783
68.02
42.007
45.78
31.31
43.59
48.06
46.97
Liama3-88
Loki
0.25
0.25
11,4499
$7.39
42.13
45.70
27.90
25.25
38, 35.69)
40.52
Loki
0.25
0.125
13.2883
48.99
42.10
45.07
22.87
12.81
30.50
34.29
Loki
0,125
0.5
6.8025
49.08
41.14
49.25
25.51
31.39
30.36
37.97
Loki
0.125
0.25
16.3507
36.39
43.62
25,09
26.21
32.99
Loki
0.125
0.125
23.61
29.18
Liamm3-708
Full Anaation
2.5653
$4.89
45.57
80.43
64.53
80.67
75.03
71.82
Loki
0.5
0.5
2.5664
$3.00
45.83
72.61
56.14
79.15
75.43
68
Loki
0.5
0.25
2.5697
79.92
46.22
62.90
48.46
78.39
71.11
64.50
Loki
0.5
0.125
2.7810
76.66
56.54
68.48
58.45
Laki
0:25
2.5742
64.14
Llama3
Loki
0.25
Laki
0.25
0.125
5.6725
41.90
47.18
47.59
23.29
5.31
26.08
Loli
0.125
0.5
2.6231
56.24
43.91
50.51
24.66
72.48
38.52
47.72
Laki
0.125
0.25
4.2512
31.01
47.30
50.43
24.57
19.71
24.73
53.12
Loki
0.125
0.125
57.6788
27.01
49.28
50.67
24,06
0.68
24.76
29.41
TinyLlama-
Foll Attention
36.
Loli
0.5
0.5
7,9979
60.17
38.14
58.33
31.57
1.90
25.12
35.87
Loki
0.5
0.25
8.0135
58.78
34.95
54.38
30.55
1.29
24.5%
34,92
Taki
0.5
0.125
8.414
57.77
38.50
54.93
30.89
1.44
24. 39
34.60
Laki
0.25
0.5
$.3190
57.35
37.87
53.83
29.69
1.67
25.13
34.26
TinyL IBB
Loki
0.25
0.25
8.5687
52.40
40.86
49.33
26.96
2.20
23.34
32,51
Loki
0.25
0.125
52.96
0.591
25.10
33.52
Eoki
0.125
24.02
31.70
Loki
0.125
Loki
0.125
0.125
11.3508
39.27
41.55
50.67
22.78
0.45
24.
29.
Mistral-BB
Full Attention
4.9140
$1.07
42.62
73.95
53.92
38.59
99.65
58.
Loki
0.5
0.5
4.9149
58.00
75.99
55.15
Laki
0.125
73.58
Loki
0.25
0.5
5.2052
71.56
40.74
56.04
38.91
26.18
45.56
46.22
Mistral-7B
Loki
0.25
0.25
0.5445
62.02
39..93
43.62
25.17
1,82
30.30
34.65
Laki
0.25
0.125
7.7609
35.51
43.67
53.20
23.65
1.06
23.86
30. 50.16
Loki
0.125
0.5
9.5167
51.73
45.44
$1.62
25.77
3.03
27.99
34,26
Loki
0,125
0.25
13.5597
34,85
16.38
50.20
22.53
0.45
23.60
29, 67
Loki
0.125
0.125
23.64
30.41
Mixrrl-Sx72
Pull Attention
3,5967
84.01
48.53
76.32
39.73
58.38
67.90
65.81
Loki
0.5
0.5
3.5970
83.24
47.32
71.22
58.53
$6.48
67.23
64.51
Loki
0.5
0.25
3.6196
$1.71
43.51
60.61
53.67
55.57
63.92
61.33
Loki
0.5
0.125
3.6635
76.18
11.63
61.72
47.78
49.28
$8.94
58.50
Loki
11.25
0.5
3.00014
57.85
63.04
Loki
3,7906
47.37
Laki
0.25
0.125
4.2566
$7.15
Loki
0.125
0.5
3.6358
72.29
45 25
90.67
33.70
55.50
47.15
50.76
Loki
0.125
0.25
4.5500
52.16
37.36
46.57
23.95
17.13
27.00
34.12
Taki
0.125
0.125
5.5250
45.93
40.33
40.72
23.55
0.91
24.78
31.04
20

---

the full attention baseline. We present results for both pre-rotary and post-rotary PCA transformations. The models evaluated include Llama2-7B. Llama2-13B, alama2-70BB, Llama3-8B, Llama3-70B. laaalllm.1111 Mistral-7B, and Mixtral-8x7B. We assess the models using various configurations of l and d for Loki.
As Ky and d decrease, the model's performance deteriorates, particularly when both are set to 0. 0.1 Notably, the impact of y on performance more pronounced than that of dy. This is evident as ky = 0.125 and dj 0.5 significantly underperform compared to ky 0.5 and d, 0.125 across nearly all models. The configurations of k 0.25 and dj = 0.25, along with ky = 0.125 and dy = 0.5, demonstrate relatively strong performance across all models, striking favorable balance between performance and accuracy, with theoretical speedup of 2.6x for both configurations. Settings with K = 0.5 maintain model quality much more effectively but do not yield significant empirical speedup.
Tables 3 and 4 present finer-grained results for each task and model.
B.2
Variable Dimensionality Analysis
Liama2-78: Fioed va, Variable de poicy 0.25)
Lamm7-11 2B Fived vs. Variatle pokiy (ky 0.25)
Llmma-88 Variable de policy (Kr- 0.25)
Accussy
Verale
Astrage Taal Aouara
gr,,
pos, Niraki)) (yr haok Drt/Wobll
Houdle urc worbied
0.3
0.5 Radi
Radi
Figure 15: Average short-context task accuracies using fixed dy vs. varying dj values across the layers for Llama2-78 7B (left), Llama2- 13B (middle) and Llama3 8B (right). For the variable policy, d, set based on per-layer explained variance (varied from 0.5 0.8). Compression ratio is the average d / across layers.
In this section, we present our experiments utilizing a variable dy policy per layer in Loki. For this experiment, we set dy based on the per-layer explained variance, varying from 0.5 to 0.8, and plotted the average short-context task accuracies (refer to Section against the compression ratio, calculated follows:
Eln d'
Compression Ratio
D
Figure 15 presents these plots for the Llama2-7B, Llama2-13B, and Llama3-8B models. We observe that the variable d policy does not yield significant improvements over the fixed dy policy used in experiments. It is possible that different layers may require distinct explained variance thresholds for optimal performance, indicating that further tuning is necessary. Nevertheless, the simplicity of the fixed d, policy, combined with its comparable performance to the variable dy policy, makes it practical choice for Loki.
21

---

C Comparison of our kernels with SparQ
As mentioned Section 4.3 we create optimized kernels in Triton efficiently compute the three matrix multiplications in Loki (lines 5, 8, and of Algorithm [ without creating temporary dense copies of subsets of the KV-cache. Initially, we planned to use the implementations developed by the authors of SparQ (26.. However, we discovered two major issues with their kernels. Let's say you are multiplying two matrices sizes m and n, then SparQ kernels parallelize compute along only the dimension. However, it is well known that one can parallelize matrix multiplications along the n dimension as well and gain more performance. Thus, we add this extra dimension of parallelism to their triton kernel. Second, their kernels cannot handle non noo-powerrs of 2 number of tokens in the KV-cache, a setting which commonly encountered in inference since we generated keys and values one at time. Therefore, we extend their kernels to handle non-powers of two number of tokens in the KV-cache successfully. In Figure T6, we compare the performance of our kernel with SparQand vanilla PyTorch based attention for attention layer in Llama2-7B for various sizes of the V-ccchhe ranging from 512 to 4096. We do this for the matmul operation of query and keys with top-k as 0.25.
Llama2-77-
Time for computing Q KT with 512 keys and k=0.25
Llama2-7B - Time for computing Q. KT with 1024 keys and k=0.25
0.1
0.2
Vanilla Attention
0.05
SparQ
0.1
Time (ms)
Our kernel
Time (ms)
0.05
0.01
0.01
2
16
8
16
Batch Size
Batch Size
Llama2-78- Time for computing Q. KT
Llama2-7B Time for computing Q.KT
with 2048 keys and k=0.25
with 4096 koys and k=0.25
0.3
Vanilla Attention
SparQ
0.2
Time (ms)
0.1
Our kernel
Time (ms)
0.1
0.05
0.05
0.01
0.01
2
16
8
16
Batch Size
Batch Size
Comparing the performance of proposed kernel for computing Q.KT. with SparQ's 26 kernel for various batch sizes and number of keys in the KV-cache.
Figure 16:
We see very high speedups over SparQ for small batch sizes. For instance, for batch size of with 4096 keys (bottom right), our kernel is faster than SparQ by nearly 2.8x! Infact, the SparQ kernel barely obtains any speedup over vanilla Py PyTorch even though it is only using 25% of the keys (1024 out of 4096). This is because SparQ only parallelizes compute across the m-dimension discussed before. attention the m-dimension is proportional to the batch size, and thus having small batch size limits the amount of parallelism SparQ kernels can exploit. Whereas, given our 2D parallel approach to matrix multiplication, we can exploit parallelism along the sequence dimension and thus still obtain significant speedups over vanilla PyTorch. Note that smaller batch sizes very common in inference and hence is extremely important to design a kernel that can function efficiently even smaller batch sizes. Our speedups overall we are always faster than both SparQ and vanilla PyTorch.
we
SparQ decrease with larger batch sizes, but
22

---

D Generation Examples
In this section, we look generation example from the Llama2-7B (Chat) model given custom prompt for different settings of k and d.
Prompt
world where technology has advanced beyond imagination, society grapples with the consequences of its own creations. The integration of artificial intelligence into everyday life has revolutionized how we live, work, and interact. However, with great power comes great responsibility, and ethical dilemmas abound. Governments struggle to regulate the use of AI, while corporations push the boundaries of what is possible in the pursuit of profit. Meanwhile, individuals navigate a landscape where privacy is increasingly scarce, and the line between human and machine blurs. Against this backdrop, a new generation of thinkers emerges, questioning the very nature of consciousness and what it means to be human. As tensions rise and alliances shift, the fate of humanity hangs in the balance, waiting to be written by those bold enough to seize the pen of destiny.
Llama2-7B Chat Full
Will create a future where AI and humanity coexist harmony, will the pursuit of power and profit lead to our downfall? The choice is ours, and the clock is ticking.
Llama2-7B Chat 1 50%, d 25%
Will we chart a path towards utopia or succumb to the dangers of unchecked progress? Can we find a balance between technological advancement and preserving the essence of humanity? The future is ours to create, and the choice is ours to make. The clock ticks on, waiting for us take the next step forward into an uncertain tomorrow, where the only constant is the endless possibilities of an age of AI
Llama2-7B Chat l 25%, 25%
Will you answer the call to shape the future of AI? Will you play part in shaping the course of human history? The possibilities are endless, and the outcome uncertain. Are you ready embark on journey that could change the world forever?
23

---

E PCAAttn: Directly Leveraging Reduced-Dimensional Attention Scores
Algorithm 2 PCAAtttn
Require: At the 7t step Input: Xs E IIIDD, KV-cache: k,s- Vis-1 E ((S-1)xd, Projection
Matrix: P.d E Dxd (first d principal components)
1: function PCA- ACA-ATTTNTIONNXS, K -1,,d V1-1.PM)
2: qs, ks,ks, Vs compute Q
3: Is,id f PsP :di ks.dd id kg idd
4: R.s.d + (onnaa((Kss-1..) ks..d)
5: V.s concat 1-,,, Vs)
6: a == softmarl (
vD
7: return aV.s :s
8: end function
One other approach we tried is to directly use the formulation in to compute the final attention scores. More specifically, we compute the PCA transformed query and key vectors, projected onto the first d principal components, and then compute the attention scores, We only store the reduced dimension key vectors in the KV-cache. We call this method PCAAtı (Algorithm 22.
Compute and Memory Analysís; When computing attention between a single query qs E RRID and the key vectors K.s RSDD the matrix multiplication Ks has complexity of O(DS). Using Att, the key and query vectors are reduced to dimensions and the complexity of the matrix multiplication is reduced to (dS)) Thus, we can get speedup of D/d in the attention dot product computation. The PCA transformation of the query and key vector generated at each step has complexity of (D2)), which small when S >>> D. The KV-cache memory requirement reduced by a factor of 0.5 D/d because we only reduce the key vectors to d dimensions and not the values. Additionally, the PCA adds significantly small memory overhead of O
Experimental Results:
Table 5: Performance of PCAAttn with various cache configurations.
Model
Method
hy
d,
Perplexity
Hellaswagt
Winogrande
MathQA 1
OpenhookOA 1
RTE
COPA 1
Mamm2-7H
Full Attention
31.4
62.8
Exact TopK
0.5
5.1191
57.2
68.9
28.3
31.2
63.9
86.0
Llama2.7B
H₂O
27.4
PCAAm
0.5
38, 9997
33.3
53.2
21.7
14.2
50.5
75
Topk
0:25
5.1799
56.9
68.6
29,4
29
66.4
76.0
Llama2-78
H₂O
0.25
5.2809
50.1
$1.6
21.1
17.8
55.2
55.0
PCAAttn
0.25
243.2631
26.9
48.5
20.5
11.4
49.1
65.0
Misrra-7B
Full Attention
4.9140
61.2
73.9
35.7
32.2
66.8
91.0
TopK
35.6
Mistral-BB
H₂O
0.5
4.9560
59.4
58.6
26.4
23.0
62.4
71.0
PCAAttn
0.5
396.8966
31.4
50.4
22.5
15.6
53.4
72.0
Tapk
60.4
73.0
35.4
30.0
65.3
Mistral-7B
H,O
0.25
5.0805
52.7
49.7
21.9
17.4
$2.0
56.0
PCAAtt
0.25
933,6016
27.2
52.2
21.6
13.6
53.0
63.0
Table shows the performance of PCAAttn on Llama2-7B and Mistral-7B models. We can see that our PCAAttn method performs poorly compared all the baselines and the H₂O method for all cache configurations We believe that this happens because the application of rotary embeddings increases the dimensionality of the key vectors and using reduced dimensionality to store the keys results in loss of information. To further investigate this, we look back at Figure TT) which shows the rank at 90% explained variance for the key vectors across all layers and heads. Even though, the average rank per layer is around 50% of the full dimensionality, the rank for some layers and especially some heads within each layer is much higher. Due to the poor performance of PCAAttn, we do not include it in the final results and decide to focus Loki instead in the main paper. Note, that we only tried the post-rotary transformations PCAAttn, and is possible that pre-rotary transformations might perform better, but we leave this for future work.
24

---

Estimate of Compute Resources Required to Replicate our Experiments
As mentioned in Section we conduct all of our experiments on Perlmutter, multi-GPU cluster with 4 A100 GPUs per node. Since we do not do any training/finettningg tuning. our experiments can be done on very small number of GPUs. For instance, all of our runs involving models with 7B and 13B parameters were done on a single A100 GPU. For models larger than this (like LLama2- 70B, Llama3-70B), we had to resort to running on four A100 GPUs (or single node) with tensor parallelism using the AxoNN parallel deep learning framework. All results for 7B and 13B sized models can be compiled within hours. For larger models like the 70B Llama-2 and as well as Mixtral models, the total times for computing all results are in the ballpark of 10 hours. Our compute benchmarking runs of Llama- 13B are very short and can be completed within 5 minutes.
25

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
The claim around investigating the intrinsic dimensionality of attention keys is addressed in Section The claims around designing Loki, providing theoretical analysis, and providing an optimized kernel are addressed in Section[4 [4] Finally. the model quality analysis and compute analysis showing the efficacy and speedup Of Loki addressed in Section 6
Justification:
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings. It is fine to include aspirational goals as motivation long it is clear that these goals are not attained by the paper.
2. Limitations
Question:
Does the paper discuss the limitations the work performed by the authors?
Answer: [Yes]
Justification: Limitations are discussed in Section
Guidelines:
The answer NA means that the paper has limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g.. independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or speech-to-text system might not be used reliably provide closed captions for online lectures because it fails handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairesss.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't acknowledged
the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory Assumptions and Proofs
26

---

Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Yes]
Justification: The paper introduces two lemmas in Section 4 and provides proof's for both in the same section.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated referenced in the statement of any theorems.
The proofs can either appear the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Question:
Answer: [Yes]
Algorithm provide a detailed description of our method to allow for repro- ducibility. Further, we detail the setup of our experiments in Section 6 to allow for the experiments to be reproduced. For our dimensionality analysis, our method is detailed in Section We also make our code available publicly with the link provided in the paper.
Justification:
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution a specific model and empirical evaluation, may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often good way to accomplish this, but reproducibility instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing of a model checkpoint, other means that appropriate to the research performed.
also be provided via detailed
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution is primarily a new algorithm, the paper should make clear how to reproduce that algorithm.
(b) If the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset instructions for how to construct the dataset).
27

---

(d)
We recognize that reproducibility may be tricky in some cases, in which case authors are welcome describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
Justification:
Our work is conducted under the MIT License and our code is released under the MIT License. The link the Github repository provided the paper. For the original paper submission, we provided our code in the supplemental material.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips.co cc/
public//uuidee//CdeSuttiii wow.eoobbssoioooooooooooooooooie for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https
//nips https///cccpllicgggeeeeeee/eeeee
for more details.
The authors should provide instructions data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: We mention the datasets, hyperparameters and metrics used in our experiments in Section and Section 6
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental material.
7. Experiment Statistical Significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [No]
28

---

Justification: For our ML experiments, we do not repeat the runs multiple times to calculate error bars, but when conducting the experiments using standard benchmarks, noticed very little variation in the perplexity and downstream task results. For our compute benchmarks, we did repeat the runs multiple times and observed a standard deviation of less than 0.06% in the results.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
It should be clear whether the error bar is the standard deviation the standard error of the mean.
It is OK to report I-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables figures symmetric error bars that would yield results that are out range (e.g. negative error rates).
error bars are reported in tables or plots, The authors should explain the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of exccution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification: We mention the compute resources used in experimental setup in Section 6 and discuss the compute resources in detail in
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make into the paper).
9. Code Of Ethics
Question: Does the research conducted the paper conform, in every respect, with the NeurlPS Code of Ethics htttps://neuriii cc. http.///pc.cceeetcccssseeeeeeeee
Answer: [Yes]
Justification:
The paper conforms the NeurIPS Code of Ethics.
Guidelines:
The answer NA means that the authors have reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special cireumstances that require deviation from the Code of Ethics.
29

---

The authors should make sure to preserve anonymity (e.g., if there is special consid- eration due to laws or regulations their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [Yes]
Justification: The paper motivates the challenging problem of LLM inference in Section and discusses the impact of our method optimizing the inference process, leading to reduced compute costs and thereby, energy consumption.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g.. disinformation, generating fake profiles, surveillance), fairesss considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification:
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
30

---

Answer: [Yes]
Justification: Throughout our paper, any exisiting methods or datasets are properly cited, and their licenses are respected. When using code from other sources for benchmarking our method, we ensure that the code is used in accordance with the license with the license being properly cited and copied. Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset used and, if possible, include URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset. For scraped data from a particular source (e.g.. website). the copyright and terms of service of that source should provided.
assets are released, the license, copyright information, and terms use in the package should be provided. For popular datasets, paperswi thcode. com com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided. If this information is not available online, the authors are encouraged to reach out to the asset's s creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [Yes]
Justification: We do not have any models datasets. Our code is released with proper documentation under the MIT License and the link to the Github repository is provided in the paper.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification:
The paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According to the NeurlPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
31

---

15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based the requirements of your country institution) were obtained?
Answer: [NA]
Justification: The paper does not involve research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with
human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent)
may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
32


Reviews:
[{"review_id": "UhnfXSCWAM", "reviewer": "Reviewer_vRA7", "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.", "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?\n2. 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?\n3. 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?\n4. 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?\n5. 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?\n6. 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?\n7. Handling of Longer Sequences: How does Loki handle sequence lengths longer than those in the evaluation datasets?\n8. Quantization Integration: Can the PCA-TopK mechanism be integrated with other optimization techniques like quantization?", "limitations": "N/A", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "Empirical Validation: Extensive evaluations demonstrate significant speedups with minimal accuracy degradation. \n\nTheoretical 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. \nMemory Footprint Consideration: While the approach excels in computation, it does not address memory footprint reduction, a limitation when compared to other sparse attention techniques.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their feedback and appreciate that the reviewer recognizes the strength of our extensive empirical evaluation backed by robust theoretical analysis.\n\n### **Weaknesses**\n**Implementation Complexity**:\nWe agree that achieving speedups with our method requires custom Triton kernels. We will open-source our kernels for wider use. Integrating with techniques like Flash Attention would involve a more complex implementation, as is the case with most modern attention methods.\n\n**Memory Footprint**:\nWe agree that our method does not reduce memory usage. It is designed as an efficient Top-K token selector without deleting tokens from the KV-Cache (which can incur a large accuracy penalty as seen with H2O). Our method can complement token-eviction methods like H2O by selecting Top-K tokens from a reduced KV-Cache.\n\n### **Questions**\n**Model Class and Size Performance Variance: How does Loki's performance vary with different model classes and sizes not included in the paper**\n\nAppendix B,  Figure 11 and Tables 3 & 4 (in the paper) show perplexity and downstream evaluation of our method on models of different model classes covering dense and MoE models. Further, for each class, we show results for models of different sizes (7B, 8B, 13B, 70B, 8x7B, 8x22B).  Appendix A, Figure 7 (in the paper) shows the dimensionality analysis on an even larger set of model classes and sizes.\n\n**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?***\n\nIntegrating PCA-TopK with Flash Attention (FA) or Ring Attention is theoretically possible, but performance challenges, especially with the Top-K operation, must be addressed to achieve speedups. While we haven't resolved all integration issues, we have a rough algorithm for FA integration:\n1. Load the query into shared memory \n2. Additional Loop:\n    - Load the first $d_f$ dimensions of the keys into shared memory \n    - Compute the approximate scores for each block of keys and update top-k indices\n3. Original FA Loop: Load blocks of top-k keys (based on the indices computed) with full dimensionality and follow the FA algorithm.\n\n**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?**\n\nMethods like H2O and SnapKV [1] delete tokens to save KV-Cache memory. In contrast, PCA-TopK selects the Top-K tokens from the entire KV-Cache to reduce memory bandwidth without deleting tokens. Hence, PCA-TopK is orthogonal to H2O/SnapKV. A combined strategy could involve deleting tokens with H2O/SnapKV, and then selecting Top-K tokens from the retained cache.\nWe will include a comparison of standalone PCA-TopK with other methods apart from H2O in future work.\n\n**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?**\n\nWe will fix Figure 6 (Left) in the subsequent versions of the paper.\n\n**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?**\n\nPCA-TopK is an approximate method and there is information loss when\nreducing the dimensionality. We try out different values of $d_f$ (reduced dimensionality) and find a good tradeoff between compute performance and accuracy. We also pick the best transform between post-rotary and pre-rotary. Figure 2 in the paper shows that RoPE indeed increases the dimensionality but it is still low enough that settings like $d_f = 0.25$ (25% dimensionality) work well enough for a good compute-accuracy tradeoff (Figure 3, 6 in the paper)\n\n\n**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?**\n\nWe discuss how our method can be used with token-eviction methods as a response to Question 3. PCA-TopK can theoretically be used with quantization, as PCA-TopK reduces dimensionality while quantization reduces bits per value. Model pruning [2] reduces model parameters, which is orthogonal to selecting tokens from the KV-Cache. We are unsure if pruning affects the low dimensionality of key vectors and leave that analysis for future work.\nWe have not quantitatively analyzed combining PCA-TopK with other approaches and are unsure of the practical performance.\n\n**Handling of Longer Sequences: How does Loki handle sequence lengths longer than those in the evaluation datasets?**\n\nFigure 1 in the rebuttal pdf shows the evaluation of PCA-TopK on LongBench [3] long-sequence benchmark for the Llama-2-7b-Chat model. We have included a discussion on the same in the Global Rebuttal.\n\n**Quantization Integration: Can the PCA-TopK mechanism be integrated with other optimization techniques like quantization?**\n\nWe discuss how our method can be used with quantization as a response to Question 6. \n\n-----\n\n### **References**\n[1] Li, Yuhong, et al. \"Snapkv: Llm knows what you are looking for before generation.\" arXiv preprint arXiv:2404.14469 (2024).\n\n[2] Han, Song et al. “Deep Compression: Compressing Deep Neural Network with Pruning, Trained Quantization and Huffman Coding.” arXiv: Computer Vision and Pattern Recognition(2015): n. Pag.\n\n[3] Bai, Yushi, et al. \"Longbench: A bilingual, multitask benchmark for long context understanding.\" arXiv preprint arXiv:2308.14508 (2023).", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "53KbwMRVus", "reviewer": "Reviewer_aGJV", "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.", "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.", "limitations": "yes", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "The study of low intrinsic dimensionality of attention keys across multiple models and datasets is insightful.\nTheoretical support is provided through lemmas and proofs.\nThe evaluation includes several models, tasks and datasets, showing it could be generalized.\nThe authors developed optimized Triton kernels for practical implementation.", "weaknesses": "The method does not seem to reduce the memory usage.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their feedback and appreciate that the reviewer finds our observation of the low intrinsic dimensionality of keys insightful and supported by theoretical analysis and extensive evaluation.\n\n### **Weaknesses**\n**The method does not seem to reduce the memory usage**\n\nWe agree that our method does not reduce memory usage. It is designed as an efficient Top-K token selector without deleting tokens from the KV-Cache (deleting tokens can incur a significant accuracy penalty, as seen with H2O). Our method can complement token-eviction methods like H2O by selecting Top-K tokens from a reduced KV-Cache.\n\n### **Questions**\n**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.**\n\nAs noted, pre-rotary transforms outperforming post-rotary transforms is an intriguing observation for which we currently lack a clear explanation. We still have not been able to develop a concrete understanding of why this is the case. A naive intuition we have is that when computing the PCA transform over the post-rotary keys, it captures the distribution of token representations occurring at specific positions in the calibration dataset. During inference, the same token can appear at any other position. The pre-rotary transform captures the distribution of only tokens with significantly less positional information so it can generalize better. \n\nWhile Lemma 4.2 (in the paper) proves that using the post-rotary transform is a good approximation, it does not show that it is the best approximation due to the variational upper bound formulation. Hence, the pre-rotary transform might provide a better approximation.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "2ob8YlqOC2", "reviewer": "Reviewer_KXhq", "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.", "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?\n2. Given the differences in Rankl@90 across layers, what is the impact of varying the policy depending on the layer?\n3. If post-training with PCA-TopK, will it yield better performance?", "limitations": "N/A", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "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.\n2. 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).\n2. 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?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their feedback and appreciate that they recognize the potential impact of our low-dimensional observation and the strength of our evaluation.\n\n### **Weaknesses**\n**There are some typos (lines 44, 150) in this article, and some figures are unclear with text overlaps (Figures 2, 6).**\n\nWe will fix these issues in the subsequent version of the paper.\n\n**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?**\n\nSPAR-Q incurs significant overhead as it needs to store the keys in KV-Cache twice – once in a row-major and once in a column-major format. This is required to get performant kernels for their implementation. Our method does not incur such a massive memory overhead (the only overhead we have is of storing the PCA projections, which are comparatively much smaller). \n\nAdditionally, our $Q.K^T$ kernels are more performant than SPAR-Q as compared in Appendix E, Figure 13 (in our paper).\n\nWe have not evaluated SPAR-Q ourselves. Their paper does not have a trade-off curve between acceleration ratio and effect. The only commonly evaluated dataset they have is TriviaQA [1] where PCA-TopK also achieves good performance (Figure 4 in our paper). Their attention latency results can be found in Figure 8 of *their* paper while Figure 9 shows end-to-end performance but evaluated on CPU. We will include end-to-end quantitative comparison with their method in subsequent versions of the paper. \n\n### **Questions**\n\n**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?**\n\nFigure 3 in the rebuttal pdf shows the dimensionality analysis for query and value vectors for Llama2-7B and Llama3-70B models. It can be seen that while query vectors exhibit low dimensionality similar to key vectors, value vectors have a significantly higher rank. We will include this analysis for more models in the updated supplementary material.\n\n**Given the differences in Rankl@90 across layers, what is the impact of varying the policy depending on the layer?**\n\nFigure 2 (Right) in the rebuttal pdf shows the evaluation of varying the $d_f$ parameter per layer based on the explained variance of the PCA components for Llama3-8B. We compare two policies: (1) Fixed $d_f$ for all layers (set to either 0.25 or 0.5), (2) Variable $d_f$ per layer. For the variable policy, we select the $d_f$ of a layer based on the explained variance threshold (ranging from 0.5 to 0.8). The compression ratio is defined as the average of $d_f/D$ across all layers, where $D$ is the full dimensionality of the vectors. We can see that using the variable policy does not show any benefit over the fixed policy. This indicates that different layers may require different variance thresholds as well, and therefore, tuning the variable policy is key to getting gains. Our evaluation on Llama2-13B also shows the same trend (plot not included due to space constraints). \n\n**If post-training with PCA-TopK, will it yield better performance?**\n\nPCA-TopK is a post-training method and only requires the PCA transforms to be calculated on a set of keys generated (in inference mode) over some calibration dataset.\nIf the reviewer is suggesting adding a fine-tuning step for better performance, we think that is a good suggestion. In that paradigm, we also feel it might be more optimal to keep the model fixed and train a good transformation that minimizes the difference between the reduced dimensional and original attention scores.\n\n----\n\n### **References**\n[1] Joshi, Mandar, et al. \"Triviaqa: A large scale distantly supervised challenge dataset for reading comprehension.\" arXiv preprint arXiv:1705.03551 (2017).", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "MTUVBd8cIS", "reviewer": "Reviewer_dYWK", "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.", "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.\n2. 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)\n3. 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.)\n4. Some figures are not sized properly. (e.g. Fig. 6.) Can you adjust the size of figures to avoid stretching the texts?\n5. I think there should be some plot that show latency-performance trade-off (ms - accuracy) Can you add this plot using some downstream tasks?\n6. 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.\n\nI 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.", "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.", "rating": 8, "confidence": 5, "soundness": 4, "presentation": 4, "contribution": 4, "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)", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their feedback and appreciate that they recognize the simplicity of the approach to get a $Q.K^T$-approximator without any gradient-based training.  \n\n### **Weaknesses**\n\n**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 , therefore, should be fine. (NOTE: I need to checkout the implementation whether this is correct or not)**\n\nWe would like to note here that the current implementation of PCA-TopK is not a fused implementation like Flash Attention. We discuss the possibility of developing a fused method in our response to Reviewer vRA7 (Question 2). In the current implementation, the projection to PCA space happens in a separate kernel. This projection operation is implemented as a standalone matrix multiplication using PyTorch. We do agree with the reviewer that a more efficient implementation would keep the PCA vectors in shared memory to minimize the GPU memory access and that is possible since the per-head PCA transformation is only $D \\times D$ dimensional, where typically $D$ = 128  (sometimes 256) for current LLM models. \n\n### **Questions**\n\n**The .gz file seems broken. Can you upload the file in other format again? I want to look into source codes, for details.**\n\nBased on the comment submitted, we assume this is resolved for now. If needed, we can submit an anonymized link to our code as an official comment to the ACs following NeurIPS guidelines.\n\n**Fig. 1. might leads mis understand that this attention mechanism 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)**\n\nWe will update the figure in the subsequent version of the paper.\n\n**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 sequence length in left.)**\n\nYes, we agree that Top-K operation is slow and while it is possible to fuse the Top-K operation with other matmuls, it still involves synchronization among threads. Currently, we are looking at ways to bypass the Top-K operation itself using a score thresholding mechanism (where the threshold is sampled offline). We leave that to future work. \n\nFigure 2 (Left & Mid) in the rebuttal pdf shows the kernel-wise breakdown for more prompt and generation lengths. It shows that the Top-K operation is almost as costly as the matmuls. A custom (or fused) kernel or a thresholding strategy might help alleviate this cost. \n\n**Some figures are not sized properly. (e.g. Fig. 6.) Can you adjust the size of figures to avoid stretching the texts?**\n\nWe will fix the figures in the subsequent version of the paper\n\n**I think there should be some plot that show latency-performance trade-off (ms - accuracy) Can you add this plot using some downstream tasks?**\n\nFigure 1 (Top-Right) in the rebuttal pdf shows a latency-accuracy trade-off curve evaluated on LongBench [1] for the Llama2-7b-Chat model (sequence length of 4096). We can see that the settings with $k_f = 0.25$ and $k_f = 0.125$ perform much better than vanilla attention. $k_f$ has a larger impact on the performance than $d_f$, which is supported by our theoretical speedup analysis in the paper. Overall, the configurations: $(k_f = 0.25, d_f = 0.25)$ and $(k_f = 0.125, d_f = 0.5)$ give a good accuracy-performance trade-off for LongBench tasks as well.  \n\t\n**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.**\n\nFigure 1 (Top-Left and Bottom) in the rebuttal pdf shows the performance of PCA-TopK on LongBench [1] tasks for the Llama-2-7B-Chat model. We have included a discussion on the same in the Global Rebuttal.\n\n### **Limitations**\n\n**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.**\n\nFor clarification, our current method does not require PCA to be performed during the training stage. We compute the PCA transforms by evaluating the trained model over a calibration dataset, storing the generated keys, and then computing the PCA. We agree that an approach that utilizes our low-dimensional observation during training or fine-tuning will be interesting to study and may improve model performance. We also agree that additional effort is required to optimise the GPU kernels further.\n\n----\n\n### **References**\n\n[1] Bai, Yushi, et al. \"Longbench: A bilingual, multitask benchmark for long context understanding.\" arXiv preprint arXiv:2308.14508 (2023).", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
We thank the reviewers for their valuable feedback.

### **PCA-TopK's performance on tasks with longer contexts**

We ran the LongBench [1] long-sequence benchmark for the Llama2-7B-Chat model with PCA-TopK and compared its performance with Full Attention in Figure 1 of the rebuttal pdf.  Figure 1 (Top-Left) shows that PCA-TopK performs well on all LongBench task categories for $(k_f = 0.25, d_f = 0.25)$, with Multidoc-QA having the most significant accuracy drop from ~22% to ~18%. Figure 1 (Bottom) shows a finer-grained evaluation on all LongBench tasks with two configurations: $(k_f = 0.25, d_f = 0.25)$ and $(k_f = 0.125, d_f = 0.5)$. For this particular model, we see that, on average, the post-rotary transform is better than the pre-rotary one, and $(k_f = 0.25, d_f = 0.25)$ is the better configuration. These observations are consistent with the evaluation of the tasks in the paper. 

To examine the accuracy vs. performance trade-off, we plot overall LongBench accuracy along with attention latencies (computed in a micro-benchmark) for different configurations of $k_f, d_f$ in Figure 1 (Top-Right). Attention times were computed with a prompt length of 3500 and a generation length of 512 to match LongBench's values for the Llama2-7B-Chat model. Due to the slow cache updates in HuggingFace, we lack an end-to-end framework but plan to integrate our attention method with an inference framework like vLLM soon. Nonetheless, our method shows potential for up to 40% attention speedups with minimal accuracy degradation in long and short-context tasks. We will include other models supported by LongBench in the camera-ready version of the paper, if accepted.

### **Memory usage**

We acknowledge the limitation pointed out by the reviewers that our method does not reduce memory usage. Our method efficiently selects Top-K tokens from the KV-Cache without deleting tokens. As we have demonstrated with H2O, deleting tokens can have significant accuracy penalties. PCA-TopK can be used in conjunction with memory reduction methods such as token-eviction or quantization, and a combined approach may lead to better memory, performance, and accuracy trade-offs.

### **Dimensionality of query/value vectors and Variable $d_f$ policy**

We analyzed the dimensionality of query and value vectors (Figure 3 in the rebuttal pdf), finding that while query and key vectors share the low-dimensional observation, value vectors do not. These findings can inspire further research into LLM properties and sparsity exploitation for improved performance. We also experiment with a variable policy for setting $d_f$ per layer instead of a fixed policy across all layers. For the variable policy, we set $d_f$ of every layer based on an explained variance threshold (varied from 0.5 to 0.8). Figure 2 (Right) in the rebuttal pdf shows the results of this evaluation on Llama3-8B, illustrating that the variable policy is no better than the simple fixed policy. We also evaluate on Llama2-13B and see a similar trend (plot not included in the pdf due to lack of space).

----

### **References**
[1] Bai, Yushi, et al. "Longbench: A bilingual, multitask benchmark for long context understanding." arXiv preprint arXiv:2308.14508 (2023).

Meta-review: This paper proposes a PCA-based sparse attention mechanism for LLMs by leveraging the low-dimensionality of key vectors in the attention block. The method  helps approximate the attention and yields inference speedup. Evaluation results demonstrate the efficacy of  this approach for multiple LLMs.

Overall, the method is interesting & authors provide a nice explanation & theoretical analysis of the technique. The reported extensive evaluations also demonstrate the generalizability of this method for different Transformer models, datasets & tasks which is impressive.

The reviewers point out a few suggestions for improvements -- the authors are encouraged to incorporate these to further improve the paper. In particular, the issue (or rather non-improvement) wrt memory footprint was raised. One suggestion — even though this is not a focus for the paper, it might be worth carving out some space to discuss this (maybe in future work or aux. section), how the compute/memory complexity tradeoff for your approach outweighs existing memory-reduction approaches & potential future improvements.
Decision: Accept (poster)