AI 也能生成逼真的《貓和老鼠》動(dòng)畫片了！長(zhǎng)度可以達(dá)到 63 秒，劇情也可以重新生成。

（來源：資料圖）

當(dāng)?shù)貢r(shí)間 4 月 8 日，來自英偉達(dá)、美國斯坦福大學(xué)、加州大學(xué)圣地亞哥分校、加州大學(xué)伯克利分校、得克薩斯大學(xué)奧斯汀分校的研究人員，基于測(cè)試時(shí)訓(xùn)練（TTT，Test-Time Training）生成了多個(gè)《貓和老鼠》的動(dòng)畫視頻。

他們?cè)陬A(yù)訓(xùn)練的 Transformer 模型中添加了測(cè)試時(shí)訓(xùn)練層并對(duì)其進(jìn)行微調(diào)，借此生成了具有強(qiáng)時(shí)間一致性的《貓和老鼠》一分鐘動(dòng)畫片。

研究人員表示，每個(gè)視頻都是由模型一次性直接生成的，沒有經(jīng)過編輯、拼接或后處理，每個(gè)故事也都是全新創(chuàng)作的。

那么，上述視頻到底是怎么生成的？研究人員表示，如今的 Transformer 模型在生成一分鐘視頻方面仍然面臨挑戰(zhàn)，因?yàn)樽宰⒁饬釉谔幚黹L(zhǎng)上下文時(shí)效率低下。諸如 Mamba 層之類的替代方案在處理復(fù)雜的多場(chǎng)景故事時(shí)表現(xiàn)不佳，因?yàn)樗鼈兊碾[藏狀態(tài)表現(xiàn)力較弱。

為此，他們使用了測(cè)試時(shí)訓(xùn)練層，其隱藏狀態(tài)本身可以是神經(jīng)網(wǎng)絡(luò)，從而具備更強(qiáng)的表達(dá)能力。當(dāng)在預(yù)訓(xùn)練的 Transformer 中加入測(cè)試時(shí)訓(xùn)練層，使其能夠根據(jù)文本分鏡腳本生成一分鐘的視頻。

研究中，他們從一個(gè)預(yù)訓(xùn)練的擴(kuò)散 Transformer（CogVideo-X 5B）開始，原本其只能以 16 幀每秒的速度生成 3 秒的短片段，或以 8 幀每秒的速度生成 6 秒的短片段。

但是，研究人員添加了從頭初始化的測(cè)試時(shí)訓(xùn)練層，并對(duì)該模型進(jìn)行微調(diào)，以便根據(jù)文本分鏡腳本生成一分鐘的視頻。

他們將自注意力層限制在 3 秒的片段內(nèi)，以便將研究成本保持在可控范圍內(nèi)。整個(gè)訓(xùn)練過程運(yùn)行在 256 個(gè)英偉達(dá) H100 上，相當(dāng)于消耗了 50 小時(shí)的訓(xùn)練時(shí)長(zhǎng)。

為了快速進(jìn)行迭代，研究人員基于約 7 小時(shí)的《貓和老鼠》動(dòng)畫片，并輔以人工標(biāo)注的分鏡腳本，整理了一個(gè)文本到視頻的數(shù)據(jù)集。

作為概念驗(yàn)證，本次研究中的數(shù)據(jù)集側(cè)重于復(fù)雜、多場(chǎng)景、長(zhǎng)鏡頭且充滿動(dòng)態(tài)動(dòng)作的敘事。而對(duì)于視覺和物理真實(shí)感方面，此前人們已經(jīng)取得顯著進(jìn)展，所以本次研究未作重點(diǎn)強(qiáng)調(diào)。

（來源：資料圖）

從宏觀層面來看，本次方法只是在預(yù)訓(xùn)練的擴(kuò)散 Transformer（Diffusion Transformer）上添加了測(cè)試時(shí)訓(xùn)練層，并在帶有文本注釋的長(zhǎng)視頻上進(jìn)行微調(diào)。

他們采用了預(yù)先添加測(cè)試時(shí)訓(xùn)練層、然后再進(jìn)行微調(diào)的方法，這種方法原則上適用于任何主干架構(gòu)。

研究中，他們選擇擴(kuò)散 Transformer 作為初步演示，因?yàn)樗且曨l生成中最流行的架構(gòu)。由于在視頻上預(yù)訓(xùn)練擴(kuò)散 Transformer 的成本過高，所以研究人員從 CogVideo-X 5B 的預(yù)訓(xùn)練檢查點(diǎn)開始。

他們將視頻設(shè)計(jì)為包含多個(gè)場(chǎng)景，每個(gè)場(chǎng)景包含一個(gè)或多個(gè) 3 秒的片段，并采用 3 秒片段作為文本與視頻配對(duì)的最小單位，之所以這樣做的原因有三個(gè)：

第一，原始預(yù)訓(xùn)練的 CogVideo-X 的最大生成長(zhǎng)度為 3 秒。第二，在《貓和老鼠》的劇集中，大多數(shù)場(chǎng)景的時(shí)長(zhǎng)至少為 3 秒。第三，以 3 秒為一段構(gòu)建多階段數(shù)據(jù)集最為方便。

在推理階段，研究人員按照以下三種格式中的任何一種，以詳細(xì)程度遞增的順序來編寫長(zhǎng)視頻的文本提示。

? 格式 1：用 5-8 句話簡(jiǎn)短概述情節(jié)。

? 格式 2：用大約 20 句話詳細(xì)描述情節(jié)，每句話大致對(duì)應(yīng) 3 秒的片段。

? 格式 3：分鏡腳本。每一個(gè) 3 秒鐘的片段都由一段 3 到 5 句話來描述，內(nèi)容包含背景顏色和鏡頭移動(dòng)等細(xì)節(jié)。一個(gè)或多個(gè)段落組成的腳本組，會(huì)被嚴(yán)格地界定為屬于某些場(chǎng)景，并且要使用關(guān)鍵詞“<場(chǎng)景開始>”和“< 場(chǎng)景結(jié)束 >”。

在微調(diào)和推理中，文本分詞器的實(shí)際輸入始終采用格式 3，格式之間的轉(zhuǎn)換由 Claude 3.7 Sonnet 按 1→2→3 的順序執(zhí)行。

為了生成時(shí)間較長(zhǎng)的視頻，研究人員針對(duì)每個(gè) 3 秒的片段都獨(dú)立使用相同的程序。

具體而言，給定一個(gè)包含 n 段落的格式 3 的分鏡腳本，首先生成 n 個(gè)序列片段，每個(gè)片段都包含從相應(yīng)段落提取的文本標(biāo)記以及后續(xù)的視頻標(biāo)記。

然后，將所有 n 個(gè)序列片段連接在一起，從而形成輸入序列，這時(shí)序列中包含交錯(cuò)排列的文本 tokens 和視頻 tokens。

對(duì)于 CogVideo-X 來說，它采用自注意力層來針對(duì)每段最長(zhǎng) 3 秒的視頻進(jìn)行全局序列處理。但是，面對(duì)長(zhǎng)視頻時(shí)，其所采用的全局注意力機(jī)制會(huì)顯著降低效率。

為了避免增加自注意力層的上下文長(zhǎng)度，研究人員將其處理范圍限定在每個(gè) 3 秒片段內(nèi)，使各 n 個(gè)序列片段能夠獨(dú)立進(jìn)行注意力計(jì)算。

由于測(cè)試時(shí)訓(xùn)練層能夠高效地處理長(zhǎng)上下文序列，因此他們采用全局方式來處理整個(gè)輸入序列。

遵循大語言模型的標(biāo)準(zhǔn)做法，研究人員將改進(jìn)架構(gòu)的上下文長(zhǎng)度分為五個(gè)階段，并將其擴(kuò)展至一分鐘。

具體來說，他們先在《貓和老鼠》的 3 秒片段上針對(duì)整個(gè)預(yù)訓(xùn)練模型進(jìn)行微調(diào)。

在此階段，新的參數(shù)特別是測(cè)試時(shí)訓(xùn)練層和門控機(jī)制中的的參數(shù)，會(huì)被賦予更高的學(xué)習(xí)率。

接下來，研究人員分別對(duì) 9 秒、18 秒、30 秒以及最終的 63 秒的視頻進(jìn)行微調(diào)。

為了避免模型在預(yù)訓(xùn)練過程中遺忘過多的世界知識(shí)，研究人員僅對(duì)測(cè)試時(shí)訓(xùn)練層、門控層和自注意力層進(jìn)行微調(diào)，且在這四個(gè)階段中使用較低的學(xué)習(xí)率。

在原始視頻的選取上，他們選擇了從 1940 年至 1948 年間發(fā)行的 81 集《貓和老鼠》，每集時(shí)長(zhǎng)約 5 分鐘，所有集數(shù)加起來時(shí)長(zhǎng)約 7 小時(shí)。

然后，研究人員在原始視頻上運(yùn)行一個(gè)視頻超分辨率模型，從而為數(shù)據(jù)集生成視覺效果更強(qiáng)的視頻，這些視頻的分辨率均為 720×480。

接著，他們讓人工標(biāo)注員將每個(gè)片段分解為場(chǎng)景，然后從每個(gè)場(chǎng)景中提取 3 秒長(zhǎng)的片段。

接下來，研究人員讓人工標(biāo)注員為每個(gè) 3 秒的片段撰寫一段詳細(xì)的描述，然后直接針對(duì)這些片段進(jìn)行微調(diào)。

為了創(chuàng)建數(shù)據(jù)，研究人員將連續(xù)的 3 秒片段拼接成 9 秒、18 秒、30 秒和 63 秒的視頻，并附上相應(yīng)的文本注釋，所有訓(xùn)練視頻的標(biāo)注均采用上文的格式 3。

（來源：資料圖）

對(duì)于 GPU 而言，要想高效地實(shí)現(xiàn)測(cè)試訓(xùn)練層-多層感知器（TTT-MLP，Test-Time Training-Multi-Layer Perceptron），就需要進(jìn)行特殊設(shè)計(jì)，以便利用其內(nèi)存層次結(jié)構(gòu)。

英偉達(dá) GPU 架構(gòu)中的核心計(jì)算單元是流式多處理器（SMs，Streaming Multiprocessors），其功能類比于 CPU 中的單個(gè)核心。

GPU 上的所有流式多處理器共享一個(gè)相對(duì)較慢但容量較大的全局內(nèi)存（即 HBM，High Bandwidth Memory），然后每個(gè)流式多處理器都有一個(gè)快速但容量較小的片上內(nèi)存（即 SMEM，Shared Memory）。

GPU 上 SMEM 與 HBM 之間的頻繁數(shù)據(jù)傳輸會(huì)顯著降低整體效率。而 Mamba 和自注意力層通過利用內(nèi)核融合技術(shù)，可以減少這類數(shù)據(jù)的傳輸。

其核心思想是將輸入和初始狀態(tài)加載到每個(gè) SMEM 中，完全在片上進(jìn)行計(jì)算，并且只將最終輸出寫回 HBM。

然而，TTT-MLP 的隱藏狀態(tài)即雙層 MLP 函數(shù) f 的權(quán)重矩陣 W(1) 和 W(2)，由于體積過大無法存儲(chǔ)于單個(gè)流式多處理器的共享內(nèi)存中。

為了減少每個(gè)流式多處理器所需的內(nèi)存，研究人員使用張量并行機(jī)制，將 W(1) 和 W(2) 在流式多處理器之間進(jìn)行分片。

（來源：資料圖）

由于大型多層感知器層可以被分片并跨多個(gè) GPU 的高帶寬存儲(chǔ)器上進(jìn)行訓(xùn)練，因此研究人員將同樣的思路用于多個(gè)流式多處理器的共享內(nèi)存中，將每個(gè)流式多處理器視為一個(gè) GPU 的類比。

研究人員利用英偉達(dá) Hopper GPU 架構(gòu)的分布式共享內(nèi)存特性，實(shí)現(xiàn)了流式多處理器間的全局歸約操作，進(jìn)而顯著提升了效率。

作為一種通用原則，如果一個(gè)模型架構(gòu) f 可以通過標(biāo)準(zhǔn)張量并行機(jī)制在 GPU 之間進(jìn)行分片，那么當(dāng) f 用作隱藏狀態(tài)時(shí)，同樣的分片策略也可以用于流式多處理器之間。

（來源：資料圖）

不過，本次研究的 TTT-MLP 內(nèi)核受限于寄存器溢出和異步指令的次優(yōu)排序。未來，通過降低寄存器壓力以及開發(fā)編譯器感知更強(qiáng)的異步操作方案，其執(zhí)行效率或能得到進(jìn)一步提升。總的來說，本次方法有望用于生成更長(zhǎng)、更復(fù)雜視頻，也許下一代兒童將能看上由 AI 生成的動(dòng)畫連續(xù)劇？目前看來，這一設(shè)想并非沒有實(shí)現(xiàn)的可能。

參考資料：

https://test-time-training.github.io/video-dit/

https://x.com/arankomatsuzaki/status/1909336661743550555

運(yùn)營(yíng)/排版：何晨龍