轉(zhuǎn)身世界就變樣?WorldMem用記憶讓AI生成的世界擁有了一致性
本文一作為肖澤琪,本科畢業(yè)于浙江大學(xué),現(xiàn)為南洋理工大學(xué)博士生,研究方向是基于視頻生成模型的世界生成和模擬,導(dǎo)師為潘新鋼。個(gè)人主頁(yè):
https://xizaoqu.github.io
近年來,基于視頻生成模型的可交互世界生成引發(fā)了廣泛關(guān)注。盡管現(xiàn)有方法在生成質(zhì)量和交互能力上取得了顯著進(jìn)展,但由于上下文時(shí)間窗口受限,生成的世界在長(zhǎng)時(shí)序下嚴(yán)重缺乏一致性。
針對(duì)這一問題,南洋理工大學(xué) S-Lab、北京大學(xué)與上海 AI Lab 的研究者提出了創(chuàng)新性的世界生成模型——WorldMem,通過引入記憶機(jī)制,實(shí)現(xiàn)了長(zhǎng)時(shí)序一致的世界生成。
WorldMem 在 Minecraft 數(shù)據(jù)集上進(jìn)行了大規(guī)模訓(xùn)練,支持在多樣化場(chǎng)景中自由探索和動(dòng)態(tài)變化,并在真實(shí)數(shù)據(jù)集上驗(yàn)證了方法的可行性。
- 論文名稱:WorldMem: Long-term Consistent World Simulation with Memory
- 項(xiàng)目主頁(yè): https://xizaoqu.github.io/worldmem
- 論文代碼:https://github.com/xizaoqu/WorldMem
- Demo:https://huggingface.co/spaces/yslan/worldmem
研究背景
世界生成模型在近期受到了廣泛關(guān)注,如谷歌的 Genie 2 [1]、阿里的 The Matrix [2]、Meta 的 Navigation World Models [4] 等。這些方法在生成質(zhì)量與交互性方面取得了顯著進(jìn)展,但長(zhǎng)時(shí)一致性問題仍未得到有效解決。
舉例:當(dāng)我們控制視角先向右轉(zhuǎn),再向左轉(zhuǎn)。
在傳統(tǒng)方法中,回看時(shí)場(chǎng)景內(nèi)容會(huì)發(fā)生顯著變化。
在 WorldMem 中,我們?cè)谑澜缟赡P椭幸胗洃洐C(jī)制,實(shí)現(xiàn)了一致的世界生成。
方法效果
WorldMem 通過引入記憶機(jī)制,實(shí)現(xiàn)了長(zhǎng)時(shí)序下世界生成的一致性。智能體可在廣闊的動(dòng)作空間中探索多樣場(chǎng)景,生成結(jié)果在視角和位置變化后仍保持良好的幾何一致性。
同時(shí),WorldMem 還支持時(shí)間一致性建模。比如在雪地中放置南瓜燈,隨著時(shí)間推移,模型不僅保留該物體,還能生成其逐漸融化周圍積雪的細(xì)節(jié),體現(xiàn)真實(shí)的事件演化過程。
方法
WorldMem 模型的主要結(jié)構(gòu)如下圖所示,包含三大核心模塊:
- 條件生成模塊
- 記憶讀寫模塊
- 記憶融合模塊
條件生成模塊——支持交互與持續(xù)生成的條件視頻生成主干
我們基于 Oasis [5] 和 Conditional DiT [6] 構(gòu)建了世界生成基模型,并采用 Diffusion Forcing [3] 訓(xùn)練策略,使模型能在有限上下文下實(shí)現(xiàn)自回歸式長(zhǎng)時(shí)生成。
盡管擴(kuò)散模型結(jié)合自回歸訓(xùn)練具備一定的長(zhǎng)時(shí)生成能力,但仍受限于上下文窗口,易出現(xiàn)遺忘問題,導(dǎo)致生成內(nèi)容逐漸失真。為此,我們引入記憶機(jī)制,以增強(qiáng)模型的長(zhǎng)期一致性。
記憶讀寫模塊——負(fù)責(zé)歷史信息的存取與精準(zhǔn)檢索
- 記憶庫(kù)(Memory Bank):構(gòu)建生成的長(zhǎng)期記憶
為緩解上下文窗口帶來的遺忘問題,我們引入了記憶機(jī)制,作為一個(gè)持續(xù)更新的外部緩沖區(qū),幫助模型回顧過去,保持場(chǎng)景在時(shí)間上的連續(xù)性。
我們?cè)O(shè)計(jì)的記憶庫(kù)用于存儲(chǔ)生成過程中的關(guān)鍵歷史信息。每個(gè)記憶單元包含圖像幀及其對(duì)應(yīng)的狀態(tài)(視角位姿與時(shí)間戳)。隨著生成推進(jìn),記憶庫(kù)不斷積累,構(gòu)建起一套可檢索的時(shí)間記錄。
- 記憶檢索(Memory Retrieve):高效選出最相關(guān)的歷史幀
由于每次生成僅能參考少量歷史幀,我們?cè)O(shè)計(jì)了一種貪心匹配算法,從龐大的記憶庫(kù)中高效篩選關(guān)鍵信息:
- 計(jì)算相似度(基于視野重疊與時(shí)間差異);
- 選取與當(dāng)前場(chǎng)景最接近的記憶單元;
- 過濾冗余,確保選出的記憶代表性強(qiáng)、信息多樣。
這一過程不僅提升了生成效率,也保障了歷史信息的有效利用。
記憶融合模塊——融合關(guān)鍵歷史幀,引導(dǎo)當(dāng)前生成
在長(zhǎng)時(shí)序視頻生成中,僅依賴當(dāng)前幀難以維持場(chǎng)景一致性。我們引入記憶融合模塊,通過融合關(guān)鍵歷史幀,引導(dǎo)當(dāng)前生成,使模型在視角或場(chǎng)景變化后,仍能還原先前內(nèi)容。
不同于 StreamingT2V [7]、SlowFast [8] 等方法主要依賴高層語義特征,我們更關(guān)注細(xì)節(jié)重建與空間一致性,因此需要更精細(xì)的歷史關(guān)聯(lián)機(jī)制。
- Memory Attention:連接過去與現(xiàn)在
我們采用跨注意力機(jī)制,實(shí)現(xiàn)當(dāng)前幀與歷史幀的動(dòng)態(tài)交互:
- 為當(dāng)前幀和記憶幀添加狀態(tài)嵌入(位姿 + 時(shí)間);
- 通過注意力計(jì)算,提取與當(dāng)前場(chǎng)景最相關(guān)的記憶信息,生成融合特征用于引導(dǎo)生成。
這種方式實(shí)現(xiàn)了歷史信息的高效利用與精準(zhǔn)檢索,顯著增強(qiáng)生成的一致性。
- 狀態(tài)嵌入設(shè)計(jì):精細(xì)表達(dá)空間與時(shí)間
為提升融合效果,我們?cè)O(shè)計(jì)了兩類嵌入:
- 位姿嵌入:采用 Plücker 坐標(biāo)表達(dá)空間位置;
- 時(shí)間嵌入:使用 MLP 映射時(shí)間戳。
二者相加構(gòu)成最終狀態(tài)特征。此外,我們引入相對(duì)嵌入機(jī)制:
- 查詢幀使用零向量,僅依賴記憶幀的相對(duì)狀態(tài);
- 并采用幀獨(dú)立檢索策略,確保每幀都能單獨(dú)提取最相關(guān)歷史信息。
整體上,記憶融合模塊顯著提升了模型的空間理解與細(xì)節(jié)保持能力,是實(shí)現(xiàn)穩(wěn)定、連貫世界生成的關(guān)鍵組成部分。
實(shí)驗(yàn)
在 Minecraft 上的結(jié)果
我們?cè)?Minecraft 基準(zhǔn)測(cè)試中評(píng)估了所提方法,結(jié)果顯示:
- 在短時(shí)生成中,傳統(tǒng)方法易出現(xiàn)一致性問題,而引入記憶機(jī)制后,生成質(zhì)量與一致性明顯提升;
- 在長(zhǎng)時(shí)推理中,傳統(tǒng)方法性能顯著下降,而我們的方法在各項(xiàng)指標(biāo)上保持優(yōu)勢(shì),展現(xiàn)出良好的長(zhǎng)期穩(wěn)定性。
- 長(zhǎng)時(shí)序生成對(duì)比
下圖展示了不同消融設(shè)置下,模型在 300 幀序列上的 PSNR 變化。結(jié)果表明:
- 缺少記憶模塊或采用隨機(jī)檢索的方法,在短時(shí)間內(nèi)即出現(xiàn)一致性下降;
- 缺少相對(duì)位置編碼的模型,在 100 幀后性能明顯退化;
- 完整方法在 300 幀以上仍保持穩(wěn)定一致性,展現(xiàn)出優(yōu)越的長(zhǎng)期建模能力。
- 可視化結(jié)果
與仿真數(shù)據(jù)(Ground Truth)相比,WorldMem 能夠基于記憶條件輸入,準(zhǔn)確建模世界場(chǎng)景,同時(shí)支持動(dòng)態(tài)變化(如降雨),并保持良好的時(shí)間一致性。
WorldMem 支持與生成世界的交互,例如在沙漠中放置干草堆或在草原上種植作物。這些事件會(huì)被寫入記憶庫(kù),并影響后續(xù)生成。當(dāng)用戶回到曾種植作物的位置時(shí),不僅能看到作物仍在,還能觀察其從發(fā)芽到生長(zhǎng)的過程,體現(xiàn)出模型對(duì)時(shí)間動(dòng)態(tài)的建模能力。
- 在真實(shí)場(chǎng)景上的結(jié)果
我們同時(shí)也將我們的方法在真實(shí)場(chǎng)景數(shù)據(jù)上做了驗(yàn)證,結(jié)果顯示,加入記憶機(jī)制后,我們的方法能夠增強(qiáng)真實(shí)世界生成的一致性。
更多定制軌跡的結(jié)果:
更多細(xì)節(jié)請(qǐng)參考我們的論文與項(xiàng)目主頁(yè)。
展望
近年來,視頻生成模型(如 WAN 2.1 [9]、Hunyuan [10] 等)展現(xiàn)出驚人的世界生成與仿真能力,驗(yàn)證了其在理解與生成復(fù)雜環(huán)境中的潛力。
我們相信,未來交互式視頻生成模型將在虛擬仿真、交互智能等領(lǐng)域發(fā)揮越來越重要的作用。
WorldMem 為世界一致性建模邁出了關(guān)鍵一步,隨著技術(shù)發(fā)展,視頻生成模型有望成為構(gòu)建真實(shí)、持久、交互式虛擬世界的核心引擎。