arXiv:2507.04416cs.CL2025-07NeurIPS被引 3

RAT用分块机制兼顾RNN效率与注意力精度,显著提速长序列处理。

RAT: Bridging RNN Efficiency and Attention Accuracy via Chunk-based Sequence Modeling

  • 分块处理:块内用循环建模局部依赖,块间用注意力捕获长程关系。
  • 10万序列下训练速度提升7倍,4千位置生成快9倍,性能接近标准注意力。
  • 适合需要高效长文本建模的场景,如超长文档理解、实时生成任务。

Transformer已成为现代大模型的核心,但其对softmax注意力的依赖在训练和推理中造成计算瓶颈。循环模型虽高效,但将全序列压缩为固定大小表示易导致长上下文中的记忆退化,并限制细粒度检索。为此,我们提出RAT,一种介于RNN效率与注意力能力之间的中间设计。RAT将输入分块,在块内使用循环建模局部依赖,块间采用softmax注意力捕捉长程交互。该设计缓解了记忆退化问题,实现对远距离词元的直接访问,同时保持计算高效。实验表明,当块大小为16时,RAT模块在10万序列长度下训练速度提升7倍,4千位置生成速度提升9倍,且性能与标准注意力相当。我们通过从零训练13亿参数模型并进行大规模评估(包括短/长上下文基准及监督微调)验证了其有效性。进一步提出一种混合架构,将RAT与局部注意力交替使用,结合高效的长程建模与强局部交互,不仅提升推理速度、降低缓存内存占用,还持续提升性能,取得最佳综合表现。代码已开源于https://github.com/CLAIRE-Labo/RAT。

原文摘要 · Abstract (English)

Transformers have become the cornerstone of modern large-scale language models, but their reliance on softmax attention poses a computational bottleneck at both training and inference. Recurrent models offer high efficiency, but compressing the full sequence into a fixed-size and holistic representation can suffer from memory degradation in long contexts and limit fine-grained retrieval. To address this, we propose RAT, an intermediate design that bridges the efficiency of RNNs and capacity of attention. RAT partitions the input into chunks, applies recurrence within each chunk for local dependencies, and softmax-based attention across chunks for long-range interactions. This design mitigates memory degradation and enables direct access to distant tokens, while retaining computational efficiency. Empirically, with a chunk size of 16, the RAT block achieves a 7$\times$ improvement in training speed for 100K sequence length and 9$times$ in generation at the 4K position, while maintaining similar performance compared to standard attention. We demonstrate this by training 1.3B parameter models from scratch and performing large-scale evaluations, including short- and long-context benchmarks, as well as supervised fine-tuning~(SFT). We further propose a hybrid architecture that interleaves RAT with local attention. By combining efficient long-range modeling with strong local interactions, this hybrid design not only improves inference speed and reduces cache memory usage, but also consistently enhances performance and shows the overall best results. Code is available at https://github.com/CLAIRE-Labo/RAT.

序列建模注意力机制RNN改进长文本处理

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