編輯:定慧 英智
「貓和老鼠」,這部創作於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分。
借用老黃那句話,未來每一個圖元都將是生成的。
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層處理長上下文的優勢。
數據集構建
並行化與效率提升
為了提高訓練效率,研究人員對非因果序列進行了並行化處理。
他們每次更新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平均比表現第二好的方法高出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效果的評估。
評估基於盲測比較(blind comparisons)中的兩兩比較偏好(pairwise preferences)。
因為直接對長視頻進行評分或一次性對許多視頻進行排序是很有挑戰性的。
具體來說,評估者會隨機獲得前述四個評估維度中的一個,以及一對共用相同故事情節的視頻,然後被要求指出在該評估維度 上哪個視頻更好。
為了收集用於評估的視頻池,研究人員首先使用Claude 3.7 Sonnet採樣了100個故事情節,然後每種方法為每個故事情節生成一個視頻。
生成視頻的方法對評估者始終是未知的。
使用Chatbot Arena中的Elo評分系統(Elo system) 來匯總兩兩比較的偏好。
TTT-MLP方法比第二好的方法平均高出34 Elo分(可回顧上述表格)。
當然,TTT-MLP生成視頻中也會有「瑕疵」。
左右滑動查看
華人共同一作
本文有4名共同一作,其中一位為華人Jiarui Xu,UCSD的五年級博士生。
在香港科技大學獲得了計算機科學學士學位。現在是FAIR實驗室的研究實習生。
之前曾在Google Research、NVIDIA Research、Microsoft Research和OpenMMLab實習。