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

亮點直擊

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

總結速覽

解決的問題：

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

提出的方案：

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

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

應用的技術：

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

達到的效果：

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

方法

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

回顧 Mixture-of-Depths

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

在原始為 LLM 設計的 MoD 模塊中， X代表文本 tokens。在本文的多模態場景中， X代表視覺 tokens。僅在視覺 tokens 上應用 MoD，因為它們占用了主要的計算負擔，并且在深層中顯示出高冗余性。

改進 Mixture-of-Depths 模塊

與從頭訓練 MoD 基于 LLMs 的模型不同，將 MoD 集成到 MLLMs 中面臨幾個挑戰。將新的 MoD 模塊插入到預訓練的 LLM 中可能會導致訓練和推理的不穩定。相對較小的多模態數據量可能不足以訓練 MoD 模塊。兩種增強 MoD 模塊的設計，以應對這些挑戰。

Tanh 門控權重歸一化

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

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

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

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

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

對稱 Token 重加權

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

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

通過這種方式，MoD 模塊可以充分利用訓練過程中的語言監督信號，并在有限的訓練數據上學習準確評估 token 的重要性。

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

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

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

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

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

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

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

實驗

設置

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

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

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

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

實現細節

主要結果

將 p-MoD 集成到 LLaVA-1.5 和 LLaVA-NeXT 基線模型中，并在 14 個基準上進行了全面評估，如下表 1 和表 2 所示。與基線模型相比，p-MoD-LLaVA-1.5 和 p-MoD-LLaVA-NeXT 在所有基準上都實現了可比甚至更好的性能，同時顯著節省了推理 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 這樣的文本豐富和基于圖表的基準需要細粒度的視覺感知和推理能力。只有在模型成功識別出與答案相關的區域并且這些區域占整個圖像的比例較小時，才能給出正確的答案。值得一提的是，p-MoD-LLaVA-NeXT 在 DocVQA 和 InfoVQA 上幾乎沒有性能下降，在 ChartQA 上甚至提高了 0.2% 的準確率，同時在時間和內存效率上也取得了顯著改進。

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

性能與效率的權衡

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

消融研究

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

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

對稱 token 重新加權

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

漸進比例衰減

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

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

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

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

可視化

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

結論與未來工作

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

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

本文轉自AI生成未來 ，作者：AI生成未來

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