復(fù)旦&通義萬(wàn)相提出ProMoE，顯式路由引導(dǎo)打破DiT MoE scaling瓶頸

混合專家架構(gòu)（Mixture-of-Experts，MoE）在擴(kuò)展模型容量的同時(shí)保持了計(jì)算效率，在大語(yǔ)言模型（LLM）的發(fā)展中發(fā)揮了巨大作用。然而，現(xiàn)有方法將 MoE 應(yīng)用于 Diffusion Transformer（DiT）時(shí)，卻發(fā)現(xiàn)收益非常有限，無法復(fù)刻 MoE 在 LLM 中的成功。

為什么同樣的架構(gòu)，跨界到視覺生成領(lǐng)域就「水土不服」了？

近期，來自復(fù)旦大學(xué)、阿里通義萬(wàn)相Wan Team、浙江大學(xué)和香港大學(xué)的研究團(tuán)隊(duì)指出，視覺 Token 的高度冗余性和功能異質(zhì)性阻礙了視覺 MoE 中專家的專業(yè)化（specialization）。為此，他們提出了 ProMoE，一種帶有顯式路由引導(dǎo)的兩步路由 MoE 框架。相關(guān)論文已被 ICLR2026 接收，第一作者為復(fù)旦大學(xué)博士衛(wèi)昱杰。

• 論文標(biāo)題：Routing Matters in MoE: Scaling Diffusion Transformers with Explicit Routing Guidance
• 論文地址：
• https://arxiv.org/abs/2510.24711
• 代碼：
• https://github.com/ali-vilab/ProMoE

視覺 Token 與語(yǔ)言 Token 的差異

為了探究 MoE 在 DiT 中收益不明顯的原因，研究團(tuán)隊(duì)發(fā)現(xiàn)，視覺 Token 具有兩個(gè)獨(dú)特的屬性，導(dǎo)致傳統(tǒng)的隱式路由分配策略產(chǎn)生次優(yōu)的效果：

• 高度空間冗余性（High Spatial Redundancy）：離散的文本 Token 語(yǔ)義高度濃縮且差異明顯，而連續(xù)的圖像 Patch（視覺 Token）在空間上高度耦合，存在大量的冗余信息，導(dǎo)致視覺 MoE 中的專家往往學(xué)到同質(zhì)化的特征。

• 功能異質(zhì)性（Functional Heterogeneity）：擴(kuò)散模型普遍依賴無分類器引導(dǎo)（CFG）技術(shù)。這就導(dǎo)致輸入 Token 天然分為兩派：條件 Token 和無條件 Token。標(biāo)準(zhǔn) MoE 范式對(duì)它們一視同仁、同時(shí)分配，忽略了它們不同的功能角色。

圖 1：（a）我們從 110 個(gè) ImageNet 類別中隨機(jī)抽取 1k 個(gè)中間層 Token，進(jìn)行 10 簇 k-means 聚類（以顏色區(qū)分）。以類別名稱 / 標(biāo)簽作為輸入時(shí)，LLM Token 形成緊湊、分離良好的簇，語(yǔ)義密度高，而視覺 Token 則較為分散。這種差異可以用類間距離與類內(nèi)距離的比值來量化（19.283 ? 0.748）。（b）我們對(duì)每個(gè) MoE 層的專家權(quán)重矩陣進(jìn)行奇異值分解，并計(jì)算由其左前 k 個(gè)奇異向量張成的子空間的平均相似度來衡量專家間的多樣性。引入路由引導(dǎo)（我們的方法）可以增強(qiáng)專家間的多樣性。

ProMoE：兩步路由與顯式語(yǔ)義路由引導(dǎo)

MoE 的核心原則是專家專業(yè)化（Expert Specialization），即確保每個(gè)專家都能獲取集中且不重疊的知識(shí)。為了在視覺模型中實(shí)現(xiàn)「專家內(nèi)一致」和「專家間多樣」，ProMoE 引入了兩步路由器（Two-Step Router）和路由對(duì)比學(xué)習(xí)（Routing Contrastive Learning）。

兩步路由器

• 第一步：條件路由（Conditional Routing）

路由器首先根據(jù) Token 的功能角色進(jìn)行硬路由分配。無條件圖像 Token（來源于 null conditioning 下的圖像 Patch）被直接分配給專門的無條件專家（Unconditional Experts）進(jìn)行處理。而條件圖像 Token 則進(jìn)入下一步，交由標(biāo)準(zhǔn)的路由專家（Routed Experts）進(jìn)行處理。這種機(jī)制實(shí)現(xiàn)了專家的功能隔離。

• 第二步：原型路由（Prototypical Routing）
  

對(duì)于條件圖像 Token，ProMoE 引入了一組可學(xué)習(xí)的「原型」（Prototypes），每個(gè)原型對(duì)應(yīng)一個(gè)特定專家。原型路由在隱空間中計(jì)算 Token 與各個(gè) Prototype 之間的余弦相似度，并選擇 identity function 作為激活函數(shù)得到路由分?jǐn)?shù)，分?jǐn)?shù)較高的 Token 會(huì)被分配給對(duì)應(yīng)的專家。

顯式語(yǔ)義路由引導(dǎo)：路由對(duì)比學(xué)習(xí)（Routing Contrastive Learning）

為了顯式增強(qiáng)原型路由的語(yǔ)義引導(dǎo)，ProMoE 提出了一種無需手動(dòng)標(biāo)注的路由對(duì)比損失（Routing Contrastive Loss，RCL）。在訓(xùn)練過程中，RCL 會(huì)產(chǎn)生兩種影響：

• 拉近：將 Prototype 拉向分配給它的 Token 集合的質(zhì)心，確保同一個(gè)專家處理的 Token 是語(yǔ)義相似的。

• 將 Prototype 推離其他專家處理的 Token 集合的質(zhì)心，鼓勵(lì)不同專家之間形成差異化，增強(qiáng)多樣性。

同時(shí)，實(shí)驗(yàn)發(fā)現(xiàn) RCL 中的「推開」操作在語(yǔ)義層面上天然起到了負(fù)載均衡的作用，比傳統(tǒng)的負(fù)載均衡損失更加靈活且有效。

圖 2：ProMoE 架構(gòu)概覽。輸入 Token 通過條件路由被分為無條件 Token 子集和條件 Token 子集。無條件圖像 Token 由無條件專家處理。條件圖像 Token 通過基于可學(xué)習(xí)原型的原型路由進(jìn)行分配。路由對(duì)比學(xué)習(xí)顯式增強(qiáng)了原型路由中的語(yǔ)義引導(dǎo)。

算法偽代碼如下：

實(shí)驗(yàn)結(jié)果

模型配置

與 Dense Model 的對(duì)比

ProMoE 在各種規(guī)模和設(shè)置下均穩(wěn)定超越了稠密模型。亮眼的是，參數(shù)量?jī)H 1.063B 的 ProMoE-L-Flow，憑借更少的激活參數(shù)，超越了計(jì)算量更大的 Dense-DiT-XL-Flow。

與 SOTA MoE Model 的對(duì)比

ProMoE 超越現(xiàn)有的視覺 MoE 方案。特別是，用 1.063B 超越了擁有 16 個(gè)專家，1.846B 的 DiffMoE。

Text-to-Image 驗(yàn)證

在 GenEval bench 中，ProMoE 在所有子任務(wù)上優(yōu)于標(biāo)準(zhǔn)的 Token-Choice MoE 模型，展現(xiàn)出一定的泛化能力。

可視化結(jié)果

Class-to-image generation

Text-to-image generation

收斂性分析

訓(xùn)練曲線顯示，ProMoE 的收斂速度明顯快于稠密模型和現(xiàn)有 MoE 模型。

Scaling 實(shí)驗(yàn)

ProMoE 展現(xiàn)出一定的擴(kuò)展?jié)摿?。隨著模型尺寸從 Base 擴(kuò)展至 XL，以及專家數(shù)量從 4 逐步增加到 16，ProMoE 的生成性能均呈現(xiàn)出穩(wěn)定的提升。

消融實(shí)驗(yàn)

總結(jié)

通過分析語(yǔ)言和視覺 Token 之間的差異，ProMoE 提出了一種帶有顯式路由引導(dǎo)的 MoE 框架。通過巧妙設(shè)計(jì)的條件路由、原型路由以及路由對(duì)比學(xué)習(xí)機(jī)制，ProMoE 用更少的激活參數(shù)超越了 Dense Model 以及現(xiàn)有 MoE 方法。這為如何在大規(guī)模擴(kuò)散模型中高效引入 MoE 架構(gòu)提供了一套可能的開源范式。

更多技術(shù)與實(shí)驗(yàn)細(xì)節(jié)，歡迎閱讀原論文。