對于目前的混合模態(tài)基礎模型，常用的架構(gòu)設計是融合特定模態(tài)的編碼器或解碼器，但這種方法存在局限：無法整合不同模態(tài)的信息，也難以輸出包含多種模態(tài)的內(nèi)容。

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

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

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

• 論文標題：MoMa: Efficient Early-Fusion Pre-training with Mixture of Modality-Aware Experts
• 論文地址：https://arxiv.org/pdf/2407.21770

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

該團隊的研究基于這一洞見：不同模態(tài)具有固有的異構(gòu)性 —— 文本和圖像 token 具有不同的信息密度和冗余模式。

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

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

模型架構(gòu)

早期融合

本文提出的新模型基于 Chameleon 的早期融合架構(gòu)，其做法是在一個統(tǒng)一 Transformer 中，將圖像和文本表示成一系列離散 token。Chameleon 的核心是一個基于 Transformer 的模型，其會在圖像和文本 token 的組合序列上應用自注意力機制。這能讓該模型捕獲模態(tài)內(nèi)和模態(tài)間的復雜關聯(lián)。該模型的訓練使用的目標是下一 token 預測目標，以自回歸方式生成文本和圖像 token。

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

借助這種方法，新模型繼承了表征統(tǒng)一、靈活性好、可擴展性高、支持端到端學習這些優(yōu)點。

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

寬度擴展：模態(tài)感知型混合專家

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

該方法基于這一洞見：不同模態(tài)的 token 有各自不同的特征和信息密度。

通過為每個模態(tài)構(gòu)建不同的專家分組，可讓模型開發(fā)出專門的處理路徑，同時維持跨模態(tài)的信息整合能力。

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

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

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

Mixture-of-Depths（MoD）

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

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

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

推理

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

為了保證推理的因果關系，受上述 MoD 論文的啟發(fā)，研究團隊引入了輔助路由器（auxiliary router），其作用是僅基于 token 的隱藏表征預測該 token 被某個專家或?qū)舆x中的可能性。

升級改造（Upcycling）

在優(yōu)化表征空間和路由機制方面，對于一個從頭開始訓練 MoE 架構(gòu)，存在一個獨特的難題。該團隊發(fā)現(xiàn)：MoE 路由器負責為每個專家劃分表征空間。但是，在模型訓練的早期階段，這個表征空間并非最優(yōu)，這就會導致訓練得到的路由函數(shù)也是次優(yōu)的。

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

具體來說，首先訓練一個每個模態(tài)都有一個 FFN 專家的架構(gòu)。經(jīng)過一些預先設定的步數(shù)之后，再對該模型進行升級改造，具體做法是：將每個特定模態(tài)的 FFN 轉(zhuǎn)換成一個專家選擇式 MoE 模塊，并將每個專家初始化為第一階段訓練的專家。這里會在保留前一階段的數(shù)據(jù)加載器狀態(tài)的同時重置學習率調(diào)度器，以確保第二階段的訓練能使用已刷新的數(shù)據(jù)。

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

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

效率優(yōu)化

為促進 MoMa 的分布式訓練，該團隊采用了完全分片式數(shù)據(jù)并行（FSDP/Fully Sharded Data Parallel）。但是，相比于常規(guī) MoE，該方法存在一些特有的效率難題，包括負載平衡問題和專家執(zhí)行的效率問題。

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

對于專家執(zhí)行的效率問題，該團隊探索了一些策略，可幫助提升不同模態(tài)的專家的執(zhí)行效率：

• 將各個模態(tài)的專家限制為同構(gòu)的專家，并禁止將文本 token 路由到圖像專家，反之亦然；
• 使用模塊稀疏性（block sparsity）來提升執(zhí)行效率；
• 當模態(tài)的數(shù)量有限時，按順序運行不同模態(tài)的專家。

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

其它優(yōu)化

為了進一步提升吞吐量，該團隊還采用了其它一些優(yōu)化技術。

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

此外，他們還針對 MoMa 開發(fā)了一些優(yōu)化技術，包括跨不同層復用模態(tài) token 索引，以最高效地同步 CPU 和 GPU 之間的設備。

實驗

設置

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

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

不同計算層級的擴展性能

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

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

模態(tài)解綁

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

擴展每個模態(tài)分組的專家數(shù)量還能進一步提升模型性能。

混合深度與專家

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

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

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

擴展專家的數(shù)量

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

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

對于不平衡的設置，圖 5b 比較了同等專家總數(shù)（8）的三種不同配置?？梢钥吹剑粋€模態(tài)的專家越多，模型在該模態(tài)上的表現(xiàn)通常就越好。

升級改造

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

結(jié)果表明，升級改造確實能進一步改善模型訓練：當?shù)谝粋€階段有 10k 步時，升級改造能帶來 1.2 倍的 FLOPs 收益；而當這個步數(shù)為 20k 時，也有 1.16 倍的 FLOPs 收益。

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

吞吐量分析

稀疏模型通常不能立即帶來性能增益，因為稀疏模型會增加動態(tài)性和相關的數(shù)據(jù)平衡問題。為了量化新提出的方法對訓練效率的影響，該團隊通?？刂谱兞繉嶒灡容^了不同架構(gòu)的訓練吞吐量。結(jié)果見表 2。

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

推理時間性能

該團隊也評估了模型在留存的語言建模數(shù)據(jù)和下游任務上的表現(xiàn)。結(jié)果見表 3 和 4。

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

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

總體而言，在混合文本和圖像兩種模態(tài)的數(shù)據(jù)上，1.4B MoMa 4t4i 模型的建模結(jié)果最好。

更多詳細內(nèi)容，請閱讀原論文。