SeDeM：面向长上下文问答的隐状态记忆选择性解压缩

# SeDeM: 长上下文问答中隐藏状态记忆的选择性解压缩
来源：https://arxiv.org/html/2608.00311
Jason CongYizhou Sun
加利福尼亚大学洛杉矶分校，美国
\{maryamhgf, cong, yzsun\}@cs.ucla.edu

###### 摘要

使用大型语言模型（LLM）进行长上下文推理成本高昂：预填充阶段的自注意力随序列长度呈二次方增长，键值（KV）缓存随处理的令牌数量增长。更大的上下文窗口也不能确保可靠地利用证据。上下文压缩降低了这一成本，但许多软压缩方法使用LLM作为压缩器，并依赖紧凑的记忆令牌来同时保存信息和为解码器提供条件。我们提出SeDeM，一种将紧凑记忆存储与解码器条件化解耦的选择性解压缩框架。LLM从选定的中间Transformer层提取隐藏状态，轻量级压缩器将它们存储为记忆块，查询条件选择器选择相关块，解压缩器仅将选定的块扩展为与中间解码器层兼容的隐藏状态。因此，解码器既避免了全上下文处理，也避免了直接从高度压缩的记忆槽生成。在四个长上下文QA基准上，SeDeM在1B和3B同主干设置中均取得比评估的压缩基线更高的QA分数，并且使用3B主干时在三个数据集上超过了全上下文微调。学习到的选择器在训练期间使用块级证据监督。相对于ICAE，它还减少了在线首令牌时间并提高了自回归解码吞吐量。

SeDeM：长上下文问答中隐藏状态记忆的选择性解压缩
Maryam Haghifam、Jason Cong和Yizhou Sun
加利福尼亚大学洛杉矶分校，美国
\{maryamhgf, cong, yzsun\}@cs.ucla.edu

参照图题
图1：SeDeM概览。模型将长上下文信息存储在紧凑的记忆库中，选择与查询相关的块，将它们选择性地解压缩为解码器兼容的隐藏状态，并将重建的状态注入解码器以生成答案。

## 1 引言

使用LLM进行长上下文推理成本高昂：预填充阶段的自注意力随序列长度呈二次方增长，KV缓存随处理的令牌数量增长。扩大上下文窗口并不能保证问答（QA）任务中可靠的答案质量：LLM可能无法使用长输入中的相关信息，尤其是当证据不在上下文开头或结尾附近时（[Liu et al., 2024a](https://arxiv.org/html/2608.00311#bib.bib26)）。因此，长上下文QA需要既高效又选择性地向解码器暴露证据的方法。上下文压缩通过在生成前用较短的表示替换完整输入来解决效率问题。硬提示压缩方法剪枝或选择离散令牌（[Li et al., 2023](https://arxiv.org/html/2608.00311#bib.bib1)；[Jiang et al., 2023](https://arxiv.org/html/2608.00311#bib.bib2), [2024](https://arxiv.org/html/2608.00311#bib.bib3)；[Pan et al., 2024](https://arxiv.org/html/2608.00311#bib.bib4)），而软压缩方法用连续记忆令牌表示上下文（[Chevalier et al., 2023](https://arxiv.org/html/2608.00311#bib.bib27)；[Ge et al., 2024](https://arxiv.org/html/2608.00311#bib.bib6)；[Li et al., 2025](https://arxiv.org/html/2608.00311#bib.bib7)；[Zhang et al., 2024](https://arxiv.org/html/2608.00311#bib.bib20)）。许多软压缩方法使用LLM本身作为压缩器：学习到的记忆或压缩令牌被插入输入，并通过跨Transformer层的自注意力进行更新。先前工作认为，这种设计可能覆盖早期层聚合的信息，导致紧凑记忆偏离局部源内容（[Ye et al., 2026](https://arxiv.org/html/2608.00311#bib.bib11)）。此外，相同的紧凑槽通常被解码器直接作为软前缀使用，因此它们既要存储源信息，又要充当解码器条件状态。在高压缩比下，这种直接条件化路径变得越来越受限。

SeDeM采取了不同的设计方向：它从中间编码器层提取隐藏状态，使用轻量级非注意力压缩器将它们存储为紧凑的段级记忆块，选择与查询相关的块，并在生成前仅将选定的块扩展为与中间解码器层兼容的隐藏状态。我们提出SeDeM，即*隐藏状态记忆的选择性解压缩*，一个用于长上下文问答的选择性解压缩框架。我们的关键思想是将上下文存储在紧凑、查询无关的记忆中，并且仅在将选定的记忆扩展为与中间解码器层兼容的表示后，才将其暴露给解码器。冻结的LLM将每个上下文段映射为隐藏状态，轻量级压缩器将每个段存储为紧凑的记忆块。在推理时，查询条件选择器选择相关的记忆块。随后，解压缩器仅将选定的块扩展为重建的隐藏状态，并在中间解码器层注入。因此，解码器避免了处理完整的原始上下文，同时也避免了直接从高度压缩的记忆槽生成。将存储与解码器条件化解耦后有如下三个结果。第一，记忆构建是查询无关的，因此同一上下文的多个查询可以复用记忆。第二，LLM计算量随选择预算而非完整上下文长度扩展。第三，该设计支持跨模型使用，即较小的编码器构建记忆，解压缩后供较大的解码器使用。这将上下文处理成本与用于答案生成的能力分离。

我们在2WikiMultiHopQA、MuSiQue、QASPER和HotpotQA-Distractor上评估SeDeM。在Llama-3.2-1B和Llama-3.2-3B同主干设置中，SeDeM在四个数据集上的报告设置中均取得比所有评估的压缩基线更高的分数。学习到的选择器使用块级证据监督进行训练。使用3B主干时，它还在2WikiMHQA、QASPER和HotpotQA-Distractor上超过了全上下文微调参考，但在MuSiQue上仍低于该参考。与ICAE相比，SeDeM减少了在线首令牌时间并提高了自回归解码吞吐量。消融实验表明，直接记忆条件化明显弱于选择性解压缩，这表明收益来自将选定记忆扩展为解码器兼容的隐藏状态，而不仅仅是压缩本身。

## 2 相关工作

##### 硬提示压缩。
硬提示压缩方法通过在生成前选择或剪枝离散输入令牌来降低长上下文成本。SelectiveContext（[Li et al., 2023](https://arxiv.org/html/2608.00311#bib.bib1)）使用语言模型惊喜度移除信息量较少的词汇单元，而LLMLingua（[Jiang et al., 2023](https://arxiv.org/html/2608.00311#bib.bib2)）、LongLLMLingua（[Jiang et al., 2024](https://arxiv.org/html/2608.00311#bib.bib3)）和LLMLingua-2（[Pan et al., 2024](https://arxiv.org/html/2608.00311#bib.bib4)）在目标预算下学习或控制令牌级压缩。这些方法保留了标准文本界面，但压缩仍与离散令牌选择绑定。

##### 软压缩和基于记忆的上下文压缩。
第二条研究路线用连续记忆表示替代令牌剪枝。Gist令牌（[Mu et al., 2023](https://arxiv.org/html/2608.00311#bib.bib5)）和AutoCompressors（[Chevalier et al., 2023](https://arxiv.org/html/2608.00311#bib.bib27)）训练语言模型将提示或段总结为软令牌。ICAE（[Ge et al., 2024](https://arxiv.org/html/2608.00311#bib.bib6)）将上下文压缩视为上下文自编码，并让LLM直接消费记忆槽，而500xCompressor（[Li et al., 2025](https://arxiv.org/html/2608.00311#bib.bib7)）将文本映射为少量特殊记忆令牌。CompLLM（[Berton et al., 2025](https://arxiv.org/html/2608.00311#bib.bib9)）独立压缩段，使压缩表示可跨查询复用；HMT（[He et al., 2025](https://arxiv.org/html/2608.00311#bib.bib10)）在长上下文处理期间检索结构化记忆组件。其中几种方法遵循LLM作为压缩器的范式，即压缩令牌通过自注意力被反复更新。ComprExIT（[Ye et al., 2026](https://arxiv.org/html/2608.00311#bib.bib11)）批评该范式会导致逐层信息覆盖和令牌容量不协调，并转而通过冻结的隐藏状态进行显式信息传输。相比之下，SeDeM将记忆视为紧凑存储而非最终解码器输入：选定的记忆块在生成前被解压缩为解码器兼容的隐藏状态。

##### KV缓存和激活压缩。
另一条研究路线通过修改解码期间存储或重用的Transformer状态来降低推理成本。StreamingLLM（[Xiao et al., 2024](https://arxiv.org/html/2608.00311#bib.bib12)）保留注意力汇聚点和最近窗口，而H2O（[Zhang et al., 2023](https://arxiv.org/html/2608.00311#bib.bib14)）、SnapKV（[Li et al., 2024](https://arxiv.org/html/2608.00311#bib.bib15)）、PyramidKV（[Cai et al., 2024](https://arxiv.org/html/2608.00311#bib.bib16)）、KIVI（[Liu et al., 2024c](https://arxiv.org/html/2608.00311#bib.bib18)）、CacheGen（[Liu et al., 2024b](https://arxiv.org/html/2608.00311#bib.bib19)）和Activation Beacon（[Zhang et al., 2024](https://arxiv.org/html/2608.00311#bib.bib20)）在长历史中减少缓存大小、内存带宽或注意力成本。SeDeM则学习可复用的段级隐藏状态记忆，选择与查询相关的块，并在解码器层条件化之前选择性地解压缩它们。

## 3 方法

设X=\(x_1,\ldots,x_N\)表示长上下文，Q表示查询，Y表示目标答案。标准的全上下文推理模型p\(Y\mid X,Q\)要求解码器处理完整上下文。SeDeM改为从上下文段构建查询无关的压缩记忆，仅选择与Q相关的记忆，并将选定的块解压缩为解码器兼容的隐藏状态。这些重建状态在解码器层\(\ell_{\mathrm{inject}}\)注入，因此答案生成以选定的隐藏状态重建为条件，而不是以完整令牌序列或仅以记忆令牌为条件。我们的方法包含三个组件：压缩器C_{\theta}、top-k选择器R_{\phi}和解压缩器D_{\psi}。长上下文首先被划分为段。对于每个段，冻结的编码器从选定的Transformer层提取隐藏状态，压缩器将这些状态映射为紧凑的记忆块。给定查询后，选择器选择相关记忆，解压缩器仅将选定的块扩展为在解码器层\(\ell_{\mathrm{inject}}\)注入的隐藏状态。

### 3.1 压缩器

对于每个上下文段X_s，我们使用LLM编码器作为特征提取器，并从选定的提取层\(\ell_{\mathrm{extract}}\)获取隐藏状态：

\(H_s^{(\ell_{\mathrm{extract}})}=\mathrm{LLM_{Enc}}_{\ell_{\mathrm{extract}}}(X_s)\in\mathbb{R}^{T\times d_{\mathrm{enc}}}.\)  (1)

这次编码器传递提供了上下文相关的表示；实际压缩由局部池化和投影执行，而不是将学习到的压缩令牌插入LLM自注意力层。为清晰起见，我们描述长度为T=N_MC的完整段的压缩器，其中N_M是每段记忆槽数量，C是每槽压缩因子。压缩器将隐藏状态划分为N_M个大小为C的连续块，因此每个记忆槽总结一个局部块。对于每个记忆槽j∈\{1,\ldots,N_M\}，我们计算局部均值池化表示：

\(\bar{h}_{s,j}^{(\ell_{\mathrm{extract}})}=\frac{1}{C}\sum_{k=1}^{C}H_s^{(\ell_{\mathrm{extract}})}[(j-1)C+k],\)  (2)

其中\(\bar{h}_{s,j}^{(\ell_{\mathrm{extract}})}\in\mathbb{R}^{d_{\mathrm{enc}}}\)。均值池化在相邻上下文隐藏状态上给出无参数的局部压缩算子。直觉是，短窗口内相邻的上下文状态携带重叠的局部信息，因此它们的均值可以作为紧凑的局部摘要。这使输入投影独立于压缩因子C：压缩器使用\(W_{\mathrm{in}}\in\mathbb{R}^{d_{\mathrm{enc}}\times d_{\mathrm{dec}}}\)，而不是从C\cdot d_{\mathrm{enc}}到d_{\mathrm{dec}}的投影。共享投影将每个池化向量映射到LLM解码器隐藏维度：

\(m_{s,j}=\bar{h}_{s,j}^{(\ell_{\mathrm{extract}})}W_{\mathrm{in}}.\)  (3)

这里\(W_{\mathrm{in}}\in\mathbb{R}^{d_{\mathrm{enc}}\times d_{\mathrm{dec}}}\)，\(m_{s,j}\in\mathbb{R}^{d_{\mathrm{dec}}}\)。段s的最终记忆块为：

\(M_s=\left[m_{s,1};\ldots;m_{s,N_M}\right]\in\mathbb{R}^{N_M\times d_{\mathrm{dec}}}.\)  (4)

记忆块保留了粗略的源顺序：\(m_{s,1},\ldots,m_{s,N_M}\)对应于X_s中连续的局部窗口。由于编码器计算\(H_s^{(\ell_{\mathrm{extract}})}\)时包含位置信息，每个池化向量都来自位置感知的上下文状态。因此，压缩器将每个段从T个状态减少为N_M个记忆向量，并使它们与解码器隐藏维度对齐。它不增加逐令牌注意力：每个记忆向量通过局部池化后接共享投影\(W_{\mathrm{in}}\)获得。此外，与压缩令牌方法不同，压缩器不会向LLM自注意力层添加额外令牌。

### 3.2 记忆块上的Top-K选择

我们的压缩是查询无关的；查询只影响评分和选择。段记忆形成全局记忆库\(M_{\mathrm{bank}}=\mathrm{Concat}_{s=1}^{S}M_s\in\mathbb{R}^{(SN_M)\times d_{\mathrm{dec}}}\)。我们将该记忆库划分为\(N_{\mathrm{blk}}\)个大小为\(N_B\)的连续记忆块，其中\(N_{\mathrm{blk}}=SN_M/N_B\)，并定义\(B_n=M_{\mathrm{bank}}[(n-1)N_B+1:nN_B]\)。每个块满足\(B_n\in\mathbb{R}^{N_B\times d_{\mathrm{dec}}}\)。在我们的主要设置中，\(N_B=N_M\)，因此每个记忆块对应一个上下文段。我们使用“选择”而非“文本检索”这一术语，因为模型选择的是潜在记忆块而非原始文本。与检索增强生成（RAG）（[Lewis et al., 2020](https://arxiv.org/html/2608.00311#bib.bib29)）不同，它不会将检索到的篇章附加到提示中；选定的记忆随后会被解压缩为解码器层隐藏状态。

为了对记忆块进行评分，我们使用查询编码器。给定查询Q，我们将其编码为查询向量\(q\in\mathbb{R}^{d_{\mathrm{dec}}}\)：

\(q=\mathrm{Pool}(\mathrm{LLM_{Query}}(Q)),\)

其中LLM_{Query}是一个轻量级查询编码器（例如，与编码器相同的LLM或较小的编码器），Pool是均值池化。然后，每个记忆块的分数为：

\(s_n=\frac{q^\top \mathrm{vec}(B_n)}{\|q\|\|\mathrm{vec}(B_n)\|},\)

其中\(\mathrm{vec}(B_n)\in\mathbb{R}^{N_B d_{\mathrm{dec}}}\)将块展平为向量。我们选择分数最高的k个块：\(B_{n_1},\ldots,B_{n_k}\)。选择过程是端到端可微的，因此选择器可以在块级证据监督下训练。

### 3.3 解压缩器

对于每个选定的记忆块\(B_n\)，解压缩器将其扩展为与解码器层\(\ell_{\mathrm{inject}}\)兼容的隐藏状态序列。由于\(B_n\in\mathbb{R}^{N_B\times d_{\mathrm{dec}}}\)，我们需要生成\(\tilde{H}_n\in\mathbb{R}^{T'\times d_{\mathrm{dec}}}\)，其中\(T'=N_B\cdot C'\)，\(C'\)是解压缩因子。解压缩器是一个轻量级神经网络\(D_{\psi}\)，其输入为记忆块\(B_n\)，输出为重建的隐藏状态：

\(\tilde{H}_n=D_{\psi}(B_n).\)

我们使用一种简单的解压缩器结构：先对每个记忆向量进行上采样（例如，通过重复或插值），然后通过一个小型前馈网络或一维卷积进行平滑处理。在训练期间，解压缩器以编码器在层\(\ell_{\mathrm{inject}}\)的隐藏状态作为监督目标，从而学习将压缩的记忆块映射回解码器兼容的表示。在推理时，仅对被选中的记忆块进行解压缩，并将重建的隐藏状态注入解码器层\(\ell_{\mathrm{inject}}\)。这样，解码器只需要处理选定块的重建状态，而不是完整的原始上下文。

### 3.4 训练目标

SeDeM的端到端训练包含两个部分。第一部分是语言建模损失，即冻结解码器在注入重建隐藏状态后生成目标答案的负对数似然：

\(\mathcal{L}_{\mathrm{LM}}=-\log p(Y\mid \tilde{H}_{n_1},\ldots,\tilde{H}_{n_k},Q).\)

第二部分是块级证据监督损失，用于训练选择器。对于每个训练样本，我们提供块级证据标签\(y_n\in\{0,1\}\)，表示记忆块\(B_n\)是否包含回答问题所需的证据。选择器的训练目标为：

\(\mathcal{L}_{\mathrm{sel}}=-\sum_{n=1}^{N_{\mathrm{blk}}} y_n\log \sigma(s_n) + (1-y_n)\log(1-\sigma(s_n)),\)

其中\(\sigma\)是sigmoid函数。总损失为：

\(\mathcal{L}=\mathcal{L}_{\mathrm{LM}}+\lambda\mathcal{L}_{\mathrm{sel}},\)

其中\(\lambda\)是平衡两个损失的系数。通过这种方式，选择器学会识别与查询相关的记忆块，而解压缩器学会将选定块重建为解码器可用的表示。