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

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

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

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

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

理解世界模型

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

動(dòng)作如何影響視頻生成？作者嘗試了兩種方式：拼接動(dòng)作與視頻表示，或利用動(dòng)作對(duì)表示進(jìn)行縮放與移位。實(shí)驗(yàn)表明后者效果更好，最終被采納。

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

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

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

TinyWorlds 最終選擇自回歸，因?yàn)樗评砀?，適合實(shí)時(shí)交互，訓(xùn)練也更高效，且實(shí)現(xiàn)更簡(jiǎn)潔。

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

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

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

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

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

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

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

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

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

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

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

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