In-Context Former: Lightning-fast Compressing Context for Large Language Model

📄 arXiv: 2406.13618v2 📥 PDF

作者: Xiangfeng Wang, Zaiyi Chen, Zheyong Xie, Tong Xu, Yongyi He, Enhong Chen

分类: cs.CL

发布日期: 2024-06-19 (更新: 2024-11-05)

备注: Accepted by EMNLP2024(Findings)

💡 一句话要点

提出IC-Former，通过线性复杂度上下文压缩加速大语言模型推理。

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 大语言模型 上下文压缩 推理加速 交叉注意力 线性复杂度

📋 核心要点

  1. 现有上下文压缩方法依赖LLM自注意力，计算复杂度高，限制了长文本处理。
  2. IC-Former利用交叉注意力与少量可学习token，直接从上下文词嵌入中提取信息。
  3. 实验表明，IC-Former显著降低计算成本，加速推理，同时保持较高的性能水平。

📝 摘要（中文）

随着基于Transformer的大语言模型（LLMs）日益普及，降低其高昂的推理成本已成为重要的研究方向。一种有效的方法是压缩长输入上下文。现有方法通常利用LLM自身的自注意力机制进行上下文压缩。虽然这些方法取得了显著成果，但压缩过程仍然涉及二次时间复杂度，限制了其适用性。为了缓解这一限制，我们提出了In-Context Former（IC-Former）。与以往方法不同，IC-Former不依赖于目标LLM。相反，它利用交叉注意力机制和少量可学习的摘要token来直接压缩来自上下文词嵌入的信息。这种方法显著减少了推理时间，在压缩范围内实现了线性时间复杂度增长。实验结果表明，我们的方法在压缩过程中仅需要基线方法1/32的浮点运算，并将处理速度提高了68到112倍，同时在评估指标上实现了超过90%的基线性能。总而言之，我们的模型有效地降低了压缩成本，并使实时压缩场景成为可能。

🔬 方法详解

问题定义：现有的大语言模型上下文压缩方法，如利用自注意力机制的方法，其计算复杂度是二次方的，这使得它们在处理非常长的上下文时效率低下，成为实际应用中的瓶颈。论文旨在解决这一问题，即如何以更低的计算成本实现有效的上下文压缩，从而加速LLM的推理过程。

核心思路：IC-Former的核心思路是解耦上下文压缩过程与目标LLM，不再依赖LLM自身的自注意力机制。取而代之的是，使用一个独立的、轻量级的模块，通过交叉注意力机制直接从上下文词嵌入中提取关键信息，并将其压缩成少量可学习的摘要token。这种方法将压缩的计算复杂度降低到线性级别。

技术框架：IC-Former的整体框架包含以下几个主要步骤：1) 将输入上下文通过词嵌入层转换为词嵌入向量；2) 将词嵌入向量输入到IC-Former模块中；3) IC-Former模块利用交叉注意力机制，将上下文词嵌入压缩成少量可学习的摘要token；4) 将压缩后的摘要token与原始输入拼接，输入到目标LLM中进行推理。IC-Former模块是整个框架的核心，它由一个交叉注意力层和一个前馈神经网络层组成。

关键创新：IC-Former的关键创新在于其解耦了上下文压缩过程与目标LLM。通过引入交叉注意力机制和可学习的摘要token，IC-Former能够以线性时间复杂度实现上下文压缩，而无需依赖LLM自身的自注意力机制。这使得IC-Former能够更高效地处理长上下文，并加速LLM的推理过程。

关键设计：IC-Former的关键设计包括：1) 使用交叉注意力机制，将上下文词嵌入作为query，可学习的摘要token作为key和value，从而实现上下文信息的压缩；2) 使用少量可学习的摘要token，以降低计算成本；3) 将压缩后的摘要token与原始输入拼接，以保留原始输入中的部分信息；4) 损失函数的设计目标是使IC-Former能够提取与下游任务相关的关键信息。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，IC-Former在上下文压缩过程中仅需基线方法1/32的浮点运算，处理速度提升68到112倍，同时在评估指标上保持了超过90%的基线性能。这表明IC-Former在显著降低计算成本的同时，能够有效地保留上下文信息，并保持较高的模型性能。

🎯 应用场景

IC-Former适用于需要处理长文本输入的大语言模型应用场景，例如：长文档摘要、问答系统、代码生成等。它可以显著降低推理延迟，提高用户体验，并降低部署成本。未来，该技术有望应用于移动设备或边缘计算等资源受限的环境中，使大语言模型能够在更广泛的场景中得到应用。

📄 摘要（原文）

With the rising popularity of Transformer-based large language models (LLMs), reducing their high inference costs has become a significant research focus. One effective approach is to compress the long input contexts. Existing methods typically leverage the self-attention mechanism of the LLM itself for context compression. While these methods have achieved notable results, the compression process still involves quadratic time complexity, which limits their applicability. To mitigate this limitation, we propose the In-Context Former (IC-Former). Unlike previous methods, IC-Former does not depend on the target LLMs. Instead, it leverages the cross-attention mechanism and a small number of learnable digest tokens to directly condense information from the contextual word embeddings. This approach significantly reduces inference time, which achieves linear growth in time complexity within the compression range. Experimental results indicate that our method requires only 1/32 of the floating-point operations of the baseline during compression and improves processing speed by 68 to 112 times while achieving over 90% of the baseline performance on evaluation metrics. Overall, our model effectively reduces compression costs and makes real-time compression scenarios feasible.