这两天,加州大学伯克利分校、斯坦福大学、英伟达等机构联合制作的《猫和老鼠》AI短片火了。
论文共同一作 Karan Dalal 的帖子收获百万观看。
我们先来欣赏下面两段视频,重温儿时的快乐。
故事梗概:In an underwater adventure, Jerry locates a treasure map and searches for the treasure while evading Tom in a chase through coral reefs and kelp forests. Jerry triumphantly discovers treasure inside a shipwreck, blissfully celebrating while Tom's pursuit leads him into trouble with a hungry shark. (在一次水下探险中,Jerry 找到了藏宝图,并在躲避 Tom 追捕的同时寻找宝藏,途中穿过珊瑚礁和海藻森林。Jerry 在一艘沉船中发现了宝藏,欣喜若狂地庆祝着,而 Tom 的追捕却让他陷入了饥饿鲨鱼的困境。)
故事梗概:Tom is happily eating an apple pie at the kitchen table. Jerry looks longingly wishing he had some. Jerry goes outside the front door of the house and rings the doorbell. While Tom comes to open the door, Jerry runs around the back to the kitchen. Jerry steals Tom's apple pie. Jerry runs to his mousehole carrying the pie, while Tom is chasing him. Just as Tom is about to catch Jerry, he makes it through the mouse hole and Tom slams into the wall.(Tom 在厨房餐桌上开心地吃着苹果派。Jerry 看起来渴望地想要吃一些。Jerry 走到屋前门外按门铃。Tom 来开门时,Jerry 绕到后面的厨房。Jerry 偷走了 Tom 的苹果派。Jerry 拿着苹果派跑到他的老鼠洞里,而 Tom 正在追赶他。就在 Tom 即将抓住 Jerry 时,他从老鼠洞里逃了出来,Tom 撞到了墙上。)
类似的《猫和老鼠》短片共五集,每集都是全新的故事。大家可以在项目主页查看对应的故事梗概和完整的提示词。
项目主页:https://test-time-training.github.io/video-dit/
效果怎么样?如果不提前告知,你能分辨出它们是 AI 生成的吗?
据论文另一位共同一作 Gashon Hussein 介绍,为了实现逼真的动画效果,他们利用 TTT(Test-time Training,测试时训练)层来增强预训练 Transformer,并进行了微调,从而生成了时间和空间上连贯性很强的《猫和老鼠》一分钟短片。
尤其值得注意的是,所有视频都是由模型一次性直接生成,没有进行任何二次编辑、拼接或后期处理。
Gashon Hussein 进一步解释了背后的技术原理。
TTT 层是专门的 RNN 层,其中每个隐藏状态代表了一个机器学习模型。此外,这些层内的更新使用梯度下降来完成。本文将 TTT 层集成到一个预训练的 Diffusion Transformer 中,随后使用文本标注对长时间视频进行微调。并且,为了管理计算复杂度,自注意力被限制在局部片段,而 TTT 层以线性复杂度来高效地处理全局上下文。
此外,为了高效地实现 TTT-MLP 内核,本文开发了一种「片上张量并行」(On-chip Tensor Parallel)算法,具体包括以下两个步骤:
- 在 GPU 流多处理器(Sreaming Multiprocessor,SM)之间划分隐藏状态模型的权重;
- 利用 Hopper GPU 的 DSMEM 功能在 SM 之间执行高效的 AllReduce 操作,显著降低全局内存(HBM)和共享内存(SMEM)之间的数据传输,确保大量隐藏状态在 SMEM 内可以有效访问。
下图 3 为方法概览,其中(左)为本文修改后的架构在每个注意力层后添加一个带有可学习门的 TTT 层,(右)为整体 pipeline 创建了由 3 秒片段组成的输入序列,这种结构允许在片段上局部应用自注意力层,在整个序列上全局应用 TTT 层。
具体实现过程是这样的:
本文研究者从一个预训练好的 DiT(CogVideo-X 5B)开始,它只能以 16 帧 / 秒的速度生成 3 秒钟的短片(或以 8 帧 / 秒的速度生成 6 秒钟的短片)。然后添加了从零开始初始化的 TTT 层,并对该模型进行微调,以便从文本故事板生成一分钟的视频。研究者将自注意力层限制在 3 秒钟的片段内,使其成本保持在可控范围内。仅通过初步的系统优化,训练运行就相当于在 256 台 H100 上花费了 50 个小时。
这项研究博得了评论区一众网友的惊呼与赞许。
研究细节
在论文《One-Minute Video Generation with Test-Time Training》中,英伟达、斯坦福等机构的研究者介绍了《猫和老鼠》短片背后的更多生成技术细节。
- 论文标题:One-Minute Video Generation with Test-Time Training
- 论文地址:https://arxiv.org/pdf/2504.05298
此前视频生成技术限制背后的根本挑战是长上下文,因为 Transformers 中自注意力层的成本随着上下文长度的增加而呈二次曲线增加。这一挑战在生成动态视频时尤为突出,因为动态视频的上下文不容易被 tokenizer 压缩。使用标准 tokenizer,每段一分钟的视频都需要 30 多万个上下文 token。基于自注意力,生成一分钟视频所需的时间要比生成 20 段每段 3 秒钟的视频增加 11 倍,而训练所需的时间也要增加 12 倍。
为了应对这一挑战,最近有关视频生成的研究将 RNN 层作为自注意力的有效替代方法,因为 RNN 层的成本随上下文长度呈线性增长。现代 RNN 层,尤其是线性注意力的变体(如 Mamba 和 DeltaNet),在自然语言任务中取得了令人印象深刻的成果。然而,我们还没有看到由 RNN 生成的具有复杂故事或动态动作的长视频。
本文研究者认为,这些 RNN 层生成的视频复杂度较低,因为它们的隐藏状态表现力较差。RNN 层只能将过去的 token 存储到固定大小的隐藏状态中,而对于 Mamba 和 DeltaNet 等线性注意力变体来说,隐藏状态只能是一个矩阵。要将成百上千个向量压缩到一个只有数千级的矩阵中,这本身就是一项挑战。因此,这些 RNN 层很难记住远处 token 之间的深层关系。
因此研究者尝试使用另一种 RNN 层,其隐藏状态本身也可以是神经网络。具体来说,研究者使用两层 MLP,其隐藏单元比线性注意力变体中的线性(矩阵)隐藏状态多 2 倍,非线性也更丰富。即使在测试序列上,神经网络的隐藏状态也会通过训练进行更新,这些新层被称为测试时间训练层(TTT)。
局部注意力机制在汤姆的颜色、杰瑞的鼠洞之间保持了一致性,并且扭曲了汤姆的身体。
TTT-MLP 在整个视频时长中表现出强大的特性和时间一致性。
研究者策划了一个文本到视频的数据集,该数据集基于大约 7 小时的《猫和老鼠》动画片,并附有人类注释的故事板。他们有意将范围限制在这一特定领域,以便快速进行研究迭代。作为概念验证,该数据集强调复杂、多场景和具有动态运动的长篇故事,此前的模型在这些方面仍需取得进展;而对视觉和物理逼真度的强调较少,因为此前的模型在这些方面已经取得了显著进展。研究者认为,尽管本文是面向这一特定领域的长上下文能力改进,但也会转移到通用视频生成上。
生成质量评估
在测评中,与 Mamba 2、Gated DeltaNet 和滑动窗口注意力层等强大的基线相比,TTT 层生成的视频更连贯,能讲述复杂的动态故事。
在 LMSys Chatbot Arena 中,GPT-4o 比 GPT-4 Turbo 高出 29 个 Elo 分数。
具体对比参考下列 demo:
当汤姆咆哮并追逐杰瑞时,Mamba 2 扭曲了汤姆的外貌。
Gated DeltaNet 在汤姆的不同角度上缺乏时间一致性。
滑动窗口注意力改变了厨房环境并重复了杰瑞偷馅饼的场景。
不过,在生成内容中,我们也能发现 TTT-MLP 一些明显的瑕疵。
1、时间一致性:盒子在同一场景的 3 秒片段之间变形。
2、运动自然性:奶酪悬停在半空中,而不是自然地落到地上。
3、美学:当汤姆转身时,厨房里的灯光变得更加明亮。