本文由華東師范大學和小紅書聯(lián)合完成，共同第一作者是華東師范大學在讀碩士、小紅書 NLP 團隊實習生黃文軒和翟子杰，通訊作者是小紅書 NLP 團隊負責人曹紹升，以及華東師范大學林紹輝研究員。
多模態(tài)大模型（MLLMs）在視覺理解與推理等領域取得了顯著成就。然而，隨著解碼（decoding）階段不斷生成新的 token，推理過程的計算復雜度和 GPU 顯存占用逐漸增加，這導致了多模態(tài)大模型推理效率的降低?，F(xiàn)有的方法通過減少預填充（prefill）階段的視覺 token 冗余來實現(xiàn)推理加速。遺憾的是，這種在預填充階段實現(xiàn)的視覺 token 稀疏化所帶來的加速優(yōu)勢，在解碼階段會逐漸減弱。當解碼輸出的文本 token 數(shù)量增多時，這些方法仍然會遇到性能瓶頸。
為了解決上述問題，團隊創(chuàng)新性地提出了一個全新的動態(tài)視覺 - 文本上下文稀疏化推理加速框架 ——Dynamic-LLaVA。該框架針對多模態(tài)大模型在不同推理模式下（包括預填充階段以及有無 KV Cache 的解碼階段），設計了定制化的稀疏化推理方案，以實現(xiàn)多模態(tài)大模型的高效推理。實驗結果表明，Dynamic-LLaVA 在幾乎不損失視覺理解和生成能力的前提下，能夠將預填充階段的計算開銷減少約 75%；在無 KV Cache 的解碼階段，計算開銷減少約 50%；在有 KV Cache 的解碼階段，GPU 顯存占用減少約 50%。Dynamic-LLaVA 為多模態(tài)大模型推理加速領域樹立了新的標桿。
- 論文標題：Dynamic-LLaVA: Efficient Multimodal Large Language Models via Dynamic Vision-language Context Sparsification
- 論文 GitHub 倉庫：https://github.com/Osilly/dynamic_llava
- 論文鏈接：https://arxiv.org/abs/2412.00876
- 研究團隊：華東師范大學、小紅書 NLP 團隊
1 引言
1.1 前置信息：預填充與解碼
本文主要圍繞以 LLaVA 為范式的多模態(tài)大模型展開研究。一個多模態(tài)大模型的推理過程可以分為預填充和解碼兩個階段：
在預填充階段，不同模態(tài)的特征被映射到與大語言模型（LLM）輸入 embedding 相同的特征分布空間中。這些多模態(tài)特征與文本 token 會一起被大語言模型處理，以生成初始輸出文本 token。以圖片理解場景為例，該階段主要處理輸入的圖片和文本格式的問題。
在隨后的解碼階段，預填充階段生成的所有 token 以及后續(xù)生成的所有輸出文本 token，將被用于自回歸生成，從而產(chǎn)生完整的輸出。同樣以圖片理解場景為例，該階段生成針對整個問題的完整回答。
1.2 多模態(tài)大模型推理加速困境
圖 1：多模態(tài)大模型生成過程（有 / 無 KV Cache）中 FLOPs（計算復雜度）和 GPU 顯存開銷的增長趨勢
現(xiàn)有的多模態(tài)大模型大多以基于解碼器架構的大語言模型（LLM）為核心，這些模型通常擁有龐大的參數(shù)規(guī)模。在生成輸出文本 token 的過程中，模型計算負擔會逐漸加重，導致對計算資源的巨大消耗。為了提升推理速度，現(xiàn)有模型通常會在解碼過程中運用 KV Cache 技術，通過存儲并復用之前計算的 KV 激活值來減少重復計算。然而，如圖 1 (B) 所示，即使使用了 KV Cache，LLaVA 在輸出 token 不斷增加時，仍會迅速面臨 GPU 顯存耗盡的問題。
與文本不同，視覺信息往往包含大量冗余。因此，許多方法嘗試通過減少視覺上下文來加速多模態(tài)大模型的推理，即對預填充階段的視覺 token 進行剪枝處理。但這種方法存在局限性：其主要提升了多模態(tài)大語言模型在預填充階段的推理效率，而在解碼階段，其效率提升會逐漸減弱。
如圖 1 (B) 和 (C) 所示，F(xiàn)astV 這種針對視覺 token 剪枝的方法，雖然相較于原始的 LLaVA 能夠節(jié)省一定的 GPU 顯存和計算開銷（FLOPs），但當輸出 token 數(shù)接近 5K 時，它仍然會遭遇計算資源瓶頸。此外，F(xiàn)astV 和原始 LLaVA 的曲線斜率基本一致，這表明在長輸出的解碼階段，這類方法并沒有顯著的推理效率優(yōu)勢。因此，僅通過減少預填充階段的視覺 token，在輸出文本 token 數(shù)量遠超視覺 token 時，難以實現(xiàn)整個推理效率的顯著提升。
1.3 邁向全階段推理加速：Dynamic-LLaVA
針對上述問題，我們認為：為了實現(xiàn)真正的全階段推理加速，不僅需要對預填充階段的視覺 token 進行剪枝，還必須對解碼階段輸出的文本 token 進行稀疏化處理，限制參與自回歸運算的 token 數(shù)量。為此，我們提出了Dynamic-LLaVA，針對多模態(tài)大模型的視覺 - 語言上下文稀疏化推理加速框架。該框架能夠集成到多模態(tài)大模型推理的不同階段中，實現(xiàn)以下目標：
- 顯著降低預填充階段計算開銷：通過優(yōu)化視覺 token 的處理方式，減少不必要的計算。
- 提升解碼階段的推理效率：無論是否使用 KV Cache，都能減少計算開銷，提高推理速度。
- 保持性能優(yōu)勢：在視覺理解任務上幾乎不損失性能；在長文本輸出場景中，生成能力也幾乎不受影響。
通過這些創(chuàng)新，Dynamic-LLaVA 為多模態(tài)大模型的高效推理提供了一種全新的解決方案。
2 方法
圖 2：Dynamic-LLaVA 整體框架
如圖 2 所示，Dynamic-LLaVA 可以集成到多模態(tài)大模型推理流程中的不同階段。具體而言，在預填充階段，該框架對視覺 token 執(zhí)行精準剪枝操作，剔除冗余信息；在不使用 KV Cache 的解碼階段，限制參與自回歸運算的視覺與輸出文本 token 數(shù)量，避免不必要的計算負擔；而在使用 KV Cache 的解碼階段，Dynamic-LLaVA 則動態(tài)調控 KV Cache，自適應判斷是否將當前輸出文本 token 的 KV 激活值納入 KV Cache，優(yōu)化資源利用效率。為了使模型適應這種全新的稀疏化推理模式，Dynamic-LLaVA 在預訓練的 LLaVA-1.5 基礎上進行了 1 個 epoch 的監(jiān)督微調（SFT），確保模型能夠高效地運行在稀疏化的推理路徑上。
2.1 預填充階段
在預填充階段，我們對輸入的視覺 token 進行稀疏化操作。如圖 2 左側部分所示，我們引入一個可訓練的輕量化的圖像預測器（Image Predictor），來判斷應當丟棄哪些視覺 token。該圖像預測器的結構如下圖：
圖 3：圖像預測器的結構示意圖
圖像預測器會對每個視覺 token 產(chǎn)生 “決策分數(shù)”，以決定對哪些視覺 token 進行保留。在端到端訓練中，視覺 token 的剪枝通過 0-1 二值化的掩碼操作實現(xiàn)（具體過程見 2.4 節(jié)）。在實際推理階段中，通過保留 “決策分數(shù)” 前 k 大的視覺 token（即圖 2 左側部分的 “Yes” 分支），實現(xiàn)視覺 token 數(shù)量減少，以實現(xiàn)推理加速。
2.2 解碼階段
不使用 KV Cache 的解碼過程：
對于視覺 token，采用和上一小節(jié)相同的做法，進行稀疏化處理。
對于輸出的文本 token，分兩類進行處理：
- 最后一個輸出的文本 token（即圖 2 中間部分的 “Last output text token”），不進行任何處理，完整輸入 LLM 的 decoder 層進行計算。這樣做的目的是保證模型的輸出內容是連貫的，產(chǎn)生新的輸出文本 token 時，始終保證自回歸運算包含上一個輸出文本 token。
- 對其他歷史的輸出文本 token 進行稀疏化操作，其形式類似于對視覺 token 的處理。引入一個結構如下圖的輸出預測器（Output Predictor），給出每個輸出文本 token 的 “決策分數(shù)”，以決定當前產(chǎn)生新的輸出內容時，應當包括哪些文本 token 進行自回歸運算。圖 2 中間部分的 “Yes” 分支，表明保留的輸出文本 token。
圖 4：輸出預測器的結構示意圖
使用 KV Cache 的解碼過程：
KV Cache 是節(jié)省冗余計算的一個關鍵推理加速技術，其思想是 “用 GPU 顯存的空間換計算時間”。顯而易見的是，KV Cache 也并非無限大，在長輸出情況下，必須丟棄一些 KV Cache 以適應有限的 GPU 顯存。目前在 LLM 領域已有大量的 KV Cache 壓縮方案，以 方法為代表，這一類方法一般基于當前 token 和歷史 KV Cache 進行重要性分數(shù)計算，以壓縮歷史 KV Cache。
與上述方法不同的是，我們對有 KV Cache 的解碼階段的設計，核心在于 “僅判斷當前新 token 的 KV 激活是否需要加入 KV Cache 中”。如圖 2 右側所示，對于當前正在處理的新 token（Last output text token），使用和上一部分結構相同的輸出預測器，以決定是否加入 KV Cache 集合中。這種 “Online KV Cache 壓縮” 方法，判斷是否保留 KV Cache 的過程計算復雜度更低，也更加適應多模態(tài)場景。在論文附錄中，我們詳細討論了我們的方法和現(xiàn)有的 LLM KV Cache 壓縮方法的區(qū)別。
需要特別說明的是，和不使用 KV Cache 的解碼階段相同，無論當前處理的 token 是否加入 KV Cache，其都會輸入 LLM decoder 層進行計算，以保證輸出的連貫性。
2.3 端到端訓練
圖 5：Dynamic-LLaVA 在端到端訓練過程中的示意圖
Dynamic-LLaVA 是一個需要訓練的多模態(tài)大模型推理加速框架。我們基于 LLaVA 進行了一個 epoch 的指令微調，以實現(xiàn)對 token 動態(tài)選擇的穩(wěn)定性，保證最終的性能。為了保證端到端訓練，在訓練階段的稀疏化操作通過 0-1 二值化掩碼實現(xiàn)（在推理中的實現(xiàn)是直接從歷史 token 序列中丟棄 token）。如圖 5 所示，上半部分表示訓練中進行 mask 的過程，在得到整個 token 序列的重要性分數(shù)后，我們選取前 k 重要的 token 進行保留，相對應的生成掩碼向量，其中 0 對應丟棄的冗余 token（不參與注意力過程的計算），1 對應保留的重要 token，進一步基于掩碼向量生成注意力過程的掩碼矩陣。掩碼矩陣用來對多頭注意力機制進行掩碼操作，以確保丟棄的 token 不參與注意力過程的計算。由于二值化操作會導致不可微問題，所以我們借助了 GumbalSoftmax 和梯度直通估計器（Straight Through Estimator, STE）來保證梯度流的正確傳播，以進行端到端的訓練，如圖 5 下半部分所示。
3 實驗
Dynamic-LLaVA 基于 LLaVA-1.5-7B 和 13B 的兩個版本進行了 1 個 epoch 的指令微調，訓練使用的數(shù)據(jù)和 LLaVA-1.5 相同。
3.1 視覺理解能力
我們首先評估了 Dynamic-LLaVA 在主要的視覺理解基準的性能，選取了目前主流的多模態(tài)大模型推理加速方法進行比較。
表 1：視覺理解基準效果對比。其中，F(xiàn)ree 表示方法是否是 Training-Free 的。Dynamic-LLaVA 的下標 "I" 和 "I | T" 分別表示僅對視覺 token 做稀疏化和同時對視覺和文本 token 都做稀疏化（該標識適用于下文所有的表格）
如表 1 所示，Dynamic-LLaVA 在大部分視覺理解任務上取得了優(yōu)越的性能。和其他對視覺內容稀疏化的方法相比，Dynamic-LLaVA 在能大幅減小計算復雜度的同時，能夠實現(xiàn)相比原始的 LLaVA-1.5 性能幾乎不下降。此外，在 SciQA、POPE、MME 和 MMBench 上，Dynamic-LLaVA 相比 LLaVA-1.5 甚至有一定的性能提升。例如，在 SciQA 任務上，Dynamic-LLaVA 的 7B 和 13B 版本，相較于 LLaVA-1.5 實現(xiàn)了 2.3% 和 0.8% 的性能提升。
表 2：與其他高效視覺 projector 的 SOTA 方法對比
值得一提的是，Dynamic-LLaVA 并沒有對 LLaVA-1.5 的視覺 projector 進行修改，就可以實現(xiàn)大幅降低預填充階段計算復雜度，同時維持模型性能。在表 2 中，和其他針對視覺 projector 做高效設計（以提高推理效率）的 SOTA 方法進行了對比。相較于其他使用了高效的視覺 projector 的方法，Dynamic-LLaVA 使用和 LLaVA-1.5 相同的 MLP 結構作為視覺 projector，實現(xiàn)了更好的性能，同時也大幅降低了預填充階段的計算復雜度。此外，Dynamic-LLaVA 也可以和其他使用高效視覺 projector 的方法集成。例如，表 2 中 Dynamic-LLaVA 使用 TokenPacker 這一高效視覺 projector 的版本，在原始的 TokenPacker 方法基礎上，進一步減少了視覺 token。相較于其他基于 TokenPacker 的推理加速方法，性能損失最少。
3.2 生成能力
現(xiàn)有的視覺理解任務中，一般只要求模型給出簡短的回復，這和現(xiàn)實世界中多模態(tài)大模型的應用場景仍然存在不小的區(qū)別。在現(xiàn)實使用中，多模態(tài)大模型多數(shù)情況下會被要求生成更長、更細致的描述。為了和現(xiàn)實世界的場景對齊，評估 Dynamic-LLaVA 在更長的輸出情況下的生成能力和推理效率。我們額外構建了兩個評估模型生成能力的基準：
- LVIS-VQA：基于 LVIS-Instruct4 數(shù)據(jù)集，選取了 1000 個回答超過 100 個單詞的單輪對話樣本構成 LVIS-VQA (single round) 和 1000 個多輪對話樣本（平均回答單詞數(shù)超過 300）構成 LVIS-VQA (multi-round)；
- ShareGPT4V-VQA：基于 ShareGPT-4V 數(shù)據(jù)集，選取了 caption 超過 300 個單詞的單論對話樣本，平均輸出 token 長度超過 1000。
我們以 PPL (Perplexity Metric) 指標評估模型生成內容的流暢度、以 METEOR (Metric for Evaluation of Translation with Explicit ORdering) 指標評估模型生成內容的質量。
表 3：生成能力基準比較。其中，解碼階段的 TFLOPs 和 Mem.（GPU 顯存占用）分別在無 / 有 KV Cache 的情況下測量得出。PPL 越低越好，METEOR 越高越好
如表 3 所示，相比 LLaVA-1.5，只進行視覺內容稀疏化的 Dynamic-LLaVA 的生成流暢度（PPL）和生成質量（METEOR）幾乎沒有變化；同時對視覺和文本進行稀疏化的 Dynamic-LLaVA，PPL 僅變高了 0.3，METEOR 甚至略有提升，而在推理效率上，在無 KV Cache 的解碼階段降低了～50% 的 TFLOPs，在有 KV Cache 的解碼階段降低了～50% 的 GPU 顯存占用。實驗結果充分表明，Dynamic-LLaVA 針對視覺和文本同時進行稀疏化，幾乎不影響實際生成能力，卻可以實現(xiàn)大幅的推理效率提升。
3.3 實際推理效率
表 4：Dynamic-LLaVA-13B 推理效率實測。其中，2K/4K 表示輸出的文本 token 數(shù)，所有結果均在一張 A100 (80G) 上測試得出，batch size 固定為 8?！啊?表示 GPU 顯存耗盡
在表 4 中，我們測試了多模態(tài)大模型實際推理的時間和 GPU 顯存占用。Dynamic-LLaVA 實現(xiàn)了更快的推理速度和更低的顯存占用。FastV 這種對預填充階段的視覺 token 進行剪枝的方法，隨著輸出長度的增長，推理效率也逐漸降低。而我們提出的 Dynamic-LLaVA，隨著輸出變長，相比于 FastV 的推理效率優(yōu)勢也逐漸顯現(xiàn)出來。
3.4 實例展示
圖 6：Dynamic-LLaVA-13B 在 LVIS-VQA (single-round) 上的推理結果展示。圖片的白色部分表示該位置的圖像塊被丟棄，文字中的灰色部分表示其在稀疏化過程中被丟棄，這表示它們不參與后續(xù)的自回歸解碼過程，但在模型的輸出中都被完整保留
圖 6 中展示了 Dynamic-LLaVA-13B 在 LVIS-VQA (single-round) 上的推理結果，以及對視覺和文本 token 的稀疏化情況。可視化結果表明，視覺 token 部分的主要信息得以保留；文本 token 中，一些不影響整體語義理解的連詞、介詞等被丟棄。這表明 Dynamic-LLaVA 能夠實現(xiàn)關鍵的視覺、語義信息的保留，從而保證了模型整體的性能。
4 總結
針對當前多模態(tài)大模型推理效率受限的問題，團隊通過分析多模態(tài)大模型推理過程中的不同階段，針對性的設計了推理加速方案。提出了 Dynamic-LLaVA—— 第一個同時稀疏化視覺和語言上下文的多模態(tài)大模型推理加速框架，將不同推理模式的推理效率優(yōu)化集成到統(tǒng)一框架中。
隨著多模態(tài)大模型技術的發(fā)展，尤其是其在復雜推理、長思維鏈領域的不斷進步。我們有理由相信，Dynamic-LLaVA 的應用場景正變得更加廣泛，其對輸出文本 token 進行稀疏化的模式，會在當前的更長輸出、更復雜推理的場景下，體現(xiàn)出更明顯的推理加速優(yōu)勢。
作者簡介
黃文軒：小紅書 NLP 團隊算法實習生，現(xiàn)碩士就讀于華東師范大學計算機科學與技術學院 2023 級。他在 ICLR、CVPR 等國際頂級會議上以第一作者身份發(fā)表了多篇學術論文，主要研究方向包括多模態(tài)大模型、大模型的高效訓練與推理等。
翟子杰：小紅書 NLP 團隊算法實習生，現(xiàn)碩士就讀于華東師范大學計算機科學與技術學院 2023 級。他在 ICML、ICLR、EMNLP 等國際頂級會議上發(fā)表過多篇學術論文，研究方向主要集中在多模態(tài)大模型、生成式搜索與推薦大模型等領域。
曹紹升：小紅書 NLP 團隊負責人，發(fā)表論文 30 余篇，授權專利 100 余項，引用近 4000 次，獲得 ICDE 2023 年最佳工業(yè)論文獎、CIKM 2015-2020 年最高引用論文、AAAI 2016 最具影響力論文。此外，還榮獲了中國發(fā)明協(xié)會創(chuàng)新成果一等獎（排名 1）、中國人工智能學會吳文俊科技進步二等獎（排名 1），連續(xù) 4 年入選世界人工智能學者榜單 AI-2000 新星榜前 100 名、Elsevier 中國區(qū)高被引學者，CCTV-13《新聞直播間》采訪報道。
葉哲宇：碩士畢業(yè)于帝國理工學院計算機專業(yè)，小紅書 NLP 團隊算法工程師，專注于大模型算法與應用方向，開源社區(qū) DMLC 成員。他在 ICLR、NAACL、EMNLP 等國際頂級會議上發(fā)表過多篇論文，研究領域涵蓋大模型應用、多模態(tài)大模型、Agent 模擬等。
林紹輝：華東師范大學計算機學院研究員，紫江青年學者，2021 年揚帆計劃獲得者，曾獲中國人工智能學會優(yōu)秀博士論文提名獎、《中國科學：技術科學》最佳審稿人。在國際頂級期刊和會議發(fā)表超過 50 篇論文，包括 TPAMI、TNNLS、TMI、CVPR、ECCV、AAAI、IJCAI 等。擔任 CVPR 2024 領域主席、IJCAI 2020 SPC 以及國際頂級期刊和會議審稿人。目前主要研究方向有計算機視覺、機器學習、圖像視頻理解、低層視覺等。
熱門跟貼