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