论文精选

重思测试时训练的表达性和效率

Rethinking Expressivity and Efficiency in Test-Time Training

精选理由

E$^2$-TTT模型在语言建模和上下文检索方面表现优异,且在长度外推方面具有明显优势,值得一看。

AI 摘要

测试时训练(TTT)通过推理过程中的连续权重更新实现长上下文处理,但现有方法难以平衡每token更新动态的表达性与分块近似硬件效率。我们提出E$^2$-TTT(表达性和高效TTT)以弥合这一差距。在标准分块近似梯度取自块起始权重的情况下,我们推导出封闭形式的态转换,精确地重现了每token递归的块末快速权重和动量状态。这使块级训练可以完全并行化,同时保留了先前块方法丢弃的更新规则的时间结构。通过从头训练高达1.3B参数的模型验证E$^2$-TTT。它在语言建模方面与先前TTT和混合注意力基线表现相当,但在上下文检索方面表现更优。其优势在长度外推方面最为明显:在标准的“针插麦堆”passkey测试中,它在8倍训练上下文长度下保留了超过90%的准确率。同时,E$^2$-TTT可以匹配高效分块方法的训练吞吐量,表明它有效地协调了表达性和效率。代码可在https://github.com/zeyun-zhong/E2-TTT获取。

原文 · arXiv cs.LG

Rethinking Expressivity and Efficiency in Test-Time Training

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.