Long Context Compression with Activation Beacon

## Abstract

Abstract Long context compression is a critical research problem due to its significance in reducing the high computational and memory costs associated with LLMs. In this paper, we propose Activation Beacon, a plug-in module for transformer-based LLMs that targets effective, efficient, and flexible compression of long contexts. To achieve this, our method introduces the following technical designs.
1) We directly compress the activations (i.e. keys and values at every layer), rather than leveraging soft prompts to relay information (which constitute a major bottleneck to encapsulate the complex information within long contexts).
2) We tailor the compression workflow, where each fine-grained input unit is progressively compressed, enabling high-quality compression and efficient computation during both training and inference.
3) We train the model through compression-based auto-regression, making full use of plain texts and instructional data to optimize the model’s compression performance.
4) During training, we randomly sample a compression ratio at each step, teaching the model to support a wide range of compression configurations.
Extensive evaluations are conducted on various long-context tasks whose lengths (e.g., 128K) may far exceed the maximum training length (20K), such as document understanding, few-shot learning, and Needle-in-a-Haystack.
Whilst existing methods struggle to handle these challenging tasks, Activation Beacon maintains a comparable performance to the uncompressed baseline across various scenarios,
achieving a 2x acceleration in inference time and an 8x reduction of memory costs for KV cache.
Our data, model, and code have been released at https://github.com/FlagOpen/FlagEmbedding/ .

## 1 Introduction

Large language models (LLMs) need to process long contexts to accomplish many important tasks, such as long-document understanding [], long-content creation [], and long-term memorization/reasoning [].
To address these needs, modern LLMs are built with extended context windows (e.g., 128K) that enable remarkable long-context processing capabilities [; ; ].
Despite their effectiveness, LLMs encounter efficiency challenges in processing long contexts.
On one hand, transformer-based LLMs incur substantial computational costs due to the quadratic complexity of self attention. On the other hand, they require tremendous GPU memory to hold the KV cache of the entire sequence for faster decoding.
Both computation and memory costs increase as the context length grows.

A wide array of studies are dedicated to alleviating efficiency issues, among which context compression is a promising direction [; ; ; ; ]. This approach aims to compress raw input into more concise representations, allowing the generation process to be conditioned on a shorter context.
Therefore, it helps to reduce both computation cost of inference and memory cost from KV cache, while also enabling the processing of longer inputs than the LLM’s built-in context window.

Despite the current progresses, it it remains a tough challenge to compress long contexts.
Specifically, existing methods usually summarize the context into a few soft tokens [; ], which constitute the major bottleneck to summarize the complex information within long contexts.
Besides, they try to compress the context “all-at-once”, lacking a fine-grained handling of the detailed information.
Moreover, these soft tokens must be re-encoded before generation, resulting in inferior efficiency in both training and inference.
Lastly, these methods are learned to compress with a fixed number of soft tokens, thus, it’s hard to customize the compression ratio for downstream tasks. While some alternamtive methods focus on deleting unimportant tokens [; ], they depend on the input question to estimate the token importance, limiting their efficiency in real-world multi-turn scenarios.

To address the above challenges, we present Activation Beacon (Figure 1), a plug-in module to transformer-based LLMs that enables effective, efficient, and flexible compression of long contexts. Activation Beacon is featured with the following technical designs.

First of all, we introduce a new special token, called the beacon token $\langle\text{b}\rangle$. The context is distilled into beacon tokens’ activations (i.e. keys and values at every layer), whose capacity are large enough to encapsulate the complex information within long contexts.

Next, we tailor the compression workflow, where each fine-grained context unit is progressively compressed.
Specifically, the long context is partitioned into equal-size chunks.
Each chunk is further split into fine-grained units of size $\alpha$ where $\alpha$ is the desired compression ratio.
A group of beacon tokens are interleaved with these units (one beacon token is dispatched to the end of every unit).
The LLM encodes one chunk at a time, distilling the chunk’s information into beacon tokens’ activations during self attention.
After encoding, the raw tokens’ activations are discarded; while the beacon tokens’ activations are accumulated and reused for encoding following chunks.
This progressive workflow brings forth several advantages:
1) It can handle inputs longer than the backbone LLM’s context window as the chunk size is small.
2) It achieves fine-grained compression since the attention scope of each beacon token is differentiated.
3) By caching and reusing activations, it facilitates contiguous gradient propagation in training, avoids re-encoding overhead in inference, and allows for incrementally updating the compression results in multi-turn scenarios.

Finally, Activation Beacon is learned with compression-based auto-regression to optimize the generation quality conditioned on the compressed context. Thanks to high sample efficiency, the model can be effectively trained with 1B plain corpus and 30K fine-tuning samples (maximum context length is 20K), which can be quickly accomplished. During training, we randomly sample the compression ratio for each chunk, enhancing the model’s flexibility to tackle different compression ratios in downstream tasks. Note that all beacon tokens share the same token embedding, one can use arbitrary number of beacon
tokens to achieve the desired compression ratio by repeating.

In our experiments, Activation Beacon is applied to Llama-2 [] and Qwen-2 []. We evaluate the resulted models on a variety of long-context tasks (whose lengths may be much longer than the training length, e.g., 128K), such as document understanding, few-shot learning, and Needle-in-a-Haystack. Whilst existing methods struggle to handle these challenging tasks, Activation Beacon maintains a comparable performance to the uncompressed baseline across various compression configurations, meanwhile achieving 2x acceleration and 8x KV cache reduction. Moreover, the LLM’s original capabilities on short context is well preserved.

Figure: Figure 1: Overview of Activation Beacon. The context is partitioned into chunks. Each chunk is further split into fine-grained units and interleaved with beacon tokens according to a compression ratio (2 in the figure). The LLM encodes one chunk at a time, compressing the context into beacon tokens’ activations, which are accumulated and reused for encoding following chunks.
Refer to caption: x1.png

## 2 Related Works

Recently, processing long context has become a fundamental capability of modern LLMs [; ; ; ].
The recipe of context window extension is roughly the same: modifying the rotary position embedding [] by extrapolation and interpolation [; ; ; ], and leveraging long-dependency data in both the pre-training and post-training stage. Despite the impressive progress in effectiveness, LLMs face significant challenges in efficiency. There is significant computational cost due to the quadratic complexity of transformer, and huge memory cost because LLMs need to hold the KV activations of the entire sequence on GPU for faster decoding. Multiple threads of research endeavour to reduce these costs, which are discussed as follows.

Sparse Attention.
Conventional sparse attention methods require re-training a model from scratch using the designated sparse patterns [; ]. However, extensive recent studies have identified that the attention pattern of LLMs are naturally sparse despite they are densely trained [; ; ; ].
They also propose to dynamically set appropriate sparse patterns for each head so that the attention mass can be largely preserved, leading to competitive performance against the full-attention method with reduced computation.
However, these methods require holding all KV activations on chip to dynamically determine the optimal sparse patterns, making them unsuitable for KV cache reduction.

KV Compression.
This line of research focuses on compressing the KV activations to reduce the attention computation as well as the cache size.
Since the KV activations are per-layer, per-head, per-token, and per-channel float numbers, they can be reduced from all the five dimensions (including the numerical dimension).
For example, CLA [] shares the KV cache across multiple layers; GQA [] compresses multiple key/value heads into a single one; MLA [] compresses the channels into fewer and more compact ones; and KIVI [] quantizes the numerical value in the activations.
The token-wise compression, as introduced in the following, is also known as context compression, which is orthogonal to the compression along other dimensions and hence can be jointly used.

Context Compression. This type of methods aim to compress the raw context into shorter yet more compact representations. Existing studies are usually tailored for compressing short context (less than 1K), which tend to be sub-optimal for long-context compression.
Specifically, Gisting [] compresses the user instruction into gist activations all at once. As a result, it cannot process context longer than the backbone LLM’s window.
ICAE [] and AutoCompressor [] alleviate this problem by segmenting the long context into chunks and compressing each chunk.
However, both of them compress the context into soft tokens, which are the major bottleneck to encapsulate the complex information in long contexts. Their compression workflow also lacks fine-grained handling of the chunked inputs, resulting in inferior compression quality.
Moreover, these soft tokens require re-encoding before generation, which introduces extra overhead.
Lastly, since the number of soft tokens are pre-defined, it is hard to flexibly assign the compression ratio for downstream tasks.
CCM [] is specifically designed for compressing conversations in online chatting, which cannot be used in general long context tasks such as long document understanding.
Another branch of methods [; ] propose to delete unimportant tokens to realize compression. However, they depend on the input question to accurately estimate the token importance, leading to low efficiency in real-world multi-turn scenarios.
Compared with existing approaches, Activation Beacon is able to achieve more effective, efficient, and flexible compression.

## 3 Methodology

LLMs accomplish arbitrary tasks in the form of next-token prediction. Formally, given the context $X=[x_{1},\dots,x_{n}]$, the LLM generates the next token based on all preceding tokens and its well-trained parameters: $\Pr(x_{n+1}\mid x_{1},\dots,x_{n};\boldsymbol{\Theta})$.
Transformer-based LLMs incur heavy computation cost due to the quadratic complexity of self attention; besides, they require tremendous GPU memory to store the KV cache of $x_{\leq n+1}$ for faster decoding []. Both the costs in computation and memory significantly expand when the context length increases.

Activation Beacon employs a new special token, namely beacon token $\langle\text{b}\rangle$, and condenses the raw context $X$ into beacon tokens’ activations $\boldsymbol{\Psi}$ (i.e. their keys and values at every layer). The next-token prediction is converted to condition on the compressed context instead of the plain one. Given $|\boldsymbol{\Psi}|<|X|$, both the computation cost and the KV cache size are reduced. Additionally, the LLM is enabled to handle context longer than its window size based on the compressed representations. We tailor the compression mechanism and the learning method of Activation Beacon towards achieving effective, efficient, and flexible compression, which will be elaborated in the following.

### 3.1 Compression Mechanism

Overview.
We propose to progressively compress each fine-grained units of long contexs.
Specifically, given the input context $X$ whose length may exceed the LLM’s context window $N$, it is first partitioned into chunks of the same size $w$ (e.g., 1024):

$$ $[x_{1},\dots,x_{n}]\xrightarrow{\mathrm{Partition}}[X_{1},\dots X_{\lceil n/w \rceil}],~{}X_{i}=[x_{(i-1)w+1},\dots,x_{iw}]=[x^{i}_{1},\dots,x^{i}_{w}].$ (1) $$

Next, for each chunk $X_{i}$, we determine a compression ratio $\alpha_{i}$ ($w$ is evenly divisible by $\alpha_{i}$). The chunk is further split into fine-grained units of size $\alpha$. Then a group of $k_{i}=w/\alpha_{i}$ beacon tokens, $B_{i}=[\langle\text{b}\rangle^{i}_{1},\dots,\langle\text{b}\rangle^{i}_{k_{i}}]$, are interleaved with these units. In other words, one beacon token is dispatched to the end of every unit:

$$ $X_{i}\xrightarrow{\mathrm{Interleave}~{}B_{i}}X_{i}^{\prime}=[x^{i}_{1},\dots, x^{i}_{\alpha_{i}},\langle\text{b}\rangle^{i}_{1},~{}\dots~{},x^{i}_{w-\alpha_ {i}+1},\dots,x^{i}_{w},\langle\text{b}\rangle^{i}_{k_{i}}].$ (2) $$

The LLM encodes these chunks one by one,
compressing the contextual information of each chunk into the corresponding beacon tokens’ activations during self attention.
After encoding $X_{i}^{\prime}$, we discard activations of all the raw tokens $X_{i}$, while we accumulate the activations of the beacon tokens $B_{i}$.
When encoding the next chunk $X_{i+1}^{\prime}$, the LLM directly conditions on the accumulated beacon activations as a proxy to the raw context $X_{\leq i}$.

This progressive workflow benefits both compression quality and running efficiency. On one hand, it enables thorough distillation of complex information within long contexts and allows for the compression of inputs that exceed the LLM’s context window. On the other hand, by caching and reusing beacon tokens’ activations, it avoids redudant computation and allows for incrementally update of the compression results in multi-turn interactions.

Figure: Figure 2: Activation Beacon performs compression during self attention while reusing all other modules of the LLM. Forward ➀: encode and compress the first chunk. Forward ➁: encode and compress the second chunk conditioned on activations of preceding beacon tokens.
Refer to caption: x2.png

Encoding and Compression.
As shown in Figure 2, Activation Beacon reuses all modules of the LLM except imposing a slight modification on self attention.
Without loss of generality, for the $i$-th chunk $X_{i}^{\prime}$, the encoding process can be written as:

$$ $\small\mathrm{LLM}(\underbrace{\langle\text{b}\rangle_{1}^{i},\dots,\langle \text{b}\rangle^{i-1}_{k_{i-1}},}_{\text{beacon activations accumulated from } X^{\prime}_{<i}}\underbrace{x^{i}_{1},\dots,x^{i}_{\alpha_{i}},\langle\text{b} \rangle^{i}_{1},~{}\dots~{},x^{i}_{w-\alpha_{i}+1},\dots,x^{i}_{w},\langle \text{b}\rangle^{i}_{k_{i}}\vphantom{\langle\text{b}\rangle^{i-1}_{k_{i-1}}}}_ {\text{the current chunk }X^{\prime}_{i}}),$ (3) $$

where the input to the LLM is a mix of the activations accumulated from previous chunks and the tokens to be encoded within the current chunk.
Let $D$ denote the LLM’s hidden size, $\boldsymbol{H}\in\mathbb{R}^{(w+k_{i})\times D}$ denote input hidden states to self attention in an arbitrary layer of the LLM. We first slice out the hidden states of raw tokens and beacon tokens:

$$ $\displaystyle\mathbb{I}^{r}=\{j\mid x^{i}_{j}\neq\langle\text{b}\rangle\}, \quad\mathbb{I}^{b}=\{j\mid x^{i}_{j}=\langle\text{b}\rangle\};\quad \boldsymbol{H}^{r}=\boldsymbol{H}[\mathbb{I}^{r}],\quad\boldsymbol{H}^{b}= \boldsymbol{H}[\mathbb{I}^{b}].$ (4) $$

Then the hidden states are projected into queries, keys, and values:

$$ $\displaystyle\boldsymbol{Q}^{r}=\boldsymbol{W}_{Q}^{r}\boldsymbol{H}^{r},$ $\displaystyle\boldsymbol{K}^{r}=\boldsymbol{W}_{K}^{r}\boldsymbol{H}^{r},$ $\displaystyle\boldsymbol{V}^{r}=\boldsymbol{W}_{V}^{r}\boldsymbol{H}^{r},$ $\displaystyle\boldsymbol{Q}^{b}=\boldsymbol{W}^{b}_{Q}\boldsymbol{H}^{b},$ $\displaystyle\boldsymbol{K}^{b}=\boldsymbol{W}^{b}_{K}\boldsymbol{H}^{b},$ $\displaystyle\boldsymbol{V}^{b}=\boldsymbol{W}^{b}_{V}\boldsymbol{H}^{b},$ (5) $$

where $\boldsymbol{W}_{*}^{r}$ are the LLM’s original projection matrices and $\boldsymbol{W}_{*}^{b}$ are the newly introduced matrices to handle beacon tokens only. Afterwards, the query/key/value states of raw tokens and beacon tokens are scattered back to acquire $\boldsymbol{Q},\boldsymbol{K},\boldsymbol{V}\in\mathbb{R}^{(w+k_{i})\times D}$:

$$ $\boldsymbol{Q}[\mathbb{I}^{r}]=\boldsymbol{Q}^{r},~{}\boldsymbol{Q}[\mathbb{I} ^{b}]=\boldsymbol{Q}^{b};\quad\boldsymbol{K}[\mathbb{I}^{r}]=\boldsymbol{K}^{r },~{}\boldsymbol{K}[\mathbb{I}^{b}]=\boldsymbol{K}^{b};\quad\boldsymbol{V}[ \mathbb{I}^{r}]=\boldsymbol{V}^{r},~{}\boldsymbol{V}[\mathbb{I}^{b}]= \boldsymbol{V}^{b}.$ (6) $$

Finally, the standard self-attention is computed over the entire input:

$$ $\displaystyle\boldsymbol{A}=\mathrm{softmax}\left(\mathrm{mask}\left(\frac{ \boldsymbol{Q}\left\{\boldsymbol{K}^{ac};\boldsymbol{K}\right\}^{T}}{\sqrt{D}} \right)\right),\quad\boldsymbol{V}=\boldsymbol{A}\left\{\boldsymbol{V}^{ac}; \boldsymbol{V}\right\}.$ (7) $$

In the above equations, $\{\cdot~{};~{}\cdot\}$ denotes matrix concatenation. $\boldsymbol{K}^{ac},\boldsymbol{V}^{ac}\in\mathbb{R}^{m_{i-1}\times D}$ are the beacon tokens’ activations accumulated from previous chunks where $m_{i-1}=\sum_{j=1}^{i-1}{k_{j}}$, and $\mathrm{mask}$ denotes the causal attention mask. During self attention, all tokens are encoded by their relative positions ($[m_{i-1},\dots,m_{i}+w-1]$ for queries and $[0,\dots,m_{i}+w-1]$ for keys).
The value states $\boldsymbol{V}$, are further processed by other modules (e.g., output projection, MLP, and LayerNorm) before passing to the next layer.
After self attention, the keys and values of beacon tokens, i.e. $\boldsymbol{K}^{b}$ and $\boldsymbol{V}^{b}$, have distilled the contextual information of $X_{i}$. They are incrementally accumulated:

$$ $\boldsymbol{K}^{ac}=\{\boldsymbol{K}^{ac};\boldsymbol{K}^{b}\},\quad \boldsymbol{V}^{ac}=\{\boldsymbol{V}^{ac};\boldsymbol{V}^{b}\}.$ (8) $$

In our default setting, the beacon tokens are interleaved with raw tokens. This leads to a differentiated attention scope for each beacon token ($\langle\text{b}\rangle_{j}^{i}$ attends to one more interval than $\langle\text{b}\rangle_{j-1}^{i}$), contributing to the fine-grained compression of the context. We also explore the setting to dispatch all beacon tokens at the end of the chunk, which results in inferior compression quality ($\S$4.6).

Note that unlike ICAE [] and LLMLingua [], Activation Beacon unifies generation and compression operations within a single forward pass of the LLM. That is to say, the hidden states of the last input token $\boldsymbol{H}[\mathbb{R}^{r}[-1]]$ is directly used to decode the next token without resorting to another decoder model.

Figure: Figure 3: Comparison of the forward FLOPs of different models using full attention and Activation Beacon (the compression ratio is annotated in the brackets).
Refer to caption: x3.png

Efficiency Analysis.
Activation Beacon reduces the KV cache by $\alpha$ times where $\alpha$ is the average compression ratio and hence the memory cost. This is because it only needs to store the compressed activations of the preceding chunks instead of the raw activations.
In terms of computation, the situation is a bit more complex. Specifically, Activation Beacon significantly reduces the computation in self attention, because each token only needs to interact with local tokens within the chunk and preceding beacon tokens, which are approximately $\alpha$ times shorter than the raw context. However, it also triggers more computation to encode the inserted beacon tokens in other modules (e.g., MLP). Formally, given an LLM with a fixed number of layers, attention heads, and hidden size, let $s$ denote the input context length, $s^{pst}$ denote the cached context length, the forward FLOPs is:

$$ $\mathrm{FLOPs}=F^{Att}(s,s^{pst})+F^{Oth}(s),$ (9) $$

where $F^{Att}$ is the computation during self attention, and $F^{Oth}$ is the computation of other modules. For full-attention models, $s=n,s^{pst}=0$. For “beaconed” models, the FLOPs is:

$$ $\mathrm{FLOPs}^{bcn}=\sum_{i=1}^{\lceil\frac{n}{w}\rceil}F^{Att}\left(\frac{( \alpha+1)w}{\alpha},\frac{(i-1)w}{\alpha}\right)+F^{Oth}(n+\lceil\frac{n}{ \alpha}\rceil).$ (10) $$

Since the implementation of $F^{Att}$ and $F^{Oth}$ depends on the actual setting of the LLM (see Appendix B), we visualize the FLOPs curve of three different LLMs in Figure 3. It can be observed that Activation Beacon consistently saves computational costs across different model settings and scales. The extent of saving amplifies as the context length grows, finally achieving more than x4 reduction at 256K context. The specific implication on latency is studied in $\S$4.3.

### 3.2 Learning Method

Compression-Based Auto-Regression. Activation Beacon is learned to optimize the generation quality conditioned on the mixture of the compressed context and the local context. Formally, the compression-based next-token prediction loss is minimized:

$$ $\min\limits_{\boldsymbol{\Theta}^{b}}.\sum_{i=2}^{\lceil N/w\rceil}\sum_{j=1}^ {w}\Pr(x^{i}_{j}\mid\langle\text{b}\rangle^{1}_{1},\dots,\langle\text{b} \rangle^{i-1}_{k_{i-1}},x^{i}_{1},\dots x^{i}_{j-1};\boldsymbol{\Theta}, \boldsymbol{\Theta}^{b}).$ (11) $$

$\boldsymbol{\Theta}$ denotes the parameters of the LLM itself, which are fixed throughout the training process. $\boldsymbol{\Theta}^{b}$ includes the projection matrices for beacon tokens at each layer $\boldsymbol{W}_{Q}^{b}$, $\boldsymbol{W}_{K}^{b}$, $\boldsymbol{W}_{V}^{b}$, and the token embedding of beacon token $\boldsymbol{e}_{\langle\text{b}\rangle}$ (we use one shared embedding for all beacon tokens).
The training loss can be obtained from all tokens except the ones in the first chunk. Such a property leads to high sample efficiency that maximizes the use of training data. Note that we exclude the beacon tokens from the above loss (setting their labels to -100) because they are solely intended for compression.

No Stop Gradients. Recurrent memory methods [; ] stop the gradients back-propagation at a given chunk number to improve the training efficiency. This is because these methods depend on the final-layer outputs of preceding chunks to encode the current chunk, which results in deepened computation graph as more chunks are involved.
In contrast, Activation Beacon only depends on the previous-layer outputs of preceding chunks (the encoding of $X_{i}^{\prime}$ at layer $l$ only conditions on the results of $X_{i-1}^{\prime}$ at layer $l-1$), which is the same as any auto-regressive LLMs. Thus, the gradients can naturally flow through all chunks to optimize the compression effect over long contexts.

Chunk-Wise Random Compression Ratio. To teach the model to flexibly support diverse compression granularities, the compression ratio $\alpha_{i}$ for the $i$-th chunk is randomly sampled from $\{2,4,8,16,32\}$ during training. At inference, one can choose one compression ratio according to the specific efficiency requirement in downstream tasks and stick to it for all chunks.

## 4 Experiments

Our experiment mainly study Activation Beacon’s effectiveness ($\S$4.2), efficiency ($\S$4.3), and flexibility ($\S$4.4) in long context compression.
Besides, we explore Activation Beacon’s impact on short-context capabilities of the backbone LLM ($\S$4.5) and the effect of each technical design ($\S$4.6).

### 4.1 Settings

Implementation. Activation Beacon is applied to Llama-2-7B (chat)(^2^22We use Llama-2 because AutoCompressor and ICAE are based on it, both of which are important baselines.) and Qwen-2-7B (instruct). The chunk size $w$ is 1024 for Llama-2 and 2048 for Qwen-2.
FlashAttention-2 [] is used to speed up attention computation.
For all our experiments, we use Huggingface framework [] and one 8xA800 (80G) machine.

Training.
The training consists of two phases. In pre-training, we use 1B tokens sampled from RedPajama []. The eos token is appended to the end of every document. In fine-tuning, we leverage LongAlpaca [], BookSum [], and synthetic data from GPT-3.5 (details in Appendix A). All the training samples are shorter than 20K. The batch size is 8. The learning rate is 5e-5 for pre-training and 1e-5 for fine-tuning, with linear decay and no warmup. As introduced, the LLM’s original parameters are frozen throughout the training process.

Baselines. We compare Activation Beacon with the uncompressed baseline (denoted as Full) and the uncompressed baseline fine-tuned with the same training data (denoted as Full-FT). Besides, we include the following context compression methods that can tackle long context for comparison, including AutoCompressors [], ICAE [], LongLLMLingua [], and SnapKV []. The first two methods only support Llama-2. To guarantee fair comparison, we fine-tune their official checkpoints using the same training data.

### 4.2 Compression Effectiveness

**Table 1: Evaluation on LongBench []. Activation Beacon maintains comparable performance to the uncompressed baseline (Full-FT), outperforming other compression methods.**
| Model | Method | Length | Single-Doc | Multi-Doc | Summ. | Few-Shot | Code |
| --- | --- | --- | --- | --- | --- | --- | --- |
| Llama-2-7B | Full | 4K | 24.7 | 22.4 | 24.6 | 63.2 | 57.7 |
| Full-FT | 32K | 34.8 | 27.5 | 23.2 | 61.8 | 57.8 |  |
| AutoCompr. | 32K | 12.9 | 16.4 | 16.3 | 23.8 | 39.4 |  |
| ICAE | 32K | 19.5 | 19.2 | 19.5 | 24.8 | 27.8 |  |
| LongLLML. | 32K | 21.5 | 18.8 | 21.7 | 49.5 | 53.2 |  |
| SnapKV | 4K | 24.2 | 22.6 | 16.3 | 60.1 | 57.7 |  |
| Ours | 32K | 34.9 | 27.5 | 25.0 | 61.4 | 57.8 |  |
| Qwen-2-7B | Full | 32K | 38.8 | 37.5 | 26.7 | 70.1 | 60.3 |
| Full-FT | 32K | 41.0 | 40.6 | 26.8 | 68.5 | 66.1 |  |
| LongLLML. | 32K | 24.7 | 20.3 | 26.3 | 55.9 | 50.1 |  |
| SnapKV | 32K | 38.7 | 37.6 | 26.2 | 67.1 | 60.3 |  |
| Ours | 32K | 40.5 | 40.3 | 26.8 | 68.4 | 66.4 |  |

Figure: Figure 4: Evaluation on Needle-in-a-Haystack. Activation Beacon can accurately retrieves the needle most of the time, despite the context is far longer than its training data.
Refer to caption: x4.png

To verify the compression effectiveness of Activation Beacon, we evaluate it on LongBench [], which consists of a variety of long-context tasks with 32K maximum length, including question answering, summarization, few-shot learning, and code completion.
Since Llama-2 has a context window of 4K, we truncate the context longer than 4K from middle before inputting to it. For compression methods implemented on Llama-2, we set adaptive compression ratio, translating to x2 compression for 4K-8K contexts, x4 compression for 8K-16K contexts, and x8 compression for 16K-32K contexts.
For methods implemented on Qwen-2, we apply a uniform compression ratio of x4.
The results are reported in Table 1. We highligh two observations in the following.

Firstly, Activation Beacon achieves superior compression quality over other compression baselines across all tasks.
Concretely, it siginificantly outperforms ICAE and AutoCompressor, which verifies that several soft tokens are not enough to encapsulate the rich information within long contexts.
LongLLMLingua also lags far behind Activation Beacon because it need to delete too many tokens given a high compression ratio (e.g., x4, x8), which may destroy the coherence of the context and lose important information.
Despite SnapKV’s top performance among baselines, it cannot compress context longer than the backbone LLM’s window. This is because it estimates the token importance based on self attention, which becomes inaccurate once the context exceeds the window size, limiting its practical usage when compressing long contexts.

Secondly, Activation Beacon achieves comparable performance to the fine-tuned uncompressed baseline (Full-FT) even though Full-FT takes in the entire context without compression. This indicates that Activation Beacon is able to compress long contexts without evident information loss, which validates its high compression quality yielded from the progressive compression workflow.
Furthermore, Activation Beacon improves upon Llama-2 by a large margin despite their context window is the same, i.e. 4K. The gain is because Llama-2 (Full) directly uses the truncated 4K context, while Activation Beacon compresses the 32K context into 4K compact activations. This implies that Activation Beacon can effectively introduce useful information from Llama-2’s unseen context. Therefore, it can be viewed as an efficient approach for context extension.

We further evaluate Activation Beacon on Needle-in-a-Haystack (NIAH) following the official settings [] to investigate whether it will lose fine-grained information. The accuracy is estimated by ChatGPT (ranges from 1 to 10). For both Llama-2 and Qwen-2, we set adaptive compression ratio as introduced above. The results are shown in Figure 4. It can be observed that Activation Beacon precisely retrieves the needle most of the time. Note that Activation Beacon conducts query-independent compression, which means it has no prior knowledge of what to compress and what not. Hence, this remarkable performance again validates our tailored compression mechanism and learning method can preserve the fine-grained contextual information. Moreover, Activation Beacon is only trained on context shorter than 20K, while its compression capability can generalize to far longer contexts (e.g., 128K).

### 4.3 Compression Efficiency

**Table 2: Evaluation on Multi-Needle-in-a-Haystack where the questions are issued one-by-one in a multi-turn conversation setting. All compression methods use a x8 compression ratio. Activation Beacon consistently outperforms other compression baselines while enjoying lower latency, especially when the context lengthens and the turn number increases.**
| Model | Length | Method | 1-Turn | 2-Turn | 3-Turn |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Acc | Latency | Acc | Latency | Acc | Latency |  |  |  |
| Llama-2-7B | 32K | Full-FT | 9.75 | 1.336 | 9.45 | 1.532 | 9.10 | 1.726 |
| AutoCompr. | 1.60 | 2.135 | 1.50 | 2.561 | 1.50 | 2.994 |  |  |
| ICAE | 2.15 | 1.182 | 2.15 | 1.476 | 2.00 | 1.805 |  |  |
| LongLLML. | 2.05 | 2.813 | 2.00 | 5.062 | 2.00 | 7.034 |  |  |
| SnapKV | 1.00 | 0.859 | 1.00 | 1.656 | 1.00 | 2.199 |  |  |
| Ours | 9.75 | 1.153 | 9.40 | 1.356 | 9.05 | 1.638 |  |  |
| Qwen-2-7B | 128K | Full-FT | 9.75 | 4.399 | 9.50 | 5.254 | 9.20 | 6.153 |
| LongLLML. | 2.00 | 10.455 | 1.55 | 19.768 | 1.50 | 27.751 |  |  |
| SnapKV | 9.45 | 3.955 | 8.95 | 7.803 | 8.85 | 10.659 |  |  |
| Ours | 9.70 | 2.445 | 9.35 | 2.773 | 9.10 | 2.981 |  |  |

We evaluate the efficiency of Activation Beacon based on the Multi-Needle-in-a-Haystack task following NeedleBench []. Specifically, we fix the context length to 32K for Llama-2 and 128K for Qwen-2, and insert 3 different needles at different positions.
The task is organized in a multi-turn conversation setting, where the model is asked to retrieve one specific needle in each turn.
The experiment is repeated 20 times for each model with distinct needle positions.
In Table 2, we report the accuracy and the end-to-end latency of compression & generation (measured in seconds).

It can be observed that Activation Beacon enjoys lower latency than other compression baselines.
Notably, it is 1.8x faster than AutoCompressor because it does not have to re-encode the soft tokens from previous chunks.
It also leads to 9.3x and 3.6x acceleration upon LongLLMLingua and SnapKV given three turns, respectively. This is because both baselines are query-dependent while Activation Beacon is not, which eliminates the need to re-compute the compression results for different input questions.
Moreover, Activation Beacon demonstrates consistent speed-up over the Full-FT baseline, achieving 2x acceleration at 128K context length. This matches our estimation in Figure 3(b) as Activation Beacon (x8) saves half of the computation.
In the meanwhile, since the compression ratio is x8, it leads to 8x reduction of the KV cache.
Lastly, Activation Beacon always attains nearly-lossless generation quality against the uncompressed baseline, which is in line with previous observations.

### 4.4 Compression Flexibility

Figure: Figure 5: Evaluation on Needle-in-a-Haystack with various compression ratios based on Llama-2. Activation Beacon achieves top compression quality across all compression configurations.
Refer to caption: x5.png

Activation Beacon is learned to support various compression ratios during training. In Figure 5, we evaluate its compression quality under different compression ratios and context lengths. According to the figure, Activation Beacon maintains top accuracy across all compression ratios, outperforming most compression baselines by a large margin.
Though SnapKV performs on par with our method at 1K and 4K context length, it fails to compress inputs longer than the LLM’s window size, which may limit its practical usage. To summarize, Activation Beacon is a flexible solution to long context compression with the support of diverse compression ratios and various context lengths. Generally, we recommend to use x8 compression ratio as it preserves most information with high efficiency.

### 4.5 Short-Context Capabilities

**Table 3: Activation Beacon preserves the short-context capabilities of the backbone LLM.**
| Model | Method | MMLU | ARC-C | BoolQ | GSM8K |
| --- | --- | --- | --- | --- | --- |
| Llama-2-7B | Full | 47.5 | 48.5 | 86.2 | 9.2 |
| Ours | 46.6 | 48.4 | 86.5 | 9.3 |  |
| Qwen-2-7B | Full | 70.1 | 62.7 | 87.1 | 76.0 |
| Ours | 69.1 | 62.7 | 87.2 | 76.2 |  |

Since Activation Beacon interleaves beacon tokens with raw tokens and is primarily trained with long-context tasks, it is intriguing to examine whether the current recipe will impair the short-context capabilities of the backbone LLM.
In Table 3, we compare Activation Beacon with the original LLM (Full) on popular benchmarks, including MMLU [], ARC-Challenge [], BoolQ [], and GSM8K [].
We can observe that Activation Beacon leads to very little performance degradation on short-context tasks. In other words, the short-context capabilities are well preserved. We conjecture that the primary reason is the LLM’s original parameters are frozen throughout the training process.

### 4.6 Ablation Studies

**Table 4: The impact of different technical factors.**
| Method | Single-Doc |
| --- | --- |
| Default | 40.5 |
| w/o Fine-Grained Compression | 35.2 |
| w/o Chunk-Wise Random Ratio | 37.7 |
| w/o Pre-training | 34.9 |
| w/o Fine-tuning | 35.5 |

We study the impact of each technical factor, including the compression of fine-grained context units, the sampling strategy of compression ratio, and training stages. The experiments are based on Qwen-2-7B and Single-Doc QA task from LongBench (32K context with x4 compression ratio). The results are shown in Table 4.
Firstly, instead of splitting the chunk into fine-grained units and interleaving beacon tokens, we append all beacon tokens at the end of the chunk so that their attention scopes are the same. It can be observed that such operation results in significant information loss after compression, which justifies the effectiveness of our fine-grained compression mechanism.
Secondly, we replace the chunk-wise random compression ratio with the instance-wise one, which randomly selects one compression ratio for each training instance rather than each chunk. We can observe that the chunk-wise setting facilitates better learning of the compression functionality. Lastly, we remove either pre-training or fine-tuning. It can be observe that both stages are useful, and the combination of both leads to the optimal performance. This also implies that the compression quality of Activation Beacon can be further enhanced given more abundant and targeted training.

## 5 Conclusion

This paper introduces Activation Beacon, a plug-in for transformer-based LLMs to enable effective, efficient, and flexible compression of long contexts. Activation Beacon is featured for several critical innovations, including the progressive compression workflow to distill the context into a small set of activations, the compression-based auto-regression to optimize the model with high sample efficiency, and the random sampling of compression ratios to support various downstream scenarios. According to extensive experimental evaluations, Activation Beacon consistently outperforms existing context compression methods across various compression configurations in terms of both effectiveness and efficiency. It even maintains comparable performance to the uncompressed baseline, meanwhile achieving 2x acceleration and 8x KV cache reduction. Moreover, the short-context capabilities of the LLM is well preserved.

## Appendix A Training Data

In the pre-training phase, we use 1B tokens from RedPajama. We add an eos token to the end of each document. The documents longer than 20480 or shorter than 2400 are removed.

In the fine-tuning phase, we use three datasets. 1) LongAlpaca [], which contains long-context QA and summarization data; 2) BookSum [], which contains chapter-level summarization of books. 3) Synthesized QA dataset.
This dataset contains 16K long-context question answering instances (13K for books and 3K for papers). Specifically, we split a given long context (a paper or a book) into short segments (a chunk with less than 4096 tokens) using the SemanticTextSplitter(^3^33https://github.com/benbrandt/text-splitter). For each segment, we prompt GPT-3.5-turbo to generate 4 question-answer pairs based on the segment.
We then group continuous segments using Template A, where we can control the resulted context length by concatenating different number of segments.
The books are randomly sampled from Books3 corpus, and the papers are sampled from Arxiv, both coming from the Pile [].
All the above fine-tuning data are formatted in the manner of multi-turn conversations and limit the context length up to 20480.
To mitigate forgetting, we also include 5000 samples from the pre-training data.

## Appendix B FLOPs Calculation

Denote the input sequence length as $s$, the cached sequence length as $s^{pst}$, query head number as $h^{q}$, key/value head number as $h^{k}$, the hidden size $D$, head dimension as $d$, intermediate size $I$, and vocabulary size $V$.

$$ $\displaystyle F^{Att}$ $\displaystyle=F^{qkv}+F^{qk}+F^{softmax}+F^{av}+F^{out}$ $\displaystyle F^{qkv}$ $\displaystyle=2\times s\times D\times d\times h^{q}+2\times 2\times s\times D \times d\times h^{k}$ $\displaystyle F^{qk}$ $\displaystyle=2\times h^{q}\times s\times(s+s^{pst})\times d$ $\displaystyle F^{softmax}$ $\displaystyle=h^{q}\times(s+s^{pst})\times(s+s^{pst})$ $\displaystyle F^{av}$ $\displaystyle=2\times h^{q}\times s\times(s+s^{pst})\times d$ $\displaystyle F^{out}$ $\displaystyle=2\times s\times d\times h^{q}\times D$ (12) $$

$$ $\displaystyle F^{Oth}$ $\displaystyle=F^{up}+F^{gate}+F^{down}+F^{lm}$ $\displaystyle F^{up}$ $\displaystyle=2\times s\times D\times 2\times I$ $\displaystyle F^{gate}$ $\displaystyle=s\times\times I$ $\displaystyle F^{down}$ $\displaystyle=2\times s\times D\times I$ $\displaystyle F^{lm}$ $\displaystyle=2\times s\times D\times V$ (13) $$

## Peer Reviews

confidence: 5
contribution: 3
initial_rating: 5
presentation: 2
questions: - How are rotary embeddings managed for the beacon tokens? Although the LLM processes a fixed chunk at a time, the relative positions of the beacon tokens vary across chunks. How are positional embeddings applied in these cases?
- Additional parameters are added and fine-tuned for self-attention projections specific to the beacon tokens. What is the impact of these added parameters on VRAM usage and latency? If the cost is significant, could LoRA fine-tuning be effective for the proposed activation beacons approach?
- What portion of time is allocated to prefilling and decoding? While the proposed method reduces some recomputation, it may require customized attention masks or iterative context processing, which could lack efficient kernel implementation or result in extra kernel calls. Please provide a latency breakdown of prefilling and decoding for specific workloads (e.g., 32/128k context, 128 decoded tokens) and compare it with the flash attention full-context baseline.
- How does the proposed approach affect fine-tuning throughput? Please compare its performance with Full-FT.

I am open to adjusting my ratings if all concerns and questions are adequately addressed.
soundness: 2
strengths: - Compressing by chunks at each layer avoids the need for recomputation and addresses gradient back-propagation challenges present in some prior baselines that rely on recursive dependencies from final-layer outputs. This design enhances both training and inference efficiency.
- The chunking approach and the interleaved insertion of beacon tokens are straightforward and intuitive.
- Evaluations on various benchmarks indicate that the proposed approach generally outperforms the KV cache compression and “soft-prompt” compression baselines, achieving notable reductions in both inference time and memory usage.
- Training with randomly sampled compression ratios enables flexible compression ratios during testing.
summary: The paper introduces “Activation Beacon,” a plug-in module to conduct long-context compression for LLMs. The proposed approach progressively compresses the activations at each layer and can be trained in the conventional auto-regressive way of language modeling. The authors demonstrate the benefits of this approach through evaluations on various long-context tasks for compression quality and inference efficiency.
weaknesses: - In addition to LongBench and NIAH, it is essential to evaluate the proposed approach on newer, more challenging benchmarks, such as RULER [1].
- Some recent context compression baselines, including CEPE [2] and LLoCO [3], are not discussed in the paper and should be included for a more comprehensive discussion or comparison.

[1] Hsieh et al. RULER: What's the Real Context Size of Your Long-Context Language Models? COLM 2024.  
[2] Yen et al. Long-Context Language Modeling with Parallel Context Encoding. ACL 2024.  
[3] Tan et al. LLoCO: Learning Long Contexts Offline. EMNLP 2024.

---

confidence: 4
contribution: 3
initial_rating: 5
presentation: 2
questions: See weaknesses
soundness: 3
strengths: 1. Activation Beacon reduces inference time by 2x and KV cache memory costs by 8x compared to the uncompressed baseline.

2. The method supports adaptive compression ratios, allowing flexibility for different tasks and contexts.

3. The proposed model maintains short-context capabilities, preserving the performance of the original LLM.
summary: The paper introduces "Activation Beacon", a compression method designed to enhance long-context processing efficiency in LLMs. The approach compresses the activations of keys and values in transformer layers, avoiding bottlenecks associated with traditional soft prompt methods. Additionally, a progressive compression workflow compresses each context unit in chunks, allowing the model to handle longer contexts than the original LLM's window. Experimental results show Activation Beacon achieves significant memory and computation savings, with minimal loss in performance.
weaknesses: 1. The performance of this method may vary with model size. Current evaluations focus on medium-sized models, lacking validation on larger-scale models, leaving its effectiveness and applicability in very large models underexplored.

2. The added complexity of managing beacon tokens and compression ratios increases implementation overhead for end-users, particularly when adapting to different tasks. In addition to actual inference latency, specific memory usage data across implementations would help clarify practical resource requirements.

---

confidence: 3
contribution: 3
initial_rating: 8
presentation: 3
questions: overall, this paper is novel and idea is well presented. please add more techniques for comparison so that users can choose different method.
soundness: 3
strengths: - The paper presents an efficient method to compress long contexts, reducing memory usage by up to 8x and speeding up inference by 2x.
- Its progressive, fine-grained compression approach maintains high compression quality, allowing the model to handle longer inputs than its built-in context window.
-It supports flexible compression ratios, preserving model performance across various long-context tasks without degrading short-context capabilities.
summary: This paper compresses activations (keys and values) rather than using soft prompts, facilitating a progressive, fine-grained compression process. Specifically, it first partition input into small chunks, interleaving special beacon tokens that accumulate contextual activations.
weaknesses: - Lack of Comparison with KIVI: The paper does not provide a direct comparison with KIVI, a relevant compression method that could offer insights into the performance trade-offs.
- GPU Time Omission: The paper does not report GPU training or inference time, leaving uncertainty around the practical computational cost and efficiency of the proposed method.
- Scalability Concerns: The method requires 8 A800 GPUs to train a 7B parameter model, raising concerns about its scalability to larger models like 70B, where computational demands could become prohibitive.
- Limited Comparative Analysis: The paper would benefit from including more baseline methods, particularly compression-based approaches like KIVI, KV-layer shared compression methods such as CacheGen, and relative-position encoding strategies like LM-Infinite.
Additional References Needed: Incorporating comparisons with relevant works, such as LM-Infinite [1] for dynamic context management, CacheGen [2] for efficient context loading, and KIVI [3] for asymmetric quantization of KV caches, would strengthen the evaluation and highlight the advantages and limitations of the proposed approach.

[1] LM-Infinite: Simple On-the-Fly Length Generalization for Large Language Models
[2] CacheGen: Fast Context Loading for Language Model Applications
[3] KIVI: A Tuning-Free Asymmetric 2bit Quantization for KV Cache

---

confidence: 3
contribution: 3
initial_rating: 8
presentation: 3
questions: My main question is in the "regarding differences with previous works" above. I want to understand if the results are improved mainly from decreasing chunk size, or if there's another difference between soft tokens and beacon tokens that explains the difference.

Also, what window size do you use? From Table 1, your model has a context length of 32k. I'm guessing you use this window size, but I don't see it explicitly stated, and line 184 suggests that 1024 would be a common window size, so I'm not sure. Since LongBench has only a few examples above 32k, I'm guessing the window logic isn't really used much (unlike for Needle In a Haystack)
soundness: 3
strengths: - The paper focuses on an impactful area (long-context efficiency for LLMs).
- The paper provides a relatively simple idea that is well-explained. I view simplicity as a plus - if a simple idea can give strong accuracy improvements, it's far better than an unnecessarily complicated idea.
- The paper demonstrates strong results. Table 2 demonstrates strong accuracy at good latency on standard benchmarks for long context. Their method is competitive with full fine-tuning and better than baselines. Table 1 provides strong accuracy as well (though latency is missing).
- The figures do a good job of explaining what's going on. Figure 1 and Figure 2 give nice overviews of the method.
- The method is computationally efficient compared to fine-tuning. Their "pretraining" (starting from an already-pretrained model) only requires 1B tokens which is very few.
- The paper ablates design choices (Table 4).
- The paper is generally well written.
summary: The paper introduces a method called Activation Beacon for efficient long-context processing. The method adds learned "beacon" tokens at regular intervals in the input query. These tokens are expected to learn "summaries" of the text. At inference time, when processing long contexts, the beacon tokens are retained and the other context tokens are discarded. Thus, the beacon tokens essentially provide a summary of the context. The authors evaluate their method in comparison with a few other recent methods for efficient long context processing. Their method significantly improves results on LongBench and Multi-Needle-in-a-Haystack. The authors also provide ablations for various design choices.
weaknesses: - In Table 1, it's not obvious whether the latencies are comparable. The compression ratio isn't mentioned.
- line 368: why do you use adaptive compression for llama-2 and uniform compression for qwen?

My main perceived weaknesses are regarding differences with previous works, and understanding why this method is performing so well:
- line 135: "ICAE and AutoCompressor... segment the long context into chunks and compress each chunk. However, both of them compress the context into soft tokens" <- how are these soft tokens different than beacon tokens? (similarly, on line 373-374, you mention soft tokens being a drawback)
- line 137: "Their compression workflow also lacks fine-grained handling of the chunked inputs, resulting in inferior compression quality" <- it seems like all they would need to do to allow "fine-grained handling of the chunked inputs" is just choose a smaller chunk size, so that the soft tokens appear more frequently. Is that right?
- - If this is true, it seems like your main contribution is the insight that soft tokens should be distributed evenly through the context. Would doing this massively improve the accuracy of ICAE and AutoCompressor? It seems like this is the main discovery, but I'm left wondering if I'm missing some more fundamental difference.

[Minor]:
line 47: "it it" -> "it"
line 53: "alternamtive" -> alternative
line 371: "highligh" -> "highlight"
line 483: "scope" -> score
Table 2: give units of "latency"

Meta-review: This paper proposes a context compression method for transformer-based LLMs. The method progressively compresses the key and value activations for all layers into beacon tokens. The paper evaluates its benefits for quality and efficiency in long-context tasks. 

All reviewers agree that the paper has certain strengths, but they also raised a few questions and requests: (1) a discussion about the scalability, (2) a report on the inference speed, and (3) adding more benchmarks and related works. During the rebuttal, the authors properly answered those questions. Thus, the AC agrees that this paper is ready to be published at this conference.

Decision: Accept (Poster)