ParallelComp: Parallel Long-Context Compressor for Length Extrapolation
作者: Jing Xiong, Jianghan Shen, Chuanyang Zheng, Zhongwei Wan, Chenyang Zhao, Chiwun Yang, Fanghua Ye, Hongxia Yang, Lingpeng Kong, Ngai Wong
分类: cs.CL
发布日期: 2025-02-20 (更新: 2025-06-09)
备注: This paper has been accepted by ICML 2025
🔗 代码/项目: GITHUB
💡 一句话要点
提出ParallelComp并行长文本压缩方法，解决LLM长文本外推的内存瓶颈和注意力衰减问题。
🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)
关键词: 长文本外推 大型语言模型 注意力机制 KV缓存 并行计算 文本压缩 注意力偏差
📋 核心要点
- 现有长文本外推方法受限于内存瓶颈和注意力衰减，难以有效处理超长文本（>128K tokens）。
- ParallelComp通过并行KV缓存淘汰机制，动态移除冗余块和不相关tokens，实现长文本压缩。
- 实验表明，ParallelComp使8B模型在128K长文本上达到GPT-4性能的91.17%，预填充阶段加速23.50倍。
📝 摘要（中文）
本文提出ParallelComp，一种并行的长文本压缩方法，旨在有效克服内存瓶颈，使80亿参数的LLM能够在单张A100 80GB GPU上，实现从8K到128K tokens的无训练外推。ParallelComp将输入分割成块，并利用并行KV缓存淘汰机制动态地移除冗余块和不相关的tokens。重要的是，我们对并行注意力中的注意力偏差（包括注意力沉没、近因偏差和中间偏差）进行了系统的理论和实证分析，并揭示了这些偏差在超长上下文设置下的独特模式。我们进一步设计了一种KV缓存淘汰技术来缓解这种现象。实验结果表明，ParallelComp使一个在8K上下文上训练的80亿参数模型在超长上下文中达到GPT-4性能的91.17%，优于Claude-2和Kimi-Chat等闭源模型。我们在chunk吞吐量上实现了1.76倍的提升，从而在预填充阶段实现了23.50倍的加速，且性能损失可忽略不计，为LLM中可扩展且鲁棒的超长上下文外推铺平了道路。代码已开源。
🔬 方法详解
问题定义：论文旨在解决大型语言模型（LLMs）在超长上下文（>128K tokens）外推时面临的内存瓶颈和注意力衰减问题。现有的无训练外推方法在处理超长文本时，由于内存限制和注意力沉没现象，其可扩展性和有效性受到严重限制。
核心思路：论文的核心思路是采用并行的长文本压缩方法，即ParallelComp。该方法通过将输入分割成块，并动态地淘汰冗余块和不相关的tokens，从而在不进行额外训练的情况下，实现对超长上下文的处理。这种设计旨在缓解内存压力，并减轻注意力偏差的影响。
技术框架：ParallelComp的技术框架主要包括以下几个阶段：1) 输入分割：将超长文本输入分割成多个块（chunks）。2) 并行处理：对各个块进行并行处理，提取关键信息。3) KV缓存淘汰：基于一定的策略，动态地淘汰KV缓存中冗余或不相关的块和tokens。4) 注意力计算：在压缩后的上下文上进行注意力计算，生成最终的输出。
关键创新：论文的关键创新在于：1) 提出了并行的长文本压缩方法ParallelComp，有效缓解了内存瓶颈。2) 对并行注意力中的注意力偏差（注意力沉没、近因偏差、中间偏差）进行了系统的理论和实证分析，并揭示了这些偏差在超长上下文设置下的独特模式。3) 设计了一种KV缓存淘汰技术，用于缓解注意力偏差。
关键设计：ParallelComp的关键设计包括：1) Chunk大小的选择：需要平衡计算效率和信息损失。2) KV缓存淘汰策略：论文设计了一种基于注意力偏差分析的淘汰策略，优先淘汰对当前预测影响较小的块和tokens。3) 并行计算框架：采用并行计算框架加速处理过程，提高吞吐量。
🖼️ 关键图片
📊 实验亮点
ParallelComp使一个在8K上下文上训练的80亿参数模型在128K超长上下文中达到GPT-4性能的91.17%，超过了Claude-2和Kimi-Chat等闭源模型。在chunk吞吐量上实现了1.76倍的提升，预填充阶段加速23.50倍，且性能损失可忽略不计。这些结果表明ParallelComp在长文本外推方面具有显著优势。
🎯 应用场景
ParallelComp可应用于需要处理超长文本的各种场景，例如长篇文档摘要、大规模代码理解、长时间对话建模等。该方法能够显著提升LLM处理长文本的能力，降低计算成本，并为开发更强大的长文本应用提供技术支持。未来，该技术有望在知识密集型任务、智能客服、内容创作等领域发挥重要作用。
📄 摘要（原文）
Extrapolating ultra-long contexts (text length >128K) remains a major challenge for large language models (LLMs), as most training-free extrapolation methods are not only severely limited by memory bottlenecks, but also suffer from the attention sink, which restricts their scalability and effectiveness in practice. In this work, we propose ParallelComp, a parallel long-context compression method that effectively overcomes the memory bottleneck, enabling 8B-parameter LLMs to extrapolate from 8K to 128K tokens on a single A100 80GB GPU in a training-free setting. ParallelComp splits the input into chunks, dynamically evicting redundant chunks and irrelevant tokens, supported by a parallel KV cache eviction mechanism. Importantly, we present a systematic theoretical and empirical analysis of attention biases in parallel attention-including the attention sink, recency bias, and middle bias-and reveal that these biases exhibit distinctive patterns under ultra-long context settings. We further design a KV cache eviction technique to mitigate this phenomenon. Experimental results show that ParallelComp enables an 8B model (trained on 8K context) to achieve 91.17% of GPT-4's performance under ultra-long contexts, outperforming closed-source models such as Claude-2 and Kimi-Chat. We achieve a 1.76x improvement in chunk throughput, thereby achieving a 23.50x acceleration in the prefill stage with negligible performance loss and pave the way for scalable and robust ultra-long contexts extrapolation in LLMs. We release the code at https://github.com/menik1126/ParallelComp.