編輯:杜偉、蛋醬
AI 圈永遠不缺「新活」。
這兩天,加州大學伯克利分校、斯坦福大學、英偉達等機構聯合制作的《貓和老鼠》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、美學:當湯姆轉身時,廚房裡的燈光變得更加明亮。


© THE END
轉載請聯絡本公眾號獲得授權
投稿或尋求報道:[email protected]