重新思考测试时训练的表现力和效率

测试时训练 (TTT) 通过在推理过程中连续权重更新来实现长上下文处理，但当前的方法很难平衡每个词元更新动态的表现力与逐块近似的硬件效率。我们提出 E$^2$-TTT（Expressive and Efficient TTT）来弥补这一差距。在块开始权重处采用梯度的标准近似下，我们导出了一个封闭形式的状态转换，它精确地再现了每个词元递归的块结束快速权重和动量状态。这使得完全并行的块级训练成为可能，同时保留了先前的块方式方法丢弃的更新规则的时间结构。我们通过从头开始训练高达 1.3B 参数的模型来验证 E$^2$-TTT。它在语言建模中的表现与之前的 TTT 和混合注意力基线相当，同时在上下文检索方面优于它们。它的优势在长度外推方面最为明显：在标准的“大海捞针”密钥测试中，它在 $8\times$ 训练上下文长度下保持了 90% 以上的准确度。同时，E$^2$-TTT 可以与高效的 chunk-wise 方法的训练吞吐量相匹配，证明它有效地协调了表达性和效率。代码可在 https://github.com/zeyun-zhong/E2-TTT 获取。