Two are better than one: Context window extension with multi-grained self-injection

📄 arXiv: 2410.19318v1 📥 PDF

作者: Wei Han, Pan Zhou, Soujanya Poria, Shuicheng Yan

分类: cs.CL, cs.AI, cs.LG

发布日期: 2024-10-25

备注: The code is available at https://github.com/Clement25/SharedLLM

💡 一句话要点

提出SharedLLM，通过多粒度自注入扩展LLM上下文窗口，降低长文本处理成本。

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

关键词: 长文本处理 上下文窗口扩展 多粒度压缩 自注入 大型语言模型

📋 核心要点

  1. 现有LLM上下文窗口有限，阻碍了其在长文本任务中的应用，而持续预训练成本高昂。
  2. SharedLLM采用双模型结构，下层模型压缩上下文，上层模型解码，实现多粒度上下文信息的自注入。
  3. 通过树形数据结构和搜索算法，SharedLLM能够高效地编码、存储和检索多粒度上下文信息。

📝 摘要（中文）

当前大型语言模型（LLM）有限的上下文窗口严重阻碍了它们在各个领域的广泛应用。虽然在长上下文数据上进行持续预训练是一种直接而有效的解决方案，但它在数据获取和计算资源方面会产生巨大的成本。为了缓解这个问题，我们提出了一种名为SharedLLM的新方法，该方法基于多粒度上下文压缩和查询感知信息检索的设计理念。SharedLLM由两个短上下文LLM（如LLaMA-2）组成，分别称为上层模型和下层模型。下层模型充当压缩器，而上层模型充当解码器。上层模型接收来自下层模型的压缩的多粒度上下文信息，并对运行文本执行上下文感知建模。压缩器和解码器之间的信息传输仅发生在最低层，以避免下层模型中的长前向路径和上层模型中的冗余交叉注意力模块。基于此架构，我们引入了一种专门的树形数据结构，以有效地编码、存储和检索文本块的多粒度上下文信息。这种结构与搜索算法相结合，能够根据输入查询快速编码和检索来自树的各个级别的信息。整个过程，其中发送者和接收者都来自同一LLM层，被称为自注入。

🔬 方法详解

问题定义：当前大型语言模型（LLM）的上下文窗口大小有限，这限制了它们在需要处理长文本的任务中的应用，例如长文档摘要、问答等。直接采用长文本数据进行持续预训练虽然有效，但需要大量的计算资源和数据，成本很高。

核心思路：SharedLLM的核心思路是利用两个短上下文的LLM，一个作为压缩器（下层模型），另一个作为解码器（上层模型），通过多粒度的上下文压缩和自注入机制，扩展LLM的有效上下文窗口，同时避免高昂的预训练成本。这样设计的目的是在不显著增加计算负担的前提下，使模型能够处理更长的文本序列。

技术框架：SharedLLM的整体架构包含两个主要部分：下层模型（压缩器）和上层模型（解码器）。下层模型负责将输入文本块压缩成多粒度的上下文表示，并将这些表示存储在一个树形数据结构中。上层模型接收当前文本块和从树形结构中检索到的相关上下文信息，然后进行解码和建模。信息在两个模型之间通过“自注入”的方式传递，即信息发送者和接收者都来自同一LLM层。

关键创新：SharedLLM的关键创新在于多粒度的上下文压缩和自注入机制。多粒度压缩允许模型保留不同层次的上下文信息，从而更好地捕捉长文本中的依赖关系。自注入机制避免了在整个模型中进行长距离的信息传递，降低了计算复杂度。此外，树形数据结构和搜索算法的结合，实现了高效的上下文信息检索。

关键设计：SharedLLM使用两个共享参数的LLM（如LLaMA-2）作为上下层模型。下层模型将输入文本块编码成多粒度的上下文表示，这些表示被组织成一个树形结构，其中每个节点代表一个不同粒度的上下文信息。上层模型使用一个查询感知的搜索算法，从树形结构中检索相关的上下文信息，并将这些信息与当前文本块一起输入到解码器中。信息传递只发生在模型的最低层，以减少计算量。具体的损失函数和网络结构细节在论文中可能有所描述，但摘要中未提及。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

摘要中未提供具体的实验数据和性能指标。但是，SharedLLM的设计目标是扩展LLM的上下文窗口，同时降低计算成本。因此，可以推断实验结果会关注模型在长文本任务上的性能提升，以及计算效率的提高。具体的性能提升幅度未知。

🎯 应用场景

SharedLLM具有广泛的应用前景，例如长文档摘要、问答系统、代码生成、机器翻译等。它可以帮助LLM更好地理解和处理长文本，提高模型在这些任务上的性能。此外，SharedLLM的设计理念也可以应用于其他需要处理长序列数据的领域，例如视频理解、语音识别等。

📄 摘要（原文）

The limited context window of contemporary large language models (LLMs) remains a huge barrier to their broader application across various domains. While continual pre-training on long-context data is a straightforward and effective solution, it incurs substantial costs in terms of data acquisition and computational resources. To alleviate this issue, we propose SharedLLM, a novel approach grounded in the design philosophy of multi-grained context compression and query-aware information retrieval. SharedLLM is composed of two short-context LLMs such as LLaMA-2, termed upper model and lower model. The lower model functions as a compressor while the upper model acts as a decoder. The upper model receives compressed, multi-grained context information from the lower model and performs context-aware modeling on the running text. Information transfer between the compressor and decoder occurs only at the lowest layers to refrain from long forward paths in the lower model and redundant cross-attention modules in the upper model. Based on this architecture, we introduce a specialized tree-style data structure to efficiently encode, store and retrieve multi-grained contextual information for text chunks. This structure, combined with a search algorithm, enables rapid encoding and retrieval of relevant information from various levels of the tree based on the input query. This entire process, wherein the sender and receiver are derived from the same LLM layer, is referred to as self-injection.