突破性的百萬級視頻和語言世界模型：Large World Model

作者：北方的郎 2024-03-26 09:58:52

來自加州大學伯克利分校的研究團隊發布了“Large World Model, LWM”，能夠同時處理百萬級長度的視頻和語言序列，實現了對復雜場景的深入理解。

本文經自動駕駛之心公眾號授權轉載，轉載請聯系出處。

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

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

博客地址：Large World Models

huggingface: LargeWorldModel (Large World Model)

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

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

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

圖 2. 針檢索任務。LWM 在 1M 上下文窗口內實現了高精度，并且性能優于 GPT-4V 和 Gemini Pro。

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

技術實現

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

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

Stage I: Learning Long-Context Language Models

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

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

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

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

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

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

Stage II: Learning Long-Context Vision-Language Models

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

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

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