本文由華東師范大學(xué)和小紅書聯(lián)合完成，共同第一作者是華東師范大學(xué)在讀碩士、小紅書 NLP 團(tuán)隊(duì)實(shí)習(xí)生黃文軒和翟子杰，通訊作者是小紅書 NLP 團(tuán)隊(duì)負(fù)責(zé)人曹紹升，以及華東師范大學(xué)林紹輝研究員。

多模態(tài)大模型（MLLMs）在視覺理解與推理等領(lǐng)域取得了顯著成就。然而，隨著解碼（decoding）階段不斷生成新的 token，推理過程的計(jì)算復(fù)雜度和 GPU 顯存占用逐漸增加，這導(dǎo)致了多模態(tài)大模型推理效率的降低。現(xiàn)有的方法通過減少預(yù)填充（prefill）階段的視覺 token 冗余來實(shí)現(xiàn)推理加速。遺憾的是，這種在預(yù)填充階段實(shí)現(xiàn)的視覺 token 稀疏化所帶來的加速優(yōu)勢，在解碼階段會逐漸減弱。當(dāng)解碼輸出的文本 token 數(shù)量增多時(shí)，這些方法仍然會遇到性能瓶頸。

為了解決上述問題，團(tuán)隊(duì)創(chuàng)新性地提出了一個全新的動態(tài)視覺 - 文本上下文稀疏化推理加速框架 ——Dynamic-LLaVA。該框架針對多模態(tài)大模型在不同推理模式下（包括預(yù)填充階段以及有無 KV Cache 的解碼階段），設(shè)計(jì)了定制化的稀疏化推理方案，以實(shí)現(xiàn)多模態(tài)大模型的高效推理。實(shí)驗(yàn)結(jié)果表明，Dynamic-LLaVA 在幾乎不損失視覺理解和生成能力的前提下，能夠?qū)㈩A(yù)填充階段的計(jì)算開銷減少約 75%；在無 KV Cache 的解碼階段，計(jì)算開銷減少約 50%；在有 KV Cache 的解碼階段，GPU 顯存占用減少約 50%。Dynamic-LLaVA 為多模態(tài)大模型推理加速領(lǐng)域樹立了新的標(biāo)桿。

• 論文標(biāo)題：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
• 研究團(tuán)隊(duì)：華東師范大學(xué)、小紅書 NLP 團(tuán)隊(duì)

1 引言

1.1 前置信息：預(yù)填充與解碼

本文主要圍繞以 LLaVA 為范式的多模態(tài)大模型展開研究。一個多模態(tài)大模型的推理過程可以分為預(yù)填充和解碼兩個階段：

在預(yù)填充階段，不同模態(tài)的特征被映射到與大語言模型（LLM）輸入 embedding 相同的特征分布空間中。這些多模態(tài)特征與文本 token 會一起被大語言模型處理，以生成初始輸出文本 token。以圖片理解場景為例，該階段主要處理輸入的圖片和文本格式的問題。

在隨后的解碼階段，預(yù)填充階段生成的所有 token 以及后續(xù)生成的所有輸出文本 token，將被用于自回歸生成，從而產(chǎn)生完整的輸出。同樣以圖片理解場景為例，該階段生成針對整個問題的完整回答。

1.2 多模態(tài)大模型推理加速困境

圖 1：多模態(tài)大模型生成過程（有 / 無 KV Cache）中 FLOPs（計(jì)算復(fù)雜度）和 GPU 顯存開銷的增長趨勢

現(xiàn)有的多模態(tài)大模型大多以基于解碼器架構(gòu)的大語言模型（LLM）為核心，這些模型通常擁有龐大的參數(shù)規(guī)模。在生成輸出文本 token 的過程中，模型計(jì)算負(fù)擔(dān)會逐漸加重，導(dǎo)致對計(jì)算資源的巨大消耗。為了提升推理速度，現(xiàn)有模型通常會在解碼過程中運(yùn)用 KV Cache 技術(shù)，通過存儲并復(fù)用之前計(jì)算的 KV 激活值來減少重復(fù)計(jì)算。然而，如圖 1 (B) 所示，即使使用了 KV Cache，LLaVA 在輸出 token 不斷增加時(shí)，仍會迅速面臨 GPU 顯存耗盡的問題。

與文本不同，視覺信息往往包含大量冗余。因此，許多方法嘗試通過減少視覺上下文來加速多模態(tài)大模型的推理，即對預(yù)填充階段的視覺 token 進(jìn)行剪枝處理。但這種方法存在局限性：其主要提升了多模態(tài)大語言模型在預(yù)填充階段的推理效率，而在解碼階段，其效率提升會逐漸減弱。

如圖 1 (B) 和 (C) 所示，F(xiàn)astV 這種針對視覺 token 剪枝的方法，雖然相較于原始的 LLaVA 能夠節(jié)省一定的 GPU 顯存和計(jì)算開銷（FLOPs），但當(dāng)輸出 token 數(shù)接近 5K 時(shí)，它仍然會遭遇計(jì)算資源瓶頸。此外，F(xiàn)astV 和原始 LLaVA 的曲線斜率基本一致，這表明在長輸出的解碼階段，這類方法并沒有顯著的推理效率優(yōu)勢。因此，僅通過減少預(yù)填充階段的視覺 token，在輸出文本 token 數(shù)量遠(yuǎn)超視覺 token 時(shí)，難以實(shí)現(xiàn)整個推理效率的顯著提升。

1.3 邁向全階段推理加速：Dynamic-LLaVA

針對上述問題，我們認(rèn)為：為了實(shí)現(xiàn)真正的全階段推理加速，不僅需要對預(yù)填充階段的視覺 token 進(jìn)行剪枝，還必須對解碼階段輸出的文本 token 進(jìn)行稀疏化處理，限制參與自回歸運(yùn)算的 token 數(shù)量。為此，我們提出了 Dynamic-LLaVA，針對多模態(tài)大模型的視覺 - 語言上下文稀疏化推理加速框架。該框架能夠集成到多模態(tài)大模型推理的不同階段中，實(shí)現(xiàn)以下目標(biāo)：

• 顯著降低預(yù)填充階段計(jì)算開銷：通過優(yōu)化視覺 token 的處理方式，減少不必要的計(jì)算。
• 提升解碼階段的推理效率：無論是否使用 KV Cache，都能減少計(jì)算開銷，提高推理速度。
• 保持性能優(yōu)勢：在視覺理解任務(wù)上幾乎不損失性能；在長文本輸出場景中，生成能力也幾乎不受影響。

通過這些創(chuàng)新，Dynamic-LLaVA 為多模態(tài)大模型的高效推理提供了一種全新的解決方案。

2 方法

圖 2：Dynamic-LLaVA 整體框架

如圖 2 所示，Dynamic-LLaVA 可以集成到多模態(tài)大模型推理流程中的不同階段。具體而言，在預(yù)填充階段，該框架對視覺 token 執(zhí)行精準(zhǔn)剪枝操作，剔除冗余信息；在不使用 KV Cache 的解碼階段，限制參與自回歸運(yùn)算的視覺與輸出文本 token 數(shù)量，避免不必要的計(jì)算負(fù)擔(dān)；而在使用 KV Cache 的解碼階段，Dynamic-LLaVA 則動態(tài)調(diào)控 KV Cache，自適應(yīng)判斷是否將當(dāng)前輸出文本 token 的 KV 激活值納入 KV Cache，優(yōu)化資源利用效率。為了使模型適應(yīng)這種全新的稀疏化推理模式，Dynamic-LLaVA 在預(yù)訓(xùn)練的 LLaVA-1.5 基礎(chǔ)上進(jìn)行了 1 個 epoch 的監(jiān)督微調(diào)（SFT），確保模型能夠高效地運(yùn)行在稀疏化的推理路徑上。

2.1 預(yù)填充階段

在預(yù)填充階段，我們對輸入的視覺 token 進(jìn)行稀疏化操作。如圖 2 左側(cè)部分所示，我們引入一個可訓(xùn)練的輕量化的圖像預(yù)測器（Image Predictor），來判斷應(yīng)當(dāng)丟棄哪些視覺 token。該圖像預(yù)測器的結(jié)構(gòu)如下圖：

圖 3：圖像預(yù)測器的結(jié)構(gòu)示意圖

圖像預(yù)測器會對每個視覺 token 產(chǎn)生 “決策分?jǐn)?shù)”，以決定對哪些視覺 token 進(jìn)行保留。在端到端訓(xùn)練中，視覺 token 的剪枝通過 0-1 二值化的掩碼操作實(shí)現(xiàn)（具體過程見 2.4 節(jié)）。在實(shí)際推理階段中，通過保留 “決策分?jǐn)?shù)” 前 k 大的視覺 token（即圖 2 左側(cè)部分的 “Yes” 分支），實(shí)現(xiàn)視覺 token 數(shù)量減少，以實(shí)現(xiàn)推理加速。

2.2 解碼階段

不使用 KV Cache 的解碼過程：

對于視覺 token，采用和上一小節(jié)相同的做法，進(jìn)行稀疏化處理。

對于輸出的文本 token，分兩類進(jìn)行處理：

• 最后一個輸出的文本 token（即圖 2 中間部分的 “Last output text token”），不進(jìn)行任何處理，完整輸入 LLM 的 decoder 層進(jìn)行計(jì)算。這樣做的目的是保證模型的輸出內(nèi)容是連貫的，產(chǎn)生新的輸出文本 token 時(shí)，始終保證自回歸運(yùn)算包含上一個輸出文本 token。
• 對其他歷史的輸出文本 token 進(jìn)行稀疏化操作，其形式類似于對視覺 token 的處理。引入一個結(jié)構(gòu)如下圖的輸出預(yù)測器（Output Predictor），給出每個輸出文本 token 的 “決策分?jǐn)?shù)”，以決定當(dāng)前產(chǎn)生新的輸出內(nèi)容時(shí)，應(yīng)當(dāng)包括哪些文本 token 進(jìn)行自回歸運(yùn)算。圖 2 中間部分的 “Yes” 分支，表明保留的輸出文本 token。

圖 4：輸出預(yù)測器的結(jié)構(gòu)示意圖

使用 KV Cache 的解碼過程：

KV Cache 是節(jié)省冗余計(jì)算的一個關(guān)鍵推理加速技術(shù)，其思想是 “用 GPU 顯存的空間換計(jì)算時(shí)間”。顯而易見的是，KV Cache 也并非無限大，在長輸出情況下，必須丟棄一些 KV Cache 以適應(yīng)有限的 GPU 顯存。目前在 LLM 領(lǐng)域已有大量的 KV Cache 壓縮方案，以 方法為代表，這一類方法一般基于當(dāng)前 token 和歷史 KV Cache 進(jìn)行重要性分?jǐn)?shù)計(jì)算，以壓縮歷史 KV Cache。

與上述方法不同的是，我們對有 KV Cache 的解碼階段的設(shè)計(jì)，核心在于 “僅判斷當(dāng)前新 token 的 KV 激活是否需要加入 KV Cache 中”。如圖 2 右側(cè)所示，對于當(dāng)前正在處理的新 token（Last output text token），使用和上一部分結(jié)構(gòu)相同的輸出預(yù)測器，以決定是否加入 KV Cache 集合中。這種 “Online KV Cache 壓縮” 方法，判斷是否保留 KV Cache 的過程計(jì)算復(fù)雜度更低，也更加適應(yīng)多模態(tài)場景。在論文附錄中，我們詳細(xì)討論了我們的方法和現(xiàn)有的 LLM KV Cache 壓縮方法的區(qū)別。

需要特別說明的是，和不使用 KV Cache 的解碼階段相同，無論當(dāng)前處理的 token 是否加入 KV Cache，其都會輸入 LLM decoder 層進(jìn)行計(jì)算，以保證輸出的連貫性。

2.3 端到端訓(xùn)練

圖 5：Dynamic-LLaVA 在端到端訓(xùn)練過程中的示意圖

Dynamic-LLaVA 是一個需要訓(xùn)練的多模態(tài)大模型推理加速框架。我們基于 LLaVA 進(jìn)行了一個 epoch 的指令微調(diào)，以實(shí)現(xiàn)對 token 動態(tài)選擇的穩(wěn)定性，保證最終的性能。為了保證端到端訓(xùn)練，在訓(xùn)練階段的稀疏化操作通過 0-1 二值化掩碼實(shí)現(xiàn)（在推理中的實(shí)現(xiàn)是直接從歷史 token 序列中丟棄 token）。如圖 5 所示，上半部分表示訓(xùn)練中進(jìn)行 mask 的過程，在得到整個 token 序列的重要性分?jǐn)?shù)后，我們選取前 k 重要的 token 進(jìn)行保留，相對應(yīng)的生成掩碼向量，其中 0 對應(yīng)丟棄的冗余 token（不參與注意力過程的計(jì)算），1 對應(yīng)保留的重要 token，進(jìn)一步基于掩碼向量生成注意力過程的掩碼矩陣。掩碼矩陣用來對多頭注意力機(jī)制進(jìn)行掩碼操作，以確保丟棄的 token 不參與注意力過程的計(jì)算。由于二值化操作會導(dǎo)致不可微問題，所以我們借助了 GumbalSoftmax 和梯度直通估計(jì)器（Straight Through Estimator, STE）來保證梯度流的正確傳播，以進(jìn)行端到端的訓(xùn)練，如圖 5 下半部分所示。

3 實(shí)驗(yàn)

Dynamic-LLaVA 基于 LLaVA-1.5-7B 和 13B 的兩個版本進(jìn)行了 1 個 epoch 的指令微調(diào)，訓(xùn)練使用的數(shù)據(jù)和 LLaVA-1.5 相同。

3.1 視覺理解能力

我們首先評估了 Dynamic-LLaVA 在主要的視覺理解基準(zhǔn)的性能，選取了目前主流的多模態(tài)大模型推理加速方法進(jìn)行比較。

表 1：視覺理解基準(zhǔn)效果對比。其中，F(xiàn)ree 表示方法是否是 Training-Free 的。Dynamic-LLaVA 的下標(biāo) "I" 和 "I | T" 分別表示僅對視覺 token 做稀疏化和同時(shí)對視覺和文本 token 都做稀疏化（該標(biāo)識適用于下文所有的表格）

如表 1 所示，Dynamic-LLaVA 在大部分視覺理解任務(wù)上取得了優(yōu)越的性能。和其他對視覺內(nèi)容稀疏化的方法相比，Dynamic-LLaVA 在能大幅減小計(jì)算復(fù)雜度的同時(shí)，能夠?qū)崿F(xiàn)相比原始的 LLaVA-1.5 性能幾乎不下降。此外，在 SciQA、POPE、MME 和 MMBench 上，Dynamic-LLaVA 相比 LLaVA-1.5 甚至有一定的性能提升。例如，在 SciQA 任務(wù)上，Dynamic-LLaVA 的 7B 和 13B 版本，相較于 LLaVA-1.5 實(shí)現(xiàn)了 2.3% 和 0.8% 的性能提升。

表 2：與其他高效視覺 projector 的 SOTA 方法對比

值得一提的是，Dynamic-LLaVA 并沒有對 LLaVA-1.5 的視覺 projector 進(jìn)行修改，就可以實(shí)現(xiàn)大幅降低預(yù)填充階段計(jì)算復(fù)雜度，同時(shí)維持模型性能。在表 2 中，和其他針對視覺 projector 做高效設(shè)計(jì)（以提高推理效率）的 SOTA 方法進(jìn)行了對比。相較于其他使用了高效的視覺 projector 的方法，Dynamic-LLaVA 使用和 LLaVA-1.5 相同的 MLP 結(jié)構(gòu)作為視覺 projector，實(shí)現(xiàn)了更好的性能，同時(shí)也大幅降低了預(yù)填充階段的計(jì)算復(fù)雜度。此外，Dynamic-LLaVA 也可以和其他使用高效視覺 projector 的方法集成。例如，表 2 中 Dynamic-LLaVA 使用 TokenPacker 這一高效視覺 projector 的版本，在原始的 TokenPacker 方法基礎(chǔ)上，進(jìn)一步減少了視覺 token。相較于其他基于 TokenPacker 的推理加速方法，性能損失最少。

3.2 生成能力

現(xiàn)有的視覺理解任務(wù)中，一般只要求模型給出簡短的回復(fù)，這和現(xiàn)實(shí)世界中多模態(tài)大模型的應(yīng)用場景仍然存在不小的區(qū)別。在現(xiàn)實(shí)使用中，多模態(tài)大模型多數(shù)情況下會被要求生成更長、更細(xì)致的描述。為了和現(xiàn)實(shí)世界的場景對齊，評估 Dynamic-LLaVA 在更長的輸出情況下的生成能力和推理效率。我們額外構(gòu)建了兩個評估模型生成能力的基準(zhǔn)：

• LVIS-VQA：基于 LVIS-Instruct4 數(shù)據(jù)集，選取了 1000 個回答超過 100 個單詞的單輪對話樣本構(gòu)成 LVIS-VQA (single round) 和 1000 個多輪對話樣本（平均回答單詞數(shù)超過 300）構(gòu)成 LVIS-VQA (multi-round)；
• ShareGPT4V-VQA：基于 ShareGPT-4V 數(shù)據(jù)集，選取了 caption 超過 300 個單詞的單論對話樣本，平均輸出 token 長度超過 1000。

我們以 PPL (Perplexity Metric) 指標(biāo)評估模型生成內(nèi)容的流暢度、以 METEOR (Metric for Evaluation of Translation with Explicit ORdering) 指標(biāo)評估模型生成內(nèi)容的質(zhì)量。

表 3：生成能力基準(zhǔn)比較。其中，解碼階段的 TFLOPs 和 Mem.（GPU 顯存占用）分別在無 / 有 KV Cache 的情況下測量得出。PPL 越低越好，METEOR 越高越好

如表 3 所示，相比 LLaVA-1.5，只進(jìn)行視覺內(nèi)容稀疏化的 Dynamic-LLaVA 的生成流暢度（PPL）和生成質(zhì)量（METEOR）幾乎沒有變化；同時(shí)對視覺和文本進(jìn)行稀疏化的 Dynamic-LLaVA，PPL 僅變高了 0.3，METEOR 甚至略有提升，而在推理效率上，在無 KV Cache 的解碼階段降低了～50% 的 TFLOPs，在有 KV Cache 的解碼階段降低了～50% 的 GPU 顯存占用。實(shí)驗(yàn)結(jié)果充分表明，Dynamic-LLaVA 針對視覺和文本同時(shí)進(jìn)行稀疏化，幾乎不影響實(shí)際生成能力，卻可以實(shí)現(xiàn)大幅的推理效率提升。

3.3 實(shí)際推理效率

表 4：Dynamic-LLaVA-13B 推理效率實(shí)測。其中，2K/4K 表示輸出的文本 token 數(shù)，所有結(jié)果均在一張 A100 (80G) 上測試得出，batch size 固定為 8?！啊?表示 GPU 顯存耗盡

在表 4 中，我們測試了多模態(tài)大模型實(shí)際推理的時(shí)間和 GPU 顯存占用。Dynamic-LLaVA 實(shí)現(xiàn)了更快的推理速度和更低的顯存占用。FastV 這種對預(yù)填充階段的視覺 token 進(jìn)行剪枝的方法，隨著輸出長度的增長，推理效率也逐漸降低。而我們提出的 Dynamic-LLaVA，隨著輸出變長，相比于 FastV 的推理效率優(yōu)勢也逐漸顯現(xiàn)出來。

3.4 實(shí)例展示

圖 6：Dynamic-LLaVA-13B 在 LVIS-VQA (single-round) 上的推理結(jié)果展示。圖片的白色部分表示該位置的圖像塊被丟棄，文字中的灰色部分表示其在稀疏化過程中被丟棄，這表示它們不參與后續(xù)的自回歸解碼過程，但在模型的輸出中都被完整保留

圖 6 中展示了 Dynamic-LLaVA-13B 在 LVIS-VQA (single-round) 上的推理結(jié)果，以及對視覺和文本 token 的稀疏化情況?？梢暬Y(jié)果表明，視覺 token 部分的主要信息得以保留；文本 token 中，一些不影響整體語義理解的連詞、介詞等被丟棄。這表明 Dynamic-LLaVA 能夠?qū)崿F(xiàn)關(guān)鍵的視覺、語義信息的保留，從而保證了模型整體的性能。

4 總結(jié)

針對當(dāng)前多模態(tài)大模型推理效率受限的問題，團(tuán)隊(duì)通過分析多模態(tài)大模型推理過程中的不同階段，針對性的設(shè)計(jì)了推理加速方案。提出了 Dynamic-LLaVA—— 第一個同時(shí)稀疏化視覺和語言上下文的多模態(tài)大模型推理加速框架，將不同推理模式的推理效率優(yōu)化集成到統(tǒng)一框架中。

隨著多模態(tài)大模型技術(shù)的發(fā)展，尤其是其在復(fù)雜推理、長思維鏈領(lǐng)域的不斷進(jìn)步。我們有理由相信，Dynamic-LLaVA 的應(yīng)用場景正變得更加廣泛，其對輸出文本 token 進(jìn)行稀疏化的模式，會在當(dāng)前的更長輸出、更復(fù)雜推理的場景下，體現(xiàn)出更明顯的推理加速優(yōu)勢。

作者簡介

黃文軒：小紅書 NLP 團(tuán)隊(duì)算法實(shí)習(xí)生，現(xiàn)碩士就讀于華東師范大學(xué)計(jì)算機(jī)科學(xué)與技術(shù)學(xué)院 2023 級。他在 ICLR、CVPR 等國際頂級會議上以第一作者身份發(fā)表了多篇學(xué)術(shù)論文，主要研究方向包括多模態(tài)大模型、大模型的高效訓(xùn)練與推理等。

翟子杰：小紅書 NLP 團(tuán)隊(duì)算法實(shí)習(xí)生，現(xiàn)碩士就讀于華東師范大學(xué)計(jì)算機(jī)科學(xué)與技術(shù)學(xué)院 2023 級。他在 ICML、ICLR、EMNLP 等國際頂級會議上發(fā)表過多篇學(xué)術(shù)論文，研究方向主要集中在多模態(tài)大模型、生成式搜索與推薦大模型等領(lǐng)域。

曹紹升：小紅書 NLP 團(tuán)隊(duì)負(fù)責(zé)人，發(fā)表論文 30 余篇，授權(quán)專利 100 余項(xiàng)，引用近 4000 次，獲得 ICDE 2023 年最佳工業(yè)論文獎、CIKM 2015-2020 年最高引用論文、AAAI 2016 最具影響力論文。此外，還榮獲了中國發(fā)明協(xié)會創(chuàng)新成果一等獎（排名 1）、中國人工智能學(xué)會吳文俊科技進(jìn)步二等獎（排名 1），連續(xù) 4 年入選世界人工智能學(xué)者榜單 AI-2000 新星榜前 100 名、Elsevier 中國區(qū)高被引學(xué)者，CCTV-13《新聞直播間》采訪報(bào)道。

葉哲宇：碩士畢業(yè)于帝國理工學(xué)院計(jì)算機(jī)專業(yè)，小紅書 NLP 團(tuán)隊(duì)算法工程師，專注于大模型算法與應(yīng)用方向，開源社區(qū) DMLC 成員。他在 ICLR、NAACL、EMNLP 等國際頂級會議上發(fā)表過多篇論文，研究領(lǐng)域涵蓋大模型應(yīng)用、多模態(tài)大模型、Agent 模擬等。

林紹輝：華東師范大學(xué)計(jì)算機(jī)學(xué)院研究員，紫江青年學(xué)者，2021 年揚(yáng)帆計(jì)劃獲得者，曾獲中國人工智能學(xué)會優(yōu)秀博士論文提名獎、《中國科學(xué)：技術(shù)科學(xué)》最佳審稿人。在國際頂級期刊和會議發(fā)表超過 50 篇論文，包括 TPAMI、TNNLS、TMI、CVPR、ECCV、AAAI、IJCAI 等。擔(dān)任 CVPR 2024 領(lǐng)域主席、IJCAI 2020 SPC 以及國際頂級期刊和會議審稿人。目前主要研究方向有計(jì)算機(jī)視覺、機(jī)器學(xué)習(xí)、圖像視頻理解、低層視覺等。