Title: Loki: Low-rank Keys for Efficient Sparse Attention
Venue: NeurIPS 2024
Accepted: yes
Keywords: Approximate Attention

Low-Dimensional Attention Keys

PCA

Top-K Attention
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

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 selfattention 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 lowerdimensional 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.

1 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 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 computational 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 or using a subset of the tokens in the KV-cache [11] (sparse attention).

In contrast to other sparse attention approaches that either permanently prune tokens from the keyvalue cache [39] or impose a 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 a critical observation: across a range of LLMs and datasets, key tensors consistently occupy a significantly lower-dimensional space than the full attention head dimension. For instance, in Figure 1 (left), we show that across various LLMs [8, 16], 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).

OFFLINE

DURING INFERENCE

Llama2-7B

Llama2-13B

Llama2-70B

Llama3-8B

Llama3-70B

Mistral-7B

Mixtral-8x7B

Mixtral-8x22B

0

20

40

60

80

100

120

128

140

Rank at 90% Explained  Variance

Dimensionality of Attention Keys

X

X

Approx.

Attention

Scores

Query

PCA Transform

PCA

Calibration

Dataset

LLM

X

Final

Attention

Scores

Generate

& S a ve

Select

T op-K T okens

Based on this observation, we introduce Loki, a sparse attention method that leverages the lowdimensional structure of key vectors to 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%) to compute approximate attention scores during inference. This dimensionality reduction, informed by our previous observation that key vectors have low effective rank, allows us to efficiently identify the topk (12.5-25%) most relevant tokens using the approximate scores. For these selected keys, we then revert to the full dimensionality to compute the final attention scores, ensuring both efficiency and accuracy. Figure 1 (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 on top of the original self attention algorithm. Thus, we implement optimized sparse matrix multiplication kernels for Loki in Triton, leading to a speedup of up to 45% over the standard HuggingFace Transformer'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: a 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 PyTorch.

Evaluation of Loki 1 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 a 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:

where q ∈ R 1 × D is the query, and K ∈ R S × D and V ∈ R S × D are the key and value caches respectively. Additionally, newer transformer models add Rotary Position Embeddings (RoPE) [29] to the keys and query, before computing the attention scores. Since every query attends to all past keys, the mechanism has a quadratic complexity O ( S 2 ) in number of input + generated tokens.

2.1 Related Work

Numerous studies have explored the low-rank structures in transformers for various purposes. Linformer [33] demonstrated that the attention score matrix is low-rank and proposed alternative low-rank attention formulations during training for linear computational complexity. LoRA [13] showed that parameter updates to a transformer model during fine-tuning reside in a 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 [5] 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 a 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 to permanently delete tokens from the KV-cache based on some heuristic. StreamingLLM [35] uses initial tokens and a rolling KV-cache for processing infinite-length 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 et al. [10] 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 a 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-based approximate approaches [14, 23] are complementary to our work and can be applied in tandem.

SparQ Attention [26] is a recent work that inspires our approach. They use high-magnitude query dimensions and corresponding key dimensions for approximate attention scoring, followed by computing the full attention scores for the topk keys. However, their method requires costly noncontiguous 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 a more efficient slicing operation.

A concurrent work, InfiniGen [19], accelerates attention by pre-fetching topk keys from CPU to GPU memory, using SVD-based low-rank approximation of the attention scores. While their lowrank 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 low-dimensional nature of attention keys applied in other contexts.

3 Dimensionality Analysis of Attention Keys

As noted in Section 1, 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 [31], Llama-3 8B/70B [8], TinyLlama-1.1B [38], Pythia-6.9B [4], Mistral-7B [15], Mixtral-8x7B/8x22B [16], and Phi3-Mini-4K [22] on three popular English language datasets: WikiText-2 [21] (Validation Split), C4 [25] (Custom Split), and BookCorpus [40] (Custom Split). Custom splits are used for datasets where the validation split is 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 on all the keys generated for each layer and head individually.

The metric we use in our analysis is the rank at which v %of the variance is explained by the principal components. We calculate this metric for each layer and head of the models as follows:

where, λ j l,h is the j th normalized eigenvalue of the covariance matrix of the keys for layer, l and head, h . We average this metric ranks across all heads of layer, l and refer to it as Rank l @ v .

3.2 Findings and Discussion

Figure 1 (left) shows the average Rank l @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 a layer-wise analysis for a few models: Llama2-7B, Llama3-70B, Mixtral-8x7B, and Phi3-Mini-4K in Figure 2. The results for the other models are similar and can be found in Appendix A.1.

0

10

20

30

0

20

40

60

80

100

120

128

140

Rank at 90% Explained  Variance

Llama2-7B Attention Keys (Post-Rotary)

0

20

40

60

80

0

20

40

60

80

100

120

128

140

Llama3-70B Attention Keys (Post-Rotary)

0

10

20

30

0

20

40

60

80

100

120

128

140

Mixtral-8x7B Attention Keys (Post-Rotary)

0

10

20

30

0

20

40

60

80

96

100

120

140

Phi3-Mini-4K Attention Keys (Post-Rotary)

0

10

20

30

Layer ID

0

20

40

60

80

100

120

128

140

Rank at 90% Explained  Variance

Llama2-7B Attention Keys (Pre-Rotary)

0

20

40

60

80

Layer ID

0

20

40

60

80

100

120

128

140

Llama3-70B Attention Keys (Pre-Rotary)

0

10

20

30

Layer ID

0

20

40

60

80

100

120

128

140

Mixtral-8x7B Attention Keys (Pre-Rotary)

0

10

20

30

Layer ID

0

20

40

60

80

96

100

120

140

Phi3-Mini-4K Attention Keys (Pre-Rotary)

Wikitext

C4

BookCorpus

Full Rank

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 Rank l @90 for a 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 a 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 l @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.1.

These results indicate the existence of the following properties: (1) The keys produced by the attention layers of transformer models lie in a significantly lower-dimensional space. (2) The lower-dimensional 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 topk (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-transformed 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 PCA-transformed space.

Lemma 4.1. Let D be the dimension of an attention head and P ∈ R D × D be the PCA projection matrix of key tensors calibrated offline on a dataset. Assuming we are generating the S th token in the sequence, let q S ∈ R 1 × D be the query vector for the S th token, K : S ∈ R S × D be the key vectors, including the past ( S -1) keys and the current key. Then, the attention scores computed using the PCA-transformed query and keys are equivalent to the attention scores computed using the original query and keys.

Proof. Let ˆ q S = q S P and ˆ K : S = K : S P be the PCA transformed query and key vectors. Focusing on the dot product term in the attention computation (Equation 1), we have:

It is important to note here that Lemma 4.1 holds for any orthogonal P .

Lemma 4.2. Let ˆ K : S, : d ∈ R S × d ( d < D ) be the reduced dimension key vectors obtained by projecting the key vectors onto the first d principal components of P . Then, the attention scores computed using ˆ K : S, : d are a good approximation of the the actual attention scores.

Proof. Let R : d ∈ R d × D be an orthogonal transformation that transforms the keys into the reduced dimension space as L : S, : d = K : S R : d . Our objective is to minimize the following expression:

Using Cauchy-Schwarz inequality, we have:

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, it follows that the optimal value of R ∗ : d = P : d , and L ∗ : S, : d = ˆ K : S, : d

Since we minimize an upper bound when proving Lemma 4.2, it is possible that some other transformation might give a 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 PCA-based Top-K Algorithm

Loki implements a PCA-based Top-K Attention approach. Previous works have shown that attention scores for a query are highly concentrated on a small subset of keys [36, 30]. This observation has motivated several methods that compute attention using only the topk most relevant keys. However, these previous works either compute the exact attention scores and then select the topk keys [12] or 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 lowerdimensional space, without any significant increase in memory requirements. Algorithm 1 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 topk keys based on the approximate attention scores. Line 8 computes the exact attention scores using the selected topk keys, directly in the transformed space (Lemma 4.1).

Compute and Memory Analysis: For vanilla attention, the complexity of computing q S K T : S 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 O ( dS ) . The complexity of selecting the topk keys (Lines 6-7) is approximately O ( Slog ( S ) + k ) (sorting followed by selection). The complexity of calculating the exact attention scores and multiplying with the values (Line 8-9)

Algorithm 1 Loki: PCA-based Top-K Attention

Require: At the S th step - Input: x S ∈ R 1 × D , KV-cache: ˆ K : S -1 , V : S -1 ∈ R ( S -1) × D , Projection Matrix: P ∈ R D × D , Configuration parameters (reduced dimensionality, topk ): d , k

1: function LOKI-ATTENTION( x S , ˆ K : S -1 , V : S -1 , P , d, k )

2: q S , k S , v S ← computeQKV ( x S )

3: ˆ q S ← q S P , ˆ k S ← k S P

4: ˆ K : S ← concat ( ˆ K : S -1 , ˆ k S ) , V : S ← concat ( V : S -1 , v S )

5: a approx ← ˆ q S, : d ( ˆ K : S, : d ) T

6: indices ← topk ( a approx , k )

7: ˆ K ′ : S ← ˆ K : S [ indices ] , V ′ : S ← V : S [ indices ]

9: return a exact V ′ : S

8: a exact ← softmax ( ˆ q S ˆ K T : S √ D )

10: end function

is O (2 Dk ) . Additionally, the complexity of projections into the PCA space (Line 3) is O (2 D 2 ) . Assuming the complexity of selecting the topk keys is small compared to the other operations, the overall complexity of the algorithm is O ( dS +2 Dk +2 D 2 ) . Then, we have:

where, d f = d/D and k f = k/S . The memory requirement of the KV-cache is the same as the original attention, with a 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 1). Standard PyTorch operations create temporary, dense copies of the KV-cache data in memory, leading to slowdowns due to expensive memory access. To alleviate this issue, we develop optimized kernels in Triton [1] for the three matrix multiplication operations in Loki. Our kernels can directly access relevant subsets of the 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 topk attention calculations. However, we identified and addressed inefficiencies in the SparQ kernels, which resulted in speedups of nearly 2 -3 × in certain scenarios. (see Appendix C).

5 Experimental Setup

We evaluate Loki on the basis of perplexity using the WikiText-2 [21] 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 metrics as the HuggingFace OpenLLM leaderboard [3]. For the LongBench tasks, we 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 topk tokens to compute the final output, and H 2 O [39], a popular token-eviction method. For these comparisons, we show the results with a budget size of k f = 0.25 and 0.125. For our method, we additionally use d f = 0.25 and 0.125. This configuration of our represents a 2.6x theoretical speedup. Table 1 provides an overview of the methods compared and the associated budget terms. H 2 O's budget was split equally between the heavy hitter and recent tokens, as per the author's recommendations. For H 2 O, we were unable to run the GSM8K task as the the author's ML benchmarking code was too memory intensive to run for that task. For the aforementioned experiments, we generate PCA transforms using the WikiText-103 dataset. For the LongBench tasks, we compare our method with the full attention baseline as we were unable to run H 2 O due to memory constraints.

For the generalizability study, we compare the results of our method with PCA transforms from different calibration datasets: WikiText-103 [21], C4 [25], and BookCorpus [40]. Additionally, we

also benchmark our triton based implementation of Loki by running an attention microbenchmark on a Llama2-13B-like 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 3, 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 f and d f values. For Llama2-7B, we see that the performance of both candidate transforms is similar. This trend is consistent across all the models except for Llama3-8B/70B 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, a significant decline not seen with the pre-rotary transform. Mistral-7B shows a similar pattern. This is a 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.

Llama2-7B

Llama2-13B

Llama3-8B

Llama3-70B

Mistral-7B

Mixtral-8x7B

0

2

4

6

8

10

12

14

Perplexity

5.11

4.57

5.57

2.57

4.91

3.60

5.22

4.61

5.66

2.59

4.92

3.66

5.20

4.61

11.45

2.86

6.54

3.79

5.50

4.82

5.74

2.63

4.93

3.64

5.36

4.70

6.80

2.62

9.52

3.64

Loki vs. Full Attention: Perplexity Evaluation

Full Attention

Loki (pre,

kf

=

0.25,

df

=

0.25)

Loki (post,

kf

=

0.25,

df

=

0.25)

Loki (pre,

kf

=

0.125,

df

=

0.5)

Loki (post,

kf

=

0.125,

df

=

0.5)

Llama2-7B

Llama2-13B

Llama3-8B

Llama3-70B

Mistral-7B

Mixtral-8x7B

0

20

40

60

80

100

Average Task Accuracy

47.67

52.11

60.26

71.82

58.30

65.81

43.50

49.67

45.58

60.32

52.63

57.17

42.28

49.76

40.52

48.98

34.66

47.37

39.76

46.22

43.48

59.00

47.96

54.49

38.46

46.71

37.97

47.72

34.26

50.76

Loki vs. Full Attention: Short-context Tasks Evaluation

Figure 4 shows the performance of Loki on the LongBench tasks for the Llama2-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 is at par or better than the full attention model. For the Single-Doc QA and Multi-Doc QA task categories, Loki performs slightly worse than the full attention model,

with the biggest drop in performance observed for HotpotQA of around 3%. Comparing different ( k f , d f ) settings, we see that using k f = 0.25 and d f = 0.25 (green), is better than using k f = 0.125 and d f = 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.

NarrativeQA

QASPER

MultiFieldQA

HotpotQA

2WikiMQA

Musique

GovReport

QMSum

MultiNews

TREC

TriviaQA

SamSum

PassageCount

PassageRetrieval

LCC

RepoBench

0

20

40

60

80

100

Task Accuracy

19.14

21.47

37.26

27.83

31.54

8.20

27.15

20.61

26.21

64.00

82.92

41.43

3.35

7.00

58.39

51.89

17.80

22.09

33.81

21.77

26.20

6.30

27.01

19.93

26.16

63.50

85.64

41.16

1.50

8.00

59.54

53.44

13.66

18.30

31.47

24.20

24.43

6.35

24.00

20.56

25.83

63.50

77.73

39.76

4.58

4.00

53.65

45.83

16.09

21.01

34.16

22.48

28.43

7.83

24.84

20.36

26.06

63.50

82.42

41.68

2.19

8.00

54.70

47.20

10.94

15.96

23.53

23.87

25.31

5.53

22.03

20.25

25.33

57.00

74.09

39.12

2.64

3.00

51.39

45.26

Single-Doc QA

Multi-Doc QA

Summarization

Few Shot Learning

Synthetic

Code Completion

Loki vs. Full Attention: LongBench English Tasks Evaluation

Full Attention

Loki (post,

kf

=

0.25,

df

=

0.25)

Loki (post,

kf

=

0.125,

df

=

0.5)

Loki (pre,

kf

=

0.25,

df

=

0.25)

Loki (pre,

kf

=

0.125,

df

=

0.5)

6.2 Comparison with Other Sparse Attention Methods

Next, we compare the performance of Loki with other methods, using k f = 0.25 for all methods and d f = 0.25 for ours. Table 2 shows the perplexity results for Llama2-7B/13B, Llama3-8B, and Mistral-7B. Loki's perplexity drop is within 0.1 of full attention across all models, a threshold considered acceptable for attention mechanism approximations [37]. In contrast, H 2 O's perplexity drop nears 0.2 for all models. Figure 5 confirms this trend on short-context evaluation. Loki performs similar to full attention for all models, except Llama3-8B, where the performance is notably worse, but still better than H 2 O. Importantly, on the challenging MMLU task, Loki degrades less than H 2 O.

TQA

Winogrande

ARC

MMLU

Avg (w/o GSM8K)

0

25

50

75

100

75.99

38.96

69.06

46.33

41.84

54.44

75.90

38.67

68.59

43.94

40.18

53.46

73.43

39.16

63.61

41.21

37.74

50.61

66.72

29.72

51.85

30.12

27.62

41.21

Llama2-7B

TQA

Winogrande

ARC

MMLU

Avg (w/o GSM8K)

0

25

50

75

100

79.38

36.90

72.22

49.15

52.06

57.94

79.24

37.22

73.01

49.66

52.26

58.28

77.83

38.48

68.90

46.16

50.32

55.72

72.27

28.97

52.25

33.53

28.64

43.13

Llama2-13B

TQA

Winogrande

ARC

MMLU

Avg (w/o GSM8K)

0

25

50

75

100

79.17

43.89

72.93

53.24

62.19

62.28

78.57

44.18

72.85

51.96

61.39

61.79

69.42

42.13

50.36

34.64

44.50

48.08

70.79

30.84

50.12

32.85

30.39

43.00

Llama3-8B

TQA

Winogrande

ARC

MMLU

Avg (w/o GSM8K)

0

25

50

75

100

81.07

42.62

73.95

53.92

59.65

62.24

80.30

43.27

73.01

53.07

59.53

61.84

77.65

42.18

62.98

46.59

53.70

56.62

70.53

30.40

49.49

33.02

29.16

42.52

Mistral-7B

Task Accuracy

Full Attention

Exact-TopK (

kf

=

0.25)

Loki (

kf

=

0.25,

df

=

0.25)

H

2

O (

kf

=

0.25)

It is important to note here that Loki is designed to be compatible with other sparse attention methods. For instance, token-eviction methods like H 2 O delete tokens to save KV-cache memory, whereas Loki reduces memory bandwidth by selecting the topk tokens without deletion, making them orthogonal. A combined approach could involve using H 2 O to delete tokens, then applying Loki to select topk tokens from the remaining cache. Similarly, Loki is theoretically orthogonal to quantization methods.

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 topk tokens. To understand why Loki works well, we examined the topk agreement between Loki's reduced dimensional attention scores and exact attention scores. Figure 6 shows the Jaccard similarity between the topk tokens selected by both methods across all layers and heads for Llama2-7B. For the settings: ( k f = 0 . 25 , d f = 0 . 25 ) and ( k f = 0 . 125 , d f = 0 . 5 ), evaluated in Figure 3, the Jaccard similarity is around 0.9, validating that the Loki is able to select the topk tokens with high accuracy.

0.125

0.25

0.5

df

0.0

0.2

0.4

0.6

0.8

1.0

Top-k Agreement b/w Exact

kf

0.125

0.25

0.5

Llama2-7B

Llama2-13B

Llama3-8B

Mistral-7B

0

1

2

3

4

5

6

7

Perplexity

5.20

4.61

5.68

4.92

5.20

4.62

5.66

4.92

5.20

4.61

5.66

4.92

Llama2-7B

Llama2-13B

Llama3-8B

Mistral-7B

0

20

40

60

80

100

Average Task Accuracy

40.85

49.09

45.82

52.09

42.34

50.08

46.33

52.93

42.28

49.76

45.58

52.63

BookCorpus

C4

Wikitext

Loki Generalization over Calibration Datasets

Loki

Vanilla

0.0

0.5

1.0

1.5

2.0

2.5

3.0

3.5

Time per Layer (s)

2.29

2.08

Llama2-13b - Attention w/

KV-cache Updates

KV-cache append

Actual compute

6.3 Generalizability

We now turn our 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 ( k f = 0 . 25 , d f = 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 it 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 6, rightmost). Regardless of the attention type (vanilla or Loki), more than 80% of the time is consumed within the Hugging Face framework for appending key-value pairs of the latest token to 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 [18] could significantly reduce this append time, but leave that exploration for future work. To isolate the impact of our optimizations, the plots in Figure 7 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 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 topk operation is nearly as expensive as the smaller matrix multiplications, which is a significant bottleneck. Replacing PyTorch's topk with a 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 k f , d f settings of Loki, with accuracy measured on LongBench and attention times from our microbenchmark. The previously evaluated settings, k f = 0 . 25 , d f = 0 . 25 and k f = 0 . 125 , d f = 0 . 5 , provide a good balance between performance and accuracy, with k f = 0 . 25 , d f = 0 . 25 favoring accuracy slightly and k f = 0 . 125 , d f = 0 . 5 favoring performance.

128

256

512

0.0

0.1

0.2

0.3

0.4

0.5

0.6

0.06

L-A

0.05

L-B

0.09

V

0.11

L-A

0.11

L-B

0.18

V

0.23

L-A

0.22

L-B

0.37

V

Prompt Length = 2048

q

.

k T

(top-

df

)

q

.

k T

(top-

kf

)

s

.

v

(top-

kf

)

top-k

other

128

256

512

0.0

0.1

0.2

0.3

0.4

0.5

0.6

0.08

L-A

0.07

L-B

0.12

V

0.15

L-A

0.14

L-B

0.25

V

0.31

L-A

0.28

L-B

0.52

V

Prompt Length = 3072

Llama2-13B: Attention Breakdown w/o KV-cache Updates

Time per Layer (s)

Number of Generated Tokens

Full Attention

(

kf

=

0.5,

df

=

0.5)

(

kf

=

0.5,

df

=

0.25)

(

kf

=

0.5,

df

=

0.125)

(

kf

=

0.25,

df

=

0.5)

(

kf

=

0.25,

df

=

0.25)

(

kf

=

0.25,

df

=

0.125)

(

kf

=

0.125,

df

=

0.5)

(

kf

=

0.125,

df

=

0.25)

(

kf

=

0.125,

df

=

0.125)

Configuration

0

20

40

60

80

100

Average Accuracy

32.62

32.55

32.45

31.68

30.94

29.18

27.58

23.75

9.73

32.65

32.65

32.11

32.62

31.95

30.69

29.55

25.25

23.21

32.71

Llama2-7B-Chat : LongBench Accuracy vs Attention Times

0

3

6

9

12

15

Time (s)

Full Attention

Loki (pre)

Loki (post)

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 in a low-dimensional space, across different models and datasets. Leveraging this insight, Loki uses attention scores computed in a lower-dimensional space to rank and select the topk most relevant tokens from the KV-cache. It then uses the full 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, we 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.

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-performance-accuracy trade-offs. Another direction involves storing the KV-cache in CPU memory and transferring only the topk keys and values to the GPU [19].

While Loki outperforms vanilla attention in our benchmarks, practical deployment would require integration with efficient attention kernels like FlashAttention [7]. As seen in our compute benchmarking, the topk 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 a variable d f 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 DDR-ERCAP0029894. Soheil Feizi was supported in part by the following grants: NSF CAREER AWARD 1942230, ONR YIP award N00014-22-1-2271, ARO's Early Career Program Award 310902-00001, Army Grant No. W911NF2120076, NSF award CCF2212458, NSF award 2229885, an Amazon Research Award and an award from Capital One.

References

Introducing triton: Open-source gpu programming for neural networks. https://openai. com/index/triton/ , 2021.

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.14508 , 2023.

Edward Beeching, Clémentine Fourrier, Nathan Habib, Sheon Han, Nathan Lambert, Nazneen Rajani, Omar Sanseviero, Lewis Tunstall, and Thomas Wolf. Open llm leaderboard. https://huggingface.co/spaces/open-llm-leaderboard-old/open_llm_ leaderboard , 2023.

Stella Biderman, Hailey Schoelkopf, Quentin Gregory Anthony, Herbie Bradley, Kyle 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.

Rewon Child, Scott Gray, Alec Radford, and Ilya Sutskever. Generating long sequences with sparse transformers. arXiv preprint arXiv:1904.10509 , 2019.

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.

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-16359, 2022.

Abhimanyu Dubey, Abhinav Jauhri, Abhinav Pandey, Abhishek Kadian, Ahmad Al-Dahle, et al. The llama 3 herd of models, 2024.

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.

Suyu Ge, Yunan Zhang, Liyuan Liu, Minjia Zhang, Jiawei Han, and Jianfeng Gao. Model tells you what to discard: Adaptive kv cache compression for llms. arXiv preprint arXiv:2310.01801 , 2023.

Suyu Ge, Yunan Zhang, Liyuan Liu, Minjia Zhang, Jiawei Han, and Jianfeng Gao. Model tells you what to discard: Adaptive kv cache compression for llms, 2024.

Ankit Gupta, Guy Dar, Shaya Goodman, David Ciprut, and Jonathan Berant. Memory-efficient transformers via top-k attention. CoRR , abs/2106.06899, 2021.

Edward J. Hu, Yelong Shen, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, Shean Wang, and Weizhu Chen. Lora: Low-rank adaptation of large language models. CoRR , abs/2106.09685, 2021.

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-arithmetic-only inference, 2017.

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.

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 El Sayed. Mixtral of experts, 2024.

Nikita Kitaev, Łukasz Kaiser, and Anselm Levskaya. Reformer: The efficient transformer. arXiv preprint arXiv:2001.04451 , 2020.

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 language model serving with pagedattention. In Proceedings of the ACM SIGOPS 29th Symposium on Operating Systems Principles , 2023.

Wonbeom Lee, Jungi Lee, Junghwan Seo, and Jaewoong Sim. InfiniGen: Efficient generative inference of large language models with dynamic KV cache management. In 18th USENIX Symposium on Operating Systems Design and Implementation (OSDI 24) , pages 155-172, Santa Clara, CA, July 2024. USENIX Association.

Zichang Liu, Aditya Desai, Fangshuo Liao, Weitao Wang, Victor Xie, Zhaozhuo Xu, Anastasios Kyrillidis, and Anshumali Shrivastava. Scissorhands: Exploiting the persistence of importance hypothesis for llm kv cache compression at test time. arXiv preprint arXiv:2305.17118 , 2023.

Stephen Merity, Caiming Xiong, James Bradbury, and Richard Socher. Pointer sentinel mixture models. CoRR , abs/1609.07843, 2016.

Microsoft. Introducing phi-3: Redefining what's possible with slms. https://azure.microsoft.com/en-us/blog/ introducing-phi-3-redefining-whats-possible-with-slms/ , 2024.

Markus Nagel, Marios Fournarakis, Rana Ali Amjad, Yelysei Bondarenko, Mart van Baalen, and Tijmen Blankevoort. A white paper on neural network quantization, 2021.

NERSC. Perlmutter system architecture. https://docs.nersc.gov/systems/ perlmutter/architecture/ .

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.

Luka Ribar, Ivan Chelombiev, Luke Hudlass-Galley, Charlie Blake, Carlo Luschi, and Douglas Orr. Sparq attention: Bandwidth-efficient llm inference, 2023.

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 '22. IEEE Computer Society, May 2022.

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.

Jianlin Su, Yu Lu, Shengfeng Pan, Ahmed Murtadha, Bo Wen, and Yunfeng Liu. Roformer: Enhanced transformer with rotary position embedding, 2023.

Mingjie Sun, Xinlei Chen, J. Zico Kolter, and Zhuang Liu. Massive activations in large language models, 2024.

Hugo Touvron et al. Llama 2: Open foundation and fine-tuned chat models. Technical report, 2023.

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.03762, 2017.

Sinong Wang, Belinda Z. Li, Madian Khabsa, Han Fang, and Hao Ma. Linformer: Self-attention with linear complexity, 2020.

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-theart natural language processing. In Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing: System Demonstrations , pages 38-45, Online, October 2020. Association for Computational Linguistics.

Guangxuan Xiao, Yuandong Tian, Beidi Chen, Song Han, and Mike Lewis. Efficient streaming language models with attention sinks. arXiv preprint arXiv:2309.17453 , 2023.

Guangxuan Xiao, Yuandong Tian, Beidi Chen, Song Han, and Mike Lewis. Efficient streaming language models with attention sinks, 2024.

Zhewei Yao, Xiaoxia Wu, Cheng Li, Stephen Youn, and Yuxiong He. Zeroquant-v2: Exploring post-training quantization in llms from comprehensive study to low rank compensation. 2023.

Peiyuan Zhang, Guangtao Zeng, Tianduo Wang, and Wei Lu. Tinyllama: An open-source small language model. Technical report, 2024.

Zhenyu Zhang, Ying Sheng, Tianyi Zhou, Tianlong Chen, Lianmin Zheng, Ruisi Cai, Zhao Song, Yuandong Tian, Christopher Ré, Clark Barrett, et al. H _ 2 o: Heavy-hitter oracle for efficient generative inference of large language models. arXiv preprint arXiv:2306.14048 , 2023.

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.

A Comprehensive Dimensionality Analysis

A.1 Complete Key-Dimensionaltiy Analysis for All Models

0

5

10

15

20

25

30

Layer ID

0

20

40

60

80

100

120

128

140

Rank at 90% Explained  Variance

Llama2-7B Attention Keys (Pre-Rotary)

Wikitext

C4

BookCorpus

0

5

10

15

20

25

30

Layer ID

0

20

40

60

80

100

120

128

140

Rank at 90% Explained  Variance

Llama2-7B Attention Keys (Post-Rotary)

Wikitext

C4

BookCorpus

0

10

20

30

40

Layer ID

0

20

40

60

80

100

120

128

140

Rank at 90% Explained  Variance

Llama2-13B Attention Keys (Pre-Rotary)

Wikitext

C4

BookCorpus

0

10

20

30

40

Layer ID

0

20

40

60

80

100

120

128

140

Rank at 90% Explained  Variance

Llama2-13B Attention Keys (Post-Rotary)

Wikitext

C4

BookCorpus

0

20

40

60

80

Layer ID

0

20

40

60

80

100

120

128

140

Rank at 90% Explained  Variance

Llama2-70B Attention Keys (Pre-Rotary)

Wikitext

C4

BookCorpus

0

20

40

60

80

Layer ID

0

20

40

60

80

100

120

128

140

Rank at 90% Explained  Variance

Llama2-70B Attention Keys (Post-Rotary)

Wikitext

C4

BookCorpus

0

5

10

15

20

25

30

Layer ID

0

20

40

60

80

100

120

128

140

Rank at 90% Explained  Variance

Llama3-8B Attention Keys (Pre-Rotary)

Wikitext

C4

BookCorpus

0

5

10

15

20

25

30

Layer ID

0

20

40

60

80

100

120

128

140

Rank at 90% Explained  Variance

Llama3-8B Attention Keys (Post-Rotary)

Wikitext

C4

BookCorpus

0

20

40

60

80

Layer ID

0

20

40

60

80

100

120

128

140

Rank at 90% Explained  Variance

Llama3-70B Attention Keys (Pre-Rotary)

Wikitext

C4

BookCorpus

0

20

40

60

80

Layer ID

0

20

40

60

80

100

120

128

140

Rank at 90% Explained  Variance

Llama3-70B Attention Keys (Post-Rotary)

Wikitext

C4

BookCorpus

0

5

10

15

20

Layer ID

0

20

40

60

80

100

120

128

140

Rank at 90% Explained  Variance

TinyLlama-1.1B Attention Keys (Pre-Rotary)

Wikitext

C4

BookCorpus

0

5

10

15

20

Layer ID

0

20

40

60

80

100

120

128

140

Rank at 90% Explained  Variance

Wikitext

C4

BookCorpus

0

5

10

15

20

25

30

Layer ID

0

20

40

60

80

100

120

128

140

Rank at 90% Explained  Variance

Wikitext

C4

BookCorpus

0

5

10

15

20

25

30

Layer ID

0

20

40

60

80

100

120

128

140

Rank at 90% Explained  Variance

Wikitext

C4

BookCorpus

0

5

10

15

20

25

30

Layer ID

0

20

40

60

80

100

120

128

140

Rank at 90% Explained  Variance

Mixtral-8x7B Attention Keys (Pre-Rotary)

Wikitext

C4

BookCorpus

0

5

10

15

20

25

30

Layer ID

0

20

40

60

80

100

120

128

140

Rank at 90% Explained  Variance

Mixtral-8x7B Attention Keys (Post-Rotary)

Wikitext

C4

BookCorpus

0

10

20

30

40

50

Layer ID

0

20

40

60

80

100

120

128

140

Rank at 90% Explained  Variance

Mixtral-8x22B Attention Keys (Pre-Rotary)

Wikitext

BookCorpus

0

10

20

30

40

50

Layer ID

0

20

40

60

80

100

120

128

140

Rank at 90% Explained  Variance

Mixtral-8x22B Attention Keys (Post-Rotary)

Wikitext

BookCorpus

0

5

10

15

20

25

30

Layer ID

0

20

40

60

80

100

120

128

140

Rank at 90% Explained  Variance

Wikitext

C4

BookCorpus

0

5

10

15

20

25

30

Layer ID

0

20

40

60

80

100

120

128

140

Rank at 90% Explained  Variance

Wikitext

C4

BookCorpus

0

5

10

15

20

25

30

Layer ID

0

20

40

60

80

96

100

120

140

Rank at 90% Explained  Variance

Phi3-Mini-4K Attention Keys (Pre-Rotary)

Wikitext

C4

BookCorpus

0

5

10

15

20

25

30

Layer ID

0

20

40

60

80

96

100

120

140

Rank at 90% Explained  Variance

Phi3-Mini-4K Attention Keys (Post-Rotary)

Wikitext

C4

BookCorpus

In this section, we present our extended dimensionality analysis results (from 3) for all the models we experimented with. Figure 8 displays the Rank l @90 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 a wide range of sizes, architecture types (dense vs. MoE), as well as older and newer architectures trained on various datasets. Despite these differences, our main observation remains robust.

An intriguing trend is the variation in Rank l @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 a more fine-grained analysis, we plot the normalized eigenvalues of the covariance matrix of the keys for a few layers and heads of Llama2-7B, Mistral-7B, and Pythia-6.9B on the WikiText-2 dataset as an example in Figure 9. Here again, we observe that the explained variance significantly

decreases after the initial principal dimensions. The results for the other models are similar to the ones shown here.

0

20

40

60

80

100

120

Feature Dimension / Rank

0.00

0.05

0.10

0.15

0.20

0.25

0.30

Normalised Eigen-values

Llama2-7B, Layer 1, Head 0

Pre-Rotary

Post-Rotary

0

20

40

60

80

100

120

Feature Dimension / Rank

0.00

0.05

0.10

0.15

0.20

0.25

0.30

Normalised Eigen-values

Mistral-7B, Layer 1, Head 0

Pre-Rotary

Post-Rotary

0

20

40

60

80

100

120

Feature Dimension / Rank

0.00

0.05

0.10

0.15

0.20

0.25

0.30

Normalised Eigen-values

Pre-Rotary

Post-Rotary

0

20

40

60

80

100

120

Feature Dimension / Rank

0.00

0.05

0.10

0.15

0.20

0.25

0.30

Normalised Eigen-values

Llama2-7B, Layer 28, Head 6

Pre-Rotary

Post-Rotary

0

20

40

60

80

100

120

Feature Dimension / Rank

0.00

0.05

0.10

0.15

0.20

0.25

0.30

Normalised Eigen-values

Mistral-7B, Layer 28, Head 6

Pre-Rotary

Post-Rotary

0

20

40

60

80

100

120

Feature Dimension / Rank

0.00

0.05

0.10

0.15

0.20

0.25

0.30

Normalised Eigen-values

Pre-Rotary

Post-Rotary

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 l @90 ) across different heads within a layer for two models: Llama2-7B and Mistral-7B. Figure 10 shows the heatmap of the Rank l @90 for the pre-rotary (top) and post-rotary (bottom) keys across all layers and heads for Mistral-7B. We observe that the Rank l @90 is considerably lower for prerotary 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 to Mistral-7B as we do not see this pattern in Llama2-7B.

Figure 11 shows the heatmap of the Rank l @90 for the pre-rotary (left) and post-rotary (right) keys across all layers and heads for Llama2-7B. We observe a 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-7B compared to Mistral-7B.

In this section, we examine the variation in the rank at which 90% of the variance is explained ( Rank l @90 ) across different heads within a layer for two models: Llama2-7B and Mistral-7B. Figure 10 shows the heatmap of Rank l @90 for the pre-rotary (top) and post-rotary (bottom) keys across all layers and heads for Mistral-7B. We observe that the Rank l @90 is significantly lower for pre-rotary keys compared to post-rotary keys. Focusing on 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, we notice a 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 is needed to understand this trend.

Figure 11 illustrates the heatmap of Rank l @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

0

5

10

15

20

25

30

Layers

0

1

2

3

4

5

6

7

Heads

3

13

50

31

41

57

47

59

59

43

46

38

61

63

50

43

61

44

33

74

54

38

44

71

76

83

84

83

63

47

37

67

1

12

84

53

54

50

58

43

62

46

63

56

28

47

41

43

48

61

51

43

69

54

48

65

58

83

43

77

61

96

47

43

2

30

70

39

41

47

61

53

53

50

50

61

37

58

53

39

59

67

65

32

45

50

51

59

77

53

62

96

68

48

43

68

1

11

40

45

44

59

57

43

37

51

71

52

60

60

57

59

42

69

43

47

55

75

87

35

52

73

60

68

50

50

49

63

2

16

55

30

62

73

47

30

46

53

53

56

47

61

49

44

54

66

60

37

46

52

66

70

73

53

47

67

46

60

55

35

1

15

30

40

43

56

40

49

52

58

62

42

45

57

40

42

45

78

44

77

46

64

75

79

56

78

74

82

62

55

78

68

3

25

70

55

34

57

42

54

55

42

61

57

45

53

45

41

46

41

64

79

57

68

64

75

68

53

62

60

43

73

54

57

5

10

58

49

60

41

66

52

62

48

56

45

67

68

49

43

48

60

51

48

48

81

61

78

88

55

61

55

70

69

60

50

20

40

60

80

Rank at 95%

0

5

10

15

20

25

30

Layers

0

1

2

3

4

5

6

7

Heads

51

47

97

81

91

86

87

88

96

93

87

88

88

89

86

80

84

74

51

92

92

63

70

84

93

91

88

90

88

90

63

90

51

47

97

81

91

86

87

88

96

93

87

88

88

89

86

80

84

74

51

92

92

63

70

84

93

91

88

90

88

90

63

90

51

47

97

81

91

86

87

88

96

93

87

88

88

89

86

80

84

74

51

92

92

63

70

84

93

91

88

90

88

90

63

90

51

47

97

81

91

86

87

88

96

93

87

88

88

89

86

80

84

74

51

92

92

63

70

84

93

91

88

90

88

90

63

90

37

60

79

79

80

95

85

93

95

86

91

84

43

88

88

78

79

90

87

82

88

88

75

90

89

90

70

90

92

70

82

79

37

60

79

79

80

95

85

93

95

86

91

84

43

88

88

78

79

90

87

82

88

88

75

90

89

90

70

90

92

70

82

79

37

60

79

79

80

95

85

93

95

86

91

84

43

88

88

78

79

90

87

82

88

88

75

90

89

90

70

90

92

70

82

79

37

60

79

79

80

95

85

93

95

86

91

84

43

88

88

78

79

90

87

82

88

88

75

90

89

90

70

90

92

70

82

79

40

50

60

70

80

90

Rank at 95%

0

5

10

15

20

25

30

0

5

10

15

20

25

30

Heads

3

26

48

37

58

64

66

56

64

69

51

61

70

76

89

58

70

46

84

74

44

71

95

65

96

42

91

97

82

71

83

71

2

37

32

33

53

60

47

79

61

65

78

50

63

70

55

77

49

67

52

79

42

41

86

76

68

88

88

86

87

86

82

72

2

12

45

67

47

39

61

59

69

64

51

36

38

54

58

71

57

57

76

76

82

89

93

92

87

94

81

73

101

61

78

68

1

7

16

61

55

48

69

50

52

69

72

61

66

79

45

70

61

87

68

69

55

80

92

83

40

54

44

87

79

93

99

68

3

5

27

70

68

52

61

43

51

80

56

63

41

56

62

61

48

71

78

49

85

42

73

83

40

88

80

78

87

76

78

68

2

6

43

68

28

51

68

58

47

68

69

40

60

77

58

49

49

87

85

79

59

41

99

93

67

77

96

85

83

64

75

83

2

18

17

67

59

48

42

31

62

67

63

64

61

80

63

68

44

87

81

56

94

89

83

96

76

91

90

82

74

87

90

72

1

3

42

36

74

60

56

83

50

69

58

62

76

55

49

74

74

89

83

52

62

85

78

43

85

68

94

72

50

86

75

67

2

2

46

46

47

64

41

70

64

52

66

63

70

65

67

55

66

85

74

52

59

86

59

40

40

83

96

99

88

80

82

67

3

3

46

52

59

62

43

65

53

64

66

52

54

46

52

74

50

54

53

56

93

58

46

66

86

87

73

84

54

47

74

63

2

17

53

63

64

59

66

27

77

55

58

71

59

66

71

58

72

56

46

42

48

66

73

79

78

94

77

84

85

63

74

79

2

4

38

60

55

53

62

64

79

71

50

54

55

51

54

71

71

49

58

79

62

85

77

93

46

62

99

99

97

83

88

72

5

7

61

62

49

43

45

37

64

59

56

68

60

68

47

42

67

59

63

61

62

89

97

94

81

81

77

79

98

86

68

61

1

2

50

75

38

57

49

56

69

81

66

54

57

68

45

70

50

46

96

85

75

62

99

97

87

94

80

96

52

45

81

75

1

9

32

65

39

57

66

54

53

67

46

37

61

82

73

41

84

90

95

50

54

84

95

96

66

84

64

79

38

82

55

69

1

30

26

66

68

45

61

75

68

63

44

33

53

72

54

53

73

66

54

39

78

68

82

82

45

99

62

71

85

63

97

45

3

27

62

63

58

64

25

64

76

80

61

72

40

60

53

73

70

48

53

81

95

45

53

73

39

61

78

73

83

71

84

47

2

10

39

42

46

59

54

68

67

49

63

41

57

56

53

61

47

70

87

52

63

100

64

68

39

46

89

88

96

97

85

79

4

39

28

56

58

55

53

66

55

60

39

41

55

56

41

61

57

47

52

62

66

78

94

83

85

99

90

93

96

73

88

90

1

30

36

68

60

59

70

74

65

64

54

57

72

52

52

55

36

80

78

48

88

94

40

79

95

97

93

89

95

43

84

69

2

3

46

71

40

47

45

58

69

71

68

58

75

63

61

63

68

75

84

91

71

94

73

42

83

90

87

88

74

88

82

79

1

6

20

51

32

59

52

48

81

65

58

69

57

62

63

92

75

90

79

51

77

92

81

95

85

42

45

100

77

52

90

68

1

51

25

61

61

48

65

59

44

65

70

44

73

57

70

56

83

41

90

94

53

77

43

98

94

57

80

100

94

89

79

14

2

14

27

72

55

42

71

56

63

45

62

47

75

54

84

60

67

89

63

77

88

84

53

91

69

81

85

99

78

84

84

73

1

39

41

73

75

63

73

71

46

66

67

52

82

82

39

54

59

71

93

82

92

77

41

94

47

96

41

82

86

92

95

58

2

18

23

58

56

59

55

63

34

51

74

73

56

61

71

63

56

51

71

65

64

77

92

97

99

86

44

94

70

83

79

47

3

8

7

55

60

66

71

67

34

72

56

52

56

59

46

92

71

49

83

70

93

54

76

98

86

76

42

97

89

48

63

67

2

38

32

42

50

50

81

57

65

70

58

74

70

81

54

53

51

55

80

68

43

41

41

94

95

91

43

81

80

77

86

71

3

32

40

41

66

66

65

49

48

65

65

52

48

68

67

76

57

74

84

86

49

49

90

90

93

71

57

91

90

75

91

68

3

5

67

54

60

58

63

63

60

51

49

58

70

83

45

56

58

92

90

89

42

59

80

87

41

100

76

71

96

79

93

73

1

6

53

55

58

57

37

74

70

51

61

54

58

57

69

59

46

50

42

85

45

44

49

80

50

98

67

101 100

69

95

73

2

6

34

34

60

41

60

47

44

62

57

59

67

53

81

58

54

47

51

65

63

53

65

40

98

89

92

87

83

88

79

73

20

40

60

80

100

Rank at 90%

0

5

10

15

20

25

30

0

5

10

15

20

25

30

Heads

60

54

54

73

67

44

86

74

77

72

80

66

77

78

86

75

81

76

83

69

62

78

96

75

101

71

97

98

89

76

85

64

63

70

50

20

71

79

67

82

73

75

89

74

78

74

74

88

70

87

70

81

63

66

88

82

80

96

96

92

87

89

83

84

52

29

56

69

65

73

80

81

68

70

78

56

55

77

79

75

77

77

81

74

83

88

98

95

93

99

85

77

100

69

76

77

35

25

66

52

72

79

62

74

76

77

62

77

73

93

73

74

74

88

83

83

74

82

97

83

63

69

68

90

83

94

100

60

64

49

60

68

69

46

73

70

79

83

77

85

65

75

79

75

75

73

78

73

79

71

70

80

70

92

76

78

95

81

87

61

41

38

65

66

65

60

71

79

68

68

73

68

76

88

77

72

74

93

82

83

76

64

100

89

79

78

96

85

91

85

81

95

60

60

52

77

63

63

70

51

90

78

78

82

75

87

75

75

67

89

66

72

101

99

85

103

74

95

87

78

81

88

87

85

54

55

59

71

77

74

84

86

77

76

80

77

89

70

74

75

76

84

75

76

74

85

75

67

84

92

96

86

77

98

75

90

53

49

62

67

67

71

56

80

80

78

73

79

78

74

73

73

75

93

84

68

78

85

73

64

65

90

100 103

90

84

90

79

63

61

53

64

63

65

60

74

71

75

71

71

74

62

74

73

72

76

73

76

98

71

64

77

88

80

78

88

62

73

86

85

47

59

57

65

65

66

73

39

83

76

78

87

80

80

79

81

77

77

73

63

77

78

81

77

80

94

77

85

77

75

71

83

58

52

57

57

66

58

73

66

68

81

75

74

79

72

74

67

69

68

72

79

79

94

71

88

69

77

104 103

97

88

101

89

60

14

56

72

79

60

82

57

67

73

81

82

75

74

71

71

63

76

77

78

77

87

102

90

84

90

84

63

102

83

82

87

62

55

62

58

64

82

68

77

77

85

79

54

75

75

65

76

74

70

96

86

80

76

101 100

99

98

78

97

79

66

71

67

57

48

61

68

58

69

76

81

75

75

66

66

77

84

76

65

87

90

96

73

78

82

94

100

78

95

79

81

71

85

77

76

35

75

46

61

64

76

82

86

73

68

69

48

77

77

76

73

70

66

74

64

78

82

86

86

66

99

80

63

84

83

95

74

62

69

59

64

74

67

42

77

62

72

76

76

69

74

76

61

73

72

75

70

101

67

75

77

55

76

89

65

87

79

89

68

61

30

57

57

61

74

82

74

78

70

61

72

71

73

74

77

71

70

90

69

75

101

71

77

61

67

88

84

90

100

79

86

58

92

54

67

66

67

74

60

77

78

70

60

77

72

58

74

72

69

73

79

76

84

98

69

71

102

85

88

98

80

96

88

44

69

62

67

71

65

62

74

75

69

78

74

85

73

74

76

54

74

81

72

90

100

64

82

94

98

98

94

99

65

91

88

62

46

66

54

61

67

75

74

67

81

78

73

68

74

75

78

83

90

82

96

75

99

79

67

88

96

84

94

80

89

79

84

55

20

40

68

50

83

76

71

91

74

78

85

75

77

80

93

90

88

69

69

80

97

72

96

90

61

75

98

84

73

86

82

54

47

44

55

57

63

62

71

66

68

74

62

87

70

78

77

85

60

95

98

72

68

65

101

99

74

81

104

98

92

84

34

67

47

43

64

60

71

71

74

81

78

76

71

80

73

86

75

74

92

79

89

93

86

77

94

70

82

93

101

76

91

91

87

53

72

60

72

56

66

71

79

71

72

84

74

86

89

66

75

69

75

98

83

100

79

63

95

73

95

73

74

94

89

94

76

33

54

30

70

67

72

71

59

57

75

62

78

82

68

89

76

80

76

75

80

77

82

93

97

104

86

68

89

74

88

80

74

65

59

27

89

51

74

77

78

50

78

79

80

78

71

73

91

78

66

83

74

100

70

64

94

85

81

64

100

77

67

73

70

59

55

48

53

59

77

70

67

76

62

73

69

69

75

71

74

72

79

72

76

67

67

65

96

91

93

69

80

83

85

95

64

55

62

62

78

59

78

78

73

63

68

75

75

76

73

71

76

74

77

72

90

68

72

92

89

91

76

73

96

91

83

94

85

55

54

30

62

60

73

60

87

67

69

78

75

73

74

72

79

74

97

88

98

66

73

87

95

62

105

82

81

95

80

98

91

52

34

52

65

69

62

63

76

69

77

79

78

73

77

70

73

65

78

66

89

71

67

76

84

72

101

84

99

103

66

98

85

62

36

66

66

67

65

73

80

70

78

72

81

67

77

87

73

75

69

74

78

72

74

78

61

102

96

94

89

98

89

91

74

20

40

60

80

100

Rank at 90%

Layers

Layers

the initial layers have lower rank values compared to the later layers. However, the same pattern in post-rotary keys that was observed in Mistral-7B is absent here, suggesting that rotary embeddings may be applied differently in 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 is low. This might indicate that some heads in 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 Rank l @90 for the queries and values, respectively, for Llama2-7B and Llama3-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 to the full dimensionality of the value vectors. This observation can intuitively be explained by the fact that both keys and queries are used

0

5

10

15

20

25

30

Layer ID

0

20

40

60

80

100

120

128

140

Rank at 90% Explained  Variance

Llama2-7B Attention Queries (Pre-Rotary)

Wikitext

C4

BookCorpus

0

5

10

15

20

25

30

Layer ID

0

20

40

60

80

100

120

128

140

Rank at 90% Explained  Variance

Llama2-7B Attention Queries (Post-Rotary)

Wikitext

C4

BookCorpus

0

20

40

60

80

Layer ID

0

20

40

60

80

100

120

128

140

Rank at 90% Explained  Variance

Llama3-70B Attention Queries (Pre-Rotary)

Wikitext

C4

BookCorpus

0

20

40

60

80

Layer ID

0

20

40

60

80

100

120

128

140

Rank at 90% Explained  Variance

Llama3-70B Attention Queries (Post-Rotary)

Wikitext

C4

BookCorpus

0

5

10

15

20

25

30

Layer ID

0

20

40

60

80

100

120

128

140

Rank at 90% Explained  Variance

Llama2-7B Attention  Values

Wikitext

C4

BookCorpus

0

20

40

60

80

Layer ID

0

20

40

60

80

100

120

128

140

Rank at 90% Explained  Variance

Llama3-70B Attention  Values

Wikitext

C4

BookCorpus

to compute the scalar attention scores and thus do not need to be high-dimensional, while values are weighted by these scores and used to compute the final output, and thus need to be high-dimensional to capture the complexity of the data.

B Comprehensive Evaluation Results

B.1 Performance of Loki on Perplexity and Short-Context Downstream Tasks

Llama2-7B

0.125

0.25

0.5

6

8

Perplexity

5.1101

8.0514

5.3044

5.1250

6.0729

5.2185

5.1223

5.4980

5.1881

5.1195

Pre-Rotary

0.125

0.25

0.5

6

7

5.1101

7.4062

5.4428

5.1241

5.5606

5.2017

5.1206

5.3601

5.1838

5.1195

Post-Rotary

Llama2-13B

0.125

0.25

0.5

5

6

7

4.568

7.6573

4.6829

4.5737

5.3912

4.6110

4.5708

4.8153

4.5979

4.5701

Pre-Rotary

0.125

0.25

0.5

5

6

4.568

6.1436

4.8082

4.5745

4.9668

4.6102

4.5737

4.7029

4.5937

4.5731

Post-Rotary

Llama2-70B

0.125

0.25

0.5

4

6

Perplexity

3.1205

6.4963

3.4571

3.1316

3.9259

3.2830

3.1293

3.8327

3.2986

3.1319

Pre-Rotary

0.125

0.25

0.5

3.5

4.0

4.5

3.1205

4.6898

3.5459

3.1457

4.7796

3.5701

3.1453

4.1427

3.4619

3.1411

Post-Rotary

Llama3-8B

0.125

0.25

0.5

7.5

10.0

5.5696

11.1097

6.0558

5.5876

6.5780

5.6648

5.5746

5.7356

5.5944

5.5703

Pre-Rotary

0.125

0.25

0.5

10

20

5.5696

22.6596

13.2883

5.7429

16.3507

11.4459

5.9343

6.8023

5.6783

5.5699

Post-Rotary

Llama3-70B

0.125

0.25

0.5

3

4

Perplexity

2.5653

4.1495

2.7577

2.5712

2.8796

2.5942

2.5665

2.6285

2.5727

2.5656

Pre-Rotary

0.125

0.25

0.5

0

25

50

2.5653

57.6788

5.6725

2.7810

4.2512

2.8593

2.5697

2.6231

2.5742

2.5660

Post-Rotary

TinyLlama-1.1B

0.125

0.25

0.5

10

15

20

7.9671

19.9290

9.4947

8.1057

11.5887

8.6352

8.0342

9.3280

8.3475

8.0040

Pre-Rotary

0.125

0.25

0.5

8

10

7.9671

11.3508

8.8956

8.0414

10.2592

8.5687

8.0135

8.9679

8.3190

7.9979

Post-Rotary

Mistral-7B

0.125

0.25

0.5

kf

5.0

5.5

Perplexity

4.914

5.7404

4.9986

4.9193

4.9636

4.9233

4.9152

4.9311

4.9185

4.9147

Pre-Rotary

0.125

0.25

0.5

kf

10

20

4.914

20.5289

7.7609

4.9317

13.5597

6.5445

4.9221

9.5167

5.2052

4.9149

Post-Rotary

Mixtral-8x7B

0.125

0.25

0.5

kf

4

6

3.5967

6.9799

3.9252

3.6201

3.8467

3.6584

3.6047

3.6417

3.6076

3.5979

Pre-Rotary

0.125

0.25

0.5

kf

4

5

3.5967

5.5250

4.2566

3.6635

4.5500

3.7906

3.6196

3.6358

3.6004

3.5970

Post-Rotary

Loki vs. Full Attention (Perplexity)

df

0.125

0.25

0.5

Llama2-7B

0.125

0.25

0.5

30

40

Average Task Accuracy

47.67

28.9800

37.0900

46.2100

34.9500

43.5000

47.1400

39.7600

46.0400

47.4500

Pre-Rotary

0.125

0.25

0.5

35

40

45

47.67

31.5300

39.9300

46.2100

35.3500

42.2800

47.1400

38.4600

44.6300

47.5800

Post-Rotary

Llama2-13B

0.125

0.25

0.5

30

40

50

52.11

30.0800

42.0000

50.7900

37.0300

49.6700

52.0700

46.2200

51.7400

52.2800

Pre-Rotary

0.125

0.25

0.5

40

50

52.11

32.0700

40.6800

49.3500

40.9200

49.7600

52.1700

46.7100

51.4600

52.2300

Post-Rotary

Llama2-70B

0.125

0.25

0.5

40

60

Average Task Accuracy

63.74

29.8700

34.0200

57.9100

33.2200

50.4100

60.4200

44.7200

55.8800

61.5600

Pre-Rotary

0.125

0.25

0.5

50

60

63.74

44.6700

57.6400

62.4900

46.1900

58.5600

62.9900

50.2000

60.0600

62.6900

Post-Rotary

Llama3-8B

0.125

0.25

0.5

40

60

60.26

28.6400

35.9000

54.3000

31.4300

45.5800

57.7100

43.4800

54.7000

60.1800

Pre-Rotary

0.125

0.25

0.5

40

60

60.26

29.1800

34.2900

51.7200

32.9900

40.5200

52.6300

37.9700

46.9700

56.7700

Post-Rotary

Llama3-70B

0.125

0.25

0.5

40

60

Average Task Accuracy

71.82

30.3600

47.1600

66.2300

44.3100

60.3200

69.5500

59.0000

68.5800

71.6000

Pre-Rotary

0.125

0.25

0.5

40

60

71.82

29.4100

32.0400

58.4600

33.1200

48.9800

64.5000

47.7200

59.0000

68.4600

Post-Rotary

TinyLlama-1.1B

0.125

0.25

0.5

30.0

32.5

35.0

36.36

29.9600

31.5100

35.4300

30.3100

33.8200

36.2300

33.3800

35.6200

36.1900

Pre-Rotary

0.125

0.25

0.5

30.0

32.5

35.0

36.36

29.8700

33.5200

34.6000

30.8200

32.5100

34.9200

31.7000

34.2600

35.8700

Post-Rotary

Mistral-7B

0.125

0.25

0.5

kf

30

40

50

Average Task Accuracy

58.3

30.1200

41.4300

55.7100

40.7600

52.6300

57.6000

47.9600

56.0400

58.2600

Pre-Rotary

0.125

0.25

0.5

kf

30

40

50

58.3

30.4100

30.1600

45.5700

29.6700

34.6600

52.8200

34.2600

46.2200

56.3100

Post-Rotary

Mixtral-8x7B

0.125

0.25

0.5

kf

40

60

65.81

31.7700

43.9800

60.8700

40.0900

57.1700

65.1000

54.4900

63.8900

65.2600

Pre-Rotary

0.125

0.25

0.5

kf

40

60

65.81

31.0400

37.1500

55.9200

34.1200

47.3700

61.3300

50.7600

59.6900

64.5100

Post-Rotary

Loki vs. Full Attention (Average Task Accuracy)

df

0.125

0.25

0.5

In this section, we provide a detailed evaluation of our method across a wide range of models and tasks. Figure 14 illustrates the performance of Loki on perplexity and downstream tasks compared to

the full attention baseline. We present results for both pre-rotary and post-rotary PCA transformations. The models evaluated include Llama2-7B, Llama2-13B, Llama2-70B, Llama3-8B, Llama3-70B, TinyLlama-1.1B, Mistral-7B, and Mixtral-8x7B. We assess the models using various configurations of k and d for Loki.

As k f and d f decrease, the model's performance deteriorates, particularly when both are set to 0.125. Notably, the impact of k f on performance is more pronounced than that of d f . This is evident as k f = 0 . 125 and d f = 0 . 5 significantly underperform compared to k f = 0 . 5 and d f = 0 . 125 across nearly all models. The configurations of k f = 0 . 25 and d f = 0 . 25 , along with k f = 0 . 125 and d f = 0 . 5 , demonstrate relatively strong performance across all models, striking a favorable balance between performance and accuracy, with a theoretical speedup of 2.6x for both configurations. Settings with k f = 0 . 5 maintain model quality much more effectively but do not yield a significant empirical speedup.

Tables 3 and 4 present finer-grained results for each task and model.

B.2 Variable Dimensionality Analysis

0.2

0.3

0.4

0.5

Compression Ratio

38

40

42

44

46

Average Task Accuracy

Llama2-7B : Fixed vs. Variable

df

policy

(

kf

=0.25)

(post, Fixed)

(post,  Variable)

(pre, Fixed)

(pre,  Variable)

0.2

0.3

0.4

0.5

Compression Ratio

42

44

46

48

50

52

Average Task Accuracy

Llama2-13B : Fixed vs. Variable

df

policy

(

kf

=0.25)

(post, Fixed)

(post,  Variable)

(pre, Fixed)

(pre,  Variable)

0.2

0.3

0.4

0.5

Compression Ratio

35

40

45

50

55

Average Task Accuracy

Llama3-8B : Fixed vs. Variable

df

policy

(

kf

=0.25)

(post, Fixed)

(post,  Variable)

(pre, Fixed)

(pre,  Variable)

In this section, we present our experiments utilizing a variable d f policy per layer in Loki. For this experiment, we set d f 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 5) against the compression ratio, calculated as follows:

Figure 15 presents these plots for the Llama2-7B, Llama2-13B, and Llama3-8B models. We observe that the variable d f policy does not yield significant improvements over the fixed d f policy used in our 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 f policy, combined with its comparable performance to the variable d f policy, makes it a practical choice for Loki.

C Comparison of our kernels with SparQ

As mentioned in Section 4.3, we create optimized kernels in Triton to efficiently compute the three matrix multiplications in Loki (lines 5, 8, and 9 of Algorithm 1) 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 of sizes m × k and k × n , then SparQ kernels parallelize compute along only the m 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-powers of 2 number of tokens in the KV-cache, a setting which is commonly encountered in inference since we generated keys and values one at a time. Therefore, we extend their kernels to handle non-powers of two number of tokens in the KV-cache successfully. In Figure 16, we compare the performance of our kernel with SparQand vanilla PyTorch based attention for an attention layer in Llama2-7B for various sizes of the KV-cache ranging from 512 to 4096. We do this for the matmul operation of query and keys with topk as 0.25.

1

2

4

8

16

0.01

0.05

0.1

Time (ms)

Llama2-7B - Time for computing Q . K T with 512 keys and k=0.25

Vanilla Attention

SparQ

Our kernel

1

2

4

8

16

0.01

0.05

0.1

0.2

Time (ms)

Llama2-7B - Time for computing Q . K T with 1024 keys and k=0.25

1

2

4

8

16

0.01

0.05

0.1

0.3

Time (ms)

Llama2-7B - Time for computing

Q

.

K T

with 2048 keys and k=0.25

Vanilla Attention

SparQ

Our kernel

1

2

4

8

16

0.01

0.05

0.1

0.2

0.5

Time (ms)

Llama2-7B - Time for computing

Q

.

K T

with 4096 keys and k=0.25

We see very high speedups over SparQ for small batch sizes. For instance, for a batch size of 1 with 4096 keys (bottom right), our kernel is faster than SparQ by nearly 2 . 8 × ! Infact, the SparQ kernel barely obtains any speedup over vanilla 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 as we discussed before. In attention the m-dimension is proportional to the batch size, and thus having a 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 are very common in inference and hence it is extremely important to design a kernel that can function efficiently even at smaller batch sizes. Our speedups over SparQ decrease with larger batch sizes, but overall we are always faster than both SparQ and vanilla PyTorch.

D Generation Examples

In this section, we look at a generation example from the Llama2-7B (Chat) model given a custom prompt for different settings of k and d .

Prompt

In a 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 we create a future where AI and humanity coexist in harmony, or will the pursuit of power and profit lead to our downfall? The choice is ours, and the clock is ticking.

Llama2-7B Chat k = 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 our humanity? The future is ours to create, and the choice is ours to make. The clock ticks on, waiting for us to take the next step forward into an uncertain tomorrow, where the only constant is the endless possibilities of an age of AI

Will you answer the call to shape the future of AI? Will you play a part in shaping the course of human history? The possibilities are endless, and the outcome uncertain. Are you ready to embark on a journey that could change the world forever?

E PCAAttn: Directly Leveraging Reduced-Dimensional Attention Scores

One other approach we tried is to directly use the formulation in 4.1 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 PCAAttn (Algorithm 2).

Compute and Memory Analysis: When computing attention between a single query q S ∈ R 1 × D and the key vectors K : S ∈ R S × D , the matrix multiplication q S K T : S has a complexity of O ( DS ) . Using PCAAttn, the key and query vectors are reduced to d dimensions and the complexity of the matrix multiplication is reduced to O ( dS ) . Thus, we can get a speedup of D/d in the attention dot product computation. The PCA transformation of the query and key vector generated at each step has a complexity of O ( D 2 ) , which is small when S >> D . The KV-cache memory requirement is 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 a significantly small memory overhead of O ( Dd ) .

Experimental Results:

Table 5 shows the performance of PCAAttn on Llama2-7B and Mistral-7B models. We can see that our PCAAttn method performs poorly compared to all the baselines and the H 2 Omethod 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 10 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 on Loki instead in the main paper. Note, that we only tried the post-rotary transformations in PCAAttn, and it is possible that pre-rotary transformations might perform better, but we leave this for future work.

F Estimate of Compute Resources Required to Replicate our Experiments

As mentioned in Section 6, we conduct all of our experiments on Perlmutter, a multi-GPU cluster with 4 A100 GPUs per node. Since we do not do any training/fine-tuning, our experiments can be done on a 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 LLama270B, Llama3-70B), we had to resort to running on four A100 GPUs (or a single node) with tensor parallelism using the AxoNN parallel deep learning framework. All results for 7B and 13B sized models can be compiled within 3 hours. For larger models like the 70B Llama-2 and 3 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.

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]

Justification: The claim around investigating the intrinsic dimensionality of attention keys is addressed in Section 3. The claims around designing Loki, providing theoretical analysis, and providing an optimized kernel are addressed in Section 4. Finally, the model quality analysis and compute analysis showing the efficacy and speedup Of Loki are addressed in Section 6.

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 as long as it is clear that these goals are not attained by the paper.

2. Limitations

Question: Does the paper discuss the limitations of the work performed by the authors?

Answer: [Yes]

Justification: Limitations are discussed in Section 7

Guidelines:

The answer NA means that the paper has no 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 a 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 on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.

The authors should reflect on the factors that influence the performance of the approach. For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to 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 fairness.

While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important 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

Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?

Answer: [Yes]

Justification: The paper introduces two lemmas in Section 4 and provides proofs 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 crossreferenced.

All assumptions should be clearly stated or referenced in the statement of any theorems.

The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a 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

Question: Does the paper fully disclose all the information needed to reproduce the main experimental 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)?

Answer: [Yes]

Justification: Algorithm 1 provide a detailed description of our method to allow for reproducibility. 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 3. We also make our code available publicly with the link provided in the paper.

Guidelines:

The answer NA means that the paper does not include experiments.

If the paper includes experiments, a 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 a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it 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 one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.

While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example

(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.

(b) If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.

(c) If the contribution is a new model (e.g., a large language model), then there should either be a way to access this model for reproducing the results or a way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).

(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to 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 to the data and code, with sufficient instructions 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 to the Github repository is provided in 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 NeurIPS code and data submission guidelines ( https://nips.cc/ public/guides/CodeSubmissionPolicy ) for more details.

While we encourage the release of code and data, we understand that this might not be possible, so 'No' is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).

The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines ( https: //nips.cc/public/guides/CodeSubmissionPolicy ) for more details.

The authors should provide instructions on 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. If only a 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 in 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, hyperparameters, 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 3 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 a 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, or as supplemental material.

7. Experiment Statistical Significance

Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?

Answer: [No]

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, we 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, confidence 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, or 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 or the standard error of the mean.

It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 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 or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).

If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.

8. Experiments Compute Resources

Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?

Answer: [Yes]

Justification: We mention the compute resources used in our experimental setup in Section 6 and discuss the compute resources in detail in F

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 as well as 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 it into the paper).

9. Code Of Ethics

Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines ?

Answer: [Yes]

Justification: The paper conforms to the NeurIPS Code of Ethics.

Guidelines:

The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.

If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.

The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in 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 1 and discusses the impact of our method in 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), fairness 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 not tied to particular applications, let alone deployments. However, if there is a 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 as 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 or 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 a system learns from feedback over time, improving the efficiency and accessibility of ML).

11. Safeguards

Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?

Answer: [NA]

Justification:

Guidelines:

The answer NA means that the paper poses no such risks.

Released models that have a 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 is challenging, and many papers do not require this, but we encourage authors to take this into account and make a 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?

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 or dataset.

The authors should state which version of the asset is used and, if possible, include a 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 be provided.

If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a 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 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 or 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 is 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 nor research with human subjects.

Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper.

According to the NeurIPS 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.

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 an equivalent approval/review based on the requirements of your country or 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 to 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.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis method propose the PCA based attention score approximation for top-k attention. Perform PCA on offline dataset, and store the PCA vectors for inference.\n\nSTRENGTHS:\nThis work easily makes a QK approximator without gradient-based training but with only simple PCA for top-k attention selection.\n\nWEAKNESSES:\nYou 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)\n\nQUESTIONS:\n1. 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.\n\nLIMITATIONS:\nPCA 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.", "TITLE": "Official Review", "RECOMMENDATION": 8, "REVIEWER_CONFIDENCE": 5, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2024-07-03"}

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

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

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

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

LIMITATIONS:
N/A
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2024-07-12

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

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

WEAKNESSES:
The method does not seem to reduce the memory usage.

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

LIMITATIONS:
yes
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2024-07-12

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\nEmpirical 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.\n\nWEAKNESSES:\nImplementation 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.\n\nQUESTIONS:\n1. 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?\n\nLIMITATIONS:\nN/A", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2024-07-13"}