arXiv:2508.17892cs.CLcs.LG2025-08被引 1

通过中间层检索压缩上下文,让大模型处理百万级长文本更快更省内存。

ILRe: Intermediate Layer Retrieval for Context Compression in Causal Language Models

  • 选一个中间解码层,只预填充到该层,大幅降低计算量。
  • 处理100万token仅需30秒内,速度提升180倍,内存降至原体积的十分之一。
  • 无需微调即可在长文本任务中表现接近全上下文模型,适合高效推理场景。

大型语言模型在众多基准上表现优异,但在长上下文场景下仍受限于有效上下文长度短、计算复杂度呈二次增长以及高内存开销。为缓解这些问题,我们提出一种新型上下文压缩管道——中间层检索(ILRe):预先离线选定一个解码器中间层,仅将输入按块流式预填充至该层,并通过查询与该层完整键缓存的注意力分数召回原始标记。特别地,我们在召回过程中引入多池化核分配策略,以保持语义完整性。该方法将预填充复杂度从 $O(L^2)$ 降低至 $O(L)$,内存占用缩减至全上下文所需体积的数十分之一,同时在长上下文场景下性能达到或超过全上下文设置。无需额外微调或算子开发,ILRe 可在华为 Ascend 910B NPU 上以不到半分钟完成单次 1M tokens 请求(加速比 ≈180×),在 RULER-$1M$ 基准上取得约 79.8 的得分(模型 Llama-3.1-UltraLong-8B-1M-Instruct)。

原文摘要 · Abstract (English)

Large Language Models (LLMs) have demonstrated success across many benchmarks. However, they still exhibit limitations in long-context scenarios, primarily due to their short effective context length, quadratic computational complexity, and high memory overhead when processing lengthy inputs. To mitigate these issues, we introduce a novel context compression pipeline, called Intermediate Layer Retrieval (ILRe), which determines one intermediate decoder layer offline, encodes context by streaming chunked prefill only up to that layer, and recalls tokens by the attention scores between the input query and full key cache in that specified layer. In particular, we propose a multi-pooling kernels allocating strategy in the token recalling process to maintain the completeness of semantics. Our approach not only reduces the prefilling complexity from $O(L^2)$ to $O(L)$ and trims the memory footprint to a few tenths of that required for the full context, but also delivers performance comparable to or superior to the full-context setup in long-context scenarios. Without additional post training or operator development, ILRe can process a single $1M$ tokens request in less than half a minute (speedup $\approx 180\times$) and scores RULER-$1M$ benchmark of $\approx 79.8$ with model Llama-3.1-UltraLong-8B-1M-Instruct on a Huawei Ascend 910B NPU.

长文本生成上下文压缩推理优化LLM

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