Long Context Compression with Activation Beacon

作者/机构: Peitian Zhang (北京人工智能研究院, 中国人民大学高瓴人工智能学院), Zheng Liu (北京人工智能研究院), Shitao Xiao (北京人工智能研究院), Ninglu Shao (北京人工智能研究院, 中国人民大学高瓴人工智能学院), Qiwei Ye (北京人工智能研究院), Zhicheng Dou (中国人民大学高瓴人工智能学院)

A1 主要贡献

本文针对大型语言模型（LLM）处理长上下文时面临的高计算和内存成本问题，提出了一种名为Activation Beacon的即插即用模块，旨在实现有效、高效且灵活的长上下文压缩。

核心问题与研究目标:
现有的大型语言模型在处理长上下文时，由于自注意力机制的二次方复杂度和存储整个序列KV缓存的需求，面临着巨大的计算和内存开销。现有的上下文压缩方法存在一些瓶颈：
1. 信息瓶颈：通常将上下文压缩为少数几个软提示（soft tokens），难以封装长上下文中的复杂信息。
2. 压缩粒度粗：采用“一次性”压缩整个上下文的方式，缺乏对细节信息的细粒度处理。
3. 效率低下：压缩后的软提示在生成前需要重新编码，导致训练和推理效率不高。
4. 灵活性差：通常被训练为固定的压缩率，难以根据下游任务定制。
5. 依赖问题：一些删除不重要token的方法依赖于输入问题来评估token重要性，限制了其在多轮对话场景中的效率。

创新点与主要贡献:
为了解决上述挑战，本文提出的Activation Beacon引入了以下几个关键技术设计：
1. 直接压缩激活值: 引入一种新的特殊令牌，称为信标令牌（beacon token）⟨b⟩。与依赖软提示不同，该方法直接将上下文信息提炼并压缩到每一层信标令牌的激活值（即键和值）中，其容量足以封装长上下文中的复杂信息。
2. 渐进式、细粒度的压缩工作流:
* 长上下文被划分为大小相等的块（chunks）。
* 每个块根据预设的压缩率α进一步分割为细粒度单元（units）。
* 在每个单元末尾插入一个信标令牌。
* LLM逐块进行编码，在自注意力过程中将块内信息提炼到信标令牌的激活值中。
* 编码后，原始令牌的激活值被丢弃，而信标令牌的激活值被累积并用于后续块的编码。
* 优势: 该流程能处理超过模型上下文窗口的输入，实现细粒度的信息压缩，并通过缓存和重用激活值避免了推理时的重编码开销，同时支持多轮场景下的增量更新。

1. 基于压缩的自回归训练: 模型通过在压缩上下文的条件下进行自回归生成任务进行训练，高效地利用了纯文本和指令数据来优化模型的压缩性能。
2. 训练中的随机压缩率: 在训练的每一步中，为每个块随机采样一个压缩率，使模型学会支持多种压缩配置，从而在下游任务中具有更高的灵活性。

核心成果:
Activation Beacon被应用于Llama-2和Qwen2模型。在各种长上下文任务（如文档理解、少样本学习和大海捞针测试，长度可达128K，远超20K的训练长度）的广泛评估中，Activation Beacon在各种压缩配置下均能保持与未压缩基线相当的性能，同时在推理时实现2倍的加速和8倍的KV缓存内存成本降低。此外，该方法能很好地保留LLM原有的短上下文处理能力。

A3 背景知识

长上下文处理的现状与挑战。最近，处理长上下文已成为现代LLM的一项基本能力（OpenAI, 2024; et al., 2024; Yang et al., 2024; DeepSeek-AI, 2024）。扩展上下文窗口的方法大致相同：通过外推和内插修改旋转位置嵌入（Su et al., 2021; Chen et al., 2023a; ntk, 2023; Peng et al., 2023; Ding et al., 2024），并在预训练和后训练阶段利用长依赖数据。尽管在效果上取得了显著进展，但LLM在效率方面面临重大挑战。由于Transformer的二次方复杂度，计算成本巨大；此外，LLM需要在GPU上保留整个序列的KV激活以实现更快的解码，导致巨大的内存成本。多项研究致力于降低这些成本，具体如下。

稀疏注意力。传统的稀疏注意力方法需要使用指定的稀疏模式从头开始重新训练模型（Zaheer et al., 2020; Beltagy et al., 2020）。然而，最近的大量研究发现，尽管LLM是密集训练的，但其注意力模式天然是稀疏的（Jiang et al., 2024a; Xiao et al., 2023; Han et al., 2023; Zhu et al., 2024）。他们还提出为每个头动态设置适当的稀疏模式，以便在减少计算量的同时，能够很大程度上保留注意力权重，从而获得与全注意力方法相媲美的性能。然而，这些方法需要将所有KV激活保留在芯片上以动态确定最优稀疏模式，这使得它们不适用于KV缓存的缩减。有些稀疏注意力方法直接丢弃中间的token（Han et al., 2023; Xiao et al., 2023）。尽管这些方法效率高且能生成无限流畅的文本，但它们无法记忆中间上下文的信息，导致在长上下文任务上的性能较差（Xiao et al., 2024）。

KV压缩。这类研究专注于压缩KV激活，以减少注意力计算和缓存大小。由于KV激活是按层、按头、按token、按通道的浮点数，它们可以从所有五个维度（包括数值维度）进行缩减。例如，CLA（Brandon et al., 2024）在多层之间共享KV缓存；GQA（Ainslie et al., 2023）将多个键/值头压缩成一个；MLA（DeepSeek-AI, 2024）将通道压缩成更少且更紧凑的通道；KIVI（Zirui Liu et al., 2023）对激活中的数值进行量化。序列级压缩（也称为上下文压缩），即Activation Beacon所属的类别，将在下文介绍。它与其他维度的压缩是正交的，不同维度压缩的互补效应可留待未来工作。此外，一些最近的研究设计了高效的KV缓存卸载和传输策略（Liu et al., 2023; Xiao et al., 2024）。它们也可以与KV压缩技术联合使用，以实现更高效的长上下文生成。

上下文压缩。这类方法旨在将原始上下文压缩成更短但更紧凑的表示。现有研究通常针对压缩短上下文（小于1K）进行定制，这在长上下文压缩方面往往不是最优的。具体来说，Gisting（Mu et al., 2023）将用户指令一次性压缩成要点激活。因此，它无法处理比骨干LLM窗口更长的上下文。CCM（Kim et al., 2024）将Gisting扩展到在线聊天中的对话压缩，但不能用于如长文档理解等通用长上下文任务。ICAE（Ge et al., 2024）和AutoCompressor（Chevalier et al., 2023）通过将长上下文分块并压缩每个块来缓解这个问题，以便压缩比骨干LLM窗口更长的上下文。CEPE（Yen et al., 2024）采用了类似的工作流程，同时引入一个独立的编码器来压缩上下文，并通过交叉注意力模块利用压缩结果。然而，这些方法将上下文压缩成软令牌，这是封装长上下文中复杂信息的主要瓶颈。它们的压缩工作流程也缺乏对分块输入的细粒度处理，导致压缩质量较差。此外，这些方法必须执行重新编码或采用额外的交叉注意力机制来利用压缩的软令牌，这引入了额外的开销。最后，由于软令牌的数量是预定义的，很难为下游任务灵活地分配压缩率。另一类方法（Jiang et al., 2023b; Li et al., 2024b）提出删除不重要的令牌来实现压缩。然而，它们依赖于输入问题来准确估计令牌的重要性，导致在现实世界的多轮场景中效率低下。与现有方法相比，Activation Beacon能够实现更有效、高效和灵活的压缩。基于上下文压缩技术，出现了一些创新框架，如LLoCO（Tan et al., 2024）。它建立在压缩器和解码器之上，上下文被离线压缩并卸载到检索系统中。然后，解码器根据检索到的压缩结果高效地响应用户输入。这两个模块都是通过领域内微调学习的。我们的工作旨在改进压缩器本身，因此与这些框架研究是正交的。

A2 方法细节

大型语言模型以下一个词元预测的形式完成任意任务。形式上，给定上下文$X = [x_1, \dots, x_n]$，LLM根据所有前面的词元及其训练良好的参数生成下一个词元：$Pr(x_{n+1} | x_1, \dots, x_n; \Theta)$。基于Transformer的LLM由于自注意力的二次方复杂度而产生高昂的计算成本；此外，它们需要巨大的GPU内存来存储$x_{\le n+1}$的KV缓存以实现更快的解码（Zhang et al., 2023）。计算和内存的成本都随着上下文长度的增加而显著扩大。

Activation Beacon采用一种新的特殊词元，即信标词元⟨b⟩，并将原始上下文$X$压缩为信标词元的激活值$\Psi$（即它们在每一层的键和值）。下一个词元的预测转为基于压缩后的上下文而非原始文本。由于$|\Psi| < |X|$，计算成本和KV缓存大小都得以降低。此外，LLM能够基于压缩表示处理比其窗口大小更长的上下文。我们为Activation Beacon量身定制了压缩机制和学习方法，以实现有效、高效和灵活的压缩，具体将在下文详细阐述。

3.1 压缩机制

渐进式细粒度压缩流程。我们提出渐进式地压缩长上下文的每个细粒度单元。具体而言，给定长度可能超过LLM上下文窗口$N$的输入上下文$X$，首先将其划分为大小为$w$（例如，1024）的等大小块：

$$[x_{1}, \ldots, x_{n}] \xrightarrow{\text { Partition }}\left[X_{1}, \ldots X_{\lceil n / w\rceil}\right], X_{i}=\left[x_{(i-1) w+1}, \ldots, x_{i w}\right]^{1}=\left[x_{1}^{i}, \ldots, x_{w}^{i}\right]$$

接着，对于每个块$X_i$，我们确定一个压缩率$\alpha_i$（$w$可被$\alpha_i$整除）。该块被进一步分割成大小为$\alpha_i$的细粒度单元。然后，一组$k_i = w/\alpha_i$个信标词元，$B_i = [\langle b \rangle_1^i, \dots, \langle b \rangle_{k_i}^i]$，被交错地插入到这些单元中。换句话说，每个单元的末尾都分配一个信标词元：

$$X_{i} \xrightarrow{\text { Interleave } B_{i}} X_{i}^{\prime}=\left[x_{1}^{i}, \ldots, x_{\alpha_{i}}^{i},\langle\mathrm{b}\rangle_{1}^{i}, \ldots, x_{w-\alpha_{i}+1}^{i}, \ldots, x_{w}^{i},\langle\mathrm{b}\rangle_{k_{i}}^{i}\right].$$

LLM逐一编码这些块，在自注意力过程中将每个块的上下文信息压缩到相应的信标词元的激活值中。在编码完$X'i$后，我们丢弃所有原始词元$X_i$的激活值，同时累积信标词元$B_i$的激活值。在编码下一个块$X'$的代理。}$时，LLM直接以累积的信标激活值为条件，作为原始上下文$X_{\le i

渐进式工作流的优势。这种渐进式工作流对压缩质量和运行效率都有益。一方面，它能够彻底地提炼长上下文中的复杂信息，并允许压缩超过LLM上下文窗口的输入。另一方面，通过缓存和重用信标词元的激活值，它避免了冗余计算，并允许在多轮交互中增量更新压缩结果。

编码与压缩的具体实现。如图2所示，Activation Beacon复用了LLM的所有模块，仅对自注意力机制进行了轻微修改。不失一般性，对于第$i$个块$X'$，编码过程可以写为：

$$\text{LLM}(\underbrace{\langle\text{b}\rangle_1^i, \dots, \langle\text{b}\rangle_{k_{i-1}}^{i-1}}_{\text{beacon activations accumulated from } X'_{<i}}, \underbrace{x_1^i, \dots, x_{\alpha_i}^i, \langle\text{b}\rangle_1^i, \dots, x_{w-\alpha_i+1}^i, \dots, x_w^i, \langle\text{b}\rangle_{k_i}^i}_{\text{the current chunk } X'_i})$$ <p>其中，输入到LLM的是从先前块累积的激活值和当前块内待编码的词元的混合体。设$D$为LLM的隐藏层大小，$H \in R^{(w+k_i) \times D}$表示LLM任意层中自注意力的输入隐藏状态。我们首先切分出原始词元和信标词元的隐藏状态：

$$\mathbb{I}^r = \{j \mid x_j^i \neq \langle \text{b} \rangle \}, \quad \mathbb{I}^b = \{j \mid x_j^i = \langle \text{b} \rangle \}; \quad H^r = H[\mathbb{I}^r], \quad H^b = H[\mathbb{I}^b].$$

然后，这些隐藏状态被投影成查询、键和值：

$$\begin{aligned} \begin{aligned} Q^r = W_Q^r H^r, \quad && K^r = W_K^r H^r, \quad && V^r = W_V^r H^r, \\ Q^b = W_Q^b H^b, \quad && K^b = W_K^b H^b, \quad && V^b = W_V^b H^b, \end{aligned} \end{aligned}$$

其中，$W^r_$是LLM原始的投影矩阵，$W^b_$是新引入的仅用于处理信标词元的矩阵。之后，原始词元和信标词元的查询/键/值状态被散布回去，以获得$Q, K, V \in 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}. $$

最后，在整个输入上计算标准的自注意力：

$$\boldsymbol{A}=\operatorname{softmax}\left(\operatorname{mask}\left(\frac{\boldsymbol{Q}\left\{\boldsymbol{K}^{a c} ; \boldsymbol{K}\right\}^{T}}{\sqrt{D}}\right)\right), \quad \boldsymbol{V}=\boldsymbol{A}\left\{\boldsymbol{V}^{a c} ; \boldsymbol{V}\right\} .$$

在上述方程中，{· ; ·}表示矩阵拼接。$K^{ac}, V^{ac} \in R^{m_{i-1} \times D}$是从先前块累积的信标词元激活值，其中$m_{i-1} = \sum_{j=1}^{i-1} k_j$，mask表示因果注意力掩码。在自注意力过程中，所有词元都按其相对位置进行编码（查询的相对位置为$[m_{i-1}, \dots, m_i + w - 1]$，键的相对位置为$[0, \dots, m_i + w - 1]$）。值状态$V$在传递到下一层之前，还会经过其他模块（如输出投影、MLP和LayerNorm）的处理。自注意力之后，信标词元的键和值，即$K^b$和$V^b$，已经提炼了$X_i$的上下文信息。它们被增量累积：

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

信标令牌交错布局的优势。在我们的默认设置中，信标词元与原始词元交错排列。这导致每个信标词元的注意力范围都不同（$\langle b \rangle_{ij}$比$\langle b \rangle_{ij-1}$多关注一个区间），有助于上下文的细粒度压缩。我们也探索了将所有信标词元分派到块末尾的设置，这导致了较差的压缩质量（§4.6）。

统一的生成与压缩前向传播。与ICAE (Ge et al., 2024) 和LLMLingua (Jiang et al., 2023b)不同，Activation Beacon在LLM的单次前向传播中统一了生成和压缩操作。也就是说，最后一个输入词元的隐藏状态$H[I^r[-1]]$被直接用于解码下一个词元，而无需借助另一个解码器模型。

效率分析。Activation Beacon通过平均压缩率$\alpha$将KV缓存减少了$\alpha$倍，从而降低了内存成本。这是因为它只需要存储先前块的压缩激活值，而不是原始激活值。在计算方面，情况稍微复杂一些。具体来说，Activation Beacon显著减少了自注意力中的计算量，因为每个词元只需要与块内的局部词元和先前的信标词元进行交互，这些词元比原始上下文短约$\alpha$倍。然而，它也触发了更多用于编码插入的信标词元的计算（例如，在MLP中）。形式上，给定一个具有固定层数、注意力头数和隐藏大小的LLM，设$s$为输入上下文长度，$s^{pst}$为缓存的上下文长度，前向FLOPs为：

$$ \mathrm{FLOPs}=F^{A t t}\left(s, s^{p s t}\right)+F^{O t h}(s), $$

其中，$F^{Att}$是自注意力期间的计算量，$F^{Oth}$是其他模块的计算量。对于全注意力模型，$s = n, s^{pst} = 0$。对于“信标化”模型，FLOPs为：

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

由于$F^{Att}$和$F^{Oth}$的实现取决于LLM的实际设置（见附录B），我们在图3中可视化了三种不同LLM的FLOPs曲线。可以观察到，Activation Beacon在不同的模型设置和规模下都能持续节省计算成本。节省的程度随着上下文长度的增长而放大，最终在256K上下文中实现了超过4倍的缩减。对延迟的具体影响在§4.3中进行了研究。

3.2 学习方法

基于压缩的自回归学习目标。Activation Beacon通过优化在压缩上下文和局部上下文混合条件下的生成质量来进行学习。形式上，最小化基于压缩的下一个词元预测损失：

$$\min_{\boldsymbol{\Theta}^b}. \sum_{i=2}^{\lceil N/w \rceil} \sum_{j=1}^{w} \operatorname{Pr}(x_j^i \mid \langle \mathrm{b} \rangle_1^1, \dots, \langle \mathrm{b} \rangle_{k_{i-1}}^{i-1}, x_1^i, \dots x_{j-1}^i; \boldsymbol{\Theta}, \boldsymbol{\Theta}^b).$$

$\Theta$表示LLM自身的参数，在整个训练过程中保持固定。$\Theta^b$包括每层信标词元的投影矩阵$W_Q^b, W_K^b, W_V^b$，以及信标词元的词元嵌入$e_{\langle b \rangle}$（我们为所有信标词元使用一个共享嵌入）。训练损失可以从除第一个块中的词元外的所有词元中获得。这种特性带来了高样本效率，最大限度地利用了训练数据。请注意，我们从上述损失中排除了信标词元（将其标签设置为-100），因为它们仅用于压缩。

无梯度截断的端到端优化。循环记忆方法（Chevalier et al., 2023; Bulatov et al., 2023）在给定的块数处停止梯度反向传播以提高训练效率。这是因为这些方法依赖于先前块的最后一层输出来编码当前块，这导致计算图随着更多块的参与而加深。相比之下，Activation Beacon仅依赖于先前块的前一层输出（在层$l$对$X'i$的编码仅以在层$l-1$对$X'$的结果为条件），这与任何自回归LLM相同。因此，梯度可以自然地流经所有块，以优化长上下文的压缩效果。

逐块随机压缩率以增强灵活性。为了让模型能够灵活支持不同的压缩粒度，在训练期间，第$i$个块的压缩率$\alpha_i$从{2, 4, 8, 16, 32}中随机采样。在推理时，可以根据下游任务的具体效率要求选择一个压缩率，并对所有块保持不变。

A4 实验环境与结果

实验环境

• 模型架构: Llama-2-7B (chat) 和 Qwen-2-7B (instruct)。在附录中还对Qwen-2.5-7B和Qwen-2.5-14B进行了实验。
• 硬件配置: 一台配备8块A800 (80G) GPU的服务器。

• 软件配置:

  • 代码实现基于Huggingface框架 (Wolf et al., 2020)。
  • 使用FlashAttention-2 (Dao, 2023) 加速注意力计算。
  • 实现细节：Llama-2的块大小w为1024，Qwen-2的块大小w为2048。

• 数据集:

  • 预训练阶段: 使用从RedPajama (Computer, 2023) 中采样的10亿个token。
  • 微调阶段: 使用LongAlpaca (Chen et al., 2023b), BookSum (Krysci´nski et al., 2022) 以及使用GPT-3.5合成的数据集（详见附录A）。所有训练样本长度均小于20K。

• 训练配置:

  • 批量大小为8。
  • 预训练学习率为5e-5，微调学习率为1e-5，均采用线性衰减且无预热。
  • 训练过程中冻结LLM的原始参数。

• 基线模型:

  • 未压缩基线 (Full)。
  • 使用相同数据微调的未压缩基线 (Full-FT)。
  • 长上下文压缩方法：AutoCompressors (Chevalier et al., 2023), ICAE (Ge et al., 2024), LongLLMLingua (Jiang et al., 2023b), 和 SnapKV (Li et al., 2024b)。

实验结果

压缩有效性评估

• 实验内容: 在LongBench基准测试上进行评估，该基准包含多种长上下文任务，最大长度为32K。对Llama-2模型，根据上下文长度设置自适应压缩率（4K-8K为x2，8K-16K为x4，16K-32K为x8）；对Qwen-2模型，统一使用x4压缩率。

• 实验结果 (Table 1):

  1. 优于其他压缩方法: Activation Beacon在所有任务上的压缩质量均优于其他基线。它显著超过了ICAE和AutoCompressor，证明了软令牌在封装长上下文信息方面的局限性。它也优于LongLLMLingua，因为后者在高压缩率下需要删除过多token，可能破坏上下文连贯性。SnapKV虽然在基线中表现最好，但无法压缩超过模型上下文窗口的文本。
  2. 与未压缩基线相当: Activation Beacon的性能与经过微调的未压缩基线（Full-FT）相当，表明其能够在不造成明显信息损失的情况下压缩长上下文。
  3. 有效扩展上下文: 对于Llama-2，Activation Beacon（上下文窗口4K）通过压缩32K上下文，性能远超直接使用截断的4K上下文的原始Llama-2模型，证明了其作为上下文扩展方法的有效性。

• 实验内容: 在大海捞针（Needle-in-a-Haystack, NIAH）任务上评估模型是否会丢失细粒度信息，上下文长度最长可达128K。

• 实验结果 (Figure 4): Activation Beacon在绝大多数情况下都能精确地检索到“针”，即使上下文长度远超其20K的训练长度。这再次验证了该方法在保留细粒度信息和泛化能力方面的出色表现。

压缩效率评估

• 实验内容: 在多针大海捞针（Multi-Needle-in-a-Haystack）任务上评估端到端延迟，该任务设置为多轮对话形式。Llama-2上下文长度为32K，Qwen-2为128K。
• 实验结果 (Table 2):
  1. 延迟更低: Activation Beacon的延迟低于所有压缩基线。比AutoCompressor快1.8倍（无需重编码软令牌），比LongLLMLingua和SnapKV分别快9.3倍和3.6倍（因其是查询无关的压缩）。
  2. 显著加速: 相较于Full-FT基线，在128K上下文长度下实现了2倍的加速，同时KV缓存减少了8倍，与理论分析（图3b）相符。
  3. 性能无损: 在加速的同时，生成质量与未压缩基线几乎无损。

压缩灵活性评估

• 实验内容: 在NIAH任务上，评估不同压缩率（x2到x16）和上下文长度（1K, 4K, 32K）下的压缩质量。
• 实验结果 (Figure 5): Activation Beacon在所有压缩率下都保持了最高的准确率，大幅领先于其他基线。这证明了其作为一种灵活长上下文压缩方案的有效性，能够支持多样化的压缩率和上下文长度。通常推荐使用x8压缩率，因为它在保持大部分信息的同时具有高效率。

短上下文能力评估

• 实验内容: 在MMLU, ARC-C, BoolQ, GSM8K等流行的短上下文基准上进行评估。
• 实验结果 (Table 3): Activation Beacon在短上下文任务上仅导致了非常轻微的性能下降，表明其短上下文能力得到了很好的保留。这主要归功于在训练过程中冻结了LLM的原始参数。

消融研究

• 实验内容: 基于Qwen-2-7B在LongBench的Single-Doc QA任务上，研究不同技术因素的影响。
• 实验结果 (Table 4):
  1. 细粒度压缩: 若不进行细粒度压缩（将所有信标令牌附加在块末尾），会导致显著的信息损失。
  2. 逐块随机压缩率: 采用逐块随机压缩率的策略比逐实例随机压缩率能更好地学习压缩功能。
  3. 训练阶段: 预训练和微调两个阶段都至关重要，二者结合才能达到最佳性能，也暗示了通过更丰富和有针对性的训练可以进一步提升压缩质量。

A5 结论

本文介绍了Activation Beacon，一个为基于Transformer的LLM设计的即插即用模块，旨在实现有效、高效且灵活的长上下文压缩。Activation Beacon的特点在于几个关键创新，包括将上下文提炼为一小组激活值的渐进式和细粒度压缩工作流，以高样本效率优化模型的基于压缩的自回归方法，以及支持各种下游场景的压缩率随机采样策略。根据广泛的实验评估，Activation Beacon在各种压缩配置下始终优于现有的上下文压缩方法。它甚至能保持与未压缩基线相当的性能，同时实现2倍的加速和8倍的KV缓存减少。此外，LLM的短上下文能力也得到了很好的保留。

A6 附录

A 训练数据

预训练数据。在预训练阶段，我们使用了来自RedPajama的10亿个token。我们在每个文档的末尾添加了一个eos token。长度超过20480或短于2400的文档被移除。

微调数据。在微调阶段，我们使用了三个数据集。1) LongAlpaca (Chen et al., 2023b)，其中包含长上下文问答和摘要数据；2) BookSum (Krysci´nski et al., 2022)，其中包含书籍的章节级摘要。3) 合成问答数据集。该数据集包含16K个长上下文问答实例（13K来自书籍，3K来自论文）。具体来说，我们使用SemanticTextSplitter将给定的长上下文（一篇论文或一本书）分割成短片段（长度小于4096个token的块）。对于每个片段，我们提示GPT-3.5-turbo根据该片段生成4个问答对。然后我们使用模板A.1将连续的片段分组，通过拼接不同数量的片段来控制最终的上下文长度。书籍是从Books3语料库中随机抽样的，论文是从Arxiv中抽样的，两者都来自Pile (Gao et al., 2020)。所有上述微调数据都以多轮对话的方式格式化，并将上下文长度限制在20480以内。为了减轻遗忘，我们还包含了来自预训练数据的5000个样本。

B FLOPs计算

FLOPs计算公式。记输入序列长度为$s$，缓存序列长度为$s^{pst}$，查询头数为$h^q$，键/值头数为$h^k$，隐藏层大小为$D$，头维度为$d$，中间层大小为$I$，词汇表大小为$V$。

$$\begin{aligned} \begin{aligned} F^{Att} & =F^{q k v}+F^{q k}+F^{s o f t m a x}+F^{a v}+F^{o u t} \\ F^{q k v} & =2 \times s \times D \times d \times h^{q}+2 \times 2 \times s \times D \times d \times h^{k} \\ F^{q k} & =2 \times h^{q} \times s \times\left(s+s^{p s t}\right) \times d \\ F^{s o f t m a x} & =h^{q} \times\left(s+s^{p s t}\right) \times\left(s+s^{p s t}\right) \\ F^{a v} & =2 \times h^{q} \times s \times\left(s+s^{p s t}\right) \times d \\ F^{o u t} & =2 \times s \times d \times h^{q} \times D \end{aligned} \end{aligned}$$ $$\begin{aligned} \begin{aligned} F^{Oth} &= F^{up} + F^{gate} + F^{down} + F^{lm} \\ F^{up} &= 2 \times s \times D \times 2 \times I \\ F^{gate} &= s \times \times I \\ F^{down} &= 2 \times s \times D \times I \\ F^{lm} &= 2 \times s \times D \times V \end{aligned} \end{aligned}$$

C 额外的效率分析

推理效率的分解分析。我们通过分解预填充（pre-filling）和解码（decoding）的延迟来进一步评估Activation Beacon的效率。具体来说，我们将上下文长度设置为32K和128K，并强制模型精确生成128个token。我们统一使用x8的压缩率。同时报告了整个生成过程中的峰值GPU内存。

根据表5，我们的方法加速了预填充和解码过程，且加速程度随着上下文变长而放大。同时，我们的方法在加速解码方面效果更好，因为它直接以信标令牌的激活值为条件，这些激活值比基线使用的原始激活值短8倍。此外，Activation Beacon显著降低了峰值GPU内存使用量，使得在消费级GPU上高效处理长上下文成为可能。

训练效率分析。最后，我们还比较了Activation Beacon与标准微调的训练效率。结果如表6所示。实验结果表明，Activation Beacon实现了与Full-FT基线相当的训练速度，同时显著降低了内存成本。

D RULER评估

在RULER基准上的评估。除了LongBench和NIAH，我们还在RULER (Hsieh et al., 2024)上评估了Activation Beacon的性能，这是一个具有挑战性的长上下文基准，包含13个合成任务。具体来说，有6个任务将常规的NIAH扩展为多键、多值和多查询的变体；1个任务（变量跟踪）旨在通过跟踪变量赋值来检验多跳推理能力；2个任务（常见/频繁词提取）旨在通过计算词频来检验聚合能力；还有2个任务在向真实段落中插入噪声上下文的同时执行问答。我们的评估基于128K上下文长度和x4的压缩率。结果如表7所示。

通过合成数据微调提升性能。可以观察到，我们最初的Activation Beacon模型（表示为“Ours”）在QA任务上保持了有竞争力的性能，但在推理/聚合任务（VT和CWE）上远远落后于全注意力基线。在Full-FT上也可以观察到类似的下降。造成这种劣势的一个可能原因是，我们当前的微调方案仅使用单跳QA数据（如附录A所述），这并未教会模型执行复杂的推理或计数，导致在这些任务上性能较差。然而，这个问题应该可以通过调整训练数据的构成来缓解。我们在训练数据中加入了200个合成样本（VT 100个，CWE 100个），最大上下文长度为20K，并对模型进行微调。新模型，表示为“Ours+Synthetic FT”，在VT和CWE上都取得了显著的提升，同时在其他任务上保持了其有竞争力的性能。这一观察进一步验证了我们的Activation Beacon的有效性：它可以在给定有限的训练数据的情况下快速学习所需的压缩能力。

E 在更大模型上的实验

在更大模型上的实验。我们将Activation Beacon应用于Qwen-2.5-7B和Qwen-2.5-14B，以考察其在更大模型上的性能。LongBench的评估结果（统一使用x4压缩率）如表8所示。可以观察到，Activation Beacon保持了其有效性，与全注意力微调基线保持了相当的性能。

更大模型的效率分析。此外，我们分别在表5和表6中报告了14B模型的训练和推理效率。Activation Beacon在训练期间显著降低了内存成本，使得能够高速高效地训练大模型。此外，它稳定地加速了推理过程，同时比全注意力基线在内存效率上高得多。