发表机构
University of Wisconsin–Madison; University of Michigan–Ann Arbor; Adobe(威斯康星大学麦迪逊分校; 密歇根大学安娜堡分校; Adobe公司)
机构由 AI 辅助整理,请以论文原文为准。AI 中文总结
提出通用测试时训练(uTTT),让所有层共享一个跨时间和深度递归的记忆,通过uTTT-MoE和uTTT-Dense实现,在语言建模和新视图合成中优于层私有方法。
AI 中文摘要
最近的测试时训练(TTT)架构将上下文压缩为快速权重,这些权重在线更新并作为记忆进行查询。现有的TTT设计将这种记忆保持为每层私有:它仅随时间递归,深度仅索引L个独立的记忆。我们认为记忆的所有权不必与深度绑定,并引入了通用测试时训练(uTTT),其中所有层读取和写入一个共享记忆,同时保留各层特定的骨干网络参数。因此,共享记忆在两个维度上递归:时间和深度,以块和层作为其单元:一个块中深层写入的内容可以被下一个块中的浅层读取。我们将这一思想实例化为uTTT-MoE和uTTT-Dense。uTTT-MoE将每个token头路由到由所有层共享的专家池中的少数专家;uTTT-Dense在每一层应用整个共享记忆而不进行路由。在语言建模中,uTTT-MoE在124M和760M参数规模下分别达到15.5和27.9的RULER准确率,在相同状态和活跃计算下比其层私有对应版本高出2.6和2.1个百分点,是测试的有界状态模型中最高的,且每token损失匹配或优于全注意力。在新视图合成中,在固定的每层计算下,共享在路由模型中使视图23对象PSNR提高0.92 dB,在密集模型中提高0.76 dB。
英文摘要
Recent Test-Time Training (TTT) architectures compress context into fast weights that are updated online and queried as memory. Existing TTT designs keep this memory private to each layer: it recurs only over time, and depth merely indexes L separate memories. We argue that memory ownership need not be tied to depth, and introduce Universal Test-Time Training (uTTT), in which all layers read and write one shared memory while retaining layer-specific backbone parameters. The shared memory thus recurs over two dimensions, time and depth, with chunks and layers as their units: a write by a deep layer in one chunk can be read by a shallow layer in the next. We instantiate this idea as uTTT-MoE and uTTT-Dense. uTTT-MoE routes each token head to a few experts in a pool shared by all layers; uTTT-Dense applies the whole shared memory at every layer without routing. In language modeling, uTTT-MoE reaches 15.5 and 27.9 RULER accuracy at 124M and 760M, 2.6 and 2.1 points above its layer-private counterpart at equal state and active compute, the highest among tested bounded-state models, with per-token loss matching or beating full attention. In novel view synthesis, sharing at fixed per-layer compute gains 0.92 dB in view-23 object PSNR in routed models and 0.76 dB in dense models.
Comments37 pages. Project page: https://zefan-cai.github.io/uTTT.github.io/ ; code: https://github.com/Zefan-Cai/uTTT