提出E²-TTT,让测试时训练兼具表达力与效率
Rethinking Expressivity and Efficiency in Test-Time Training

- 基于分块梯度推导出闭式状态转移,保持时间动态结构
- 1.3B模型训练效果媲美旧方法,在长文本检索上更优
- 支持并行化训练,长度外推性能显著提升
测试时训练(TTT)通过推理阶段持续更新权重实现长上下文处理,但现有方法难以兼顾每个词的动态表达能力与分块近似带来的硬件效率。本文提出E²-TTT(表达力与效率兼备的TTT),在标准分块起始权重求梯度的假设下,推导出闭式状态转移公式,可精确复现逐词迭代下的分块结束时的快速权重与动量状态。该方法实现完全并行的分块级训练,同时保留了先前分块方法所丢弃的时间依赖结构。我们从零训练最大达1.3B参数的模型进行验证:在语言建模任务中表现与以往TTT及混合注意力基线相当,但在上下文检索任务中更优;其优势在长度外推中尤为突出——在标准的“针在草堆”传票测试中,训练上下文长度的4倍时仍保持超过90%准确率。同时,E²-TTT的训练吞吐量可匹配高效分块方法,有效调和了表达力与效率的矛盾。代码已开源。
原文摘要 · Abstract (English)
Test-Time Training (TTT) enables long-context processing via continuous weight updates during inference, but current methods struggle to balance the expressivity of per-token update dynamics with the hardware efficiency of chunk-wise approximations. We propose E$^2$-TTT (Expressive and Efficient TTT) to bridge this gap. Under the standard approximation of taking gradients at the chunk-start weights, we derive a closed-form state transition that exactly reproduces the chunk-end fast-weight and momentum states of the per-token recurrence. This enables fully parallelized chunk-level training while preserving the temporal structure of the update rule that prior chunk-wise methods discard. We validate E$^2$-TTT by training models up to 1.3B parameters from scratch. It performs on par with previous TTT and hybrid attention baselines in language modeling while outperforming them on in-context retrieval. Its advantage is most pronounced in length extrapolation: on the standard ``Needle in a Haystack'' passkey test, it retains over 90% accuracy at $8\times$ the training context length. Meanwhile, E$^2$-TTT can match the training throughput of efficient chunk-wise methods, demonstrating that it effectively reconciles expressivity with efficiency. The code is available at https://github.com/zeyun-zhong/E2-TTT.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。