英偉達華人AI版《貓和老鼠》爆火全網,60秒經典淚目!僅靠Transformer新增一層
「貓和老鼠」,這部創作于1940年的動畫,現在已經85歲了。
如今,只需一個提示,不用任何剪輯,AI便可「無限續杯」童年經典,繼續延續它的魅力。
你沒有看過的《貓和老鼠》情節畫面,由AI制作
來自英偉達、斯坦福、UCSD、UC伯克利以及UT Austin研究團隊,通過Test-Time Training構建了一個「一分鐘視頻」生成器。
只需在預訓練Transformer中嵌入TTT層 ,僅5B參數模型可理解復雜提示,生成長達1分鐘視頻。
圖片
論文地址:https://test-time-training.github.io/video-dit/
基于《貓和老鼠》的動畫進行測試后,相較于Mamba2等基線,TTT生成視頻Elo要領先34分。
借用老黃那句話,未來每一個像素都將是生成的。
圖片
首批1分鐘無剪輯AI版「貓和老鼠」
用下面的故事概要,生成「全新的一集」貓和老鼠。
杰瑞高興地在整潔的廚房里吃奶酪,直到湯姆頑皮地把奶酪拿走,逗他。生氣的杰瑞收拾好行李離開了家,拖著一個小行李箱。后來,湯姆注意到杰瑞不在了,感到難過,并沿著杰瑞的小腳印一路追蹤到舊金山。杰瑞無精打采地坐在一條小巷里,湯姆找到了他,溫柔地用奶酪作為道歉。杰瑞原諒了湯姆,接受了奶酪,兩人一起回家,友誼得到了恢復。
在一個陽光明媚的紐約早晨,湯姆,一只攜帶公文包的藍灰色貓,來到了他在世界貿易中心的辦公室。當他安頓下來時,他的電腦突然關機了——杰瑞,一只淘氣的棕色老鼠,咬斷了電纜。一場追逐開始了,最終湯姆撞到了墻上,而杰瑞逃進了自己的鼠洞。湯姆決心滿滿地沖進了一個辦公室的門,不小心打斷了由憤怒的斗牛犬斯派克主持的會議,斯派克生氣地把他趕走了。在舒適的鼠洞里,杰瑞對這場混亂大笑起來。
湯姆正在廚房的桌子旁高興地吃著蘋果派。杰瑞眼巴巴地看著,希望自己也能吃一些。杰瑞走到房子的前門外面按了門鈴。當湯姆來開門時,杰瑞繞到房子后面跑進了廚房。杰瑞偷走了湯姆的蘋果派。杰瑞拿著派跑向他的鼠洞,而湯姆在后面追趕他。就在湯姆快要抓到杰瑞的時候,杰瑞成功鉆進了鼠洞,湯姆撞到了墻上。
在一次水下冒險中,杰瑞找到了一張藏寶圖,并在躲避湯姆的追逐中穿過珊瑚礁和海帶森林尋找寶藏。杰瑞最終在一處沉船內發現了寶藏,欣喜若狂地慶祝,而湯姆的追逐卻讓他遇到了一條饑餓的鯊魚,陷入了麻煩。
湯姆和杰瑞參觀了一個熱鬧的嘉年華,湯姆急切地嘗試了一個投球游戲,但反復失手,這讓杰瑞嘲笑他。感到沮喪的湯姆變得非常堅定,但盡管他盡了最大努力,還是沒能擊倒罐子。杰瑞自信地走上前去,輕松成功,擊倒了罐子并贏得了一座閃亮的金質獎杯。杰瑞高興地慶祝他的勝利,而湯姆則驚訝、嫉妒且生氣,當他們一起離開嘉年華時——杰瑞自豪地拿著他的獎品,而湯姆則悶悶不樂地跟在后面。
TTT層:視頻生成的新希望
TTT層的核心思想,是對RNN層隱藏狀態更新方式進行了創新。
TTT層的隱藏狀態不再是簡單的矩陣,而是神經網絡,具體來說是兩層MLP。
它比線性注意力變體里的線性隱藏狀態多了2倍的隱藏單元,表達能力大大增強。
圖片
TTT層通過自監督學習壓縮歷史上下文,把這些信息變成隱藏狀態的權重。
在這個過程中,有個很重要的更新規則,就是對自監督損失進行梯度下降:
其中是學習率。
輸出標記z_t是模型根據更新后的權重對x_t的預測,也就是
圖片
為了讓學習過程更具挑戰性,研究人員會先把x_t處理成損壞的輸入然后優化損失函數
這和去噪自編碼器有點類似,都需要挖掘x_t各個維度之間的相關性,才能從部分信息中重建它。
在實際操作過程中,TTT層還有一些關鍵設計。學習自監督任務時,不再是手動設計任務,而是采用端到端的學習方式。
TTT層和其他RNN層、自注意力層的接口是一樣的,可以很方便地替換到各種網絡架構里。
基于TTT層的視頻生成方法
接下來看看研究人員是如何用TTT層生成視頻的。
他們用預訓練的CogVideo-X 5B作為基礎模型,在這個模型里加上TTT層,再進行微調。
架構調整
在架構設計上,有幾個很巧妙的地方,首先是門控機制。
TTT層是隨機初始化的,如果直接加到預訓練網絡里,會讓模型的預測性能大幅下降。
研究人員用一個可學習向量來控制TTT層的輸出,公式為
初始值設置為0.1,這樣在微調初期,TTT層既能發揮作用,又不會過度影響原來的模型。
其次是雙向機制。
由于擴散模型是非因果的,為了讓TTT層更好地工作,研究人員通過
讓TTT層可以逆序掃描輸入序列,輸出結果依然是按時間順序排列的。
他們還對整個架構的序列建模塊進行了修改,加上了TTT層和門控
讓模型更好地處理長視頻。
圖片
整體流程優化
研究人員把視頻分成多個場景,每個場景又包含一個或多個3秒的片段。
這么做是因為原來的CogVideo-X的最大生成長度是3秒,而且《貓和老鼠》里大多數場景都至少有3秒,用3秒片段構建數據集也更方便。
文本提示有三種格式:格式1是用5-8句話簡單概括劇情;格式2是用大約20句話詳細描述,每句話對應一個3秒片段;格式3是分鏡腳本,每個3秒片段都用3-5句話描述,包含了各種細節。
在實際處理的時候,輸入到文本分詞器的都是格式3,格式之間的轉換由Claude 3.7 Sonnet完成。
從文本到序列的轉換過程也很有意思。
CogVideo-X會把文本標記和帶噪聲的視頻標記連接起來,形成輸入序列。生成長視頻的時候,就對每個3秒片段分別處理,然后把這些片段連接起來。
在處理過程中,自注意力層只在每個3秒片段內局部處理,而TTT層則對整個輸入序列進行全局處理。
既控制了計算成本,又能充分發揮TTT層處理長上下文的優勢。
數據集構建
研究人員基于1940年到1948年的81集《貓和老鼠》建了數據集。
他們先對原始視頻進行超分辨率處理,提升畫質,將分辨率統一為720×480。
然后讓注釋者把劇集分解成場景,提取3秒片段,并為每個片段編寫詳細的腳本。
為了訓練不同時長的視頻,他們還把3秒片段連接起來,形成9秒、18秒、30秒和63秒的視頻。
并行化與效率提升
為了提高訓練效率,研究人員對非因果序列進行了并行化處理。
他們每次更新b個標記的W(這里b=64),公式是
然后用W_ib生成小批量i中所有時間步的輸出標記。
這樣不僅實現了并行計算,還通過平均梯度減少了方差,讓訓練過程更加穩定。
由于TTT-MLP的隱藏狀態太大,無法在單個SM的SMEM中存儲,研究人員使用了片上張量并行技術。
把W^(1)和W^(2)在多個SM之間分片存儲,就好像把一個大任務拆分成小份,讓多個助手一起幫忙。
圖片
研究人員用NVIDIA Hopper GPU架構的DSMEM特性實現SM之間的全規約操作,還采用了多階段流水線、異步預取等技術,減少數據傳輸時間,進一步提高效率。
一分鐘視頻效果評估
以如下故事概要作為視頻生成詞。
湯姆正高高興興地在廚房桌旁吃著蘋果派。杰瑞則滿懷渴望,羨慕地看著,希望自己也能嘗一口。杰瑞走到屋前,按響了門鈴。當湯姆去開門時,杰瑞繞到了后門進入廚房。杰瑞偷走了湯姆的蘋果派,拿著蘋果派快步跑向他的鼠洞,而湯姆緊追不舍。就在湯姆即將抓住杰瑞的時候,杰瑞順利鉆進了鼠洞,湯姆則撞到了墻上。
視頻幀對比了TTT-MLP與Gated DeltaNet以及滑動窗口注意力——這些都是人工評估中表現領先的基線方法。
TTT-MLP通過在場景轉換過程中保留細節展現了更好的場景一致性,并且通過準確描繪復雜動作實現了更自然的運動效果。
圖片
作為對比:
- TTT-MLP:在場景變化和不同角度之間保持時間一致性,生成流暢、高質量的動作。
- 滑動窗口注意力:改變了廚房環境,改變了房屋顏色,并復制了杰里偷派的情節。
- 門控 DeltaNet:在不同角度的Tom中缺乏時間一致性,但在后續幀中保持了廚房環境。
- Mamba2:扭曲了湯姆的外觀,當他咆哮并追逐杰瑞時,但整個視頻中保持了相似的廚房環境。
圖片
一分鐘視頻的人工評估結果顯示,TTT-MLP平均比表現第二好的方法高出34 Elo分(1033 VS 999)。
在提升最多的評估維度上,場景一致性提高了38分,動作平滑性提高了39分。
作為對比,在Chatbot Arena中,GPT-4比GPT-3.5 Turbo高出46 Elo分,而GPT-4o比GPT-4Turbo 高出29 Elo分。
圖片
對于63秒的視頻,全注意力(處理超30萬個token)的推理時間比局部注意力長11倍,而訓練時間長12倍。
TTT-MLP分別需要2.5倍和3.8倍——雖然比全注意力高效得多,但仍不及,例如Gated DeltaNet,其在推理和訓練上均只比局部注意力長1.8倍。
評估維度與評估方法
從MovieGen提出的六個評估維度中,選取了其中四個與研究領域相關的,用于TTT效果的評估。
- 文本一致性(Text following):與所提供提示的對齊度 ,指生成內容與輸入提示或要求的匹配程度。
- 動作自然度(Motion naturalness): 自然的肢體動作、面部表情,以及對物理定律的遵守。看起來不自然或怪異的動作將被扣分。
- 美學質量(Aesthetics):有趣且引人入勝的內容、燈光、色彩以及鏡頭效果。
- 時間一致性(Temporal consistency): 場景內部以及跨場景的一致性。
評估基于盲測比較(blind comparisons)中的兩兩比較偏好(pairwise preferences)。
因為直接對長視頻進行評分或一次性對許多視頻進行排序是很有挑戰性的。
具體來說,評估者會隨機獲得前述四個評估維度中的一個,以及一對共享相同故事情節的視頻,然后被要求指出在該評估維度 上哪個視頻更好。
為了收集用于評估的視頻池,研究人員首先使用Claude 3.7 Sonnet采樣了100個故事情節,然后每種方法為每個故事情節生成一個視頻。
生成視頻的方法對評估者始終是未知的。
使用Chatbot Arena中的Elo評分系統(Elo system) 來匯總兩兩比較的偏好。
TTT-MLP方法比第二好的方法平均高出34 Elo分(可回顧上述表格)。
圖片
當然,TTT-MLP生成視頻中也會有「瑕疵」。
- 時間一致性 (Temporal consistency): 物體有時會在3秒片段的邊界處變形,這可能是因為擴散模型 在不同片段間從不同的模式 (modes) 中采樣。
- 動作自然度 (Motion naturalness): 物體有時會不自然地漂浮,因為重力效果未能被正確建模。
- 美學質量 (Aesthetics): 除非在提示中明確要求,否則光線變化與動作并不總能保持一致。復雜的攝像機運鏡(例如視差效果 parallax)有時會被不準確地描繪。
圖片
圖片
圖片
華人共同一作
圖片
本文有4名共同一作,其中一位為華人Jiarui Xu,UCSD的五年級博士生。
在香港科技大學獲得了計算機科學學士學位。現在是FAIR實驗室的研究實習生。
之前曾在Google Research、NVIDIA Research、Microsoft Research和OpenMMLab實習。
圖片
參考資料: