SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Transformer in NLP. Transformers [67] as a popular option have been frequently adopted by plenty of natural language processing (NLP) applications with prevailing successes [68, 69, 70, 71, 72, 46, 73, 13, 74, 75]. Roughly, modern transformer-based networks can be categorized into two 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 to capture the intrinsic relationship between words and their context. Notable examples include BERT [76], RoBBERTa [69] and T5 [77]. (2) Decoder-only (i.e., GPT-style models [78]). Usually, this group of transformers adopts the Casual Language Modeling task, which is optimized to generate the next word-token in a sequence based on the preceding words/tokens. Such an autoregressive manner is highly preferred by downstream tasks like text generation and question answering. GPT-3 [79], OPT [39], PaLM [13], and BLOOM [80] are representative architectures within this huge family.
Training of Transformer. Training a gigantic transformer-based model is not trivial. It notoriously suffers from various issues such as overfitting, instability, etc. [81, 82, 83] analyze these bottlenecks from the optimization perspective. To address the issues, a great amount of pioneering effort is devoted, including data augmentations [84, 85], a better initialization [86, 83, 87, 88], customized optimizers [89], improved normalization [90, 91], weight decay [92], and early stopping. However, there is still a long way to go before we can fully clarify the mystery of transformer training.
B.2 Discussions and Limitations
Previous Attempts. During our experiments, we find several noteworthy observations. In  $\mathsf{H}_2\mathsf{O}$ , 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 a higher number of attention scores, resulting in a higher accumulated attention score and, consequently, a greater likelihood of being retained. To address this concern, we conducted an additional experiment utilizing
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  $\mathsf{H}_2$  occurrences at the beginning of sentences. This finding suggests that the initial tokens play a 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 is a significant bottleneck in LLMs. The proposed Heavy Hitter Oracle  $(\mathsf{H}_2\mathsf{O})$  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 a source of inspiration for future advanced algorithm designs. We envision  $\mathsf{H}_2\mathsf{O}$  as a foundational framework that could facilitate further innovation in this area. Moreover, long content generation is an area of growing importance that currently grapples with several efficiency issues. We hope our work that supports the generation of very long sequences will support further research in this direction, particularly in terms of enhancing consistency, devising superior evaluation methods, and establishing robust benchmarks.
Furthermore, another contribution of this study is the formulation of a dynamic submodular framework. 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 a submodular problem with slight variations at each time.
Limitations. Furthermore, despite the notable advancements in throughput of our  $\mathsf{H}_2\mathsf{O}$ , implementing LLMs for generative inference remains challenging due to the immense parameter count. As a substantial portion of these parameters is encompassed within the MLP blocks, building upon our observations of  $\mathsf{H}_2$  occurrences in the MLP blocks, future research efforts can be directed towards leveraging the characteristics of  $\mathsf{H}_2$  to devise an offloading policy. Such a policy can potentially enhance the efficiency of LLM inference even further.
C Extended Experiments
In this section, our goal is to demonstrate that  $\mathsf{H}_2\mathsf{O}$  can improve the diversity of generated text (mentioned in Section 5.3), throughput is further improved when combined with quantization (mentioned 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  $\mathsf{H}_2$  are reported, including experiments under zero-shot/one-shot inference regime;  $\mathsf{H}_2\mathsf{O}$  can also enhance the "Top-K" baseline; extra results of  $\mathsf{H}_2$  in attention blocks; the emergence of  $\mathsf{H}_2$  in MLP blocks as well as its properties.
C.1 Increased Diversity of Generated Text
Model Input
In the year 2087, humanity has achieved remarkable technological advancements and established colonies on multiple planets within the Milky Way galaxy. Interstellar travel has become commonplace, with faster-than-light spacecraft enabling people to explore distant star systems. Earth has undergone significant changes due to sustainable development efforts, such as harnessing renewable energy sources and implementing widespread ecological restoration projects. However, alongside these triumphs, new challenges have emerged, including the rise of artificial intelligence, ethical dilemmas surrounding genetic engineering, and interplanetary political tensions. Against this backdrop, a team of intrepid scientists embarks on a mission to uncover the secrets of an ancient alien civilization, hidden deep within an uncharted exoplanet. As they navigate treacherous terrains and encounter otherworldly phenomena, they must confront their own fears and reconcile humanity's thirst for knowledge with the potential consequences of uncovering secrets that were better left buried. The fate of both their mission and the future of humanity hang in the balance.
OPT-6.7B Full Cache Output
The game is set in 2087, a few years after the events of the first game. The player takes the role of a team of scientists who are sent to an uncharted planet to investigate a mysterious signal. The team is composed of a pilot, a scientist, a mechanic, and a technician.
OPT-6.7B Local 20% Cache Output
OPT-6.7B  $H_2O$  20% Cache Output
The game is set in the year 2087, and is a first-person exploration game. The player takes on the role of a scientist who has been sent to a distant planet to investigate a mysterious signal. The planet is inhabited by a race of intelligent beings known as the "Titans," who have been
Figure 6: Visualization of one generation example with OPT-6.7B. Results are compared between the baseline model with full cache, our  $\mathsf{H}_2\mathsf{O}$ , 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  $\mathsf{H}_2\mathsf{O}$ , and the "Local" strategy. Results are reported in Figure 6 and 7. Even with less KV cache budget, our  $\mathsf{H}_2\mathsf{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  $\mathsf{H}_2\mathsf{O}$  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 to 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  $\mathsf{H}_2\mathsf{O}$  describes both the people and the environment.
Moreover, We conducted a quantitative comparison via diversity metric(Self-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  $\mathsf{H}_2\mathsf{O}$  and the local method achieve 0.0051 and 0.0436, respectively. The comparatively lower Self-BELU score of  $\mathsf{H}_2\mathsf{O}$  indicates the slightly increased diversity of generated text.
In a small, bustling cafe nestled in the heart of a vibrant city, a serendipitous event unfolded, leaving a lasting impression on all who witnessed it. As the patrons sat sipping their coffees and engaging in animated conversations, a talented street musician entered the cafe, carrying a weathered guitar and radiating an aura of creativity.
LLaMA-7B Full Cache Output
He began to play, and the patrons were captivated. The musician's performance was so moving that the patrons began to 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 Local 20% Cache Output
He ((((( ((((( ( [)), 1999, 2000, 2001, 1, and, and, and, and, and, and, and, and, and, and,
LLaMA-7B  $\mathsf{H}_2\mathsf{O}$  20% Cache Output
He began 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 was a young
Figure 7: Visualization of one generation example with LLaMA-7B. Results are compared between the baseline model with full cache, our  $\mathsf{H}_2\mathsf{O}$ , 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
H2OW. Quant-4bit	84.00	43.20	78.80
C.2 Throughput Improvement of  $\mathsf{H}_2\mathsf{O}$  Combining with Quantization
Table 6 demonstrates  $\mathsf{H}_2\mathsf{O}$  are capable of quantization and perverse the full model's performance with even 4-bit quantization. To further explore the compatibility of our  $\mathsf{H}_2\mathsf{O}$  with quantization with respect to throughput improvement, we conduct an additional evaluation with the quantization method implemented in FlexGen (Note that [94] employed a different 4-bit quantization method). The corresponding results are presented in Table 7. Notably, for OPT-6.7B, we observed extra performance enhancements in  $\mathsf{H}_2\mathsf{O}$  when utilizing quantization compared to the vanilla version. This improvement results from the GPU memory freed by weight quantization, which allows for a 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  $\mathsf{H}_2\mathsf{O}$  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 an optimized CUDA kernel.
Table 7: Generation throughput (token/s) with different systems without offloading. We use  $\mathsf{H}_2\mathsf{O}$ -c denotes the  $\mathsf{H}_2\mathsf{O}$  with 4-bit weights compression. In the sequence length row, we use  $512 + 32$  to denote a prompt length of 512 and a generation length of 32. "OOM" means out-of-memory. The gray text in the bracket denotes the batch size. We run OPT-6.7B on a 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	OOM	OOM	OOM
FlexGen	20.2 (2)	16.8 (1)	16.9 (1)
H2O (20%)	35.1 (4)	51.7 (4)	52.1 (4)
H2O-c (20%)	50.5 (70)	72.5 (52)	62.3 (44)
C.3 Effectiveness on Zero-shot and One-shot Inference
When facing zero-shot and one-shot inference tasks,  $\mathsf{H}_2\mathsf{O}$  can successfully mitigate the memory requirements of the KV cache by up to  $5\times$  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.
Figure 8: Comparison results on zero-shot and one-shot inference.
C.4 Comparison with StreamingLLM
Recent works, such as StreamLLM [52] 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:
Multi-document question answering [95]: In this task, each test sample comprises ten documents, followed by a question. The key information to answer this question is stored in one of the documents. We rearranged the crucial document's position and found that the eviction strategy in StreamLLM or LM-Infinite can not perform well when the key document has been dropped.
Text Summarization Task (XSUM and CNN-DailyMail): Text summarization tasks require models to generate concise summaries of lengthy texts. Effective summarization demands a comprehensive understanding of the entire document, making it challenging when crucial information is dispersed throughout the input. In particular, summarization often relies on long-context attention, and critical information may not be effectively captured within the limited local tokens.
The results are reported in Figure 9, illustrating a consistent decline in the performance of StreamLLM. Since StreamLLM will always maintain the first few tokens as well as the local tokens, regardless of various input content, such a strategy will inevitably result in the loss of crucial information and subsequently lead to a decrease in performance. In contrast, our  $\mathsf{H}_2\mathsf{O}$  delivers markedly superior performance. Also, with  $\mathsf{H}_2\mathsf{O}$ , The model can successfully stream to four million tokens.
Figure 9: Comparison results of StreamLLM [52] and our  $\mathsf{H}_2\mathsf{O}$  on generalization tasks. The number in each method represents the KV Cache budget of the start/heavy-hitter tokens and the local tokens, respectively. For example, H2O-256-256 means maintaining 256 Heavy-Hitters and 256 local tokens.
C.5 Enhancing the "Top-K" Baseline
We find  $\mathsf{H}_2$  can further enhance another strong baseline with a "Top-K" strategy. The results are reported in Table 8. After combing with  $\mathsf{H}_2$ , the "Top-K" method achieves an extra improvement with up to  $2.00\%$  accuracy across 4 different tasks.
Table 8: Results of the "Top-K" method w. or w.o.  $\mathsf{H}_2$ . 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-Hitters and the local tokens. And we can observe Heavy-Hitters contribute more to maintaining the performance of models.
C.7 Performance of Inference with Different Number of Shots.
Table 10 demonstrates our  $\mathsf{H}_2\mathsf{O}$  achieves consistent improvements across different number of shots during inference and maintains the full model's performance with  $20\%$  memory budget.
Table 9: Ablation study of  ${\mathrm{H}}_{2}\mathrm{O}$  across different tasks.
Tasks	Models	Full	w. Local	w. H2	w. Local + H2
PiQA	OPT-13B	77.37	54.62	76.12	77.26
OPT-30B	78.51	55.82	67.25	78.45
OpenBookQA	OPT-13B	41.40	25.60	30.40	41.20
OPT-30B	43.20	25.20	26.60	43.00
MathQA	OPT-13B	26.67	22.04	23.82	26.93
OPT-30B	26.23	20.87	21.98	26.87
Winogrande	OPT-13B	68.59	49.96	51.85	67.32
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	5-shots	10-shots
OPT-30B	OPT-66B	OPT-30B	OPT-66B
OpenBookQA	Full	43.20	44.40	43.00	44.80
Local	25.20	30.60	26.60	38.80
H2O	43.00	44.20	42.80	44.80
COPA	Full	85.00	83.00	86.00	85.00
Local	48.00	59.00	60.00	76.00
H2O	84.00	82.00	85.00	86.00
MathQA	Full	26.23	27.87	26.67	27.00
Local	20.87	25.49	21.11	23.08
H2O	26.87	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 10. We can observe that  $\mathsf{H}_2$  broadly exists in each layer.
C.9 Comparison with SpAtten
Compared with SpAtten [32], the main differences of  $\mathsf{H}_2\mathsf{O}$  are i) They accumulate attention scores across attention heads and layers, while in our algorithm, each token can be kept or evicted independently 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  $\mathsf{H}_2$  tokens can effectively enhance other sparse attention strategies; iii) we formulate the KV cache eviction as a 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 11.
Table 11: Comparison between SpAtten [32] and  $\mathsf{H}_2\mathsf{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
H2O	84.00	43.00	78.45
C.10 Heavy-Hitter in MLP Blocks
Besides the attention blocks, the presence of Heavy-Hitters  $(\mathsf{H}_2)$  is observed within the MLP blocks of LLMs. We utilize the Wiki-Text-103 dataset as the input and record the activated frequency of neurons in the hidden layer of MLP blocks. As depicted in Figure 11, the activated frequency of neurons follows a power-law distribution, wherein a small number of neurons are activated by nearly all input tokens (with a  $100\%$  frequency) while the majority of other neurons are rarely activated.
Subsequently, a thorough examination of various characteristics pertaining to  $\mathsf{H}_2$  in MLP blocks is conducted, encompassing the following aspects: (1) The elimination of  $\mathsf{H}_2$  leads to a substantial decline in performance, although such degradation can be easily recovered even with a mere  $1\%$  of the training data; (2)  $\mathsf{H}_2$  exhibits a significant degree of overlap across different type of input content; (3) The emergence of  $\mathsf{H}_2$  occurs early in the training process, thus exhibiting an "early-bird" characteristic, and their positions undergo gradual changes during subsequent training phases.
Figure 10: The distribution of accumulated attention scores with respect to the corresponding word. The x-axis represents the word index in the vocabulary, and the y-axis represents the accumulated attention score. Results are obtained from OPT-1.3B.
Figure 11: The emergence of  $\mathsf{H}_2$  in MLP blocks of 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.
Elimination of  $\mathsf{H}_2$ . 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.,  $\mathsf{H}_2$ ). This pruning operation leads to a substantial decline in performance, as evidenced by an increase in perplexity from 19.32 to 31.78. The results emphasize the criticality of  $\mathsf{H}_2$  in preserving the functionality of the model. To assess the recoverability of the discarded information, we conduct a 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 of the pre-trained model. In contrast, when training the model from scratch using only  $1\%$  of the training data, the resulting model achieves a perplexity of 554.12 only. These findings demonstrate that the knowledge encoded in  $\mathsf{H}_2$  can be easily restored.
Table 12: . Perplexity on the test-set of Wiki-Text-3 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 a 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  $\mathsf{H}_2$  is depicted in Figure 12, revealing significant concurrence across multiple datasets.
Early-Bird Property. Furthermore, our investigation reveals that  $\mathsf{H}_2$  displays an "early-bird" characteristic, as illustrated in Figure 13. By visualizing the distribution of activation frequencies across various checkpoints throughout the training process, we observe the emergence of a power-law behavior at an initial stage, specifically as early as a training budget of  $4\%$ . Subsequently, the positions of  $\mathsf{H}_2$  exhibit gradual and minimal changes.
Figure 13: The distribution of activated frequency during training. Experiments are conducted with different checkpoints of OPT-2.7B during training.
Figure 12: The distribution of activated frequency across diverse input content. The x-axis represents the index of neurons, which is ordered by the activated frequency from Wiki-Text-103.
D Theoretical Analysis
Recently, a number of works have studied the attention scheme in LLMs from a theoretical perspective [96, 97, 98, 99, 100, 101, 102, 103, 104, 105, 106, 107, 108, 109, 110, 111, 112, 113, 114, 115, 116, 117, 118, 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  $\mathsf{H}_2$ , which is a modification of dynamic submodular maximization. Furthermore, assuming the attention scheme to be submodular, we establish that constructing the set  $S_{i}$  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 1. We also offer a mathematical formulation for sparsity preservation that is observed in Section 3 and proposed an algorithm (Algorithm 4) to solve the problem.