通过分块调度让循环记忆Transformer实现长文本并行推理
Diagonal Batching Unlocks Parallelism in Recurrent Memory Transformers for Long Contexts
- 提出对角分批调度，让RMT在保持精确递推的同时并行处理序列段
- 在13万词长序列上比标准Transformer快3.3倍，比串行RMT快1.8倍
- 无需重训练即可部署，适合需要高效长文本生成的场景
Transformer模型在长序列推理时因时间复杂度为二次、内存复杂度为线性而面临挑战。循环记忆Transformer（RMT）通过将渐近成本降至线性时间与恒定内存，提供了可行方案。然而其记忆更新机制导致执行串行化，形成性能瓶颈。本文提出对角分批（Diagonal Batching）调度策略，可在不破坏精确递推的前提下，解锁RMT中各段间的并行性。该方法完全基于运行时计算重排，无需重新训练即可适配现有RMT模型。在LLaMA-1B ARMT模型上，该方法在131,072个标记的序列上实现了比标准全注意力模型3.3倍的加速，以及比串行RMT实现1.8倍的加速。通过对角分批，显著降低了推理延迟与成本，使RMT成为实际长上下文应用中的高效解决方案。
原文摘要 · Abstract (English)
Transformer models struggle with long-context inference due to their quadratic time and linear memory complexity. Recurrent Memory Transformers (RMTs) offer a solution by reducing the asymptotic cost to linear time and constant memory usage. However, their memory update mechanism leads to sequential execution, causing a performance bottleneck. We introduce Diagonal Batching, a scheduling scheme that unlocks parallelism across segments in RMTs while preserving exact recurrence. This approach eliminates the sequential constraint, enabling efficient GPU inference even for single long-context inputs without complex batching and pipelining techniques. Because the technique is purely a run-time computation reordering, existing RMT models adopt it with no retraining. Applied to a LLaMA-1B ARMT model, Diagonal Batching yields a 3.3x speedup over standard full-attention LLaMA-1B and a 1.8x speedup over the sequential RMT implementation on 131,072-token sequences. By removing sequential bottleneck, Diagonal Batching reduces inference cost and latency, thereby strengthening RMTs as a practical solution for real-world, long-context applications.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。