Stacked from One: Multi-Scale Self-Injection for Context Window Extension
作者: Wei Han, Pan Zhou, Shuicheng Yan
分类: cs.CL, cs.AI
发布日期: 2026-03-05
💡 一句话要点
提出SharedLLM,通过多尺度自注入扩展LLM上下文窗口至128K tokens。
🎯 匹配领域: 支柱九:具身大模型 (Embodied Foundation Models)
关键词: 长文本建模 上下文扩展 自注入学习 多粒度压缩 大型语言模型 高效推理 树形数据结构
📋 核心要点
- 现有LLM上下文窗口受限,阻碍其在多样化领域的应用,直接采用长上下文数据进行持续预训练成本高昂。
- SharedLLM通过多粒度上下文压缩和查询感知信息获取,利用自注入机制,在短上下文模型基础上扩展长上下文处理能力。
- 实验表明,SharedLLM在长上下文任务上表现优异,且推理速度更快、内存占用更少,实现了效率与性能的平衡。
📝 摘要(中文)
本文提出SharedLLM,一种基于多粒度上下文压缩和查询感知信息获取的新框架,旨在解决大型语言模型(LLM)有限的上下文窗口问题。SharedLLM由两个堆叠的短上下文LLM组成:一个下层模型作为压缩器,一个上层模型作为解码器。下层模型将长输入压缩成紧凑的多粒度表示,然后传递给上层模型进行上下文感知处理。为了最大化效率,信息传递仅发生在最底层,避免了冗长的正向传播和多余的交叉注意力操作。整个过程,上下层模型都源自相同的底层LLM层,被称为“自注入”。为了支持这种架构,一种专门的基于树的数据结构实现了上下文信息的高效编码和查询感知检索。尽管仅在8K tokens的序列上训练,SharedLLM有效地泛化到超过128K tokens的输入。在全面的长上下文建模和理解基准测试中,SharedLLM实现了优于或可与强基线相媲美的性能,在效率和准确性之间取得了最佳平衡。此外,这些设计选择使SharedLLM能够显著减少内存占用,并产生显著的推理加速(比流式传输快2倍,比编码器-解码器架构快3倍)。
🔬 方法详解
问题定义:现有大型语言模型(LLM)的上下文窗口大小有限,无法有效处理长序列输入,限制了其在需要长程依赖的任务中的应用。直接采用长序列数据进行持续预训练虽然有效,但数据获取和计算成本非常高昂,难以负担。
核心思路:本文的核心思路是利用两个共享底层参数的短上下文LLM,通过“自注入”的方式,将一个LLM作为压缩器,另一个作为解码器,实现对长上下文信息的有效压缩、传递和利用。下层模型负责将长输入压缩成多粒度表示,上层模型则基于这些压缩表示进行推理,从而扩展模型的有效上下文窗口。
技术框架:SharedLLM包含两个堆叠的短上下文LLM:下层模型(压缩器)和上层模型(解码器)。长输入首先被下层模型处理,生成多粒度上下文表示。这些表示随后被传递到上层模型,用于后续的推理和预测。为了支持高效的上下文信息编码和检索,论文还设计了一种基于树的数据结构。整个过程的关键在于“自注入”,即上下层模型共享相同的底层LLM参数,从而降低了训练成本。
关键创新:SharedLLM最重要的创新点在于其“自注入”机制和多粒度上下文压缩方法。通过共享底层参数,模型可以在短序列上进行训练,然后泛化到长序列,避免了昂贵的长序列预训练。多粒度上下文压缩则允许模型在不同尺度上捕获上下文信息,从而更有效地利用长程依赖。
关键设计:为了最大化效率,信息传递仅发生在模型的最低层,避免了冗长的正向传播和冗余的交叉注意力操作。论文还设计了一种专门的树形数据结构,用于高效地编码和查询上下文信息。具体的参数设置、损失函数和网络结构等细节在论文中进行了详细描述,但摘要中未明确提及。
🖼️ 关键图片
📊 实验亮点
SharedLLM在仅使用8K tokens序列训练的情况下,能够有效处理超过128K tokens的输入。在长上下文建模和理解基准测试中,SharedLLM的性能优于或可与强大的基线模型相媲美。此外,SharedLLM显著减少了内存占用,并实现了显著的推理加速(比流式传输快2倍,比编码器-解码器架构快3倍)。
🎯 应用场景
SharedLLM可应用于需要处理长文本的各种场景,如长文档摘要、代码生成、法律文本分析、金融报告解读等。该方法降低了长上下文建模的计算成本,使得在资源受限的环境中部署长上下文LLM成为可能,具有广泛的应用前景。
📄 摘要(原文)
The limited context window of contemporary large language models (LLMs) remains a primary bottleneck for their broader application across diverse domains. Although continual pre-training on long-context data offers a straightforward solution, it incurs prohibitive data acquisition and computational costs. To address this challenge, we propose~\modelname, a novel framework based on multi-grained context compression and query-aware information acquisition. SharedLLM comprises two stacked short-context LLMs: a lower model serving as a compressor and an upper model acting as a decoder. The lower model compresses long inputs into compact, multi-grained representations, which are then forwarded to the upper model for context-aware processing. To maximize efficiency, this information transfer occurs exclusively at the lowest layers, bypassing lengthy forward passes and redundant cross-attention operations. This entire process, wherein the upper and lower models are derived from the same underlying LLM layers, is termed~\textit{self-injection}. To support this architecture, a specialized tree-based data structure enables the efficient encoding and query-aware retrieval of contextual information. Despite being trained on sequences of only 8K tokens, \modelname~effectively generalizes to inputs exceeding 128K tokens. Across a comprehensive suite of long-context modeling and understanding benchmarks, \modelname~achieves performance superior or comparable to strong baselines, striking an optimal balance between efficiency and accuracy. Furthermore, these design choices allow \modelname~to substantially reduce the memory footprint and yield notable inference speedups ($2\times$ over streaming and $3\times$ over encoder-decoder architectures).