arXiv:2511.07343cs.LGcs.AI2025-11被引 7

TNT让测试时记忆RNN训练快17倍且更准,突破效率与性能的矛盾。

TNT: Improving Chunkwise Training for Test-Time Memorization

  • 分两阶段训练:先用大块并行预训练,再用小块微调精度。
  • 训练速度提升17倍,同时模型准确率更高。
  • 适合想高效训练长序列记忆RNN的研究者。

具有深层测试时记忆模块的循环神经网络(如 Titans 和 TTT)代表了一种有前景的线性扩展范式,与 Transformer 不同。尽管这些模型尚未达到顶尖 Transformer 的性能,但其潜力受限于极慢的训练速度和低硬件利用率。现有并行化方法受块大小超参数制约:大块加速但性能下降,需在速度与精度间做出次优妥协。为此,我们提出 TNT,一种新训练范式,通过两阶段过程解耦训练效率与推理性能。第一阶段为效率导向的预训练,采用层次化记忆结构:全局模块处理大块以获取长程上下文,多个并行局部模块处理细粒度细节。关键在于定期重置局部记忆状态,打破序列依赖,实现大规模上下文并行。第二阶段为短时微调,仅调整局部记忆模块至更小、高分辨率的块大小,以最小开销最大化精度。在 Titans 与 TTT 模型上评估,TNT 实现训练速度提升至最准确基线配置的 17 倍,同时提高模型准确率。这一改进消除了关键可扩展性障碍,为开发高表达力 RNN 提供了实用基础,并推动未来工作缩小与 Transformer 的性能差距。

原文摘要 · Abstract (English)

Recurrent neural networks (RNNs) with deep test-time memorization modules, such as Titans and TTT, represent a promising, linearly-scaling paradigm distinct from Transformers. While these expressive models do not yet match the peak performance of state-of-the-art Transformers, their potential has been largely untapped due to prohibitively slow training and low hardware utilization. Existing parallelization methods force a fundamental conflict governed by the chunksize hyperparameter: large chunks boost speed but degrade performance, necessitating a fixed, suboptimal compromise. To solve this challenge, we introduce TNT, a novel training paradigm that decouples training efficiency from inference performance through a two-stage process. Stage one is an efficiency-focused pre-training phase utilizing a hierarchical memory. A global module processes large, hardware-friendly chunks for long-range context, while multiple parallel local modules handle fine-grained details. Crucially, by periodically resetting local memory states, we break sequential dependencies to enable massive context parallelization. Stage two is a brief fine-tuning phase where only the local memory modules are adapted to a smaller, high-resolution chunksize, maximizing accuracy with minimal overhead. Evaluated on Titans and TTT models, TNT achieves a substantial acceleration in training speed-up to 17 times faster than the most accurate baseline configuration - while simultaneously improving model accuracy. This improvement removes a critical scalability barrier, establishing a practical foundation for developing expressive RNNs and facilitating future work to close the performance gap with Transformers.

RNN训练加速记忆机制

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