BumbleBee: Dynamic KV-Cache Streaming Submodular Summarization for Infinite-Context Transformers

Authors: lilly  kumari
shengjie  wang
tianyi  zhou
nikhil  sarda
anthony  rowe
jeff  bilmes
Venue: colmweb_org_COLM_2024_Conference
Published: 2024-03-23 05:32:40

Abstract: Transformer-based Large Language Models (LLMs) have shown tremendous advancements across various domains. However, their need to maintain key-value representations (a KV cache) of previously seen tokens in the GPU memory leads to a significant memory overhead that scales linearly with the sequence length and batch size. With the advent of extremely long context LLMs, efficiently modeling long-range dependencies becomes challenging. In this work, we focus on the problem of long context summarization by formulating it as a subset selection problem. Specifically, we propose a novel submodular optimization framework called BumbleBee that uses a mixture of submodular functions to balance the diversity amongst the context tokens in the key embedding space and their importance computed using accumulated attention attributed to them across different input tokens. Our framework can work for both the LLM prefill and decoding phases, utilizing offline or online versions of our submodular algorithm respectively. While the context sizes grow to be as large only as the summary size, the temporal extent of the contexts may grow unboundedly, justifying the moniker ‘‘Infinite-Context Transformers.’’ Empirically, we validate the effectiveness of our framework across 13 different datasets using the LLaMA 7B and 13B models. Our results show that BumbleBee improves accuracy compared to state-of-the-art techniques at comparable context reduction ratios.

Main paper: 
 introduction: 
 The multi-headed self-attention (Vaswani et al., 2017) serves as the building block of several state-of-the-art transformer-based models for various tasks such as language understanding and generation (Raffel et al., 2020) , image recognition (Dosovitskiy et al., 2021) , and recommendations (Wu et al., 2023) using Large Language Models (LLMs). It enables models to learn long-range dependencies and complex patterns by focusing on multiple sequence sections simultaneously, regardless of their proximity. However, deploying LLMs is challenging for at least two main reasons: (1) quadratic scaling of attention: the attention mechanism scales quadratically with sequence length, leading to high computational costs and memory requirements for processing longer sequences; (2) autoregressive decoding wherein the LLM generates the output sequence token by token conditioned on the previously generated results and the input sequence, which requires accessing or recomputing the key and value representation of all previous tokens.Existing approaches to tackle the first problem include input sequence truncation (Devlin et al., 2018; Lewis et al., 2019) , sliding window-based sequence chunking (Li et al., 2020; Gong et al., 2020) , memory augmentation methods (Dai et al., 2019; Rae et al., 2019; Martins et al., 2021; Zemlyanskiy et al., 2021) , and retrieval-based methods (Khandelwal et al., 2019; Wu et al., 2022; Borgeaud et al., 2022; Zhong et al., 2022) . However, conventional techniques like chunking suffer from well-known problems such as context fragmentation (Dai et al., 2019) . Similarly, retrieval-based non-parametric methods to augment an input query with relevant context require constructing an external memory bank that is much larger than the dataset size. In addition, a retriever index must be constructed on top of the memory bank.A common approach to address the second problem is to store the previously generated key and value vectors in a KV cache (Pope et al., 2023) which helps avoid redundant computation at each generation step. This presents a memory bottleneck, as the KV cache size grows linearly with the generated sequence length and the batch size for a specific model. Additionally, as model size increases, the memory footprint of a single token increases as well. While it is possible to offload the KV cache to the host DRAM, this will incur a host-to-device latency for each inference call. Given the growth in model size and context lengths of publicly available models (Reid et al., 2024; Anthropic, a) , reducing the memory overhead of the KV cache while retaining model accuracy becomes even more important.Existing system-level techniques such as FlexGen (Sheng et al., 2023) , PagedAttention (Kwon et al., 2023) , and FlashAttention (Dao et al., 2022) improve the utilization of GPU resources and throughput when dealing with the attention mechanism and KV caches. However, they do not consider the impact of ever-increasing KV cache sizes. Modeling techniques such as multi-query (Shazeer, 2019) and group-query attention (Ainslie et al., 2023) help mitigate the size of the KV cache by removing unnecessary heads but they require expensive retraining/fine-tuning.Given the wide deployment of accelerators and models already in the field, we need to leverage inference time techniques that reduce the size of the available context at all layers (KV cache) without impacting accuracy. Existing related works such as Heavy Hitters (Zhang et al., 2023) , Scissorhands (Liu et al., 2023) , KeyFormer (Adnan et al., 2024) and FastGen (Ge et al., 2023) rely on heuristic-based techniques to preserve only the important key-value attention states while evicting the non-important attention states. They employ modular scores such as attention scores attributed to different tokens accumulated over different time steps. This does not take into account the importance of keeping a particular KV attention state in the context of other states already present in the KV cache and thus does not efficiently capture long-term dependencies between tokens. We hypothesize that KV cache selection should be framed as a subset selection problem where we evaluate the utility of different key-value pairs as a set instead of independently evaluating each key-value attention state.In this work, we present a novel yet simple approach called BumbleBee 1 for KV cache summarization, thus enabling existing LLMs to be used for longer contexts without any additional fine-tuning, CPU offloading (Sheng et al., 2023) , or context truncation (Beltagy et al., 2020) in case of longer sequences (Fig. 1 ). The temporal span of the context (which is defined as the difference between the time of the latest and the time of the earliest context token) may be unboundedly long in a BumbleBee model. This is true even though the total number of tokens in a BumbleBee context does not grow unboundedly. Indeed, BumbleBee draws inspiration from the following aspects of human psychology:1. Selective attention (Chun & Turk-Browne, 2007; Uncapher & Rugg, 2009) allows us to focus on relevant information and filter out distractions or irrelevant details. This plays a critical part in determining what subset of information should get successfully encoded into our memory. Similarly, in BumbleBee, we maintain a memory of diverse, representative, and important time points of the sequence observed so far by utilizing online submodular summarization. 2. Humans process information in an online and dynamic fashion over an extended period, possibly all the way to very early in life. This involves relying on our (possibly extremely long) past memory to make sense of incoming information through associations, selectively attending to relevant information, and organizing our existing memory in light of new information. This top-down processing style (Gazzaley & Nobre, 2012) enables us to prioritize and encode pertinent information for future recall.Similarly, in the case of BumbleBee, we process an incoming segment or chunk by contextualizing it with the online summary (or memory) obtained thus far and update the summary/memory based on the segment's informativeness and its association with the current memory. Like human memory, BumbleBee's online summary may include tokens anywhere from the earliest point in the stream all the way to the present.To realize the above inspiration, BumbleBee utilizes online submodular summarization to maintain a diverse subset of important key-value attention states on the fly that are representative of long-term global history. Incorporating a fixed-size summary of the entire historical sequence allows the BumbleBee to capture dependencies having arbitrarily long range without also unboundedly increasing the memory consumption or the length of input fed to the transformer. That is, the memory grows no larger than the size of the summaries. However, since the temporal extent itself between the oldest and newest token present in the summary-sized context windows can grow unboundedly, BumbleBee deservedly may be seen as an instance of an 'Infinite-Context" transformer. We further combine the global submodular summary with the most recent local context as shown in Fig. 1 , so BumbleBee's predictions are dependent on both the latest context as well as the global wide-ranging patterns observed so far in the temporal history. Thus, our approach boosts the transformer's capabilities to capture both the arbitrary-long coarse-grained long-term and the detailed short-term contextual dependencies while maintaining similar (and thus feasible) memory and computational costs for a fixed-size input sequence length. 
motivation: 
 LLMs are pre-trained using a fixed context window and prior work has shown their limited generalization on sequences significantly longer than the pre-trained context window (Kazemnejad et al., 2024) . Certain recent closed sourced models have a context window longer than 128k tokens, for example, GPT-4 (OpenAI, 2023) has a context window of 128k, Claude 2.1 (Anthropic, b) can process 200k tokens while Gemini 1.5 pro (Reid et al., 2024) can handle a 1M input context. To utilize such a long context, caching of key-value states is utilized to keep the latency associated with the inference phase low. However, the KV cache grows linearly with the sequence length. For example, for a LLaMA-13B model (Touvron et al., 2023b) , the KV cache for 128k tokens would roughly require 40 (num. of decoder layers)×5120 (hidden size)×2 (kv pair)×2 (fp16)×128000 = 105 GB. To address the memory costs associated with the KV cache, we focus on the following:Is it possible to maintain the LLM's performance on downstream tasks without storing every observed token representation in the KV cache?We conduct experiments on the WikiText-103 (Merity et al., 2016) dataset to support our hypothesis that self-attention is selective and keeping only a set of both diverse and important tokens is sufficient to maintain performance. Specifically, we use a LLaMA-7B model (with a pre-trained context window size of 2048) for a next-token prediction task on randomly sampled articles from wikitext-103 and visualize the attention scores attributed to different tokens present in the context window of size 2048 in Figure 2 . On the x-axis, each token represents the query used for next-token prediction, while each vertical slice (column) illustrates the logarithm of normalized attention scores for the 2048 in-context tokens (or keys). As we move from the bottom row to the topmost row along the y-axis, the relative distance of the in-context keys from the query token increases from 0 to 2047. The anti-diagonal pattern shown in Figure 2 shows that there is a small subset of tokens that are strongly attended to while they are present in the context window.However, the constraint of the context window restricts tokens in longer sequences to only attend to nearby local tokens, limiting their ability to capture broader contextual information. In this work, we show that by using a submodular function to summarize the entire historical contextual information in the KV cache across different attention heads, one can reduce the memory footprint associated with the KV cache while maintaining acceptable performance compared to using the entire cache. 
related work: 
 Attention speedup: Self-attention is a critical component of the transformer (Vaswani et al., 2017) mechanism powering modern language models but suffers from quadratic complexity. Techniques such as Linformer (Wang et al., 2020) , Performer (Choromanski et al., 2020) , Linear Transformers (Katharopoulos et al., 2020) and Reformer (Kitaev et al., 2020) aim to reduce the time complexity of self-attention via low-rank approximations or hashing techniques. Longformer (Beltagy et al., 2020) introduces sliding-window attention to reduce the computational overhead of self-attention. FlashAttention (Dao et al., 2022) implements an IO-aware kernel to compute self-attention and uses tiling to reduce memory overhead. Keyformer (Adnan et al., 2024) speeds up attention computation by exploiting the empirical observation that 90% of attention weights focus on a small subset of tokens. For workloads that share similar inputs, prompt caching (Gim et al., 2023) and prefix sharing (Ye et al., 2024) have emerged to reduce the computation and memory overhead of self-attention.KV cache compression and management: H2O (Zhang et al., 2023) identifies "heavy-hitter tokens" via a modular score function (i.e., simply the sum of individual token scores) and keeps those while evicting the rest, reducing the size of the KV cache. In Mu et al. (2024) , prompts are compressed to gist tokens to reduce the size of the KV cache. FastGen (Ge et al., 2023) uses profiling information from the attention computation to determine which tokens to evict from the KV cache.Additionally, several post-training quantization techniques have been proposed to reduce the memory overhead of the KV cache. WKVQuant (Yue et al., 2024) proposes a pastonly quantization mechanism to enable higher precision for self-attention computation.In Yang et al. (2024) , important tokens are retained at full precision while the rest are quantized. Quality-aware Quantization (Dong et al., 2024) uses an attention-aware approach to selectively quantize elements in the KV cache. GEAR (Kang et al., 2024) approximates quantization error via a low-rank matrix and uses a sparse matrix to correct them at inference. Sparq (Ribar et al., 2024) drops tokens according to attention sparsity scores and incorporates the error of the pruned value cache. Unlike the lossy compression techniques above, Loma (Wang & Xiao, 2024) proposes a lossless method for compressing the KV cache. However, it requires further fine-tuning of the underlying model.Finally, efficient management of the KV cache, such as PagedAttention (Kwon et al., 2023) improves the throughput of LLM serving by reducing fragmentation and redundant dupli-cation.Applications of submodularity to LLMs: INGENIOUS (Renduchintala et al., 2023 ) is a technique that uses submodular optimization for selecting representative subsets of the training data such that the language models trained thereof achieve comparable performance to models trained on the full dataset. SMART (Renduchintala et al., 2024) proposes a data mixture strategy for instruction tuning, leveraging a submodular function for importance score assignment to tasks that are used to determine the mixture weights. Submodular functions have also been used to augment LLMs for multi-document summarization (Kurisinkel & Chen, 2023). Div-S3 (Kumari et al., 2024) summarizes via the submodular span (Kumari & Bilmes, 2021) based on the conditional submodular gain function. This approach facilitates the selection of diverse and relevant exemplars for in-context learning with LLMs in a data-efficient manner. To the best of our knowledge, ours is the first work to apply submodularity for KV cache summarization. 
background: 
 In this section, we provide a concise background about self-attention used in transformerbased LLMs. We also briefly discuss submodular function optimization as well as the KV cache mechanism used to avoid re-computations during the decoding stage.We denote a sequence of tokens as x = {x 1 , x 2 , . . . x n } where x t ∈ R d is the t-th token embedding. In a decoder-only transformer model, we denote the number of decoder layers present as n l and the number of attention heads present in each layer as n h . The self-attention mechanism of head h in layer (l + 1), utilizing distinct weight matrices for query (W q ), key (W k ), and value (W v ), operates on the hidden states from the preceding layer, represented as H l . Specifically, we get the following query, key, and value embeddings after the linear projections:Q l+1 = H l W q , K l+1 = H l W k , and V l+1 = H l W v .Row t in the key embedding matrix denotes the key vector corresponding to the t-th token, and this holds for both query and value matrices as well. Given query q t , its attention output o t using scaled dot product attention is computed as shown in Eq. 1. Here, we use S t to denote the KV cache accumulated at step t. From here onward, we omit notation pertaining to the specific layer and head for the sake of simplicity.EQUATIONKV Cache: Decoder-only transformers operate in an autoregressive manner, predicting tokens sequentially based on previously generated (and observed) ones. By caching previously computed representations, specifically the embeddings for keys and values corresponding to observed tokens, the model can avoid redundant computations in each decoding step. Specifically, the KV cache initially computes the attention states for an input prompt, represented by S 0 = {(k i , v i )|i ≤ n}, and caches them in memory. For every subsequent step j ≤ k, the model reuses the cached states S j = {(k i , v i )|i < n + j} to compute the attention state (k n+j , v n+j ) of the new token s n+j . This significantly reduces the floating-point computations used for matrix operations to compute the new attention states. After each step, the newly computed attention states are appended to the cache for subsequent use, such thatS j = S j-1 ∪ {(k n+j , v n+j )}.Note that the KV cache is an approximation and trades off precision for speed. The attention state computation for token s j is limited to the sequence available at step j, namely {s i |i < n + j} as opposed to over the entire sequence {s i |i < n + k}. KV caches significantly reduce latency and result in a minimal loss in accuracy making it an essential ingredient in practical LLM deployment.Since the KV cache is updated with a new token at each decoding step, its size grows linearly with the overall sequence length (including both input prompt and generated tokens) and the batch size. This becomes a major bottleneck when dealing with longer sequences whose KV cache cannot fit in the GPU's high bandwidth memory (HBM). In this work, therefore, we study how to summarize the KV cache by keeping only a small subset of important and diverse and thus representative key-value embeddings while discarding the rest.Submodularity: A submodular (Bilmes et al., 2022) function f : 2 V → R defined on the ground set V always has a diminishing return property: f (v|A) ≥ f (v|B) for any v / ∈ B and A ⊆ B ⊆ V, where f (v|A) := f (A ∪ {v})f (A). Intuitively, the gain of an item v, i.e., f (v|A), diminishes as the conditioning set grows from A to B. A submodular valuation f (A) expresses the diversity and representativeness of the input set A ⊆ V of items. There are many useful such functions, one of them being the facility location (FL) function which is similar to a k-medoids objective (Kaufman & Rousseeuw, 1987) . The FL function utilizes similarity scores sim(v, v ′ ) computed over every pair v, v ′ ∈ V of items. A valuation is then the sum of similarities from any item in the ground set V to its closest representative in the given set A as shown in Eq. ( 2).f FL (A) = ∑ v∈V max v ′ ∈A sim(v, v ′ ). (2) c(A) = ∑ u∈U ϕ u ( ∑ v∈A m u (v)).(3)A subset A with a high function value indicates that for every item in the ground set, there exists an item in A that is very similar, or in other words, A is representative of the ground set V. Another widely used submodular function is the feature-based function, and it has the form shown in Eq. ( 3). Here, we have ϕ u (•) as a monotone non-decreasing non-negative concave function, and m u (•) is a non-negative weight associated with the u-th feature of every item v ∈ A. Due to the diminishing property of ϕ u (•), to have a large function valuation for a set A, we would require the sum of every feature across items to be uniformly large, thus inducing diversity and fairness over a feature representation of the selected subset. 
bumblebee: 
 In Section 2, we saw that some keys despite their distance from the query are heavily attended to, showing that it is important to preserve these keys in the overall context even if they are quite distant. While the importance of keys captured by the attention scores is one aspect that should be considered for context (or KV cache) selection, the diversity and representativeness of the selected keys are equally critical. That is, we require the context summary (which has a fixed size) to be both relevant to and representative of the entire context. This is how efficiency is achieved -amongst sets of a given fixed size, a diverse set means a non-redundant set, while a redundant set means that certain concepts are inefficiently over-represented while some concepts are poorly represented. To capture these properties (diversity and importance), our final scoring function g λ is a convex combination Algorithm 1 Offline Submodular KV cache Summarization during Prefill/Encoding Phase 1: Input: Submodular functions capturing diversity f FL in the key embeddings space and importance c via attention frequency for layer l and attention head h; mixture functiong λ (•) = λ f FL (•) + (1 -λ)c(•); a set of n KV attention states K n = {(k i )} n i=1 , V n = {(v i )} ni=1 corresponding to the n prompt tokens; budget τ s . 2: Output: A final summary S n such that S n ⊆ {(k i , v i )} n i=1 and |S n | ≤ τ s . 3: Initialize: S n = ∅; compute accumulated attention score vectors a n for each key k ∈ {k i } n i=1 . a i n denotes accumulated attention scores attributed to key k i across all n query tokens. 4: for j = 1 to τ s do 5: k imp ← argmax e∈K n \S n g λ (S n ∪ e) -g(S n ) 6: S n ← S n ∪ {(k imp , v imp )}S t ← S t-1 ∪ {(k t , v t )} 8: Append a(q t , k t , S t ) to a t s.t. |a t | = |S t | 9: else 10: Let S ′ t = S t-1 ∪ {(k t , v t )}; k discard ← argmin k i ∈S ′ t g λ (k i |S ′ t \ k i ) 11: S t ← S ′ t \ {(k discard , v discard )} 12: if k discard ̸ = k t then 13:Evict a j t (the accumulated attention score for the discarded key k discard ) from a t .14:Append a(q t , k t , S t ) to a t 15:end if 16:end if 17: end for of the facility location (Eq. 2) and the feature-based (Eq. 3) functions.g λ (A) = λ f FL (A) + (1 -λ)c(A) (4)Here, λ ≥ 0 is a hyperparameter that controls the trade-off between representativeness and relevance. Both component functions are monotone, non-negative, submodular, and assumed to be normalized i.e., f FL (∅) = 0 and f FL (V) = 1 (and the same for c(•)). This normalization ensures the compatibility among the mixture components and the resulting mixture function g λ inherits these properties.In Alg. 1, we present the offline KV cache summarization pseudo-code to compute a summary of the KV cache under a cardinality constraint τ s . This algorithm is used for the KV cache summarization in each self-attention head present in different decoder layers. In Line 3, we first compute the attention scores a n for all keys accumulated over different time steps and then use that to instantiate our feature-based function c(•). In our current setting, we only experiment with one feature function, meaning |U| = 1 in Eq. 3. The non-negative weight m u (k i ) in Eq. 3 for key k i is a i n , which represents the accumulated attention k i receives from the observed input queries. Using the key embeddings, we instantiate the facility location function f FL (•) on a similarity matrix computed using pairwise cosine similarities followed by the ReLU transformation.In Lines 4-7, we use the greedy algorithm (Nemhauser et al., 1978; Minoux, 1978; Mirzasoleiman et al., 2015) to perform a cardinality-constrained submodular maximization. Thanks to submodularity, the resultant set is within a factor of (1 -1/e) from the optimal summary (Nemhauser et al., 1978) . This offline routine is suitable for KV cache summarization after the prefill stage of LLMs, particularly in serving systems where prompts are shared across user requests, and their KV embeddings are pre-computed and cached (Kwon et al., 2023; Gim et al., 2023) .In Alg. 2, we outline the KV cache summarization algorithm suitable for a streaming setting where we do not have prior knowledge of the complete sequence, restricting us from using existing offline summarization algorithms to obtain a global summary of the sequence. In light of this limitation, we propose to summarize the sequential data observed so far and produce an online summary of a fixed size that serves as its representative set. We remark that the computation at each step is at most O(τ 2 s ) and which is fixed regardless of the length n of the sequence. Further discussion is given in Appx. E. We use the same mixture function as Alg. 1 as our final scoring function.In Lines 6-8, we keep caching the (key, value) pairs for different attention heads unconditionally until the summarization budget is exhausted. Once the KV cache is full, we utilize the key embeddings in S t and the accumulated attention score vector a t to update the submodular function components f FL (•) and c(•) of the final mixture function g λ (•). In Line 10, we create set S ′ t that includes the newest incoming (k t , v t ) pair and evaluate the conditional gain of keeping an item around in the context of the remaining summary set as shown in Line 11. The element with the least conditional gain is removed from S ′ t and the accumulated attention vector a t is modified in Lines 13-15 to include the attention score associated with the newest key when the discarded key is one of the keys in the previous summary set S t-1 .The streaming summarization algorithm (Algo. 2) is suitable for memory-constrained settings where the KV cache for the entire context cannot be maintained in the GPU memory. Also, in multi-turn dialogue systems (Duan et al., 2023; Maharana et al., 2024) , after a certain number of interactions, it can be challenging to keep track of the entire conversation context. In such scenarios, BumbleBee can enable modern LLM-based serving systems to maintain a representative yet important summary of past interactions while minimizing the KV cache-based memory utilization. 
datasets & tasks: 
 The datasets studied in this paper are derived from three benchmarks: lm-eval-harness (Gao et al., 2023) , HELM (Liang et al., 2022) , and LongBench (Bai et al., 2023) . Following Heavy-Hitters (Zhang et al., 2023) , we select the following six few-shot datasets from lm-eval-harness: OpenbookQA (Mihaylov et al., 2018) , COPA (Roemmele et al., 2011) , RTE (Wang et al., 2018) , MathQA (Amini et al., 2019) , PiQA (Bisk et al., 2020) , and Winogrande (Sakaguchi et al., 2021) . From the HELM benchmark, we choose the single document summarization dataset XSUM (Narayan et al., 2018) . From the LongBench benchmark (Bai et al., 2023) , which is meant for evaluating the long-context understanding of LLMs, we select four tasks and their associated datasets: (1) Single document question answering: Qasper (Dasigi et al., 2021) and MultiFieldQA, (2) Multi-document question answering: HotpotQA (Yang et al., 2018) and 2WikiMultihopQA (Ho et al., 2020) , (3) Summarization: QMSum (Zhong et al., 2021) , and (4) Few-shot learning: TREC (Li & Roth, 2002) .All submodular functions and their optimizations are implemented with an upcoming to-be-open-sourced optimized C++-based software system called Submarine (Bilmes, 2024) . All submodular computation, which is non-SIMD-style mixed integer-floating point, is done on the multi-threaded CPU side since that is a natural, cost-effective, efficient and thus still performant platform for such computation than vector-capable GPUs.We use LLaMA 7B and 13B models (Touvron et al., 2023a) for tasks belonging to the lm-evalharness benchmark. For the XSUM dataset, we use LLaMA2 7B and 13B models (Touvron et al., 2023b) . On the LongBench selected datasets, we use the Llama-2-Chat 7B fine-tuned model and the LongChat-32k 7B model (Li et al., 2023) .To show the effectiveness of BumbleBee for the context (KV cache) summarization task, we compare it to the following baselines: (1) All: we use the entire KV cache and do not perform any cache reduction (2) Local: only the most recent context x% tokens are maintained in the KV cache and the remaining old KV states are evicted (3) Random + Local: randomly selected tokens along with the most recent tokens are retained in the KV cache (4) Attention sinks + Local: the first four tokens known as attention sinks (Xiao et al., 2023) along with the most recent tokens are kept in the KV cache (5) H2 + Local: only the tokens that are most frequently attended to, referred to as Heavy Hitters (Zhang et al., 2023) are maintained in the KV cache along with the most recent tokens. Table 1 : Results on the few-shot tasks from the lm-eval-harness benchmark using LLaMA 7B and 13B (Touvron et al., 2023a ). In the above methods (except All), we do a 10× context reduction, so our KV cache summarization budget is 0.1× the input sequence length. Table 2 : Results on six datasets from the LongBench benchmark using Llama2-7B-chat-4k (Touvron et al., 2023b) and LongChat-v1.5-7B-32k (Li et al., 2023) . * indicates that the reported numbers are sourced directly from Bai et al. (2023) .In the case of BumbleBee, we test two concave function choices for Eq. ( 3): (♥) log-based: ϕ(x) = log(1 + x) and (♦) power-based: ϕ(x) = g -1 (x) where g(y) = αy 1/α + βy. Herein, we set β and α as 1 and 0.04 respectively to ensure that the function saturation curve is compatible with the facility location function, the first submodular function component of g λ (•) in Eq. 4, thus ensuring that neither component from Eqn (3) or Eqn (2) ever always dominates the other in the optimization. Selecting the form and curvature of the concave function in the submodular mixture is a hyperparameter that is just as important for good BumbleBee performance as is the mixture parameter λ in Eqn. (4). Hyperparameters were tuned on a development and tested on a test set (see Appx. D). The extreme cases (λ ∈ {0, 1}), corresponding to using only one component from Eqns (3) and (2), yield suboptimal performance.In Table 1 , we compare BumbleBee to existing KV cache reduction methods on five-shot learning tasks from the lm-eval-harness benchmark (Gao et al., 2023) . Across both LLaMA (Touvron et al., 2023a) variants we used for inference, BumbleBee consistently outperforms other baselines in terms of accuracy, showing a performance comparable to the best-case no-compute-constraint setting when All of the KV cache is used without any summarization/reduction.In Table 2 , we report the results on six different datasets from Long-Bench (Bai et al., 2023) . For all datasets excluding QMSum and TREC, we use F1 score for evaluation. For QMSum and TREC, we use Rouge-L and accuracy respectively as the evaluation metric. When using the LLaMA-7B-chat-4k model pre-trained with a 4k context window, we truncate the middle part of the input contexts if their length exceeds 4k as suggested in Bai et al. (2023) . As can be seen in Table 2 , the offline version of BumbleBee outperforms other context reduction methods such as H2 on four datasets even when the summarization budget is 20% of the entire context size.However, when using the LongChat-7B-32k which was trained further to generalize to longer sequences, we adopt a Sliding Window (SW) strategy to process smaller chunks/segments of longer contexts, aggregate their KV embeddings, and the attention scores accumulated for keys in different chunks. This is done to process the entire context/sequence in a memorymanageable way while avoiding context truncation. This strategy however is sub-optimal as tokens in faraway segments cannot attend to tokens at the beginning of the original sequence. Despite this, we observe that BumbleBee outperforms the state-of-the-art KV cache reduction method H2 by strong margins (1.6%-8.5% in absolute terms). In Appx. C, we present qualitative results that illustrate how BumbleBee effectively retains a representative subset of keys within the reduced KV cache.(a) LLaMA-13B (b) LLaMA-7BFigure 3 : ROUGE-based results on XSUM dataset, a few-shot summarization task from the HELM benchmark (Liang et al., 2022) using two different LLaMA models (Touvron et al., 2023b) . To reduce the pressure on the context window across all decoder layers, we perform a 5x KV cache reduction for each of the above methods except All. XSUM Summarization Task: In Fig. 3 , we compare different context reduction methods to the full cache setting. We use LLaMA2 models (Touvron et al., 2023b) to assess the downstream summarization performance in a 3-shot learning setting. BumbleBee outperforms other SOTA cache reduction techniques and even performs better than the All cache setting when using LLaMA-7B.We use the XSUM dataset to show the overall performance sensitivity of BumbleBee to the convex mixture weight hyperparameter λ in Eq. 4. In the case of LLaMA-13B model, λ ∈ [0.4, 0.8] performs comparably showing that a representative subset of the KV cache is preferred for the downstream task. In LLaMA-7B, we see that for λ > 0.2, the performance starts to drop as we increase the mixture weight λ corresponding to the Facility location function f (•). However, the first submodular component with its relative weight of 0.2 still outperforms the H2 + Local method showing that both representativeness and relevance of the selected cache subset are desirable to maintain a performance comparable to the entire cache setting. 
conclusions & future work: 
 For future work, we plan to evaluate BumbleBee on reasoning (Sawada et al., 2023) and problem-solving datasets (Cobbe et al., 2021) and investigate the impact of hyperparameter tuning and alternative function mixtures in our formulation. Additionally, we plan on modifying sliding-window attention with a submodular formulation to better capture long-term dependencies. More generally, regarding human episodic memory and selective attention (Chun & Turk-Browne, 2007; Uncapher & Rugg, 2009) , it will be interesting to consider and study the dynamic "bounds", "limits", and "finite capacity" of episodic human memory such as (Brehmer et al., 2004; Dings & McCarroll, 2022) . With BumbleBee, this corresponds to the size of the historical summary that is retained. It may be that in human memory, this bound grows slowly over time. A corresponding BumbleBee variant could be produced where the submodular summary size grows very slowly over time as well (e.g., log(1 + log(1 + log(1 + . . . n)))). The attention layers should be able to support this. 
 
 Appendix: f ground-truth (summ) 8.2689 × 10 -4 f ground-truth (local) 8.3301 × 10 -4Table 3 : Performance comparison of the generated ground-truth model on the entire dataset when only the summaries are used and when only the local context is used a. f summ,local learnt -trained using both summaries and local context tokens.b. f local learnt -trained using only the local context tokens.c. f summ learnt -trained using only the summary tokens. For the entire training setting, unless mentioned otherwise, we perform an 80:10:10 train, validation, and test split respectively (random seed initialization same across all three settings), and train for 50 epochs. The results are evaluated on the held-out test set. From Table 4 and Fig. 4 , it is clear that f summ,local learnt performs significantly better than f local learnt and f summ learnt , exhibiting the importance of both summary and local context input tokens for the final predictions. In Fig. 5 , we provide qualitative results showing the predicted values using different learned models. Note that the plot in (c) looks granular in the vertical dimension since the summaries are not changing at every time point, i.e., time is segmented into regions where the summary is fixed, meaning S t = S t-1 for certain regions.We observe the stripe pattern in Fig. 4c because the input summary tokens remain constant across multiple episodic periods due to the design of the data sampling process. We further demonstrate that using a large local context does not necessarily yield comparable results using the same training setting. To achieve this, we simply increase the size of the local context window by 10 times the original to obtain, f local×10 learnt . We compare the performance of these models on metrics like MAC (Multiply and Accumulate), runtime memory, and MSE. The MAC and runtime memory metrics are computed for a single training example in a mini-batch. For one of the lm-eval-harness (Gao et al., 2023) tasks namely COPA, we analyze how the KV cache across different self-attention heads gets updated as new queries are processed in a streaming setting. In Figures 6 & 7 , we show such visualization for two samples selected from the test set of COPA (Roemmele et al., 2011) . Since each incoming query attends to its key vector, the offset diagonal line simply indicates that property. Overall, we see that H2 is heavily biased towards selecting the initial set of keys (close to position 0) and maintaining them in the KV cache. This pattern holds across most of the attention heads. However, BumbleBee summaries appear more time-diverse across the majority of heads showing that our framework is capable of balancing both representativeness and relevance in the final summary of the KV cache.Next, we visualize how good the subsets selected by Heavy Hitters (H2) and BumbleBee are on one of the samples chosen from the 2WikiMultihopQA (Ho et al., 2020) dataset from the LongBench task. We visualize the t-SNE embeddings of the keys across randomly selected self-attention heads in certain layers along with the selected subset of the KV cache. Fig. 8 shows that BumbleBee can select a representative subset of keys from the entire KV cache when using the offline greedy algorithm, thus explaining its strong performance on various LongBench datasets.Table 6 compares the decoding speeds (in ms/token) for two different context reduction ratios, demonstrating faster decoding speed as the context reduction ratio becomes higher.We implement BumbleBee in PyTorch. To update the KV cache in the streaming setting, we modify the LlamaAttention class from the huggingface library. For the submodular optimization, we use an internal highly optimized submodular toolkit, and we plan to Table 6 : Decoding speed (in ms/token) for two KV cache reduction ratios (5:1 and 10:1) and the baseline KV cache method using the entire context (1:1) across all heads. All experiments are performed on an A100 80GB GPU using the LongChat-7B-32k with a batch size of 1.open-source the integrated codebase shortly.On the LongBench dataset, we first compute the key embeddings for all the context tokens. This is performed for each self-attention head in each decoder layer. We then use stochastic greedy algorithm (Mirzasoleiman et al., 2015) with ϵ = 1e-5 to compute the offline KV cache summaries for each of h × l heads in a parallel fashion. Here, h denotes the number of attention heads present in one decoder layer, and l denotes the number of decoder layers.For similarity matrix computation, we use ReLU truncated cosine similarity to ensure pairwise similarities sim(i, j) ≥ 0. We have also experimented with other cosine similaritybased metrics such as 1 + cos(i, j) and |cos(i, j)| but find that ReLU(cos(i, j)) works the best.Mixture weight λ: on the tasks from the lm-eval-harness benchmark (Gao et al., 2023) , λ ∈ {0.2, 0.3} perform the best across both evaluated models, i.e., LLaMA 7B and 13B. However, we did not perform a more fine-grained search/tuning for λ in this range. On LongBench (Bai et al., 2023) , we report the results for λ = 0.3 across all six datasets. For the XSUM (Narayan et al., 2018) few-shot summarization task, we use a held-out validation set to tune the mixture weight and set λ = 0.2 when using LLaMA2-7B and λ = 0.4 for LLaMA-13B.Compute: we use an NVIDIA-A100 GPU to perform our inference-based experiments and a multi-threaded CPU for all submodular computation using Submarine (Bilmes, 2024) . For the experiments involving training on synthetic data in Sec. B, we use an NVIDIA RTX 2080.Compute costs: In BumbleBee (Alg. 2), we maintain a global summary of size τ s . So, constructing the pairwise similarity matrix for f FL has a time complexity of O(τ 2 s × d). Identifying the item with the least conditional gain (Line 11 of Alg. 2) requires O(τ 2 s ). However, if we cache the similarity matrix, we only need to compute the similarity of the incoming item to the others in the summary, resulting in an overall complexity of O(τ s × d + τ 2 s ) for each new token. Memory Costs: Caching the similarity matrix incurs O(τ 2 s ) memory costs.  


Decision: This paper presents a method to compress the KV cache in LLMs using a mixture of submodular functions to select a diverse and representative set of keys from the context. The authors have done a good job during rebuttal. After rebuttal, the paper received scores of 6777. All the reviewers are happy about the paper, commenting that (1) the proposed method is novel and interesting, (2) the paper is technically sound and well written, and (3) experimental results are comprehensive and sufficient. Therefore, the AC would like to recommend acceptance of the paper.

Reviews: review: paper_summary: paper_summary: This paper introduces a novel approach to manage the memory overhead of Transformer-based LLMs by summarizing the key-value (KV) cache dynamically. The proposed BumbleBee framework uses a mixture of submodular functions to optimize the selection of key-value pairs, balancing diversity and importance based on accumulated attention scores. This method enables efficient memory utilization without sacrificing accuracy, even with long context sequences. BumbleBee operates both offline and online, accommodating various operational phases of LLMs. The effectiveness of this approach is validated through experiments on multiple datasets, showing improved performance over existing techniques with similar context reduction ratios.
strength_weakness: strength_weakness: 1. **Lack of Theoretical Analysis**: While the empirical results are promising, the paper lacks a theoretical analysis of the convergence properties or bounds of the submodular optimization framework, which would strengthen the validation of the method.

2. **Impact of Hyperparameters**: The paper provides a brief discussion on the impact of the mixture weight hyperparameter $\lambda$, but it lacks a comprehensive analysis of other critical hyperparameters that may influence the framework's performance. Additionally, the study does not include ablation experiments, such as evaluations using only $f(A)$ or only $c(A)$, which could offer insights into the individual contributions of different components of the submodular function to the overall effectiveness of the BumbleBee framework.
questions: questions: 1. Could the authors elaborate on the theoretical underpinnings of the optimization problem utilized in the BumbleBee framework? Specifically, are there any provable guarantees regarding the convergence rate or the quality of the summarization achieved through your submodular optimization approach?

2. The experimental section lists two different concave functions—log-based and power-based—for the submodular formulation in Tables 1 and 2. Could the authors clarify the rationale behind using these two distinct function choices? Additionally, how do these choices influence the model's performance and why were they not discussed earlier in the methodology section?
paper_summary: paper_summary: This paper presents a technique to compress the KV cache in LLMs using a mixture of submodular functions to select a diverse and representative set of keys from the context. Experiments on a series of well-studied LLM benchmark datasets indicate that this technique typically outperforms prior work and is occasionally competitive with full-context models.
strength_weakness: strength_weakness: * The computation of $g_\lambda(\cdot)$ appears fairly expensive, especially the FL function (eq. 2) in the streaming setting (algorithm 2), which seems to require $O(n^2 \tau_s d)$ time when computed naively. This may be a reasonable tradeoff for the memory savings, but the paper does not include an evaluation on the impact on latency. A highly-optimized toolkit for computing such functions is briefly mentioned in Appendix C but more information on runtime costs would be helpful.

* The relevance of the particular submodular functions chosen here to the task at hand is not conveyed clearly. It appears the feature-based function (eq. 3) may be computing the sum of attention scores for each key—similar to H2—but this is not explicitly described.

* The notation is overloaded and confusing. $V$ is used to represent both a grounding set and the set of value embeddings, even though the actual grounding set is based on keys rather than values. Similarly, $f$ is used to represent both a generic submodular function as well as the specific FL function (eq. 2). Accumulated attention scores in Algorithm 2 are denoted by $a_t$, thereby not using the set notation or indicating that they are specific to each key.

* The evaluation on the lm-eval-harness benchmark is based on the 15-month-old Llama models, and it is unclear whether the results would generalize to newer, stronger models that may use context more effectively.

* Although a number of papers on KV cache compression are cited, the evaluation only compares against H2 and Attention Sinks as baselines from prior work.
questions: questions: * Is $m_u$ in eq. 3 defined as the attention weight for each key?

* Modern large-scale LLMs typically use multi-query or group-query attention. How does this approach need to be modified—if at all—in order to accommodate these techniques? For instance, would attention weights be accumulated separately for each query index?

* There are no results for BumbleBee 💙 in the LongChat-7B 32k experiments, and it is not clear which flavor of BumbleBee is used for Figs. 3, 6-8.
paper_summary: paper_summary: This paper proposes a submodular optimization framework called BumbleBee to compress the KQ cache in a long context. They aim to maintain the most representative KQ embeddings in history evaluated based on similarity and representativeness. The KV cache is iteratively updated at each decoding step. They validate BumbleBee's effectiveness on extensive tasks with impressive cache budget reductions.
strength_weakness: strength_weakness: 1. Updating the summary KV cache summarization is time consuming. Step 11 in Algorithm 2 and step 5 in Algorithm 1 need to calculate the loss of removing each KV cache embedding. I wonder if this process can be optimized or computed in parallel.
questions: questions: 1. What is the value of \tau_s and \tau_l? Which is more important, global summary or local history?

2. What is the cost of BumbleBee memory and time? Do you save embeddings similarity (as shown in Equation 2) and the past attention weights used in equation (3) in the cache?

3. What are the results of "ALL" for LongChat-7b-32k in Table 2?

4. Why are the attention map values in Figure 2 negative? The colors in this figure are all red and orange, making it difficult to observe the region that is most attended.
paper_summary: paper_summary: This paper proposes BumbleBee, a KV cache summarization method. KV cache is a technique to speed up the decoding procedure by caching key-value pairs computed in the previous steps. KV cache can avoid redundant computations, however, the number of cache entries grows linearly with the number of tokens and can be huge for recent LLMs dealing with long sequences. The proposed method reduces the amount of KV cache by selecting appropriate subsets of key-value pairs. The paper formulates the subset selection problem as an online submodular function maximization problem under cardinality constraints and proposes a simple greedy procedure to select appropriate subsets of key-value pairs for each generation step. 

The paper combines BumbleBee with multiple LLMs and evaluates the performance in multiple tasks. Experimental results show that the proposed method outperforms baseline KV cache reduction methods.
strength_weakness: strength_weakness: 1. The paper does not report the decoding speed of the proposed method. Since the KV cache is introduced to speed up the decoding procedure, I think it is important to show that BumbleBee is faster than the baseline KV cache method.

2. The paper reports experimental results for fixed budget sizes. The paper employs different budget sizes for LM-eval-harness tasks and LongBench tasks, but the reason why doing so is not explained. Budget size is an important parameter that can directly trade the performance with the GPU memory and running time. Therefore, it would be nice to show that the proposed method outperforms the baseline methods if we change the budget size.
questions: questions: 1. Is there any theoretical guarantee for the online submodular maximization algorithm of Algorithm 2?

2. How do we compute $m_u(\cdot)$? 

3. Tab. 5 and Fig. 5 seem not to be referenced from the main body of the paper.

Comments: title: Thank you
Thank you
Thank you
Thank you
comment: Thank you sincerely for your detailed and helpful review! We hope our responses have effectively addressed your concerns and would greatly appreciate the opportunity to further integrate your suggestions to enhance our work. Please let us know if you have any additional questions or comments.
We thank the reviewer for carefully reviewing our rebuttal and taking the time to revise their assessment. We are pleased that our response addressed the main concerns. Your detailed feedback and suggestions throughout the entire process have been invaluable, and we will incorporate them into our updated manuscript. 


Regarding Q3, we provide the LongChat results using the log-based concave function shown in the table at https://imgur.com/a/NTon3cO and will update the corresponding table in our revised manuscript.
We thank the reviewer for carefully reviewing our rebuttal. We are pleased that our response addressed the main concerns and sincerely appreciate your detailed feedback and suggestions throughout the review process.
We thank the reviewer for carefully reviewing our rebuttal and taking the time to revise their assessment. We are glad that our response addressed the main concerns, and we greatly appreciate your detailed feedback and suggestions throughout the review process.
Thank you for the response. In conjunction with the additional results (Tables 1-4) provided in response to the other reviews, most of my concerns have been addressed. I will raise my score accordingly.

It would be helpful to discuss asymptotic compute and memory costs in Section 5. In addition, I think a more intuitive justification for the two specific submodular functions proposed here—and why they complement each other—would be useful to include in the main paper. Readers may also have questions about generalization to different tasks and experimental settings, especially as Fig. 3 suggests that performance can be fairly sensitive to $\lambda$ and model size, so the new ablation experiments could be helpful in the appendix.
Thank you for the response. Your response has addressed most of my concerns. Good luck!
Thank you for the response! I will raise my score since the authors address my concerns appropriately.