arXiv:2502.14317cs.CL2025-02ICML被引 18

提出并行压缩方法,让小模型在单卡上实现超长文本推理。

ParallelComp: Parallel Long-Context Compressor for Length Extrapolation

  • 分块处理输入,动态剔除冗余内容,结合并行缓存机制降低内存占用。
  • 8B模型在单张A100上实现8K到128K上下文扩展,性能达GPT-4的91.17%。
  • 解决注意力陷阱问题,提升吞吐与推理速度,适合资源受限场景部署。

超长上下文(文本长度>128K)的外推仍是大语言模型的重大挑战。现有无训练外推方法受内存瓶颈和注意力塌陷限制,难以实用。本文提出ParallelComp,一种并行长上下文压缩方法,在无需训练的前提下,使8B参数模型在单张A100 80GB GPU上实现从8K到128K token的外推。该方法将输入分块,通过并行键值缓存淘汰机制动态移除冗余块和无关词元。我们系统分析了并行注意力中的注意力偏差(包括注意力塌陷、近期性偏差、中间偏差),发现其在超长上下文下呈现特定模式,并设计相应缓存淘汰策略以缓解。实验表明,该方法使8B模型(训练于8K上下文)在超长上下文下的性能达到GPT-4的91.17%,优于Claude-2和Kimi-Chat等闭源模型。预填充阶段加速达23.50倍,吞吐提升1.76倍,性能损失可忽略,为大模型超长上下文外推提供了高效可靠的解决方案。代码已开源。

原文摘要 · Abstract (English)

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.

长上下文压缩推理加速KV缓存

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。