首個開源世界模型！百萬級上下文，長視頻理解吊打GPT-4，UC伯克利華人一作
首個開源「世界模型」來了！
來自UC berkeley的研究人員發布并開源了LWM（LargeWorldModel）系列模型：
論文地址：https://arxiv.org/pdf/2402.08268.pdf
代碼地址：https://github.com/LargeWorldModel/LWM
LWM采用了一個包含各種視頻和書籍的大型數據集，利用RingAttention技術對長序列進行可擴展的訓練，最終將上下文長度增加到1M token。
在超長上下文的加持下，LWM系列模型可以輕松完成各種多模態任務，
比如文本圖像生成（LWM 基于文本提示以自回歸方式生成圖像）：
文本視頻生成：
煙花在天空中爆炸
海浪拍打著海岸
倫敦熙熙攘攘的街道，背景是紅色的電話亭和大本鐘
桌子上的芒果切片，相機從左向右移動
一個球在空中飛舞
慢動作花瓣落在地上
森林中燃燒的篝火
一艘在波濤洶涌的海洋上航行的船
基于圖像的對話：
用戶：如果你是一名攝影師，想要捕捉這個地點的精髓，你會選擇一天中的什么時間和天氣條件？描述您選擇背后的原因。
LWM：我會選擇日落時分的時間。這張照片的特點是一個木制碼頭，延伸到平靜的湖面，是寧靜而風景如畫的場景。夕陽在湖面上投下溫暖的金色光芒，增強了場景的寧靜。此外，天氣條件似乎是涼爽和多云的，這與湖泊和木制碼頭的自然美景相得益彰。這些元素的結合營造出一種視覺上吸引人且平靜的氛圍，將是理想的時機。
用戶：描述一下不同觀眾如何以不同的方式解釋經典藝術作品的這種創造性轉折。
LWM：在對經典藝術作品的創造性轉折中，女人的臉變成了狗的臉，創造了視覺上引人注目和意想不到的表現。這件藝術品可以看作是對人類理解的局限性，以及人類與動物之間更深層次的聯系和理解的評論。
觀看長視頻（超過一小時），并回答問題：
即使最先進的商業模型GPT-4V和Gemini Pro也都失敗了。
雖說「世界模型」還是個概念股吧，但LWM展現出的多模態能力是相當優秀的。
關鍵是，人家是開源的（基于Llama2 7B），于是受到廣大開發者的熱烈歡迎，僅僅不到兩周的時間，就在GitHub上斬獲了6.2k stars。
LWM在博客開頭就展示了自己的優勢區間，除了上面提到的長視頻理解，下圖比較了幾個模型的事實檢索能力：
LWM在1M上下文窗口內實現了高精度，性能優于GPT-4V和Gemini Pro。
我們可以看到，LWM在在不同的上下文大小和位置上都保持了高精度（全綠）。
目前，LWM放出了一系列不同上下文大小（從32K到1M）的模型，包括純語言版本和視頻語言版本。其中視覺語言模型僅在Jax中可用，純語言模型在PyTorch和Jax中都可用。
開源技術細節
上圖展示了LWM的多模態訓練。
第一階段是上下文擴展，重點是使用Books3數據集擴展上下文大小，從32K增長到1M。
第二階段，視覺語言培訓，重點是對不同長度的視覺和視頻內容進行培訓。餅圖詳細說明了訓練數據的分布情況，包括495B的文本-視頻數據，以及33B的文本數據。
圖中還展示了模型的交互功能。
語言模型階段
這個階段首先開發LWM-Text和LWM-Text-Chat，通過使用RingAttention逐步增加序列長度數據進行訓練，并修改位置編碼參數以考慮更長的序列長度。
由于計算的二次復雜度所施加的內存限制，對長文檔的訓練非常昂貴。
為了解決計算限制，研究人員使用RingAttention，利用具有序列并行性的塊計算在理論上擴展到無限上下文，僅受可用設備數量的限制。
作者使用Pallas進一步將RingAttention與FlashAttention融合在一起，以優化性能。通常，如果每個設備有足夠大的token，RingAttention期間的通信成本與計算完全重疊，并且不會增加任何額外的開銷。
訓練步驟
模型以LLaMA-2 7B為基礎，分5個階段逐步增加模型的有效上下文長度：32K、128K、256K、512K和1M。對于每個階段，使用來自The Pile的Books3數據集的不同過濾版本進行訓練。
上表詳細介紹了每個訓練階段的信息，例如token數量、總時間和Books3數據集過濾約束。每個階段以前一個階段作為初始化。
研究人員還構建了一個簡單的QA數據集，用于學習長上下文聊天能力。將Books3數據集中的文檔分塊成1000個token的固定塊，將每個塊提供給短上下文語言模型，并提示它生成一個關于該段落的問答對。
對于聊天模型的微調，研究人員在UltraChat和自定義QA數據集上訓練每個模型，比例約為7：3。
作者發現將UltraChat數據預打包到訓練序列長度至關重要，而且需要與自定義的QA數據示例分開。
聊天模型并沒有采用漸進式訓練，而是從各自的預訓練模型以相同的上下文長度進行初始化。
視覺模型階段
第二階段旨在有效地聯合訓練長視頻和語言序列。
LWM和LWM-Chat 的架構修改
LWM是數百萬長標記序列上的自回歸變換器。視頻中的每一幀都用VQGAN產生256個token。這些token與文本token連接起來，饋送到Transformer中，以自回歸方式預測下一個token。
輸入和輸出token的順序反映了不同的訓練數據格式，包括圖像-文本、文本-圖像、視頻、文本-視頻和純文本格式。
LWM本質上是使用多種模式以任意到任意方式進行訓練的。為了區分圖像和文本token，以及進行解碼，這里采用特殊的分隔符。在視覺數據中，也會處理視頻的中間幀和最終幀。
這里使用來自aMUSEd的預訓練VQGAN，將256 × 256個輸入圖像標記為16 × 16個離散token。
模型使用視覺和文本token的交錯串聯進行訓練，并進行自回歸預測。
不同序列長度的訓練
以LWM-Text-1M文本模型為初始化，對大量組合的文本-圖像和文本-視頻數據執行漸進式訓練過程，這里沒有額外擴展RoPE θ，因為它已經支持高達1M的上下文。
上表顯示了每個訓練階段的詳細信息，每個模型是從先前較短的序列長度階段初始化的。
對于每個階段，根據以下數據進行訓練：
LWM-1K：在大型文本圖像數據集上進行訓練，該數據集由LAION-2Ben和COYO-700M混合組成。數據集被過濾后僅包含至少256分辨率的圖像——總共大約1B個文本圖像對。
在訓練過程中，將文本-圖像對連接起來，并隨機交換模態的順序，以對文本-圖像生成、無條件圖像生成和圖像標題進行建模。這里將文本-圖像對打包為1K個token的序列。
LWM-8K：在WebVid10M和3M InternVid10M示例的文本視頻數據集組合上進行訓練。與之前的工作類似，每種模態使用相同的比例聯合訓練圖像和視頻。
這里將圖像打包成8K token序列和30幀視頻，速度為4FPS。與圖像訓練類似，隨機交換每個文本-視頻對的模態順序。
LWM-Chat-32K/128K/1M：在最后3個階段，研究人員對每個下游任務的聊天數據組合進行訓練：
文本圖像生成
圖像理解
文本視頻生成
視頻理解
通過對預訓練數據的隨機子集進行采樣，并用聊天格式進行增強，構建了文本-圖像和文本-視頻聊天數據的簡單版本。為了理解圖像，這里使用來自ShareGPT4V的圖像聊天指示。
最后，對于視頻理解聊天數據，使用Valley-Instruct-73K和Video-ChatGPT-100K指令數據的組合。對于所有短上下文數據（圖像生成、圖像理解、視頻生成），將序列打包到訓練上下文長度。
在打包過程中，研究人員發現關鍵是要掩蓋注意力，以便每個文本視覺對只關注自己，以及重新加權損失，以使計算與非打包+填充訓練方案中的訓練相同。
對于視頻理解數據，如果視頻太長，會統一采樣最大幀數，以適應模型的訓練上下文長度。在訓練期間，4 個下游任務等比例平均分配。
盡管視覺語言模型可以攝取長視頻，但由于上下文長度有限，通常是通過對視頻幀執行大型時間子采樣來完成的。
例如，Video-LLaVA被限制為從視頻中均勻采樣8幀，無論原始視頻有多長。因此，模型可能會丟失更細粒度的時間信息，而這些信息對于準確回答有關視頻的任何問題非常重要。
相比之下，本文的模型是在1M令牌的長序列上訓練的，因此，可以同時處理數千幀視頻，以在短時間間隔內檢索細粒度信息。在上圖的示例中，LWM正確回答了有關由500多個獨立剪輯組成的1小時長YouTube視頻的問題。
不過作者也承認，LWM生成的答案可能并不總是準確的，并且該模型仍在努力解決需要對視頻有更高層次理解的更復雜的問題。希望LWM將有助于未來的工作，開發改進的基礎模型，以及長視頻理解的基準。
