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

亮點直擊

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

總結(jié)速覽

解決的問題：

• 高計算成本：盡管多模態(tài)大語言模型（MLLMs）在多任務(wù)中表現(xiàn)出色，但其龐大的訓(xùn)練和推理成本嚴(yán)重阻礙了其發(fā)展。計算的主要負(fù)擔(dān)來自于通過LLM處理的大量視覺token。
• 視覺token冗余：在處理高分辨率圖像和視頻時，視覺token數(shù)量急劇增加，導(dǎo)致計算成本暴漲，限制了MLLMs的進(jìn)一步應(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策略：漸進(jìn)式比例衰減（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的保留比例，進(jìn)一步提高計算效率。

達(dá)到的效果：

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

方法

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

回顧 Mixture-of-Depths

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

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

改進(jì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 層中被正（或負(fù)）權(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í)的部分）來準(zhǔn)確且穩(wěn)健地評估視覺 tokens 的重要性并分配相應(yīng)的權(quán)重是具有挑戰(zhàn)性的。權(quán)重預(yù)測器模塊的梯度來自于 token 重加權(quán)過程。原始 MoD 僅對選中的 tokens 進(jìn)行 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 可以進(jìn)一步修改為：

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

漸進(jìn)式比例衰減（Progressive Ratio Decay）

在通過 TanhNorm 和 STRing 模塊升級 MoD 后，成功地訓(xùn)練了基于 MoD 的 MLLMs。然而，模型的性能與效率之間的權(quán)衡仍遠(yuǎn)未令人滿意。當(dāng) 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 上進(jìn)行了一系列探索性實驗，在所有層中設(shè)置 70% 的 token 保留比例。如下圖 3a 所示，首先將層劃分為若干組。在每個實驗中，僅對其中一組層應(yīng)用較低的 token 保留比例（15%），而保持其他層不變，并在 14 個基準(zhǔn)上評估該模型的表現(xiàn)。圖 3b 中的結(jié)果強烈支持了我們的假設(shè)：將較低的 token 保留比例應(yīng)用于越深的層，性能越好。這表明，視覺 tokens 在較深的層中具有更高的冗余性，這與之前的研究 [4, 45] 的觀察結(jié)果一致。

本文設(shè)計了一個漸進(jìn)式比例衰減（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ù)進(jìn)行微調(diào)，但 LLaVA 團(tuán)隊發(fā)布的 LLaVA-NeXT 訓(xùn)練數(shù)據(jù)包含 779k 數(shù)據(jù)。為了公平比較，我們使用 LLaVA 團(tuán)隊發(fā)布的 779k 數(shù)據(jù)和官方 LLaVA-NeXT 代碼庫訓(xùn)練 LLaVA-NeXT 模型。

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

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

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

主要結(jié)果

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

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

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

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

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

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

消融研究

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

1. 零中心化歸一化：我們將 TanhNorm 的零中心化特性與 Softmax 歸一化函數(shù)進(jìn)行比較，包括標(biāo)準(zhǔn)和偏移形式：
   從上表 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 基準(zhǔn)測試中分別提高了 1%，且平均性能提高了 0.3%。這表明 STRing 通過更有效地利用語言監(jiān)督信號來訓(xùn)練一個更好的權(quán)重預(yù)測器，尤其是在訓(xùn)練數(shù)據(jù)有限的情況下。

漸進(jìn)比例衰減

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

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

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

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

可視化

下圖 5 展示了 p-MoD-LLaVA-NeXT 在不同層次上選擇的 tokens。通過漸進(jìn)比例衰減策略的引導(dǎo)，MoD 層處理的 tokens 數(shù)量逐層逐漸減少。從中可以觀察到，選中的視覺 tokens 在圖像中逐步集中于具有豐富語義信息的關(guān)鍵區(qū)域。具體而言，在上面的圖中，位于不同位置的文本和圖形所對應(yīng)的視覺 tokens 幾乎都在第 26 層以較低的比例被選中。在下面的圖中，選中的 tokens 漸漸集中于標(biāo)尺上的度量標(biāo)記和其他物體。這些結(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ù)測模塊。最后，模型采用了漸進(jìn)比例衰減策略，逐層減少 token 保留比例。p-MoD 模型在 14 項不同基準(zhǔn)測試中，取得了與基線模型相當(dāng)或更優(yōu)的結(jié)果，并且在訓(xùn)練和推理效率上有了顯著提升。希望 p-MoD 能成為未來高效 MLLMs 研究的強大且高效的基線。

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

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

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