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

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

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

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

在探索如何讓AI更好地理解世界方面，最近的一項突破性研究引起了廣泛關(guān)注。來自加州大學(xué)伯克利分校的研究團(tuán)隊發(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)實世界中的許多場景，如長篇書籍、電影或電視劇，都包含了豐富的信息，需要更長的上下文來進(jìn)行深入理解。為了應(yīng)對這一挑戰(zhàn)，LWM團(tuán)隊采用了環(huán)形注意力(RingAttention)技術(shù)，成功擴(kuò)展了模型的上下文窗口，使其能夠處理長達(dá)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開始，逐步擴(kuò)展到1M tokens。這一過程不僅有效降低了訓(xùn)練成本，還使模型能夠逐步適應(yīng)更長序列的學(xué)習(xí)。在訓(xùn)練過程中，研究人員還發(fā)現(xiàn)，混合不同長度的圖像、視頻和文本數(shù)據(jù)對于模型的多模態(tài)理解至關(guān)重要。具體包括：

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

Stage I: Learning Long-Context Language Models

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

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

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

任意對任意長序列預(yù)測。RingAttention 能夠使用非常大的上下文窗口進(jìn)行跨視頻-文本、文本-視頻、圖像-文本、文本-圖像、純視頻、純圖像和純文本等多種格式的訓(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ù)集上進(jìn)行訓(xùn)練，比例為7:3。

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

Stage II: Learning Long-Context Vision-Language Models

架構(gòu)修改：在第一階段的基礎(chǔ)上，對LWM和LWM-Chat進(jìn)行修改，使其能夠接受視覺輸入。具體來說，使用預(yù)訓(xùn)練的VQGAN將256x256的輸入圖像轉(zhuǎn)換為16x16的離散token，對視頻進(jìn)行逐幀的VQGAN編碼并將編碼連接起來。此外，引入了特殊的標(biāo)記符號和來區(qū)分文本和視覺token，以及和來標(biāo)記圖像和視頻幀的結(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é)果：在長視頻理解方面，模型能夠處理長達(dá)1小時的YouTube視頻并準(zhǔn)確回答問題，相較于現(xiàn)有模型具有明顯優(yōu)勢。在圖像理解和短視頻理解方面，模型表現(xiàn)一般，但通過更嚴(yán)格的訓(xùn)練和更好的分詞器，有潛力改進(jìn)。在圖像和視頻生成方面，模型可以從文本生成圖像和視頻。Ablation研究表明，屏蔽序列填充對于圖像理解等下游任務(wù)至關(guān)重要。

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

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

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

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

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

訓(xùn)練損失曲線：圖10和圖11展示了第一階段語言模型和第二階段視覺-語言模型的訓(xùn)練損失曲線。可以看出，隨著訓(xùn)練進(jìn)行，損失持續(xù)下降。

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