Test-Time Training Done Right

测试时间训练（TTT）通过在推理过程中调整模型权重的一部分（称为快速权重）来建模上下文依赖性。这种快速权重类似于循环神经网络（RNNs）中的循环状态，用于存储当前序列中过去标记的临时记忆。现有的 TTT 方法在处理长上下文数据时效果不佳，主要原因是它们在现代 GPU 上效率低下。在许多这些方法中，TTT 层以极低的浮点运算利用率运行（通常低于 5%），因为它们有意使用小的在线小批量更新（例如，每 16 或 64 个标记更新一次快速权重）。此外，小批量更新意味着数据中存在细粒度的块级因果依赖关系，这不适合超出一维有序序列的数据类型，例如集合或 N 维网格（如图像或视频）。相比之下，我们采取了相反的方向，使用极大的块更新策略，范围从 2K 到 1M 标记，适用于不同模态的任务，我们称之为大块测试时间训练（LaCT）。这种方法将硬件利用率提高了几个数量级，并且更重要的是，支持非线性状态大小的扩展（最多可达模型参数的 40%），从而显著提升状态容量，同时无需复杂的、容易出错的内核实现。此外，它还允许轻松集成复杂的优化器，例如用于在线更新的 Muon 优化器。我们在多种模态和任务上验证了我们的方法，包括基于图像集的新视角合成、语言模型以及自回归视频扩散模型。我们的方法可以扩展到多达 140 亿参数的自回归视频扩散模型，并处理长达 56K 标记的序列。在我们最长序列的实验中，我们进行了具有 1 百万上下文长度的新视角合成任务。我们希望这项工作能够激发并加速长上下文建模和测试时间训练领域的研究进展。网站：https://tianyuanzhang.com/projects/ttt-done-right

论文试图解决测试时训练（Test-Time Training, TTT）方法在处理长上下文数据时效率低下的问题，尤其是现有方法在现代GPU上硬件利用率低、更新频率受限以及难以扩展到多维数据的问题。这是一个重要但尚未完全解决的问题。

论文提出了一种名为Large Chunk Test-Time Training（LaCT）的方法，通过使用极大规模的块（chunk）更新（从2K到1M个token），显著提高了硬件利用率，并允许非线性状态规模扩大至模型参数的40%。这种方法与传统TTT方法形成对比，后者通常采用小批量更新策略，导致FLOPs利用率极低。

该方法支持多种模态任务，包括图像集的新视角合成、语言建模和自回归视频扩散等。实验展示了LaCT能够扩展到高达14B参数的自回归视频扩散模型，并在最长序列实验中处理了1百万token的上下文长度。代码和项目网站已开源（https://tianyuanzhang.com/projects/ttt-done-right），便于后续研究者复现和改进。未来可进一步探索更复杂的优化器和更大规模的数据集。

相关研究包括传统的TTT方法（如Fast Weights和Recurrent Memory Mechanisms），这些方法主要针对短上下文和一维序列数据。此外，还有其他关于长上下文建模的工作，例如Transformer-XL、Longformer和BigBird等，它们通过改进注意力机制来处理长序列。同时，MuP/MuTransfer等在线优化技术也为LaCT提供了理论支持。