基于大模型的上下文压缩及自编码器

本文提出了一种基于大型语言模型（LLM）的上下文内自编码器（ICAE），通过将长上下文压缩为紧凑记忆槽来优化模型性能。ICAE采用两阶段训练：先在文本数据上通过自编码和语言建模目标进行预训练，再通过指令数据进行微调。实验表明，基于Llama模型的ICAE仅增加1%参数即可实现4倍上下文压缩，显著降低推理延迟和内存消耗，同时保持了良好的信息保留能力。研究还发现ICAE的记忆机制与人类相似，并展示了其

誉鏐

369人浏览 · 2025-11-14 12:42:09

誉鏐 · 2025-11-14 12:42:09 发布

【前言】本文基于以下论文讲解《用于大语言模型中上下文压缩的上下文自编码器》，为了方便大家理解，大家可以点开下面论文链接，直接查看原文。

https://arxiv.org/abs/2307.06945https://arxiv.org/abs/2307.06945

【摘要】

本文提出了上下文内自编码器（ICAE），它利用大型语言模型（LLM）的能力，将长上下文压缩成简短的紧凑记忆槽，这些记忆槽可以直接被目标LLM用于各种目的的条件生成。

ICAE首先在海量文本数据上，通过自编码和语言建模目标进行预训练，使其能够生成准确且全面代表原始上下文的记忆槽。随后，通过指令数据对其进行微调，以针对各种提示生成理想的响应。实验表明，轻量级的ICAE（仅引入约1%的额外参数）基于Llama模型有效实现了4倍的上下文压缩，在推理延迟和GPU内存成本方面均带来优势，并在记忆能力方面展现出有趣的洞察以及可扩展性的潜力。

1、INTRODUCTION

对于基于Transformer架构的大型语言模型而言，长上下文建模因其固有的自注意力机制而成为一个根本性挑战。先前的大量研究试图通过LLM的架构创新来解决长上下文问题。尽管这些方法显著降低了计算和内存复杂度以处理长上下文，但它们往往难以克服长上下文上性能的显著下降。

与这些努力相反，我们从一个新颖的角度——上下文压缩——来探讨长上下文问题。

正如上图所示：

当被大型语言模型用作条件输入以响应给定提示时，不同的上下文长度（例如2572字符、512词、128个记忆槽）能够实现相同的功能。

1.1上下文压缩的动机

一段文本可以在LLM中以不同长度表示，同时传递相同的信息。

本文探索了这一问题并提出ICAE——利用LLM的能力实现上下文的高倍压缩。ICAE包含两个模块：

1）可学习编码器（基于LLM通过LoRA技术适配而来），负责将长上下文编码为少量记忆槽；

2）固定解码器（即LLM本体），通过将代表原始上下文的记忆槽作为条件输入，与各类提示交互以完成不同任务。

1.1.1 Piepline

首先通过自编码和语言建模双重目标对ICAE进行预训练，使其学会生成能让解码器（即LLM）重建原始上下文或执行文本续写的记忆槽。基于海量文本数据的预训练使ICAE获得良好泛化能力，其所生成的记忆槽能更准确全面地表征原始上下文。

随后，通过指令微调使预训练后的ICAE能适应实际场景，增强其生成记忆槽与各类提示的交互能力。

【注】基于Llama架构的ICAE通过本文中所提到的训练方法，可实现4倍上下文压缩的有效记忆槽生成。

1.1.2 ICAE

上下文自编码器（ICAE）——一种利用LLM能力实现上下文压缩的创新方法。

除上下文压缩外，ICAE为探索LLM的记忆机制提供了途径。我们观察到，在预训练阶段进行的广泛自监督学习能有效提升ICAE将原始上下文编码为压缩记忆槽的能力。这一预训练过程与人类通过记忆训练提升大脑编码能力的机制存在相似性。

2、IN-CONTEXT AUTOENCODER

2.1 模型架构

如图所示：

ICAE的编码器是一个经过LoRA适配的LLM，用于将原始上下文 c=(w1,w2,…,wL)c=(w1,w2,…,wL) 编码为少量记忆槽 (m^1,…,m^k)(m^1,…,m^k)。

ICAE的解码器则是目标LLM本身，它可以基于编码器生成的记忆槽执行各种任务（例如本图中的自编码任务）。e(⋅)e(⋅) 表示目标LLM中的词嵌入查找，em(⋅)em(⋅) 表示用于生成记忆槽的可学习记忆令牌嵌入查找。"[AE]"是一个特殊标记，用于指示自编码预训练任务。

当编码长度为 LL 的上下文 c=(w1,…,wL)c=(w1,…,wL) 时，我们首先在上下文 cc 后追加 kk (k<<Lk<<L) 个记忆令牌 (m1,…,mk)(m1,…,mk)，并将其输出 (m^1,…,m^k)(m^1,…,m^k) 作为该上下文的记忆槽。因此ICAE编码器非常轻量——相较于目标LLM仅增加了LoRA适配器和记忆令牌的嵌入查找层。

如前所述，我们希望记忆槽 (m^1,…,m^k)(m^1,…,m^k) 能够替代原始上下文 cc 作为目标LLM的条件输入。因此，我们使用未经改动的目标LLM作为ICAE的解码器，以确保记忆槽在目标LLM内的兼容性。

2.2 预训练

2.2.1 自编码

与典型自编码器类似，ICAE的预训练目标之一是从其生成的长度为 k 的记忆槽 (m1~,…,mk~)(m1,…,mk) 中还原原始长度为 LL 的输入文本 cc：

为标识自编码任务，我们在解码器中的记忆槽 (m1~,…,mk~)(m1,…,mk) 后添加特殊标记"[AE]"，由于该预训练目标无需任何额外标注，我们可以利用海量文本数据来训练此上下文自编码器。

2.2.2 文本续写

虽然自编码预训练为上下文编码提供了直接的学习目标，但其固有的简单性及对单一目标的专注可能导致泛化能力不足。为解决此问题，我们在预训练阶段引入了额外目标：文本续写。这一自监督任务被广泛认为有助于语言模型学习更具泛化性的表征：

其中 o=(wL+1,…,wL+N)o=(wL+1,…,wL+N) 表示上下文 cc 的后续内容。该目标有助于提升泛化能力，并避免对自编码任务的过度依赖与过拟合。

2.3 指令微调

对于LLM而言，提供上下文的目的不仅仅是为了机械记忆或文本续写；更常见的使用场景是将所提供的上下文作为基础，以准确恰当地响应各种提示，最终完成我们期望其执行的任务。

为增强ICAE生成的记忆槽与多样化提示的交互能力，我们使用PwC数据集（带上下文的提示数据集）对ICAE进行进一步微调。该数据集由本文引入，包含数千个（上下文，提示，响应）样本（如图1所示）。

形式上，ICAE通过微调学习将上下文编码为记忆槽，基于这些记忆槽，解码器（即目标LLM）能够根据给定提示 p1…pmp1…pm 生成符合期望的响应 r1……rn

3、EXPERIMENTS

3.1 实验设置

数据我们使用Pile数据集对ICAE进行预训练。在指令微调阶段，采用第2.3节介绍的PwC数据集，该数据集包含24万组（上下文，提示，响应）训练样本和1.8万组测试样本。测试样本的上下文长度分布如图10所示。实验默认设置中，训练时ICAE编码器与解码器的最大令牌长度（不含记忆槽）均为512。

模型配置 我们采用LlaMa作为目标LLM来验证ICAE的上下文压缩性能。在ICAE编码器中，LoRA适配器应用于LLM多头注意力机制的查询和值投影层。默认配置下，记忆槽长度k设为128，LoRA秩r设为128。最终构建的ICAE仅在目标LLM基础上增加约1%的可学习参数。

3.2 实验结果

3.2.1 预训练ICAE评估

我们首先评估未经指令微调的预训练ICAE在自编码任务中的表现，通过以下三个指标衡量其从记忆槽还原原始上下文的能力：BLEU、精确匹配（EM）² 和交叉熵损失。

如图所示：基于Llama-7b的ICAE自编码结果。模型整体损失值低于0.05，表明生成的记忆槽几乎完整保留了原始上下文信息。当上下文长度在300以内时，ICAE能近乎完美地重建原文，BLEU与EM分数均接近100%。当上下文长度超过400后，两项指标开始下降，说明128长度的记忆槽容量渐显不足。但即使上下文长度达到500，其中位数BLEU值仍超过0.98，EM中位数接近0.6（例如能完美重建512词上下文中的前300词），展现出卓越性能。

我们进一步分析了记忆槽容量kk对结果的影响。如上图所示，随着记忆槽长度k减小，模型记忆长样本的能力显著下降。

【总结】

图4：基于Llama-7b、记忆槽长度k=128的ICAE自编码结果。横轴表示测试样本的原始上下文长度（例如横轴标值100对应上下文长度在95至105之间的测试样本）。图示显示：

随着上下文长度增加，BLEU和精确匹配(EM)分数逐渐下降
交叉熵损失始终保持在较低水平(<0.05)
在300词以内上下文长度时，模型能近乎完美重建原文

图5：不同记忆槽长度k下的BLEU分数和损失曲线对比。结果表明：

记忆槽容量与压缩性能正相关：k=128时，500词长上下文仍能保持95%以上BLEU值
压缩能力存在临界点：当k降至64和32时，性能急剧下降，说明实现4倍以上无损压缩具有挑战性
损失曲线验证了相同趋势，k值减小会导致重建质量显著降低

核心结论：ICAE在k=128设置下能有效实现4倍上下文压缩，但超过此压缩比将导致信息丢失。

同样，表1所示的文本续写评估结果也表明，更高的压缩比往往会导致语言建模性能更显著的损失。

表2展示了ICAE执行文本重建的一个具体案例："large pretrained language model"被重建为"large pretrained model"，而"The results prove"被重建为"The experimental evidence proves"。这些重建错误类似于人类在记忆相同文本时会犯的错误。这表明，模型在记忆过程中会基于自身理解选择性强调或忽略某些信息。

结论就是：LLM能力越强，需要记忆的内容就越少，因此记忆负担越小。

为深入探究记忆机制，我们测试了ICAE使用128个记忆槽对不同类型512词文本的重建性能，以验证其记忆能力在不同内容类型间是否一致。根据表3结果，与能较好重建的常规文本相比，压缩和重建非常见文本（如随机文本）变得极具挑战性，这体现在显著恶化的损失值和BLEU分数上。所有这些结果都强有力地支持了我们的直觉：LLM的记忆模式与人类高度相似。

3.2.2 微调后的ICAE

为评估微调后ICAE的性能，我们在PwC测试集上进行测试。

我们使用GPT-4对比两个系统的输出，以判断哪个表现更优或是否表现相当。

表4展示了基于记忆槽和原始上下文的LLM结果对比。对于Llama-7b（微调ICAE），由于没有官方指令微调的Llama-1模型，我们选择与Alpaca和StableLM-tuned-alpha-7b进行比较。使用128个记忆槽的Llama-7b（ICAE）显著优于可访问原始上下文（约512个标记）的Alpaca和StableLM，胜率分别达到56.7%和74.1%，胜率+平局率合计达73%～81%。然而，与GPT-4（我们将其视为黄金标准）相比仍存在显著差距，约70%的情况表现不及GPT-4，胜率+平局率仅约30%。

当我们将基础模型切换为Llama-2-chat时，观察到ICAE的性能相比基于Llama-1的版本有显著提升：当k=128k=128时，其胜率+平局率再次达到约75%；尽管由于有损压缩，其表现仍落后于使用原始上下文的对比模型。随着kk值增大，胜率+平局率进一步提升，但压缩率相应降低。我们在Llama-2-13b-chat上进行相同对比研究，观察到ICAE更优的结果，这支持了我们在第3.2.1节的假设：ICAE在更大规模LLM上能获得更多收益。

我们研究了记忆槽长度对结果的影响。表5展示了不同记忆槽长度ICAE模型间的两两比较。更高压缩比确实会增加保证响应质量的难度，但更大压缩比并不总是导致更差性能。表5突出显示：经过预训练的8×8×压缩（k=64k=64）ICAE，其表现可与未预训练的4×4×压缩（k=128k=128）ICAE相媲美。在相同压缩比下，经过预训练的ICAE性能远优于未预训练的模型。

【总结】

未经预训练的模型，这强调了预训练的重要性。通过对比经过预训练和未经预训练的ICAE所生成的输出，我们发现经过预训练的ICAE出现幻觉的情况要少于未经预训练的模型（示例见附录D中的表9）。我们推测，ICAE的预训练改善了LLM的工作记忆，因为这与人类通过大量记忆训练来提升大脑记忆编码能力，从而增强记忆容量存在某种相似性。我们还检验了预训练目标，发现结合使用自编码和语言建模目标比单独使用其中任何一个都能产生更好的结果（见表5第4行）。

表5的最后一行将ICAE的128长度记忆槽与限制在128个标记（约100词）内的摘要进行了比较。在相同上下文长度下，记忆槽的表现显著优于摘要，其胜/负比约为2:1，这证明记忆槽比自然语言更具紧凑性和信息丰富性。

3.3 分析

3.3.1 可扩展性

如前所述，ICAE在使用更强大的目标LLM时应能实现更好的压缩性能。为验证此假设，我们在表6中比较了ICAE在三个目标LLM上的表现：Llama-7b、Llama-2-7b和Llama-2-13b。结果完全符合预期——更强大的目标LLM能够实现更优的上下文压缩比。

3.3.2 延迟优化

我们通过实验评估了ICAE的4倍上下文压缩对推理效率的影响。在此效率测试中，我们将上下文长度固定为512或2048，生成长度固定为128。表7显示，ICAE的上下文压缩有助于提升LLM的推理效率，实现了超过2倍的加速效果。

在计算密集型场景下（例如8批次×2048长度和32批次×512长度），其加速效果更为显著——可达约3.5倍。

考虑到压缩后的记忆槽可以预先缓存（适用于教科书、政府报告或法律条文等常用文本），在这些场景中ICAE可带来超过7倍的推理加速效果。

3.3.3 多段记忆槽

至此，我们主要讨论了单段记忆槽。本节将探讨多段记忆槽的应用。如图6（左）所示，我们可以将长上下文分割为 NN 个片段，分别进行压缩后拼接，以此表征原始长上下文。然而，这种方法在初始阶段并未奏效，因为模型在训练过程中从未接触过多段拼接的模式。幸运的是，我们可以通过在训练中加入少量多段拼接样本，使模型具备处理拼接记忆槽的能力——这与OpenAI通过特定训练为GPT引入"中间填充"能力的思路一致。

图6（右）结果显示，在输入总长度相同的情况下，ICAE的记忆槽能实现更优的性能——因为记忆槽可表征4倍于原始上下文长度的信息。

ICAE的这一能力展现了处理长上下文的巨大潜力：它能在不修改现有LLM的前提下，显著节省处理长上下文时所需的GPU内存。如图6（右）所示，2048长度的记忆槽即可达到与4096词元上下文相当的效果。这意味着使用2048个记忆槽替代原始4096个上下文词元，可在保证质量损失最小的同时，节省约20GB的GPU内存。

4、结论与未来工作

我们提出了上下文内自编码器，通过利用LLM的能力实现上下文的高倍压缩。通过生成紧凑且信息丰富的记忆槽来表征原始上下文，ICAE使LLM能够在相同上下文长度下获取更多信息，或以更短上下文表示相同内容，从而增强模型处理长上下文的能力，并在检索增强生成和高级提示方法等实际应用场景中降低推理时的计算与内存开销。此外，ICAE为探索LLM的记忆机制提供了视角，为理解LLM与人类记忆之间的联系提供了新思路，并为LLM上下文管理的未来研究指明了方向。

受计算资源所限，我们的实验仅在130亿参数规模的Llama模型上进行。如文中所述，ICAE有望在更强大的LLM上获得更大收益，实现更显著的压缩比。未来我们期望获得充足计算资源，在更大规模LLM上验证ICAE的有效性。同时，我们计划探索ICAE在多模态LLM中的应用——鉴于图像、视频和音频的上下文通常更长且具有更大压缩潜力，采用离散记忆槽将有助于在LLM/AGI时代实现跨模态紧凑表征的统一。