首個Mamba+Transformer混合架構多模態大模型來了,實現單卡千圖推理

机器之心發表於2024-09-21

AIxiv專欄是機器之心釋出學術、技術內容的欄目。過去數年,機器之心AIxiv專欄接收報導了2000多篇內容,覆蓋全球各大高校與企業的頂級實驗室,有效促進了學術交流與傳播。如果您有優秀的工作想要分享,歡迎投稿或者聯絡報導。投稿郵箱:liyazhou@jiqizhixin.com;zhaoyunfeng@jiqizhixin.com

本文作者來自於香港中文大學深圳和深圳大資料研究院。其中第一作者為香港中文大學深圳博士生王熙棟和研究助理宋定傑,主要研究方向分別為醫療AGI和多模態學習;博士生陳舒年研究方向為多模態學習,博士生張辰研究方向為高效語言模型。通訊作者為香港中文大學深圳資料科學學院王本友教授。

擴充套件多模態大語言模型(MLLMs)的長上下文能力對於影片理解、高解析度影像理解以及多模態智慧體至關重要。這涉及一系列系統性的最佳化,包括模型架構、資料構建和訓練策略,尤其要解決諸如隨著影像增多效能下降以及高計算成本等挑戰。

該團隊將模型架構調整為 Mamba 和 Transformer 塊的混合體,在資料構建中考慮多個影像之間的時間和空間依賴性,並採用漸進式訓練策略。提出了首個混合架構多模態大語言模型 LongLLaVA,在效率和效能之間實現了更好的平衡。

圖片



LongLLaVA 不僅在各種基準測試中取得了有競爭力的結果,還保持了高吞吐量和低視訊記憶體消耗,其可以在單個 A100 80GB GPU 上處理近千張影像,展現出了廣闊的應用前景。

圖片



  • 論文地址:https://arxiv.org/abs/2409.02889
  • 專案地址:https://github.com/FreedomIntelligence/LongLLaVA

1. 簡介

多模態大語言模型(MLLMs)的快速進步展示了它們在各個應用領域中的顯著能力。然而,多影像理解場景仍然是一個重要但尚未充分探索的方面。特別是,將 MLLMs 的應用場景擴充套件到理解更長的影片、更高解析度的影像以及基於更多歷史資訊的決策,對於提升使用者體驗和進一步擴充 MLLMs 的應用範圍至關重要。

然而,將 MLLM 的上下文長度擴充套件以提高其可用性,面臨著處理更多影像時效能下降和計算成本高昂的挑戰。一些研究專注於構造包含多個影像的長上下文訓練資料,以增強效能。其他研究探索了創新性的訓練策略,以減輕效能下降。關於高計算成本的問題,LongVILA 透過降低通訊成本在提高多節點效率方面取得了進展。然而,在管理更長的上下文時,加速計算這個問題本身仍有待解決。

為了解決上述挑戰,該研究提出了 LongLLaVA 系統解決方案,採用混合架構進行加速。該解決方案在三個維度上進行了全面最佳化:多模態架構、資料構建和訓練策略

  • 對於多模態架構,採用結合 Transformer 和 Mamba 的混合架構,並提出了一種高效影像表示方法,該方法對影像 Token 應用 2D 池化以降低計算成本同時保持效能。
  • 對於資料構建,為不同的任務設計了獨特的格式,使模型能夠區分影像之間的時間和空間的依賴關係。
  • 訓練策略方面,採用了一種三階段的多模態自適應方法 —— 單影像對齊、單影像指令調整和多影像指令調整 —— 以逐步提升模型處理多模態長上下文的能力。

實驗結果表明,LongLLaVA 在高效理解多模態長上下文方面表現卓越。它在VNBench的檢索、計數和排序任務中領先,並在單張 80GB GPU 上對 1000 張影像進行大海撈針評估時達到了近 100% 的準確率。從保證研究可復現和促進社群發展出發,團隊將開源所有與 LongLLaVA 相關的模型、程式碼和資料集。

2. LongLLaVA: 將 LLaVA 擴充套件到更長的上下文

為了解決上述挑戰並提高模型對長文字和多影像場景的適應性,團隊從三個角度進行了改進:多模態模型架構,資料構造和訓練策略。

圖片



2.1 多模態架構

LongLLaVA 基於 LLaVA 的三個核心元件:視覺編碼器、 對映器 和大語言模型

視覺資訊處理。團隊使用 CLIP 作為視覺編碼器來編碼視覺資訊,並採用兩層 MLP 作為對映器,將視覺特徵對映到適合 LLM 的文字嵌入空間。在對映之前,應用2D池化,有效地節省了訓練和推理時間,同時保持了影像塊之間的基本空間關係。

混合 LLM 架構。LongLLaVA 採用了一種混合 LLM 架構,將 Transformer 和 Mamba 層以 1:7 的比例整合,如圖 2 所示。在每一層中還採用了混合專家(MoE)方法,使用 16 個專家,併為每個 Token 選擇前兩個專家。在層之間使用 RMSNorm 來增強歸一化,但省略了位置嵌入。該模型整合了分組 Query 注意力(GQA)和 SwiGLU 啟用函式,與其他大型語言模型相似。模型的總體引數數量為 530 億,推理過程中的啟用引數總數為 130 億。

圖片


2.2 資料處理協議

為確保模型在多影像場景中有效地區分影像之間的時序和空間依賴關係,並在各種任務中表現良好,團隊細緻地區分了不同場景下的特殊字元。如圖 3 所示,這些特殊字元全面處理了不同情境下影像之間的各種關係,從而增強了模型對不同任務的適應性。

  • 常規單圖和多圖輸入:使用 <img></img> 幫助模型區分影像和文字 Token 。
  • 影片:在不同的幀之間新增 <t>,以表示它們之間的時間依賴性。
  • 高解析度影像:使用換行符 "\n" 來區分主影像與其子影像。 對於子影像的排列,透過從左上角遍歷到右下角的方式進行分割,在分割行之間新增"\n"以保留子影像的相對空間位置。

2.3 訓練策略

團隊逐步實現單模態和多模態的適配,將預訓練語言模型轉變為多模態長上下文模型。

純文字指令微調。首先提升預訓練語言模型在純文字場景中遵循不同長度指令的能力。這是透過使用包含來自 Evol-instruct-GPT4、WildChat 和 LongAlign 的 278k 條純文字條目的資料集實現的。

圖片


在多模態自適應方面,在 LLaVA 中 “單影像對齊” 和 “單影像指令微調” 階段之後,團隊引入了 “多影像指令微調” 階段,逐步增強模型的多模態長上下文能力。採用漸進式訓練不僅是為了更好地控制變數,也是為了增加模型的可重用性。具體的資料集使用情況如圖 4 所示。

第一階段:單影像對齊。這一階段是為了將視覺模態特徵與文字模態進行對齊。團隊使用了 ALLaVA-Caption 和 ShareGPT4V 等資料集,這些資料集包含大約 600K 個高質量的影像 - 字幕對。在此階段,僅訓練對映器,同時凍結視覺編碼器和 LLM 的引數

第二階段:單影像指令微調。這個階段的目的是賦予模型多模態指令遵循能力。團隊使用了 LLaVA-1.5 和 Manti-Single 等資料集,總共有約 932K 個高質量的問答對。在此過程中,只凍結了視覺編碼器,而對映器和 LLM 部分進行訓練。

第三階段:多影像指令微調。在這一階段,模型被訓練以在多模態長文字場景中遵循指令。團隊分別從 Mantis、VideoChat2 和 ShareGPT4Video 中取樣 200K、200K 和 50K 資料項。為了保留模型的單影像理解和純文字對話能力,團隊將來自單影像指令微調和純文字指令微調階段的額外 200K 和 50K 資料項作為 Replay 部分。此外,為了提高模型解釋複雜單影像(分割成多個子影像)的能力,團隊從單影像指令微調階段取樣 50K 條資料,進行填充和分割,將原始影像分割成尺寸為 336x336 的子影像作為 SubImage 部分

3. 評估結果

3.1 主要結果

如表 2 所示,LongLLaVA 在 MileBench 上表現出色,甚至超過了閉源模型Claude-3-Opus,尤其在檢索任務方面表現出色。突顯其在處理多影像任務方面的強大能力。

圖片


LongLLaVA 在涉及中等至長影片的任務中表現出色,超越了傳統的影片模型,如 Video-LLaMA2 和 VideoChat2。在取得了這些令人印象深刻結果的同時,LongLLaVA 的 FLOPs 比其他模型少一個數量級

3.2 長上下文大型語言模型的診斷評估

考慮到以前的評估不能充分捕捉 MLLM 在長語境下的能力,團隊採用了一個新的診斷評估集 VNBench,以進一步分析模型在長語境下的原子能力。VNBench 是一個基於合成影片生成的長上下文診斷任務框架,包括檢索、排序和計數等任務。

圖片


結果顯示如表 3 所示,LongLLaVA 在跨語境檢索、排序和技術能力等任務中的表現與領先的閉源模型相當,甚至在某些方面超過了 GPT-4V。在開源模型中,LongLLaVA 也展現出其卓越的效能。展示了 LongLLaVA 在管理和理解長上下文方面的先進能力。

3.3 消融實驗

圖片


表 4 中顯示,使用具有相同資料的混合 LLM 架構,在評估集中都觀察到了顯著的改進,證明了其在多模態場景中的潛力。對於 Token 壓縮,選擇了 2D 池化,這顯著減少了計算負載,同時將效能下降控制在可接受範圍內。與 1D 池化相比,2D 池化方法得到更好的結果。在資料構建方面,在訓練團隊的單影像資料後,模型在 SEEDBench 上的準確率提高了 1.5%,在 MileBench 上提高了 12.3%。隨後的多影像訓練使得 MileBench 上的準確率進一步提高了 7.4%,驗證了資料集構建的有效性。

4. 更多分析

為了解 LongLLaVA 的內部工作原理和跨模態長文字處理能力,該團隊進行了進一步分析。

4.1 關於混合架構的動機

圖片


團隊探討了不同架構在 ICL 能力和推理效率方面的優缺點,強調了混合架構的平衡優勢。

ICL 分析。團隊評估了在 VL-ICL 基準測試中對多模態情境學習中匹配影像任務的效能。該任務的輸入包含一個影像對,輸出表示是否存在特定的關係。MLLM 需要從示例中學習關係。如表 5 所示,混合架構和 Transformer 架構隨著示例數量的增加表現出快速的效能提升,而 Mamba 架構的提升較少,證實了其在情境學習方面的不足。

效率分析。團隊關注三個方面:預填充時間(首次推理延遲)、吞吐量(每秒生成的下一個 Token 數)和記憶體使用。團隊將輸入文字長度控制在 100K,並測量生成 1 個 Token 和 1000 個 Token 的輸出所需的時間和最大記憶體使用。吞吐量計算為圖片。為了更好地模擬實際應用場景,使用 vLLM 框架和 Int8 量化評估了 Transformer 和混合架構。如表 5 所示,Mamba 架構具有最快的預填充時間,最高的吞吐量。與具有相似推理引數的 Transformer 架構相比,混合架構實現了 2.5 倍的吞吐量,75% 的預填充時間,並減少了記憶體使用。

4.2 影像數量的縮放定律

隨著可處理影像數量的增加,模型能夠支援更多影像塊以進行高解析度影像理解,以及使用更多影片幀進行影片理解。為了探索增加子影像和影片幀數量的影響,團隊分別在 V* Bench 和 Video-MME 基準測試上評估了 LongLLaVA。

增加子影像數量。V* Bench 評估了一個模型在大型影像中定位小目標的能力。如圖 5 所示,最初增加子影像的數量顯著提高了模型效能,表明模型對影像細節的理解更好。然而,團隊也發現,進一步增加子影像的數量略微降低了效能,這表明過多的子影像可能會干擾在此任務上的效能。

圖片



增加幀數規模。影片多模態編碼器是一個測試模型從影片中提取資訊能力的基準。從圖 6 中可以看到,隨著取樣幀數的增加,模型在基準測試中的效能顯著提高,當提取 256 幀時達到峰值。這表明模型能夠有效地理解和利用額外取樣幀中包含的資訊,以提供更好的響應。

圖片


5. 進一步將影像數量擴大到 1000

利用 LongVA 中提出的 V-NIAH 評估框架,團隊進行了 “大海撈針” 測試來評估模型效能。考慮到模型的訓練序列長度限制為 40,960 個 token,採用 token 池化技術將原始 token 數量從 144 個減少到 36 個。這種調整能夠高效地從大量資料集中檢索相關資訊。如圖 7 所示,模型在 1000 張影像集上實現了近 100% 的檢索準確率,而無需額外的訓練。

圖片


然而,當增加測試影像數量超過 1,000 張時,團隊觀察到檢索準確率下降。這種效能下降可能是因為超出了模型的訓練序列長度,這可能會影響其保持更多影像準確性的能力。在未來的工作中團隊將延長訓練序列長度至 140,000 Token,即 LongLLaVA 進行單卡推理的極限長度,以進一步釋放模型潛力。

6. 結論

LongLLaVA(長上下文大型語言和視覺助手)這一創新性混合架構模型,在長上下文多模態理解方面表現出色。該模型整合了 Mamba 和 Transformer 模組,利用多個影像之間的時空依賴性構建資料,並採用漸進式訓練策略。

LongLLaVA 在各種基準測試中表現出競爭性的效能,同時確保了效率,為長上下文多模態大型語言模型(MLLMs)設定了新的標準。

相關文章