H2O: Heavy-Hitter Oracle for Efficient Generative Inference of Large Language Models
Authors: Zhenyu Zhang
Ying Sheng
Tianyi Zhou
Tianlong Chen
Lianmin Zheng
Ruisi Cai
Zhao Song
Yuandong Tian
Christopher Re
Clark Barrett
Zhangyang Wang
Beidi Chen
Venue: NeurIPS 2023 poster (2023)
Keywords: Large Language Models; Efficient Generative Inference

Abstract: Large Language Models (LLMs), despite their recent impressive accomplishments, are notably cost-prohibitive to deploy, particularly for applications involving long-content generation, such as dialogue systems and story writing. Often, a large amount of transient state information, referred to as the $\mathsf{KV}$ $\mathsf{cache}$, is stored in GPU memory in addition to model parameters, scaling linearly with the sequence length and batch size. In this paper, we introduce a novel approach for implementing the  $\mathsf{KV}$ $\mathsf{cache}$ which significantly reduces its memory footprint.  Our approach is based on the noteworthy observation that a small portion of tokens contributes most of the value when computing attention scores. 
 We call these tokens Heavy Hitters ($\mathsf{H_2}$). Through a comprehensive investigation, we find that ($i$) the emergence of $\mathsf{H_2}$ is natural and strongly correlates with the frequent co-occurrence of tokens in the text, and ($ii$) removing them results in significant performance degradation. Based on these insights, we propose Heavy Hitter Oracle ($\mathsf{H_2O}$), a $\mathsf{KV}$ $\mathsf{cache}$ eviction policy that dynamically retains a balance of recent and $\mathsf{H_2}$ tokens.
 We formulate the $\mathsf{KV}$ $\mathsf{cache}$ eviction as a dynamic submodular problem and prove (under mild assumptions) a theoretical guarantee for our novel eviction algorithm which could help guide future work. We validate the accuracy of our algorithm with OPT, LLaMA, and GPT-NeoX across a wide range of tasks. Our implementation of $\mathsf{H_2O}$ with 20\% heavy hitters improves the throughput over three leading inference systems DeepSpeed Zero-Inference, Hugging Face Accelerate, and FlexGen by up to $29\times$, $29\times$, and $3\times$ on OPT-6.7B and OPT-30B. With the same batch size, $\mathsf{H_2O}$ can reduce the latency by up to $1.9\times$.

H20: Heavy-Hitter Oracle for Efficient Generative Inference of Large Language Models
Zhenyu Zhang', Ying Sheng", Tianyi Zhou', Tianlong Chen', Lianmin Zheng", Ruisi Cai', Zhao Song". Yuandong Tian", Christopher Ré², Clark Barrett?, Zhangyang Wang Beidi Chen "University of Texas at Austin, "Stanford University, UUniversity of California, San Diego, Unniversity of California, Berkeley, "Adobe Research, Meta AI (FAIR), "Carnegie Mellon University fzhenyu. 2hang, tianlong. chen, ruisi. cai, ruisi.cai,atlasvang aasssaannauteeaa.....aa @utexas. .ddu, yingl1 ord edu, (chrismre, barrett) stanford. edu, tEzhou@ucsd. edu, Liarminzheng@gmaill zsong@adobe. com, yuandong@meta. com, beidicaandrew cmu, edu
Abstract
Large Language Models (LLMs), despite their recent impressive accomplishments. are notably cost-prohibitive to deploy, particularly for applications involving long- content generation, such as dialogue systems and story writing. Often, large amount transient state information, referred to as the KV cache, is stored GPU memory in addition to model parameters, scaling linearly with the sequence length and batch size. In this paper, we introduce a novel approach for implementing the KV cache which significantly reduces its memory footprint. Our approach is based on the noteworthy observation that a small portion of tokens contributes most of the value when computing attention scores. We call these tokens Heavy Hitlers (H2). Through a comprehensive investigation, we find that (i) the emergence of H2 is natural and strongly correlates with the frequent co-occurrence of tokens in the text, and (ii) removing them results in significant performance degradation. Based on these insights, we propose Heavy Hitter Oracle (H₂O), KV cache eviction policy that dynamically retains a balance of recent and H2 tokens. We formulate the KV cache eviction as a dynamic submodular problem and prove (under mild assumptions) theoretical guarantee for our novel eviction algorithm which could help guide future work. We validate
accuracy of our algorithm with OPT, LLaMA, and GPT-NeoX across wide range of tasks. Our implementation of H2O with 20% heavy hitters improves the throughput over three leading inference systems DeepSpeed up to 29x, 29x, and 3x on OPT-6.7B and OPT-30B. With the same batch size, H2O reduce the latency by up 1.9x. The code is available at https: //github. com/ FMInf #MInference/H20
Hugging Face Accelerate, and FlexGen by
Introduction
Large Language Models (LLMs) have demonstrated remarkable proficiency in a wide range of natural language processing applications such as content creation, summarization, and dialogue systems [ 2 4 However, their deployment is very costly. In addition the widely-studied bottlenecks of model size and the quadratic cost of attention layers, the problem of the size of KV cache, which stores the intermediate attention key and values during generation to avoid re-computation, is becoming increasingly prominent (5). For instance, 30 billion-parameter model with an input batch size of 128 and a sequence length of 1024 results in 180GB of KV cache. A natural approach is to limit its maximum size is done in classical software hardware caches (6]. However, it is challenging to reduce KV cache memory footprints in LLMs without accuracy drops.
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

Dynamic Sparsity
Static Sparsity (Strided)
Static Sparsity (Local)
Static Sparsity Hl.O
Childine
langhed
played
92
UJ
94
52
Spureis (Stiaal)
0J1
0.5
Stik Sparals (Loxl)
H₂O O Query
Statk Sparsiyy w. Hij
20
20
80
120
Marsey Redastom
Figure 1: Upper plots illustrate symbolic plots of an attention map deploying different KV cache policies in LLM generation. Lower right: contrasts their accuracy-memory tradc-off. Left: the overview of H2O framework.
While there exists substantial literature on sparse attention approximation in training, they have not seen wide adoption for alleviating KV cache bottleneck. First, most existing methods, e.g., Reformer 77) and Flash Attention [8) are designed to overcome the quadratic memory required by attention mechanisms when modeling long sequences but still require large cache size. Second, variants like sparse transformer [9]. low-rank based transformers [10 I1 or multi-query attention [12
13 [5] can reduce the cache size, but directly applying them on pre-trained LLMs for generation
results in high miss rates and degrades the accuracy as shown in Figure [ Finally, some recent advances such as gisting tokens [4] can learn to compress the KV cache for documents, but their expensive eviction policies are difficult to deploy during generation.
Therefore, ideal KV cache should have (i) small cache size reduce memory footprint, (ii) low miss rate to maintain the performance and long-content generation ability of LLMs, and (iii) low-cost eviction policy to reduce the wall-clock time during generation. However, there are three technical challenges. First, it is not immediately clear whether the size of the KV cache can be restricted each decoding step might, in principle, require access to all previous attention keys and values. Second, identifying an optimal eviction policy that maintains generation accuracy is combinatorial problem Finally, even if an optimal policy can be brute-forced, it is infeasible for deployment on real-world applications.
Fortunately, preliminary exploration has yielded intriguing observations about the empirical properties LLMs. These findings pave the way for the potential design of an efficient KV cache.
Sparsity for small cache size; We observe that even when trained densely, the attention matrices of LLMs are over 95% sparse at inference time (shown in Figure 2 This holds for wide range of re-trained LLMs. Therefore, only 5% of the KV cache is sufficient for decoding the same output token at each generation step, which suggests it may be possible to have up to a 20 reduction in KV cache size without accuracy drop.
Heavy-Hitters for low miss rate: We discover that the accumulated attention scores of all tokens in attention blocks adhere to a power-law distribution. It suggests that there exists small set of influential tokens that are critical during generation, named heavy-hitters (H2). H2 provides an opportunity to step away from the combinatorial search problem and identify eviction policy that maintains accuracy.
Greedy algorithm for low-cost policy: we surprisingly find that retaining the H2 based on local statistics each decoding seep-summing the attention scores of only the preceding tokens-2a as effective as considering the attention of future tokens (shown in Figure 2.
Based on the above, first rigorously define the generative process of L.LMs operating with size-constrained KV cache in Section 2.1 Then we propose Heavy-Hitter Oracle (H2O), a framework
Belady's ' Algorithm is optimal for standard cache, but necessarily for KV cache.

---

that exploits the properties of LLMs and uses simple, low-cost eviction policies that retrain the quality of LLMs throughout the generation process. Specifically,
In Section 3 we explore the emergence of H2 in attention, revealing their fundamental and critical roles: (i) H2 exhibit strong correlation of frequently co-occurring words in textual data; and (ti) removing H2 completely damages the model's functionality. We demonstrate that H2 can largely lower the cache miss rate of the existing policies mentioned above. Theoretically, assuming the attention scheme is submodular, H2 corresponds to greedy algorithm and is therefore near-optimal.
In Section 4 by the accumulated attention score at each decoding step. We formulate the eviction policy with greedy H2 variant of dynamic submodular maximization. The analysis shows that it results in a similar generative process as the one using the H2 eviction policy.
we present greedy but low-cost variant of H2 which is dynamically determined
We perform extensive experiments on OPT, LLaMA, and GPT-NeoX on a single NVIDIA A100 (80GB) GPU to evaluate H₂O across a range of tasks from Im-eval-harness [T5] and HELM [T6]. We implement H2O on top of FlexGen that can easily adapt different cache eviction techniques to produce a system with high-throughput inference. Performance experiments show our framework achieves 29 x, 29x, 3x higher throughputs compared to three leading inference systems, DeepSpeed
Zero-Inference [7]. Hugging Face Accelerate [18), and FlexGen [19] respectively. With the same batch size, H2O achieves up to 1.9x lower latency compare to FlexGen.
2 Related Work and Problem Setting
The substantial parameter counts of large language models (LLMs) present significant challenges for inference. To overcome this limitation, previous efforts have employed model compression techniques with specific designs to achieve efficient LLM inference,
Efficient Inference of LLMs.
such as the method described 20 21 [22] which employs one-shot pruning on LLMs, resulting in
negligible performance degradation even without retraining. Additionally, alternative approaches explore quantization methods specifically tailored to LLMs, as discussed [3] 24 25 26 27 28).
Also, CoLT5 [29] employs ooken-wise conditional computation strategy to reduce the overall computation cost. These methods address efficient inference from orthogonal perspectives and can be organically integrated. The techniques investigated in this study are closely associated with pruning or sparsity but focus on a distinct inference bottleneck, namely, KV cache. One closely related work (30] utilizes learnable mechanism that determines necessary tokens during inference but requires an extra fine-tuning process, which makes it less practical.
Sparse, Low-rank Attention Approx. The quadratic computational complexity of attention mod- ules is one of the major bottlenecks of transformer inference B1. Various efforts are devoted to addressing this challenge 7 10 For example, Reformer 7 reduces the computational cost from quadratic to superlinear complexity via locality-sensitive hashing. Performer [I0] employs positive orthogonal random features to approximate attention kernels. One relevant work, Sparse Transformer [9) introduces sparsity to reduce KV cache memory footprint and achieve efficient attention mechanism, considered as our baseline in this paper. Moreover, SpAtten (32) utilizes accumulated attention scores to select important tokens for efficient attention inference while they don't consider the variance of token importance across attention heads and layers. Comparison with SpAtten is detailed in Appendix C.9
Caching. Caching, which plays a pivotal role in optimizing system performance, entails the devel- opment of effective eviction policies to handle frequently accessed data. Conventional approaches such as Least Recently Used and Least Frequently Used (33 34] prioritize the recency and frequency of data access. And the design of cache encounters many similar challenges traditional caching.
LLM Inference Breakdown. The generative procedure of LLMs encompasses two distinct phases: (i) the prompt phase, in which an input sequence is utilized to produce the KV cache (consisting of the key and value embeddings), similar to the forward pass employed during LLM training; and (ii) the token generation phase, which leverages and updates the KV cache to generate new tokens incrementally. Each generation step relies on the previously generated tokens. The primary focus of this paper is to enhance the efficiency of the KV cache in attention during the token generation phase, thereby accelerating LLM inference.
3

---

OpceBook@M
MatkQA
Baacte No.. Hong tikk
73
PADA
COPA
1011
Wanggande
OpecBadDA PI4
8.0%
Taye baier
(8)
(b)
(c)
(d)
Figure 2: (a) Attention Sparsity in vre-trained LLMs. (b) The distribution of accumulated attention scores with respect to the corresponding word (red scatter) and the co- occurrence times of words in the data (gray curve). The x-axis represents the word index in the vocabulary. (c) The performance comparison between the baseline model with full KV and the model w.o. heavy hitter. (d) Comparison hetween the baseline model with full KV, HzO with the local statistic, H2O with the global statistic, and model with only the most recent KV (local).. Apart from the baseline model, each model is evaluated with 20% KV cache budget.
2.1 Problem Formulation
We formally define the generative process with limited KV cache size. Denote attention query matrix as Q E Rakd and key matrix as E R=dd Qi represents the i-th row of Q and Ksi, represents the first rows of K. Let & denote the budget of space and h < n. For simplicity, Ks. (E Rixd) denotes sub-matrix of K which selects S, rows from K. (For the non-selected rows [i],,,, we put all zeros in that row) Eviction policy is defined as:
Definition 2.1 (Eviction Policy, informal). Let S denote the source set. Let St denote the target set. We defined the eviction policy g S - S; such that
|S:| = k (KV cache size is not changing over the time)
| S 1 or equivalently |S; 11 2 k (we can evict at most KV in the KV cache)
Then, we define the generative process with our eviction policy.
Definition 2.2 (The generative process with eviction policy, informal). Let is denote the size of the KV cache. For each [n], for the i-th token, we have
Let S C [m] denote the tokens in KV cache when predicting the i-th token.
The information we have is a length-i vector 0; == D exp((i;, (Ksp(Q) (Ks.s)T)
(normalized attention)
scalar D; == (K
(,yss)) 1; (the evicted KV is set to 0, and we need to subtract them when computing the normalization)
Replacing S, by in the above definition o; and D; leads to standard generative process.
The eviction policy (Definition 2.1 updates S; based S;-1 and their corresponding information.
Remark 2.3. Our goal is to find a KV cache eviction policy such that the output of the generative process is similar comparable to the original one without limiting the cache size.
3
Observations
We present two key empirical insights LLMs that inspire the design of H2O, as follows.
3.1 Sparsity for Small Cache Size
Inspired by previous literature, which reveals the existence of attention sparsity in DistilIBERT [35] and bounded-nowd mounded-norm self-attention heads [36]. We first show an observation on the sparsity of attention in pre-trained LLMs. Then we discuss how it can potentially unlock the possibility of reducing KV cache size without an accuracy drop. Given the normalized attention score Softmax(QK matrix that is calculated by the query matrix Q and the key matrix K, set the threshold one percent of the maximum value in each row and calculates the corresponding sparsity.
Observation. We conduct zero-shot inference with the pre-trained OPT model on the validation set of Wiki-Text- 103. We plot the layer-wise sparsity within attention blocks and visualize the normalized attention score matrix. The results are presented in Figure 2 (a). We observe that although the LLMs are densely trained, the resulting attention score matrices are highly sparse, with a sparsity over 95% in almost all layers.

---

The attention blocks' sparsity suggests that access to all previous key and value embed- dings unnecessary for generating the next token. This suggests it is possible evict unessential KV embeddings and reduce the requirement of KV cache during generation.
Insights.
3.2 Heavy-Hitters for Low Miss Rate
The previous section showed the sparse nature of attention blocks in pre-trained LLMs, which provides the opportunity for designing small KV cache size while still maintaining the performance of LLMs. However, determining the best eviction policy that preserves generation accuracy presents a combinatorial challenge. Although Belady's 's Algorithm (37 is optimal and easy to compute for standard cache (offline), not applicable for KV cache design. Because once evicting important KVs, it could destroy the performance of LLMs due to the sequential dependency of LLM generation.
Observation. Fortunately, in the early stage of our exploration, we find that the accumulated attention scores of all the tokens within attention blocks follow a power-law distribution, as shown in Figure 2 This suggests the existence of small set of tokens that are critical during generation. We denote those tokens as heavy- sattteitterss (H2). In order to verify the importance of these tokens, we compare the quality LLM generation after masking heavy hitters with that of the original model. Not surprisingly, as shown in Figure the accuracy drops drastically, confirming the importance of those tokens. Additionally, we can see the accumulated attention score of each word (in red dots) have a high correlation with their co-occurrences in the data (gray curve).
Analysis. First, based on H2, we see an opportunity to side-step from the combinatorial search problem and design KV cache eviction policy that preserves the LLM generation quality. We conduct an empirical study implementing a KV cache eviction policy that retains only the H2 and the recent KV embeddings in the cache. The intuition is that recent words typically exhibit stronger correlations with current tokens. We assess the effectiveness of this eviction policy through pre- trained OPT-30B and six downstream tasks. The outcomes of these evaluations are illustrated in It is obvious that the H2 based eviction policy can largely reduce the KV cache size without
Figure 2 degrading the performance of OPT-30B.
Moreover, during the post analysis, inspired by 338, we find that H2 based policy is related to the classical greedy algorithm (a polynomial-time algorithm with provable guarantees) under the assumption that the attention schema is submodular. We present details in Appendix D
Lemma 3.1 (informal). Assuming the attention scheme is submodular, then greedily constructing the set S; (without cache size limitation) satisfies the near-optimal property in terms of submodular.
4 Heavy-Hitter Oracle
The goal of this section is to propose the greedy algorithm using the H2- H2-based policy and to show the provable guarantees. We first present the Ha-based policy called H₂O cache eviction policy and formulate its deployment in LLM generation as a variant of submodular maximization problem, named dynamic submodular. Then we present H₂O in the generative process, followed by a practical example of deploying our proposal. Finally, we provide theoretical guarantees for H2O and show our efficient system implementation.
4.1
Greedy Algorithm for Low-Cost Policy
We have shown a simple yet effective KV cache policy based on H2. However, it is impractical to deploy such an algorithm because do not have access to the future-generaie tokens. Fortunately, we empirically observe that local H2, which is calculated using local statistics at every decoding step by summing up the attention scores of the previous tokens, is equally effective as taking into account the attention future tokens (Figure 2). In the following, we formally define this dynamic attention score computation (with space limitation) as a novel dynamic submodular type problem.
Definition 4.1 (Dynamic submodular framework, informal). Define function F 2lm x 2|m] IR, then for any set Z C [n)], assume that F(Z,.) 2 R is submodular function w.r.t. Z, i.e.
For all sets X, C [n] satisfy that Z C X Y.
For all element x E [n] satisfy that x E [n]]Y,
have f(X U {x}) f(x) 2 f(Y U f(Y), where f()) = F(Z..).
5

---

Remark 4.2. We provide practical insights of Definition 4.1 X denotes the existing words in the KV cache. Y any superset of X. z can be viewed as a "word" which is either newly added to KV cache or existing deleted from KV cache. An example I can be attention score, L.e,, see Algorithm 7
load the sequence of S1, "., Su (we promise that [Sa] k and |SAA,-1 1) into Definition 4.1 i.e. for each E [n], we choose Z = Sü, then it becomes a particular instance of the dynamic submodular problem.
Next, we provide a formal description our algorithm, followed by an example.
Definition 4.3 (H20 Eviction Policy). Let Fawore 2 -> IR. denote certain score function. Let S denote the source set. Let S; denote the target set. We defined the eviction policy S - S; s.1.
| k (KV cache size is not changing over the time)
| < 1 or equivalently |St S5-1 2 k (we can evict at most KV in the KV cache)
We construct S + (Si-1 U (u) as u +- arg max,e(S, 20(()) Fscore U (i}(t
To describe our algorithm (Algorithm D. we choose a particular instantiation of the function F Fscore. i.e., the summation of that sets in the attention matrix.
plone
Algorithm H2 Eviction Algorithm
1:
procedure H2_ EVICTION( guxd, &e N)
2: Let à denote the budget size of cache
Step
3: So
4:
for
do i k then
5:
6:
S,-1 U (()
7:
else
8:
D. - (KxplQ.. 1.+) 114180-1)
Decoding Stup
9:
DJ (exp(Qs, (Ks; 1.x) 1(118,_1)
10:
(T) = Exero
11:
G. 9--1 Us}
12:
arg max Focoro (8:-1 (1M(o) veG;
13:
(8 (
14:
end if end for end procedure
Eviction n. Glataal Statitic
15:
16:
Figure 3: Illustration of Algorithm Srrrihhmmdurnnng during two consecutive decoding steps.
Figure presents an illustrative example of our H2 Eviction Algorithm. We assume that the budget size of KV cache is 3. Following the completion of the fourth decoding step, the KV embeddings associated with the third token are evicted based on the accumulated attention score. Consequently, these evicted KV embeddings become inaccessible in the subsequent decoding steps.
4.2 Theoretical Guarantee and System Implementation
We state theoretical result as follows. The proofs and more details provided in Appendix D
Theorem 4.4 (informal). Under the mild assumption, let denote the budget of space limitation. If for each token, we greedily compute the attention score based top-k choice, then the set S; generate each for token satisfy that f(S.) (1 - a] ma J(S) B.
show
where B >> O are parameters.
Remark 4.5. We remark the above theorem provides theoretical explanation of why can hope our greedy algorithm (with cache limitation) can provide good solution to the problem.
Implementation Details. We provide a general framework that can support any KV cache eviction algorithm and enhance throughput and reduce the latency of L.IM generation with careful implemen- tation. For example, to ensure I/O efficiency, we do not swap memory when stored KV is evicted, but directly fill with newly-added KV. More details included in Appendix A

---

XSUM, aT
XSUM, 1LMAA-33R
XSUM, 1-AMA-
XSUM, GPT-NoxX-208
76-
16-
BOTTEC
PsB
Honry Onadle
Heery-Sttle Sttte Orade
Heery-itter Orade
Loal
20-
Fall
Fill
Fill
100
20
192
20
20
106
60
40
KV Cache Hadgd (h))
KV Cadee Bedar 65
KV Cocle Balget (h)
KV Cocke Budge (5)
CNN/Daall Mail, LLaMA 73
CNN--Dail Mail. GPT NeoX 200
OpenBookQA OPT-30B
OpenBookQA OPT ddB
45.0
42.0
60
Hlanyfitttr Choclo
Local
33.1
Pull
Herry Sttte Oradz
Hony Hitter Cruck
Reary -itter Grade
Local
27.-
Tece
31.1
Local
50
Full
FxB
25.0-
215-
Fdi
100
KV Cucle (5d
KV Cucee Bedar (5))
Clobe Bodgd (5)
KV Cucke Badad
COPA, OPT-10R
COPA, GPT-66B
MaiQAA OPT-308
RTE, OPT-663
27-
10-
255
25
10.
HeayyElitr
66-
Pull
Heery-ltte trads
Licey-llire Orac
Hieryditner ttrde
Loal
Lavel
Local
56-
Pull
21-
Fell
Pell
102
10
10
40
20
20
100
40
KV Cache (ho
KV Cacha Bedar: (4)
KV Cade Stdgeet 196)
KV Cacha Bedan:
Figure 4: Comparsion results between the baseline model with full cache, our H₂O, and the "Local" strategy that utilizes the most recent KV embeddings.
5 Empirical Evaluation
In this section, our goal is to demonstrate that H2O, remarkably simple KV cache eviction policy is capable of enhancing end-to-end throughput and reducing latency in wall-clock while maintaining generation quality across broad spectrum of domains and tasks.
we show that H2O can reduce the memory footprint of KV cache by up to 5x without accuracy degradation on wide range of model architectures (OPT, LLaMA, GPT--NeX),, sizes (from 6.7B to 175B) and evaluation benchmarks (HELM and Im-eval-harness), More importantly, can enhance the performance of existing KV cache sparsification techniques.
Section 5.1
In Section 5.2 we demonstrate that H2O increase the inference throughput by up 3x, 29x, 29 compared to the state-of-the-art inference engine FlexGen, DeepSpeed and the widely used Hugging Face Accelerate without compromising model quality.
Section 5.3 we present extensive ablation studies to show the effectiveness of H2O under different sequence lengths, especially the input with infinite sequence length and its compatibility with quantization.
All details (hyperparameters, data splits, etc.), along with additional experiments, are in Appendix A
5.1 End-to-End Results
We demonstrate that H2O can reduce KV cache memory footprint by 5-10x while achieving compa- rable accuracy on majority of tasks.
Setup. Our experiments are based on three representative model families of LLMs, including the OPT
30 with model sizes, L.LaMA 400, and GPT--eeX-2200 [41] We sample eight tasks from two
popular evaluation frameworks (HELM 16 and on-eval-harness (55):: COPA (42, MathQA 43).
OpenBookQA 44 PiQA 455, RTE 46) Winogrande [47). XSUM [88] CNN/Daily Mail [49]. Also,
we evaluate our approach on recent generation benchmarks, AlpaceEval [50] and MT-bench [5]],
and the details are included in Appendix. We use NVIDIA A100 80GB GPU.
Since H2O evenly assigns the caching budget to H2 and the most recent KV, except for full KV cache, we consider the "Local" strategy as baseline method. In addition, we also provide two different variants of Sparse Transformers (strided and fixed) strong baselines. Also, the full
Baselines.
7

---

KV cache with fewer shots (0/1-shot) prompts are considered as the baseline, which has similar sequence length of the 5-shot tasks with 20% KV cache budget.
Main Results. We evaluate LLMS with KV cache budget ranging from 4% to 100% on 5-
Table I: Quantatively comparison between H₂O with Full methods of different number of shots.
shot downstream tasks. Results are summarized
in Figure 4 and Table D 2 The following ob-
Methods
PIQA
COPA
OpenbookQA
Winogrande
servations can be drawn: (I) With different KV
Full
80.09
81.00
44,80
71.51
cache budgets, our H2O demonstrates consistent
O-shot Full
78.89
76.00
41,40
70.00
and significant improvements against the "Lo- cal" strategy across various model sizes, model types, and downstream tasks. We can draw sim- ilar conclusions comparing H2O with other baselines like Sparse Transformer; (2) Mcanwhile, with
1-shot Full
79.11
76.00
43.60
70.24
Local
57.94
56,00
28.40
$1.30
H₂O 79.22
85.00
43.80
71.67
less than 20% KV cache budget(i.e.. more than 5x memory reduction), H2O achieves comparable performance as the model with full KV embeddings; (3) H2O with 20% KV cache budget approxi- mately uses 1.2 samples per input and show consistent improvement over zero-shot and one-shot full model that use and 2 samples, respectively. (4) Our H₂O shows consistent effectiveness in the more challenging long sequence generation tasks, XSUM, and CNN/Daily Mail.
Analysis. Since the evicted KV will not be seen in the future steps. dropping certain critical KV embeddings can cause a se- vere functional collapse, result- ing in significant performance degradation, e.g., (LLaMA- 13B, XSUM} (LLaMA-7B. CN- N/Daily Mail], the "Local" strat- egy collapses at 60% budgets while our H₂O can still match the full cache performance with 20% budgets. In some tasks, our methods even surpass the baseline models, which demonstrates regularization effect of our H2O. For example, in (OPT-66B, RTEJ, {OPT-30B, MathQAJ and (GPT-NeoX- -20B XSUM), our H2O achieves an extra performance improvement of 0. 037, 0.64% and 0.18 with 20% KV cache budget, respectively. These consistent results validate the effectiveness of our H₂O framework.
Table 2: Results of different sparsification methods w. or wo H2- Experiments are conducted with OPT-30B with 20% KV cache budget.
Models
COPA
OpenBookQA
PIQA
Winogrande
Full
85.00
43.20
70.24
Local vco. Hs
48.00
23.20
55.82
49.17
Local w: Hz
84.00
43.00
78.45
69.06
Sparse Transformer (strided)
50.00
24.60
56.20
47.59
Sparse Transformer (strided) w. H2
83.00
42.60
78.24
69.61
Sparse Transformer (fixed) w.n.
Sparse Transformer (lixed) w. H3 76.00
41.40
77.80
64.96
Enhancing Baseline Techniques. Importantly, we observe other sparsification baselines fail under an extremely low cache budget while combining the most recent KV embeddings with the ones of heavy hitters successfully achieves comparable performance as using full KV embeddings. From Table 2] we can observe that both "strided" and "fixed" sparse attention fail under 20% KV cache budgets, encountering a significant performance drop (up to 35% compared with the full cache). After combining with H2, both approaches reach similar performance using full KV embeddings.
5.2 Heavy Hitter for High- Throughut Generative Inference
Table 3: Generation throughput (token/s) T4 GPU with different systems. In the sequence length row, we use "512 + 32" to denote prompt length of 512 and a generation length of 32, "OOM" means out-of-memory. The gray text in the bracket denotes the effective batch size and the lowest level of the memory hierarchy that the system needs for offloading, where "C" means CPU and "G" means GPU.
Seq. length
512+32
512+512
512+1024
Model size
6.7B
30B
6.7B
30B
6.7B
30B
Accelerate
20.4 (2. G)
0.6 (8, C)
15.5 (1. G)
0.6 (8, C)
5.6 (16, C)
0.6 (8, C)
DeepSpeed
10.2 (16, C)
0.6 (4, C)
9.6 (16, C)
0.6 (4, C)
10.1 (16, C)
0.6 (4, C)
FlexGen
20.2 G)
8.1 (144, C)
16.8 (1, G)
8.5 (80. C)
16.9 (1, G)
7.1 (48, C)
H2O (20%)
35.1 (4. G)
12.7 (728, C)
51.7 (4. G)
18.83 (416, C)
52.1 (4. G)
13.82 (264. C)
8

---

Table Generation throughput and latency on A GPU. In sequence length row, we use "7000 1024" to denote prompt length of 7000 and a generation length of 1024. "OOM" means out memory. out-of-memory.
Seq. length
Model size
Batch size
Metric
FlexGen
H2O (20%)
7000+1024
30B
1
latency (s)
57.0
50.4
5000+5000
13B
4
latency (s)
214.2
155.4
2048+2048
6.7B
24
latency (s)
99.5
53.5
2048+2048
6.7B
24
throughput (token/s)
494.1
918.9
2048+2048
6.7B
64
throughput (token/s)
OOM
1161.0
We implement our KV cache eviction policy in a state-of-the-art inference engine, FlexGen [19], and report the throughput and latency improvements. H2O is orthogonal to existing optimizations in Flex- Gen, such as offloading and quantization, so they can be combined to achieve better performance.
Table 4: Results of generation throughput (token/s) a T4 GPU with different systems real-world datasets, XSUM.
Model size
6.7B
30B
Accelerate
11.98 (1, G)
0.23 (2, C)
Setup We conducted experiments on two GPUs: an NVIDIA T4 (16GB) GPU and an NVIDIA A100 (80GB) GPU. On the T4 GPU, we evaluate the gen- eration throughput following the settings in the Flex- Gen paper. The evaluated models are OPT-6.7B and OPT-30B. When the model and KV cache do not fit into a single GPU, we turn on CPU offloading. The results of both pure GPU and GPU with CPU offloading are reported. All the speedup results are tested in an end-to-end setting, including both the pre-filling and generation phases. And it includes the time for constructing the H2O KV cache. We use synthetic datasets where all prompts are padded to the same length. The system then required to generate the same number of tokens for each prompt. We test different combinations of prompt and generation lengths. We also test our method on real-world datasets (XSUM) for further assessment. The evaluation metric is generation throughput, which is the number of generated tokens / (prompt time + decoding time). We use DeepSpeed ZeRO-Inference [17, Hugging Face Accelerate [8] and FlexGen [19] as baselines. On the A GPU, with more GPU memory, we evaluate the performance of the systems with sequence lengths up to 10K. Although OPT is only trained on 2K sequence length, we benchmark the throughput and latency performance to show the potential of H2O for better models in the future.
DeepSpeed
3.52 (6, C)
0.31 (2, C)
FlexGen
10.80 (1, G)
3.29 (44, C)
H₂O (20%)
30.40 (1, G) 6.70 (180, C)
Results. Table 3 shows the generation throughput of all systems on the T4 GPU. With our KV cache eviction policy, the memory usage is reduced, which brings two advantages: 1) we can use much larger batch size: 2) we can make setting from requiring offloading to not requiring offloading. As shown in Table & 4 H2O with a 20% budget improves the generation throughput over FlexGen, DeepSpeed, and Accelerate by up to 3x, 29x, and 29x, respectively, across both synthetic and real-world dataset. The results the A100 GPU with sequence lengths from 4K to 10K are listed in T With the same batch size, H2O can reduce the latency by 1.1 1.9x compared to FeexGen. Additionally, H2O saves memory so it allows larger batch size, which brings 2.3x improvement on generation throughput for OPT-6.7B.
5.3 Ablation Results
We present extensive ablation studies of H2O on (1) infinite-length input, (2) different number of shots, (3) compatibility with quantization methods on KV cache, and (4) dissecting the effectiveness of different components. We find surprising property of H2O it not only improves the efficiency of LLMs, but also increases the diversity of the generated text.
Q1: Can H2O empower LLMs to process infinite-enggth inputs? Al: Effective genera- tion with sequence length up to four million tokens. Some recent works [52] 53) demon- strate the possibility of handling infinite-length inputs, a notable challenge in current LLMs. These methods employ an attention sink that retains the first few tokens and applies position rolling in the KV cache, empowering LLMs to process infinite-length inputs. Inspired by this progress, we further implement our H₂O for infinite-length inputs. Figure 5 showeases the pos- itive results of H2O, i.e., H2O can empower LLMs to tackle input with length up to four mil-

---

lion tokens, achieving a better performance (lower perplexity) than the original StreamLLM method 12 across various cache size. Further comparisons are reported Appendix C.4
Q2: Does the number of shots during inference effects the effectiveness of H2 A2: Effectiive across zero-shot to ten-shots inference. We further examine H2O under different numbers of shots during inference, and the re- sults are reported in Table I0 and Figure 8 With different shots inference, our H2O achieves matching performance (difference less than 1.00%) as the full model across dif- ferent downstream tasks. The "Local" strategy encounters significant performance degradation (up to 37.00%. Such results demonstrate the effectiveness of our H2O under different inference scenarios. More details about zero-shot and one-shot inference reported in Appendix C.3
Streaming with H2O to 4 Million Tokens
2M4
Input Length
5.4
StremLLLM
Perplexity
62-
Stezam
5.0
5.-
Q3: Compatible with Quatization? A3: Yes. To pur- sue further efficiency, we show the compatibility of H2O with another orthogonal approach, i.e., quantization in Ta- ble 6 We use OPT-30B as our base model and COPA, OpenBookWA, and PiQA as evaluation tasks. Intuitively sparsity and quantization are highly related so combin- ing them might introduce larger errors. Surprisingly the combination almost always achieves better accuracy than H2O or quantization alone. Experiments about throughput improvement are detailed in Appendix C.2
Cache Size
(Upper) streaming with H₂O to handle inputs with sequence lengths of four million tokens. (Bottom) Perplexity compari- son between the original StreamLLM method and our HzO, results are collected on the first text sample PG-19 19 (54)
Figure 5:
Q4: When does H₂O match the baseline with full KV
embeddings? A4: With both H2 and the recent tokens. We investigate the separate effects of KV embeddings of H2 and the local tokens. We conduct experiments on 4 tasks with OPT-13B and ОPT-30B. For each task, we compare the performance of three KV cache eviction policies, including only the KV embeddings of H2, only the ones of local tokens, and our H2O that keep both. As shown in Table 9 only retaining the embeddings of H2 or local tokens can't maintain similar performance as the model using full embeddings, with a performance degradation from 2.85% to 22.75%. Incorporating both components, our H2O successfully retains the baseline performance with full embeddings. Besides, the model with only H2 shows a consistent improvement against the one with only local tokens, which indicates H2 might contribute more to maintaining the performance.
Q5: Extra benefits from H₂O? A5: Increased diversity of generated text. Besides all the benefits of our H2O, we also observe an bonus introduced by H2O. i.e., the improved diversity of generated content. The results are reported in Appendix [C.]] Given the same prompts, we visualize the generated text of the models with different KV cache budgets. Compared with the model of full KV cache, our H2O can generate sentences with fewer repeated words and more creativity.
Conclusion and Discussion
In this paper, we study one of the key bottlenecks of LLM deployment, KV cache, particularly for long-content and large-batch generation applications. We propose H2O, simple KV cache eviction policy for significantly reducing its memory footprint. The main insight of our approach is the recognition of subset of tokens, known as Heavy Hitters, which contribute the most value when computing attention scores. We formulate the KV cache eviction as a dynamic submodular problem and provide the theoretical guarantees for our algorithm. Through extensive evaluations, we demonstrate that H₂O can significantly improve end-to-end throughput and decrease latency in wall-clock time, without compromising the generation quality of L.LMs across variety of tasks.
7 Acknowledgement
Ying Sheng and Clark Barrett are partly supported by NSF-2110397 and the Stanford Center for Automated Reasoning. Z. Wang is in part supported by a Google Research Scholar Award and the NSF AI Institute for Foundations of Machine Learning (IFML).
10

---

References
[1] Romal Thoppilan, Daniel De Freitas, Jamie Hall, Noam Shazeer, Apoorv Kulshreshtha, Heng- Tze Cheng, Alicia Jin, Taylor Bos, Leslie Baker, Yu Du, et al. Lamda: Language models for dialog applications. arXiv preprint
0tt::2011082.32 2022.
[2] Ann Yuan, Andy Coenen, Emily Reif, and Daphne Ippolito. Wordcraft: story writing with large language models. In 27th International Conference on Intelligent User Interfaces, pages 841-852, 2022.
[3] Jason Wei, Yi Tay, Rishi Bommasani, Colin Raffel, Barret Zoph, Sebastian Borgeaud, Dani Yogatama, Maarten Bosma, Denny Zhou, Donald Metzler, et al. Emergent abilities of large language models. arXiv preprint arXiv:2206.0766 2022.
[4] Tianyi Zhang, Faisal Ladhak, Esin Durmus, Percy Liang, Kathleen McKeown, and Tatsunori Hashimoto. Benchmarking large language models for news summarization. arXiv preprint arXiv:2301. 13848, 2023.
[5] Reiner Pope. Sholto Douglas, Aakanksha Chowdhery, Jacob Devlin, James Bradbury, Anselm Levskaya, Jonathan Heek, Kefan Xiao, Shivani Agrawal, and Jeff Dean. Efficiently scaling transformer inference. arXiv preprint wwxxx:2221.005000 2022.
[6] Laszlo A Belady, Robert Nelson, and Gerald S Shedler. An anomaly in space-time char- acteristics of certain programs running in a paging machine. Communications of the ACM, 12(6)334,-353, 1969.
[7] Nikita Kitaev, Lukasz Kaiser, and Anselm Levskaya. Reformer: The efficient transformer. arXiv preprint arXiv:2001. 04451, 2020.
[8] 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:1 13:16444-16359 2022.
[9] Rewon Child, Scott Gray, Alec Radford, and llya Sutskever. Generating long sequences with sparse transformers. arXiv preprint 1904::9904.1009 2019.
[10] Krzysztof Choromanski, Valerii Likhosherstov, David Dohan, Xingyou Song, Andreca Gane, Tamas Sarlos, Peter Hawkins, Jared Davis, Afroz Mohiuddin, Lukasz Kaiser, al. Rethinking attention with performers. arXiv preprint arXiv:2009, 14794, 2020.
[11] Angelos Katharopoulos, Apoorv Vyas, Nikolaos Pappas, and François Fleuret. Transformers are mns: Fast autoregressive transformers with linear attention. In International conference on machine learning, pages 5165. PMLR, 2020.
[12] Noam Shazeer. Fast transformer decoding: One write-head all you need. arXiv preprint arXiv:1911.0211 2019.
[13] Aakanksha Chowdhery, Sharan Narang, Jacob Devlin, Maarten Bosma, Gaurav Mishra, Adam Roberts, Paul Barham, Hyung Won Chung, Charles Sutton, Sebastian Gehrmann, et al. Palm: Scaling language modeling with pathways. arXiv preprint 20iv:2204.00311, 2022.
[14] Jesse Mu, Xiang Lisa Li, and Noah Goodman. Learning to compress prompts with gist tokens. arXiv preprint arXiv: 230423304 2023.
[15] Leo Gao, Jonathan Tow, Stella Biderman, Sid Black, Anthony DiPofi, Charles Foster, Laurence Golding, Jeffrey Hsu, Kyle McDonell, Niklas Muennighoff, Jason Phang, Laria Reynolds, Eric Tang, Anish Thite, Ben Wang, Kevin Wang, and Andy Zou. A framework for few-shot language model evaluation. In Zenodo. 2021.
September
[16] Percy Liang, Rishi Bommasani, Tony Lee, Dimitris Tsipras, Dilara Soylu, Michihiro Yasunaga, Yian Zhang, Deepak Narayanan, Yuhuai Wu, Ananya Kumar, et al. Holistic evaluation of language models. arXiv preprint arXiv:2211.09110, 2022.

---

[17] Reza Yazdani Aminabadi, Samyam Rajbhandari, Minjia Zhang, Ammar Ahmad Awan, Cheng Li. Du Li, Elton Zheng, Jeff Rasley, Shaden Smith, Olatunji Ruwase, et al. Deepspeed inference: Enabling efficient inference of transformer models at unprecedented scale. arXiv preprint wwXiv:2207.000:2 2022.
[18] HuggingFace. Hugging face accelerate. http:////gggggggaeccooooooo
index
[19] Ying Sheng, Lianmin Zheng, Binhang Yuan, Zhuohan Li, Max Ryabinin, Daniel Y Fu, Zhiqiang Xie, Beidi Chen, Clark Barrett, Joseph E Gonzalez, al. High-throughput generative inference of large language models with single gpu. arXiv preprint arXiv:2303.06863 2023.
[20] Elias Frantar and Dan Alistarh. Massive language models can be accurately pruned in one-shot. arXiv preprint arXiv:2301 007744 2023.
[21] Mingjie Sun, Zhuang Liu, Anna Bair, and J Zico Kolter. A simple and effective pruning approach for large language models. arXiv preprint arXiv:2306. 1695, 2023.
[22] Lu Yin, You Wu, Zhenyu Zhang, Cheng- Yu Hsieh, Yaqing Wang. Yiling Jia, Mykola Pech- enizkiy, Yi Liang, Zhangyang Wang, and Shiwei Liu. Outlier weighed layerwise sparsity (owl); A missing secret sauce for pruning Ilms to high sparsity, arXiv preprint arXiv:2310.05175, 2023.
[23] Elias Frantar, Salch Ashkboos, Torsten Hoefler, and Dan Alistarh. Gptq: Accurate post-training quantization for generative pre-trained transformers. arXiv preprint arXiv:2210. 221.:22300. 2022.
[24] Guangxuan Xiao, Ji Lin, Mickael Seznec, Julien Demouth, and Song Han. Smoothquant: Accurate and efficient post-training quantization for large language models. arXiv preprint arXiv:2211.10438 2022.
[25] Zhewei Yao, Reza Yazdani Aminabadi, Minjia Zhang, Xiaoxia Wu, Conglong Li, and Yuxiong He. Zeroquant: Efficient and affordable post-training quantization for large-scale transformers. arXiv preprint arXiv:2206.01861, 2022.
[26] Tim Dettmers, Mike Lewis, Younes Belkada, and Luke Zettlemoyer. Gpt3. int8 O: 8-bit matrix multiplication for transformers at scale. In Advances in Neural Information Processing Systems, 2022.
[27] Tim Dettmers, Mike Lewis, Younes Belkada, and Luke Zettlemoyer. Llm. int8 0: 8-bit matrix multiplication for transformers at scale. arXiv preprint arXiv:2208.07339 2022.
[28] Ji Lin, Jiaming Tang, Haotian Tang, Shang Yang, Xingyu Dang, and Song Han. Awq: Activation-awan aware weight quantization for IIm compression and acceleration. arXiv preprint arXiv:2306.0097. 2023.
[29] Joshua Ainslie, Tao Lei, Michiel de Jong, Santiago Ontañón, Siddhartha Brahma, Yury Zemlyanskiy, David Uthus, Mandy Guo, James Lee-Thorp. Yi Tay, et al. Colt5: Faster long-range transformers with conditional computation. arXiv preprint 2rX3v:2303.09752 2023.
[30] Sotiris Anagnostidis, Dario Pavllo, Luca Biggio, Lorenzo Noci, Aurelien Lucchi, and Thomas Hoffmann. Dynamic context pruning for efficient and interpretable autoregressive transformers. arXiv preprint arXiv:2305. 2023.
[31] Yi Tay, Mostafa Dehghani, Dara Bahri, and Donald Metzler, Efficient transformers: A survey. arXiv preprint arXiv:2009.067:3 2020.
[32] Hanrui Wang, Zhekai Zhang, and Song Han. Spatten: Efficient sparse attention architecture with cascade token and head pruning. In 202/ IEEE International Symposium on High- Performance Computer Architecture (HPCA), pages 97-110. IEFE, 2021.
[33] Elizabeth neil, Patrick E neil, and Gerhard Weikum. The Iru-k page replacement algorithm for database disk buffering, Aem Sigmod Record, 22(2):297-306, 1993.
12

---

[34] Donghee Lee, Jongmoo Choi, Jong-Hun Kim, Sam H Noh, Sang Lyul Min, Yookun Cho, and Chong Sang Kim. Lrfu: spectrum of policies that subsumes the least recently used and least frequently used policies. IEEE transactions on Computers, 50(12):1352-136) 2001.
[35] Valerii Likhosherstov, Krzysztof Choromanski, and Adrian Weller. On the expressive power self-attention matrices. arXiv preprint arXiv:2 036443 2021.
[36] Benjamin L. Edelman, Surbhi Goel, Sham Kakade, and Cyril Zhang. Inductive biases and variable creation in self-attention mechanisms. In International Conference on Machine Learning, pages 5793-5831. PMLR, 2022.
[37] Laszlo A. Belady. study of replacement algorithms for virtual-storage computer. IBM Systems journal, journal, 5(2):78-101, 1966.
[38] Simeng Han, Xiang Lin, and Shafiq Joty. Resurrecting submodularity for neural text generation. arXiv preprint arXiv: 191 1,03014, 2019.
[39] Susan Zhang, Stephen Roller, Naman Goyal, Mikel Artetxe, Moya Chen, Shuohui Chen, Christopher Dewan, Mona Diab, Xian Li, Xi Victoria Lin, et al. Opt: Open pre-trained transformer language models. arXiv preprint arXiv:2205.01065 2022.
[40] Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Maric-Ami Anne Lachaux, Timo- thée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, et al. Llama: Open and efficient foundation language models. arXiv preprint wrXiv:2302.13977 2023.
[41] Sid Black, Stella Biderman, Eric Hallahan, Quentin Anthony, Leo Gao, Laurence Golding, Ho- race He, Connor Leahy, Kyle McDonell, Jason Phang. Michael Pieler, USVSN Sai Prashanth, Shivanshu Purohit, Laria Reynolds, Jonathan Tow, Ben Wang, and Samuel Weinbach. GPT- NeoX-20B: An open-source autoregressive language model. In Proceedings of the ACL Workshop on Challenges & Perspectives in Creating Large Language Models, 2022.
[42] Melissa Roemmele, Cosmin Adrian Bejan, and Andrew S Gordon. Choice of plausible alternatives: An evaluation of commonsense causal reasoning. In AAAI spring symposium: logical formalizaion of commonsense reasoning, pages 90-95, 2011.
[43] Aida Amini, Saadia Gabriel, Shanchuan Lin, Rik Sencel-Kedzionste Yejin Choi, and Han- naneh Hajishirzi. MathQA: Towards interpretable math word problem solving with operation- based formalisms. In Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume I (Long and Short Papers), pages 2357-2367, Minneapolis, Minnesota, June 2019, Association for Computational Linguistics.
[44] Todor Mihaylov, Peter Clark, Tushar Khot, and Ashish Sabharwal. Can a suit of armor conduct electricity? a new dataset for open book question answering. In EMNLP, 2018.
[45] Yonatan Bisk, Rowan Zellers, Ronan Le Bras, Jianfeng Gao, and Yejin Choi. Piqa: Reasoning about physical commonsense natural language. In Thirty-Fourrt AAAI Conference on Artificial Intelligence, 2020.
[46] Alex Wang, Amanpreet Singh, Julian Michael, Felix Hill, Omer Levy, and Samuel R Bowman. Glue: A multi-task benchmark and analysis platform for natural language understanding. arXiv preprint 1rxiv:/804.0744 2018.
[47] Keisuke Sakaguchi, Ronan Le Bras, Chandra Bhagavatula, and Yejin Choi. Winogrande: An adversarial winograd schema challenge scale. Communications of the ACM, 64(9):99-106, 2021.
[48] Shashi Narayan, Shay B Cohen, and Mirella Lapata. Don't give me the details, just the summary! topic-aware convolutional neural networks for extreme summarization. arXiv preprint arXiv: /808.08458 2018.
[49] Ramesh Nallapati, Bowen Zhou, Caglar Gulcehre, Bing Xiang, et al. Abstractive text sum- marization using requence-to-sequeeee rnns and beyond. arXiv preprint arXiv: (602006022 2016.
13

---

[50] Xuechen Li, Tianyi Zhang, Yann Dubois, Rohan Taori, Ishaan Guirajani, Carlos Guestrin, Percy Liang, and Tatsunori B. Hashimoto. Alpacaeval: An automatic evaluator of instruction- following models. https://giis httpst///githbb.... www.uub.-oo/ttttt-aaaaaa habb/alpca__eee evall 2023.
[51] Lianmin Zheng, Wei-Lin Chiang, Ying Sheng, Siyuan Zhuang, Zhanghao Wu, Yonghao Zhuang, Zi Lin, Zhuohan Li, Dacheng Li, Eric. Xing, Hao Zhang, Joseph E. Gonzalez, and Ion Stoica. Judging Ilm-as-a-judge with mt- bench and chatbot arena, 2023.
[52] Guangxuan Xiao, Yuandong Tian, Beidi Chen, Song Han, and Mike Lewis. Efficient streaming language models with attention sinks. arXiv preprint arXiv:2309 17453, 2023.
[53] Chi Han, Qifan Wang, Wenhan Xiong, Yu Chen, Heng Ji, and Sinong Wang. Lm- infinite: Simple on-the-fly length generalization for large language models. arXiv preprint arXiv:2308.1611 2023.
[54] Jack W Rae, Anna Potapenko, Siddhant M Jayakumar, and Timothy P Lillicrap.. Compressive transformers for long-range sequence modelling. In The International Conference on Learning Representations (ICLR), 2020.
[55] Song Han, Huizi Mao, and William Dally. Deep compression: Compressing deep neural networks with pruning, trained quantization and huffman coding. arXiv preprint arXiv:1510.00144 2015.
[56] 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 annegr-arithmmeicoonn) inference. In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 2704-2713, 2018.
[57] Markus Nagel, Mart Baalen, Tijmen Blankevoort, and Max Welling. Data-free quantization through weight equalization and bias correction. In Proceedings of the JEEE/CVF International Conference on Computer Vision, pages 1325-1333 2019.
[58] Ritchie Zhao, Yuwei Hu, Jordan Dotzel, Chris De Sa, and Zhiru Zhang. Improving neural network quantization without retraining using outlier channel splitting. International conference machine learning, pages 7543-7552. PMLR, 2019.
[59] Pavlo Molchanov, Stephen Tyree, Tero Karras, Timo Aila, and Jan Kautz. Pruning convo- lutional neural networks for resource efficient inference. arXiv preprint 1rXiv:1611.0644 2016.
[60] Zhuang Liu, Mingjie Sun, Tinghui Zhou, Gao Huang, and Trevor Darrell. Rethinking the value of network pruning. arXiv preprint arXiv: 1810.05270, 2018
[61] Yang He, Ping Liu, Ziwei Wang, Zhilan Hu, and Yi Yang. Filter pruning via geometric median for deep convolutional neural networks acceleration. In Proceedings of the JEFF/CVF conference computer vision and pattern recognition, pages 4340 43000-4349400 2019.
[62] Torsten Hoefler, Dan Alistarh, Tal Ben-Nun, Nikoli Dryden, and Alexandra Peste. Sparsity in deep learning: Pruning and growth for efficient inference and training in neural networks. Mach. Learn. Res., 22(241):1-124 2021.
[63] Geoffrey Hinton, Oriol Vinyals. Jeff Dean, et al. Distilling the knowledge in neural network. arXiv preprint arXiv:1503.0253: 2(7), 2015.
[64] Jang Hyun Cho and Bharath Hariharan. On the efficacy of knowledge distillation. In Pro- ceedings of the IEEE/CVF international conference on computer vision, pages 4794-4802, 2019.
[65] Raphael Tang. Yao Lu, Linqing Liu, Lili Mou, Olga Vechtomova, and Jimmy Lin. Dis- tilling task-specific knowledge from bert into simple neural networks. arXiv preprint arXiv:/903. 12136, 2019.
[66] Hugo Touvron, Matthieu Cord, Matthijs Douze Francisco Massa, Alexandre Sablayrolies, and Hervé Jégou. Training data-efficient image transformers distillation through attention. In International Conference on Machine Learning, pages 10347-10357. PMLR, 2021.
14

---

[67] Ashish Vaswani, Noam M. Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In NIPS, 2017.
[68] Zhilin Yang, Zihang Dai, Yiming Yang, Jaime Carbonell, Russ R Salakhutdinov, and Quoc V
Le. Xlnet: Generalized autoregressive pretraining for language understanding. Advances in neural information processing systems, 32, 2019.
[69] Yinhan Liu, Myle Ott, Naman Goyal, Jingfei Du, Mandar Joshi, Danqi Chen, Omer Levy, Mike Lewis, Luke Zeytlemoyer, and Veselin Stoyanov. Roberta: robustly optimized bert pretraining approach. arXiv preprint arXiv:1907.11693 2019.
[70] Alon Talmor, Jonathan Herzig, Nicholas Lourie, and Jonathan Berant. Commonsenseqa:
A question answering challenge targeting commonsense knowledge. arXiv preprint 2018.
[71] Ajay Jaiswal, Liyan Tang, Meheli Ghosh, Justin Rousseau, Yifan Peng, and Ying Ding. Radbert-cl: Factually-awarr contrastive learning for radiology report classification. Proceed- ings of machine learning research, 158:196-208 2021.
[72] Wei Yang, Yuqing Xie, Aileen Lin, Xingyu Li, Luchen Tan, Kun Xiong, Ming Li, and Jimmy Lin. End-to-end open-domain question answering with bertserini. arXiv preprint arXiv: 1et::0/02000111 2019.
[73] Ming Ding, Chang Zhou, Qibin Chen, Hongxia Yang, and Jie Tang. Cognitive graph for multi-hop reading comprehension at scale. arXiv preprint
2019.
[74] Jason Wei, Xuezhi Wang, Dale Schuurmans, Maarten Bosma, Ed Chi, Quoc Le, and Denny Zhou. Chain of thought prompting elicits reasoning in large language models. arXiv preprint arXiv:2201.11903 2022.
[75] Jingfeng Yang, Hongye Jin, Ruixiang Tang, Xiaotian Han, Qizhang Feng, Haoming Jiang, Bing Yin, and Xia Hu. Hamesssing the power of Ilms in practice: A survey on chatgpt and beyond. arXiv preprint arXiv:2304.1377 2023.
[76] Jacob Devlin, Ming-Wei Wei Chang, Kenton Lee, and Kristina Toutanova, Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint arXiv:1810. 810.04800 2018.
[77] 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. The Journal of Machine Learning Research, 21 (1):5485-55511 2020.
[78] Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever, et al. Language models are unsupervised multitask learners. OpenA/ blog, 1(8):9, 2019.
[79] Tom B Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhari- wal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell et al. Language models are few-shot learners. arXiv preprint arXiv:2005. 14165, 2020.
[80] Teven Le Scao, Angela Fan, Christopher Akiki, Ellie Pavlick, Suzana Ilié, Daniel Hesslow, Roman Castagné, Alexandra Sasha Luccioni, François Yvon, Matthias Gallé, al. Bloom: A 176b-parameter open-access multilingual language model. arXiv preprint arXiv:2211.05100, 2022.
[81] Jingzhao Zhang, Sai Praneeth Karimireddy, Andreas Veit, Seungyeon Kim, Sashank J Reddi, Sanjiv Kumar, and Suvrit Sra. Why (adam) beats (sgd) for attention models, 2020.
[82] Liyuan Liu, Haoming Jiang, Pengeheng He, Weizhu Chen, Xiaodong Liu, Jianfeng Gao, and Jiawei Han. On the variance of the adaptive learning rate and beyond. arXiv preprint AxXiv:900.03320 2019.
[83] Liyuan Liu, Xiaodong Liu, Jianfeng Gao, Weizhu Chen, and Jiawei Han. Understanding the difficulty of training transformers. arXiv preprint arXiv:2004.08249 2020.
15

---

[84] Dušan Variš and Ondřej Bojar. Sequence length is domain: Length- transformer models. arXiv preprint arXiv:2109.0727,, 2021.
Length-based based overfitting in
[85] Wancong Zhang and leshan Vaidya. Mixup training leads to reduced overfitting and improved calibration for the transformer architecture. arXiv preprini arXiv:2102./ 1402, 2021.
[86] Xiaodong Liu, Kevin Duh, Liyuan Liu, and Jianfeng Gao. Very deep transformers for neural machine translation. arXiv preprint arXiv: 2020.
[87] Peng Xu, Dhruv Kumar, Wei Yang, Wenjie Zi, Keyi Tang, Chenyang Huang, Jackie Chi Kit Cheung, Simon JD Prince, and Yanshuai Cao. Optimizing deeper transformers small datasets. arXiv preprint arXiv:2012. 2020.
[88] Chen Zhu, Renkun Ni, Zheng Xu, Kezhi Kong. W Ronny Huang, and Tom Goldstein. Gradinit: Learning to initialize neural networks for stable and efficient training. Advances in Neural Information Processing Systems, 34:16410-16422, 2021.
[89] Jeremy M Cohen, Behrooz Ghorbani, Shankar Krishnan, Naman Agarwal, Sourabh Medapati, Michal Badura, Daniel Suo, David Cardoze, Zachary Nado, George E Dahl, et al. Adaptive gradient methods at the edge of stability. arXiv preprint arXiv:2207.14488. 2022.
[90] Hongyu Wang, Shuming Ma, Li Dong, Shaohan Huang, Dongdong Zhang, and Furu Wei. Deepnet: Scaling transformers 1,000 layers. arXiv preprint urXiv:2203.00555, 2022.
[91] Qiming Yang, Kai Zhang, Chaoxiang Lan, Zhi Yang, Zheyang Li, Wenming Tan, Jun Xiao, and Shiliang Pu. Unified normalization for accelerating and stabilizing transformers. arXiv preprint arXiv:2208.00 01313, 2022.
[92] Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization, arXiv preprint arXiv:1711.05100,,20 2017.
[93] Yaoming Zhu, Sidi Lu, Lei Zheng, Jiaxian Guo, Weinan Zhang, Jun Wang, and Yong Yu. Texygen: A benchmarking platform for text generation models. The 41st international ACM SIGIR conference on research & development in information retrieval, pages 1097-1100, 2018.
[94] Tim Dettmers and Luke Zettlemoyer. The case for 4-bit precision: k-bit inference scaling laws. arXiv preprint arXiv:2212.09770, 2022.
[95] Nelson F Liu, Kevin Lin, John Hewitt, Ashwin Paranjape, Michele Bevilacqua, Fabio Petroni, and Percy Liang. Lost in the middle: How language models use long contexts. arXiv preprint 0rXiv:2307.03172, 2023.
[96] Pranjal Awasthi and Anupam Gupta. Improving length- generalization in transformers via task hinting. arXiv preprini arXiv:2310.00726 2023.
[97] Amir Zandieh, Insu Han, Majid Daliri, and Amin Karbasi. Kdeformer: Accelerating trans- formers via kernel density estimation. arXiv preprint arXiv:2302.02451, 2023.
[98] Josh Alman and Zhao Song. Fast attention requires bounded entries. arXiv preprint S.Ii::2302.122144 2023.
[99] Yichuan Deng, Zhao Song, and Tianyi Zhou. Superiority of softmax: Unveiling the perfor- mance edge over linear attention. arXiv preprint arXiv:2310. /685, 2023.
[100] Jan van den Brand, Zhao Song, and Tianyi Zhou. Algorithm and hardness for dynamic attention maintenance in large language models. arXiv preprint arXiv:2304.02207 2023.
[101] Yeqi Gao, Zhao Song, and Xin Yang. Differentially private attention computation. arXiv preprint 20X.v:2305.04770 2023.
[102] Clayton Sanford, Daniel Hsu, and Matus Telgarsky. Representational strengths and limitations of transformers, arXiv preprint arXiv:2306.022 028966 2023.
16

---

[103] Yichuan Deng, Sridhar Mahadevan, and Zhao Song, Randomized and deterministic attention sparsification algorithms for wver-parameterized feature dimension. arxiv preprint: arxiv 2304.03426, 2023.
[104] Zhihang Li, Zhao Song, and Tianyi Zhou, Solving regularized exp, cosh and sinh regression problems. arXiv preprint, 2303. 15725, 2023.
[105] Yichuan Deng, Zhihang Li, and Zhao Song. Attention scheme inspired softmax regression. arXiv preprint arXivv2304. 2023.
[106] Praneeth Kacham, Vahab Mirrokni, and Peilin Zhong. Polysketchformer: Fast transformers via sketches for polynomial kernels. arXiv preprint arXiv:2310.0166 2023.
[107] Yeqi Gao, Zhao Song, and Junze Yin. An iterative algorithm for rescaled hyperbolic functions regression. arXiv preprint wrxiv:2305.0066/ 2023.
[108] Insu Han, Rajesh Jarayam, Amin Karbasi, Vahab Mirrokni, David P Woodruff, and Amir Zandieh. Hyperattention: Long-context attention in near-linear time. arXiv preprint arXiv:2310.05869 2023.
[109] Timothy Chu, Zhao Song, and Chiwun Yang. How protect copyright data in optimization of large language models? arXiv preprint arXiv:2308. 0 12247, 2023.
[110] Ritwik Sinha, Zhao Song, and Tianyi Zhou. A mathematical abstraction for balancing the trade- off between creativity and reality in large language models. arXiv preprint vrXiv:2306.02299 2023.
[111] Yeqi Gao, Zhao Song, Weixin Wang, and Junze Yin. A fast optimization view: Reformu- lating single layer attention Ilm based on tensor and svm trick, and solving matrix multiplication time. arXiv preprint arXiv:2309.074/1 2023.
[112] Gary Marcus, Ernest Davis, and Scott Aaronson. A very preliminary analysis of dall-e 2. arXiv preprint arXiv:2204. 13807, 2022.
[113] Yeqi Gao, Zhao Song. Xin Yang, and Ruizhe Zhang. Fast quantum algorithm for attention computation. arXiv preprint arXiv:2307.0880 000400045, 2023.
[114] Puncesh Deora, Rouzbeh Ghaderi, Hossein Taheri, and Christos Thrampoulidis. On the optimization and generalization of multi-head attention. arXiv preprint arXiv:2310.12680, 2023.
[115] Yeqi Gao, Zhao Song, and Junze Yin. Gradientcoin: A peer-to-peer decentralized large language models. arXiv preprint 2301::2308.10002, 2023.
[116] Yichuan Deng, Zhao Song, Shenghao Xie, and Chiwun Yang. Unmasking transformers: A theoretical approach to data recovery via attention weights. arXiv preprint arXiv:2310.124: 12462, 2023.
[117] Yichuan Deng, Zhihang Li, Sridhar Mahadevan, and Zhao Song. Zero-th order algorithm for softmax attention optimization. arXiv preprint urXiv:2307.08352. 2023.
[118] Josh Alman and Zhao Song. How capture higher-order correlations? generalizing matrix softmax attention to kronecker computation. arXiv preprint arXiv:2310.04066 2023.
[119] Yichuan Deng, Zhao Song, and Shenghao Xie. Convergence of two-layer regression with nonlinear units. arXiv preprint arXiv:2308.08333 2023.
[120] Timothy Chu, Zhao Song, and Chiwun Yang. Fine-tune language models to approximate unbiased in-context learning. arXiv preprint arXiv:2310.03333 2023.
[121] Abhishek Panigrahi, Sadhika Malladi, Mengzhou Xia, and Sanjeev Arora. Trainable trans- former in transformer. arXiv preprint arXiv:2307.00 00 189, 2023.
[122] Haoyu Zhao, Abhishek Panigrahi, Rong Ge, and Sanjeev Arora. Do transformers parse while predicting the masked word? arXiv preprint arXiv:2303.08117, 2023.
17

---

[123] Alexander Schrijver. Combinatorial optimization: polyhedra and efficiency, volume 24. Springer, 2003.
[124] Kasper Green Larsen, Jelani Nelson, Huy L Nguyen, and Mikkel Thorup. Heavy hitters via cluster-preserving clustering. In 2016 IEEE 57th Annual Symposium on Foundations of Computer Science (FOCS), pages 61-70. IEEE, 2016.
[125] Vasileios Nakos and Zhao Song. Stronger 12/12 compressed sensing: without iterating. In Proceedings of the 51st Annual ACM SIGACT Symposium on Theory of Computing, pages 289 297, 2019.
[126] Vasileios Nakos, Zhao Song, and Zhengyu Wang. (nearly) sample-optimal sparse fourier transform in any dimension; ripless and filterless. In 2019 IEEE 60th Annual Symposium Foundations of Computer Science (FOCS), pages 1568-1577. IEEE, 2019.
[127] Andreas Krause and Carlos Guestrin. Beyond convexity: Submodularity in machine learning. ICML Tutorials, 2008.
[128] Jeff Bilmes. Submodularity in machine learning applications. In 7wenty-Ninth Conference on Artificial Intelligence, AAAI-15 Tutorial Forum, 2015.
[129] Simeng Han, Xiang Lin, and Shafiq Joty. Resurrecting submodularity for neural text generation. arXiv preprint arXiv:1911.03014 2019.
[130] George L. Nemhauser, Laurence A Wolsey, and Marshall Fisher. An analysis of approxima- tions for maximizing submodular set functions -i. Mathematical programming, 14(1):265- 294, 1978.
[131] Lianke Qin, Zhao Song, and Yitan Wang. Fast submodular function maximization. arXiv preprint arXiv:2305. 08367, 2023.
[132] Junda Wu, Tong Yu, Rui Wang, Zhao Song, Ruiyi Zhang, Handong Zhao, Chaochao Lu, Shuai Li, and Ricardo Henao. Infoprompt: information-theoreet soft prompt tuning for natural language understanding. arXiv preprint arXiv:2306.0492: 2023.
[133] Shuai Li, Zhao Song, Yu Xia, Tong Yu, and Tianyi Zhou. The closeness of in-context learning and weight shifting for softmax regression. arXiv preprint, 2023.
[134] Michael B Cohen, Yin Tat Lee, and Zhao Song. Solving linear programs in the current matrix multiplication time. In STOC, 2019.
[135] Yin Tat Lee, Zhao Song, and Qiuyi Zhang. Solving empirical risk minimization in the current matrix multiplication time. In Conference Learning Theory, pages 2140-2157. PMLR, 2019.
[136] Haotian Jiang, Tarun Kathuria, Yin Tat Lee, Swati Padmanabhan, and Zhao Song. A faster interior point method for semidefinite programming. In 2020 IEEE 6lst annual symposium on foundations of computer science (FOCS), pages 910-918. IEEE, 2020.
[137] Zhao Song and Zheng Yu. Oblivious sketching- central path method for linear pro- gramming. In International Conference on Machine Learning, pages 9835-9847. PMLR, 2021.
[138] Shunhua Jiang, Zhao Song, Omri Weinstein, and Hengjie Zhang. A faster algorithm for solving general 1ps. In Proceedings of the 53rd Annual ACM SIGACT Symposium on Theory of Computing, pages 823-832, 2021.
[139] Baihe Huang, Shunhua Jiang. Zhao Song. Runzhou Tao, and Ruizhe Zhang, Solving sdp faster: A robust ipm framework and efficient implementation. 2022 IEEE 63rd Annual Symposium Foundations of Computer Science (FOCS), pages 233-244 IEEE, 2022.
[140] Yuzhou Gu and Zhao Song. arXiv:2211.0003,, 2022.
faster small treewidth sdp solver. arXiv preprint
18

---

[141] Yuzhou Gu, Zhao Song, and Lichen Zhang. A nearly-linear time algorithm for structured support vector machines. arXiv preprint arXiv:2307.07735, 2023.
[142] Lianke Qin, Zhao Song, Lichen Zhang, and Danyang Zhuo An online and unified algorithm for projection matrix vector multiplication with application to empirical risk minimization. In AISTATS, 2023.
[143] Zhao Song, Mingquan Ye, and Lichen Zhang. Streaming semidefinite programs: O(vn) passes, small space and fast runtime. arXiv preprint arXin:2300.055/3, 2023.
[144] Haotian Jiang, Yin Tat Lee, Zhao Song, and Lichen Zhang. Convex minimization with integer minima in Onn*) time. In ACM-SIAM Symposium on Discrete Algorithms (SODA), 2024.
[145] S Cliff Liu, Zhao Song. Hengjie Zhang, Lichen Zhang, and Tianyi Zhou. Space-efficient interior point method, with applications to linear programming and maximum weight bipartite matching. In ICALP. arXiv preprint arXiv:2009.061 2020.
19

---

Appendix
Table of Contents
A More Implementation Details
21
B Extended Related Works, Discussions, and Limitations
22
Extended Related Works
22
Discussions and Limilations
22
C
Extended Experiments
24
C.I
Increased Diversity of Generated Text
24
Throughput Improvement of H2O Combining with Quantization
25
C
Effectiveness on Zero-shot and One-shot Inference
26
C.4
Comparison with StreamingLLM
26
Enhancing the "Top-K" Baseline
27
Separate Effects of Heavy- and Local Tokens
27
C.7
Performance of Inference with Different Number of Shots
................................
27
C88
Heavy Hitter Attention Blocks
28
C.9
Comparsion with SpAtten
28
C. 10 Heavy-Hitter in MLP Blocks
28
D
Theoretical Analysis
33
Notations
................................ ................................
Submodular
................................
34
D.3
Dynamic Submodular
35
D7.4
Static Attention
35
D15
Recursive Atlention Definition
35
D.6
Eviction Policy
37
D.7
Explaining Submodular Diminishing Return Property in Attention Scheme
38
Submodular: High Level Ideas]
39
D.9
Robust Greedy with Error Propagation
39
D.10
Robust Submodular and Adding Items
40
D.
Universal Conditions]
41
D
12
Induction Lemma for Exact Function
41
D.13
Induction Lemma for Approximate Function
42
14
heoretical Resull
43
D. 15
Extended Related Work for Theoretical Attention Probems
44
D.16
Sparsity Preserving
44
D. 17 Definition of Loss Function
45
D. 18 Gradient
46
D.19
Hessian
47
D.20 Hessian is Positive Definite
47
D.21 Hessian is Lipschitz
48
D.22 Greedy Type Algorithm
49
20

---

A More Implementation Details
In this section, our goal is to provide the details of system implementation (mentioned in Section 4.2) and experiment settings (mentioned in Section 5). well as the pseudocode.
System Details. We implement H2O on top of FlexGen. FlexGen a white-box implementation of OPT models, and we have done some surgery on handling the KV cache. Specifically, for the given parameter K, we always maintain a list of KV cache with the first K entries as heavy hitters, and the last K entries most recent tokens. In order avoid data movement in memory, the memory for KV cache is preallocated. We use a circular queue to update the last K entries efficiently.
Experiment Details. Our study involves the evaluation with varying sizes of KV cache that encom- pass 4%, 10%, 20%, and 60% of the prompt's 's length. We select two tasks (XSUM and CNN/Daily Mail) from the HELM framework [16] and present the performance based on 1000 test samples. Additionally, we employ the Im-eval-harness framework [[5] for six other tasks (COPA, MathQA, OpenBookQA, PiQA, RTE, and Winogrande). For the tasks derived from the HELM framework, report the performance for zero-shot inference, while for the tasks from Im-eval-harness, we default to conduct five-shot inference.
Pseudocode. We show a simplified pseudocode below to demonstrate our implementation skele- ton. The function generation_ _) is the base loop in FlexGen that controls prefetch and overlap the I/O streams and computation. Then in function compute _areee___err__ Q), the function attention attention_forrardd () will be called for attention layers. During prefill iterations, the function compute_ attention on () will return K heavy hitters and K recent tokens if the prompt has a length compute_ttentioon
compute_attentil aat.tntitton () greater than equal to 2K, otherwise return all the KV cache. The function compute
will return new KV cache and the indices for evicted entries during decoding iterations, which would be the place implement customized eviction strategy. (If the current number of stored KV cache is less than 2K, the eviction will be ignored.) During each decoding iteration, the oldest one among the last K tokens (if we have no less than K tokens' KV cache stored) will be removed from the reserved last K entries, one heavy hitter will be evicted for each head, and the newest token will be added to the KV cache with a position the last K entries. This happens in the function store_ store_cachel ().
def generation_les _oop (...):
# Prologue
# Generate
for i in range
for j in range (nuu_layers): for k in range hww.aagpuaaaccheeeeeeeeeee load_ weight jt1, k) cache (i, j. k+1)
store Soore_hiddee(( j, k-1)
foad_hiddent (i, J. k+1)
conpute_layer (1, k)
store_cache b. k-1) sync()
# Epilogue
h the hidden states (activations)
def irtention_forvaree (h, ...):
# the read/vrite buffer are intermediate stops for prefetching if prefill:
h, new_k_cache, neche_cache select K heavy hitters and recent tokens
compute_attentione (h, ...)
new_k new_k_cache, new_v_cache = select (new_ new_v_cache, K) store (neu_k_cache. www.sttorreee..cccco.
new_v_cache)
else:
cache v_cache cache_read_but pop
evict_ids track the entries that will be evicted h, nev_k_cache, nev_v_cache, evict_ids =
compute (t k_cache, v_cache, ...)
21

---

store (nev_k_ nex_v_cache, evict_ids)
return h
def (ocre_cache(...)
if prefill:
store cache directly
else:
k_new, v_new, evict_ids wtttbbb pop
# circular queue for the last к entries extract the index for the oldest token i-th iteration
oldest ((i 1) k)
update the KV cache (k_home and v_home)
tache_replae (e (k_home, evict_ids: ids, k_new, K, oldest) cache replaceev.oove (v_home, evict_ids, v_new, K, oldest)
B Extended Related Works, Discussions, and Limitations
The goal of this section is to first introduce more background and related works for Section 2 then describe some previous attempts in experiments well as discuss the social impact and limitations this work.
B.1 Extended Related Works
Quantization, Pruning, Distillation for Inference. Previously, model compression algorithms have been extensively investigated as a viable approach for mitigating the computational resource requirements of model inference. These algorithms be broadly categorized into three groups: (1) quantization [55] 56, 57 (58), which involves mapping model parameters or activations from high-precision data types low-precision counterparts, such using 8-bit integers instead of the which aims
commonly employed 32-bit floating point format; (2) pruning or sparsity [59 60 61 62
to eliminate unnecessary neurons weights within the models; (3) and distillation
63
64
65 [66
where predictions from larger models are utilized as supervised information to train smaller models.
Transformer in NLP. Transformers 667 a popular option have been frequently adopted by
plenty of natural language processing (NIP) applications with prevailing successes [68] 69 70, 71 722
46
Roughly, modern transformer-basee networks be categorized into groups: (1) Encoder-Decoder or Encoder-only (i.e., BERT-style models (76)). This type of transformers commonly leverages the Masked Language Modeling task which encourages models capture the intrinsic relationship between words and their context. Notable examples include BERT [76). RoBBERTa [69] and TS [77]. (2) Decoder-only (i.e., GPT-style models (7)). Usually, this group of transformers adopts the Casual Language Modeling task, which is optimized to generate the next word/token sequence based the preceding words/tokens. Such an autoregressive manner is highly preferred by downstream tasks like text generation and question answering. GPT-3 [79]. OPT MaI.M 13). and BLOOM are representative architectures within this huge family.
73
3
74 75,
Training of Transformer. Training gigantic transformer- based model is not trivial. It notoriously suffers from various issues such overfitting, instability, etc. [8] 82 83] analyze these bottlenecks
from the optimization perspective. To address the issues, great amount of pioneering effort is
devoted, including data augmentations optimizers [89), improved normalization there is still a long way to go before we can fully clarify the mystery of transformer training.
85 better initialization [6
87 88), customized
[90, 11], weight decay (92), and early stopping. However,
B.2 Discussions and Limitations
Previous Attempts. During our experiments, we find several noteworthy observations. In HzO, employing the accumulated attention score to evict KV embeddings can lead to a potential bias favoring the least recent tokens. This bias arises because most previous tokens have higher number of attention scores, resulting in a higher accumulated attention score and, consequently, greater likelihood of being retained. To address this concern, conducted additional experiment utilizing
22

---

the averaged attention score to determine which KV embeddings should be retained. However, this alternative approach resulted in performance degradation. Additionally, we observed a significant proportion of H2 occurrences at the beginning of sentences. This finding suggests that the initial tokens play substantial role in subsequent generation tasks.
Social Impact. Our work represents an initial effort in designing a KV Cache policy, a realm that has been relatively unexplored and yet a significant bottleneck LLMs. The proposed Heavy Hitter Oracle (H2O) provide a solution to improve the efficiency of LLM generation, which can save energy cost and contribute to green AI. Besides, our approach also serves as source inspiration for future advanced algorithm designs. We envision H2O as a foundational framework that could facilitate further innovation in this area. Moreover, long content generation an area of growing importance that currently grapples with several efficiency issues. We hope work that supports the generation of very long sequences will support further research this direction, particularly in terms of enhancing consistency, devising superior evaluation methods, and establishing robust benchmarks.
Furthermore, another contribution of this study the formulation of dynamic submodular frame- work. We believe that this theoretical framework possesses the potential to be applicable beyond specific domains of interest. For instance, there may exist numerous other dynamic problems where the task involves solving submodular problem with slight variations at each time.
Limitations. Furthermore, despite the notable advancements in throughput of our H2O, implement- ing LLMs for generative inference remains challenging due to the immense parameter count. As substantial portion of these parameters is encompassed within the MLP blocks, building upon our observations of H2 occurrences in the MLP blocks, future research efforts can be directed towards leveraging the characteristics of H2 to devise an offloading policy. Such a policy can potentially enhance the efficiency of LLM inference even further.
23

---

C Extended Experiments
In this section, our goal is to demonstrate that HzO can improve the diversity of generated text (mentioned in Section 5.3) throughput is further improved when combined with quantization (men- tioned in Section 5.3) and superior performance to handle extremely long length inputs (up to four middle tokens, mentioned in Section 5.3). Moreover, additional investigations about H2 are reported, including experiments under zero-shot/one-sho inference regime; H2O can also enhance the "Top-K" baseline; extra results of H2 in attention blocks; the emergence of H2 in MLP blocks as well as its properties.
C.1 Increased Diversity of Generated Text
Model Imput
In the year 2087, humanity has achieved remarkable technological advancements and established colonies multiple plancts within the Milky Way galaxy. Interstellar travel has hecome commonplace, with faster-than-light spacecraft enabling people to explore distant star systems. Earth has undergone significant changes due to sustainable development efforts, such harnessing renewable energy sources and implementing widespread ccological restoration projects. However, alongside these triumphs, new challenges have emerged, including the rise of artificial intelligence, cthical dilemmas surrounding genetic enginecring, and interplanctary political tensions. Against this backdrop, team of intrepid scientists embarks ancient alien civilization. hidden deep within an uncharted exoplanct. As they navigate treacherous terrains and encounter otherworkdly phenomena, they must confrent their own fears reconcile humanity's 's thirst for knowledge with the potential consequences of uncovering sccrets that were better left buried. The fate of both their mission and the future of humanity hang in the balance.
mission uncover the secrets of an
OPT-6, 7B Full Cache Output
The game 2087, few years after the events of the first game. The player takes the role of team of scientists who are sent to an uncharted planet to investigate a mysterious signal. The team is composed of pilot, a scientist, a mechanic, and technician.
OPT-6. 7B Locul 20% Cache Output
OPT-6.7B H2O 20% Cache Output
The game is set in the year 2087, and is a first-person exploration game. The player takes the role of scientist who has been to distant planct investigate mysterious signal. The planct is inhabited by a intelligent beings known as the "Titans," who have been
Figure 6: Visualization of one generation example with OPT-6.7B 7B. Results are compared between the baseline model with full cache, our HzO, and the "Local" strategy that utilizes the most recent KV embeddings.
Given the same prompt text, we visualize the generated text with OPT-6. 7B and LLaMA-7B across different methods, including the baseline model with full cache, our H2O, and the "Local" strategy. Results are reported in Figure and 7 Even with less KV cache budget, our H₂O can generate more diverse content. Specifically. with the OPT-6. 7B, the full model generates some redundant works, like "a few years after the events of the first game" while our H2O describes "the game is a first-person exploration game". As a comparison, when all KV cache budget is assigned to the most recent tokens, the model fails generate meaningful text and only repeats the word "." and ".". Similar observations can also be drawn from the results of LLaMA-7B, in which the full model repeatedly says "so moving that", "so moved that", and "began to cry" while our H2O describes both the people and the environment.
Moreover, We conducted quantitative comparison via diversity metric(SSlf-BELU 93 We randomly sample 100 instances from the XSUM dataset, as prompts. Subsequently, we employed the LLaMa-7B to generate text of equal length. The full model reaches a Self-BELU score of 0.0057 while H2O and the local method achieve 0.0051 and 0.0436, respectively. The comparatively lower Self-BELU score of H2O indicates the slightly increased diversity of generated text.
24

---

Model Input
In small, bustling cafe nestled in the heart of vihrant city, serendipitous event unfolded, leaving lasting impression on all who witnessed it. As the patrons sat sipping their coffees and engaging in animated conversations, talented street musician entered the cafe, carrying weathered guitar and radiating aura of creativity.
LLaMA-7B Full Cache Output
Hc began to play, and the patrons were captivated. The musician's 's performance was so moving that the patrons began applaud, and the musician was so moved that he began to cry. The patrons were so moved that they began to cry, and the musician was so
LLaMA-7B Locul Cache Output
He (e(CCCC((((((( D) 1999, 2000, 2001, 1, and, and, and, and, and, and, and, and, and, and, and,
LLaMA-7B O₂O 20% Cache Output
hegan to play, and the room was filled with the sound of his music. The patrons of the cafe were enthralled by the music, and the atmosphere was electric. The cafe was packed with people, all of whom were enjoying the music. The musician young
Figure Visualization one generation example with LLaMA-7B. Results are compared between baseline model with full cache, our HzO, and the "Local" strategy that utilizes the most recent KV embeddings.
Table 6: Compatibility with Quantization.
Models
COPA
OpenBookQA
PIQA
Full
85.00
43.20
78.51
H2O
84.00
43.00
78.45
Quant-4bit
84.00
43.28
78.67
H2O w. Quant-4bit 84.00
43.20
78.80
C.2 Throughput Improvement of H₂O Combining with Quantization
Table 6 demonstrates H₂O are capable of quantization and perverse the full model's performance with even 4-bit quantization. To further explore the compatibility of our H2O with quantization with respect to throughput improvement, we conduct an additional evaluation with the quantization method implemented in FlexGen (Note that (94] employed different 4-bit quantization method). The corresponding results are presented in Table 7 Notably, for OPT-6.7B, we observed extra performance enhancements in H₂O when utilizing quantization compared to the vanilla version. This improvement results from the GPU memory freed by weight quantization, which allows for significant increase in batch size. However, it should be emphasized that the quantization method employed in FlexGen is not implemented most efficiently, resulting in considerable computational overhead. Despite the batch size being enlarged by 20 times, the actual throughput improvement is less than 2 times. Nevertheless, it is important to acknowledge the potential benefits of combining H2O with quantization, as exemplified by the ability to increase the batch size further. For instance, the implementation of 4-bit quantization could be accelerated by optimized CUDA kernel.
Table 7: Generation throughput (token/s) with different systems without offloading. We use H2O-c denotes the H₂O with 4-bit weights compression. In the sequence length use "512 32" to denote prompt length of 512 and generation length of 32. "OOM" means out-of-memory. The gray text in the bracket denotes the batch size. run OPT-6.7B on single T4 GPU.
Seq. length
512+32
512+512
512+1024
Model size
6.7B
6.7B
6.7B
Accelerate
20.4 (2)
15.5 (1)
OOM
DeepSpeed
00M
00M
00M
FlexGen
20.2 (2)
16.8 (1)
16.9 (1)
H₂O (20%)
35.1 (4)
$1.7
52,1 (4)
O2O-c (20%)
50.5 (70)
72.5 (52)
62.3 (44)
25

---

C.3 Effectiveness on Zero-shot and One-shot Inference
When facing zero-shot and one-shot inference tasks, H2O can successfully mitigate the memory requirements of the KV cache by up to 5x under both one/zero-shot inference across different tasks, achieving matching performance as the model with full KV cache while the local method fails. However, unlike 5-shot, we also found that some tasks are more difficult, and the model requires more information to generate the correct content, resulting in a higher KV cache budget (30-40%) This might be expected since 0/1 shots in our benchmarks have shorter sequence lengths ranging from 100 to 300.
PIQA, LLaMA-7B
COPA, LLaMA-7B
OpenBookQA, LLaMA-7B
45.0
75
42.5
75
70-
40.0
1-shot
70-
Bleavy-Ifitte Oracle
65
37.5
Local
65
60
32.5
60
55-
Local
Heavy Hitter Oracle
50.0
Heavy-Hitter Oracke
Local
27.5
50-
Full
55
25.0
100
80
60
50
40
30
20
10
100
80
60
50
40
30
20
10
100
50
50
50
40
30
20
10
KV Cache Budget (%)
KV Cache Budget (86)
Cache Budget
80-
42.5
75
40.0
75
70-
57.5
O-shot
70
65
Heavy- Hiner Oraele
55.0
65-
60-
Fail
Local
50.0
52.5
60
Hesvy-Hiter
55-
27.5
Heavy-Hiitter Oracic
55
Full
Local
50-
25.0
22.5
Full
Local
45
40
20
100
80
60
40
20
10
100
80
50
40
50
20
10
Cache Budget (%)
KV Clche Budget (36)
KV Cacho Budget
(%)
Winogrande, LLaMA-7B
MathQA, LLaMA-7B
XSUM, LLaMA-7B
70
20-
65
1-shot
Accuracy
26-
Rougo L
Herv Hiteer Oraele
60
24-
Local
Full
55
Hotty-Hitter Oracle
Heavy Hister Oracle
50
Local
Full
18-
Local
5-
16-
50
20
100
50
60
50
40
30
20
10
100
80
50
50
40
30
20
10
KV Cache Budget (%)
XV Cache Budget (95)
KV Cache Budget (56)
70
30
10
65
O-sbot
Accurscy
60
Heavy Titter Oracle
Local
24
Rosgee
6-
8-
Local
Heavy-Hitte Oracle
55
Full
20-
Heavy-Hille Ocacle
Full
18-
Local
Full
100
50
30
20
10
100
30
60
40
30
20
10
100
30
10
Cache Bragget (%))
KV Cache Budget (%))
KV Cache Budget (%)
Figure 8: Comparsion results on zero-shot and one-shot inference.
C.4
Comparison with Streaming
Recent works, such as StreamLLM [2 and LM-Infinite [53], have shown promising potential in enabling Language Models (LLMs) to handle input of infinite length. They achieve this by only retaining the initial tokens and a limited local context. However, this approach may pose challenges for certain tasks where vital information lies within the middle of the input and would be lost using this strategy. We investigate it through two specific types of tasks:
26

---

Multi-document question answering [95]: In this task, each test sample comprises ten documents, followed by a question. The key information answer this question is stored in one of the documents. We rearranged the crucial document's 8 position and found that the eviction strategy in StreamI.LM or IM-Infinite can not perform well when the key document has been dropped.
Text Summarization Task (XSUM and CNN-DailyMai)) Text summarization tasks require models generate concise summaries of lengthy texts. Effective summarization demands comprehensive understanding of the entire document, making it challenging when crucial information is dispersed throughout the input. particular, summarization often relies on long-conten, attention, and critical information may not be effectively captured within the limited local tokens.
The results reported in Figure illustrating consistent decline in the performance of StreamL Since Streaml .IM will always maintain the first few tokens as well as the local tokens, regardless of various input content, such strategy will inevitably result in the loss of crucial information and subsequently lead to a decrease in performance. contrast, our H2O delivers markedly superior performance. Also, with H2O, The model can successfully stream four million tokens.
10- Document Question Answering
Summarization Task
StreamL LM-45000
0.25
StreumL.LM-4-252
1--2566266
H2O-128-1 128
40
H20-256-256
EM Score
Rouge-L
0.15
20
0.05
0.00
1th
2nd Answer Locations
3rd
XSUM
CNN-DailyMail
Figure 9: Comparison results of StreamLLM 32 and H₂O generization tasks. The number in each method represents the KV Cache budget the start/heavy-hitter tokens the local tokens, respectively. example, H20-256-250 means maintaining 256 Heavy-Hitter and 256 local tokens.
C.5 Enhancing the "Top-K" Baseline
We find H2 can further enhance another strong baseline with "Top-K" strategy. The results are reported in Table & After combing with H2, the "Top-K" method achieves an extra improvement with up to 2.00% accuracy across 4 different tasks.
Table Results of the "Top-K" method or w.o. H2. Experiments are conducted with OPT-30B with 20% KV cache budget.
Models
COPA
OpenBookQA
PiQA
Winogrande
Full
85.00
43.20
78.51
70.24
TopK w.o. H2
80.00
41.40
76.96
65.35
TopK w. H2
82.00
42.80
77.96
66.48
C.6 Separate Effects of Heavy-Hitter and Local Tokens
Table 9 demonstrates the separate effects of Heavy- -itters and the local tokens. And we can observe Seavy-Hitters contribute more to maintaining the performance of models.
C.7 Performance of Inference with Different Number of Shots.
Table 10 demonstrates our HzO achieves consistent improvements across different number of shots during inference and maintains the full model's performance with 20% memory budget.
27

---

Table Ablation study H2Q across different tasks.
Tasks
Models
Full
Local
w:
w: Local
OPT-13B
77:37
54.62
76.12
77.26
PIQA
OPT-30B
78.51
55.82
67.25
78.45
OPT-13B
41.40
25.60
30.40
41.20
OpenBookQA
OPT-30B
43.20
25.20
26.60
43.00
OPT-13B
26.67
22.04
23.82
26.93
MatQQA
OPT-30B
26.23
20.87
21.98
26.87
OPT-138
G8. 59
49.96
G1.85
67.32
Winogrande
OPT-30B
70.24
49.17
47.36
69.06
Table 10: Results under different sequence length of OPT-30B with 20% KV cache budget.
Tasks
Methods
OPT-308
Snabots
DPL-638
OPT:300
10-ahots
OPT-GGB
Full
43.20
44.40
43.00
44.80
OpenBookQA
Local
25.20
50.60
26.60
38.80 35.80
COPA
Local
48.00
59.00
60.00
76.00
HgO
84.00
92.00
85.00
86.00
Full
26.25
27.87
26.
27. 27.00
MathQA
Local
20.57
25.40
21 11
25. 08
H₂O
27.67
26.47
27,30
C.8 Heavy-Hitter in Attention Blocks
The distribution of accumulated attention scores of all the tokens within attentions blocks is illustrated in Figure I0 We can observe that H2 broadly exists in each layer.
C.9
Comparsion with SpAtten
Compared with SpAtten 32 the main differences of H2O are i) They accumulate attention scores across attention heads and layers, while in our algorithm, each token can be kept evicted indepen- dently across heads and layers, providing more flexibility for selecting critical KV embeddings; ii) We also use KV of the most recent tokens during generation and demonstrate that such H2 tokens can effectively enhance other sparse attention strategies; iii) we formulate the KV cache eviction as dynamic submodular problem and prove a theoretical guarantee (under mild assumptions) for our novel algorithms. Moreover, we also provide a quantitative comparison with SpAtten, and the results are reported in Table
Table 11: Comparison between SpAtten
32
and H₂O across various tasks with OPT-30B.
Models
COPA
OpenBookQA
PiQA
Full
85.00
43.20
78.51
SpAtten
82.00
41.90
77.06
H₂O
84.00
43.00
78.45
C.10 Heavy-Hitter in MLP Blocks
Besides the attention blocks, the presence of Heavy-Hitten (H2) is observed within the MLP blocks of LLMs. We utilize the Wiki- Miki-Text-Iie 103 dataset as the input and record the activated frequency of neurons in the hidden layer of MLP blocks. As depicted in Figure II] the activated frequency of neurons follows a plwer-law distribution, wherein a small number of neurons are activated by nearly all input tokens (with 100% frequency) while the majority of other neurons are rarely activated.
Subsequently, a thorough examination of various characteristics pertaining H2 in MLP blocks is conducted, encompassing the following aspects: (1) The elimination of H2 leads to substantial decline in performance, although such degradation can easily recovered even with a mere 1% of the training data; (2) H2 exhibits significant degree of overlap across different type of input content; (3) The emergence H2 occurs early the training process, thus exhibiting an "irrdy- characteristic, and their positions undergo gradual changes during subsequent training phases.
28

---

Lasa
Lasri
Laye3
Leyo4
1.5
Alcestiadd Stare
2.0
1.2
1.55
0.5
10
19060
20660
50060 4206
Laver
Lawer
Layerh
2.5
2.5
3.0
2.0
1.5
1.3
0-0
8M
0.
10060
2000 3000 43033 30302
10030 20003 30065 40600 50000
10000 20003 30000 42000 20000
10000 20000 30000 4306 50300
85-
0.0
10060
2000
2000 4303) 50301
10030 21603 70060 46610 50000
10000 20300 30030 42000 58000
18066 20000 70060 4306 50330
4.3
2.0
2.5
10
11860 20000 43031 50333
2160)
(0000) 2000) 4300) 18000
18586 2000
Lays:
Laye 15
Layer
Laye:
Actstelie (tadii Sunco
Lager:
Layer 33
10
Laye:
Wadlaaa
Were lnke
Wedledos
Figure 10: The distribution of accumulated attention scores with respect the corresponding word. The x-axis represents the word index in the vocabulary, and the y-axis represents accumulated attention score. Results are obtained from OPT.L.3B.
29

---

48-
Laye
Lxpes
1ayer
Lone
28
150 930 P10
2001182 1010017911
2100 98 11 1EEIomE 250 1000
259 018 (100) VIE
12-
25
Lave:
Lepor 15
12
20
158 500 7500 210110012220000 23255330
299118E 1011002515
2300 10 100E 255
250 11001001111000111
9
02
03
Layer
Laper
100
300
100
25
30-
226 5000 250110002223000
250 211000001000020055000
Propeme EN)
02
03
75
20
228 500 2500 00000000000000 200.10002200000 250200 Nomea
2000286 220010
250 11
Kraam
Figure 11: The emergence of H2 in MLP blocks OPT-6.7B. The x-axis represents the index of neurons in the hidden layers of MLP blocks, and the y-axis represents the activated frequency.
30

---

Elimination of H2. We first train a GPT-2 using Wiki-Text-103 dataset and subsequently identify and prune the neurons exhibiting an activation frequency exceeding 20% (i.e., H2). This pruning operation leads to substantial decline in performance, from 19.32 to 31.78. The results emphasize the criticality of H2 preserving the functionality the model. To assess the recoverability of the discarded information, we conduct few-shot fine-tuning experiment, and the results are summarized in Table 12 The pruned model is fine-tuned with varying ratios of training data for 500 iterations, and it successfully regains performance levels equivalent to those the pre-trained model. In contrast, when training the model from scratch using only 1% of the training data, the resulting model achieves perplexity of 554.12 only. These findings demonstrate that the knowledge encoded in H2 can be easily restored.
evidenced by increase in perplexity
Table 12:
Perplexity on the test-set of Text- with GPT-2.
Settings
1%
10%
40%
100%
Pretrained Model
19.32
Remove H2
31.78
Fine-tuning
19.86
19.84
19.76
19.83
Overlap across Diverse Input Contents. Moreover, we conduct comparative analysis of the activation frequencies acquired from various input contents. Specifically, utilizing the pretrained OPT-1.3B model, we evaluate three datasets, namely Wiki-Text-103, Penn Treebank, and Amazon
Review. The positioning of H2 depicted in Figure 12 revealing significant concurrence across
multiple datasets.
Early-Bird Property. Furthermore, our investigation re- veals that H2 displays an "early-bird" characteristic, as illus- trated in Figure [3] By visualizing the distribution of activa- tion frequencies across various checkpoints throughout the training process, we observe the emergence of power-law behavior at an initial stage, specifically as early as train- ing budget of 4%. Subsequently, the positions of H2 exhibit gradual and minimal changes.
100
Checkpoint
Activated Frequency (%)
Checkpoint
80
36% Checkpoint
60% Checkpoint
100% Checkpoint
40
2000
4000
6000
8000
10000
Neuron Index
Figure 13: The distribution of activated frequency during training. Experiments are conducted with different checkpoints of OPT-2.7B during training.
31

---

Layer
Layer2
Layer 5
Layer4
Antivea (oeqteyr Ph
100
Wici-Teat-10h
100
Wike-Teat-11
100
Wia--Tas-105
100
Wike Teas-01
Arnosuot
Astaout
Parss Treebank
Pess Treeback
Pam Treebask
Pess Treduak
60-
50
60
60-
20
20-
20
20-
1000
5000
2003
6000
Liger
Leyer
Layss
Layer
Arthomee Papurry (N))
100
100
Wike-Teu
50-
Wike Tea:
Aunatio
Reviow
Assazut
Parm Treebank
Peze Thocback
Pam Treehask
Peas Tredhakk
60-
50-
40-
20
20-
20
2000
4000
5000
5000
2003
4030
0300
8900
2000
5002
5000
2002
40:0
6200
5020
Layer
Artivaad Fequrryy (6)
100
1ea--Teas-103
100
Wile-Texi--n
100
100
Wilk Te0i-10
Amroort Resiew
Amaron Reviox
Amroort Review
Arrzcon Bovicw
Poen Treebork
Pets Traebark
Pas Treebaik
Peza Treebrk
60
20
20
2000
4100
5000
$000
2003
4030
6300
8900
2000
4100
5003
$000
2002
40:00
6200
5020
Artivnaa (neqwary(
100
Witi-Tes--03
100
Wile-Texs-112
100
Wilks Teat-103
Amnovn Restew
Amazon Revios
Amvovr Revirw
Arrzcon Review
60
Pas Treekaak
Pess Teabbark
60
Pau Treebark
40-
Peza Treebakk
20
20
20-
0-
2000
4100
5000
$000
2003
4000
6000
1000
2000
4100
5000
$100
2002
40:00
6000
5020
100
80-
Amnovo Restew
Arazznn Reviow
50
Anvonu Revirw
50-
Artazon Review
Peu Treebaak
Pezs Trocbank
Pau Treebark
Pezz Trcehark
00-
40
40-
20
0-
2000
$500
5000
$000
2003
0030
0200
5000
2000
4040
5002
5000
2002
5030
6200
5000
Antoned Vnoquery (S5)
Will-Tea
100
Biti-Tex--10
Tour: 163
80
Armoore Review
50
Amtacnn Reviow
50
Armoors Reeviww
50-
Artacon Hevrew
Prem Traebook
Traabark
Porm Trrebonic
Pern Tradwerk
40
40-
20
2000
4000
5000
8000
2003
4000
0200
5900
2000
4000
5002
5000
2002
4000
6200
5020
Figure 12: The distribution of activated frequency across diverse input content. The x-axis represents the index neurons, which is ordered by the activated frequency from Wiki-Text-100
32

---

D Theoretical Analysis
Recently, number of works have studied the attention scheme in LLMs from a theoretical perspective
196
98
99
100
[0T
[02
[03] [04 [O5] [06 [07]]088 [109 110 II [12
[14 [1$] [16
117
119
120
121
122.. In this work, we provide a different and novel angle compared to the
previous work. We present the concept of the submodular property and propose an eviction policy, known as greedy H2. which is a modification of dynamic submodular maximization. Furthermore, assuming the attention scheme to be submodular, we establish that constructing the set S; without any cache size limitation satisfies the near-optimal property in terms of submodularity. We provide theoretical guarantees for our robust and approximate greedy eviction policy algorithm (Algorithm 2). Due to space limitation, we only give informal description of algorithm (Algorithm (2 in Section 4.1 In Section D.6 we give an algorithm (Algorithm 2 which has full and complete implementation details for Algorithm I. We also offer mathematical formulation for sparsity preservation that is observed in Section 3 and proposed an algorithm (Algorithm 4 to solve the problem.
Specifically, in Section D.I. we provide several basic definitions and notations. In Section D.2 we briefly the definition of submodular function. In Section [D3] we define the dynamic submodular framework, which gives the formal version of Definition 4. In Section [e.4] we briefly review the static attention computation problem. In Section [D.]] we formulate the attention computation in recursive fashion. In Section we briefly review our eviction policy, which gives the formal version of Definition 4.3 In Section D.7 we discuss the diminishing return for submodular. In Section D.S we discuss the high-level ideas ideas for submodular. In Section D.9] we analyze the robust greedy algorithm error propagation. In Section D.10, we explain how to add items into sets via approximate function. In Section D.11 provide several definitions related to dynamic properties. In DocctionD.12) we prove an induction lemma for the exact function. In Section D.13 we prove an induction lemma for the approximation function. In Section D.14 provide theoretical guarantees for both the full-knowledge version (formal version of Lemma 3.1) and the limited-cache-sizz
(D155) version (formal version of Theorem 4.4). In Section D.15] we provide more detailed discussion of theoretical work about attention computation and regression-relates problems. In Section 16 we provide a mathematical formulation for sparsity preserving. In Section D.17] we provide the definition of loss function which can potentially generate sparse (heavy hitter type attention sore). In explain how compute the gradient of the loss function. In Section (.19) we show
Section D.18 how to compute the Hessian of the loss function. In Section D.20] we show that Hessian is positive definite. In Section [.21] we prove the Lipschitz property for the Hessian matrix. In Section [D22 we show that using a gradient-type algorithm is sufficient to optimize that (heavy hitter type) loss function.
D.1 Notations
For a positive integer n, let [n] := 21, ... nn].
For a vector x E RR, let Vi E R denote the vector with the i-th entry being v,, and diag(x) e
Raxe denote the diagonal matrix with the i-th digonal entry being z;. For two matrices A, W E
RTx=, let |||| == (Σ 1= Σj=1 W,g4?,," 11/2 and W a denote the matrix where (W o A)k = W A.y.A.. For matrix W E Raxn let Dw. == diag(W,.)) with E [n].
For two vectors x E R2 and w E RD let ||~||w == C w,a?)1/2, For a vector x, its Eg norm is defined ||<||| == (E 2//11/2 and its é, norm is defined || = (Ein |.; P))// For square matrix A, we denote ir [A] as the trace of matrix A.
For a matrix E R (suppose n 2 k), we use ||^||| to denote its spectral norm, i.e., |||||| We use |||| to denote its Frobenius norm || == (Ei Σ A2,) 1/2
sup, Ax ||11||
Suppose matrix A E R has SVD decomposition UEVT where U E Rank (this matrix has orthonormal columns), Σ E R is a diagonal matrix, and V E R We call columns of U singular vectors. We use At E B to denote the Moore-Penrose pseudoinverse, then At = V Suppose Σ E k sorted diagonal matrix, let .., oh denote the diagonal entries of Σ. Then we call o, the i-th singular value of the matrix, and we write it as o, (A).
33

---

For any symmetric matrix B E R we denote its eigenvalue decomposition as UAUT, where A a diagonal matrix. Let AI. Ds denote the entries on diagonal of A R We say X the i-th eigenvalue. Usually, we write it as (B).
The connection between eigenvalues and singular values is
( (A
We the notation A = 0 to denote that matrix A is positive semidefinite (psd). Mathematically, ž 0 means for all vectors I, we have xT Ax 2 0.
Similarly, for two squarer matrices A and B, use A = B to denote the case where for all vectors J. x Ax 2 x BB..
Let Pr] and El] denote the probability and expectation. We define the maximum between and b as max aa b}. We denote min(a, b} (resp. max(a, b)) as the minimum (reps. maximum) between and b.
Throughout, for non-negative real numbers and the notation (1 + c)b if [11 c)b, (1 + c)bj.
D.2 Submodular
We provide the standard definition of the submodular function.
Definition D.1 (Submodular function (13)). For finite set Ω, submodular function is defined as
f: 2 + R. where 2 denotes the power set of Qt. is characterized by the fulfillment of any of the following equivalent criteria:
Condition 1.
For S,T,C Q with S T and every a E Q.T we have that S(SU (a}) - S(S) 2 (aT - f(T)
Condition 2.
For every S,T S we have that f(S)+ f(T) 2 f(SUT)
&oondiiinn 3.
-or every S C Ω and a1, 0g S such that ai # 02 we have that f(S {ai)) f(SU {a2}) 2 f(S {ai, a23) + f(S).
For convenience discussion, this paper, we always choose f2 [n] when we want discuss the submodular function.
Next, provide some examples/types of submodular functions. One important class is called monotone,
Definition D.2 (Monotone). set set function f is monotone for every T S we have f(T) S f(S).
Here a number of monotone submodular functions
Linear (Modular) functions
linear function can be represented as S(S) = Eies wi If all the weights w, nonnegative, then the function f is considered monotone.
Budget-additive functions
budget-additive function has the form f(S) = min B, Eies WA where each weight wi and the budget B are nonnegative.
Coverage functions
We have set -= (E1, E2,... En where each E, is a subset of a broader set S The coverage function can expressed as S(S) = Uree Egj for any S Q. This function can be generalized by assigning non non-negative weights to the elements.
34

---

D.3 Dynamic Submodular
The standard submodular function is mapping 2 to real. Here we need more general definition that maps 2|0| 2 real.
Definition D.3 (Strong submodular). We define function F 2/m] x 2 IR, then for any
Z C [n.]. assume that F(Z,.) 2(x) + IR is submodular function.
In fact, the above definition stronger than we want, what we really need can be written as follows. We remark that in Definition 4.1 we provide informal definition. Here provide more detailed version of the definition.
Definition D.4 (Dynamic submodular framework, formal version of Definition
Define function
F: 2 x [n] 2 + R. Then for any index i E [n), any set Z S ii 1]. we assume that
F(Z, i, 2 +R R
a submodular functioo w.r.t. to Z, i.e.,
For all sets X, Y C. [n]] satisfy that Z c X CY,
For all element x E [n]] satisfy that E [n]]Y,
we have
I2,(X)
(z, (Y
J2,(Y),
where [z,(()) = i.
Remark D.S. We remark that Desinition D.4 is weaker version of Desinition D.3 We also like mention that the informal definition (see Definition 4.7 only contains two inpul parameters, but in fact we need three input parameters (see Definition D.4)
In the later, when we use fz.((), we will replace by for convenient of aanalysis, for example see Definition D.23 Definition D.24 Definition D.25 and Definition D.26
D.4
Static A Attention
Before we describe the recursive attention computation, we will first describe the static version of
attention computation as follows (for examples, see Definition 1.1 in 988 and others 197 100 101
106
103
118.
108); )
Definition D.6. Given three matrices K, V Bdxd the goal is compute
Att(Q, = D-
where square matrix A E RRax can rewritten as follows
= exp(QKT)
and diagonal matrix D paxn can be written as follows
D = diag(A1)
Here we apply exp() to matrix entry-wisely to matrix. 1, is a length-n vector where all the entries are ones. The operator diag() is turning vector into diagonal matrix.
D.5 Recursive Attention Definition
We first provide some definitions.
Definition D.7. Given query matrix Q E Raxd, we use Qi, to denote length-d row vector that represents the i-th row of Q for each i [n].
Definition D.8. Given key matrix K Raxdd we use Kei, to denote x d matrix that selects the first rows from K.
35

---

cat
an
apple today
cat
an
apple today
(Kr
(K2,-) (Kg..) (Ka,.)
[K (K2,a) (Kg,.) (Ka,a)
1 (Q1.))
01,1
I (Q1.))
01,1
cat (Q2.-)
0211
022
ent (Q2.)
02,2
an (Qs,,)
09.1
03.
03,3
(Qa,-)
03,1
03,2
03,3
apple (Q+,.)
04,1
04,2
04,3
04
apple (Q+,,)
01.1
04,2
04,3
01,4
today
today
(a)
(b)
Figure 14: Exact version of the attention computation. Here, our example about the language modeling task.
At this stage, the model predicts the word 'apple' and computes exact attention vector 04. (b) Approximate version of the attention computation. Let the budget of the number of tokens we can track be 3. We truncate the key matrix for the first token K and only keep K2,,, K3,, and K4,, in the memory. Compared with the exact version, we don't compute 04.1-
Next, we show the exact version of computing attention. Definition D.9 (Exact Version). Given Q, K Raxd
For each E [n], we use o, E R' to denote a length-i vector.
For each j E [1]. we use oz, to denote the j-th coordinate of vector o, E
For the first layer, we have
01,1 =
For the second layer, we have
length-2 vector 02 == Dz eexp(Q2. (Ks2-)) T
scalar
dx2
scalar D2 exx
1x2
For each i E for the i-th layer, we have
length-i vector o = D. explQii
(K
scalar
1xd
dxi
scalar D; := exp(Qi.,
1;
1xi
Now, we show the approximate version computing attention. Instead computing the entire attention o,, only compute the attention of the tokens that are being tracked.
Definition D.10 (Approximate Version), Given Q, K E RRaxd
For each E [n), we use o, R$ denote a length-i vector.
For each j E [6,,
oú denote the j-th coordinate of vector E R
Lett ls [n] be the budget of the number of tokens we can track (due to the memory issue).
36

---

For the first layer, we have
01,1
For the second layer, we have
length-2 vector 02 = D2
exp(Q22 (Ks2.)
scalar
1xd
dx2
scalar D2 exp (K
12
2x1
For each i [nd], for the i-th token, we have
Let St C. [n] denote the tokens we'r 're tracking when we are predicting the i-th token.
|Si| = k
|SASS & or equivalently |S; SS
length-i vector o;
DJ
exp
Qi,
(
scalar
1xd
dxi
scalar D; == (exp(Qu,, (Ks,))
1Avs)) 1
ix1
In certain sense, the above definition is related to finding heavy hitters in compressive sensing (for more background on compressive sensing, we refer readers to 124 125 126)).
D.6 Eviction Policy
The goal of this section is to our eviction policy under space limitation. We start by giving process without having space limitations. We denote the attention query matrix as Q E Rexd and the key matrix as K E Rexd Qi represents the i-th row of Q and K represents the first i rows of K. For simplicity, Ks,,* denotes sub-matrix of K which selects S, rows from K.
Definition D.11 (The generative process with Full knowledge). For each E [n)], for the i-th token, we have
length-i vector 0; == D exp(Q...(Ksi.))) which denotes the attention of the i-th word.
scalar D; = exp(Qi.. (K 1, which denotes the normalization farmatizationaactorr
In the above process, since we have no space limitation, then we can keep all the scores. In the following process, we show how to handle things when there space limitation. In this case, need to dynamically maintain set S; such that [S < k (Compared to the full knowledge case, can think of [S; = i).
Definition D.12 (H2O Eviction Policy, Formal version of Definition 4.3 Let k denote the budget of space and is n. Let Fscore 2(m) + R denote certain score function. Let S,-1 denote the source set. Let St denote the target set. defined the eviction policy g S. - S S.I.
[S|] < 1: (KV cache size is not changing over the time)
I.. 1SA S (we can evict at most KV in the KV cache)
We construct S; + (8)-1 as u <- arg max,g(S...1()) FS
Remark D.13. We remake that, Definition 4.3 we introduce a simplified notation where we state
that [Sa] = k in the first bullet point, but in general, consider the case where [Si] S k for the first tokens. accommodate this more general scenario, we present Definition D. 12 which provides broader definition that handles the situation where |S,) k.
37

---

Algorithm H2 Eviction Algorithm, Query matrix Q E Ruxd key matrix K 1aax budget size of KV cache k E N. Formal and detailed version of Algorithm
procedure Hg 1v__EVICTOON( Rddd K Ruxd k N)
So - 0, 0o f- 0 3: for i
2:
Initialization the algorithm
do
b Linear scan each token
4:
if i < k then
> If the KV cache is not full
5:
S - S [
Expand the set directly D If the KV cache is full
6:
else
7:
Dj +- (exp(Qi. (K) 11115_.)) 1, P. Compute the normalization factor
8:
Compute score vector Accumulate the scores (Remark D.15 Let score function be Fscore (T) == h1 2 oz,a) b> Score function (Remark D.14
Q: - exp(e exp(Qi, (Ks, 110) 11+)
9:
0 f- 01-1+
10:
11:
" < arg max Sscore (Si U (1M{
& Find an entry to swap
12:
S + (S--1 (f))fu}
> Construct S;
13:
end if
14:
end for
15:
end procedure
Remark D.14. We remark the above function Escore can have multiple choices. In the later analysis. provide two choices. If we use the exact function, then Fscore is fs.,; (see Definition D.23 and
Definition D.25 If we use the approximate function, then Fscose if Jss,a (see Definition D.24 and Definition D.26. Let h R R (the h being used Line 10 in Algorithm 2) denote function. We hope to choose h such that it gives the submodular property for the function Fscore in the sense of our dynamic framework. For example, h(z) is usually chosen to be some non-decreasing concave function such as V and log(= + 1). For simplicity, in Algorithm (which is the informal version of Algorithm 2. we choose h[(z) = z for the purpose of providing readers with an intuitive understanding and due to space limitation.
Remark D.15. We remark that, in Line 10 Algorithm 2 the og,s the accumulation of atlention score for the token s in set T. In our Algorithm [ represent that accumulation of the attention score for simplicity.
we only use Os in Line T0 (see Algorithm 7 to
In Algorithmm 2 for simplicity of the presentation, we can create multiple vectors ... -.. On (see Line 2 Line 9 Line 10, However, every time at i-th, we only need to use the information for o; which has at most length n size. Thus, a straightforward and better implementation for only using one length-n to store accumulative score reduce n² usage n.
D.7 Explaining Submodular Diminishing Return Property in Attention Scheme
In the standard submodular problem [123], we are given a ground set [x] and a function f: 2(x] + IR. The goal is to find set of elements S such that f(S) is maximized.
We say function is submodular (Recall the formal definition in Definition D.I. if for every X, Y c [r]] with X C Y and every x E [rt] we have
f(XU f(x) f(Y u() f(Y)
Submodular functions can represent the cost items, due to their diminishing returns property 127 [18]. It suggests that the increase in information obtained from selecting candidate object, such as a word or sentence, becomes smaller as more objects have already been chosen for the summary.
For instance, we introduce a new token, denoted as "w",, into two sets, S and So, where the concepts covered by So are a subset of those covered by S. By intuition, the information added to Sg by "w" should be larger compared to adding it to S, as the new concepts carried by "w" might have already been covered by the concepts present in S but not in Sg. This property is known as the diminishing return property. Hence, we propose that the neural coverage function [129] should exhibit desirable characteristic called submodularity.
38

---

D.8 Submodular: High Level Ideas
We formalize the submodular function maximization problem with cardinality constraint in this section. Informally, set function is submodular if it has decreasing marginal increment.
Definition D.16 (Submodular function). We denote f 2( 7 R as a set function. The discrete derivative A defined as follows:
A () U (i)) f(S).
A fuunctin f is submodular if. for any S C T and E [n][T, the following inequality holds;
( (T) < 4 (
For simplicity, we present the problem of maximizing submodular function with cardinality constraint Tb. Our goal is to solve the optimization problem efficiently.
max
f(S)
(1)
S.t. | S k
Representation of f(S) One challenge in designing the algorithm is determining the representation of input instances. Since the constraint in optimization problem ( is straightforward, we need to
decide how to represent f(S). Suppose = i1, 12, ... Im C [2] we can decompose f(S) into sum of increments follows:
S(S) = S(So) + -
+ij. - Without loss of generality, we assume f(0) 0. By the definition of A Si] we have S(S,) ((S)-1) = A Therefore, the decomposition (2) can be simplified follows:
where So 0 and Sj S
f(S) = _ (1 j=1
(3)
To introduce our advanced data structure later, we further represent ( in the form of
A ( u. h(S)
(4)
where u, E Rd is d-dimensional vector and I(S)) Rxd is d-by-d matrix.
In practice, a significant subclass submodular functions is the monotone submodular functions, i.e. functions satisfying S(A) < S(B) for all A C B C [n]. When f is monotone, we could restrict all ((S) to be positive semidefinite (PSD) matrixes. When the matrix h(S) positive semidefinite (PSD), it makes us achieve faster acceleration.
D.9 Robust Greedy with Error Propagation
The procedure for the greedy selection algorithm initiates with empty set Sy For each iteration in the main loop, the algorithm chooses the element that maximizes the marginal increment to add to the set. When the size of the set eventually reaches k;, the algorithm return that set S. Specifically, for iteration E (1,2.... k}, we let
S. + U
where the element in the singleton is je arg max ja f(S,11 (i)). The greedy strategy is effective the sense that the approximation error of is 1/e.
Theorem D.17 130). For monotone submodular function the greedy algorithm (Algorithm guarantees to output a set S satisfying
S(S) 2 (1 (1-11/6 1/e)
max (S(T))
-
Corollary D.18 (131). Given
39

---

Algorithm 3 Greedy algorithm benchmark
1: procedure GiREDDYAL.CORTTIM(SSbbnnddll function f)
2:
So 8 0 3: for f = 0 do
> Initialize an empty So
4: - arg arg u {i))
Find the element [n]] that maximize f(S, (i))
5: + (i}
> Append element to set S, to get S
6: end for
7: return S
8:
end procedure
accuracy parameter >
integer k 2
Let O denote an oracle that takes an arbitrary set S S [n] and [n]],S, returns value
(SS, i) with guarantee thut (iis)) - < O(S,i) I (iiS)) +
Let A denote an algorithm that each time step t 1,2, ... k. it selects j. arg max, (O(S and lets St + S u Lic}.
Then this algorithm A
returns set Sk
the S₁ satisfy that
f(Sk) 2 (1 - 1/e) max (f(T)) - k(2 1/e)c.
T=
D.10 Robust Submodular and Adding Items
We first propose a lemma that shows the robustness of the submodular. Lemma D.19. Suppose that for each [n], we have
for all St [n] with [S.] $ k
for all X C [m]. X S;U (i) and [X (8a).
(x)s,,(X)) - (S.) (fs,.(X) fs,, S c/(2m)
Let opt; denote the optimal cost that can be achieved by using f in i-th iteration. If the greedy algorithu use f to find the solution has performance at least (1 11-1p 1/e) opt, then using f. we can obtain a solution that has performance least (1 1/e) opt,
Proof. The proof follows from Lemma D.I8
Next, we explain how to add items into sets based on exact values. Definition D.20 (Expanding items based on exact value). If the following conditions hold
Let S: S 1].
Let fs..t : 2 + 12
Then, we can define S.+1 as follows
Sf = k then S:11 St U (i) where arg sagmaxeSSufi fs. U (i)1)) fs.((S.)
If|S. < then S = S (i).
40

---

Remark D.21. We remark that u arg m - is the same as = arg ( (s,,(S; U (w). For the convenience of discussion and analysis, we will switch using both cases in different places.
Here, we explain how to add items into sets via approximate values. Definition D.22 (Expanding items based on approximate value). the following conditions hold
Let Let S; S 1).
. Let Ss 2. - R.
Then, we can define Ss+1 as follows
|S = k then Si41 = Si U (uf where = arg max fs. U (iP1)) fs.((S.)
If|S:] < k, then 3
D.11
Universal Conditions
We state several definitions for both the exact function and approximation function.
Definition D.23 (Universal Monotone Condition (for exact function)). We say f has universal monotone condition, if for all E for all S; [ 1), have
(x) 2 SY,,((Y, VY
Definition D.24 (Universal Monotone Condition (for approximate function)). We say f has universal monotone condition, if for all for all S, [i 11. we have
ix) 2 (Y), cX
Definition D.25 (Universal Dynamic Condition 1(for exact function)). We say f has universal dynamic condition I(for exact function), if for all for all S. S li 1], S 2]. |SAS 1, we have
(S.) 2 (1 (1.0))fs....11(3)
Definition D.26 (Universal Dynamic Condition (for approximate function)). We say F has universal
dynamic condition !foor approximate function), for all i for all S; E [i 1). S:-1 S ii 2].
|S, S we have
(S.) 2 (1-00)fs....ee
We define opt, as follows: Definition D.27. Let denote the budget length. For each [n]], we define opt,
max(Jx;(Y)) x - Y [1),X
& k,IVi Y 1).
Definition D.28 (Universal Dynamic Condition 2). For i E [n], define opt, Definition D.27 We say it has universal dynamic condition for each E [n], we have
opt: (1-7) -pp;-11
D.12 Induction Lemma for Exact Function
The goal of this section to prove Lemma D.29
Lemma D.29 (Induction Lemma), For fixed i, suppose the following conditions hold
Set Condition. S, [i-1]
Budget Condition. [S|] < k
41

---

Value Condition. f 2 (1 1/e) (11e)-(1-1 (1- -))
Universal Dynamic Condition 1 (for exact funcciion. (See Definition D.25 (S.)(S()
-1-0)-fst-ii-11
Universal Dynamic Condition 2. (See Definition (.28) opt, 2 -7 opt
Universal Monotone Condition (for exact function). (See Definition D.23) (X) fs,,(Y)
) for all Y C
Then if we construct S Definition D.20 then we have
Set Condition. S.+1 S [1]
Budget Condition. |St+1|
Value Condition, fs, (S 2 (1 (1- (1-0))+1 - opt;+1
Proof. Proof of Set Condition.
Note that S; S 1]. by using the way we construct S;+1, then it is obvious that Sc+1 S [t].
Proof of Budget Condition.
Note that S < k, by using the way we construct S441, then straightforward that |Si+1| S k.
Proof of Value Condition.
We show that
[s
2 (S)
2 (1 0) (S
(1 8) ((1 ([1-1/1) ([1-)/e)[1- 0)((1 opt i co)
(1 000--0))11-11110 ((1-1/e)(1
(1 1/e) (1)))(111)))) (1)) (1-0))+((--))0pppp__
where the first step follows from Universal Monotone Condition, the second step follows from Universal Dynamic Condition 1, the third step follows from Value Condition, the forth step follows from Universal Dynamic Condition 2, and the last step follows from simple algebra.
Thus, complete the proof.
D.13 Induction Lemma for Approximate Function
The goal of this section is prove Lemma D.30
Lemma D.30 (Induction Lemma). For fixed i, suppose the following conditions hold
Set Condition. S, [
Budget Condition. |S$| S k.
Value Condition. fs._,,i-1(5)) 2 (1 1/e) otttot i cp
Universal Dynamic Condition (for approximate function). (see Definition D.26 (s,)(8)) 2 (1-0) fs....-1(S))
Universal Dynamic Condition 2. opt, 2 (1 (1) opt,+1
Universal Approximate Condition. (See Definition D.28) fs. fx. for
all X
Universal Monotone Condition (for approximate function). (see Desinition D.24 (s..(X) 2 (s) for all (Y)(Yo)or
X
Then if we construct S Definition D.22 then we have
42

---

Set Condition. S,+1 - [t]
Budget Condition. |5|+1| S
Value Condition, (S(+1))))) (1-1/e) (1 7) opt;41 + (i 1) c0
Proof. Proof of Set Condition.
Note that S, S - 1]. by using the way we construct S4+1, then is obvious that S1+1 S [i].
Proof of Budget Condition.
Note that | S by using the way construct S441, then it is straightforward that |St+1| $ k.
Proof of Value Condition.
We can show that
(8.11)(41) 2 (S+1()
2 (8)) -E
2 (1 0) .Ss 1,1-1 (S,) 60
2 (1 0) ((1 1/e) 0) -)'ppp, -i. e0) 60
2 - 0) ((1 (11-1/4 1/e)(1 0)(( opt,+1 cp) cg
N (1 - 1/e) (1 - 0)** (1 (1-7)) opt
opt.+1 - (i 1) ED.
where the first step follows from Universal Approximate Condition, the second step follows from Universal Monotone Condition, the third step follows from Universal Dynamic Condition the forth step follows from Value Condition, the fifth step follows from Universal Dynamic Condition and the last step follows from simple algebra.
Thus, we complete the proof.
D.14 Theoretical Result
We first give the guarantee of our full full-knowledge version (without cache size limitation).
Lemma D.31 (Formal version of Lemma B.]]. Under the mild assumption, let k denote any target size. If we greedily compute the attention score based on full information, then we can find the set S; such that
f(S.) (1 1/e) a1-a)opt.
where a E (0, 1) are parameters.
Proof. The proof follows from using Theorem D.17 Corollary D.18 Lemma D.29 with choosing 0 = a/ (10m).
Next, we show the guarantee for our robust and approximate greedy eviction policy algorithm (Algorithm 2)
Theorem D.32 (Formal version Theorem 4.4p. Under the mild assumption, let denote the budget of space limitation. for each token, greedily compute the attention score based on top-k choice, then can show the St we generate each for token i [ra] satisfy that
((5.) -1 1/e) (1 aaoppt - B,
where a E (0, (0,1).3> are parameters.
Proof. The proof follows from using Theorem D.17 Corollary D.18 Lemma D.30 with choosing 6o (1000)) and 0 7 a1(1).).
43

---

D.15 Extended Related Work for Theoretical Attention Problems
The static attention computation is asking the following question that given Q, K, V E R the goal is to D exp (QK VV where D = diag( (xxp (up(QQ"]) )1,). [98] studied the static attention computation from both algorithm and hardness. On the positive, they provide an almost linear time algorithm to approximately compute the attention matrix. On the negative side, assuming strong exponential time hypothesis (SETH), they prove a hardness result. Their hardness result is, unless SETH fails, there is algorithm that runs in truly subquadratic time to approximately compute the attention matrix. Further, [100] considers the dynamic of attention computation problem. They also provide both algorithmic results and hardness results. In the work of [003, they consider attention matrix construction.
the sparsification d > n,, and show how to sparsify the columns of matrix Q. [103] provides two algorithms, one randomized algorithm, and the other is a deterministic algorithm. Differential privacy a famous and textbook topic in graduate school, recently the work of [IOI] shows how to give differentially private algorithm for computing the attention matrix. For a given A E Rn and vector b E IR", [04] formulates and studies exponential regression min, exp(Axx b1[2- Then 105 considers the normalization factor in exponential regression and defines the softmax regression problem min, (exp(Axx) 1,,) l,))-eexp(Aa exp(Ax) - bl]2. 2- [107] moves the scaling factor from exp( (Ar) to b and defines a rescaled softmax regression problem min, || [xp((r) (exp(Ax) 1,) bll2.
particular, they assume that situation that
D.16 Sparsity Preserving
Recall that in Figure we observe that even when trained densely, the attention matrices of LLMs are over 95% sparse at inference time. Only 5% of the KV cache sufficient for decoding the same output token at each generation step. Here, we provide some formal formulations for sparsity. Definition D.33. Suppose the following conditions
Let So C [m]]
Let = |Sp]-
Let E 10,1) denote a threshold for truncating the value.
Let E (0, 1) denote a fraction of mass (larger than 7) outside So-
Let mapping D Rdd
Rm
20'
For each x Rd, D(x) E Ihm is vector that has length m.
We say the distribution D is (as, T, -)-good if the following conditions hold
For all E R4. So (D(p)(
For all Rd, sDpp, S a
Claim D.34. Suppose sample points {I1, T2. Tn Rd from (a, T, good distribution uniformly at random, then we have
. So S (hee (u:p)(x;)
. |[(Uee()) supp,
(Dp)) (D(x)))) akn
Proof. Since for all i E [n], we have So C (((s)) thus
So Muem (r:),
Therefore we proved the first property.
We know that for all i [n),, we have (D(z)) akn. Thus
[(Useta)
S
(P(= S n ok
Therefore, finish the proof for the second property.
44

---

D.17 Definition of Loss Function
In this section, we follow the theoretical softmax regression literature T05 and define number of functions to make the calculations of gradient and Hessian convenient. We also proposed new penalty term (&1 type sparsity penalty, see Definition D.41 into the final loss function, which is not We first provide some function definitions.
studied in previous work [104] 105
107
110
132
133
Definition D.35 (Function u, 105 as follows
Given malrix
E R let function u Rd - R2 be defined
u(x) := e=xx(Az)
Definition D.36 (Function o, see Definition 5.4 in 105 as an example). We define u(x) as Defini-
tion D.35
Then we define a : IRd + R as follows
(xx)
(u(z), 1,)
Definition D.37 (Function f. see Definition 5.1 in 105 example). Provided that the following conditions are true
We define u(x) Definition D335
We define a(x) Definition D.36
Let function f FR IR" be defined follows
f(x) == a(z)-tu(a))
Definition D.38 (Function c, see Definition 5.5 [05 as an example), Provided that the following conditions true
Given a vector E R
Let f(x) be defined as Definition D.38
Then, let function e IRd + RR" defined as follows
c(x) ; f(z) -
Definition D.39 (Loss function Texp: see Definition 5.3 in [105 as example). We define Lexp
Rd - R
(xxp(x) = 0.5 |/c())]]
Definition D.40 (Loss function Lreg). Given A Roxd
Let function Treg Rdd > IR be defined as follows
Treg(a)) = 0.5 (dlag(w) Axll :=)
define a novel penalty function
Definition D.41 (Implicitly controlling the sparsity). Given Ra We define
Laparse (x) exp(Ax))l1
Then it is obvious that we have Claim D.42. Given A E Raad Let u be defined as Definition D.35 Let a be defined as Defini- tion D.36
We have
Leparse (x) (exp(AAr), 1n)
Leparse (x) (u(x), 1n)
45

---

(eparse(()) = a(r)
Proof. The proof is trivially following from the definition of u(x) (see Definition D.35 and (
(see Definition D.36
The final loss function can defined as follows. Intuitively, we can write attention D e==exp(QKT) into subproblems where each subproblem can be viewed as softmax problem.
Definition D.43. If the following conditions hold
We define Lesp as Definition D.39
We define Lieg as Definition D.40
We define Issarree Desinition D.47
Then we define L fumctin
L(x) == Lexp(x) + (xparso(x) + Loog (z).
D.18 Gradient
Next, we show the gradient of Lexp-
Lemma D.44 (Gradient, Lemma 5.6 in 105)). Provided that the following conditions are true
Given matrix A E Raad and vector b R.
Let Av,i R2 denote the i-th column of matrix A, for every E [d].
We define a(x) Definition D.36
We define f(x) Definition D.37
We define c(x) as Definition
D.38
We define Leep (x) Definition
D.39
Let denote hadamard product.
For every i [d]. we have
Part 1.
d eexp(Ax) exp(Axt) A., dx
Part
d<exp(Ax). 1n)
(exp(Ar), Ar,i)
dæ;
Part 3,
dx(xx)1 - A.,)) dx;
Part
df(x) dc(x) dx;
(f()), f(x) f(x) A.s
dx,
PPrrt 5.
(Txxp()) = AT, (f(z)(((x) -b) f(x) (iag(f(=))(f(z)) - -b)) dr,
46

---

D.19 Hessian
Here, we compute the Hessian for several functions.
Lemma D.45 (Hessian of u(x), Lemma 5.9 in (105)). Provided that the following conditions
true
Given matrix A Ruxdd
For every E [d], let A.₃i IR" denote the i-th column of matrix A.
Let denote hadamard product.
Then, we have, for each [d]
Part 1.
d² exp((Ar) dxt
A.; o our))(()) (())) A..i
Part 2.
d² exp(Ax) Ar, o u(xx) A.i drıdaj
Lemma D.46 (Lemma 5.10 in [105)). Provided that the following conditions are
We define axx) as Definition D.36
For every E [d], let As, IR" denote the i-th column of matrix
Let denote hadamard product.
Then, we have
Part
dFa(z) dx?
(u(x), Ar, An,))
Part 2.
d?a(x) (u(z), An, Ans)) dr,dæ,
Part
da(()) A diag(u(z))
dr²
D.20 Hessian Positive Definite
It is well known that in literature [044 105
1071. the Hessian H of loss function be written as A AT(B(z) + W2). A for some matrix function B(x) E RRxx (for example see explanation in Section 5.10 in (105).. In this section, show that Hessian positive definite.
Lemma D.47. If the following conditions hold
Given matrix A E paxd
We define Laparso (x) as Definition D.41
We define I Iasparsa (x: Definition D.40
We define L Lsparse (x) Definition D.39
Let L(x) Leep(2) (I) Leparse (x) + Lreg (z).
47

---

Let ( be the Hessian of L(x)
Let W diag(w) E lgnx". Let W E = denote the matrix that i-th diagonal entry is "
Let minn denote the minimum singular value of A.
Let > O denote a scalur.
Then, we have
Part 1. Ifall i [n], w 2 20 1//0mi (A)².
then
L
dx²² = I
Ia
Part 2 Ifall i [m]. w 2 200 exp((() + 1/0 (A)?, then
(1-1/101 1/10) (B(x) W²) w 3 + 1/10) (B(x) + W²)
Proof. The entire proof framework follows from 104 mainly explain the difference.
107.. in the next few paragraphs, we
The ((x) based on Lsparso is (i))( Note that it is obvious that (iag(u(>))) > O.
From the upper bound size, we know that
diag(u(x)) 3 ||([x) In
3 exp((??)
where the last step follows from Proof of Part Lemma 7.2 in [105..
To prove Part 1, following from [104 105], we only use the lower bound of diag(u(z)). By putting things together, we get our results.
To prove Part 2, we follow from 107 and use both the upper bound and lower bound diag(u(x))
L
D.21 Hessian Lipschitz
In this section, we show Hessian is Lipschitz.
Lemma D.48. following conditions hold
Let (()) AT (B((z) + W2)
denote the Hessian of L.
Then, we have
[H(x)) - H()ll < n² n exp(4002))|== - -
Proof. The entire proof framework follows from 104 mainly explain the difference.
105
107
in the next few paragraphs, we
Note that the B(x) based on Lexp + Lyog, have been proved by 105
We only need to prove (xx) based on Leparse and add them together.
Note that (xx) based on Leparse is in fact (t))( Using Lemma 7.2 in [105,, we know that
diag(u(x)) diag(u()))) (1)) < [lu(x) u(y)l12
22 u112
where the last step follows from Part in Lemma 7.2 in 105 Thus, putting things together, complete the proof.
48

---

D.22 Greedy Type Algorithm
In this section, we propose a greedy-type algorithm (based on the approximate Newton method) to solve the optimization problem.
Algorithm 4 A greedy type algorithm.
1: procedure OURITTRRATIVEEEETHOOAA E Raxd, E IR" w E R
2: Initialize xо
3: T f log( x l22/E)
D Let I' denote the number of iterations.
4: for 0 + T' do
5:
D - Baiag + diag(w w)
6:
D SUBSAMPI.((D, A, @11), 6
7:
Compute gradient exactly
Get the approximate Hessian H by computing A DA
8:
9:
Update I by using the Newton step + H
10:
end for
11: I 4- 2f++1
12: return T
13:
end procedure
Theorem D.49. Given matrix A E Ruxd, b E RR", and w E IR..
We use x" to denote the optimal solution of
min Lexp + Lipprrse + Lreg ZER
that
-g(x") -,, where g denotes the gradient function. - -1<x* 112 S R
Suppose that R 2 10, M exxx(((p0@ + log n)), and > 0.
Assume that |||| < R. Here ||A|| denotes the spectral norm matrix A.
Suppose that b O,, and ||b|| < 1. Here On denotes a length-n vector where all the entries are zeros. (Here > O,, denotes 6; 2 0 for all i E [m])
Assume that wi 200 exp(R²) + smallest singular value of matrix A.
m/m for all i [ni]. Here (A) denotes the
Let uo denote an starting/initial point such that M| - x² S 0.11
We use to E e.0,0.1) represent our accuracy parameter.
We use 5 E 00,0.1) represent failure probability. There is randomized algorithm that
runs log[[ x= 112/e) iterations
spend
(nnz(A)) +d)) poly(log(n/6))
time per iteration,
and finally outputs vector ă E dd such that
Pr| x - 1
49

---

Proof. The proof framework follows from approximate Newton (second order method) literature
134
135
136
137
138
139
140
104
[05]
107
110
142 143
[45].
Following from Lemma
D.47
we know the Hessian of the loss function is positive definite.
Following from Lemma D.48
we know the Hessian of L is Lipschitz.
Following Section
[05 by running Algorithm 4 complete the proof.
We remark that w denotes the exponent of matrix multiplication (i.c., n² is the time of multiplying an n x n matrix with another n x n matrix). The most naive algorithm gives w 3. Currently, the ofate-of-the-ar algorithm gives 2.373.
50


Reviews:
[{"review_id": "xkyNrZ5CZO", "reviewer": "Reviewer_AZfb", "summary": "Paper addresses inefficiency in key-value matrix caching during attention computations of generation with transformer models by introducing a new pruning strategy to discard less useful information without significant degradation of performance. This strategy is based on the reliance of transformers' on observable token cooccurrence patterns, where only a subset of tokens significantly contribute to performance. They frame the problem as dynamic submodular optimization, and show empirically that the strategy considerably reduces memory overhead and inference time while preserving most of the original performance.", "questions": "Presumably this approach is implemented uniformly across layers, though this is not made entirely clear in the writing. Given the variation in sparsity across layers, what is the effect of varying the policy based on the layer?", "limitations": "The authors do address limitations", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "- The authors introduce a simple, intuitive approach that is empirically demonstrated to preserve performance in the zero-shot setting across model architectures/benchmarks. Most of the paper's claims are well-validated. \n\n- They tackle an important issue of LLM inference efficiency (both memory and speed-up), and their approach has many potentially critical practical implications for research engineering work involving LLMs.", "weaknesses": "- Given the effort put into the rest of the work, it is disappointing that the authors' claim in 5.3 that H20 leads to more diverse generations is only weakly substantiated through a small qualitative evaluation (assessed by the authors). This should be demonstrated with a large human study or through diversity metrics over an evaluation set (e.g. unique ngrams, self-bleu). \n\n- While they test more model architectures/datasets for the final results, the sparsity findings are only considered for OPT model generations. This is also reported in isolation, without reference to prior work on sparsity in attention matrices, even though this has been widely studied, and gives the impression that the authors are claiming this as a novel observation and sole motivation for their \"heavy-hitters\" strategy. I would suggest adding citations to previous work noting sparsity, e.g:\n\n1. Edelman, B., Goel, S., Kakade, S.M., & Zhang, C. 2022. Inductive Biases and Variable Creation in Self-Attention Mechanisms. ICML. \n2. Likhosherstov, V., Choromanski, K., & Weller, A. 2021. On the Expressive Power of Self-Attention Matrices. ArXiv, abs/2106.03764.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We are glad that the reviewer appreciates our work and efforts, and acknowledges we **tackle an important issue** and our approach **has many potentially critical practical implications**. We provide detailed responses below to address the questions. \n\n**[Q1: More Evaluation of Diverse Generation]** Thanks for the suggestion that helps strengthen our observations! We conduct the additional more formal evaluation on the finding of increased diversity by Self-BELU[1] as suggested. \n\n**Result**: The full model reaches a Self-BELU score of **0.0057** while H$_2$O and Local method achieve **0.0051** and **0.0436**, respectively. The comparatively **lower** Self-BELU score of H$_2$O indicates the slightly **increased diversity** of generated text.\n  \n**Details**: We randomly sample 100 instances from the XSUM dataset and use them as prompts. Subsequently, we employed the LLaMa-7B model to generate text of equal length. \n\n**[Q2: References and More Observation of Sparsity Finding]**: Thanks for pointing us to the related work that supports our sparsity observations. We have added it to the discussion. Below we present our extended similar sparsity observations in more model families:  \n\n**Results**: Figure R1 of the uploaded PDF shows that attention sparsity also exists in LLaMA and GPT-NeoX, where nearly all layers exhibit sparsity levels exceeding **90%**. \n\n**Discussion**: Additionally, we include the missing two references as the following in Section 3.1 in the updated version: “Inspired by previous literature, which reveals the existence of attention sparsity in DistillBERT[3] and bounded-norm self-attention heads[2], we first show an observation that such sparsity of attention also exists in pre-trained LLMs.”\n\n**[Q3: Layer-wise H2O Policy Results]** This is a great suggestion! Indeed, our approach assigns the same KV cache budget for each layer in LLMs. We will clarify this setting in the updated manuscript. Furthermore, we present additional preliminary studies suggested by the reviewer about how the layer-wise cache budget can further help the final performance below. \n\n**Results**: We conducted experiments with three layerwise sparsity-aware strategies. The results in Table R1 suggested that one of them can further boost H$_2$O performance by up to $1.00%$ accuracy and the other two strategies fail and cause a performance degradation ranging from $0.07%$ to $13.20%$. This is a very exciting signal that suggests that if we can find a low-cost algorithm to predict the sparsity level, we can exploit it for a higher compression ratio. We have included the results in our manuscript. \n\nTable R1: Results of layer-wise KV cache budget. Experiments are conducted on LLaMA-7B across different tasks. And the overall KV cache budget is 20%.\n\n||PiQA|COPA|OpenbookQA|Winogrande|\n|---|---|---|---|---|\n|Full| 80.09 | 81.00 | 44.80 | 71.51|\n|Local|57.94|56.00|28.40|51.30|\n|H2O|79.22|85.00|43.80|71.67|\n|H2O Sparsity-Aware|72.03|78.00|41.20|62.27|\n|H2O Random|72.15|78.00|30.60|64.64|\n|H2O Sparsity-Aware-Custom|79.33|86.00|44.60|72.30|\n\n\n\n\n\n\n**Details**: Three strategies under the same overall KV cache budget (20%): i) we assign layer-wise budget which is proportional to $1-s_i$, where $s_i$ is the attention sparsity levels of layer $i$ (Sparsity-Aware); ii) we first define sensitive and non-sensitive layers based on whether its sparsity ratio $s_i$ is smaller than the model’s average sparsity. Then we assign more KV cache budget to the sensitive layers where we make the budget of sensitive layers $60%$ larger than the ones of non-sensitive layers(Sparsity-Aware-Custom); iii) we assign layer-wise budgets following a uniform distribution (Random). \n\n\n[1] Texygen: A Benchmarking Platform for Text Generation Models\n\n[2] Edelman, B., Goel, S., Kakade, S.M., & Zhang, C. 2022. Inductive Biases and Variable Creation in Self-Attention Mechanisms. ICML.\n\n[3] Likhosherstov, V., Choromanski, K., & Weller, A. 2021. On the Expressive Power of Self-Attention Matrices. ArXiv, abs/2106.03764.\n\nThanks for reading our response! Please let us know if the above addresses the questions and concerns."}, {"review_id": "CJSFMTC8UZ", "reviewer": "Reviewer_CXTK", "summary": "This paper first identified an interesting observation that only a small portion of tokens contributes most of\nthe value when computing attention scores. And based on this observation, this paper proposed Heavy Hitter Oracle (H2O), a KV cache eviction policy that dynamically retains a balance of recent and H2 tokens. The results show significant improvement in memory reduction and inference speedup.\n", "questions": "See weakness.", "limitations": "yes", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "- The paper is well-written and easy to follow.\n- Strong performance in terms of decoding latency and memory reduction. \n- The idea of evicting $(K,V)$ of non-important tokens from the cache is intuitive and works well.", "weaknesses": "Minor:\n- I wonder what would happen if the current approach is applied during training, although it might be slow. Will it result in better performance since the model can have a better inductive bias to learn from important tokens?\n- Would be interesting if we can see the same method work in the encoder-decoder framework.\n- I wonder if the same approach can be applied to reduce the memory size in Memorizing Transformer ([link](https://arxiv.org/pdf/2203.08913.pdf)), or essentially any similar soft external memory mechanism in LLM.\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for strongly supporting our work and providing valuable feedback, which helps improve our work. We present pointwise responses for your constructive suggestions as the following:\n\n\n[**Q1: Results on Encoder-Decoder Framework**] Thanks for the suggestion. We further evaluate our method with representative encoder-decoder models, including T5-Large, T5-3B, and T5-11B.\n\n**Results**: As shown in Table R1, with a 20% KV cache budget, H$_2$O can match the performance of the full model across different model sizes, which demonstrates our approach is also effective in the Encoder-Decoder framework.\n\nTable R1: Results of T5 models on the XSUM task. \n|Rouge-L|T5-Large|T5-3B|T5-11B|\n|---|---|---|---|\n|Full|14.6831|14.8753|15.7553|\n|Ours (20%)|14.0567|14.8074|15.5190\n\n**Details**: Experiments are conducted on the XSUM task, and we use the Rouge-L score to evaluate the model’s performance.\n\n\n[**Q2: Reduce the Memory Size in Memorizing Transformer**] This is a good suggestion. We investigate the potential applications of our approach for improving memorizing transformers. \n\n**Results**:  The results are reported in Table R2. With the same memory size, the heavy-hitter strategy can effectively improve the final performance while restoring only the non-heavy-hitter or random tokens causing performance degradation. This positive result demonstrates the possibility of reducing the external memory size of memorizing transformers by introducing heavy-hitter tokens. And we will systematically study it in future work.\n\nTable R2: Results of memorizing transformers with heavy-hitter tokens. \n||Baselines|w. Heavy-Hitters|w. Non-Heavy-Hitters|w. Random Tokens|\n|---|---|---|---|---|\n|Perplexity|4.014|3.962|4.178|4.049|\n\n**Details**: Experiments are conducted with 8-layer memorizing transformers on enwik8 dataset. Four methods are compared, including the vanilla memorizing transformers (baseline), with only heavy-hitter, non-heavy-hitter or random tokens.\n\n\n\n[**Q3: Apply H2O during Training**] This is indeed a natural question! We conducted a preliminary investigation about incorporating heavy-hitter into training.\n\n**Results**: With our approach, the final model reaches 67.68 Perplexity (PPL) on the validation set, while the baseline model only reaches 69.02 PPL. The performance improvement demonstrates the potential of applying our approach during training, and we will investigate more in the future.\n\n**Details**: As illustrated in Figure 2(b),  heavy-hitter tokens exhibit high co-occurrences within the dataset. Based on that, we first calculate the co-occurrences statistic in the training data and identify the heavy-hitter tokens. Then, we manually increase the attention scores of heavy-hitter tokens, encouraging the model to focus more on the heavy-hitter tokens. Experiments are conducted with a GPT-2 model. And we train the model from scratch on 10% data from Wiki-Text-103.\n\nThanks for reading our responses! Please let us know if the above addresses the questions and concerns."}, {"review_id": "cVqE96Cxw4", "reviewer": "Reviewer_N7jZ", "summary": "The paper presents H2O, a technique to identify key elements in the sequence whose KV values play the most important role in attention scores for generation of future tokens and proposes to only keep these in cache to improve upon memory requirement and latency during inference.\n\nThe proposed algorithm is then formulated as a dynamic submodular problem through which theoretical guarantees are provided for the effectiveness of the algorithm.", "questions": "- Do you have any analysis/thought that incorporates this during training as well thereby removing the inference & training skew?\n- Do you have any comparison with similar alternative methods like ColT5? https://arxiv.org/abs/2303.09752", "limitations": "None", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "The algorithm in itself is pretty impactful. Given the amount of work done in trying to improve the efficiency of serving large language models, this algorithm presents a pretty impactful way to retain most of the quality while increasing memory requirement, throughput and latency by huge margins.\n\nThe proposed algorithm is well analyzed through theoretical guarantees, experiments, and ablations.\n\nThe paper explores an important problem that has huge impact on the usability of large models by making them less resource intensive.", "weaknesses": "- Some of the equations in the paper are difficult to follow because of the choice of notations. Especially, [i] and 1_{[i]}\n- In Algorithm 1, line 9, should the tokens that are already discarded from the cache ([i]/S_{[i]}) be masked out similar to what has been done on line 8?\n- In Algorithm 1, K_{S_{i}} is used in the equations (line 8 and 9), while S_{i} is initialized only at line 13.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Many thanks to the reviewer for acknowledging that the problem we explored is “important”, “has huge impact”, and also our algorithm is “impactful”. We provide detailed responses below to address the constructive feedback.\n\n**[Q1: Inconsistent Notations]** Thanks for pointing it out. To fix the inconsistent notations, we have modified line 8, 9 of Algorithm 1 in the updated manuscript.\n\n**[Q2: Comparison with Alternative Methods]** Thanks for the question. \nCoLT5 and H$_2$O address a similar problem but with different approaches. We first present positive results on a preliminary study of combining H$_2$O with the backbone of CoLT5, i.e., T5, which suggests that although densely trained, similar to auto-regressive models, T5 also has sparse structures at inference time. This explains why it precisely supports why CoLT5 can improve training and inference speed by routing tokens to different branches.  Finally, we conclude with a discussion of the main difference between CoLT5 and our approach.\n\n**Results**: As shown in Table R1, our approach can preserve the performance with only a 20% KV cache budget across various sizes of T5 models.\n\nTable R1: Results of T5 on the XSUM task. \n|Rouge-L|T5-Large|T5-3B|T5-11B|\n|---|---|---|---|\n|Full|14.6831|14.8753|15.7553|\n|Ours (20%)|14.0567|14.8074|15.5190\n\n\n\n**Discussion**: CoLT5 employs a token-wise conditional computation strategy to reduce the overall computation cost, which requires training the model from scratch while our H$_2$O can be directly plugged into pretrained models without retraining. Additionally, CoLT5 uses a routing network to predict which tokens are more important and need to be processed by the heavy branch. But our approach leverages a simple but effective heuristic (accumulated attention scores) to identify the “heavy-hitter” tokens at inference time. \n\nFurthermore, our method can naturally be incorporated with CoLT5 to further reduce the KV cache requirement in its heavy branch. Since there aren’t open-sourced pretrained CoLT5 models and it’s too difficult to train one model on our own in such a short period of time, we only present some preliminary results of applying H$_2$O on the backbone of CoLT5 to demonstrate how our method can be introduced in the heavy branch of CoLT5 above. We will add the reference and conduct the experiment of applying H$_2$O on CoLT5 in the future once the pre-trained CoLT5 models are released.\n \n[1] COLT5: Faster Long-Range Transformers with Conditional Computation.\n\n**[Q3: Apply H$_2$O into Training]** Thanks for the good and natural question! We conduct an additional experiment of applying H$_2$O into training and observe positive results. \n\n**Results**: After training from scratch, the resulting model from our approach reaches 67.68 Perplexity (PPL) on the validation set, while the baseline model only shows 69.02 PPL. The performance improvement demonstrates the possibility of introducing H$_2$O into training and we will comprehensively study it in the future. \n\n**Details**: Based on the results shown in Figure 2(b), heavy-hitter tokens have a high co-occurrence property in the data. With this property, we first calculate the co-occurrences statistic in the training data and identify the heavy-hitter tokens. Then, we manually modify the attention scores during training to encourage the model to focus more on the heavy-hitter tokens. Experiments are implemented with a GPT-2 model on 10% training data from Wiki-Text-103.\n\nThanks for reading our responses! Please let us know if the above addresses the questions and concerns."}, {"review_id": "oVJA3s5yp9", "reviewer": "Reviewer_nR7P", "summary": "The author proposes to reduce the size of KV cache to accelerate the inference of the Transformer models. The size reduction is conducted through a mechanism called heavy hitter oracle, which treats attention score as the importance score. The authors conduct experiments with various tasks and observe consistent improvements.", "questions": "What is the reason for emitting the throughput improvements for experiments in Section 5.1?", "limitations": "As the key focus of this paper is efficiency improvement, elaborations on this aspect would significantly improve the quality of this study. For example, in Table 2, I would recommend reporting the throughput of H2O under both GPU offloading and CPU offloading.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The studied problem is important. The proposed method is simple. The improved throughput in Section 5.2 seems significant. The theoretical analyses provide some bonus on the greedy algorithm.", "weaknesses": "The submission lacks several crucial empirical evidence to support the efficiency claim (i.e., 'reduce KV cache memory footprint by 5-10x while achieving comparable accuracy on a majority of tasks'). \n\nMy major concern for this study is the gap between the evaluation in 5.1 and 5.2. Specifically, the evaluation conducted in Section 5.2 is on the synthetic/toy data, while the major results in Section 5.1 are evaluated on real-world datasets. This inconsistency left doubt on what would be the throughput improvement of the proposed method on real-world datasets. Moreover, if the throughput improvement on real-world datasets is not as significant as in section 5.2, the results in section 5.2 would be a bit misleading. \n\nFurthermore, the baseline considered in the experiments appears too simplistic. Comparisons with the simple Top-K baseline are only briefly discussed in the appendix. The main body of the submission only contains comparisons with a very naive baseline called Local. Since It is observed that some attention heads only attend to the first token or special token, I think the bad performance of Local is as expected. Since most experiments are conducted in a few-shot setting, I recommend directly using 1-shot or 0-shot full models as baselines. Note that this approach differs from simply dropping the cache for middle tokens, as Transformers utilize positional encoding. \n\nAlso, for experiments in Section 5.1, I suggest also to report the zero/one-shot performance. Otherwise, it is unclear whether the proposed method only applies to scenarios with lots of redundancy on model inputs (e.g., 5-shot in-context learning).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for really insightful and detailed questions and suggestions. These help us strengthen experimental results and make more precise claims. We present details of our implementation as well as additional experiment results that we have added to the updated draft below.\n\n**[Q1: Throughput on Real-World Datasets]** This is a great catch! Firstly, we’d like to clarify the reason why we comprehensively assess the throughput improvements of our method with different simulated  sequence length/batch size in a separate section (section 5.2). \n- First, the throughput improvements of H$_2$O depend mainly on the sequence length of input examples, the length of the generated text, and the predefined KV cache budget, hardware spec, but do not depend on the specific content of the data. Since it is hard to find benchmarks that cover all the cases, especially the popular long-context generation which is known to be very difficult to evaluate, we use simulated data to show different system performance in a wide-range of settings.\n- Second, independent of the eviction policy proposed in our current work, section 5.2 also aims to provide a good evaluation on our system implementation which can support more future work on different or even better policies. For example, this interesting work [1] that shows up after the submission deadline presents a trainable KV cache compression policy and can potentially be deployed in our system with slight modification on the policy.       \n\n**Results**: However, for consistency and precision of scientific claims, we agree that it is very necessary to provide the throughput improvement on real-world datasets. As shown in Table R1, with our KV cache eviction policy, the overall throughput is improved by up to 3x, 29x, and 22x over FlexGen, Accelerate, and DeepSpeed, respectively. The results demonstrate that our method can provide significant throughput improvement in real-world applications. And we will include more extensive results for the real benchmarks and revise our claims in the updated version.\n\nTable R1: Results of generation throughput (token/s) on a T4 GPU with different systems on real-world datasets, XSUM. The number in the bracket denotes the effective batch size and the lowest level of the memory hierarchy that the system needs for offloading, where “C” means CPU and “G” means GPU.\n\n| Model size | 6.7B | 30B |\n| --- | --- | --- |\n|Accelerate |11.98 (1, G)| 0.23 (2, C) |\n|DeepSpeed|3.52 (6, C) | 0.31 (2, C) |\n| FlexGen |10.80 (1, G)  | 3.29 (44, C) |\n| H2O (20%) | 30.40 (1, G) | 6.70 (180, C) |\n\n[1]  Anagnostidisµ et al. Dynamic Context Pruning for Efficient and Interpretable Autoregressive Transformers. \n\n**[Q2: More Competitive Baselines]** Thanks for helping us come up with this solid baseline. We have conducted additional experiments of zero/one-shot inference with full cache size and use it as a more competitive baseline. \n\n**Results**: We show in Table R2 that H$_2$O consistently outperforms both the zero-shot and one-shot baselines, with performance enhancements of up to $9%$ across diverse tasks.  When H$_2$0 employs a 20\\% KV cache for 5-shot inference, it approximates using $1.2$ samples per input, while the zero-shot and one-shot full model use $1$ and $2$ samples, respectively. The observed performance gain further validates the effectiveness of our approach.\n\nTable R2: Results of LLaMA-7B across different tasks. Experiments are conducted under 20\\% KV cache budget under the 5-shot inference scenario, except for the two baselines methods (0-shot/1-shot full model)\n\n​​|PiQA|COPA|OpenbookQA|Winogrande|\n|---|---|---|---|---|\n|Full| 80.09 | 81.00 | 44.80 | 71.51|\n|0-shot Full| 78.89 | 76.00 | 41.40 | 70.00 |\n|1-shot Full|79.11|76.00|43.60|70.24|\n|Local|57.94|56.00|28.40|51.30|\n|Ours|79.22|85.00|43.80|71.67|\n\n\n\n\n\nFurthermore, we have also conducted a comparative analysis between our H$_2$O and another competitive baseline, namely, sparse transformer, detailed in Section 5.1, with results reported in Table 1. Both variants of the sparse transformer method (“strided” and “fixed”) encounter a notable performance degradation of up to $35$% compared to the full cache scenario.\n\n\n**[Q3: Additional Experiments of One/Zero-Shot Inference]** Thanks for pointing this out! We want to clarify that our original incentive of using 5-shot (long context) is simulating the setting that KV-cache memory becomes the bottleneck. However, we agree that there is more redundancy in 5-shots and thus conduct additional experiments to assess our approach under zero/one shot scenarios.\n\n**Results**: Figure R2 of the uploaded PDF shows that H$_2$O can successfully mitigate the memory requirements of the KV cache by up to $5\\times$ under both one/zero-shot inference (COPA, LLaMA-7B, one-shot inference; MathQA, LLaMA-7B, zero-shot inference), achieving matching performance as the model with full KV cache while the local method fails. However, unlike 5-shot, we also found that some tasks are more difficult and the model requires more information to generate the correct content, resulting in a higher KV cache budget (30-40%). This might be expected since 0/1 shots in our benchmarks have shorter sequence length (100-300). \n\nIn general, H$_2$O achieves consistent effectiveness across various tasks. And we have incorporated all the results in updated manuscripts.\n\n\n\n\n**[Q4: Throughput of H2O under both GPU/CPU Offloading]** \nWe have tested on both pure GPU and GPU/CPU offloading in Table 2 (30B models are run with CPU offloading). We will clarify this setting in the updated manuscript to make it easy to follow.\n\nThanks for reading our responses! Please let us know if the above addresses the questions and concerns.\n"}]

Author rebuttal:
We thank reviewers [AZfb,CXTK,N7jZ,nR7P] for their thoughtful and highly supportive feedback! We were glad that the reviewers found the problem **significant and impactful** [AZfb,CXTK,N7jZ, nR7P], the observations and theory **interesting and well-analyzed** [CXTK, nR7P], the methods **simple and intuitive** [AZfb,CXTK,N7jZ,nR7P], the presentation **easy** to follow [CXTK], experimental results **strong and impressive** [AZfb,CXTK,N7jZ,nR7P].

We have updated the paper to incorporate constructive suggestions, which will show in the revision. We summarize the major changes below:

1. [**nR7P**] **Throughput on Real-World Datasets**:
    - We add  H$_2$O throughput on real-world benchmarks that validate our efficiency claims. 
    - We add a discussion to explain why we use simulated settings to assess the throughput improvements, (1) long-form generation benchmarks are rare and still we want to have a wider coverage of system evaluation, (2) our evaluation is independent of existing datasets and could be useful as a unique metric to evaluate future kv cache policies.   

2. [**nR7P**] **Stronger Baseline**:
    - We add 0/1 shot as an additional strong baseline, and show that 20%- H$_2$O with 5-shot inputs has up to $9$% better performance across tasks than this baseline. 

3. [**AZfb**] **Diversity of Generation Content**:
    - We add a more formal evaluation on the increased diversity by Self-BELU and show that the generated output by  H$_2$O are more diverse than full model, with **lower** self-similarity in terms of BELU (0.0051 vs. 0.0057). 

4. [**CXTK, N7jZ**] **Application to Encoder-Decoder / COLT5**:
    - We add H2O results on T5 and show that using **20%** KV Cache can also preserve the performance, similar to auto-regressive cases. 
    - We present a discussion on the relation between COLT5 and  H$_2$O, (1) they are different approaches to solve a similar problem (improved efficiency by focusing on important tokens), and conceptually H2O can be applied on the top of COLT5 to further improve its performance; (2) while we do not have access to COLT5, we test  H$_2$O on T5 and show improvement, supports our analysis. 

5. [**AZfb**] **Existence of Sparsity and Layer-wise Customization**:
    - We add related work that supports attention sparsity observation and extended similar sparsity (more than 90%) observations on LLaMA and GPTNeo-X.
    - We present preliminary results that layer-wise KV Cache eviction policy according to sparsity can potentially further boost  H$_2$O accuracy by 1 point.  

6. [**CXTK, N7jZ**] **Extension to Training**:
    - We add preliminary results for deploying H2O in training on WikiText-103 and show that it could potentially lead to better generalization. 

7. [**nR7P**] **Zero / one shot results**:
    - We add zero/one shot results and show that some tasks maintain accuracy under 20% budget and others under 30-40% budget. Generally our  H$_2$O is effective compared to other baselines in this setting. 

8. [**CXTK**] **Memorizing transformer**:
    - We add preliminary results on deploying  H$_2$O in memorizing transformers and the positive signals show the possibility of reducing the external memory size of memorizing transformers by introducing heavy-hitter tokens.  


Meta-review: The paper addresses the issue of improving the efficiency of LLMs during inference by introducing a novel pruning strategy called H2O. Firstly, the authors present a simple and intuitive approach that is empirically demonstrated to preserve performance across various model architectures and benchmarks in the zero-shot setting. This approach has the potential for significant practical implications in research engineering work involving LLMs. The paper's clear and well-structured presentation is commended, making it easy to follow.

On the other hand, reviewers have highlighted certain weaknesses that need to be addressed. The claim of increased diversity in generations due to H2O is not adequately supported by a small qualitative evaluation and should be reinforced by more rigorous quantitative measures or human studies. Furthermore, the paper should provide citations to related work on sparsity in attention matrices, as the observation of sparsity is not presented as a novel finding. The inconsistency in evaluation between synthetic/toy data and real-world datasets raises concerns about the practical applicability of the proposed method. To address this, the paper should include throughput results on real-world datasets to validate the efficiency claims. Additionally, reviewers have raised questions about the application of H2O during training, its potential use in encoder-decoder frameworks, and its applicability to reducing memory size in Memorizing Transformers. 

In response to these comments, the authors have made important revisions to the paper. They have added throughput results on real-world benchmarks, introduced a stronger baseline, conducted formal evaluations of increased diversity, explored the application of H2O in training, and provided preliminary results for its use in Memorizing Transformers. These changes address many of the concerns raised by the reviewers and strengthen the paper's overall contribution.
Decision: Accept (poster)