機器之心報道

編輯：+0

還記得 DeepMind 的 Genie 3 世界模型嗎？它首次讓世界模型真實地模擬了真實世界。

最近，X 博主 anandmaj 在一個月內(nèi)復(fù)刻 Genie 3 的核心思想，開發(fā)出了TinyWorlds，一個僅300 萬參數(shù)的世界模型，能夠?qū)崟r生成可玩的像素風(fēng)格環(huán)境，包括 Pong、Sonic、Zelda 和 Doom。

帖子附帶演示視頻，展示了模型通過用戶輸入實時生成視頻幀的過程。

博主還分享了從架構(gòu)設(shè)計到訓(xùn)練細節(jié)的完整經(jīng)驗，并開源了代碼倉庫。

• 代碼：https://github.com/AlmondGod/tinyworlds

理解世界模型

世界模型是一類神經(jīng)網(wǎng)絡(luò)，它們通過生成視頻來模擬物理世界。

DeepMind 在Genie 3上展示了這一理念的潛力：當(dāng)世界模型在大規(guī)模視頻數(shù)據(jù)上訓(xùn)練時，會出現(xiàn)類似 LLM 中的「涌現(xiàn)能力」。例如：

• 可控性：按下方向鍵，鏡頭會隨之平移。
• 一致性：離開房間再返回，墻上的新油漆依舊存在。
• 質(zhì)量：水坑中的倒影清晰可見。

在 Genie 出現(xiàn)之前，研究者普遍認為要擴展世界模型，必須依賴帶動作標(biāo)注或包含三維結(jié)構(gòu)的數(shù)據(jù)。

然而 DeepMind 發(fā)現(xiàn)，只要足夠規(guī)?；赜?xùn)練原始視頻，這些高級行為便會自然涌現(xiàn)，就像語言模型會自然習(xí)得語法和句法一樣。

挑戰(zhàn)在于：世界模型的訓(xùn)練通常需要逐幀的動作標(biāo)簽（例如「按下右鍵 → 鏡頭右移」）。這意味著我們無法直接利用互聯(lián)網(wǎng)中龐大的未標(biāo)注視頻。

Genie 1 給出的解決方案是先訓(xùn)練一個動作分詞器，自動推斷幀間的動作標(biāo)簽。這樣一來，就可以把海量未標(biāo)注視頻轉(zhuǎn)化為可用的訓(xùn)練資源。

這也是 Genie 3 能夠擴展至數(shù)百萬小時 YouTube 視頻，并解鎖上述涌現(xiàn)能力的關(guān)鍵所在。

受此啟發(fā)，anandmaj 從零實現(xiàn)了一個最小化版本的世界模型：TinyWorlds。

構(gòu)建數(shù)據(jù)集

在開始訓(xùn)練 TinyWorlds 前，作者首先要決定模型能夠生成怎樣的游戲世界。模型訓(xùn)練時接觸的環(huán)境，決定了它未來的生成范圍。

因此，TinyWorlds 的數(shù)據(jù)集由處理過的 YouTube 游戲視頻構(gòu)成，包括：

• Pong：經(jīng)典的雅達利雙人游戲
• Sonic：二維橫版動作平臺
• Zelda：鳥瞰式冒險
• Pole Position：3D 像素賽車
• Doom：3D 第一人稱射擊

構(gòu)建時空變換器

與只需處理一維文本的大語言模型不同，視頻理解需要處理三維數(shù)據(jù)（高度 × 寬度 × 時間）。TinyWorlds 的核心是一個時空變換器（Space-time Transformer），它通過三層機制來捕捉視頻信息：

• 空間注意力：同一幀內(nèi)部的 token 相互關(guān)聯(lián)。
• 時間注意力：token 關(guān)注前幾個時間步的信息。
• 前饋網(wǎng)絡(luò)：token 經(jīng)過非線性處理以提取更高層次特征。

動作如何影響視頻生成？作者嘗試了兩種方式：拼接動作與視頻表示，或利用動作對表示進行縮放與移位。實驗表明后者效果更好，最終被采納。

同時，TinyWorlds 也借鑒了大語言模型的優(yōu)化技巧：SwiGLU 加速學(xué)習(xí)，RMSNorm 提升穩(wěn)定性，位置編碼則用于指示 token 在圖像中的位置。

架構(gòu)設(shè)計與分詞策略

在生成方式上，作者比較了擴散模型與自回歸模型。

TinyWorlds 最終選擇自回歸，因為它推理更快，適合實時交互，訓(xùn)練也更高效，且實現(xiàn)更簡潔。

最終架構(gòu)由三個模塊組成：

• 視頻分詞器：將視頻壓縮為 token。
• 動作分詞器：預(yù)測兩幀之間的動作。
• 動力學(xué)模型：結(jié)合歷史視頻和動作，預(yù)測未來幀。

視頻分詞器通過有限標(biāo)量量化（FSQ），將圖像劃分為立方體，并用這些立方體表示圖像塊。這樣產(chǎn)生的小 token 信息密集，減輕了動力學(xué)模型的預(yù)測負擔(dān)。

動作分詞器的任務(wù)是從原始視頻中自動生成幀間動作標(biāo)簽，使模型可以在未標(biāo)注數(shù)據(jù)上訓(xùn)練。

在訓(xùn)練初期，它容易忽略動作信號。為解決這一問題，作者引入了掩碼幀（迫使模型依賴動作）和方差損失（鼓勵編碼器覆蓋更多可能性）。

在小規(guī)模實驗中，動作 token 尚未完全映射到具體操作（如「左」「右」），但通過擴大模型或引入少量監(jiān)督標(biāo)簽，這一問題有望改善。

訓(xùn)練世界生成器

動力學(xué)模型是整個系統(tǒng)的「大腦」，負責(zé)結(jié)合視頻與動作預(yù)測未來幀。訓(xùn)練中它通過預(yù)測掩碼 token 學(xué)習(xí)時序關(guān)系，推理時則根據(jù)用戶輸入動作生成下一幀。最初由于模型過小，性能停滯且輸出模糊；擴大規(guī)模后效果顯著提升。

盡管 TinyWorlds 只有 300 萬參數(shù)，它依然能夠生成可交互的像素風(fēng)格世界：

• 駕駛《Pole Position》中的賽車
• 在《Zelda》的地圖上探索
• 進入《Doom》的 3D 地牢

雖然生成的畫面仍顯模糊、不連貫，但已經(jīng)具備可玩性。

作者認為，若擴展至千億級參數(shù)并引入擴散方法，生成質(zhì)量會有巨大提升。這正是「苦澀的教訓(xùn)」的再一次印證：規(guī)模與數(shù)據(jù)往往勝過技巧。

https://x.com/Almondgodd/status/1971314283184259336