arXiv:2506.01215cs.CLcs.LG2025-06NeurIPS

通过分阶段压缩与重计算,高效处理超长文本输入。

Compress, Gather, and Recompute: REFORMing Long-Context Processing in Transformers

  • 分两阶段处理:逐步压缩并构建跨层上下文嵌入,结合早停策略提升效率。
  • 在100万上下文长度下,比基线性能提升超52%和34%。
  • 适合需要长文本理解的复杂任务,如代码分析、多轮对话等。

随着大语言模型在实际应用中的普及,处理远超模型预训练上下文长度的超长文本已成为关键挑战。现有高效长上下文处理方法中,基于循环压缩的方法存在信息丢失问题,而随机访问方法则需大量内存。本文提出REFORM,一种新型推理框架,采用两阶段策略:第一阶段逐步处理输入块,维护压缩的键值缓存(KV cache),构建跨层上下文嵌入,并采用早停机制提升效率;第二阶段通过相似性匹配识别关键令牌,选择性地重新计算KV缓存。相比基线,REFORM在RULER和BABILong数据集上分别实现超过52%和34%的性能提升(100万上下文长度)。同时,在Infinite-Bench、RepoEval和MM-NIAH上也优于基线,展现跨任务与领域的适应性。此外,推理时间减少30%,峰值内存降低5%,兼顾高效与高性能。

原文摘要 · Abstract (English)

As large language models increasingly gain popularity in real-world applications, processing extremely long contexts, often exceeding the model's pre-trained context limits, has emerged as a critical challenge. While existing approaches to efficient long-context processing show promise, recurrent compression-based methods struggle with information preservation, whereas random access approaches require substantial memory resources. We introduce REFORM, a novel inference framework that efficiently handles long contexts through a two-phase approach. First, it incrementally processes input chunks while maintaining a compressed KV cache, constructs cross-layer context embeddings, and utilizes early exit strategy for improved efficiency. Second, it identifies and gathers essential tokens via similarity matching and selectively recomputes the KV cache. Compared to baselines, REFORM achieves over 52% and 34% performance gains on RULER and BABILong respectively at 1M context length. It also outperforms baselines on Infinite-Bench, RepoEval, and MM-NIAH, demonstrating flexibility across diverse tasks and domains. Additionally, REFORM reduces inference time by 30% and peak memory usage by 5%, achieving both efficiency and superior performance.

长文本处理Transformer优化推理加速上下文压缩

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