arXiv:2605.11744cs.CLcs.LG2026-05中稿 · ICML

让长文本生成的训练和推理用相同分段方式,提升效率与一致性。

Training-Inference Consistent Segmented Execution for Long-Context LLMs

  • 训练时只允许前一段的键值状态参与梯度传播,保持与推理一致
  • 在128K上下文长度下,峰值预填充内存降低约6倍
  • 适合需要高效处理超长文本的部署场景

基于Transformer的大语言模型在长上下文生成中面临严重的可扩展性挑战,主要源于全上下文注意力带来的计算与内存开销。在实际计算与内存约束下,许多高效的推理方法仅在推理阶段采用有界上下文或分段执行,而训练时仍使用全上下文注意力,导致训练与推理的执行机制及状态转移语义不一致。针对此问题,我们提出一种训练-推理一致的分段生成框架,使训练与推理均遵循相同的分段前向执行语义。训练过程中,通过限制梯度仅沿前一段的键值(KV)状态回传,强制与推理的一致性;同时在前向过程中允许头级访问历史KV状态,但不将其纳入梯度传播。在多个长上下文基准上,该方法性能接近全上下文注意力,且在延迟-内存权衡上优于强基线,尤其在极长上下文长度下显著提升可扩展性(如128K上下文时,相比FlashAttention的全上下文注意力,峰值预填充内存降低约6倍)。

原文摘要 · Abstract (English)

Transformer-based large language models face severe scalability challenges in long-context generation due to the computational and memory costs of full-context attention. Under practical computation and memory constraints, many inference-efficient long-context methods improve efficiency by adopting bounded-context or segment-level execution only during inference, while continuing to train models under full-context attention, resulting in a mismatch between training and inference execution and state-transition semantics. Based on this insight, we propose a training-inference consistent segment-level generation framework, in which training and inference follow the same segment-level forward execution semantics. During training, consistency with inference is enforced by restricting gradient propagation to KV states carried over from the immediately preceding segment, while permitting head-specific access to past KV states during the forward pass without involving them in gradient propagation. Across long-context benchmarks, our approach achieves performance comparable to full-context attention, while achieving competitive latency-memory trade-offs against strong inference-efficient baselines, and substantially improving scalability at very long context lengths (e.g., approximately 6x lower peak prefill memory at 128K compared to full-context attention with FlashAttention).

长上下文分段执行推理优化

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