文章鏈接：https://arxiv.org/pdf/2412.04449
項目鏈接：https://github.com/MCG-NJU/p-MoD

亮點直擊

• MoD機制的引入：通過選擇性處理重要視覺token，顯著減少計算負擔(dān)，提升模型效率。
• 創(chuàng)新的TanhNorm與STRing設(shè)計：解決了MoD模塊在多模態(tài)訓(xùn)練中的穩(wěn)定性和性能問題，即使在有限數(shù)據(jù)下也能精準評估token的重要性。
• PRD策略：通過逐層減少保留token的比例，針對視覺token的冗余問題提供了有效的解決方案，大幅提高計算效率。
• 全面的實驗驗證：在多個基準測試中，p-MoD模型不僅保證了高效性，還表現(xiàn)出了超越傳統(tǒng)基線的性能。

總結(jié)速覽

解決的問題：

• 高計算成本：盡管多模態(tài)大語言模型（MLLMs）在多任務(wù)中表現(xiàn)出色，但其龐大的訓(xùn)練和推理成本嚴重阻礙了其發(fā)展。計算的主要負擔(dān)來自于通過LLM處理的大量視覺token。
• 視覺token冗余：在處理高分辨率圖像和視頻時，視覺token數(shù)量急劇增加，導(dǎo)致計算成本暴漲，限制了MLLMs的進一步應(yīng)用。

提出的方案：

• MoD機制：通過引入深度混合（Mixture-of-Depths, MoD）機制，選擇每個Transformer層要處理的最重要的視覺token，同時跳過冗余的token，從而提高計算效率。
• 創(chuàng)新設(shè)計：

• TanhNorm：用于權(quán)重歸一化的tanh門控方法，確保MoD模塊的正確初始化，增強訓(xùn)練穩(wěn)定性并解決推理時的數(shù)值穩(wěn)定性問題。
• STRing：對稱token重加權(quán)模塊，利用有限的訓(xùn)練數(shù)據(jù)和語言監(jiān)督信號來精確評估token的重要性。
• PRD策略：漸進式比例衰減（Progressive Ratio Decay, PRD）策略，基于移位余弦調(diào)度逐層減少保留token的比例，優(yōu)化視覺token的冗余問題。

應(yīng)用的技術(shù)：

• MoD機制：最初用于大語言模型的MoD機制被改造為適用于多模態(tài)語言模型。
• TanhNorm與STRing：用于解決MoD模塊的初始化與訓(xùn)練穩(wěn)定性問題，提升有限數(shù)據(jù)下的token重要性預(yù)測。
• PRD策略：通過逐層減少視覺token的保留比例，進一步提高計算效率。

達到的效果：

• 高效推理：在推理時，p-MoD模型的TFLOPs僅為基線模型的55.6%，KV緩存存儲僅為53.8%。
• 高效訓(xùn)練：訓(xùn)練時，p-MoD僅使用77.7%的GPU時間，顯著減少了計算成本。
• 優(yōu)異表現(xiàn)：在14個基準測試中，p-MoD模型的表現(xiàn)與LLaVA-NeXT基線相當，甚至超越了基線，證明了高效性和性能的平衡。

方法

本節(jié)介紹 p-MoD 模型。如下圖 2 所示，本文的模型由 p-MoD 層組成，這些層通過 tanh 門控權(quán)重歸一化（TanhNorm）和對稱 token 重加權(quán)（STRing）來升級 MoD 架構(gòu)。關(guān)鍵的設(shè)計是漸進式比例衰減（PRD）策略，它控制不同層之間的 token 保留比例。接下來的部分中，將首先簡要回顧 MoD，然后詳細解釋p-MoD 模型的每個組件。

回顧 Mixture-of-Depths

在經(jīng)過 transformer 層處理后，選中的 tokens 會根據(jù)其對應(yīng)的權(quán)重進行縮放。通過這種方式，權(quán)重預(yù)測器參與到梯度路徑中，使其參數(shù)能夠通過反向傳播進行更新。這個過程稱為 token 重加權(quán)。

在原始為 LLM 設(shè)計的 MoD 模塊中， X代表文本 tokens。在本文的多模態(tài)場景中， X代表視覺 tokens。僅在視覺 tokens 上應(yīng)用 MoD，因為它們占用了主要的計算負擔(dān)，并且在深層中顯示出高冗余性。

改進 Mixture-of-Depths 模塊

與從頭訓(xùn)練 MoD 基于 LLMs 的模型不同，將 MoD 集成到 MLLMs 中面臨幾個挑戰(zhàn)。將新的 MoD 模塊插入到預(yù)訓(xùn)練的 LLM 中可能會導(dǎo)致訓(xùn)練和推理的不穩(wěn)定。相對較小的多模態(tài)數(shù)據(jù)量可能不足以訓(xùn)練 MoD 模塊。兩種增強 MoD 模塊的設(shè)計，以應(yīng)對這些挑戰(zhàn)。

Tanh 門控權(quán)重歸一化

原始的 MoD 機制是為從頭訓(xùn)練 MoD 基于 LLMs 的模型設(shè)計的。在多模態(tài)訓(xùn)練場景中，需要將新的 MoD 模塊插入到預(yù)訓(xùn)練的原始 LLM 中。一種常見的方法是交替訓(xùn)練，其中在第一階段僅訓(xùn)練新插入的模塊，在第二階段同時訓(xùn)練所有模塊。然而，這種方法不適合 MoD，因為 MoD 模塊在 token 重加權(quán)過程中通過其權(quán)重來縮放 tokens，而后續(xù)被凍結(jié)的 LLM 層無法處理被縮放的 tokens。本文的實驗結(jié)果一致表明，訓(xùn)練 MoD 模塊并凍結(jié) LLM 層無法收斂。因此，唯一可行的方法是直接解決這一挑戰(zhàn)：將 MoD 模塊插入到 LLM 中，并同時訓(xùn)練它們。

這里  是一個超參數(shù)，用于控制歸一化權(quán)重的范圍。雖然很簡單，但 TanhNorm 確保了以下幾點：

這些特性帶來了以下好處：

• 歸一化的權(quán)重在訓(xùn)練開始時接近零，這確保了通過插入 MoD 模塊后，LLM 的正確訓(xùn)練初始化保持不變。
• 零中心的權(quán)重分布降低了某些 tokens 在所有 MoD 層中被正（或負）權(quán)重重復(fù)縮放的風(fēng)險。這確保了訓(xùn)練的穩(wěn)定性，并減輕了推理過程中數(shù)值穩(wěn)定性的問題。

上述兩個好處有助于提高模型性能。在第 4.4 節(jié)中，我們驗證了 TanhNorm 的有效性。

對稱 Token 重加權(quán)

與用于訓(xùn)練 MoD 基于 LLMs 的文本數(shù)據(jù)相比，MLLMs 是在規(guī)模小幾個數(shù)量級的多模態(tài)數(shù)據(jù)上訓(xùn)練的。訓(xùn)練一個權(quán)重預(yù)測器（MoD 模塊中唯一可學(xué)習(xí)的部分）來準確且穩(wěn)健地評估視覺 tokens 的重要性并分配相應(yīng)的權(quán)重是具有挑戰(zhàn)性的。權(quán)重預(yù)測器模塊的梯度來自于 token 重加權(quán)過程。原始 MoD 僅對選中的 tokens 進行 token 重加權(quán)，如方程 2 所示。這樣，語言監(jiān)督信號僅監(jiān)督權(quán)重預(yù)測器為選中的 tokens 分配高權(quán)重，而對跳過的 tokens 的權(quán)重預(yù)測過程沒有監(jiān)督（即沒有梯度）。

為了充分利用有限的訓(xùn)練數(shù)據(jù)，通過對選中和跳過的 tokens 對稱地應(yīng)用 token 重加權(quán)過程來增強 MoD 機制。通過我們對稱的 token 重加權(quán)（STRing）模塊，方程 3 可以進一步修改為：

通過這種方式，MoD 模塊可以充分利用訓(xùn)練過程中的語言監(jiān)督信號，并在有限的訓(xùn)練數(shù)據(jù)上學(xué)習(xí)準確評估 token 的重要性。

漸進式比例衰減（Progressive Ratio Decay）

在通過 TanhNorm 和 STRing 模塊升級 MoD 后，成功地訓(xùn)練了基于 MoD 的 MLLMs。然而，模型的性能與效率之間的權(quán)衡仍遠未令人滿意。當 token 保留比例設(shè)置低于 70% 時，模型的性能急劇下降。

原始的 MoD 基于 LLMs 在所有 MoD 層中采用固定的 token 保留比例。這一策略假設(shè) tokens 在不同層中的冗余程度相同。在多模態(tài)場景中這一假設(shè)是不成立的。在每個 transformer 層中的自注意力作用下，視覺 tokens 會逐漸相互聚合信息，文本 tokens 也會從視覺 tokens 中獲取相關(guān)信息。因此，隨著層數(shù)的加深，視覺 tokens 應(yīng)該變得越來越冗余。

為了驗證本文的假設(shè)，在基于 MoD 的 MLLM 上進行了一系列探索性實驗，在所有層中設(shè)置 70% 的 token 保留比例。如下圖 3a 所示，首先將層劃分為若干組。在每個實驗中，僅對其中一組層應(yīng)用較低的 token 保留比例（15%），而保持其他層不變，并在 14 個基準上評估該模型的表現(xiàn)。圖 3b 中的結(jié)果強烈支持了我們的假設(shè)：將較低的 token 保留比例應(yīng)用于越深的層，性能越好。這表明，視覺 tokens 在較深的層中具有更高的冗余性，這與之前的研究 [4, 45] 的觀察結(jié)果一致。

本文設(shè)計了一個漸進式比例衰減（PRD）策略。如前圖 2 右側(cè)所示，PRD 按照 shifted cosine 調(diào)度逐層減少 token 保留比例。假設(shè)模型有 L 層，第 l 層的 token 保留比例可以表示為：

這里， 是一個平移因子，可以通過垂直平移余弦衰減曲線來靈活控制模型的整體計算成本。

在實驗中，證明了使用 PRD 策略的 p-MoD 模型在相同計算預(yù)算下，顯著優(yōu)于使用固定保留比例的模型。它還優(yōu)于其他類型的比例調(diào)度器。

實驗

設(shè)置

模型
為了評估 p-MoD 的有效性，選擇了兩個代表性的開源 MLLM：LLaVA-1.5  和 LLaVA-NeXT，作為實驗中的基線模型。將對應(yīng)的 p-MoD 模型命名為 p-MoD-LLaVA-1.5 和 p-MoD-LLaVA-NeXT。LLaVA-1.5 將輸入圖像調(diào)整為固定分辨率，這與其 CLIP 視覺編碼器的訓(xùn)練設(shè)置對齊，將圖像編碼為 576 個 token。LLaVA-NeXT 將高分辨率圖像分成多個片段，這些片段由視覺編碼器獨立處理。這一策略增強了視覺感知能力，但也導(dǎo)致了更多的視覺 token（最多 2880 個）和更高的計算成本。

訓(xùn)練數(shù)據(jù)
遵循基線方法的訓(xùn)練設(shè)置。LLaVA-1.5 和 LLaVA-NeXT 都使用 LCS-558K 作為預(yù)訓(xùn)練數(shù)據(jù)。在指令調(diào)優(yōu)階段，LLaVA-1.5 使用 LLaVA-mix-665。原始的 LLaVA-NeXT 模型使用了 760k 數(shù)據(jù)進行微調(diào)，但 LLaVA 團隊發(fā)布的 LLaVA-NeXT 訓(xùn)練數(shù)據(jù)包含 779k 數(shù)據(jù)。為了公平比較，我們使用 LLaVA 團隊發(fā)布的 779k 數(shù)據(jù)和官方 LLaVA-NeXT 代碼庫訓(xùn)練 LLaVA-NeXT 模型。

基準
在 14 個基準上進行了全面的實驗：GQA、OK-VQA、AI2D 和 ScienceQA-IMG 是傳統(tǒng)的視覺問答基準；DocVQA、TextVQA、ChartQA 和 InfographicVQA 聚焦于細粒度的視覺問答；POPE 評估 MLLMs 的幻覺；SEED-Bench（圖像）、RealWorldQA、MME、MMBench 和 MMMU 是為 MLLMs 定制的綜合性基準。為了確保本文的結(jié)果可以方便地復(fù)現(xiàn)，我們使用 LLaVA 團隊開發(fā)和維護的 lmms-eval 評估框架評估本文的模型。

評估指標
本文報告所有基準的性能，并結(jié)合反映訓(xùn)練和推理效率的多種指標。我們報告了 TFLOPs，它衡量推理計算復(fù)雜度，以及 KV 緩存存儲，這是推理過程中的主要內(nèi)存瓶頸。此外，我們還報告了模型訓(xùn)練中的 GPU 小時數(shù)和推理延遲。

實現(xiàn)細節(jié)

主要結(jié)果

將 p-MoD 集成到 LLaVA-1.5 和 LLaVA-NeXT 基線模型中，并在 14 個基準上進行了全面評估，如下表 1 和表 2 所示。與基線模型相比，p-MoD-LLaVA-1.5 和 p-MoD-LLaVA-NeXT 在所有基準上都實現(xiàn)了可比甚至更好的性能，同時顯著節(jié)省了推理 TFLOPs 和 KV 緩存存儲。

表 1 包含了一般的 VQA 和 MLLM 基準，衡量了 MLLM 在各個方面的綜合能力。我們的 p-MoD-LLaVA-NeXT 模型在顯著減少推理 TFLOPs 和 KV 緩存存儲的同時，達到了可比的性能。我們的 p-MoD-LLaVA-1.5 模型甚至超過了 LLaVA-1.5 基線模型，推理 TFLOPs 為 58.7%，KV 緩存存儲為 53.8%。

在表 2 中，值得注意的是，像 DocVQA、ChartQA、TextVQA 和 InfoVQA 這樣的文本豐富和基于圖表的基準需要細粒度的視覺感知和推理能力。只有在模型成功識別出與答案相關(guān)的區(qū)域并且這些區(qū)域占整個圖像的比例較小時，才能給出正確的答案。值得一提的是，p-MoD-LLaVA-NeXT 在 DocVQA 和 InfoVQA 上幾乎沒有性能下降，在 ChartQA 上甚至提高了 0.2% 的準確率，同時在時間和內(nèi)存效率上也取得了顯著改進。

以上結(jié)果表明，本文的方法在保持性能的同時，顯著提高了效率。

性能與效率的權(quán)衡

總之，表 3 展示了一個趨勢：隨著計算成本的降低，p-MoD 的性能線性下降，下降幅度較淺，提供了卓越的性能-效率權(quán)衡。

消融研究

Tanh-gated 權(quán)重歸一化通過與不同的權(quán)重歸一化方法進行比較來驗證 TanhNorm 的有效性，如下表 4 所示。首先，確認使用沒有權(quán)重歸一化模塊會導(dǎo)致訓(xùn)練過程中出現(xiàn)溢出，如第一行所示。然后，探討 TanhNorm 的兩個重要特性：

1. 零中心化歸一化：我們將 TanhNorm 的零中心化特性與 Softmax 歸一化函數(shù)進行比較，包括標準和偏移形式：
   從上表 4 的最后一行可以看出，TanhNorm 顯著優(yōu)于原始和偏移 softmax 函數(shù)。這證明了零中心化是 TanhNorm 強大性能的關(guān)鍵，因為盡管偏移 softmax 函數(shù)與 TanhNorm 具有相同的范圍，但由于缺乏零中心化特性，表現(xiàn)較差。

對稱 token 重新加權(quán)

對稱 token 重新加權(quán)（STRing）策略的影響，如前表 4 中的最后兩行所示。STRing 在 DocVQA 和 RWQA 基準測試中分別提高了 1%，且平均性能提高了 0.3%。這表明 STRing 通過更有效地利用語言監(jiān)督信號來訓(xùn)練一個更好的權(quán)重預(yù)測器，尤其是在訓(xùn)練數(shù)據(jù)有限的情況下。

漸進比例衰減

為了驗證漸進比例衰減（PRD）策略的有效性，進行了四個不同調(diào)度函數(shù)的實驗，這些函數(shù)控制解碼器層中的 token 保留比例。這些函數(shù)包括：

• 常數(shù)函數(shù)：對所有層使用相同的比例。
• 交替函數(shù)：在標準 transformer 層和具有較低保留比例的 MoD 層之間交替，原 MoD 論文中采用此方法。
• 階梯衰減函數(shù)：以階梯方式減少保留比例。
• 線性衰減函數(shù)：逐漸減少保留比例。

在本文的實驗中，將所有層的平均保留比例固定在約 54%。下表 5 中的結(jié)果得出以下結(jié)論：

• 非衰減函數(shù)，如常數(shù)和交替函數(shù)，顯著低于衰減函數(shù)。這證實了視覺 token 在較深層中存在更高的冗余性，支持漸進衰減的必要性。
• 連續(xù)衰減函數(shù)，如余弦和線性衰減，顯著優(yōu)于不連續(xù)的階梯衰減函數(shù)。在這些方法中，余弦衰減函數(shù)取得了最佳結(jié)果，證明了我們采用的 shifted cosine 調(diào)度在優(yōu)化層間 token 保留比例方面的有效性。

可視化

下圖 5 展示了 p-MoD-LLaVA-NeXT 在不同層次上選擇的 tokens。通過漸進比例衰減策略的引導(dǎo)，MoD 層處理的 tokens 數(shù)量逐層逐漸減少。從中可以觀察到，選中的視覺 tokens 在圖像中逐步集中于具有豐富語義信息的關(guān)鍵區(qū)域。具體而言，在上面的圖中，位于不同位置的文本和圖形所對應(yīng)的視覺 tokens 幾乎都在第 26 層以較低的比例被選中。在下面的圖中，選中的 tokens 漸漸集中于標尺上的度量標記和其他物體。這些結(jié)果表明，我們的模型通過保留有信息的 tokens 有效地壓縮了視覺信息，從而在顯著降低訓(xùn)練和推理成本的同時，保持了良好的性能。

結(jié)論與未來工作

本文探討了通過 Mixture-of-Depths 機制構(gòu)建高效的 MLLMs（多模態(tài)語言模型）。提出命名為 p-MoD的模型，具有三個關(guān)鍵設(shè)計。首先，應(yīng)用了 tanh-gated 權(quán)重歸一化機制，以確保預(yù)測 token 權(quán)重在訓(xùn)練初始化和推理過程中保持穩(wěn)定。其次，采用了對稱 token 重新加權(quán)策略，充分利用語言監(jiān)督信號，通過有限的訓(xùn)練數(shù)據(jù)訓(xùn)練出更優(yōu)的權(quán)重預(yù)測模塊。最后，模型采用了漸進比例衰減策略，逐層減少 token 保留比例。p-MoD 模型在 14 項不同基準測試中，取得了與基線模型相當或更優(yōu)的結(jié)果，并且在訓(xùn)練和推理效率上有了顯著提升。希望 p-MoD 能成為未來高效 MLLMs 研究的強大且高效的基線。

本研究的一個限制是 p-MoD 僅在 LLaVA-1.5 和 LLaVA-NeXT 模型上進行了實驗，這些模型專注于單圖像理解任務(wù)。認為p-MoD 在處理更多視覺 token 的任務(wù)中，如高分辨率、多圖像和長視頻理解任務(wù)中，可能會取得更顯著的結(jié)果。將把 p-MoD 應(yīng)用到其他視覺任務(wù)中的探索留待未來的研究。

本文轉(zhuǎn)自AI生成未來 ，作者：AI生成未來

原文鏈接:??https://mp.weixin.qq.com/s/DhCPVbneK0adJVPx15Q3HQ??