這兩天，加州大學(xué)伯克利分校、斯坦福大學(xué)、英偉達(dá)等機(jī)構(gòu)聯(lián)合制作的《貓和老鼠》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 在一艘沉船中發(fā)現(xiàn)了寶藏，欣喜若狂地慶祝著，而 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 撞到了墻上。）

類似的《貓和老鼠》短片共五集，每集都是全新的故事。大家可以在項目主頁查看對應(yīng)的故事梗概和完整的提示詞。

項目主頁：https://test-time-training.github.io/video-dit/

效果怎么樣？如果不提前告知，你能分辨出它們是 AI 生成的嗎？

據(jù)論文另一位共同一作 Gashon Hussein 介紹，為了實現(xiàn)逼真的動畫效果，他們利用 TTT（Test-time Training，測試時訓(xùn)練）層來增強(qiáng)預(yù)訓(xùn)練 Transformer，并進(jìn)行了微調(diào)，從而生成了時間和空間上連貫性很強(qiáng)的《貓和老鼠》一分鐘短片。

尤其值得注意的是，所有視頻都是由模型一次性直接生成，沒有進(jìn)行任何二次編輯、拼接或后期處理。

Gashon Hussein 進(jìn)一步解釋了背后的技術(shù)原理。

TTT 層是專門的 RNN 層，其中每個隱藏狀態(tài)代表了一個機(jī)器學(xué)習(xí)模型。此外，這些層內(nèi)的更新使用梯度下降來完成。本文將 TTT 層集成到一個預(yù)訓(xùn)練的 Diffusion Transformer 中，隨后使用文本標(biāo)注對長時間視頻進(jìn)行微調(diào)。并且，為了管理計算復(fù)雜度，自注意力被限制在局部片段，而 TTT 層以線性復(fù)雜度來高效地處理全局上下文。

此外，為了高效地實現(xiàn) TTT-MLP 內(nèi)核，本文開發(fā)了一種「片上張量并行」（On-chip Tensor Parallel）算法，具體包括以下兩個步驟：

• 在 GPU 流多處理器（Sreaming Multiprocessor，SM）之間劃分隱藏狀態(tài)模型的權(quán)重；
• 利用 Hopper GPU 的 DSMEM 功能在 SM 之間執(zhí)行高效的 AllReduce 操作，顯著降低全局內(nèi)存（HBM）和共享內(nèi)存（SMEM）之間的數(shù)據(jù)傳輸，確保大量隱藏狀態(tài)在 SMEM 內(nèi)可以有效訪問。

下圖 3 為方法概覽，其中（左）為本文修改后的架構(gòu)在每個注意力層后添加一個帶有可學(xué)習(xí)門的 TTT 層，（右）為整體 pipeline 創(chuàng)建了由 3 秒片段組成的輸入序列，這種結(jié)構(gòu)允許在片段上局部應(yīng)用自注意力層，在整個序列上全局應(yīng)用 TTT 層。

具體實現(xiàn)過程是這樣的：

本文研究者從一個預(yù)訓(xùn)練好的 DiT（CogVideo-X 5B）開始，它只能以 16 幀 / 秒的速度生成 3 秒鐘的短片（或以 8 幀 / 秒的速度生成 6 秒鐘的短片）。然后添加了從零開始初始化的 TTT 層，并對該模型進(jìn)行微調(diào)，以便從文本故事板生成一分鐘的視頻。研究者將自注意力層限制在 3 秒鐘的片段內(nèi)，使其成本保持在可控范圍內(nèi)。僅通過初步的系統(tǒng)優(yōu)化，訓(xùn)練運(yùn)行就相當(dāng)于在 256 臺 H100 上花費(fèi)了 50 個小時。

這項研究博得了評論區(qū)一眾網(wǎng)友的驚呼與贊許。

研究細(xì)節(jié)

在論文《One-Minute Video Generation with Test-Time Training》中，英偉達(dá)、斯坦福等機(jī)構(gòu)的研究者介紹了《貓和老鼠》短片背后的更多生成技術(shù)細(xì)節(jié)。

• 論文標(biāo)題：One-Minute Video Generation with Test-Time Training 
• 論文地址：https://arxiv.org/pdf/2504.05298

此前視頻生成技術(shù)限制背后的根本挑戰(zhàn)是長上下文，因為 Transformers 中自注意力層的成本隨著上下文長度的增加而呈二次曲線增加。這一挑戰(zhàn)在生成動態(tài)視頻時尤為突出，因為動態(tài)視頻的上下文不容易被 tokenizer 壓縮。使用標(biāo)準(zhǔn) tokenizer，每段一分鐘的視頻都需要 30 多萬個上下文 token?；谧宰⒁饬?，生成一分鐘視頻所需的時間要比生成 20 段每段 3 秒鐘的視頻增加 11 倍，而訓(xùn)練所需的時間也要增加 12 倍。

為了應(yīng)對這一挑戰(zhàn)，最近有關(guān)視頻生成的研究將 RNN 層作為自注意力的有效替代方法，因為 RNN 層的成本隨上下文長度呈線性增長?，F(xiàn)代 RNN 層，尤其是線性注意力的變體（如 Mamba 和 DeltaNet），在自然語言任務(wù)中取得了令人印象深刻的成果。然而，我們還沒有看到由 RNN 生成的具有復(fù)雜故事或動態(tài)動作的長視頻。

本文研究者認(rèn)為，這些 RNN 層生成的視頻復(fù)雜度較低，因為它們的隱藏狀態(tài)表現(xiàn)力較差。RNN 層只能將過去的 token 存儲到固定大小的隱藏狀態(tài)中，而對于 Mamba 和 DeltaNet 等線性注意力變體來說，隱藏狀態(tài)只能是一個矩陣。要將成百上千個向量壓縮到一個只有數(shù)千級的矩陣中，這本身就是一項挑戰(zhàn)。因此，這些 RNN 層很難記住遠(yuǎn)處 token 之間的深層關(guān)系。

因此研究者嘗試使用另一種 RNN 層，其隱藏狀態(tài)本身也可以是神經(jīng)網(wǎng)絡(luò)。具體來說，研究者使用兩層 MLP，其隱藏單元比線性注意力變體中的線性（矩陣）隱藏狀態(tài)多 2 倍，非線性也更豐富。即使在測試序列上，神經(jīng)網(wǎng)絡(luò)的隱藏狀態(tài)也會通過訓(xùn)練進(jìn)行更新，這些新層被稱為測試時間訓(xùn)練層（TTT）。

局部注意力機(jī)制在湯姆的顏色、杰瑞的鼠洞之間保持了一致性，并且扭曲了湯姆的身體。

TTT-MLP 在整個視頻時長中表現(xiàn)出強(qiáng)大的特性和時間一致性。

研究者策劃了一個文本到視頻的數(shù)據(jù)集，該數(shù)據(jù)集基于大約 7 小時的《貓和老鼠》動畫片，并附有人類注釋的故事板。他們有意將范圍限制在這一特定領(lǐng)域，以便快速進(jìn)行研究迭代。作為概念驗證，該數(shù)據(jù)集強(qiáng)調(diào)復(fù)雜、多場景和具有動態(tài)運(yùn)動的長篇故事，此前的模型在這些方面仍需取得進(jìn)展；而對視覺和物理逼真度的強(qiáng)調(diào)較少，因為此前的模型在這些方面已經(jīng)取得了顯著進(jìn)展。研究者認(rèn)為，盡管本文是面向這一特定領(lǐng)域的長上下文能力改進(jìn)，但也會轉(zhuǎn)移到通用視頻生成上。

生成質(zhì)量評估

在測評中，與 Mamba 2、Gated DeltaNet 和滑動窗口注意力層等強(qiáng)大的基線相比，TTT 層生成的視頻更連貫，能講述復(fù)雜的動態(tài)故事。

在 LMSys Chatbot Arena 中，GPT-4o 比 GPT-4 Turbo 高出 29 個 Elo 分?jǐn)?shù)。

具體對比參考下列 demo：

當(dāng)湯姆咆哮并追逐杰瑞時，Mamba 2 扭曲了湯姆的外貌。

Gated DeltaNet 在湯姆的不同角度上缺乏時間一致性。

滑動窗口注意力改變了廚房環(huán)境并重復(fù)了杰瑞偷餡餅的場景。

不過，在生成內(nèi)容中，我們也能發(fā)現(xiàn) TTT-MLP 一些明顯的瑕疵。

1、時間一致性：盒子在同一場景的 3 秒片段之間變形。

2、運(yùn)動自然性：奶酪懸停在半空中，而不是自然地落到地上。

3、美學(xué)：當(dāng)湯姆轉(zhuǎn)身時，廚房里的燈光變得更加明亮。