本文經(jīng)自動駕駛之心公眾號授權(quán)轉(zhuǎn)載，轉(zhuǎn)載請聯(lián)系出處。

在探索如何讓AI更好地理解世界方面，最近的一項突破性研究引起了廣泛關(guān)注。來自加州大學(xué)伯克利分校的研究團隊發(fā)布了“Large World Model, LWM”，能夠同時處理百萬級長度的視頻和語言序列，實現(xiàn)了對復(fù)雜場景的深入理解。這一研究無疑為未來AI的發(fā)展開啟了新的篇章。

論文地址：World Model on Million-Length Video And Language With RingAttention

博客地址：Large World Models

huggingface: LargeWorldModel (Large World Model)

在傳統(tǒng)方法中，AI模型往往只能處理較短的文本或視頻片段，缺乏對長時間復(fù)雜場景的理解能力。然而，現(xiàn)實世界中的許多場景，如長篇書籍、電影或電視劇，都包含了豐富的信息，需要更長的上下文來進行深入理解。為了應(yīng)對這一挑戰(zhàn)，LWM團隊采用了環(huán)形注意力(RingAttention)技術(shù)，成功擴展了模型的上下文窗口，使其能夠處理長達100萬個令牌(1M tokens)的序列。例如實現(xiàn)超過 1 小時的問答視頻：

圖1.長視頻理解。LWM 可以回答有關(guān)超過 1 小時的 YouTube 視頻的問題。

超過 1M 上下文的事實檢索:

圖 2. 針檢索任務(wù)。LWM 在 1M 上下文窗口內(nèi)實現(xiàn)了高精度，并且性能優(yōu)于 GPT-4V 和 Gemini Pro。

圖 3. 針檢索任務(wù)。LWM 對于上下文窗口中不同的上下文大小和位置實現(xiàn)了高精度。

技術(shù)實現(xiàn)

為了訓(xùn)練和評估LWM，研究人員首先收集了一個包含各種視頻和書籍的大型數(shù)據(jù)集。然后，他們逐步增加了訓(xùn)練的上下文長度，從4K tokens開始，逐步擴展到1M tokens。這一過程不僅有效降低了訓(xùn)練成本，還使模型能夠逐步適應(yīng)更長序列的學(xué)習(xí)。在訓(xùn)練過程中，研究人員還發(fā)現(xiàn)，混合不同長度的圖像、視頻和文本數(shù)據(jù)對于模型的多模態(tài)理解至關(guān)重要。具體包括：

模型訓(xùn)練分兩個階段：首先通過訓(xùn)練大型語言模型擴展上下文大小。然后進行視頻和語言的聯(lián)合訓(xùn)練。

Stage I: Learning Long-Context Language Models

擴展上下文：利用RingAttention技術(shù)，可以無近似地擴展上下文長度到數(shù)百萬個token。同時，通過逐步增加訓(xùn)練序列長度，從32K tokens開始，逐步增加到1M tokens，以減少計算成本。此外，為了擴展位置編碼以適應(yīng)更長的序列，采用了簡單的方法，即隨上下文窗口大小增加而增加RoPE中的θ。

上下文擴展和視覺語言訓(xùn)練。使用 RingAttention 將書籍上的上下文大小從 4K 擴展到 1M，然后對長度為 32K 到 1M 的多種形式的視覺內(nèi)容進行視覺語言訓(xùn)練。下面板顯示了理解和響應(yīng)有關(guān)復(fù)雜多模式世界的查詢的交互功能。

訓(xùn)練步驟：首先從LLaMA-2 7B模型初始化，然后在5個階段逐步增加上下文長度，分別是32K、128K、256K、512K和1M tokens。每個階段都使用不同過濾版本的Books3數(shù)據(jù)集進行訓(xùn)練。隨著上下文長度的增加，模型能夠處理更多tokens。

任意對任意長序列預(yù)測。RingAttention 能夠使用非常大的上下文窗口進行跨視頻-文本、文本-視頻、圖像-文本、文本-圖像、純視頻、純圖像和純文本等多種格式的訓(xùn)練。請參閱LWM 論文了解關(guān)鍵功能，包括屏蔽序列打包和損失加權(quán)，它們可以實現(xiàn)有效的視頻語言訓(xùn)練。

對話微調(diào)：為了學(xué)習(xí)長上下文的對話能力，構(gòu)建了一個簡單的問答數(shù)據(jù)集，將Books3數(shù)據(jù)集的文檔分割成1000 token的塊，然后利用短上下文語言模型為每個塊生成一個問答對，最后將相鄰的塊連接起來構(gòu)造一個長上下文的問答示例。在微調(diào)階段，模型在UltraChat和自定義問答數(shù)據(jù)集上進行訓(xùn)練，比例為7:3。

語言評估結(jié)果：在單針檢索任務(wù)中，1M上下文的模型可以在整個上下文中近乎完美地檢索出隨機分配給隨機城市的數(shù)字。在多針檢索任務(wù)中，模型在檢索一個針時表現(xiàn)良好，在檢索多個針時性能略有下降。在短上下文語言任務(wù)評估中，擴大上下文長度并沒有降低性能。在對話評估中，增加對話交互能力可能會降低系統(tǒng)檢索具體信息或“針”的精度。

Stage II: Learning Long-Context Vision-Language Models

架構(gòu)修改：在第一階段的基礎(chǔ)上，對LWM和LWM-Chat進行修改，使其能夠接受視覺輸入。具體來說，使用預(yù)訓(xùn)練的VQGAN將256x256的輸入圖像轉(zhuǎn)換為16x16的離散token，對視頻進行逐幀的VQGAN編碼并將編碼連接起來。此外，引入了特殊的標記符號和來區(qū)分文本和視覺token，以及和來標記圖像和視頻幀的結(jié)束。

訓(xùn)練步驟：從LWM-Text-1M模型初始化，采用與第一階段類似的逐步增加序列長度的訓(xùn)練方法，首先在1K tokens上訓(xùn)練，然后是8K tokens，最后是32K、128K和1M tokens。訓(xùn)練數(shù)據(jù)包括文本-圖像對、文本-視頻對以及下游任務(wù)的聊天數(shù)據(jù)，如文本-圖像生成、圖像理解、文本-視頻生成和視頻理解。在訓(xùn)練過程中，逐步增加下游任務(wù)的混合比例。

視覺-語言評估結(jié)果：在長視頻理解方面，模型能夠處理長達1小時的YouTube視頻并準確回答問題，相較于現(xiàn)有模型具有明顯優(yōu)勢。在圖像理解和短視頻理解方面，模型表現(xiàn)一般，但通過更嚴格的訓(xùn)練和更好的分詞器，有潛力改進。在圖像和視頻生成方面，模型可以從文本生成圖像和視頻。Ablation研究表明，屏蔽序列填充對于圖像理解等下游任務(wù)至關(guān)重要。

文本到圖像。LWM 根據(jù)文本提示以自回歸方式生成圖像。

文本到視頻。LWM 根據(jù)文本提示以自回歸方式生成視頻。

第二階段通過逐步增加序列長度并在大量文本-圖像和文本-視頻數(shù)據(jù)上訓(xùn)練，成功擴展了第一階段的語言模型，使其具備視覺理解能力。這一階段的模型可以處理長達1M tokens的多模態(tài)序列，并在長視頻理解、圖像理解和生成等方面展現(xiàn)出強大的能力。

技術(shù)細節(jié)（Further Details）

訓(xùn)練計算資源：模型使用TPUv4-1024進行訓(xùn)練，相當于450個A100 GPU，使用FSDP進行數(shù)據(jù)并行，并通過RingAttention支持大上下文。

訓(xùn)練損失曲線：圖10和圖11展示了第一階段語言模型和第二階段視覺-語言模型的訓(xùn)練損失曲線?？梢钥闯?，隨著訓(xùn)練進行，損失持續(xù)下降。

訓(xùn)練超參數(shù)：附錄F提供了詳細的訓(xùn)練超參數(shù)，包括參數(shù)量、初始化模型、序列長度、RoPE參數(shù)、每批tokens數(shù)、總tokens數(shù)、訓(xùn)練步驟數(shù)、學(xué)習(xí)率計劃、學(xué)習(xí)率預(yù)熱步數(shù)、最大學(xué)習(xí)率和最小學(xué)習(xí)率、計算資源等。

推斷擴展：實現(xiàn)了RingAttention用于解碼，支持對長達數(shù)百萬tokens的序列進行推斷，需使用至少v4-128 TPU，并進行32路tensor并行和4路序列并行。

量化：文檔指出，模型使用單精度進行推斷，通過量化等技術(shù)可以進一步提高擴展性。

一些例子

基于圖像的對話。

圖 6. 圖像理解。LWM 可以回答有關(guān)圖像的問題。

超過 1 小時的 YouTube 視頻視頻聊天。

圖 7. 長視頻聊天。

即使最先進的商業(yè)模型 GPT-4V 和 Gemini Pro 都失敗了，LWM 仍能回答有關(guān) 1 小時長的 YouTube 視頻的問題。每個示例的相關(guān)剪輯位于時間戳 9:56（頂部）和 6:49（底部）。