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: COLM (2024)
Keywords: transformers, KV cache selection, online summarization, Submodular Optimization, context summarization
TL;DR: A framework using a mixture of submodular functions for dynamic context (KV cache) summarization in LLMs for efficient inference and infinite-context transformers.

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.

Published as a conference paper at COLM 2024
BumbleBee: Dynamic KV-Cache Streaming Submodular Summarization for Infinite-Context Transformers
Lilly Kumari
Shengjie Wang Jeff Bilmes
Tianyi Zhou
Nikhil Sarda
Anthony Rowes
University of Washington, Seattle, *NYU Shanghai, *University of Maryland,
#Google, Carrnegie Mellon University {1kumari, bilmesJ@um @uw.edu
Abstract
The need for Transformer-based Large Language Models (LLMs) to main- tain key-value representations (a KV cache) of previously seen tokens in the GPU memory leads to a significant 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 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 con- text tokens in the key embedding space and their importance computed using accumulated attention attributed to them across different input to- kens. Our framework can work for both the LLM prefill and decoding phases, utilizing offline online versions of our submodular algorithm respectively. While the context sizes grow to be as large only as the sum- mary size, the temporal extent of the contexts may grow unboundedly, justifying the moniker "Infinite-Contet 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.
1 Introduction
The multi-headed self-attention (Vaswani et al., 2017) serves the building block of several state-of-the-an transformer-based models for various tasks such as language understanding and generation (Raffel et al., 2020), image recognition (Dosovitskiy et al., 2021), and recom- mendations (Wu et al., 2023) using Large Language Models (LLMs). enables models to learn long-range dependencies and complex patterns by focusing on multiple sequence sec- tions 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 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 context fragmentation (Dai et al., 2019). Similarly, retrieval-based non-parametric methods to augment an input query with relevant context require constructing external memory bank that is much larger than the

---

Published as a conference paper at COLM 2024
Figure 1: Illustra- tion of the attention mechanism in a self- attention head of one of the BumbleBer's decoder layers. i and denote the local context length and the limited global summary length respectively.
T2
t-
t
Online Submodular Summarization
Global Summary
Local history
Time-point
0
2
3
5
6
Tj
dataset size. In addition, retriever index must be constructed on top of the memory bank.
common approach to address the second problem is to store the previously generated key and value vectors in 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 single token increases as well. While it is possible 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.
offload the KV cache to the host DRAM, this will incur
Existing system-level techniques such as FlexGen (Sheng et al., 2023), PagedAtten- tion (Kwon et al., 2023), and HlashAttention (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 lever- age inference time techniques that reduce the size of the available context at all layers (KV cache) without impacting accuracy. Existing related works such 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 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 subset selection problem where we evaluate the utility of different key-value pairs set instead of independently evaluating each key-value attention state.
In this work, we present novel yet simple approach called BumbleBee' 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 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:
Named after our favorite transformer who is known for being "efficient", highly "acute", and "adaptable".
2

---

Published as a conference paper at COLM 2024
1. Selective attention (Chun & Turk-Browne, 2007; Uncapher & Rugg, 2009) allows to focus on relevant information and filter out distractions or irrelevant details. This plays 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 far by utilizing online submodular summarization.
2. Humans process information in online and dynamic fashion over extended period, possibly all the way to very early in life. This involves relying on our (possibly extremely long) past memory
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 /emory based on the segment's informativeness and its association with the current memory. Like human memory, BumbleBee's ss 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 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 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-Conte. combine the global submodular summary with the most recent local context as shown in Fig. 1, so BumbleBee's 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 '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.
transformer. We further
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 context window of 128k, Claude 2.1 (Anthropic, b) can process 200k tokens while Gemini 1.5 pro (Reid et al., 2024) can handle 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 LLaMA-13B model (Touvron et al., 2023b), the KV cache for 128k tokens would roughly require 40 (num. of decoder layers) x 5120 (hidden size) x2 (kv pair) 2 (fp16) x 128000 105 GB.
Layer
Layer 32 fredd
-158
-05
3049
23-0
52:40 Qua'y pestticn A75 the orig aecunnce
Quarr pooition et the oig sequence
(a) Test sample 1
(b) Test sample
Figure 2: Attention maps for two different Wik articles using LLaMA-7B model.
3

---

Published as a conference paper at COLM 2024
To address the memory costs associated with the KV cache, we focus the following:
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 set of both diverse and im- portant tokens is sufficient to maintain performance. Specifically, we use a LLaMA-7B model (with pre-trained context window size of 2048) for next-token prediction task on randomly sampled articles from wikitext-1033 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 2047. The anti-diagonal pattern shown in Figure 2 shows that there is 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 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.
3 Related Work
Attention speedup: Self-attention is 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. MlshAAtenttonn (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 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 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 past- only quantization mechanism to enable higher precision for self-attention computation. In Yang 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 Paged (Kwon et al., 2023) improves the throughput of L.L serving by reducing fragmentation and redundan dupli-
4

---

Published as a conference paper at COLM 2024
cation.
Applications of submodularity to LLMs: INGENIOUS (Renduchintala et al., 2023) is a tech- nique 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 mix- ture 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 facil- itates the selection of diverse and relevant exemplars for in-context learning with L.MSs in a data-efficient manner. To the best of our knowledge, ours is the first work to apply submodularity for KV cache summarization.
4 Background
In this section, we provide a concise background about self-attention used in transformer- based LLMs. We also briefly discuss submodular function optimization as well as the KV cache mechanism used to avoid re-computations during the decoding stage.
Notation: We denote sequence of tokens as x {x1, X.... xn) where x; E Rd is the 1-th token embedding. In decoder-only transformer model, we denote the number of decoder layers present as Ny and the number of attention heads present in each layer as ng. The self-attention mechanism of head /t in layer (1+ 1), utilizing distinct weight matrices for query (W₂), key (W₂), and value (Wa), operates on the hidden states from the preceding layer, represented as H'. Specifically, we get the following query, key, and value embeddings after the linear projections: Q++1 H' Hg' kl+1 H' Ws, and vl+1 H',,. Row in the key embedding matrix denotes the key vector corresponding to the 1-th token, and this holds for both query and value matrices as well. Given query q1, its attention output o, using scaled dot product attention is computed as shown in Eq. 1. Here, we use S, 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.
exp (kf
0
0j
Σ (kpe))c
a(g Si) 0j
(1)
IyEss, exp (k/ 9t))
(S
KV Cache: Decoder-only transformers operate in an autoregressive manner, predicting to- kens 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, rep- resented by So (i), < n), and caches them in memory. For every subsequent step j i& k, the model reuses the cached states S = {(k, Di < n + j} to compute the attention state (kn+j, Pn+z) of the new token Sn+j. This significantly reduces the floating-point compu- tations 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 that 5 = S, U ((kn+j }.
Note that the KV cache is an approximation and trades off precision for speed. The attention state computation for token sj is limited the sequence available at step j, namely {s/li <
n + i} as opposed to over the entire sequence {sili < n + k). KV caches significantly reduce latency and result in minimal loss in accuracy making it essential ingredient in practical LLM deployment.
Since the KV cache is updated with 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 major bottleneck when dealing with longer sequences whose
5

---

Published as a conference paper at COLM 2024
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: submodular (Bilmes al., 2022) function f 2V + R. defined on the ground set V always has a diminishing return property: f(ø) A) 2 f(o(B) for any v g B and A C S V, where f() := fAA - Intuitively, the gain of an item v, i.e, 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 C V of items. There are many useful such functions, one of them being the facility location (FL. function which similar k-medoids objective (Kaufman & Rousseeuw, 1987). The FL function utilizes similarity scores sim (v, o') computed over every pair v, V of items. A valuation is then the sum of similarities from any item in the ground set to its closest representative in the given set A as shown in Eq. (2).
(A) = Σ max sim(), vVV FEA
(2)
C(A) = Σ HEW
Lu
mu mu(v)).
(3)
DEA
A subset A with high function value indicates that for every item in the ground set, there exists an item in A that is very similar, in other words, 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 (. as a monotone non-decreasing non-negative concave function, and m.(-) is a non-negative weight associated with the u-th feature of every item E A. Due to the diminishing property (.),,, to have large function valuation for 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, 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 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 8A is convex combination
amongst sets of given fixed size, diverse
Algorithm 1 Offline Submodular KV cache Summarization during Prefill/Encodde Phase
1: Input: Submodular functions capturing diversity fFL in the key embeddings space and importance e via attention frequency for layer and attention head h;; mixture function (-) = (-) (1 A)c(-); set of KV attention states Kn {k;)) Vn
{(o;) != corresponding to the n prompt tokens; budget T.-
2: Output: A final summary S,, such that S,, S { [ki, 2;) }i= and |Sm| T. T..
3: Initialize: S, = Ø; compute accumulated attention score vectors an for each key
E [k,) 1- a, denotes accumulated attention scores attributed to key k, across all и query tokens.
4: for for 1 to T, do
5: kimp f argmax ss, 81 (S₁ Ue) -
6: Sn + Sn U (kimp, Vimp)) } where Vimp is the value embedding associated with kimp
7: end for

---

Published as a conference paper at COLM 2024
Algorithm 2 BumbleBee: Streaming Submodular KV cache Summarization for Transformers
1: Input: Submodular functions for diversity fEL in the key embeddings space and impor- tance c w.r.t. attention frequency resp. for layer I and attention head h; mixture function (x(+) (-F + (1 - A)c(-); stream of QKV attention states {q; k; Di) i
budget T,- 2: Output: A running summary St of for every time step f such that S; C v:),
3: Initialize: So Ø Ag where at E R|ša denotes the accumulated attention scores corresponding to keys present in S across t time steps.
4: fort l,..,, do
5:
Update a; for each k E S by adding a aQq,, k, S Kk)
6:
if T, then
7:
S₁ t S, 1 v:)}
8:
Append a(qi, kr, Si) to as s.t. la:1 = [Srl
9:
else
10:
Let = U
kdiscard < argmink)esp 8A (k:/s) k;)
1 (kdiscardd 2discard) if Kdiscard o ki then
St 4-4+, S
12:
Evict af (the accumulated attention score for the discarded key kdiscard) from ar- Append a(q,, k,, Si) to a
14:
15:
end if
16:
end if
17:
end for
of the facility location (Eq. 2) and the feature-based (Eq. 3) functions.
ga (A) = AfFL(A) + (1 A))
(4)
Here, A 2 a hyperparameter that controls the trade-off between representativeness and relevance. Both component functions are monotone, non-negative, submodular, and assumed to be normalized ie., (2) = and (V) = 1 (and the same for c(-)). This normalization ensures the compatibility among the mixture components and the resulting mixture function 8a inherits these properties.
In Alg. 1, we present the offline KV cache summarization pseudo-code to compute summary of the KV cache under cardinality constraint T,. 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 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| = in Eq. 3. The non-negative weight my (k;) in Eq. 3 for key k; is aly which represents the accumulated attention k; receives from the observed input queries. Using the key embeddings, we instantiate the facility similarity matrix computed using pairwise cosine similarities
location function (-)(.) followed by the ReLU transformation.
In Lines 4-7, use the greedy algorithm (Nemhauser et al., 1978; Minoux, 1978; Mirza- soleiman et al., 2015) to perform a ardinality-constraimne Thanks to submodularity, the resultant set is within factor of (1 - 1/e) from the optimal summary (Nemhauser et al., 1978). This offline routine is suitable for KV cache summa- rization 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).
submodular maximization.
In Alg. 2, we outline the KV cache summarization algorithm suitable for streaming setting where we do not have prior knowledge of the complete sequence, restricting us from using existing offline summarization algorithms to obtain 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 fixed size that serves as its representative set. We remark that the computation at each step is at most O 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 Alg. as our final scoring function.
7

---

Published as a conference paper at COLM 2024
In Lines 6-8, we keep caching the (key, value) pairs for different attention heads uncondition- ally until the summarization budget is exhausted. Once the KV cache is full, we utilize the key embeddingg in St and the accumulated attention score vector a to update the submod- ular function components frL (.) and c(-) of the final mixture function gx (.). In Line 10, we create set S that includes the newest incoming (k, 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 removed from S/ and the accumulated attention vector a; is modified in Lines 13-15 to include the attention score associated with the newest key when the discarded key one of the keys in the previous summary set S
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 certain number of interactions, it can be challenging to keep track of the entire conversation context. In such scenarios, BumbleBee can enable modern LLM--ased serving systems to maintain a representative yet important summary of past interactions while minimizing the KV cache-based memory utilization.
Experiments
6.1 Datasets & Tasks
The datasets studied in this paper are derived from three benchmarks: Im-evul-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 Im-eval-harness: OpenbookQA (Mihaylov et al., 2018), COPA (Roemmele et al., 2011), RTE (Wang et al., 2018), MathQA (Amíni et al., 2019), PiQA (Bisk et al., 2020), and Wino- grande (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, 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-sourcee 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 performant platform for such computation than vector-capable GPUs.
CPU side since that is natural, cost-effective, efficient and thus still
6.2 Evaluated Models
We use LLaMA 7B and 13B models (Touvron et al., 2023a) for tasks belonging to the Im-eval- harness 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).
6.3 Baselines & Methods
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 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.
8

---

Published as a conference paper at COLM 2024
Model
Methods
OpenBookQA
COPA
RTE
MathQA
PIQA
Winogrande
LLMMA-13B
47.4
85
73.28
31.86
80.36
75.60
Local
28.4
64
53.43
23.25
58.32
Random Local
27.6
58
54.63
21.76
54.13
50.64
Attn Sinks Local
44.4
80
67.51
29.78
79.22
70.48
H2 + Local
44.2
83
64.98
29.71
79.49
70.32
BumbleBee
71.48
31.02
71.98
BumbleBee
46.6
83
67.15
30.82
79,49
73.01
LLMMA-BB
All Local Random Local
44.6
81
68.95
29.85
80.03
71.51
28.4
56
30,90
23.02
58.22
51.38
28.0
63
21.76
49.30
Attn Sinks Local
41.6
82
58.12
27.40
78.07
67.80
H2 + Local BumbleBee
41.4
78
63.54
27.50
77.31
65.82
43.2
79
68.95
27.74
78.24
68.75
BumbleBee
43.2
79
63.90
28.51
78.56
68.19
Table 1: Results on the few-shot tasks from the Im-eval-harness benchmark using LLaMA 7B and 13B (Touvron et al., 2023a). In the above methods (except All), we do 10x context reduction, so our KV cache summarization budget is 0.1 the input sequence length.
Model
Method
Qasper
MultiFieldQAAenn
HotpotQA
2WikiMQA
QMSum
TREC
LLaMA-7B-chat
All*
19.20
36.80
25.40
20.80
61.5
4k
All (self)
21.60
36.76
27.55
31.58
20.78
64.0
Attn Sinks + Local
14.74
22,93
22.08
29.73
19.25
56.0
(20%)
26.28
BumbleBee (20%)
19.37
27.73
26.14
27.67
20.68
61.5
BumbleBee (20%)
19.59
28.99
21.05
LongChat-7B 32k
H2 (SW, 20%)
21.64
30.72
14.07
15.10
18.11
40.5
BumbleBee (SW, 20%)
23.27
33.16
22.52
17.58
20.27
44.5
Table Results on six datasets from the LongBench benchmark using Llama2-7B-chan 4k (Touvron et al., 2023b) and LongChat-v1.5-2 -1.--7B--32 (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: (xx) = log(1 + x) and (+) power-based: ((x) = -x) where g(y) = ay 1/1 /a + By. Herein, we set B and as 1 and 0.04 respectively to ensure that the function saturation curve is com- patible with the facility location function, the first submodular function component of ga(-) in Eq. 4, thus ensuring that neither component from Eqn (3) Eqn (2) ever always domi- nates the other in the optimization. Selecting the form and curvature of the concave function hyperparameter that is just as important for good BumbleBee
in the submodular mixture performance as is the mixture parameter À in Equ. (4). Hyperparameters were tuned on development and tested on test set (see Appx. D). The extreme cases (A {0,1)), corre- sponding to using only one component from Eqns (3) and (2), yield suboptimal performance.
6.4 Results
LMMeval-harness Tasks: In Table we compare BumbleBee to existing KV cache reduction methods on five-shot learning tasks from the Im-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-constee without any wwwammariizaiion/edduiiiii
setting when All of the KV cache is used
LongBench Tasks: 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 Fl score for evaluation. For QMSum and TREC, we use Rouge-L and accuracy respectively as the evaluation metric. When using the LLaMA-7B-chat-44 model pre-trained with a 4k context window, we truncate the middle part of the input contexts their length exceeds 4k as suggested in Bai et al. (2023). As can be outperforms other context reduction methods such as H2 on four datasets even when the summarization budget is 20% of the entire context size.
in Table 2, the offline version of BumbleBee

---

Published as a conference paper at COLM 2024
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/ /esgments 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/ seeouteee in a memory- manageable 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.
LLaMA-138 model
LL model
10 tocal
Eurdeedee
DUnibiices
ROUGE B
27.75
25
04 Of 0F
04 monorre weight.
mimtttrree on ght.4
mixture Weight
mimttture wiighti
(a) LLaMA-13B
(b) LLaMA-7B
Figure 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 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 3-shot learning setting. BumbleBee outperforms other SOTA cache reduction techniques and even performs better than the All cache setting when using LLaMA-7B.
Sensitivity analysis of mixture weight A: We use the XSUM dataset to show the overall performance sensitivity of BumbleBee to the convex mixture weight hyperparameter A in Eq. 4. In the case of LLaMA-13B model, A E [0.4, 0.8) performs comparably showing that representative subset of the KV cache is preferred for the downstream task. In LLaMA- 7B, we see that for A > 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 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 + .)). The attention layers should be able to support this.
Ackmowledgments: We thank Harshil Dadlani for running synthetic experiments. This ma- terial is based upon work supported by the National Science Foundation under Grant No. IIS- 2106937 and is supported in part by funds from federal agency and industry partners as spec- ified in the Resilient & Intelligent NexiG Systems (RINGS) program Grant No. IIS-2148367.
10

---

Published as a conference paper at COLM 2024
References
Muhammad Adnan, Akhil Arunkumar, Gaurav Jain, Prashant J. Nair, Ilya Soloveychik, and Purushotham Kamath. Keyformer: Kv cache reduction through key tokens selection for efficient generative inference, 2024.
Joshua Ainslie, James Lee-Thorp, Michiel de Jong, Yury Zemlyanskiy, Federico Lebrón, and Sumit Sanghai. Gqa: Training generalized multi-query transformer models from multi-head checkpoints, 2023.
Aida Amini, Saadia Gabriel, Shanchuan Lin, Rik Loncel-Kedziorske Yejin Choi, and Han- naneh Hajishirzi. Mathqa: Towards interpretable math word problem solving with operation-based formalisms. In Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume (Long and Short Papers), pp. 2357-2367, 2019.
Anthropic. Long context window tips docs. atohroochhoo...comm.. com//claude/docs/looog wttp://oooootttttwwooooooottteee a. [Accessed 29-03-2024). 2-1, b. Accessed: 2024-
anthrooic..............
Anthropic. Claude 2.1. https://ww.anthropic.ccmmm 03-25.
Yushi Bai, Xin Lv, Jiajie Zhang, Hongchang Lyu, Jiankai Tang, Zhidian Huang, Zhengxiao Du, Xiao Liu, Aohan Zeng, Lei Hou, et al. Longbench: A bilingual, multitask benchmark for long context understanding. ar preprint ar Xit:2300 14008, 2023.
Iz Beltagy, Matthew E Peters, and Arman Cohan. Longformer: The long-document trans- former. Xiv preprint or Xrio:2004.00150, 2020.
Jeff Bilmes. Submarine: SUBModularity for ARtificial INtelligencE and machine learning. Online Software System, 2024. https://submarit page.
Jeff Bilmes et al. Submodularity in Machine Learning and Artificial Intelligence. CoRR, abs/2202.00132, 2022. URL httpss///aaxiv..rg/ab////// 220p:/222222000222222222222
Jeff A. Bilmes. What HMMs can do. JEICE Transactions on Information and Systems, E89-D (3):869-899 891, March 2006. ISSN 0916-8532. doi: 10.1093
/a//9/////ett////// celiisyyy889 c89-d.33.869.
Yonatan Bisk, Rowan Zellers, Jianfeng Gao, Yejin Choi, et al. Piqa: Reasoning about physical commonsense in natural language. In Proceedings of the AAAI conference on artificial intelligence, volume 34, pp. 7432-7439, 2020.
Sebastian Borgeaud, Arthur Mensch, Jordan Hoffmann, Trevor Cai, Eliza Rutherford, Katie Millican, George Bm Van Den Driessche, lean-Baptiste Lespiau, Bogdan Damoc, Aidan Clark, et al. Improving language models by retrieving from trillions of tokens. In International conference on machine learning, pp. 1206-2240. PMLR, 2022.
Yvonne Brehmer, Viktor Müller, Timo von Oertzen, and Ulman Lindenberger. Episodic memory in childhood and old age: The role of cortical coherence. In Axel Mecklinger, Hubert Zimmer, and Ulman Lindenberger (eds.), Bound in Memory: Insights from Behavioral and Neuropsychological Studies, Berichte aus der Psychologie, pp. 69-91. Shaker Verlag, Aachen, 2004. ISBN 3-8322-2871-3
Kyunghyun Cho, Bart Van Merriënboer, Dzmitry Bahdanau, and Yoshua Bengio. On the properties of neural machine translation: Encoder-decoder approaches. arXiv preprint ar 4rX1::1409.1259, 2014.
Krzysztof Choromanski, Valerii Likhosherstov, David Dohan, Xingyou Song, Andreea Gane, Tamas Sarlos, Peter Hawkins, Jared Davis, Afroz Mohiuddin, Lukasz Kaiser, et al. Rethinking attention with performers. arviv preprint wxXiv:2009.14790 2020.
Marvin M Chun and Nicholas Turk- Interactions between attention and memory. Current opinion in neurobiology, 17(2):177-184, 2007.
11

---

Published as a conference paper at COLM 2024
Karl Cobbe, Vineet Kosaraju, Mohammad Bavarian, Mark Chen, Heewoo Jun, Lukasz Kaiser, Matthias Plappert, Jerry Tworek, Jacob Hilton, Reiichiro Nakano, Christopher Hesse, and John Schulman. Training verifiers to solve math word problems, 2021.
Jerome T Connor, R Douglas Martin, and Les E Atlas. Recurrent neural networks and robust time series prediction. IEEE transactions on neural networks, 5(2):240-254, 1994.
Zihang Dai, Zhilin Yang, Yiming Yang, Jaime Carbonell, Quoc V Le, and Ruslan Salakhut- dinov. Transformer-xl: Attentive language models beyond a fixed-length context. Xiv preprint ar 190: X9t::190.002600, 2019.
Tri Dao, Daniel Y. Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. Fltenttt:: Fast and memory-efficient exact attention with 10-awareness. In Advances in Neural Information Processing Systems, 2022.
Pradeep Dasigi, Kyle Lo, Iz Beltagy, Arman Cohan, Noah Smith, and Matt Gardner. A dataset of information-seeking questions and answers anchored in research papers. In Proceedings of the 2021 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, pp. 4599 2021.
Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-
training of deep bidirectional transformers for language understanding. arXiv preprint ArXiv:1810.08 04800,, 2018.
Roy Dings and Christopher Jude McCarroll. The complex phenomenology of episodic memory: Felt connections, multimodal perspectivity, and multifaceted selves. Journal of Consciousness Studies, 29(11-12):29-55 2022. doi: 10. /0.55776/00511220....1..00 URL https:
//www. www.inngetaconnecc.. com/ content/imp. http:///s/22220000000/0/20/////2
Shichen Dong, Wen Cheng, Jiayu Qin, and Wei Wang. Qaq: Quality adaptive quantization for Ilm kv cache, 2024.
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszkoreit, and Neil Houlsby. An image is worth 16x16 words: Transformers for image recognition at scale. In 9th International Conference on Learning Representations, ICLR 2021, Virtual Event, Austria, May 3-7, 2021. OpenReview.net, 2021. URL https: wwwp:evwwwwnneeouuniiniroee/
Haodong Duan, Jueqi Wei, Chonghua Wang, Hongwei Liu, Yixiao Fang, Songyang Zhang, Dahua Lin, and Kai Chen. Botchat: Evaluating IIms' capabilities of having multi-turn dialogues. ar preprint MiKi:23100136500 2023.
Leo Jonathan Tow, Baber Abbasi, Stella Biderman, Sid Black, Anthony DiPofi, Charles Foster, Laurence Golding, Jeffrey Hsu, Alain Le Noac'h, Haonan Li, Kyle McDonell, Niklas Muennighoff, Chris Ociepa, Jason Phang, Laria Reynolds, Hailey Schoelkopf, Aviya Skowron, Lintang Sutawika, Eric Tang, Anish Thite, Ben Wang, Kevin Wang, and Andy Zou. A framework for few-shot language model evaluation, 12 2023. URL https://zenodo. org/records//12555
Adam Gazzaley and Anna C Nobre. Top-down modulation: bridging selective attention and working memory. Trends in cognitive sciences, 16(2):129-135, 2012.
Suyu Ge, Yunan Zhang, Liyuan Liu, Minjia Zhang, Jiawei Han, and Jianfeng Gao. Model tells you what to discard: Adaptive kv cache compression for Ilms, 2023.
In Gim, Guojun Chen, Seung-seob Lee, Nikhil Sarda, Anurag Khandelwal, and Lin Zhong. Prompt cache: Modular attention reuse for low-latency inference. arv preprint ar axXiv:2311.04934 2023.
Hongyu Gong, Yelong Shen, Dian Yu, Jianshu Chen, and Dong Yu. Recurrent chunking mechanisms for long-text machine reading comprehension. arXiv preprint XxXit2205500805, 2020.
12

---

Published as a conference paper at COLM 2024
Xanh Ho, Anh-Khoa Duong Nguyen, Saku Sugawara, and Akiko Aizawa. Constructing multi-hop qa dataset for comprehensive evaluation of reasoning steps. In Proceedings of the 28th International Conference Computational Linguistics, pp. 6609-6625 2020.
Sepp Hochreiter and Jürgen Schmidhuber. Long short-term memory. Neural computation, (8):1735-1780, 1997.
Sepp Hochreiter, Yoshua Bengio, Paolo Frasconi, Jürgen Schmidhuber, et al. Gradient flow in recurrent nets: the difficulty of learning long-term dependencies, 2001.
Hao Kang, Qingru Zhang, Souvik Kundu, Geonhwa Jeong, Zaoxing Liu, Tushar Krishna, and Tuo Zhao. Gear: An efficient kv cache compression recipe for near-lossless generative inference of Ilm, 2024.
Angelos Katharopoulos, Apoorv Vyas, Nikolaos Pappas, and François Fleuret. Transformers are mnss: Fast autoregressive transformers with linear attention. In International Conference on Machine Learning, pp. 5156-5165, PMLR, 2020.
Leonard Kaufman and Peter J. Rousseeuw. Clustering by means of medoids. 1987. URL https://api. semanticscholarr hrr.aaarr0rr//rrroo1::55566622
Amirhossein Kazemnejad, Inkit Padhi, Karthikeyan Natesan Ramamurthy, Payel Das, and Siva Reddy. The impact of positional encoding on length generalization in transformers. Advances in Neural Information Processing Systems, 36, 2024.
Urvashi Khandelwal, Omer Levy, Dan Jurafsky, Luke Zettlemoyer, and Mike Lewis. Gen- eralization through memorization: Nearest neighbor language models. arXiv preprint arXiv:1911.0017. 2019.
Nikita Kitaev, Lukasz Kaiser, and Anselm Levskaya. Reformer: The efficient transformer. arXiv preprint ar arXi:2000.00455. 2020.
Lilly Kumari and Jeff Bilmes. Submodular span, with applications to conditional data summarization. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 35, pp. 12344-12352, 2021.
Lilly Kumari, Shengjie Wang, Arnav Das, Tianyi Zhou, and Jeff Bilmes. An end-to-end submodular framework for data-efficient in-context learning. In Findings of the Association for Computational Linguistics: NAACL 2024, pp. 3293-3308, 2024.
Litton J Kurisinkel and Nancy Chen. Llm based multi-document summarization exploiting main-event biased monotone submodular content extraction, 2023.
Woosuk Kwon, Zhuohan Li, Siyuan Zhuang, Ying Sheng, Lianmin Zheng, Cody Hao Yu, Joseph Gonzalez, Hao Zhang, and Ion Stoica. Efficient memory management for large language model serving with pagedattention. In Proceedings of the 29th Symposium on Operating Systems Principles, SOSP 23, pp. 611-626, New York, NY, USA, 2023. Association for Computing Machinery. ISBN 9798400702297. doi: 10.1145/300006.3611166 URL https://doi. org/ 10. 10.0..145//10 114511100000...... 3613165.
Mike Lewis, Yinhan Liu, Naman Goyal, Marjan Ghazvininejad, Abdelrahman Mohamed, Omer Levy, Ves Stoyanov, and Luke Zettlemoyer. Bart: Denoising sequence-to-sequeece pre-training for natural language generation, translation, and comprehension. arv preprint X1t:1110.1111 134611,.11111... 2019.
Canjia Li, Andrew Yates, Sean MacAvaney, Ben He, and Yingfei Sun. Parade: Passage representation aggregation for document reranking. arXiv preprint ar Xiv:2000 Stt::2008800093,, 2020.
Dacheng Li, Rulin Shao, Anze Xie, Ying Sheng, Lianmin Zheng, Joseph Gonzalez, Ion Stoica, Xuezhe Ma, and Hao Zhang. How long can context length of open-source LLMs truly promise? In NeurlPS 2023 Workshop on Instruction Tuning and Instruction Following, 2023. URL
www.eeewwwweeeeooomiiiiiiyyo FFNWV.
13

---

Published as a conference paper at COLM 2024
Xin Li and Dan Roth. Learning question classifiers. In COLING 2002: The 19th International Conference on Computational Linguistics, 2002.
Percy Liang, Rishi Bommasani, Tony Lee, Dimitris Tsipras, Dilara Soylu, Michihiro Yasunaga, Yian Zhang, Deepak Narayanan, Yuhuai Wu, Ananya Kumar, et al. Holistic evaluation of language models. Xiv preprint ar Xix:2211.009100 2022.
Zichang Liu, Aditya Desai, Fangshuo Liao, Weitao Wang, Victor Xie, Zhaozhuo Xu, Anas- tasios Kyrillidis, and Anshumali Shrivastava. Scissorhands: Exploiting the persistence of importance hypothesis for Ilm kv cache compression at test time. arXiv preprint axXit:2305.17118, 2023.
Adyasha Maharana, Dong- Ho Lee, Sergey Tulyakov, Mohit Bansal, Francesco Barbieri, and Yuwei Fang. Evaluating very long-term conversational memory of Ilm agents. arXiv preprint ar arXie:2402.17753 2024.
Pedro Henrique Martins, Zita Marinho, and André FT Martins. co-former: Infinite memory transformer. Xrv preprint ar #xie:2109.00301. 2021.
Stephen Merity, Caiming Xiong, James Bradbury, and Richard Socher. Pointer sentinel mixture models, 2016.
Todor Mihaylov, Peter Clark, Tushar Khot, and Ashish Sabharwal. Can suit of armor conduct electricity? new dataset for open book question answering, In Proceedings of the 2018 Conference Empirical Methods in Natural Language Processing, pp. 2381-2391, 2018.
Tomas Mikolov, Martin Karafiát, Lukas Burget, Jan Cernocký, and Sanjeev Khudanpur. Recurrent neural network based language model. In Interspeech, volume pp. 1045-1048. Makuhari, 2010.
Michel Minoux. Accelerated greedy algorithms for maximizing submodular set functions. In Optimization Techniques: Proceedings of the 8th IFIP Conference on Optimization Techniques Würzhurg, September 5-9, 1977, pp. 243. Springer, 1978.
Baharan Mirzasoleiman, Ashwinkumar Badanidiyuru, Amin Karbasi, Jan Vondrák, and Andreas Krause. Lazier than lazy greedy. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 29, 2015.
Jesse Mu, Xiang Lisa Li, and Noah Goodman. Learning to compress prompts with gist tokens, 2024.
Shashi Narayan, Shay B Cohen, and Mirella Lapata. Don't give me the details, just the summary! topic-aware convolutional neural networks for extreme summarization. In Proceedings of the 2018 Conference on Empirical Methods Natural Language Processing, pp- 1797-1807, 2018
George Nemhauser, Laurence A Wolsey, and Marshall L Fisher. An analysis approxi- mations for maximizing submodular set functions-2 Mathematical programming, 14(1): 265-294, 1978.
OpenAL GPT-4 technical report. CoRR, abs abs/2333.007774, 2023.
Reiner Pope, Sholto Douglas, Aakanksha Chowdhery, Jacob Devlin, James Bradbury, Jonathan Heek, Kefan Xiao, Shivani Agrawal, and Jeff Dean. Efficiently scaling trans- former inference. Proceedings of Machine Learning and Systems, 5, 2023.
Jack W Rae, Anna Potapenko, Siddhant M Jayakumar, and Timothy P Lillicrap. Compressive transformers for long-range sequence modelling. arXio preprint ar xrXio:1911.055000 2019.
Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J. Liu. Exploring the limits of transfer learning with unified -ext transformer. Mach. Learn. Res., 21)), jan 2020. ISSN 1532-4435.
14

---

Published as a conference paper at COLM 2024
Machel Reid, Nikolay Savinov, Denis Teplyashin, Dmitry Lepikhin, Timothy Lillicrap, Jean- baptiste Alayrac, Radu Soricut, Angeliki Lazaridou, Orhan Firat, Julian Schrittwieser, et al. Gemini 1.5: Unlocking multimodal understanding across millions of tokens of context. ar arxiv preprint ar Xrix:44033005530 2024.
HSVNS S Kowndinya Renduchintala, Krishnateja Killamsetty, Sumit Bhatia, Milan Aggar- wal, Ganesh Ramakrishnan, Rishabh lyer, and Balaji Krishnamurthy. Ingenious: Using informative data subsets for efficient pre-training of language models, 2023.
HS V N S Kowndinya Renduchintala, Sumit Bhatía, and Ganesh Ramakrishnan. Smart: Submodular data mixture strategy for instruction tuning, 2024.
Luka Ribar, Ivan Chelombiev, Luke Hudlass-Galley, Charlie Blake, Carlo Luschi, and Douglas Orr. Sparq attention: Bandwidth-efficien ilm inference, 2024.
Melissa Roemmele, Cosmin Adrian Bejan, and Andrew S Gordon. Choice of plausible alternatives: An evaluation of commonsense causal reasoning. In 2011 AAAI Spring Symposium Series, 2011.
David E Rumelhart, Geoffrey E Hinton, and Ronald J Williams. Learning representations by back pack-propggg errors. nature, 523(688)) 1986.
Keisuke Sakaguchi, Ronan Le Bras, Chandra Bhagavatula, and Yejin Choi. Winogrande: An adversarial winograd schema challenge at scale. Communications the ACM, 64(9):999-106 2021.
Tomohiro Sawada, Daniel Paleka, Alexander Havrilla, Pranav Tadepalli, Paula Vidas, Alexander Kranias, John Nay, Kshitij Gupta, and Aran Komatsuzaki. Arb: Advanced reasoning benchmark for large language models, 2023.
Noam Shazeer. Fast transformer decoding: One write-head is all you need, 2019.
Ying Sheng, Lianmin Zheng, Binhang Yuan, Zhuohan Li, Max Ryabinin, Beidi Chen, Percy Liang, Christopher Re, Ion Stoica, and Ce Zhang. Flexgen: inference of large language models with single gpu. 2023.
generative
Martin Sundermeyer, Ralf Schlüter, and Hermann Ney. Lstm neural networks for language modeling. In Thirleenth annual conference of the international speech communication association, 2012.
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. Llama: Open and efficient foundation language models, 2023a.
Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Nikolay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, al. Llama Open foundation and fine-tuned chat models. ar preprint Sxiv:2307.092888 2023b.
Melina R Uncapher and Michael D Rugg. Selecting for memory? the influence of selective attention on the mnemonic binding of contextual information. Journal of Neuroscience, 29 (25)88270-8279, 2009.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances neural informa- tion processing systems, 30, 2017.
Alex Wang, Amanpreet Singh, Julian Michael, Felix Hill, Omer Levy, and Samuel Bowman. Glue: A multi-task benchmark and analysis platform for natural language understanding, In Proceedings of the 2018 EMNLP Workshop BlackboxNLP; Analyzing and Interpreting Neural Networks for NLP, pp- 353 355, 2018
Sinong Wang, Belinda Z Li, Madian Khabsa, Han Fang, and Hao Ma. Linformer: Self- attention with linear complexity. arxiv preprint Xix: X00420006004688, 2020.
15

---

Published as a conference paper at COLM 2024
Yumeng Wang and Zhenyang Xiao. Loma: Lossless compressed memory attention, 2024.
Likang Wu, Zhi Zheng, Zhaopeng Qiu, Hao Wang, Hongchao Gu, Tingjia Shen, Chuan Qin, Chen Zhu, Hengshu Zhu, Qi Liu, Hui Xiong, and Enhong Chen. A survey on large language models for recommendation, 2023.
Yuhuai Wu, Markus N Rabe, DeLesley Hutchins, and Christian Szegedy. Memorizing transformers. Xiv preprint ar arXit:2203.08913, 2022.
Guangxuan Xiao, Yuandong Tian, Beidi Chen, Song Han, and Mike Lewis. Efficient stream- ing language models with attention sinks. arXio preprint ar XrXo:2309.1745.3 2023.
June Yong Yang, Byeongwook Kim, Jeongin Bae, Beomseok Kwon, Gunho Park, Eunho Yang, Se Jung Kwon, and Dongsoo Lee. No token left behind: Reliable kv cache compression via importance-aware mixed precision quantization, 2024.
Zhilin Yang, Peng Qi, Saizheng Zhang, Yoshua Bengio, William Cohen, Ruslan Salakhutdi- nov, and Christopher D Manning. Hotpotqa: A dataset for diverse, explainable multi-hop question answering. In Proceedings of the 2018 Conference on Empirical Methods in Natural Language Processing, pp. 2369- 2380, 2018
Lu Ye, Ze Tao, Yong Huang, and Yang Li. Chunkattention: Efficient self-attention with prefix-aware kv cache and two-phase partition, 2024.
Yuxuan Yue, Zhihang Yuan, Haojie Duanmu, Sifan Zhou, Jianlong Wu, and Liqiang Nie. Wkvquant: Quantizing weight and key /value cache for large language models gains more, 2024.
Yury Zemlyanskiy, Joshua Ainslie, Michiel de Jong, Philip Pham, llya Eckstein, and Fei Sha. Readtwice: Reading very large documents with memories. Xrv preprint ar Xri::2100.04441, 2021.
Zhenyu Zhang, Ying Sheng, Tianyi Zhou, Tianlong Chen, Lianmin Zheng, Ruisi Cai, Zhao Song. Yuandong Tian, Christopher Ré, Clark Barrett, Zhangyang Wang, and Beidi Chen. H20: Heavy-hitter oracle for efficient generative inference of large language models, 2023.
Ming Zhong, Da Yin, Tao Yu, Ahmad Zaidi, Mutethia Mutuma, Rahul Jha, Ahmed Hassan, Asli Celikyilmaz, Yang Liu, Xipeng Qiu, et al. Qmsum: new benchmark for query- based multi-domain meeting summarization. In Proceedings of the 2021 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, pp. 5905-5921, 2021.
Zexuan Zhong, Tao Lei, and Danqi Chen. Training language models with memory augmen- tation. Xiv preprint arXit:2205.12674, 2022.
A Supplementary Notes to the Introduction
We remark that certain deep neural networks such as Recurrent Neural Networks (RNN) (Rumelhart et al., 1986) and their variants such as Long-Short Term Memory net- works (LSTM) (Hochreiter & Schmidhuber, 1997) and Gated Recurrent Units (GRU) (Cho et al., 2014) have been widely used for sequential data (Connor et al., 1994; Mikolov et al., 2010; Sundermeyer et al., 2012). Despite their strong performance, RNNs suffer from vanish- ing and exploding gradient problems (Hochreiter et al., 2001) as well property similar to dynamic Bayesian networks, where the future and past are separated by (or independent given) the present, thereby limiting their ability to capture extremely long-tern term-termmo dependencies without extremely large state spaces (Bilmes, 2006). Thanks to its transformer underpinnings, BumbleBee model does not have any such Markov property even in the "infinite-context" case.
inherent Markov
16

---

Published as a conference paper at COLM 2024
B Experiments on Synthetic Data
We are given a dataset having inputs defined as x = X2,.,,, where x E R denotes the t-th input token and dn is the input feature dimensionality. Similarly, we have the outputs defined as y y2.......... .n. where y; E Roout and dout is the dimension of the output embedding. We hypothesize that BumbleBee can achieve significant improvements in the task involving estimating y/: E y over existing methods across different dimensions if y: is function of both of the following:
x1-1} S x where (a) The local context window, local, a sequence {x₁ ty 11--1) (,, is the local context window length.
(b) The summary sequence, summi, comprising relevant vectors chosen from { x x2,...... ................................ x S Specifically, let S, C {1,2,..., h - 1)} and say S; = ($1, 92, ,k,,,, is ordered by time (or position) in the original sequence so that $1 $2 <... < $k. Then we have that summ, (Xs Xsy X
Also, assume in this discussion that 2 h + k to avoid any negative indices.
To understand the context where expect BumbleBee to function well, we define (undesirable) conditional independence property as follows:
yt 1I. summ, llocal
(5)
This property states that the predicted variable y/ is independent of the summary given the local context this means that the local context is sufficient to predict y/ and once we have this local context, the summary is not needed. This property is precisely what we must not have in the data for BumbleBee to function properly. That is, this property should not hold for any of the estimation tasks under consideration.
In terms of conditional mutual information, the property that we would like to have is that symm; llocal;) > 0, meaning that even if we have the local context, the summary is still informative of yt- This property being true means BumbleBee should be effective.
To demonstrate the effectiveness of the BumbleBee on such tasks, we define an ideal experi- mental setup as follows:
We obtain x ~ (O, 1), for each E {(k + h), (k + / + 1),..., NN, we sample local, E using sliding window mechanism and summ, uniformly random based on the window size and summary budget size defined previously. We define generative process to obtain y as follows:
y f ground truth (summ, local)
fground-truth is chosen in such a way that the summary and local context tokens both equally important to y. This is done by obtaining,
yocal f fground-truth (local)
summ gground-truth (ornnntttutSuuunnnn
where the above model fground-truth is achieved via standard transformer masking/ padding (setting the corresponding attention values in the matrix to zero), and then comparing their Mean Squared Error (MSE) scores with respect to to be comparable as shown in Table Note that, we are not training the ground-truth model, but rather initializing with seed for which the predictions made using (a) only the local tokens and (b) only using the summary tokens receive similar MSE values. In terms of the architecture, fground truth is a single-layer encoder-decoder transformer model with linear projection layer applied to the outputs of the decoder. Here, h 32, k 32, di, = 16, dout
masding/padding
80,000.
To demonstrate our initial hypothesis, we now perform training and obtain three different
learnt models, as follows:
17

---

Published as a conference paper at COLM 2024
Model
MSE (Mean Squared Error)
ground-truth (summ)
8.2689 10
ground-truth (local)
8.3301 10-4
Table 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
paummlood trained using both summaries and local context tokens. learnt
plocal Jearnt
trained using only the local context tokens.
c. Jearnt trained using only the summary tokens.
For the entire training setting, unless mentioned otherwise, we perform 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 performs significantly better than flamd and psumon
Jeuryt exhibiting the importance of both summary and local context input tokens for the final predictions. In Fig. 5, provide qualitative results showing the predicted values using different learned models.
[jaz-
064
334
002
342
bx2
010
038
0.00
03663
06/0
1.00
3:0
0190
00015
6405
0.50
0003
0015
0.00
C
(a)
(b)
(c)
Figure 4: Performance comparison on subset of the held-out set for different models trained using (a) summaries and local context; (b) only local context; (c) only summaries. 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 segmented into regions where the summary is fixed, meaning S, = S, 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.
Model
MSE (Mean Squared Error)
MAE (Mean Absolute Error)
ssummllccal
9.14 10-7
7.7279 10-4
learnt
flocal
4.21 10 10-5
5.09 x 10 3
fourar
4.83 x 10-4
1.685 x 10 10-2
Table 4: Model performance on held-out test set
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, floal We compare the performance of these models on metrics like MAC (Multiply and Accumulate), runtime memory, and MSE. The MAC and runtime memory metrics computed for single training example in a mini-batch.
Table clearly shows the effectiveness of having mechanism to summarize the historical
18

---

Published as a conference paper at COLM 2024
Groand-guull wolaeePPollillens
Ground muth raluesPrediciiens
0.75
0.70
0.70
0.65
0.65
Poolution mnsle traincl
Aocal parnes
Productien caudel maned land
Pooductonn trannel the
Preductions
local
0.60
Tramh
ssung only
0.60
Cioond-Trrsh
Predictioss using securaaties
5000 10000
12000
56000
SWINK
600000
64000
Time
Time
Figure 5: Plot for error comparison with respect ground-truth data of different trained models for two different time windows
context over naively fitting longer inputs in the transformer's context in terms of metrics such as runtime memory used, MSE, etc.
Model
MSE (Mean Squared Error)
MAC
Runtime Memory
summ,loca learnt
9.14 x 10 7
6.33 x 105
3.43 MB
xlocalx10
3.43 x 10 4
3.12 x 106
5.98 MB
learnt
Table 5: Performance comparison of model trained using a naive large local context BumbleBee model
C Detailed Analysis of Results
For one of the Im-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 streaming setting. In Figures & 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 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 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, 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 ms/token) for two different context reduction ratios, demonstrating faster decoding speed as the context reduction ratio becomes higher.
D Reproducibility
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, use internal highly optimized submodular toolkit, and we plan to
19

---

Published as a conference paper at COLM 2024
H2
BumbleBee
H2
BumbleBee
Layer 1, Head
Layer 1, Head
Layer 21, Head
Layer 21. Head
Query stoo
31-
31
Quary tine
51
51
91
91
LLaer 5, Head 32
Layer 5, Head
Layer 25, Head
Layer 25, Head 40
Quary tine-atoo
11
11
51
22
71
Layer 10, 31
Layer 10, 31
Layer 30, Head
Layer 30, Head
Query wwyytin--dee slop
11
31
51-
Query time
91-
91
Layer 13, Head 10
Layer 13, Head 10
LLyer 36, Head
Layer 36, Head
11
11
Quary time-ateee
51
51
Qvery
n1
Layer
Layer
Layer 40,
Layer 40,
Quary time-ttp
21
331
51
91
100
100
100
300
Figure 6: Test sample 1: Visualization of the keys selected at different time-steps (y-axis) when using H2 and BumbleBee for online KV cache reduction and summarization respectively. The x-axis represents the keys, and the y-axis represents the queries. We show the evolution of the selected KV cache (marked by purple points) as we process an incoming query across different attention heads and determine how to update the KV cache under the budget constraint.
Context reduction ratio
Original Context Length
16k
100k
59.30
OOM
47.49
4.16
71.50 0.10
10:1
39.74
1.31
48.16 + 0.09
Table 6: Decoding speed (in ms 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 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 e = le-5 to compute the offline KV cache summaries for each of heads in a parallel fashion. Here, h denotes the number of attention heads present in one decoder layer, and / denotes the number of decoder layers.
For similarity matrix computation, pairwise similarities sim(i,)) 2 0. We have also experimented with other cosine similarity- based metrics such as + ji) and [cos(i,)) i)l but find that ReLU(cos(i,))) j)) works the best.
use ReLU truncated cosine similarity to ensure
Mixture weight A: the tasks from the Im-eval-harness benchmark (Gao et al., 2023),
20

---

Published as a conference paper at COLM 2024
H2
Bumbleßee
H2
BumbleBee
Layer 1, Head
Layer 1, Head
Layer 21, Head
Layer 21. Head
Query stop
31-
31
Quary tine
51
51
91
91
LLayer 5, Head 32
Layer 5, Head
Layer 25, Head
Layer 25, Head 40
Quary time-atee
11
11
J1-
51
22
Query
71
Layer 10, 31
Layer 10, 31
Layer 30, Head
Layer 30, Head
Query wwyytin--dee slop
11
31
51-
Query time
91-
91
Layer 13, Head 10
Layer 13, Head 10
LLyer 36, Head
Layer 36, Head
Quary time-sten
11
11
51-
51
Query
71
Layer
Layer
Layer 40,
Layer 40,
Quary time-tte
331
57
91
100
100
100
100
Figure 7: Test sample 2: Visualization of the keys selected at different time-steps (y-axis) when using H2 and BumbleBee for online KV cache reduction and summarization respectively. The x-axis represents the keys, and the y-axis represents the queries. We show the evolution of the selected KV cache (marked by purple points) as we process an incoming query across different attention heads and determine how to update the KV cache under the budget constraint.
E (0.2,0.3) perform the best across both evaluated models, iee, LLaMA 7B and 13B. However, we did not perform more fine ane-ggrined
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 held-out validation set to tune the mixture weight and set À = 0.2 when using LLaMA2-7B and A = 0.4 for LLLMA-13B.
Compute: we use an NVIDIA-A100 GPU to perform our inference-based experiments and 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.
E Complexity Analysis
In BumbleBee (Alg. 2), we maintain global summary of size Ty- So, constructing the pairwise similarity matrix for fri has time complexity of O x d). Identifying the item with the least conditional gain (Line 11 of Alg. 2) requires O 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 OT,, d + ?32) for each new token.
Compute costs:
Memory Costs: Caching the similarity matrix incurs O memory costs.
21

---

Published as a conference paper at COLM 2024
Layer 2, Head 16
H2 selected KV cache
BumbleBee's selected KV cache
50
50
25
2000
L-SNE 2
t.SNE
Time steps
1500
1500
-25
1000
1000
-50
-50-
-75
500
-75
500
-100
-50
50
100
-100
-50
50
100
t-SNE 1
t-SNE
Layer 11, Head 15
selected cache
BumbleBee's selected cache
75
3000
3000
50
2500
2500
1-SNE
2000
1-SNE 2
2000
Time steps
0
1500
1500
-25
-25
1000
1000
-75
500
-75
500
-100
-50
-50
50
t:SNE
t:SNE
Layer 30, Head 29
H2 selected KV cache
BumbleBee's selected KV cache
100
300D
3080
2500
2500
25
2000
-SNE7 2
2000
S-SN
Time steps
0-
1500
1500
-25
-25
1000
1000
-50
-50
-75
-75
t:SNE
t-SNE
Figure 8: t-SNE visualization of the key embeddings for different attention heads in the LLaMA-7B-chat-4 model. The keys selected by H2 and BumbleBee as part of the final KV cache summary are marked by A.
22


Reviews:
[{"review_id": "c0Ku7fiRbF", "reviewer": "Reviewer_v71g", "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.", "questions": "", "limitations": "", "rating": 7, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": true, "final_justification": "", "rebuttal": "Thank you for your insightful feedback and for appreciating the importance of our work.\n\n**RTR1 & Q1**: Existing streaming submodular optimization algorithms like Sieve-Streaming (Badanidiyuru et al., 2014) and Sieve-Streaming++ (Kazemi et al., 2019) evaluate the final summary only after processing the entire data stream. However, real-time LLM decoding requires continuously maintaining optimal summaries, equivalent to generating the highest possible minimum-quality summary. Integrating these into transformers complicates theoretical proofs concerning performance and convergence due to the transformer's non-linear nature. Analyzing facility location function in a streaming setting is also challenging as the ground set changes with the stream. \n\nLine 11 in Alg. 2 is exact at each step - the reason, given the knowledge of the sequence observed so far at time t and running summary $S_{t-1}$ of fixed size $\\tau_s$, it discards the item with least marginal gain, ensuring the resultant summary's function valuation is maximum among all $(\\tau_s+1)$ candidate summaries.\n\n\n**RTR2**: Besides mixture weight, $\\lambda$, our framework involves another hyperparameter: similarity measure in Eq 2. Since self-attention head similarities between query and key embeddings use dot products, we also use a dot-product-based similarity measure. Truncating negative entries to get non-negative similarity values (to preserve submodularity) worked best across three tasks (https://imgur.com/a/QR42zP0).\n\nPlease find ablation results at https://imgur.com/a/vKUX9QT. Using mixture function with non-zero $\\lambda$ to capture both representativeness and token importance performs the best.\n\n\n**Q2**: The choice of the concave function $\\phi_u$ used to define the feature-based function is another hyperparameter. We experimented with two functions and found both perform equally well except on RTE dataset. These functions were selected for their similar diminishing returns behavior to the facility location function to ensure one did not dominate the other. The chosen $\\phi_u$ determines how the marginal gain of adding a new key diminishes in the context of the current summary. For larger models with grouped-query attention, one can selectively choose the concave function $\\phi_u$ in Eq. 3 to control the degree of diminishing returns for different query heads in a query block, weighting them based on importance such that $c(A) = \\sum_{u \\in U} w_u \\phi_u(\\sum_{v\\in A} m_u(v))$ where $w_u > 0$.", "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. **Novelty**: The approach of using a submodular optimization framework to balance the diversity and importance of key-value pairs in the memory cache of Transformers is novel and addresses a significant challenge in handling long contexts with LLMs effectively.\n\n2. **Technical Soundness**: The paper is technically sound with a clear explanation of the submodular function utilized.\n\n3. **Experimental Validation**: The paper presents comprehensive experiments across multiple datasets and compares the BumbleBee framework with several baselines, showing improvements in reducing memory overhead while retaining or even improving model accuracy.", "reasons_to_reject": "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.\n\n2. **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_to_authors": "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?\n\n2. 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?", "extra_scores": {}}, {"review_id": "kwgX8ILApn", "reviewer": "Reviewer_wh5U", "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.", "questions": "", "limitations": "", "rating": 7, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": true, "final_justification": "", "rebuttal": "Thank you for your detailed feedback! \n\n**RTR1**: In BumbleBee, we maintain a running summary of fixed size $\\tau_s$, so our runtime complexity will not be a function of sequence length n. For FL function, as each key-value is observed, we compute its similarity to items in the KV cache summary incurring $\\tau_s d$ cost. Computing $g_\\lambda$ for each item in line 11 of Alg. 2 takes $\\mathcal{O}(\\tau_s^2)$ including FL evaluations. Multiple functions for different heads can be evaluated in parallel (as in our implementation). \n\n\n**RTR2**: H2 is a specific case of our formulation with $\\lambda$ set to 0 and $\\phi_u$ as an identity function. Our submodular objective uses the FL function to capture diversity in key embedding space and a concave function to transform modular attention scores into a strictly submodular objective. For larger models using grouped-query attention GQA, the concave function $\\phi_u$ in Eq. 3 can be chosen selectively to control the rate of diminishing returns for different query heads in a query block, weighting them based on the head's importance: $c(A) = \\sum_{u \\in U} w_u \\phi_u(\\sum_{v\\in A} m_u(v))$ where $w_u > 0$.\n\n\n**RTR3**: Thank you for pointing this out. We will fix this in our updated manuscript.\n\n\n**RTR4**: We tailored our analysis on lm-eval-harness datasets to align with the published H2 results that used LLaMA-1 models. For additional datasets, we employed LLaMA-2 models. As models become more proficient in utilizing the input context, we expect the mixture functions $f$ and $c$ in Eq 4, which are based on key embeddings and accumulated attention scores, to more effectively capture key-space representativeness and context relevance.\n\n\n**RTR5**: We benchmark against baselines suitable for streaming. FastGen's need for targeted profiling across all attention heads over the entire sequence makes it unsuitable for our setting. Similarly, gisting, which trains an LLM to condense input prompts into gist tokens, is incompatible with our setting where no further finetuning is needed.\n\n**Q1**: $m_u(v)$ represents the total/accumulated attention the key $v$ receives from the observed input queries.\n\n\n**Q2**: Our current formulation (Eq 4) can be easily applied to GQA. See our response to RTR2.\n\n\n**Q3**: Fig 3, 6, and 7 use online BumbleBee algorithm, while Fig. 8, based on the LongBench dataset, uses the offline BumbleBee variant. We are working on LongChat results using the log-based concave function and will post them soon.", "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 paper is well written and the proposed approach is clearly motivated.\n\n* Submodular formulations are an intuitive choice to encourage diversity in the compressed KV cache.\n\n* Experiments on a variety of benchmarks and models suggest the performance of this approach is fairly consistent.", "reasons_to_reject": "* 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.\n\n* 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.\n\n* 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.\n\n* 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.\n\n* 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_to_authors": "* Is $m_u$ in eq. 3 defined as the attention weight for each key?\n\n* 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?\n\n* 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.", "extra_scores": {}}, {"review_id": "9ZrKXcVbYW", "reviewer": "Reviewer_aHg4", "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.", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the review and insightful comments! Please see our responses below:\n\n**RTR1**: This step to determine the key-value item to discard has a time complexity of $\\mathcal{O}(\\tau_s^2)$ (including FL evaluations). This process, however, can run concurrently with attention output computations in the subsequent decoder layer. Also, the summary updates for different self-attention heads within a layer are performed in parallel. Hence, the overall method is scalable. We will clarify this in the next version.\n\n\n**Q1**: On lm-eval-harness and XSUM datasets, both $\\tau_l$ and $\\tau_s$ are set to 10% of the input sequence's length. For LongBench datasets with offline BumbleBee, only tokens from the test query/input are kept in the local window. The associated context (e.g., passages, transcripts) is summarized to 20% of its length at each attention head, ensuring the summarized KV cache generalizes across different test inputs. The ratio between  $\\tau_l$ and $\\tau_s$ should be task-dependent, for example, for context summarization tasks such as XSUM, a larger portion of the KV cache budget should be dedicated to the summary to cover diverse and representative aspects across the provided context (https://imgur.com/a/YWZ2aNK).\n\n\n**Q2**: In our current setup, we cache both accumulated attention scores and similarity matrix. Also, the per-layer summary update is overlapped with the attention output computations in the subsequent decoder layer. \n\n*Memory costs*: $\\mathcal{O}(\\tau_s)$ for each head when computing similarity matrix on-the-fly. $\\mathcal{O}(\\tau_s^2)$ when caching the similarity matrix.\n*Compute costs*: $\\mathcal{O}(\\tau_s^2 \\times d + \\tau_s^2)$ when similarity matrix is computed real-time. Here, $d$ is the key-embedding dimensionality and the 2nd term corresponds to the summary update costs (Line 11 of Alg. 2) . We cache the similarity matrix and thus only need to compute the similarity of the incoming item to others in the summary. In this case, therefore, it costs overall $\\mathcal{O}(\\tau_s \\times d + \\tau_s^2)$ for each new token.\n\n\n**Q3**: We do not report the results for ALL due to encountering OOM error on our GPU (A100 80 GB) when using a longer context (>32k tokens).\n\n\n**Q4**: In Fig. 2, we plot the logarithm of the attention scores as mentioned in the paper. For better visualization, we will add another figure that uses histogram equalization to improve contrast (similar to https://imgur.com/a/4bph47f).", "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. A novel method to compress the KV cache.\n2. Large reduction in KV cache memory.", "reasons_to_reject": "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_to_authors": "1. What is the value of \\tau_s and \\tau_l? Which is more important, global summary or local history?\n\n2. 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?\n\n3. What are the results of \"ALL\" for LongChat-7b-32k in Table 2?\n\n4. 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.", "extra_scores": {}}, {"review_id": "w6OqzXITrx", "reviewer": "Reviewer_zX3y", "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. \n\nThe 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.", "questions": "", "limitations": "", "rating": 7, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your positive feedback and questions!\n\n**RTR1**: Please see https://imgur.com/a/1Bo8d8v\n\n\n**RTR2**: For LongBench datasets with offline BumbleBee, only tokens from the test input are kept in local window. The associated context (passages, transcripts) is summarized to 20% of its length at each attention head, ensuring the summarized KV cache generalizes across different test inputs. On lm-eval-harness, we provide results for the smallest summarization budget where performance loss is minimal compared to the no KV cache reduction setting.\n\n\n**Q1**: The current algorithm does not have a theoretical guarantee w.r.t. the non-linear transformer. The reason for this is multi-faceted: (1) Algorithms like Sieve-Streaming (Badanidiyuru et al. 2014), Sieve-Streaming++ (Kazemi et al. 2019), etc. assess the final summary only after the entire data stream has been processed, which is not conducive to the needs of real-time LLM decoding where continuously optimal summaries are essential (2) Integrating these into transformers further complicates theoretical proofs concerning performance or convergences due to non-linear nature of transformers (3) Unlike simpler feature-based functions, analyzing facility location functions in a streaming setting is not viable as the ground set changes as the stream is updated.\n \nOur summary update algorithm (Line 11, Alg 2) is, however, exact at each step - the reason, given the knowledge of the sequence observed so far at time t and running summary $S_{t-1}$ of fixed size $\\tau_s$, it discards the single item with the least marginal gain ensuring that the function valuation of resultant summary is maximum amongst all $(\\tau_s+1)$ candidate summaries.\n\n\n**Q2**: $m_u(v)$ represents the accumulated attention the key $v$ receives from observed input queries. The subscript $u$ represents different features that one might want to model to capture context relevance, for example, in grouped-query attention (GQA), each $u$ could correspond to a specific query head meaning $|U|$ would be the size of each query group.\n\n\n**Q3**: Thank you for pointing this out. We will fix this in our updated manuscript. Table 5 corresponds to the synthetic data experiment where we naively increase the size of the local context window by 10 times the original and compare it to the setting where the model uses the summarized context. Fig. 5 supports the findings of Table 4 by qualitatively showing the predicted values using different learned models.", "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. The idea of formulating a KV cache reduction problem as a subset selection problem with a submodular objective function is interesting. I agree that employing submodular objective functions is a reasonable choice for the KV-cache reduction.\n\n2. The proposed method is simple, and we do not need additional complex procedures like fine-tuning. \n\n3. Experimental results show that the proposed method outperforms baseline KV-cache reduction methods.\n\n4. The paper is well-organized and easy to follow.", "reasons_to_reject": "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.\n\n2. 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_to_authors": "1. Is there any theoretical guarantee for the online submodular maximization algorithm of Algorithm 2?\n\n2. How do we compute $m_u(\\cdot)$? \n\n3. Tab. 5 and Fig. 5 seem not to be referenced from the main body of the paper.", "extra_scores": {}}]

Author rebuttal:

Meta-review: 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.
Decision: Accept