Context Compression for Auto-regressive Transformers with Sentinel Tokens
Authors: Siyu Ren
Qi Jia
Kenny Q. Zhu
Venue: EMNLP 2023 Main (2023)
Keywords: context compression
key-value cache compression
TL;DR: context compression

Abstract: The quadratic complexity of the attention module makes it gradually become the bulk of compute in Transformer-based LLMs during generation. Moreover, the excessive key-value cache that arises when dealing with long inputs also brings severe issues on memory footprint and inference latency. In this work, we propose a plug-and-play approach that is able to incrementally compress the intermediate activation of a specified span of tokens into compact ones, thereby reducing both memory and computational cost when processing subsequent context. 
Experiments on both in-domain language modeling and zero-shot open-ended document generation demonstrate the advantage of our approach over sparse attention baselines in terms of fluency, n-gram matching, and semantic similarity. At last, we comprehensively profile the benefit of context compression on improving the system throughout. Code is available at \url{https://github.com/DRSY/KV_Compression}.

Context Compression for Auto-regressive Transformers with Sentinel
Tokens
Siyu Ren Qi Jia
Kenny Q. Zhu"
Shanghai Jiao Tong University, China {roy0702, Jia_qi jesjiutssjuueeeee edu.cn
University of Texas at Arlington, USA kenny www.uta.euaaeeeee
Abstract
CUDA kernels typically required for these mod- els achieve performance comparable to off-the- shelf LLMs and fulfill real efficiency gains.
The quadratic complexity of the attention mod- ule makes it gradually become the bulk of com- pute in Transformer- ation. Moreover, the excessive key-value cache that arises when dealing with long inputs also brings severe issues on memory footprint and inference latency. this work, we propose plug-an -nd--ndd-ll approach that is able to in- crementally compress the intermediate activa- tion of specified span of tokens into compact ones, thereby reducing both memory and com- putational cost when processing subsequent context. Experiments both in-domain lan- guage modeling and zero-shot open-ended doe- ument generation demonstrate the advantage of approach over sparse attention baselines in terms of fluency, n-g n-gram matching, and se- mantic similarity. At last, we comprehensively profile the benefit of context compression improving the system throughout. Code is available https:/ //gi thub. com/DRSY/KV. Compression.
this work, we aim to improve the efficiency of existing Transformer-basee LLMs without any architectural changes. Specifically, we focus on the key-value cache, which accounts for the major- ity of memory footprint and data movement (1/O) cost when dealing with increasingly long input ing LLMs. We propose plug-and-play approach to incrementally compress the key-value memory of a contiguous span of tokens into compact ones. Specifically, we introduce a pair of special sentinel tokens <CL> and <CR> into the vocabulary LLMs and use them to mark the boundary of the span to be compressed. During training, we modify the causal attention mask such that future tokens after <CR> precluded from attending tokens between <CL> and <CR>, By continually training LLMs with the next token prediction objective, the model learns to extract and condense task-relevant information of the bounded span into the ending sentinel token. The reduced context length allevi- ates both memory and computational costs when processing subsequent tokens, thereby improving system throughput with larger batch sizes and faster decoding speed during inference.
LLMss during gener-
1 Introduction
The Transformer architecture (Vaswani et al., 2017) has become the underpinning component of mod- ern large language models (LLMs) (Radford et al., 2019; Devlin al., 2019; Zhang et al., 2022; Tou- vron al., 2023) in recent years. However, the quadratic computational complexity and memory footprint of the attention mechanism have largely limited applying Transformers to increasingly long contexts with constrained computing resources.
We conduct experiments on the WikiText-2 lan- guage modeling benchmark and show that our approach is generalizable to LLMs with various sizes (from 1.3B to 3B) and position encoding schemes such as absolute position embedding (De- vlin et al., 2019) and rotary position encoding (Su et al., 2021). Compared to sparse attention base- lines, our approach is able to effectively compress historical key-value cache with significantly duced degradation in perplexity. Moreover, we demonstrate that our approach outperforms sparse attention in zero-shot open-ended document gener- ation across different compression ratios as eval- uated by perplexity, ROUGE (Lin, 2004), and
To mitigate such issues, prior works (Child et al., 2019; Beltagy et al., 2020; Zaheer et al., 2020; Ki- taev et al., 2020) have explored assortment of efficient Transformer variants, mostly by replacing the original quadratic attention operation with var- ious forms of linearized approximation. Though promising, large-scale scale pre-training and specialized
TTee corresponding author.

---

BERTSScrr (Zhang et al., 2019). Finally, we em- pirically demonstrate that context compression is able to confer considerable improvement in system throughput for text generation.
A
cat
<CL>
jumps over
<Gl>
carpet
cat
<CL>
jumps
2 Background and Related Work
<CR>
In this section, we present necessary background knowledge as well drawing distinctions of our approach existing literatures.
carpet
Figure I: Modified attention mask for context compres- sion. Green and grey boxes indicate positions that are allowed and disallowed be attended (attention di- rection from row - column).
Complexity of Attention Mechanism The self- attention mechanism of LLMs produces a query vector for each token to select and retrieve informa- tion from all previously computed key and value vectors. Thus, key-value memory is required at runtime store past information. Given pre- trained LLM with M Transformer layers, H atten- tion heads, per-head hidden dimension of dhead- batch size and current context length L, the model stores key and value vectors computed so far as tensor of shape (M, 2, b, H, L, dread). Dur- ing auto-regressive decoding, the size of key-value cache grows linearly with the context length L, leading to significant increase in memory footprint and latency.
for a compression scheme with more flexible com- pression choices and larger compression ratios.
Approach
In this section, we present a plug-and-play ap- proach to compress the full-length contextual rep- resentations into shorter ones while preserving in- formation necessary accomplish the end task.
3.1 Context Compression with Sentinel Tokens
Efficient Transformers Various efficient vari- ants of Transformers have been proposed to address the problematic complexity of attention. Sparse Transformer (Child et al., 2019) limits the recep- tive field of each token local window. Long- former (Beltagy al., 2020) and BigBird (Zaheer al., 2020) introduce additional randomly and globally accessible tokens to compensate for the information loss. Linear Transformer (Katharopou- los et al., 2020) reformulates the self-attention as a linear dot- dot-product of kernel feature maps and make use of the associativity property of matrix products to reduce the complexity from O(N²) to O(N). Nevertheless, these approximations were shown degrade the expressiveness of full-attention and have been upstaged by GPT-like LLMs recently. In this work, however, we tackle this problem by com- pressing lengthy input into more compact represen- tations, which is orthogonal to on-going efforts to architectural optimization.
To alleviate the computational and memory inten- sity of key-value memory when processing increas- ingly long context, we introduce two sentinel to- kens <CL> and <CR> into the vocabulary of LLM, that are placed around contiguous span of tokens of which the corresponding key-value memory will be compressed.
In order to consolidate information tokens surrounded by <CL> and <CR>, we design a modified causal attention mask to facilitate such behaviour. An illustrative example is shown in Figure 1. Specifically, <CL> serves start-of- compression symbol and can only attend to it- self. Normal tokens (tokens except for <CL> and <CR>) have access all previous <CR> and nor- mal tokens. This ensures that the contextual repre- sentations are built upon both compressed (lossy) and complete (lossless) past information. <CR> then act as a form of information selection and retrieval from contextual representations of sur- rounded tokens. This modified masking scheme, combined with task-specific fine-tuning of LLM, encourages distilling task-relevant information of potentially long token sequences into compact representation, thereby reducing the size of key- value memory required for subsequent processing.
Gisting Tokens Mu et al. (2023) explored using gisting tokens compress textual prompts during instruction tuning. They showed that verbose task instructions can be compressed into much shorter ones. However, their approach is only applicable to prefix text around 20 tokens. In contrast, we aim

---

Compression Ratio (r)
Model
Method
0.2
0.4
0.7
0.9
Scattered Attention
15.0
18.2
22.6
28.2
35.8
47.2
59.6
81.0
106.4
151.6
OPT-1.3B
Local Attention
15.0
15.1
15.2
15.7
16.3
17.2
18.2
19.9
23.0
29.8
KV Compression
15.0
15.0
15.2
15.6
15.9
17.3
17.8
18.0
18.1
18.3
Scattered Attention
13.1
16.2
19.8
25.5
32.4
41.7
56.3
75.7
101.4
142.3
OPT-2.7B
Local Attention
13.1
13.1
13.4
13.9
14.4
15.3
16.1
17.5
20.1
26.6
KV Compression
13.1
13.2
13.5
13.7
14.0
15.3
15.8
16.0
16.1
16.3
Scattered Attention
11.2
13.1
16.5
20.7
26.6
34,9
47.2
65.7
91.6
135.9
RedPajama-3B
Local Attention
11.2
11.3
11.5
11.8
12.2
12.9
13.6
14.8
17.2
22.9
KV Compression
11.2
11.4
11.5
11.8
11.9
12.3
13.5
13.7
13.8
14.0
Table 1: Perplexity (the lower the better) of three LLMs WikiText-2 language modeling benchmark.
Input Transformation Given an input æ with tokens, we define compression ratio r as the percentage of the tokens to be enclosed by one or more pairs of <CI> and <CR>, and max compres- sion length 1 as the largest number of tokens to be compressed by a single pair of <CL> and <CR> tokens. We repeatedly sample span of tokens with length conforming to a uniform distribution (2, 1) as compression target, until ratio r reached. But these spans must not overlap. We tie the position encoding of <CI.> and <CR> to their previous to- ken, such that they don't occupy valuable positions for LLMs using absolute position embedding.
Compression Ratio (r)
Method
0.7
0.8
0.9
Local Attention
23.1
25.2
KV Compression
21.8
22.0
Table 2: Perplexity (the lower the better) of RedPajama- 3B on WritingPromppss test set.
4 Experiments
We mainly evaluate our approach on language mod- eling task and then explore its zero-shot general- ization ability on open-ended document generation. Finally, quantitatively measure the effect of context compression on system throughput.
Training We opt for lighht procedure for adapting LLMs to downstream tasks through fine- tuning. Specifically, all parameters of the LLM are frozen except for the embeddings of <CL> and <CR> and LoRA (Hu et al., 2021) modules applied to all attention layers. This not only eases the train- ing process of models with billions of parameters but also ensures the inherent language modeling ability and parametric knowledge in feedforward networks are intact (Geva et al., 2021).
4.1 Language Modeling
Benchmark We mainly conduct experiments the Wikitext-2 (Merity et al., 2016) dataset consist- ing of Wikipedia articles, which has been widely used for the evaluation of language modeling. We report token-level (excluding <CL> and <CR> for KV Compression) average perplexity on the test set as the evaluation metric. We also report results on the WritingPrompts (Fan et al., 2018) dataset to investigate the generality of our method.
Models To demonstrate the generality of our approach, dubbed KV Compression, we adopt OPT-1.3B, OPT-2.7B (Zhang et al., 2022), and RedPajama- (Computer, 2023) three LLMs with different sizes and position encoding methods.
Manipulating Key-Value Memory for Inference- time Compression Given a text piece as prefix, we apply the same input transformation strategy de- fined above to reach a specified compression ratio. To realize perceivable memory and computation reduction, the key-value cache of tokens enclosed by sentinel tokens is freed from GPU in progres- sive manner (e.g., using for-loop over blocks of tokens) if the original prefix is too long fit into GPU. Otherwise we just feed the transformed pre- fix through the model and free the kayuv cache marked tokens in one go.
We compare our approach with two sparse attention baselines: (1) Scattered Attention which samples positions that permits attention from Bernoulli distribution B (1-r); and (2) Local At- tention (Beltagy et al., 2020) which restricts the attention of token 2; to be within ([t rr,t). t). For KV Compression, we set as 25 throughout the ex-
Baselines

---

17
76.0
tYoughouss Took-aRRe send
250
75.0
47.5
15
74.5
42.5
Local Aitention KV Comprestininuro
160
73.5
35.0
140
75.0
0.0
0.7
0.4
0.8
0.8
0.2
11.4
0.8
Compressin
Compression
Figure 2: RedPajama-3B on open-ended generation on 200 sampled C4 documents. Generation quality is measured by fluency (perplexity), n-gram matching (ROUGE-L), and semantic similarity (BERTScore). We report system throughput as the number of tokens generated per second with different maximum GPU VRAM.
periments. Our proposed approach, alongside the sparse attention baselines have selected, shares common goal: enhancing the ratio of compu- tation to memory access by curtailing the stored key-value cache. For detailed training setup please refer to Appendix A.
open-ended document generation. Given prefix text p. nucleus sampling (Holtzman et al., 2019) is used to generate completion for it.
Baselines Because Scattered Attention performs poorly according to Table 1, we only compare KV Compression with Local Attention with compres- sion ratio r ranging from 0.0 to 0.9 applied to the prefix p using input transformation defined in Sec- tion 3 for KV Compression and restricted attention defined in Section 4. For Local Attention to achieve inference-time compression, amounts main- taining FIFO queue to store the key-value cache: as the time step during generation increases, old key-value memory in the queue is popped out and newly generated key-value memory is pushed in.
The results WikiText-2 shown in Table 1. As the compression ratio gets larger, all three methods result in an increase perplex- ity either due to: (I) important tokens are out-of- scope for attention, or (2) the capacity of single <CR> token is insufficient to encapsulate the full information of compressed token span. We observe that Local Attention substantially outperforms Seat- tered Attention, suggesting the importance of local context for language modeling. KV Compression achieves the best perplexity across various com- pression ratios. Notably, at high compression ra- tios, e.g.. r 2 0.7, KV Compression incurs signif- icantly fewer degradation perplexity compared to Local Attention, demonstrating its advantage in compressing scattered local information mean- while keeping coherent global information. In Ta- ble 2, report the perplexity of RedPajama-3B on WritingPrompts dataset using Local Attention and the proposed KV Compression, with compres- sion ratio from (0.7, 0.8, 0.9}. KV Compression achieves consistently lower perplexity compared to Local Attention, indicating its superiority as Domain-generalizzn method for context compres- sion.
Results
Evaluation Metries We evaluate the quality of generated completions from three aspects: (1) fluency, (2) n-gram matching with ground-truth continuation, and (3) semantic similarity with ground-truth continuation. Fluency is evaluated by perplexity computed from Pythia- 1B (Biderman et al., 2023) model pre-trained using C4. N-gram matching and semantic similarity are measured by ROUGE-L and BERTScore (Zhang et al., 2019) re- spectively. To account for the randomness induced by nucleus sampling, for each prefix, we generate completions and report the average results.
Results Figure 2 shows that, as r increases, the generated completions of Local Attention tend to diverge from the original topic, leading to de- creased HOUGE-L/BERTSSE
4.2 Zero-shot Open-ended Generation
and increased per- plexity. Again, KV Compression excel at pre- serving relevant information and can still generate decent-quality continuations up to 0.5 compres- sion ratio. Notably, KV Compression can gener- ate fluent text even when the prefix extremely compressed. The reason is that, during training, Local Attention receives rapidly faded information
Data We randomly select 200 documents from C4 (Raffel et al., 2020) validation set for evaluation. We use the leading 128 words as prefixes and treat the next 64 words as ground-truth references.
Models We directly take the RedPajama-3B model trained on Wikitext-2 to perform zero-shot

---

from the distant past, making the discourse struc- ture for subsequent generations incoherent. On the contrary, KV Compression better preserve such in- formation by consolidating it into sentinel tokens.
put transformation procedure theoretically includes text spans of all possible linguistic structures, it may still benefit from an elaborately designed strat- egy/algorithn for selecting compression targets in a given context.
4.3 Throughput Gain from Context Compression
References
With KV Compression, the key-value cache corre- sponding to tokens enclosed by sentinel tokens can be freed from memory. In this way, it permits a larger batch size and improves the system through- put in return. To quantitatively measure the impact of context compression on throughput, we conduct experiments on open-ended generation by testing with different maximum available GPU VRAM. The full length of the dummy input prefix is set to 800 and we use RedPajama-3B with nucleus sam- pling to generate continuation for it. The through- put of the text generation system is defined as the number of tokens generated per second. The re- sults are shown in the bottom right of Figure 2. We can see that, at extreme compression ratios (e.g.. 0.8), context compression offers more than 1.5x throughout improvement with 12GB GPU VRAM and a slightly smaller 1.3x improvement with 24GB GPU VRAM. At moderate compres- sion ratio (e.g., r Re 0.5), KV compression still able to deliver 1. 1.2x-1.4x 1.4x throughout improvement while only suffering from mild quality drop (Sec- tion 4.2). More visualized comprrsoon ratio correlation is deferred to Appendix D.
Iz Beltagy, Matthew E. Peters, and Arman Cohan. 2020. Longformer: The long-document transformer.
Stella Biderman, Hailey Schoelkopf, Quentin Anthony, Herbie Bradley, Kyle O' O'Brien, Eric Hallahan, Mo- hammad Aflah Khan, Shivanshu Purohit, USVSN Sai Prashanth, Edward Raff, Aviya Skowron, Lintang Sutawika, and Oskar van der Wal. 2023. Pythia: A suite for analyzing large language models across training and scaling.
Rewon Child, Scott Gray, Alec Radford, and Ilya Sutskever. 2019. Generating long sequences with sparse transformers. CoRR, abs/1904. 10509.
Together Computer. 2023. Redpajama: An open source recipe to reproduce Ilama training dataset.
Jacob Devlin, Ming-Was Wei Chang, Kenton Lee, and Kristina Toutanova. 2019. BERT: Pre-training of deep bidirectional transformers for language under- standing. Pruceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguisties: Human Language Tech- nologies, Volume (Long and Short Papers), pages 4171-4186, Minneapolis, Minnesota. Association for Computational Linguistics.
Angela Fan, Mike Lewis, and Yann Dauphin. 2018 Hierarchical neural story generation. In Proceedings of the 56th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pages 889 898, Melbourne, Australia. Association for Computational Linguistics.
5 Conclusion
In this work, we propose simple yet effective approach that enables LLMs summarize the key- value memory of specified span of tokens. Experi- ments on language modeling and open-ended gen- eration demonstrate that our approach substantially outperforms sparse attention baselines in terms of information compression and generation quality.
Mor Geva, Roei Schuster, Jonathan Berant, and Omer Levy. 2021. Transformer feed forward layers are key- value memories. In Proceedings of the 2021 Confer- ence Empiricul Methods in Natural Language Pro- cessing, pages 5484-5495, Online and Punta Cana, Dominican Republic. Association for Computational Linguistics.
Ari Holtzman, Jan Buys, Maxwell Forbes, and Yejin Choi. 2019. The curious case of neural text degener- ation. CoRR, 199751.
Limitations
In the current evaluation setup, we apply the same strategy used for training to select spans of tokens to be enclosed by <CL> and <CR> for inference- time context compression. However, different text pieces may display different degrees of importance for downstream tasks. For instance, grammatical and semantic complete noun phrase can be more compressible than an ungrammatical one that con- tains only partial linguistic units. Though our in-
Edward J. Hu, Yelong Shen, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li. Shean Wang. and Weizhu Chen. 2021. Lora: Low-rank adaptation large language models. CoRR, abs/2106.090 006809.55...
Angelos Katharopoulos, Apoorv Vyas, Nikolaos Pap- pas, and François Fleuret. 2020. Transformers are mns: Fast autoregressive transformers with linear attention. In International Conference on Machine Learning, pages 5156-5165. PMLR.

---

Nikita Kitaev, Lukasz Kaiser, and Anselm Levskaya. 2020. Reformer: The efficient transformer. CoRR, 04/200110445/1
Susan Zhang, Stephen Roller, Naman Goyal, Mikel Artetxe, Moya Chen, Shuohui Chen, Christopher De- Mona Diab, Xian Li, Xi Victoria Lin, Todor Mi- haylov, Myle Ott, Sam Shleifer, Kurt Shuster, Daniel Simig, Punit Singh Koura, Anjali Sridhar, Tianlu Wang, and Luke Zettlemoyer. 2022. Opt: Open pre- trained transformer language models.
Chin-Yew Lin. 2004. ROUGE: package for auto- matic evaluation of summaries. In Text Summariza- tion Branches Out, pages 74 81, Barcelona, Spain. Association for Computational Linguistics.
Tianyi Zhang, Varsha Kishore, Felix Wu, Kilian Q Weinberger, and Yoav Artzi. 2019. Bertscore: Evaluating text generation with BERT. CoRR, abs/1904.09675.
Ilya Loshchilov and Frank Hutter. 2017. Fixing weight decay regularization in adam. CoRR,
abs/1711.05
Stephen Merity, Caiming Xiong, James Bradbury, and Richard Socher. 2016. Pointer sentinel mixture mod- els.
Jesse Mu, Xiang Lisa Li, and Noah Goodman. 2023. Learning to compress prompts with gist tokens. arXiv preprint arXiv:2304.08467.
Alec Radford, Jeff Wu, Rewon Child, D. Luan, Dario Amodei, and Ilya Sutskever. 2019. Language models unsupervised multitask leamerr.
Colin Raffel, Noam Shazeer. Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter Liu. 2020. Exploring the limits of transfer learning with unified text-to-text trans- former. The Journal of Machine Learning Research, 21(1):5485-551)
Jianlin Su, Yu Lu, Shengfeng Pan, Wen, and Yunfeng Liu. 2021. Roformer: Enhanced transformer with rotary position embedding. CoRR, abs/2 04.09864.
Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière. Naman Goyal. Eric Hambro, Faisal Azhar, Aurelien Rodriguez, Armand Joulin, Edouard Grave, and Guillaume Lample. 2023. Llama: Open and efficient foundation language models.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan Gomez, ukasz Kaiser, and Illia Polosukhin. 2017. Attention all you need. In Advances Neural Information Pro- cessing Systems, volume 30. Curran Associates, Inc.
Thomas Wolf, Lysandre Debut, Victor Sanh, Julien Chaumond, Clement Delangue, Anthony Moi, Pier- ric Cistac, Tim Rault, Rémi Louf, Morgan Funtowicz, Joe Davison, Sam Shleifer, Patrick von Platen, Clara Ma, Yacine Jernite, Julien Plu, Canwen Xu, Teven Le Scao, Sylvain Gugger, Mariama Drame, Quentin Lhoest, and Alexander M. Rush. 2020. Transform- ers: State--f-the-art natural language processing, In Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing: System Demonstrations, pages 38 38,4,, Online. Association for Computational Linguistics.
Manzil Zaheer, Guru Guruganesh, Avinava Dubey, Joshua Ainslie, Chris Alberti, Santiago Ontañón, Philip Pham, Anirudh Ravula, Qifan Wang, Li Yang, and Amr Ahmed. 2020. Big bird: Transformers for longer sequences. CoRR, hbb/2007.1406.

---

A Implementation Details
of tokens surrounded by sentinel tokens takes a specified ratio of the total number of tokens. Here we study the generalization performance of dif- ferent train-time compression ratios and explore the best practice. We visualize the perplexity of RedPajama-33 trained with different compression ratios at varying test-time compression ratios.
Language Modeling For language modeling ex- periments, we train OPT-1.3B, OFT-2.7B, and RedPajama-3B with the next token prediction ob- jective. The loss is computed on all tokens except for newly added <CI,> and <CR> tokens. For tokens right before <CL> and <CR>, adjust their ground-trut truth label to be tokens right after the corresponding <CL> and <CR> tokens.
Figure 3 illustrates the results. We can see that training with a higher compression ratio always sults in better generalization ability across various time compression ratios. The reason is that, at small compression ratio, language models can sometimes exploit limited local context and can still reliably predict the next token. In this case, the sentinel tokens are not guaranteed to acquire the desired context compression ability. When the majority of context is not allowed to be attended to, the sentinel tokens are forced cultivate enough compression ability in order to minimize the loss function.
The trainable parameters for all LLMs include two token embeddings of <CL> and <CR>, and LoRA modules applied to all attention layers. The rank of weight matrices of LoRA is set to 16. The percentage of trainable parameters takes about 5% of the total parameters of the LLM. We use A (Loshchilov and Hutter, 2017) optimizer with 2e-5 learning rate. The batch size is set to 12 and the maximum sequence length during training is set to 256 due to the limited computation bud- get. The implementation is based on Huggingface transformers (Wolf et al., 2020) library and all ex- periments are conducted using a single RTX 3090 GPU.
C Memory Usage with Varied Compression Ratio
Open-ended Document Generation For open- ended document generation, we nucleus sam- pling with top P-0.9, Due the randomness in- duced by sampling, we perform 8 times nucleus sampling for each compressed prefix and report the average evaluation metrics adopted in the main paper. This helps reduce variance and ensures a more reliable evaluation result.
16
Posk GPU
12
B5-5,
BS-1,
1-10
30
B Generalization Ability of KV Compression
0.0
0.1
0.2
0.3
0.4
0.5
0.7
0.8
0.9
Compression
Figure Peak cached memory profiled using Pytorch when using RedPajama-3B 3B to produce 100-token con- tinuation for variable-length prefixes.
9.4
We have shown that context compression con- fers decent improvement in system throughput, es- pecially for moderate-sized GPUs. Here we re- port detailed memory usage assuming a practical scenario similar to multi-turn dialogue: the prefix length (length of historical dialogue) gradually in- creases, and the model is asked to output response of roughly 100 tokens. To maximize the through- put of dialogue service, we assume the model is simultancously generating responses for multiple instances, i.e., a batch size larger than one.
0.5
0.6
Perploito
15.0
14.5
0.2
0.3
0.4
0.5 Compression Ratio
0.6
0.7
0.8
Figure 3: Generalization performance of different train- time compression ratios.
In our proposed KV Compression, an LLM is trained on transformed data in which the percent
The visualized memory usage is shown in Figure Context compression is able to reduce more than

---

Attention
Attention KV Compressiontouns
17
KV Compressiontours)
KV Compresioniours
16
BERTScore
15
73
14
72
45
13
40
12
11
35
0.0
0.2
0.4
0.6
0.8
0.0
0.2
0.4
0.6
0.8
0.0
0.2
0.6
0.8
Compression Ratio
Compression Ratio
Compression Ratio
Figure 5: OPT-1B on open-ended generation on 200 sampled C4 documents. Generation quality measured by fluency (perplexity), n-gram matching (ROUGE-L), and semantic similarity (BERTScore).
3GB of peak cached GPU memory. In practice, this translate to the feasibility of generating more tokens for a single instance and enabling more in- stances in parallel.
D Open-ended Generation Results of OPT
Figure 5 summarizes the results of OPT-IB on open-ended document generation. Compared to that of Redpajama-3B (Figure 2), the generation quality of OPT- IB is substantially lower. Com- paring different context compression methods, KV Compression performs uniformly better across all three evaluation metrics.


Reviews:
[{"review_id": "QypN8K2Lmh", "reviewer": "Reviewer_xYfS", "summary": "", "questions": "A. Why were these baseline chosen? Aren't there baselines to compare to from previous papers?\\\nB. Not clear how the context compression is done at inference time? Are random spans of tokens are chosen?\\\nC. Same question as B regarding training.\\\nD. Can this method play well with Flash Attention to improve inference even further? If no, how your method compare with flash attention?", "limitations": "", "rating": 3, "confidence": 3, "soundness": 3, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We appreciate the valuable feedback provided by the reviewers. Our responses to Reviewer's questions are listed below:\n\n**Q**: \"A. Why were these baseline chosen? Aren't there baselines to compare to from previous papers?\"  \n**A**: The inference process of LLMs is inherently memory-bound. In the context of auto-regressive decoding, the predominant latency bottleneck is not attributed to actual computations, but rather to the repetitive movement of an increasingly expansive key-value cache. This cache is transferred from the High Bandwidth Memory (HBM) to the on-chip computation unit. Our proposed approach, alongside the sparse attention baselines we have selected, shares a common goal: enhancing the ratio of computation to memory access by curtailing the stored key-value cache.\n\nWhile Multi-query attention primarily concentrates on diminishing the count of attention heads, an endeavor that complements our approach, FlashAttention, introduces IO-aware modifications to the query-key computation scheme. Notably, both these methods run parallel to our aim of compressing the key-value pairs. In a subsequent version of our work, due credit will be given to these two aforementioned papers through proper citations.\n\n**Q**: \"Not clear how the context compression is done at inference time? Are random spans of tokens are chosen?\"  \n**A**: During training, as stated in Section 3.1, we employ a random strategy to select spans of tokens to be compressed. We acknowledge that though more systematic span selection strategy may exist, randomly choosing compression target has sufficient flexibility and variability, thereby enabling language model to  compress diverse range of contextual information. At inference time, we currently adopt the same random strategy as training time.\n\n**Q**: \"Can this method play well with Flash Attention to improve inference even further? If no, how your method compare with flash attention?\"  \n**A**: The original Flash Attention implementation requires the attention mask matrix to be lower triangular, meaning that it is not directly compatible with our method due to the sparse attention mask. To combine our method with efficient attention implementation, we need to use more advanced kernels that support sparsity patterns in attention masks. One such implementation is sparse flash attention[1]. Xformers[2]' implementation of memory-efficient attention also supports an arbitrary attention mask, one can directly use it to replace the default attention implementation in HuggingFace.  \nReferences:  \n[1] Faster Causal Attention Over Large Sequences Through Sparse Flash Attention.  \n[2] xFormers: A modular and hackable Transformer modeling library.\n\n**Q**: \"No mention of open-sourcing implementation\"  \n**A**: We have uploaded the source code in the supplementary material, including the main Python files and Bash scripts. We will open-source it after the paper is accepted.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "1. Easy to follow, mostly well explained and written paper.\n2. Introduce a novel method to reduce memory consumption and increase throughput of existing pre-trained LLMs with a simple step of fine-tuning\n3. Present empirical data to support the claims validity of the presented method.\n4. Show actual benefit with current hardware and not only theoretical benefits.", "reasons_to_reject": "1. Unclear baseline choice make it hard to compare this work to previous works, authors weren't clear why was those baselines chosen\n2. Limited experimentation, the method was tested on a single benchmark and another generation task which (to my knowledge) is not widely used. Furthermore, relatively small and unpopular models were used for the evaluation making it hard to assume if this method can scale beyond the narrow experimentation presented.\n3. No mention of open-sourcing implementation.", "questions_to_authors": "A. Why were these baseline chosen? Aren't there baselines to compare to from previous papers?\\\nB. Not clear how the context compression is done at inference time? Are random spans of tokens are chosen?\\\nC. Same question as B regarding training.\\\nD. Can this method play well with Flash Attention to improve inference even further? If no, how your method compare with flash attention?", "paper_topic_and_main_contributions": "The paper presents a novel method to compress the memory consumption of KV cache and reduce computation time of attention in decoder Transformer models through the introduction of 2 sentinel tokens and fine-tune the model to compress the information between the span of 2 tokens to a single token allowing us to save only 1 token instead of the entire span of tokens. The authors show that using this method we can reduce the KV cache significantly with moderate accuracy loss, use less memory and increase throughput of the model.", "excitement": 3, "reproducibility": 2, "ethical_concerns": "No", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"Typos_Grammar_Style_and_Presentation_Improvements": "1. Line 078: \"we present *the* necessary background\"\n2. Figure 2 could use bolder lines to make the figure easier to read"}, "extra_text": {"Missing_References": "1. [Fast Transformer Decoding: One Write-Head is All You Need](https://arxiv.org/abs/1911.02150) presented a different attention mechanism that can compress KV cache by 2 order of magnitude\n2. [FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness](https://arxiv.org/abs/2205.14135) presented a different computation algorithm to compute attention that saves memory and can be a drop-in replacement to the widely used attention algorithm."}}, {"review_id": "Z9G82SCOw0", "reviewer": "Reviewer_dqav", "summary": "", "questions": "A. Rationale for model selection?\n\nB. Why is the code not open source?", "limitations": "", "rating": 4, "confidence": 1, "soundness": 4, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We appreciate the valuable feedback provided by the reviewers. Our responses to Reviewer's questions are listed below:\n\n**Q**: \"Rationale for model selection?\"  \n**A**: Our primary deliberation revolves around encompassing various model scales, ranging from 1.3 billion to 2.7 billion parameters. Additionally, we take into account diverse position encoding schemes: absolute position encoding, as featured in OPT, and rotary position encoding, as employed in RedPajama, to demonstrate the generality of our approach. Extending to other LLMs is straightforward.\n\n**Q**: \"Why is code not open source?\"  \n**A**: We have uploaded the source code in the supplementary material, including the main Python files and Bash scripts. We will open-source it after the paper is accepted.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "The motivation is interesting.\n\nThe performance is significant.", "reasons_to_reject": "Insufficient model selection", "questions_to_authors": "A. Rationale for model selection?\n\nB. Why is the code not open source?", "paper_topic_and_main_contributions": "This paper introduces a method for reducing memory and computational costs by compressing intermediate activations in auto-regressive Transformers using sentinel tokens. The proposed approach is shown to outperform sparse attention baselines in terms of fluency, n-gram matching, and semantic similarity through experiments on various tasks. Additionally, the paper provides a comprehensive analysis of the positive impact of context compression on system throughput.", "excitement": 4, "reproducibility": 2, "ethical_concerns": "No", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "hjIfbHIa6c", "reviewer": "Reviewer_UZc8", "summary": "", "questions": "1 - Do you think it would be easy to write an efficient CUDA kernel for your specific method?\n\n2 - Did you run experiments to verify how the randomness of the selection process impacts the variability of the performance?\n\n3 - Is a given compression ratio consistently implying a given throughput across methods and implementations? In other words, do you think that for a given type of implementation (optimized CUDA kernels, fast attention, ...) a given compression ratio can yield a similar GPU throughput between Local Attention and KV compression? If there is such a relation, I think it would be interesting to mention it in your paper. If not, I think it would be interesting to compare *achievable throughput* and generation performance.", "limitations": "", "rating": 3, "confidence": 3, "soundness": 3, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We appreciate the valuable feedback provided by the reviewer. Our responses to Reviewer's questions are listed below:\n\n**Q**: \"relation between compression ratio and throughput\"  \n**A**: Among methods that reduce the length of key-value cache, the notable contribution of our proposed KV Compression is that it offers the flexibility to compress selected spans of tokens(in contrast to all past information beyond the window size in the Local Attention method), therefore resulting in better language modeling ability(reflected by PPL) and generation quality(reflected by PPL, ROUGE, and BERTScore). In Section 4.3, we profile the gain of throughput(the number of tokens generated per second) by compressing the key-value cache of ***prefix*** tokens, which essentially decreases the Memory Access Cost(MAC) of repeatedly reading from/writing to HBM. Specifically, we use different methods to encode the prefixing 800 tokens and only keep (1-compression_ratio) of the key-value cache. Therefore, the change of reported throughput only reflects the impact of reduced MAC and is the same of different methods under the same compression ratio.\nThe specific CUDA kernels for efficient attention implementation contribute to the inference efficiency ***from another aspect***. our proposed KV Compression can be readily combined with existing implementation, e.g., xformers' memory_efficient_attention, which we observe empirically offers more than 2x speed up than vanilla attention implementation.\n\n**Q**: \"Do you think it would be easy to write an efficient CUDA kernel for your specific method?\"  \n**A**:  One can implement a specific CUDA kernel for our proposed method based on the existing Flash Attention kernel. The modification involves the outer-loop logic, i.e., for token spans enclosed by sentinel tokens, subsequent tokens except the ending sentinel token can be skipped.\n\n**Q**: \"Did you run experiments to verify how the randomness of the selection process impacts the variability of the performance?\"  \n**A**:  Yes. To examine how the randomness in the selection process influences the generation quality, we ran our experiments on open-ended generation five times with different random seeds. The standard deviations for ppl, Rouge-L, and BERTScore across different compression ratios are 1.3, 0.6, and 0.5 respectively.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "This paper presents a novel method, and displays extensive experiments that convincingly demonstrate its relevance for open-ended text generation. The core idea is elegant and paves the way for future work. It could be impactful for the community as it allows more time and memory-efficient generation using LLMs at a low engineering and performance cost. The writing of the paper is compact and straightforward, which is satisfying.", "reasons_to_reject": "This paper is missing an important discussion in my opinion, which is the relation between compression ratio and throughput *across different methods and implementations of these methods*. I am not sure whether it would affect final conclusions positively or negatively, but I think it is a crucial point to demonstrate the actual usefulness of the approach. For instance, it is clear that local attention throughput can be increased by specific CUDA kernels, and thus I feel that it would be interesting to compare *Generation capabilities vs. Throughput* across methods in the same fashion as in Figure 2. After rebuttal, the authors addressed this concern for the inference part, but I believe the clarity of the paper could be improved on that matter.\n\nAnother potential drawback - that is addressed in the limitations section - is the randomness of the span selection process. It seems to me that it might affect the robustness of the generation performance. It would have been interesting to try to answer the following question: \"How often the randomness of span selection hurts generation capabilities?\" by studying worst-case scenarios, and variability of the performance across seeds/samples.\n\nOverall, while this paper **is relevant as is**, it would benefit from more thorough analysis and I think a long paper may be more suited for this work.", "questions_to_authors": "1 - Do you think it would be easy to write an efficient CUDA kernel for your specific method?\n\n2 - Did you run experiments to verify how the randomness of the selection process impacts the variability of the performance?\n\n3 - Is a given compression ratio consistently implying a given throughput across methods and implementations? In other words, do you think that for a given type of implementation (optimized CUDA kernels, fast attention, ...) a given compression ratio can yield a similar GPU throughput between Local Attention and KV compression? If there is such a relation, I think it would be interesting to mention it in your paper. If not, I think it would be interesting to compare *achievable throughput* and generation performance.", "paper_topic_and_main_contributions": "This paper introduces a novel method to compress self-attention maps during LLM-based text generation. The method consists in using two sentinel tokens that are randomly placed along the sequence around spans that will then be considered as single entities in the self-attention operations. Experiments show that after a short and efficient finetuning of the LLMs adapted with the sentinel mechanism, models are able to maintain satisfying generation capabilities across compression ratios, which is not the case for compared compression approaches. The authors finally expose the throughput increase allowed by their method, hence demonstrating the relevance of their method.", "excitement": 3, "reproducibility": 5, "ethical_concerns": "No", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"Typos_Grammar_Style_and_Presentation_Improvements": "L159 => span**s** of tokens"}, "extra_text": {}}, {"review_id": "ly5hDUmuxY", "reviewer": "Reviewer_BNcU", "summary": "", "questions": "- The caption in figure 1 does not clearly indicate the color encodings of the attention mask.\n- The proposed method is compared against two methods only, any reason for not considering more methods?\n- The proposed method is only experimented on one dataset, Wikitext-2. Do you expect a different result if you conduct your experiment on prompts dataset, like chatgpt-prompts? Do you think your result will be comparable with the work of “Learning to Compress Prompts with Gist tokens” by Jesse Mu, Xiang Lisa Li, Noah Goodman. ", "limitations": "", "rating": 3, "confidence": 4, "soundness": 3, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We appreciate the valuable feedback provided by the reviewer. Our responses to Reviewer's questions are listed below:\n\n**Q**: \"The caption in Figure 1 does not clearly indicate the color encodings of the attention mask.\"  \n**A**: We apologize for the confusion and will fix it in the revision. In Figure 1, the color in $i$-th row indicates whether one token can be attended by the $i$-th token or not. Green-colored tokens are allowed to be attended to while grey-colored tokens are disallowed. For example, for the token 'carpet', there are five tokens that it can attend to: 'A', 'cat', '<CR>', 'a', and 'carpet'.\n\n**Q**: \"The proposed method is compared against two methods only, any reason for not considering more methods?\"  \n**A**: The end goal of our proposed method is to curtail the length of the key-value cache so that both memory access cost and memory footprint are ameliorated. Therefore, our method is directly comparable to sparse attention-based methods that restrict the receptive field of each token to a smaller subset of tokens, thus reducing the size of the key-value cache required to be stored. We didn't compare against methods that reduce key-value cache from orthogonal perspectives, e.g., decrease the number of attention heads.\n\n\n**Q**: \"The proposed method is only experimented on one dataset, Wikitext-2. Do you expect a different result if you conduct your experiment on prompts dataset, like chatgpt-prompts? Do you think your result will be comparable with the work of “Learning to Compress Prompts with Gist tokens” by Jesse Mu, Xiang Lisa Li, Noah Goodman.\"  \n**A**: In addition to using WikiText-2 for language modeling experiments, we also used C4 to evaluate the effectiveness of different context compression methods on open-ended generation using C4. Language modeling dataset like WikiText contains data with presumably more inherent variations than prompts dataset. Thus, we expect our method to work comparably well on domains beyond the general language modeling corpus. To validate this point, we report here additional results(ppl) on WritingPrompts dataset using Redpajama-3B under compression ratios of 0.7, 0.8, and 0.9:\n\n| Method/Compression Ratio |  0.7 |  0.8 |  0.9 |\n|:------------------------:|:----:|:----:|:----:|\n|      Local Attention     | 23.1 | 25.2 | 29.5 |\n|      KV Compression      | 21.8 | 22.0 | 22.5 |\n\nThe gist token method proposed by Mu et at. is restricted to compressing a moderate-length prefix, e.g., prompt/instructions for supervised instruction tuning. When the target compression ratio is high, such prefix compression methods can hardly retain information needed for subsequent processing. In contrast, our method is more flexible and can retain much coherent information flow even under high compression ratio.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "The authors proposed a novel key-value context compression method and conducted comparisons against one dataset, two methods and three models. The proposed approach shows improvement compared with other sparse compression methods.\n", "reasons_to_reject": "Two major reasons to reject this paper.\n- The proposed approaches are combinations of well-known algorithms and the metrics are already popular in the community. The Author's contributions are marginal at best.\n- No theory behind. Formal definition and theoretical background of the proposed method are missing in the paper.\n", "questions_to_authors": "- The caption in figure 1 does not clearly indicate the color encodings of the attention mask.\n- The proposed method is compared against two methods only, any reason for not considering more methods?\n- The proposed method is only experimented on one dataset, Wikitext-2. Do you expect a different result if you conduct your experiment on prompts dataset, like chatgpt-prompts? Do you think your result will be comparable with the work of “Learning to Compress Prompts with Gist tokens” by Jesse Mu, Xiang Lisa Li, Noah Goodman. ", "paper_topic_and_main_contributions": "In this paper, the authors proposed a key-value context compression approach that places two sentinel tokens around a contiguous span of tokens and defines the position that are allowed and disallowed to be attended to using causal attention masking.", "excitement": 3, "reproducibility": 3, "ethical_concerns": "No", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"Typos_Grammar_Style_and_Presentation_Improvements": "- Typos, Abstract section line 013  “flucency\" --> \"fluency\"\n- Typos, Abstraction section line 016 “thoughout” → “throughput”\n", "Justification_for_Ethical_Concerns": "Justification for Ethical Concerns"}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper presents a novel method aimed at compressing the memory consumption of Key-Value (KV) cache and reducing the computation time of attention in decoder Transformer models. The approach leverages the use of two sentinel tokens, <CL> and <CR>, which encapsulate spans of tokens to compress. By subsequently fine-tuning the model, it compresses the information between these two tokens into one, enabling memory savings. 

In general, the reviewers were in consensus about the novelty of the and were generally positive on the empirical results compared to prior work.

There were some methodological concerns related to how the improvement of the approach compared to existing work in practical terms. For instance, while KV Compression was compared against Local Attention in terms of compression ratio and performance metrics (e.g., perplexity), they were not compared in terms of throughput. Instead, throughput is provided only for the proposed method.

Additionally, one reviewer felt that the datasets and models evaluated in the initial submission were limited. The authors provided additional results that satisfied the reviewer.
Decision: Accept-Main