能感知多模態分情況行事，Meta提出模態感知型專家混合

機器之心報道

混合專家，也得術業有專攻。

對於目前的混合模態基礎模型，常用的架構設計是融合特定模態的編碼器或解碼器，但這種方法存在侷限：無法整合不同模態的信息，也難以輸出包含多種模態的內容。

爲了克服這一侷限，Meta FAIR 的 Chameleon 團隊在近期的論文《Chameleon: Mixed-modal early-fusion foundation models》中提出了一種新的單一 Transformer 架構，它可以根據下一個 token 的預測目標，對由離散圖像和文本 token 組成的混合模態序列進行建模，從而在不同模態之間進行無縫推理和生成。

在約 10 萬億混合模態 token 上完成預訓練之後，Chameleon 表現出了適應廣泛的視覺和語言能力，能很好地處理多種不同的下游任務。Chameleon 在生成混合模態長回答任務的表現尤其亮眼，它甚至打敗了 Gemini 1.0 Pro 和 GPT-4V 等商用模型。然而對於 Chameleon 這樣各種模態會在模型訓練的早期混合起來的模型，想要拓展它的能力，需要投入大量算力。

基於以上問題，Meta FAIR 的團隊對路由式稀疏架構（routed sparse architecture）進行了一番研究探索，提出了 MoMa：模態感知型專家混合架構。

論文標題：MoMa: Efficient Early-Fusion Pre-training with Mixture of Modality-Aware Experts

論文地址：https://arxiv.org/pdf/2407.21770

之前已有研究表明，這類架構可以有效地擴展單模態的基礎模型的能力，也可以增強多模態對比學習模型的性能。但是，將其用於較早將各種模態融合的模型訓練還是一個機遇與挑戰並存的課題，還少有人研究。

該團隊的研究基於這一洞見：不同模態具有固有的異構性 —— 文本和圖像 token 具有不同的信息密度和冗餘模式。

在將這些 token 整合成統一的融合架構的同時，該團隊也提出通過整合針對具體模態的模塊來進一步優化該框架。該團隊將這一概念稱爲模態感知型稀疏性（modality-aware sparsity），簡稱 MaS；其能讓模型更好地捕獲每個模態的特徵，同時還能通過部分參數共享和注意力機制維持強大的跨模態整合性能。

之前的 VLMo、BEiT-3 和 VL-MoE 等研究已經採用了混合模態專家（MoME/mixture-of-modality-experts）方法來訓練視覺 - 語言編碼器和掩碼式語言建模，來自 FAIR 的研究團隊更進一步將 MoE 的可用範圍又推進了一步。

模型架構

早期融合

本文提出的新模型基於 Chameleon 的早期融合架構，其做法是在一個統一 Transformer 中，將圖像和文本表示成一系列離散 token。Chameleon 的核心是一個基於 Transformer 的模型，其會在圖像和文本 token 的組合序列上應用自注意力機制。這能讓該模型捕獲模態內和模態間的複雜關聯。該模型的訓練使用的目標是下一 token 預測目標，以自迴歸方式生成文本和圖像 token。

在 Chameleon 中，圖像的 token 化方案採用了一個學習型圖像分詞器，它將基於大小爲 8192 的 codebook 將 512 × 512 的圖像編碼成 1024 個離散 token。對於文本的分詞將使用一個詞表大小爲 65,536 的 BPE 分詞器，其中包含圖像 token。這種統一的分詞方法可以讓模型無縫處理圖像和文本 token 交織錯雜的任意序列。

藉助這種方法，新模型繼承了表徵統一、靈活性好、可擴展性高、支持端到端學習這些優點。

在此基礎上（圖 1a），爲了進一步提升早融合模型的效率和性能，該團隊還引入了模態感知型稀疏性技術。

寬度擴展：模態感知型混合專家

該團隊提出了一種寬度擴展方法：將模態感知型模塊稀疏性集成到前向模塊中，從而擴展標準混合專家（MoE）架構。

該方法基於這一洞見：不同模態的 token 有各自不同的特徵和信息密度。

通過爲每個模態構建不同的專家分組，可讓模型開發出專門的處理路徑，同時維持跨模態的信息整合能力。

圖 1b 展示了這種模態感知型專家混合（MoMa）的關鍵組件。簡單來說，先是對各個特定模態的專家進行分組，然後實現分層路由（分爲模態感知型路由和模態內路由），最後選擇專家。詳細過程參見原論文。

總體來說，對於一個輸入 token x，MoMa 模塊的形式化定義爲：

在 MoMa 計算之後，該團隊又進一步使用了殘差連接和 Swin Transformer 歸一化。

Mixture-of-Depths（MoD）

之前也有研究者探索將稀疏性引入深度維度，他們的做法要麼是隨機丟棄某些層，要麼就是使用可學習的路由器。

該團隊的做法參考了第二種方法，同時整合了近期提出的混合深度（MoD）技術。有關 MoD 的更多介紹可參閱機器之心報道《DeepMind 升級 Transformer，前向通過 FLOPs 最多可降一半》。

具體而言，如下圖所示，該團隊的做法是在每個 MoD 層中，在混合專家 （MoE）路由之前都集成 MoD，從而確保在模態分離之前，整批數據都能應用 MoD。

推理

在推理階段，我們不能直接使用 MoE 的專家選擇路由或 MoD 的層選擇路由，因爲在一批數據中進行 top-k（選擇前 k 個）選擇會破壞因果關係。

爲了保證推理的因果關係，受上述 MoD 論文的啓發，研究團隊引入了輔助路由器（auxiliary router），其作用是僅基於 token 的隱藏表徵預測該 token 被某個專家或層選中的可能性。

升級改造（Upcycling）

在優化表徵空間和路由機制方面，對於一個從頭開始訓練 MoE 架構，存在一個獨特的難題。該團隊發現：MoE 路由器負責爲每個專家劃分表徵空間。但是，在模型訓練的早期階段，這個表徵空間並非最優，這就會導致訓練得到的路由函數也是次優的。

爲了克服這一侷限，他們基於 Komatsuzaki 等人的論文《Sparse upcycling: Training mixture-of-experts from dense checkpoints》提出了一種升級改造方法。

具體來說，首先訓練一個每個模態都有一個 FFN 專家的架構。經過一些預先設定的步數之後，再對該模型進行升級改造，具體做法是：將每個特定模態的 FFN 轉換成一個專家選擇式 MoE 模塊，並將每個專家初始化爲第一階段訓練的專家。這裡會在保留前一階段的數據加載器狀態的同時重置學習率調度器，以確保第二階段的訓練能使用已刷新的數據。

爲了促進專家更加專業，該團隊還使用了 Gumbel 噪聲來增強 MoE 路由函數，從而使得新的路由器能以可微分的方式對專家進行採樣。

這種升級改造方法加上 Gumbel-Sigmoid 技術，可克服學習到的路由器的侷限性，從而提升新提出的模態感知型稀疏架構的性能。

效率優化

爲促進 MoMa 的分佈式訓練，該團隊採用了完全分片式數據並行（FSDP/Fully Sharded Data Parallel）。但是，相比於常規 MoE，該方法存在一些特有的效率難題，包括負載平衡問題和專家執行的效率問題。

對於負載平衡問題，該團隊開發了一種平衡的數據混合方法，可讓每臺 GPU 上的文本 - 圖像數據比例與專家比例保持一致。

對於專家執行的效率問題，該團隊探索了一些策略，可幫助提升不同模態的專家的執行效率：

將各個模態的專家限制爲同構的專家，並禁止將文本 token 路由到圖像專家，反之亦然；

使用模塊稀疏性（block sparsity）來提升執行效率；

當模態的數量有限時，按順序運行不同模態的專家。

由於實驗中每臺 GPU 處理的 token 都足夠多，因此即使使用多個分批次矩陣乘法，硬件利用率也不算大問題。因此，該團隊認爲對於當前規模的實驗環境而言，按順序執行的方法是比較好的選擇。

其它優化

爲了進一步提升吞吐量，該團隊還採用了其它一些優化技術。

其中包括降低梯度通信量、自動化的 GPU 核融合等一般優化操作，研究團隊還通過 torch.compile 實現了圖優化。

此外，他們還針對 MoMa 開發了一些優化技術，包括跨不同層複用模態 token 索引，以最高效地同步 CPU 和 GPU 之間的設備。

實驗

設置

實驗中使用的預訓練數據集和預處理過程與 Chameleon 一樣。爲了評估擴展性能，他們訓練模型使用的 token 數量超過 1 萬億。

表 1 給出了密集和稀疏模型的詳細配置情況。

不同計算層級的擴展性能

該團隊分析了不同模型在不同計算層級上的擴展性能，這些計算層級（FLOPs）相當於三種大小的密集模型：90M、435M 和 1.4B。

實驗結果表明，一個稀疏模型僅使用總 FLOPs 的 1/η 就能比肩同等 FLOPs 的密集模型的預訓練損失（η 表示預訓練加速因子）。

模態解綁

引入特定模態的專家分組可提高不同規模模型的預訓練效率，這對圖像模態尤其有益。如圖 3 所示，使用 1 個圖像專家和 1 個文本專家的 moe_1t1i 配置顯著優於相應的密集模型。

擴展每個模態分組的專家數量還能進一步提升模型性能。

混合深度與專家

該團隊觀察到，當採用 MoE 和 MoD 以及它們的組合形式時，訓練損失的收斂速度會得到提升。如圖 4 所示，向 moe_1t1i 架構添加 MoD（mod_moe_1t1i）可大幅提升不同模型大小的模型性能。

此外，在不同的模型大小和模態上，mod_moe_1t1i 能媲美甚至超過 moe_4t4i，這表明在深度維度上引入稀疏性也能有效提升訓練效率。

另一方面，還能看到堆疊 MoD 和 MoE 的收益會逐步下降。

擴展專家的數量

爲了研究擴展專家數量的影響，該團隊進行了進一步的消融實驗。他們探索了兩種場景：爲每種模態分配同等數量的專家（平衡）以及爲每種模態分配不同數量的專家（不平衡）。結果見圖 5。

對於平衡的設置，從圖 5a 可以看到，隨着專家數量提升，訓練損失會明顯下降。但文本和圖像損失表現出了不同的擴展模式。這表明每種模態的固有特性會導致出現不同的稀疏建模行爲。

對於不平衡的設置，圖 5b 比較了同等專家總數（8）的三種不同配置。可以看到，一個模態的專家越多，模型在該模態上的表現通常就越好。

升級改造

該團隊自然也驗證了前述的升級改造的效果。圖 6 比較了不同模型變體的訓練曲線。

結果表明，升級改造確實能進一步改善模型訓練：當第一個階段有 10k 步時，升級改造能帶來 1.2 倍的 FLOPs 收益；而當這個步數爲 20k 時，也有 1.16 倍的 FLOPs 收益。

此外，還能觀察到，隨着訓練推進，經過升級改造的模型與從頭開始訓練的模型之間的性能差距會不斷增大。

吞吐量分析

稀疏模型通常不能立即帶來性能增益，因爲稀疏模型會增加動態性和相關的數據平衡問題。爲了量化新提出的方法對訓練效率的影響，該團隊通常控制變量實驗比較了不同架構的訓練吞吐量。結果見表 2。

可以看到，相比於密集模型，基於模態的稀疏性能實現更好的質量 - 吞吐量權衡，並且能隨專家數量增長展現出合理的可擴展性。另一方面，儘管 MoD 變體取得了最好的絕對損失，但由於額外的動態性和不平衡性，它們的計算成本往往也更高。

推理時間性能

該團隊也評估了模型在留存的語言建模數據和下游任務上的表現。結果見表 3 和 4。

如表 3 所示，通過使用多個圖像專家，1.4B MoMa 1t1i 模型在大多數指標上都優於相應的密集模型，只有在 COCO 和 Flickr 上的圖像到文本條件困惑度指標例外。進一步擴展專家數量也能提升性能，其中 1.4B MoE 8x 在圖像到文本性能上達到了最佳。

此外，如表 4 所示，1.4B MoE 8x 這個模型還非常擅長文本到文本任務。1.4B MoMa 4t4i 在所有條件圖像困惑度指標上表現最佳，而其在大多數基準上的文本困惑度也非常接近 1.4B MoE 8x。

總體而言，在混合文本和圖像兩種模態的數據上，1.4B MoMa 4t4i 模型的建模結果最好。

更多詳細內容，請閱讀原論文。