压缩、聚合与重计算：重构 Transformer 中的长上下文处理
计算与语言
2025-11-18 v2 机器学习
摘要
随着大语言模型在实际应用中日益普及，处理通常超出模型预训练上下文限制的超长上下文已成为一项关键挑战。虽然现有的高效长上下文处理方法展现出前景，但基于循环压缩的方法在信息保留方面存在困难，而随机访问方法则需要大量内存资源。我们引入了 REFORM，一种通过两阶段方法高效处理长上下文的新型推理框架。首先，它增量处理输入块，同时维护压缩的 KV 缓存，构建跨层上下文嵌入，并利用提前退出策略以提高效率。其次，它通过相似度匹配识别并聚合关键 token，并选择性地重计算 KV 缓存。与基线相比，REFORM 在 1M 上下文长度下的 RULER 和 BABILong 上分别实现了超过 52% 和 34% 的性能提升。它在 Infinite-Bench、RepoEval 和 MM-NIAH 上也优于基线，展示了在多种任务和领域上的灵活性。此外，REFORM 将推理时间减少了 30%，峰值内存使用量减少了 5%，实现了效率与卓越性能的双重目标。
关键词
引用
@article{arxiv.2506.01215,
title = {Compress, Gather, and Recompute: REFORMing Long-Context Processing in Transformers},
author = {Woomin Song and Sai Muralidhar Jayanthi and Srikanth Ronanki and Kanthashree Mysore Sathyendra and Jinwoo Shin and Aram Galstyan and Shubham Katiyar and Sravan Babu Bodapati},
journal= {arXiv preprint arXiv:2506.01215},
year = {2025}
}
备注
NeurIPS 2025