arXiv:2505.23884cs.LGcs.CL2025-05被引 117

大块测试时训练让模型长序列记忆更强,效率更高。

Test-Time Training Done Right

  • 用超大块(2K到100万词元)更新快速权重,突破传统小批量限制。
  • 在56K词元长序列上实现140亿参数视频生成,最长达百万上下文。
  • 无需复杂内核,兼容先进优化器,适合图像、视频等多模态任务。

测试时训练(TTT)通过在推理阶段调整部分模型权重(快权重)来建模上下文依赖,这些快权重类似于RNN中的递归状态,临时存储当前序列中过往词元的记忆。现有方法在处理长上下文数据时表现不佳,因现代GPU利用率极低(常低于5%),主要由于其刻意采用极小的在线小批量(如每16或64词元更新一次)。此外,小批量导致细粒度的块状因果依赖,不适用于一维序列以外的数据,如集合或二维网格(如图像、视频)。相反,本文提出大块测试时训练(LaCT),采用从2000到100万词元不等的大块更新,显著提升硬件利用率,并支持高达模型参数40%的非线性状态容量,大幅提升记忆能力,且无需繁琐的内核实现。同时支持如Muon等复杂优化器的轻松集成。我们在多种模态和任务中验证了该方法,包括基于图像集的新视角合成、语言模型与自回归视频扩散模型。最大实验达到140亿参数的自回归视频扩散模型,在长达56,000词元的序列上运行;最长达100万词元的上下文用于新视角合成。本工作有望推动长上下文建模与测试时训练领域的发展。

原文摘要 · Abstract (English)

Test-Time Training (TTT) models context dependencies by adapting part of the model's weights (referred to as fast weights) during inference. This fast weight, akin to recurrent states in RNNs, stores temporary memories of past tokens in the current sequence. Existing TTT methods struggled to show effectiveness in handling long-context data, due to their inefficiency on modern GPUs. The TTT layers in many of these approaches operate with extremely low FLOPs utilization (often <5%) because they deliberately apply small online minibatch sizes (e.g., updating fast weights every 16 or 64 tokens). Moreover, a small minibatch implies fine-grained block-wise causal dependencies in the data, unsuitable for data beyond 1D ordered sequences, like sets or N-dimensional grids such as images or videos. In contrast, we pursue the opposite direction by using an extremely large chunk update, ranging from 2K to 1M tokens across tasks of varying modalities, which we refer to as Large Chunk Test-Time Training (LaCT). It improves hardware utilization by orders of magnitude, and more importantly, facilitates scaling of nonlinear state size (up to 40% of model parameters), hence substantially improving state capacity, all without requiring cumbersome and error-prone kernel implementations. It also allows easy integration of sophisticated optimizers, e.g. Muon for online updates. We validate our approach across diverse modalities and tasks, including novel view synthesis with image set, language models, and auto-regressive video diffusion. Our approach can scale up to 14B-parameter AR video diffusion model on sequences up to 56K tokens. In our longest sequence experiment, we perform novel view synthesis with 1 million context length. We hope this work will inspire and accelerate new research in the field of long-context modeling and test-time training. Website: https://tianyuanzhang.com/projects/ttt-done-right

测试时训练长序列建模视频生成大块更新

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