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

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

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

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

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

之前已有研究表明，這類架構(gòu)可以有效地?cái)U(kuò)展單模態(tài)的基礎(chǔ)模型的能力，也可以增強(qiáng)多模態(tài)對(duì)比學(xué)習(xí)模型的性能。但是，將其用于較早將各種模態(tài)融合的模型訓(xùn)練還是一個(gè)機(jī)遇與挑戰(zhàn)并存的課題，還少有人研究。

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

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

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

模型架構(gòu)

早期融合

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

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

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

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

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

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

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

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

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

總體來(lái)說(shuō)，對(duì)于一個(gè)輸入 token x，MoMa 模塊的形式化定義為：

在 MoMa 計(jì)算之后，該團(tuán)隊(duì)又進(jìn)一步使用了殘差連接和 Swin Transformer 歸一化。

Mixture-of-Depths（MoD）

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

該團(tuán)隊(duì)的做法參考了第二種方法，同時(shí)整合了近期提出的混合深度（MoD）技術(shù)。有關(guān) MoD 的更多介紹可參閱機(jī)器之心報(bào)道《DeepMind 升級(jí) Transformer，前向通過(guò) FLOPs 最多可降一半》。

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

推理

在推理階段，我們不能直接使用 MoE 的專家選擇路由或 MoD 的層選擇路由，因?yàn)樵谝慌鷶?shù)據(jù)中進(jìn)行 top-k（選擇前 k 個(gè)）選擇會(huì)破壞因果關(guān)系。

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

升級(jí)改造（Upcycling）

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

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

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

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

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

效率優(yōu)化

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

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

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

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

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

其它優(yōu)化

為了進(jìn)一步提升吞吐量，該團(tuán)隊(duì)還采用了其它一些優(yōu)化技術(shù)。

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

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

實(shí)驗(yàn)

設(shè)置

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

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

不同計(jì)算層級(jí)的擴(kuò)展性能

該團(tuán)隊(duì)分析了不同模型在不同計(jì)算層級(jí)上的擴(kuò)展性能，這些計(jì)算層級(jí)（FLOPs）相當(dāng)于三種大小的密集模型：90M、435M 和 1.4B。

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

模態(tài)解綁

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

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

混合深度與專家

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

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

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

擴(kuò)展專家的數(shù)量

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

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

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

升級(jí)改造

該團(tuán)隊(duì)自然也驗(yàn)證了前述的升級(jí)改造的效果。圖 6 比較了不同模型變體的訓(xùn)練曲線。

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

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

吞吐量分析

稀疏模型通常不能立即帶來(lái)性能增益，因?yàn)橄∈枘Ｐ蜁?huì)增加動(dòng)態(tài)性和相關(guān)的數(shù)據(jù)平衡問(wèn)題。為了量化新提出的方法對(duì)訓(xùn)練效率的影響，該團(tuán)隊(duì)通常控制變量實(shí)驗(yàn)比較了不同架構(gòu)的訓(xùn)練吞吐量。結(jié)果見(jiàn)表 2。

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

推理時(shí)間性能

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

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

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

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

更多詳細(xì)內(nèi)容，請(qǐng)閱讀原論文。