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

Thank you for your positive feedback and questions!

RTR1: Please see https://imgur.com/a/1Bo8d8v

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.

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.

Our 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.

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.

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.

Thank you for the review and insightful comments! Please see our responses below:

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.

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).

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.

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. 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.

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).

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).

Thank you for your detailed feedback!

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).

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$.

RTR3: Thank you for pointing this out. We will fix this in our updated manuscript.

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.

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.

Q1: $m_u(v)$ represents the total/accumulated attention the key $v$ receives from the observed input queries.

Q2: Our current formulation (Eq 4) can be easily applied to GQA. See our response to RTR2.

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.

Thank you for your insightful feedback and for appreciating the importance of our work.

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.

Line 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.

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).

Please 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.

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$.