GitHub熱榜第一：百萬token上下文，還能生成視頻，UC伯克利出品
今日GitHub熱榜榜首，是最新的開源世界模型。
上下文窗口長度達到了100萬token，持平了谷歌同時推出的王炸Gemini 1.5，伯克利出品。
強大的模型，命名也是簡單粗暴——沒有任何額外點綴，直接就叫LargeWorldModel（LWM）。
LWM支持處理多模態信息，能在100萬token中準確找到目標文本，還能一口氣看完1小時的視頻。
網友看了不禁表示，這種大海撈針般的測試，LWM能完成的如此出色，而且還開源，實在是令人印象深刻。
那么，LWM的表現到底有多強呢？
百萬上下文窗口，可看1小時視頻
在測試過程中，研究人員用多段一個多小時的視頻檢驗了LWM的長序列理解能力，這些視頻由YouTube上不同的視頻片段拼接而成。
他們將這些視頻輸入LWM，然后針對其中的細節進行提問，涉及的片段位于整個視頻的不同位置，同時研究者還將LWM與GPT-4V等模型做了對比。
結果GPT-4V是一問一個不吱聲，閉源強者Gemini Pro和開源強者Video-LLaVA都給出了錯誤的答案，只有LWM回答對了。
在另一段視頻的測試中，其他模型都說找不到有關信息，只有LWM找到了答案，而且完全正確。
不僅是理解細節，LWM也能把握視頻的整體內容，做出歸納總結。
在理解的基礎之上，LWM也可以結合自有知識進行推理，比如分析視頻中不符合常理的地方。
Benchmark測試結果顯示，LWM在MSVD-QA等三個數據集上的評分僅次于Video-LLaVA。
LWM不僅能理解長短視頻，在超長文本任務上的表現同樣優異。
在1百萬token窗口的“插針”檢索測試中，LWM取得了單針檢索全綠的成績。
多針檢索時，表現也同樣優異：
語言任務數據集的測試結果表明，LWM在32k到1M的窗口長度上表現不輸甚至超過只有4k窗口的Llama2-7B。
除了多模態信息理解，LWM還支持圖像和視頻的生成，至于效果，還是直接上圖感受一下吧。
那么，研究人員又是怎樣訓練出這樣一款世界模型的呢？
循序漸進，分而治之
LMW的訓練過程，大致可分為兩個階段。
第一階段的目標是建立一個能夠處理長文本序列的語言模型，以理解復雜的文檔和長文本內容。
為實現這一目的，研究人員采取了漸進式的訓練方式，使用總計33B Token、由圖書內容組成的Books3數據集，從32k開始訓練，逐步將窗口擴增至1M。
而為了增強LWM的長文本處理能力，開發者應用了RingAttention機制。
RingAttention是該團隊去年提出的一種窗口擴增方式，入選了ICLR 2024。
它運用了“分而治之”的思想，將長文本分成多個塊，用多個計算設備做序列并行處理，然后再進行疊加，理論上允許模型擴展到無限長的上下文。
在LWM中，RingAttention還與FlashAttention結合使用，并通過Pallas框架進行優化，從而提高性能。
在文本能力的基礎上，研究人員又用模型生成了部分QA數據，針對LWM的對話能力進行了優化。
第二階段則是將視覺信息（如圖像和視頻）整合到模型中，以提高對多模態數據的理解能力。
在此階段，研究人員對LWM-Text模型進行了架構修改，以支持視覺輸入。
他們使用VQGAN將圖像和視頻幀轉換為token，并與文本結合進行訓練。
這一階段同樣采用循序漸進的訓練方法， LWM首先在文本-圖像數據集上進行訓練，然后擴展到文本-視頻數據集，且視頻幀數逐步增多。
在訓練過程中，模型還會隨機交換文本和視覺數據的順序，以學習文本-圖像生成、圖像理解、文本-視頻生成和視頻理解等多種任務。
性能方面，研究人員在TPUv4-1024（大致相對于450塊A100）上訓練，批大小為8M、全精度（float32）的條件下，花費的時間如下表所示，其中1M窗口版本用了58個小時。
目前，LWM的代碼、模型都已開源，其中多模態模型為Jax版本，純文本模型有Jax和PyTorch兩個版本，感興趣的話可以到GitHub頁面中了解詳情。
論文地址：https://arxiv.org/abs/2402.08268。
GitHub：https://github.com/LargeWorldModel/LWM。
