今日GitHub熱榜榜首，是最新的開源世界模型。

上下文窗口長度達到了100萬token，持平了谷歌同時推出的王炸Gemini 1.5，伯克利出品。

強大的模型，命名也是簡單粗暴——沒有任何額外點綴，直接就叫LargeWorldModel（LWM）。

LWM支持處理多模態(tài)信息，能在100萬token中準確找到目標文本，還能一口氣看完1小時的視頻。

網(wǎng)友看了不禁表示，這種大海撈針般的測試，LWM能完成的如此出色，而且還開源，實在是令人印象深刻。

那么，LWM的表現(xiàn)到底有多強呢？

百萬上下文窗口，可看1小時視頻

在測試過程中，研究人員用多段一個多小時的視頻檢驗了LWM的長序列理解能力，這些視頻由YouTube上不同的視頻片段拼接而成。

他們將這些視頻輸入LWM，然后針對其中的細節(jié)進行提問，涉及的片段位于整個視頻的不同位置，同時研究者還將LWM與GPT-4V等模型做了對比。

結(jié)果GPT-4V是一問一個不吱聲，閉源強者Gemini Pro和開源強者Video-LLaVA都給出了錯誤的答案，只有LWM回答對了。

在另一段視頻的測試中，其他模型都說找不到有關(guān)信息，只有LWM找到了答案，而且完全正確。

不僅是理解細節(jié)，LWM也能把握視頻的整體內(nèi)容，做出歸納總結(jié)。

在理解的基礎之上，LWM也可以結(jié)合自有知識進行推理，比如分析視頻中不符合常理的地方。

Benchmark測試結(jié)果顯示，LWM在MSVD-QA等三個數(shù)據(jù)集上的評分僅次于Video-LLaVA。

LWM不僅能理解長短視頻，在超長文本任務上的表現(xiàn)同樣優(yōu)異。

在1百萬token窗口的“插針”檢索測試中，LWM取得了單針檢索全綠的成績。

多針檢索時，表現(xiàn)也同樣優(yōu)異：

語言任務數(shù)據(jù)集的測試結(jié)果表明，LWM在32k到1M的窗口長度上表現(xiàn)不輸甚至超過只有4k窗口的Llama2-7B。

除了多模態(tài)信息理解，LWM還支持圖像和視頻的生成，至于效果，還是直接上圖感受一下吧。

那么，研究人員又是怎樣訓練出這樣一款世界模型的呢？

循序漸進，分而治之

LMW的訓練過程，大致可分為兩個階段。

第一階段的目標是建立一個能夠處理長文本序列的語言模型，以理解復雜的文檔和長文本內(nèi)容。

為實現(xiàn)這一目的，研究人員采取了漸進式的訓練方式，使用總計33B Token、由圖書內(nèi)容組成的Books3數(shù)據(jù)集，從32k開始訓練，逐步將窗口擴增至1M。

而為了增強LWM的長文本處理能力，開發(fā)者應用了RingAttention機制。

RingAttention是該團隊去年提出的一種窗口擴增方式，入選了ICLR 2024。

它運用了“分而治之”的思想，將長文本分成多個塊，用多個計算設備做序列并行處理，然后再進行疊加，理論上允許模型擴展到無限長的上下文。

在LWM中，RingAttention還與FlashAttention結(jié)合使用，并通過Pallas框架進行優(yōu)化，從而提高性能。

在文本能力的基礎上，研究人員又用模型生成了部分QA數(shù)據(jù)，針對LWM的對話能力進行了優(yōu)化。

第二階段則是將視覺信息（如圖像和視頻）整合到模型中，以提高對多模態(tài)數(shù)據(jù)的理解能力。

在此階段，研究人員對LWM-Text模型進行了架構(gòu)修改，以支持視覺輸入。

他們使用VQGAN將圖像和視頻幀轉(zhuǎn)換為token，并與文本結(jié)合進行訓練。

這一階段同樣采用循序漸進的訓練方法， LWM首先在文本-圖像數(shù)據(jù)集上進行訓練，然后擴展到文本-視頻數(shù)據(jù)集，且視頻幀數(shù)逐步增多。

在訓練過程中，模型還會隨機交換文本和視覺數(shù)據(jù)的順序，以學習文本-圖像生成、圖像理解、文本-視頻生成和視頻理解等多種任務。

性能方面，研究人員在TPUv4-1024（大致相對于450塊A100）上訓練，批大小為8M、全精度（float32）的條件下，花費的時間如下表所示，其中1M窗口版本用了58個小時。

目前，LWM的代碼、模型都已開源，其中多模態(tài)模型為Jax版本，純文本模型有Jax和PyTorch兩個版本，感興趣的話可以到GitHub頁面中了解詳情。

論文地址：https://arxiv.org/abs/2402.08268。

GitHub：https://github.com/LargeWorldModel/LWM。