Routing Matters in MoE: Scaling Diffusion Transformers with Explicit Routing Guidance
Authors: Yujie Wei
Shiwei Zhang
Hangjie Yuan
Yujin Han
Zhekai Chen
Jiayu Wang
Difan Zou
Xihui Liu
Yingya Zhang
Yu Liu
Hongming Shan
Venue: ICLR 2026 Poster (2026)
Keywords: Image Generation
Mixture-of-Experts
Diffusion Transformer
Primary area: generative models
TL;DR: We present ProMoE, a Mixture-of-Experts framework featuring a two-step router with explicit routing guidance, promoting expert specialization in MoE-based DiT models.

Abstract: Mixture-of-Experts (MoE) has emerged as a powerful paradigm for scaling model capacity while preserving computational efficiency. Despite its notable success in large language models (LLMs), existing attempts to apply MoE to Diffusion Transformers (DiTs) have yielded limited gains. We attribute this gap to fundamental differences between language and visual tokens. Language tokens are semantically dense with pronounced inter-token variation, while visual tokens exhibit spatial redundancy and functional heterogeneity, hindering expert specialization in vision MoE. To this end, we present $\textbf{ProMoE}$, an MoE framework featuring a two-step router with explicit routing guidance that promotes expert specialization. Specifically, this guidance encourages the router to $\textit{first}$ partition image tokens into conditional and unconditional sets via conditional routing according to their functional roles, and $\textit{second}$ refine the assignments of conditional image tokens through prototypical routing with learnable prototypes based on semantic content. Moreover, the similarity-based expert allocation in latent space enabled by prototypical routing offers a natural mechanism for incorporating explicit semantic guidance, and we validate that such guidance is crucial for vision MoE. Building on this, we propose a routing contrastive loss that explicitly enhances the prototypical routing process, promoting intra-expert coherence and inter-expert diversity. Extensive experiments on ImageNet benchmark demonstrate that ProMoE surpasses state-of-the-art methods under both Rectified Flow and DDPM training objectives. Code and models will be made publicly available.

Published as a conference paper at ICLR 2026
ROUTING MATTERS IN MoE: SCALING DIFFUSION TRANSFORMERS WITH EXPLICIT ROUTING GUIDANCE
Yujie Wei', Shiwei Zhang"; Hangjie Yuan", Yujin Han", Zhekai Chen Jiayu Wang?,
Difan Zou', Xihui Liu 4.5 Yingya Zhang", Yu Liu?, Hongming Shan If
Fudan University Tongyi Lab, Alibaba Group "The University of Hong Kong SMMLab yjwei2228 fudan, audan.edu..cmm en, zhangjin. wwwsalibabaaaaa inc. com,
*Zhejiang University
hmshan edua
ABSTRACT
Mixture-of-Experts (MoE) has emerged as a powerful paradigm for scaling model capacity while preserving computational efficiency. Despite its notable success in large language models (LLMs)), existing attempts to apply MoE to Diffusion Transformers (DiTs) have yielded limited gains. We attribute this gap to fun- damental differences between language and visual tokens. Language tokens are semantically dense with pronounced inter-token variation, while visual tokens exhibit spatial redundancy and functional heterogeneity, hindering expert spe- cialization in vision MoE. To this end, we present ProMoE, an MoE frame- work featuring two-step router with explicit routing guidance that promotes expert specialization. Specifically, this guidance encourages the router to first partition image tokens into conditional and unconditional sets via conditional routing according to their functional roles, and second refine the assignments of conditional image tokens through prototypical routing with learnable proto- types based on semantic content. Moreover, the similarity-based expert alloca- tion in latent space enabled by prototypical routing offers a natural mechanism for incorporating explicit semantic guidance, and validate that such guid- ance crucial for vision MoE. Building on this, we propose a routing con- trastive loss that explicitly enhances the prototypical routing process, promoting intra-expert coherence and inter-expert diversity. Extensive experiments on Im- ageNet benchmark demonstrate that ProMoR surpasses state-of-the-art methods under both Rectified Flow and DDPM training objectives. Code is available at https:/ /gii thub. , com/ ali-liiiiii seiivvvlabbee
INTRODUCTION
Diffusion models (Ho et al., 2020) have made substantial advances for visual synthesis (Rombach al., 2022b; Yang al., 2024; Wei al., 2024a; Wan al., 2025). Driven by the growing demand for higher fidelity and quality, research has focused on scaling up diffusion models (Esser et al., 2024b) and propelled an architectural shift from U-Net (Ronneberger al., 2015) backbones the now-prevalent Diffusion Transformers (DiTs) (Peebles & Xie, 2023). Despite the proven effective- ness of DiT-based models (Esser et al., 2024a), their dense activation of all parameters irrespective of task or input incurs substantial computational overhead, thereby hindering further scalability.
To scale toward larger and more capable models, the large language model (LLM) community has widely adopted the Mixture-oo--Eppeeee (MoE) (Jacobs et al., 1991; Shazeer et al., 2017) paradigm, which expands model capacity while maintaining computational efficiency. Conceptu- ally, an MoE layer dispatches input tokens to specialized "expert" sub-networks via a router and
Project Leader 1 Corresponding Author Yujie Wei and Hongming Shan are with Institute of Science and Technology for Brain-inspiee Intelligence, MOE Frontiers Center for Brain Seience, Key Laboratory of Computational Neuroscience and Brain-Inspired Intelligence, and State Key Lahoratory of Brain Function and Disorders, Fudan University, Shanghai, China

---

Published as a conference paper at ICLR 2026
LLM (Llama-3 8B)
Visual DIT (DIT-XL/2)
Expert Subspace Similarity (Lower Implies Higher Diversity)
Anglnner das dasesod)
Anggeen - ther decanne)
19.283
Angliteer chas dedaned) Aegglain
6774
DrF-MoB (Joken-Choice)
Ouss
0.1-
13
Layer
(a) I-SNE visualization of language and visual tokens
(b) Comparison of expert diversity across MoE layers
Figure 1: (a) We randomly sample 1 intermediate-layer tokens from 110 ImageNet classes for 10-cluster k-means clustering (differentiated by color). With class names/labels as inputs, LLM tokens form compact, well-separated clusters with high semantic density, whereas visual tokens are diffuse. This disparity is quantified by the ratio of inter- to intra-class distance (19. (19.283 >> 0.748). (b) We measure inter-expert diversity using singular value decomposition on each MoE layer's expert weight matrices and computing the mean similarity of the subspaces spanned by their top-k left singular vectors (Hu et al., 2021). Incorporating routing guidance (Ours) enhances expert diversity.
returns a weighted sum of the selected experts' outputs. Despite MoB's profound success in lan- guage modeling (Jiang et al., 2024; Dai et al., 2024), recent efforts to integrate it into DiT models have not yielded the significant gains observed in LLMs. Specifically, DiT-MoE (Fei et al., 2024), which employs token-to-exper routing, often underperforms dense counterparts despite activating the same number of parameters. In contrast, EC-DiT (Sun et al., 2024), which assigns each expert fixed quota of tokens, delivers only marginal gains even with extended training. More recently, DiftMoE (Shi et al., 2025), which introduces a global token-distribution routing scheme, still re- ports relatively limited improvements. This pronounced gap between MoE's transformative impact in LLMs and its modest returns in DiT models motivates a fundamental question: What are the underlying factors that impede the effectiveness of MoE in DiT models?
To answer this question, we examine how linguistic and visual inputs differ in models and highlight the following two distinctive properties of visual inputs. 1) High Spatial Redundancy. Unlike discrete text tokens, which are semantically dense with salient inter-token differences, visual tokens (i.e., image patches) continuous, spatially coupled, and substantially redundant (Fig. I(a)). The high correlation between patches often leads experts to learn homogeneous features. 2) Functional Heterogeneity. The practice of classifier-free guidance (Ho & Salimans, 2022) in diffusion mod- els inherently introduces two functionally distinct input types: conditional and unconditional. A naive MoE treats them uniformly with undifferentiated routing, ignoring their different roles. These properties collectively impede effective expert diversity and specialization (Fig. ((b)).
Motivated by these observations, we revisit the foundational principle of MoE design: expert spe- cialization, in which each expert acquires focused and non-overlapping knowledge (Dai et al., 2024; Cai et al., 2025). We decompose this objective into two criteria: Intra-Expert Coherence, which ensures that an expert consistently processes similar patterns, maintaining stable functional role; and Inter- Exter-Experttt Diversity, which encourages different experts to specialize in distinct tasks to achieve functional differentiation. In language modeling, the semantic density and separability of discrete text tokens provide a potent inductive bias that naturally fosters expert specialization, sat- isfying both criteria. In contrast, for visual inputs, the combination of intrinsic redundancy and extrinsic functional heterogeneity makes expert specialization non-trivial. Therefore, in this paper, we move beyond implicit expert allocation, and introduce explicit routing guidance designs to pro- mote both intra-expert coherence and inter-expert diversity.
To this end, we present ProMoE, Mixture-of-Expeer framework featuring a two-step router with explicit routing guidance to promote expert specialization. Specifically, this guidance provides two distinct routing signals: the token's 's functional role and its semantic content. Guided by these signals, the router implements two steps: conditional routing and prototypical routing. First, conditional routing addresses functional heterogeneity by partitioning visual tokens into unconditional and con- ditional sets. Unconditional image tokens, derived from image patches under null conditioning (e.g., empty labels or texts), are processed by dedicated unconditional experts. In contrast, condi- tional image tokens, obtained from patches under specific conditioning, are dispatched to standard experts (routed experts). This hard routing mechanism enforces functional segregation, fostering specialization across unconditional and routed experts. Second, prototypical routing further assigns conditional image tokens using a set of learnable prototypes, each associated with a specific expert, by computing cosine similarity between token embeddings and the prototypes in latent space.
2

---

Published as a conference paper at ICLR 2026
While prototypical routing is flexible and effective, it still relies implicit learning from token semantics. Fortunately, its similarity- based allocation in latent space provides a natural mechanism for injecting explicit semantic routing guidance. We validate the importance of semantic guidance in systematic experiments (Sec. 4.2), where both explicit claasificattonnbaselee and implicit (clustering- based) guidance yield clear improvements. Building on this, we propose routing confrastive loss that explicitly enhances the prototypical routing process by assigning semantically similar tokens to the same expert while preserving distinct token distributions across experts. Compared with alternative guidance strategies, the proposed contrastive loss requires no manual labels and is more robust, promoting intra-expert coherence and inter-expert diversity in vision MoE.
Extensive experimental results demonstrate ProMoE's superior performance and effective scalability on both Flow Matching and DDPM paradigms. Notably, ProMoE achieves significant gains over dense models despite using fewer activated parameters, and surpasses state-of-the arrt methods that have 1.7x more total parameters than ours.
In summary, our contributions are fourfold: By analyzing differences between language and vi- sual tokens, we present ProMoE, an MoE framework with explicit routing guidance for DiT models. 2) We design a two-step router, where conditional routing first partitions image tokens by functional roles, and prototypical routing then refines assignments using learnable prototypes based on seman- tic content. 3) We propose routing contrastive loss that enhances prototypical routing, explicitly enforeing intra-expert coherence and inter-expert diversity. 4) Extensive experiments demonstrate that ProMoE outperforms dense models and state-of of-t-offthhe-art MoE methods across diverse settings.
2 RELATED WORK
Diffusion Models. Diffusion models (Nichol & Dhariwal, 2021; Wu et al., 2025b; Deng et al., 2026) have made remarkable progress in visual synthesis. Early work (Rombach et al., 2022a; Podell et al., 2023; Wei et al., 2024b) primarily use U-Net trained with the DDPM objective (Ho et al., 2020; Song et al., 2020). Recent models (Chen et al., 2023; Ma et al., 2024; Hatamizadeh et al., 2024; Chu et al., 2024; Wu et al., 2025a; Wei et al., 2025; Tang et al., 2025) have shifted to Transformer architecture, offering superior scalability and generative quality, and many are trained with the more effective Rectified Flow (RF) (Liu al., 2022), flow-matching formulation (Lipman et al., 2022) that constructs straight-line path between data and noise distributions. In this work, we adopt a standard DiT backbone and train with both DDPM and RF objectives, demonstrating the effectiveness and scalability of our approach across different training paradigms.
Mixture of Experts. Bfxture-of-Expertt (MoE) (Jacobs et al., 1991; Shazeer et al., 2017: Lep- ikhin et al., 2020) are designed to expand model capacity while minimizing computational overhead by sparsely activating sub-networks for distinct inputs. Inspired by MoE successes in LLMs (Dai et al., 2024; Liu et al., 2024; Li et al., 2025; Muennighoff et al., 2024), recent work has integrated MoE to scale diffusion models to improve generative quality (Riquelme et al., 2021). Early MoE applications in U-Net-based diffusion models (Lee al., 2024; Balaji al., 2022; Feng al., 2023; Xue et al., 2023; Park et al., 2023; 2024; Zhao et al., 2024) often assign experts by diffusion timestep ranges, showing strong scaling potential. However, adapting Moï DiT architecture (Shen al., 2025; Sehwag et al., 2025; Cheng et al., 2025) faces several limitations. Token-choice routing meth- ods (e.g., DiT-MoE (Fei al., 2024)) suffer poor expert specialization due to imbalanced token as- signments. whereas expert-choice methods (e.g., EC- (Sun et al., 2024)) that fix token quotas per expert yield only marginal gains. More recently, DiffMoB (Shi al., 2025) and Expert Race (Yuan et al., 2025) explore batch-level global token selection and mutual expert- token routing, yet still rely on implicit expert learning and struggle with limited expert specialization due to the redundancy and functional heterogeneity of visual tokens. In contrast, analyze language-vision token differences and introduce explicit routing guidance to the MoE router based on the token's 's functional role and its semantic content. We further enhance the routing process through the proposed routing contrastive loss, promoting intra-expert coherence and inter-expert diversity.
3 PRELIMINARIES
Diffusion Models. Diffusion models are generative models that learn data distributions by re- versing a forward noising process. The continuous-time forward process can be formulated as
3

---

Published as a conference paper at ICLR 2026
ProMoE Block
Routing Soores
Scale
ProMoL
Unconditional Expert
Expert
Expert
Expert
Shared
Export
Block
Scale, Shift
Layer Norm
Two-Step Router
Stap2:
Semantic Routing
Prototypical Routing
Scale
Uncondinooal Image Tokens
Multi-Ilead
Self-Altention
Scale, Shift
MLP
Layer Norm
Condinionl inuge Tokons
Routing Contrastive Learning
Step1: Conditional Routing
Conditioning
Input Tokens
Input
Diffusion Transformer Block
Tokens
Uncunditional Image Tokens
Conditional Image Takens
Leamable Prototyyps For Routing
Figure 2: Overview ProMoE architecture. The input tokens are split by conditional routing into unconditional and conditional subsets. Unconditional image tokens are processed by unconditional experts. Conditional image tokens are assigned by prototypical routing with learnable prototypes. The routing contrastive learning explicitly enhances semantic guidance in prototypical routing.
X; = 0xX0 + dtC, with t E 4((0 1) and e ~ I), and are monotonically decreasing and increasing functions of f, respectively. For the reverse process, denoising network Fg is trained to predict the target y at each timestep f, conditioned on c (e.g., class labels text prompts):
= Exn.c,e.1 ||y
Fg (x,, c, 1)
(1)
where the training target y can be the Gaussian noise for DDPM models (Ho et al., 2020), or the vector field (c Xo) for Rectified Flow models (Liu et al., 2022).
Mixture of Experts. Mixture of-peeetss (MoE) an architectural paradigm that scales model capacity while preserving computational efficiency by selectively activating subset of "experts" sub-networks. A standard MoE layer comprises NE experts and a trainable router R. Each expert E; is implemented as feed-forward network (FFN). Given input X E BxLxD where B is the batch size, L is the token length, D is the hidden dimension, the router R maps the input X to token-expert affinity scores S E IR BxLLx NE via activation function A:
S = A(R(X))
(2)
At each forward pass, the router activates the top-K highest-scoring experts and dispatches the input to them. The final output is the weighted sum of the activated experts' outputs with gating function:
G - {
S E TopK(S) Otherwise
MoE(X)
= Σ +
(3)
1=1
where G E BB is the final gating tensor. There are two common routing paradigms in MoE: Token-Choice (TC) and Expert-Choice (EC). TC, each token independently selects top-K experts; in EC, each expert selects a fixed number of top-K tokens.
4 PROMOE
In this section, we present ProMoE, an MoE framework for DiTs that integrates a two-step router with explicit routing guidance. The overall pipeline is depicted in Fig. 2. We first detail the two-step router in Sec. 4.1. We then validate the importance semantic routing guidance in visual MoEs in Sec. 4.2 and further propose routing contrastive learning to enhance semantic guidance in Sec. 4.3.
4

---

Published as a conference paper at ICLR 2026
4.1 TwO-STE STOP ROUTER
The ProMoE router operates in two steps: conditional routing based on the token's functional role, followed by fine-grained prototypical routing based on token semantics.
Conditional Routing, Unlike LLMs, diffusion models typically employ classifier-free guidance (CFG) (Ho & Salimans, 2022) at inference to enhance sample quality. Specifically, CFG steers the generation process by combining the model's conditional and unconditional noise predictions. This paradigm naturally defines two functionally heterogeneous tokens: 1) unconditional image tokens, derived from image patches under null conditioning (e.g.. empty labels or texts); and 2) conditional image tokens, obtained from patches under specific conditioning (e.g., class labels or texts).
e
To handle different token types, the first step of the ProMoE router employs hard routing based on input conditioning. Specifically, unconditional image tokens are deterministically assigned to N₂ unconditional experts, each implemented a feed-forward network (FFN), analogous to routed experts, Conversely, conditional image tokens are passed to the second step for fine-grained routing among routed experts. This explicit partitioning encourages experts learn the functional disparity between token types, facilitating the specialization of unconditional and routed experts.
Prototypical Routing. The second step of our ProMoE router is dispatch conditional image to- kens for fine- fine-graine grained expert allocation. Concretely, we introduce a novel prototypical routing mech- anism where the routing weights are parameterized by a set of learnable prototypes P e x as illustrated in Fig. 2. Each prototype p; corresponds to an expert E, and is trained to represent the shared characteristics of a cluster of semantically similar tokens.
Compared with standard MoE token assignment, which computes pre-activation scores Z E RBxLxNe via a linear layer, we assign tokens using cosine similarity, which more effective and naturally suited for measuring semantic similarity in latent space between tokens and prototypes:
x.P |x: P
Z [
where x; and Py are the i-th token in X and the j-th prototype in P. a is a scaling factor.
Then, the activation function A transforms the pre-activation scores Z into token-expert affinity scores S. Instead of softmax, which is computationally expensive and sensitive to sequence length, we opt for simple monotonic function that preserves relative rankings. We evaluate both sigmoid and identity functions, finding that the identity A(Z)
performs best in practice. as shown in Table 8. We argue that the identity activation enables direct top-K selection and provides stable training, thus improving performance. Consequently, we adopt identity activation as S Z. Finally, each conditional image token is routed the top-K experts with gating scores G, as in Eq. (3).
Forward Process. Besides unconditional and routed experts, we also incorporate N, shared ex- perts that process all tokens to learn shared knowledge (Dai et al., 2024; Cheng et al., 2025). For each token, the output of our MoE block is defined as the sum of the shared experts' output and a selective output determined by the token type (conditional unconditional):
NE
Σ
Gj+ Ex(x) x
N.
MoE(x) Σ
+
(5)
No
(x)
ifxe
Shared
where F.S, Ey,, and F. are the shared, routed, and unconditional experts, respectively. X, and X,, conditional and unconditional image token sets, respectively, and X, UX, X.
To maintain a constant number of activated parameters, MoE models often employ fine-grained expert segmentation (Dai et al., 2024), where the inner hidden dimension of each expert is divided by the number of activated experts. In our most settings, each forward pass of our model activates exactly two experts: the single shared expert and one expert selected from the combined pool of routed and unconditional experts. Therefore, to match the computational cost of a dense model, we divide the hidden dimension of each expert's intermediate layer by factor of two.
5

---

x X, The RCL loss is then computed over the Published as a conference paper at ICLR 2026
4.2 SEMANTIC ROUTING GUIDANCE
Due to the inherent high spatial redundancy of visual tokens, a naive MoE router fails to sufficiently distinguish tokens for effective routing, leading experts learn homogeneous features. Conse- quently, additional semantic routing guidance is required to promote intra-expert coherence and inter-expert diversity. To validate this, we conduct experiments by augmenting the MoE router with two guidance types: 1) Explicit Routing Guidance and 2) Implicit Routing Guidance.
Explicit Routing Guidance. We design a routing classification loss that uses class labels to ex- plicitly guide token assignment. Specifically, we manually partition the IK ImageNet classes into Ne superclasses based on coarse labels in (Feng Patras, 2023), and allocate one expert per su- perclass. Since labels are sample-level, we instantiate the router as a classifier C: we average-pool the input X over the token length dimension obtain X, feed X into C to produce sample- expeet affinity scores S E RBxN, and assign the expert with highest score. During training, we supervise the routing process with a cross-entropy loss Lets = CE(SS,) c), where C is the superclass label.
Implicit Routing Guidance. We replace the standard MoE router with k-means clustering, signing all tokens a cluster to a single expert. Unlike the routing classification loss that provides explicit supervision, this design offers implicit guidance by measuring token similarity. encouraging semantically similar tokens to be co-assigned. Concretely, we initialize Ng cluster centroids by ran- domly sampling tokens. At each forward pass, compute each token's ''s distances to all centroids to obtain distance-based token-expert affinity scores. Each token is then assigned to its nearest cen- troid and thus routed to the corresponding expert. During training, centroids are updated iteratively by replacing each with the mean of their currently assigned tokens.
Results for both routing guidance are reported in Table 1, with all MoE models having the same activated parameters and comparable to- tal parameters to ensure fairness. On the base model size, DiT-MoE. (Fei et al., 2024) and DiffMoE (Shi et al., 2025) yield limited perfor- mance improvements. In contrast, adding either explicit or implicit guidance produces substan- tial gains. Notably, for both guidance strategies, disable the load-balancing loss to isolate its routing effects; despite its importance for TC routing, guidance alone still markedly improves per- formance. These findings highlight the pivotal role of semantic routing guidance in vision MoBs.
Table I: Comparison results under Rectified Flow ImageNet (256x256) after 500K training steps, evalu- ated with CFG=1.5.
Model (500K)
FIDSOK +
IS f
Dense-DiTe F-Flow DiT-MoE-B-Flow DifMMoE-B-Flow
9.02
131.13
8.94
131.66
8.22
137.46
Classification-busedd Routing
5.91
165.45
K-Means-based Routing
6.24
159.77
4.3 ENHANCING SEMANTIC ROUTING GUIDANCE VIA ROUTING CONTRASTIVE LEARNING
While the routing guidance strategies in Sec. 4.2 are effective, they suffer from key limitations: 1) The classificatoo-bbased routing loss is defined at the sample level, restricting token-level flexibil- ity and requiring costly manual annotations, hindering generalization. 2) Clustering-basee routing supports only top-1 assignment, and struggles to scale top-K, methods like k-means rely on disjoint clusters, making multi-centroid assignment difficult. Moreover, k-means is sensitive to the number of clusters and the cluster initialization (Arthur & Vassilvitskii, 2006), reducing robustness.
To address these limitations, we propose the Routing Contrastive Loss (RCL), as illustrated in Fig. 2, to explicitly enhance semantic guidance in prototypical routing. Given a mini-batch of conditional image tokens, RCL encourages semantically similar tokens to be routed to the same expert and pushes dissimilar tokens toward different experts, prompting expert specialization in MoF. Con- cretely, for each prototype p; associated with expert F,, tokens assigned to p; form the positive set, representing a cluster of semantically similar tokens, while tokens dispatched to other prototypes constitute the negative sets, comprising multiple clusters with semantics different from p..
Next, RCL pulls each prototype p: toward the centroid of its positive token set to enforce intra- expert coherence, while pushing it away from the centroids of negative sets to encourage inter- expert diversity. Let X, denote the tokens assigned to expert E, in a mini-batch. its centroid m is computed as the token mean: m;

---

#Head Published as a conference paper at ICLR 2026
Table 2: Model configurations of ProMoE with different model sizes, aligning with DiT (Peebles & Xie, 2023). "E14AISIU1" denotes that total of 14 experts are used, with I expert activated for each token, expert shared by all tokens, and unconditional expert for unconditional image tokens.
Model Config #Activated Parums.
#Total Params.
#Experts
#Blocks I.
#Hidden dim.
ProMoE-S
33M
75M
E14AIS1U1
12
384
ProMoF-B
130M
300M
EI4AISIU1
12
768
12
ProMoE-L
458M
1.063B
E14AISIU1
24
1024
16
ProMoE-XL 675M
1,568B
E14AISIU1
28
1152
16
Table 3: Quantitative comparison with Dense DiTs under Rectified Flow on ImageNet
(256x256) after 500K training steps, evaluated with CFG scales 1.0 and 1.5.
Activated
Total
cfg=1.0
cfg=1.5
Model (500K)
# Params.
# Params.
FID5OK 4
IS 1
FID50K
IS
Dense-DiT-B--loww ProMoF-B-Flow
130M
130M
30.61
49.89
9.02
131.13
130M
300M
24.44
60.38
6.39
154.21
Dense-DiT--L-Flow
458M
458M
15.44
84.20
3.56
209.03
ProMoF-L--Flw
458M
1.063B
11.61
100.82
2.79
244.21
Conse-DiT-XL-Flor
675M
675M
13.38
91.57
3.23
227.05
ProMoB-XL-Flow
675M
1.568B
9.44
114.94
2.59
265.62
prototypes of Na experts that are assigned tokens an online manner:
1
log exp(sim(p:. m///r)
LRCL
No
Dj=1 exp(sim(p) p., mz)/r)
where sim(a, b) thermore, we empirically find that the push-away operation in RCL acts as a load-balancing regu- larizer based on token semantics, and is more effective than traditional bal-baaan loss (Shazeer et al., 2017) (see Appendix E. The final training loss of ProMoE is the combination of Eqs. (1) and (6), weighting LRCI. by factor ÀRCI.-
|| denotes cosine similarity, and t is a temperature hyperparameter. Fur-
ab
5 EXPERIMENT
5.1 EXPERIMENTAL SETUP
Baseline and model configurations. We compare against Dense-DiT (Peebles & Xie, 2023) and MoE baselines, including DiT-MoE (Fei et al., 2024), EC-DIT (Sun al., 2024), and DiffMoE (Shi et al., 2025). For fair comparison, all MoE models are evaluated with equivalent activated parame- ters to the dense model and comparable total parameters, training with both DDPM (Ho et al., 2020) and Rectified Flow (Esser et al., 2024a) objectives. We scale ProMoE across four sizes (S/B/L/XL) to align with established DiT benchmarks, as shown Table Models are named as: [Model]- (Trziigg Type]. with additional expert configuration. For instance, expert configuration E14A1S1U1 denotes 14 total experts (E1 4), top-1 activation (A1) over routed experts, shared expert (S1), and 1 unconditional expert (U1). More details are provided in Appendix A.
Implementation details. We conduct experiments on class-conditional image generation using the ImageNet (Deng al., 2009) dataset, which contains 1,281,167 training images across 1,000 classes. Following (Peebles & Xie, 2023), we train all models with the AdamW optimizer with learning rate of le-4. The batch size is 256, and weight decay is O. We horizontal flips as the only data augmentation, and a pretrained VAE from Stable Diffusion (Rombach et al., 2022b) encode and decode images. We also maintain an exponential moving average (EMA) of model parameters during training with a decay rate of 0.9999, and all reported results use the EMA mode.
Evaluation metrics. We evaluate image generation quality of all methods using Fréchet Inception Distance (FID) (Heusel et al., 2017; Dhariwal & Nichol, 2021), calculated over 50K generated samples with 250 DDPM or Flow Matching Euler sampling steps. We also report Inception Score (IS) (Salimans et al., 2016) to measure the diversity of generated images.

---

Published as a conference paper at ICLR 2026
Table Quantitative comparison with Dense DiTs under DDPM ImageNet (256x256) after
500K training steps, evaluated with CFG scales of 1.0 and 1.5.
Activated
Total
cfg=1.0 1.0
cfg=1.5
Model (500K)
# Params.
# Params.
FID50K
IS
FID5OK +
IS 1
Mense-DiT-B-DDPP ProMoE-B B-DDPM
130M
130M
41.19
35.94
18.61
78.71
130M
300M
40.37
37.84
17.90
82.65
Miine-DiT-L-DDPEA
458M
458M
20.81
65.51
6.29
148.38
ProMoE-L-DDPMM
458M
1.063B
18.75
73.07
5.12
168.91
Dons--DiT-XI--DDDRP
675M
675M
17.67
74.05
5.07
165.81
ProMoF-XI-DDIPM
675M
1.568B
15.87
81.90
4.11
187.86
4..
Sereorrriiio flo
ofg
30
nesoc 16
FID550
100
200
10C
400
50c
100
200
500
400
500
100
200
300
400
500
Taining
Training
(a) Comparison with mudel
SOTAs
Scaliin model
(c) Scaling number experts
Figure 3: Comparisons and scaling results across diverse settings.
Table 5: Quantitative comparison with MoE baselines under Rectified Flow on ImageNet (256x256) after 500K training steps, evaluated with CFG scales of 1.0 and 1.5.
Activated
Total Params.
cfg=1.0
cfg=1.5
Model (500K)
#Fxperts
# Params.
GFLOPs
FID50K 4
1S
FIDSOK
IS f
DiT-MoF-L--Flow
ESAISOU0
458M
1.163B
77.50
16.57
80.25
4.10
199.05
EC-DiT-L-Flow
ESA1SOU0
458M
1.163B
77.50
15.58
84.11
3.65
209.06
DifMoo--L-Flov
ESAISOU0
458M
L.95BB
82.53
14.46
87.55
3.51
212.78
DiffMoE-L--Fion
ESGA
458M
18846B
90.03
13.55
92.33
3.30
222.40
LiMoB-L--Flow
E14AISIUI
458M
1,063B
77.72
11.61
100.82
2.79
244.21
5.2 MAIN RESULTS
Comparison with Dense DIT. The results in Fig. 3(a) and Tables and 4 draw three conclusions: 1) ProMoE consistently surpasses dense counterparts at equivalent activated parameters across all sizes, objectives, and CFG settings, demonstrating strong effectiveness, scalability, and generaliza- tion. 2) Gains are more pronounced under Rectified Flow, the current dominant training paradigm, highlighting ProMoE's ability to scale modern diffusion models. Compared to dense models, with- out CFG, ProMoo---Flow reduces FID by 24.8% and increases by 19.7%; at the largest scale, ProMoE-XL-Flow reduces FID by 29.4%. With CFG=1.5, ProMoE-B-Flor reduces FID by 29.2% while ProMoE-Xe -L-FFlow reduces FID by 19.8%. 3) ProMoE is notably parameter-efficientt it uses fewer activated parameters yet outperforms dense models with more. Specifically, ProMoE-L-Flow achieves FID 11.61/2.79 at CFG 1.0/1.5, versus 13.38/3.23 2.23 for Dense-DiT-XL-FFE
Comparison with MoE SOTAs. The results in Fig. 3(a), Tables 5 and 10 show that ProMoE outperforms all baselines across both objectives at equivalent activated parameters, with and without CFG. Without CFG, ProMoF-LL-Flow reduces FID by 19.7% and increases IS by 15.2% relative DillooE-L-Flow; with CFG=1.5, it reduces FID by 20.5% and increases IS by 14.8%. Notably, ProMoE-L-Flow (1.063B params) surpasses the larger DifMooE--Fllw with 16 experts (1.846B params), despite fewer total parameters, underscoring the effectiveness of our method.
Comparison of computational cost and efficiency. Tables and 14 show that ProMoE achieves lower inference time and fewer GFLOPs than the SOTA MoE method DiffMoE, and maintains GFLOPs comparable to other MoF baselines while delivering substantially higher performance. This demonstrates that the performance gains primarily stem from our superior methodological de- sign. More details about these experimental settings are provided in the Appendix A.
Visualization Results. Fig. 4 shows the samples generated by ProMoE-Xt ProMoE-XL-Flow on ImageNet (256x 256) after training steps with CFG=4.0; see Appendix C.3 for more analyses and results.
8

---

Published as a conference paper at ICLR 2026
Figure 4: Samples generated by ProMoE-XL-Flow after 2M iterations with cfg=4.0.
Table 6: Comparisons on GenEval benchmark under Rectified Flow after 400K training steps.
Model (400K)
#Experts
Activated
Params.
Params.
Total
Single Obj.
Two Obj.
Counting
GenEval
Colors
Position
Color Attr.
Overall
Dense
Tokon-Choice MoB
ESAISOUO ESAIS0U1
3B
12B
0.856
0.320
0.354
0.627
0.157
0.207
0.417
ProMoE
3B
12B
0.884
0.371
0,418
0.675
0.212
0.217
0.463
Comparison of training losses. Fig. 5 shows training loss curves for our method, the dense model, and MoE baselines. At the I. scale, ProMoE achieves lower loss and faster convergence than both MoE and dense baselines. This advantage persists at the largest XL scale, where it continues to converge faster than the dense model, even with extended training up to 1.2M steps.
0.77
0.76
Profec
eC How
0.76
DIINoc fow
0.75
0.72
200
400 500
(K)
Figure 5: Training loss curve comparisons.
Comparison on the text-to-image task. We conduct text-to-image totttoo-mmae experiments to further demon- strate the generalization ability of ProMoF. 'he detailed experimental setup is provided in Ap- pendix A, and we evaluate on the GenEval benchmark (Ghosh et al., 2023). As shown in Table 6, ProMoE significantly outperforms both the dense baseline and the Token-Choice MoE across the overall metric and all sub-tasks. These results strongly demonstrate ProMoE's s robust generalization capabilities in challenging generation scenarios. Furthermore, we also provide qualitative results in Fig. 14, further verifying ProMoE's s capability in high-quality text-to-image generation.
Quantitative analysis of expert utilization.
We compare the expert utilization of ProMoE against the DiT-MoE baseline using averaged token-per-expert ratios computed on two dis- joint class subsets, each consisting of 200 ran- domly sampled classes and 10,000 generated images. As shown in Fig. 6, DiT-MoE ex- hibits very similar tions across different class subsets, and shows minimal variation in token proportions among experts within each subset, indicating poor spe- cialization and failure to induce diversity. In contrast, ProMoE demonstrates clear ex- pert specialization, producing distinct utiliza- tion patterns across the disjoint subsets. Within each subset, experts exhibit varied usage fre- quencies without suffering from starvation or overuse. These results confirm that ProMoE achieves meaningful expert specialization and good load balancing.
distribu-
DIT-MoL-L-Flom
DIT DIT-MME-L-FFow
miier
Fared ('8matee
ProMoE-L-Flon (Class Subset
(Class Subset f)
Figure Analysis expert utilization.
5.3 SCALING BEHAVIOR
Scaling the model size. As shown in Fig. 3(b), ProMoE exhibits consistent performance improve- ments over Dense-DiT when scaling from base (B) to large (L) to XL, with 130M, 458M, and 675M activated parameters, respectively, thereby validating the scalability of our method.

---

Published as a conference paper at ICLR 2026
Table 7: Ablation study of each component on ImageNett (256x256) after 500K training steps, trained with Rectified Flow and evaluated with CFG scales of 1.0 and 1.5.
cfg=1.0
cfg=1.5
Model (500K)
FIDSOK
IS t
FIDSOK 4
IS f
Dense- Dense-DiT-B DiT-B-Flow (Dense)
30.61
49.89
9.02
131.13
Prototypical Routing
27.93
53.35
7.92
140.86
Routing Contrastive Learning
24.97
58.59
6.75
150.15
Conditional Routing
24.44
60.38
6.39
154.21
Table 8: Ablation of activation functions, trained with Rectified Flow on ImageNet (256x256) for 500K steps.
Table 9: Ablation of conditional routing in
K-Means-based Routing, trained with Rectified Flow on ImageNet (256x256) for 500K steps.
Activation (500K)
cfg=1.0
cfg=1.5
K-Means-based
cfg=1.0
cfg=1.5
FID4
1St
FID
ISt
Routing (500K)
FID
1St
FID4
IS1
Softmax
25.74
58.04
6.92
149.11
w/o Cond.
30.12
50.47
8.75
133.14
Sigmoid
25.49
58.51
150.94
w/ Cond.
25.61
59.76
6.24
159.77
Identity
24.44
60.38
6.39
154.21
Scaling the number of experts. Fig. 3(c) shows monotonic gains as the expert number increases from 4 to 16, with shared and unconditional expert per setup. For fair comparison, we maintain comparable total parameters MoE baselines and 14 experts across settings.
5.4 ABLATION STUDIES
Ablation each component. Table 7 shows that using prototypical routing improves perfor- mance and surpasses DiT-MoE-B-Flow and DiffMoE-B-Flow (see Table 1). Adding routing con- trastive learning (RCL) yields substantial gains, reducing FID by 10.6% and increasing IS by 9.8%, highlighting the importance of semantic routing guidance. These results validate the effectiveness of each component.
Ablation on score activation function. Since our prototypical routing computes similarities in latent space, choosing an appropriate activation to map similarities into routing scores crucial. As shown in Table 8, the identity mapping yields the best performance, sigmoid is second-best. and softmax performs worst. Consequently, we adopt the identity function as the score activation.
Ablation on conditional routing within the K-Means-based Routing. We emphasize that the proposed conditional routing general, method-agnostic component that can benefit other routing schemes. We ablate it within the K-Means-basse Routing method in Sec. 4.2. As shown in Table 9, removing conditional routing significantly degrades performance, underscoring its importance.
6 CONCLUSION
In this paper, we present ProMoE, Mixture-of-Experts framework featuring a two-step router with explicit routing guidance to promote expert specialization. We analyze differences between lan- guage and vision tokens: discrete text tokens are semantically dense, whereas visual tokens exhibit high spatial redundancy and functional heterogeneity, hindering the effectiveness of MoE in DiT models. To address this, we introduce routing guidance based on the token's ''s functional role and se- mantic content, yielding two-step router comprising conditional routing and prototypical routing. Furthermore, we propose a routing contrastive loss that enhances semantic guidance in prototypical routing, explicitly promoting intra-expert coherence and inter-expert diversity. Extensive experi- ments demonstrate that ProMoE outperforms dense DiT and existing MoE SOTAs, even with fewer activated or total parameters, providing robust solution for applying MoE to DiT models. Limitations. While we follow standard evaluation protocols and report FIDSOK and IS, these metrics may not fully capture fine-grained perceptual quality semantic faithfulness.
10

---

Published as a conference paper at ICLR 2026
Acknowledgements This work was supported in part by National Natural Science Foundation of China (No. 62471148), STI2030-Major Projects (No. 20212D0200204), Shanghai Center for Brain Science and Brain-inspire Technology, and Alibaba Group.
ETHICS STATEMENT
Our method achieves substantial improvements on the ImageNet benchmark over dense DiT and state-of-the-art MoE. methods, providing an effective solution for scaling DiT with MoE. Nonethe- less, it inherits common risks of generative models, such as the potential to create fake data. Robust image forgery detection may help mitigate these concerns. In addition, we adhere to ethical guide- lines in all experiments.
REPRODUCIBILITY STATEMENT
We make the following efforts to ensure the reproducibility of ProMoE: (1) All experiments are conducted on the publicly available ImageNet-IK IK benchmark. (2) Our code and trained model weights will be made publicly available. (3) We provide implementation details in Sec. 5.1 and Appendix
REFERENCES
David Arthur and Sergei Vassilvitskii. k-means+4: The advantages of careful seeding. Technical report, Stanford, 2006.
Shuai Bai, Keqin Chen, Xuejing Liu, Jialin Wang, Wenbin Ge, Sibo Song, Kai Dang, Peng Wang, Shijie Wang, Jun Tang, al. Qwen2. 5-vl technical report. arXiv preprint 2rX2::2502.1392. 2025.
Yogesh Balaji. Seungjun Nah, Xun Huang, Arash Vahdat, Jiaming Song, Qinsheng Zhang, Karsten Kreis, Miika Aittala, Timo Aila, Samuli Laine, al. ediff-i: Text-to-image diffusion models with an ensemble of expert denoisers. arXiv preprint arXiv:2211.01324 2022.
Weilin Cai, Juyong Jiang, Fan Wang, Jing Tang, Sunghun Kim, and Jiayi Huang. A survey on mixture of experts in large language models. IEEE Transactions Knowledge and Data Engi- neering, 2025.
Junsong Chen, Jincheng Yu, Chongjian Ge, Lewei Yao, Enze Xie, Yue Wu, Zhongdao Wang, James Kwok, Ping Luo, Huchuan Lu, et al. Pixart-a: Fast training of diffusion transformer for photore- alistic text-to-image synthesis. arXiv preprint arXiv:2310.00426, 2023.
Kun Cheng, Xiao He, Lei Yu, Zhijun Tu, Mingrui Zhu, Nannan Wang, Xinbo Gao, and Jie Hu. Diff-moe: Diffusion transformer with time-aware and space-adaptive experts. In Forty-second International Conference on Machine Learning, 2025.
Xiangxiang Chu, Jianlin Su, Bo Zhang, and Chunhua Shen. Visionllama: unified lama backbone for vision tasks. In European Conference on Computer Vision, pp. 1-18. Springer, 2024.
Damai Dai, Chengqi Deng, Chenggang Zhao, RX Xu, Huazuo Gao, Deli Chen, Jiashi Li, Wangding Zeng, Xingkai Yu, Yu Wu, et al. Deepseekmoe: Towards ultimate expert specialization in mixture- of-experts language models. arXiv preprint S.kiv:2401.06066 2024.
Haoyou Deng, Keyu Yan, Chaojie Mao, Xiang Wang. Yu Liu, Changxin Gao, and Nong Sang. Densegrpo: From sparse dense reward for flow matching model alignment. arXiv preprint arXiv:2601.202/8. 2026.
Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A large-scale hi- erarchical image database. In 2009 IEEE conference on computer vision and pattern recognition, pp. 248-255, Icee, 2009.
Prafulla Dhariwal and Alexander Nichol. Diffusion models beat gans on image synthesis. Advances in neural information processing systems, 34:8780-8794 2021.

---

Published as a conference paper at ICLR 2026
Abhimanyu Dubey, Abhinav Jauhri, Abhinav Pandey, Abhishek Kadian, Ahmad Al-Dahle, Aiesha Letman, Akhil Mathur, Alan Schelten, Amy Yang, Angela Fan, et al. The llama 3 herd of models. arXiv e-prints, pp. arXiv-2407, 2024.
Patrick Esser, Sumith Kulal, Andreas Blattmann, Rahim Entezari, Jonas Müllrr, Harry Saini, Yam Levi, Dominik Lorenz, Axel Sauer, Frederic Boesel, et al. Scaling rectified flow transformers for high-resolution image synthesis. In Forty-first International Conference on Machine Learning, 2024a.
Patrick Esser, Sumith Kulal, Andreas Blattmann, Rahim Entezari, Jonas Müllrr, Harry Saini, Yam Levi, Dominik Lorenz, Axel Sauer, Frederic Boesel, al. Scaling rectified flow transformers for high-resolution image synthesis. In Forty-first international conference on machine learning. 2024b.
Zhengcong Fei, Mingyuan Fan, Changqian Yu, Dehang Li, and Junshi Huang. Scaling diffusion transformers to 16 billion parameters. arXiv preprint arXiv:2407. 1633, 2024.
Chen Feng and loannis Patras. Maskcon: Masked contrastive learning for coarse-labelled dataset. In Proceedings of the IEEE/CVF Conference Computer Vision and Pattern Recognition, pp- 19913-19922, 2023.
Zhida Feng. Zhenyu Zhang. Xintong Yu, Yewei Fang, Lanxin Li, Xuyi Chen, Yuxiang Lu, Jiaxi- ang Liu, Weichong Yin, Shikun Feng, et al. Erni--vlg 2.0: Improving text-to-image diffusion model with knowledge-enhanced Conference on Computer Vision and Pattern Recognition, pp. 10135-10145, 2023.
In Proceedings of the IEEE/CVF
Dhruba Ghosh, Hannaneh Hajishirzi, and Ludwig Schmidt. Geneval: An object- focuut- framework for evaluating text-to-image alignment. Advances in Neural Information Processing Systems, 36: 52132-52152, 2023.
Ali Hatamizadeh, Jiaming Song, Guilin Liu, Jan Kautz, and Arash Vahdat. Diffit: Diffusion vision transformers for image generation. In European Conference on Computer Vision, pp. 37-55. Springer, 2024.
Martin Heusel, Hubert Ramsauer, Thomas Unterthiner, Bernhard Nessler, and Sepp Hochreiter. Gans trained by a two time-scale update rule converge to local nash equilibrium. Advances in neural information processing systems, 30, 2017.
Jonathan Ho and Tim Salimans. arXiv:2207.12598 2022.
Classifier-free
diffusion guidance.
arXiv preprint
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in neural information processing systems, 33:6840-6851, 2020.
Edward Hu, Yelong Shen, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, Shean Wang, Lu Wang, and Weizhu Chen. Lora: Low-rank adaptation of large language models. arXiv preprint arXiv:2106.09669 2021.
Robert A Jacobs, Michael Jordan, Steven J Nowlan, and Geoffrey E Hinton. Adaptive mixtures of local experts. Neural computation, 3(1):79-87, 1991.
Albert Q Jiang, Alexandre Sablayrolles. Antoine Roux, Arthur Mensch, Blanche Savary, Chris Bam- ford, Devendra Singh Chaplot, Diego de las Casas, Emma Bou Hanna, Florian Bressand, et al. Mixtral of experts. arXiv preprint arXiv:2401.04m0 2024.
Yunsung Lee, Jin Young Kim, Hyojun Go, Myeongho Jeong. Shinhyeok Oh, and Seungtaek Choi. aaltittrchhtecture multi-expert diffusion models. In Proceedings the AAAI Conference Artificial Intelligence, volume 38, pp. 13427-13436, 2024.
Dmitry Lepikhin, HyoukJoong Lee, Yuanzhong Xu. Dehao Chen. Orhan Firat, Yanping Huang, Maxim Krikun, Noam Shazeer, and Zhifeng Chen. Gshard: Scaling giant models with conditional computation and automatic sharding. arXiv preprint arXiv:2006. 16668, 2020.
12

---

Published as a conference paper at ICLR 2026
Aonian Li. Bangwei Gong, Bo Yang, Boji Shan, Chang Liu, Cheng Zhu, Chunhao Zhang, Congchao Guo, Da Chen, Dong Li, et al. Minimax-01: Scaling foundation models with lightning attention. arXiv preprint wwXhv:2501008/33 2025.
Yaron Lipman, Ricky TQ Chen, Heli Ben- Ben-Hamu, Maximilian Nickel, and Matt Le. Flow matching for generative modeling. arXiv preprini arXiv:2210.02747, 2022.
Aixin Liu, Bei Feng, Bing Xue, Bingxuan Wang, Bochao Wu, Chengda Lu, Chenggang Zhao, Chengqi Deng, Chenyu Zhang, Chong Ruan, et al. Deepseek-v3 technical report. arXiv preprint 2XX1:22412.1937,, 2024.
Xingchao Liu, Chengyue Gong, and Qiang Liu. Flow straight and fast: Learning to generate and transfer data with rectified flow. arXiv preprint arXiv:2209.03003 2022.
Nanye Ma, Mark Goldstein, Michael S Albergo, Nicholas M Boffi, Eric Vanden-Eijnden and Sain- ing Xie. Sit: Exploring flow and diffusion-based generative models with scalable interpolant transformers. In European Conference on Computer Vision, pp. 23 40. Springer, 2024.
Niklas Muennighoff, Luca Soldaini, Dirk Groeneveld, Kyle Lo, Jacob Morrison, Sewon Min, Wei- jia Shi, Pete Walsh, Oyvind Tafjord, Nathan Lambert, al. Olmoe: Open mixture-of-experts language models. arXiv preprint arXiv:2409.0206) 2024.
Alexander Quinn Nichol and Prafulla Dhariwal. Improved denoising diffusion probabilistic models. In International conference on machine learning, pp. 8162-8171. PMLR, 2021.
Byeongjun Park, Sangmin Woo, Hyojun Go, Jin- Young Kim, and Changick Kim. Denoising task routing for diffusion models. arXiv preprint arXiv:2310.07138, 2023.
Byeongjun Park, Hyojun Go, Jin-Young Kim, Sangmin Woo, Seokil Ham, and Changick Kim. Switch diffusion transformer: Synergizing denoising tasks with sparse mixture-of-experts In European Conference on Computer Vision, pp. 461-477. Springer. 2024.
William Peebles and Saining Xie. Scalable diffusion models with transformers. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pp. 4195-4205 2023.
Dustin Podell, Zion English, Kyle Lacey, Andreas Blattmann, Tim Dockhor, Jonas Müller, Joe Penna, and Robin Rombach. Sdxl: Improving latent diffusion models for high-resolution image synthesis. arXiv preprint arXiv:2307.01952. 2023.
Carlos Riquelme, Joan Puigcerver, Basil Mustafa, Maxim Neumann, Rodolphe Jenatton, André Susano Pinto, Daniel Keysers, and Neil Houlsby. Scaling vision with sparse mixture of experts. Advances in Neural Information Processing Systems, 34:8583-8500 8595, 2021.
Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Bjöm Ommer. High- resolution image synthesis with latent diffusion models. In Proceedings of the IEEE/CVF confer- ence on computer vision and pattern recognition, pp. 10695, 2022a.
Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Bjöm Ommer. High- resolution image synthesis with latent diffusion models. In Proceedings of the IEEE/CVF Con- ference on Computer Vision and Patlern Recognition, pp. 10684-10695, 2022b.
Olaf Ronneberger, Philipp Fischer, and Thomas Brox. U-net: Convolutional networks for biomedi- cal image segmentation. In International Conference on Medical image computing and computer- assisted intervention, pp. 234-241. 234-ier..pringer, 2015.
Tim Salimans, Ian Goodfellow, Wojciech Zaremba, Vicki Cheung, Alec Radford, and Xi Chen. Improved techniques for training gans. Advances in neural information processing systems, 29, 2016.
Christoph Schuhmann, Romain Beaumont, Richard Vencu, Cade Gordon, Ross Wightman, Mehdi Cherti, Theo Coombes, Aarush Katta, Clayton Mullis, Mitchell Wortsman, al. Laion-5b: An open large-scale dataset for training next generation image-text models. Advances in neural in- formation processing systems,
2022.
13

---

Published as a conference paper at ICLR 2026
Vikash Sehwag, Xianghao Kong. Jingtao Li, Michael Spranger, and Lingjuan Lyu. Stretching each dollar: Diffusion training from scratch on a micro-budget. In Proceedings of the Computer Vision and Pattern Recognition Conference, pp. 28596-28608, 2025.
Noam Shazeer, Azalia Mirhoseini, Krzysztof Maziarz, Andy Davis, Quoe Le, Geoffrey Hinton, and Jeff Dean. Outrageously large neural networks: The sparsely-gan -ated mixture-of-expper layer. arXiv preprini 1XX1v:700.066: 2017.
Ying Shen, Zhiyang Xu, Jiuhai Chen, Shizhe Diao, Jiaxin Zhang, Yuguang Yao, Joy Rimchala, Is- mini Lourentzou, and Lifu Huang. Latte-flow: Layerwise timestep-expert flow-based transformer. 2025.
arXiv preprint 2rXiv:2506.066
Minglei Shi, Ziyang Yuan, Haotian Yang, Xintao Wang, Mingwu Zheng, Xin Tao, Wenliang Zhao, Wenzhao Zheng, Jie Zhou, Jiwen Lu, et al. Diffmoe: Dynamic token selection for scalable diffusion transformers. arXiv preprint arXiv:2503. 144877 2025.
Yang Song, Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. arXiv preprint arXiv:2011.13456 2020.
Haotian Sun, Tao Lei, Bowen Zhang, Yanghao Li, Haoshuo Huang, Ruoming Pang, Bo Dai, and Nan Du. Ec-dit: Scaling diffusion transformers with adaptive expert-choice routing. arXiv preprint N.X::22100.02098. 2024.
Longxiang Tang, Ruihang Chu, Xiang Wang, Yujin Han, Pingyu Wu, Chunming He, Yingya Zhang, Shiwei Zhang, and Jiaya Jia. Exploiting discriminative codebook prior for autoregressive image generation. arxiv preprint arXiv:2508. 10719, 2025.
Team Wan, Ang Wang, Baole Ai, Bin Wen, Chaojie Mao, Chen-Wei Xie, Di Chen, Feiwu Yu, Haiming Zhao, Jianxiao Yang, et al. Wan: Open and advanced large-scale video generative models. arXiv preprint 2.Xi:22503.20314. 2025.
Yujie Wei, Shiwei Zhang, Zhiwu Qing, Hangjic Yuan, Zhiheng Liu, Yu Liu, Yingya Zhang, Jingren Zhou, and Hongming Shan. Dreamvideo: Composing your dream videos with customized sub- ject and motion. In Proceedings of the IEEE/CVE Conference on Computer Vision and Pattern Recognition, pp. 2537-6549, 2024a.
Yujie Wei, Shiwei Zhang, Hangjie Yuan, Xiang Wang. Haonan Qiu, Rui Zhao, Yutong Feng, Feng Liu, Zhizhong Huang, Jiaxin Ye, al. Dreanvideo-2: Zero-shot subject-driven video customiza- tion with precise motion control. arXiv preprint arXiv:2410.138800 2024b.
Yujie Wei, Shiwei Zhang. Hangjie Yuan, Biao Gong, Longxiang Tang, Xiang Wang, Haonan Qiu, Hengjia Li, Shuai Tan, Yingya Zhang, et al. Dreamrelation: Relation-centric video customization. arXiv preprint ar arXiv:2503.07600 2025.
Pingyu Wu, Kai Zhu, Yu Liu, Longxiang Tang, Jian Yang. Yansong Peng. Wei Zhai, Yang Cao, and Zheng-Jun Zha. Alitok: Towards sequence modeling alignment between tokenizer and autore- gressive model. arXiv preprint arXiv:2506.0522 05288,, 2025a.
Pingyu Wu, Kai Zhu, Yu Liu, Liming Zhao, Wei Zhai, Yang Cao, and Zheng-Jun Zha. Improved video vae for latent video diffusion model. In Proceedings of the Computer Vision and Pattern Recognition Conference, pp. 18124-18133, 2025b.
Enze Xie. Junsong Chen, Junyu Chen, Han Cai, Haotian Tang, Yujun Lin, Zhekai Zhang, Muyang Li, Ligeng Zhu, Yao Lu, et al. Sana: Efficient high-resolution image synthesis with linear diffu- sion transformers. arXiv preprint arXiv:2410. 10629, 2024.
Zeyue Xue, Guanglu Song, Qiushan Guo, Boxiao Liu, Zhuofan Zong, Yu Liu, and Ping Luo. Raphael: Text-to-image generation via large mixture of diffusion paths. Advances in Neural Information Processing Systems, 36:41693 2023.
Zhuoyi Yang. Jiayan Teng, Wendi Zheng, Ming Ding, Shiyu Huang, Jiazheng Xu, Yuanming Yang, Wenyi Hong, Xiaohan Zhang, Guanyu Feng, et al. Cogvideox: Text-to-video diffusion models with expert transformer. arXiv preprint 2:Xx:2448.00072. 2024.
14

---

Published as a conference paper at ICLR 2026
Yike Yuan, Ziyu Wang, Zihao Huang, Defa Zhu, Xun Zhou, Jingyi Yu, and Qiyang Min. Expert race: A flexible routing strategy for scaling diffusion transformer with mixture of experts. arXiv preprint 4rX1v:2503.16007, 2025.
Wangbo Zhao, Yizeng Han, Jiasheng Tang, Kai Wang, Yibing Song, Gao Huang, Fan Wang, and Yang You. Dynamic diffusion transformer. arXiv preprint arXiv:2410.0344 2024.
Wenliang Zhao, Lujia Bai, Yongming Rao, Jie Zhou, and Jiwen Lu. Unipe: A unified predictor- corrector framework for fast sampling of diffusion models. Advances Neural Information Processing Systems, 36:49842-49869, 2023.
15

---

Published as a conference paper at ICLR 2026
APPENDIX
A EXPERIMENTAL SETUP
Baselines. We compare against open-source state-of-the-arr DiT-based MoE methods, with acti- vated parameters equivalent to the dense model and comparable total parameter counts. We im- plement DiT-MoE (Fei et al., 2024), EC-DiT (Sun et al., 2024), and DiffMoE (Shi al., 2025) following their original papers and referring to the open-source repository'. All methods are trained with the same training settings, including learning rate, batch size, and data augmentation.
Implementation details. We train ProMoE with two objectives: the standard DDPM objec- tive (Ho et al., 2020), and Rectified Flow with Logit-Normal sampling from SD3 (Esser et al., 2024a) to align better with modern DiT training paradigms (e.g., Sana (Xie et al., 2024)). Besides the hyperparameters in Sec. 5.1, we provide additional details as follows. In prototypical routing, in Eq. (4) is set to 1. For routing contrastive learning, the temperature is set to 0.07, and the loss weight ARCI. is I unless stated otherwise.
Evaluation metrics. We follow the standard DìT evaluation protocol (Peebles & Xie, 2023), com- puting FID and IS on 50,000 generated images* at classifier-free guidance scales of 1.0 and 1.5.
Computational cost evaluation settings. We report the computational cost of our model in terms of training time, inference time, and FLOPs, and compare it against dense and MoE baselines. The inference time is measured for the core denoising process (excluding VAE decoding) with batch size of 128, a classifier-free free guidance (CFG) scale of 1.5, and 250 sampling steps. For a fair comparison, all experiments are conducted on 4 NVIDIA A800 GPUs under identical hardware settings.
Text-to-image experimental setup. We conduct new experiments on the text-to-image generation task to demonstrate the generalization of our ProMoE. Below, we detail the experimental setup. 1) For the model architecture, we adopt an MM-DiT-styl architecture, where text and visual tokens concatenated and jointly processed via self-attention for effective conditioning. The dense baseline uses 36 layers, 16 attention heads, hidden size of 11,008, and has 3B parameters. We build on this backbone by extending it to an MoE architecture, where each expert is visual FFN. As comparison baseline, we adopt the standard MoE paradigm commonly used in LLMs, denoted as Token-Choice MoE, which uses 5 routed experts and has 12B total parameters with 3B activated parameters. Our ProMoE uses the same overall capacity (12B total, 3B activated parameters) and also employs 5 experts (1 unconditional and 4 routed), ensuring a fair comparison under matched parameter and activation budgets. 2) For implementation details, we train all models on 2M-image subset of LAION-5B (Schuhmann et al., 2022) to enable rapid validation. We use the Adam W optimizer with learning rate of le-4, a global batch size of 384, and a weight decay of 0.02. Images encoded into the latent space using the pretrained Wan2.5- Preview VAE (Wan et al., 2025), and text is encoded with Qwen2.5-VI. (Bai et al., 2025). The training objective is Rectified Flow. Both the baselines and our ProMoE All experiments are conducted on 16 A800 GPUs for 400K training steps. For inference, we use the EMA model and adopt the UniPC (Zhao et al., 2023) sampler with 50 sampling steps, with classifier-free free guidance scale of 5.0.
trained under the same settings to ensure fair comparison.
B IMPLEMENTATION ALGORITHMS
The implementation algorithm of ProMoE is provided in Algorithm 1. In the algorithm, input class labels are used solely to distinguish conditional image tokens from unconditional image tokens. During inference, if classifier-free guidance (CFG) is disabled, all tokens are dispatched to the routed experts and the shared expert; the unconditional expert is not used. If CFG is enabled, class labels are replaced with a batch-level binary mask indicating which samples receive conditioning (i.e. treated as conditional image tokens). No routing contrastive loss is computed during inference.
16

---

Published as a conference paper at ICLR 2026
C MORE RESULTS
C.I MORE T-SNE VISUALIZATIONS OF LANGUAGE AND VISUAL TOKENS
To further validate the findings on differences between language and visual tokens in Sec. 1, we extend the visualization results in Fig. (a). Figs. and 12 present t-SNE visualizations of token embeddings from DiT-XL/2 (Peebles & Xie, 2023) and Llama-3 8B (Dubey et al., 2024) across different layers; for DiT-XL/2, we also visualize tokens at different diffusion timesteps. To facilitate comparison, cluster token embeddings into 10 groups using k-means. For model inputs, randomly sample 110 ImageNet classes, feed the corresponding class labels to DiT-XL/2 and the class names Llama-3 8B, and randomly select 1K intermediate-laye tokens for visualization. The results in Figs. 11 and 12 further confirm that language tokens are semantically dense with high inter-token differences, whereas visual tokens exhibit high spatial redundancy.
C.2 T-SNE VISUALIZATIONS OF TOKEN ASSIGNMENTS
To assess the impact of visual-token redundancy on MoE expert selection, as indicated by Fig. I(a) and Figs. 11 and 12, we visualize intermediate-layer token assignments of ProMoE-L-Flow and DiT- MoE-L-Flow 500K training steps without classifier-free guidance, as shown Fig. 7. Following Sec. C.I. we randomly sample 10 ImageNet classes, feed the corresponding class labels to both ProMoE and DiT-MoE, and randomly select 2,560 tokens from an intermediate-laye MoE block to visualize the expert selection of each token. Compared with token-choice MoE methods such as DiT-MoF, our approach assigns experts according to token semantics, producing well-formed clusters in the token-embedding space: semantically similar tokens form compact clusters and are routed to the same expert, whereas clusters assigned to different experts are clearly separated. These results further corroborate the importance of explicit routing guidance for visual MoE, and our method achieves effective intra-expert coherence and inter-expert diversity.
ProMoE-L-Flow
DIT-MoE-L-Flow
Figure 7: t-SNE visualization results of ProMoE and DiT-MoE on expert allocation (token assignment). Each color corresponds to a single expert.
C.3 MORE VISUALIZATION RESULTS
We provide additional generation results Figs. 13 and 14. Our method produces high-quality images across both simple and challenging categories.
C.4 MORE COMPARISON RESULTS USING THE DDPM OBJECTIVE
We provide additional comparisons with dense models and MoE SOTAs. Besides the FID results in Fig. 3(a), we also report Inception Score comparisons, as shown in Fig. 8. In addition, we present quantitative comparisons with MoE SOTAs under the DDPM objective in Table 10. Across both training objectives and CFG settings, our method consistently outperforms the dense model and existing MoE. SOTAs, demonstrating its effectiveness.
17

---

Published as a conference paper at ICLR 2026
Table 10: Quantitative comparison with MoE baselines under DDPM on ImageNet (256x256)
after 500K training steps, and evaluated at CFG scales of 1.0 and 1.5.
Activated
Total Params.
cfg=1.0
cfg=1.5
Model (500K)
#Experts
# Params.
FID5OK
IS 1
FIDSOK 4
IS
DiT-MoE-L-DDPM EC-DiT-L-DDPM R8AISOU0
ESAISOUU
458M
1.163B
23.12
60.08
7.55
133.63
458M
1.163B
20.76
64.77
6.48
146.49
DifMooB-L-DDPM
ESAISOU0
458M
1.095B
19.45
70.93
5.47
158.30
ProMoE-L-DDPN
E14AIS1U1
458M
1.063B
18.75
73.07
5.12
168.91
250
Dense LIT-
MOE
cfg-1.5 1.5
200
EC-DIT-L-- ProfAc---FTow
150
100
cf 1.0
50
Training Step (K)
Figure 8: Comparison with dense model and MoE SOTAs Inception Score.
C.5 MORE COMPARISON RESULTS WITH EXTENDED TRAINING STEPS
We provide comparison results between our ProMoE and Dense DiT on more training steps in Tables 11 and 12. We observe that ProMoE-L--Flow at 500K steps surpasses Dense-DiT-L-Flow at at 1M steps
1M steps on FID, and ProMoE-XL-Flow at 500K steps surpasses on both FID and IS. With longer training, ProMoE-L-Flow at 1M steps outperforms Dense-DiTe DiT-L- Flow at 2M steps and Dense-DiT-XL--Flww at IM steps. These findings are consistent with those in Sec. 5.2, demonstrating faster convergence and scalability of our method.
We also evaluate I. and XI. ProMoE models against dense and SOTA DiffMoR baselines 2M steps with CFG scales of 1.0, 1.5, and 4.0. As shown in Table 12, our ProMoE models consistently outperform baselines for both L/XL sizes under various settings. At CFG=4.0, for L-size, ProMoE reduces FID by ~21 % over the dense model and 10% over DiffMoE: for XL-size, ProMoE reduces FID by ~30% over the dense model and 16% over DiffMoE, while also achieving the best IS. At lower CFG values, ProMoE still yields consistent improvements over both dense and DiffMoE baselines. Furthermore, ProMoB-L 2M steps surpasses Dense/DiffMoO--X. activated parameters at CFG scales of 4.0 and 1.5. These results further solidify ProMoE's better performance and parameter efficiency under prolonged training and stronger guidance.
FID with fewer
Table Quantitative comparison with Dense DiTs under Rectified Flow on ImageNet (256x 256) after more training steps, evaluated with CFG scales of 1.0 and 1.5.
Activated
Total
cfg=1.0
cfg=1.5
Model (Training Steps)
Params.
Params.
FID50K
IS
FID50K
IS f
Dense Dense-DiT-L- DiT-L-Flow (1M)
458M
458M
12.21
100.97
2.97
245.63
ProMoE-L-Floo (500K)
458M
1.063B
11.61
100.82
2.79
244.21
ProMoF-L.-Flow (1M)
458M
1.063B
9.88
118.91
2.75
278.22
DiT-XL-Flow (1M)
675M
675M
10.67
107.68
2.82
260.61
ProMoE-XL XL--llow (500K)
675M
1.568B
9.44
114.94
2.59
265.62
ProMoF-XXI-Floo (1M)
675M
1.568B
8.34
128.58
2.53
292.38
C.6 INCREASING THE NUMBER OF ACTIVATED EXPERTS
As discussed in Sec, 4.2, classification- and clustering based routing inherently do not support top-k assignment, permitting only top-1. In contrast, ProMoR is more flexible and scalable, and supports
18

---

Published as a conference paper at ICLR 2026
Table 12: Quantitative comparison under Rectified Flow after 2M training steps on ImageNet
(256x 256), evaluated with CFG scales of 1.0, 1.5, and 4.0.
Activated
Total
cfg=1.0
cfg=1.5 5
cfg=4.0
Model (Training Steps)
Params.
Params.
FIDS0K L
IS
FIDSOK
IS
FID50K
I8
Dense-Die Dense-DiT---Floww (2M)
458M
458M
10.55
112.55
2.81
266.24
17.71
464.25
DifMMoE-L-Flow (2M)
458M
1.163B
10.09
122.01
2.57
276.75
15.70
465.60
ProMoE-LLFlow (2M)
458M
1,03BB
125.88
2.22
290.61
14.05
466.96
Deuse-DiTXXL-Flow (2M)
675M
675M
9.30
121.92
2.55
282.59
17.05
474.39
DifMooE-LL-Flow (2M)
675M
1.735B
8.83
134.37
2.49
297.25
14.25
475.13
ProMoE-XL-Flow (2M)
675M
1,568B
8.16
139.66
2.08
304.26
11.95
478.60
top-k assignment. To validate this, we increase the number of activated routed experts from 1 to 3, which raises the activated parameters while keeping the total parameter count unchanged. As shown in Table 13, this increase yields improved performance, confirming the effectiveness and scalability of our method.
Table 13: Results of increasing the number of activated routed experts on ImageNet (256x256)
after 500K steps, trained with Rectified Flow and evaluated at CFG scales 1.0 and 1.5.
Activated
Total
cfg=1.0
cfg=1.5 1.5
Model (500K)
#Experts
# Params.
Params.
FIDSOK 4
1S 1
FID5OK |
1S
ProMoE-L-Flow ProMoE-L-Flow
E14AISIUI
458M 558M
1.063B
11.61 11.40
100.82
2.79
244.21
E14A3SIUI
1.063B
103.78
2.72
246.96
C.7 FULL DETAILS OF COMPUTATIONAL COST AND EFFICIENCY
We provide comprehensive analysis of computational costs, including detailed comparisons of training time, inference time, and GFLOPs. These results verify the inherent efficiency of our ap- proach and demonstrate that our significant performance gains are primarily attributable to superior methodological design rather than increased computational overhead.
Table 14: Comparison of Training Time, Inference Time, and FLOPs with Dense Model and MoE baselines under Rectified Flow on ImageNet (256x256) after 500K training steps.
Activated Params.
Total
Inference Time (cfg=1.5)
Model (500K)
#Experts
# Training Time
GFLOPs
Params.
Dense-DiT-L DiT-L--low
458M
458M
166.67 GPU Hours
1.25 sec/sample
77.50
DiT-MoE-LL-lom -loww
E8A ISOU0
458M
1.163B
333.33 GPU Hours
1.49 sec/sample
77.50
EC-DiT-L-Flow
ESOIO
458M
1.163B
277.78
GPU Hours
1.49 scc/sample
77.50
DiffMoF-L--Foo
ESAI
458M
1.095B
333.33
GPU Hours
1.64 sec/sample
82.53
ProMoE-L-Flow
E14A1SIU1
458M
1.063B
333.33 GPU Hours
1.53 scc/sample
77.72
C.8 COMPARISON WITH OTHER UNSUPERVISED CLUSTERING METHODS
We conduct additional experiments comparing our ProMoE with a GMM-based routing baseline on the large (L) model. The results in Table 15 show that ProMoE consistently outperforms the GMM- based routing baseline across all metrics (FID and IS) and CFG scales under both top-1 (E14AI) 1) and top-3 (E14A3) activation settings. We argue that GMM, like other clustering- based routing schemes, relies purely on implicit learning and lacks explicit guidance, which is insufficient for robust expert specialization in vision MoE, especially for large-scale models. In contrast, ProMoE achieves better performance, highlighting the robustness and effectiveness our explicit semantic routing guidance strategies.
19

---

Published as a conference paper at ICLR 2026
Table 15: Quantitative comparison with GMM-based routing baseline under Rectified Flow ImageNet (256x256) after 500K training steps, evaluated with CFG scales of 1.0 and 1.5.
Activated Params.
Total Params.
cfg=1.0
cfg=1.5
Model (500K)
#Experts
FIDSOK
FIDS5K
IS 1
GMM-bused Routing
E14AISIU1
458M
1.063B
15.56
84.94
3.76
206.05
ProMoE-LL-Flow
E14AISIU1
458M
1.063B
11.61
100.82
2.79
244.21
GMM-based Routing
E14A3S1U1
558M
1.063B
15.44
86.09
3.72
208.15
ProMoF-LL.FFow
E14A3STU1
558M
1.0638
11.40
103.78
2.72
246.96
cl
112
Usese Stutoo
250
cla
100
Dines 2022
Donce
200
40
200
500
40B
500
100
200
303
400
503
100
200
300
400
Training Stop (K)
Training Step (K)
Training Stop (1)
Figure 9: More scaling results on model size.
100
clg-1.0
250
oo do
digm1.3
Police Letoo E1
200
Roke
1.0
#1050K
150
60
125
40
100
200
500
250
400
200
400
500
Thaining Step (K)
Training Step 00
Traiing Step (K)
Figure 10: More scaling results on the number of experts.
D MORE RESULTS ON SCALING BEHAVIOR
D.1 SCALING MODEL SIZE
Fig. 3(b) shows FID results for model size scaling at CFG=1.0. We additionally report FID results CFG=1.5 and Inception Score at CFG=1.0 and 1.5, as shown in Fig. 9. ProMoE consistently outperforms its dense counterparts, and ProMoE-L-Flow surpasses Dense-XL-Flos in terms of FID and Inception Score both CFG=1.0 and 1.5, despite using fewer activated parameters. These observations are consistent with those in Sec. 5.2.
D.2 SCALING THE NUMBER OF EXPERTS
We report Inception Score for scaling the number of experts at CFG=1.0 and 1.5, and FID at CFG=1.5, as shown in Fig. 10. Performance of ProMoE improves as the number of experts in- creases, demonstrating the scalability of our approach.
D.3 SCALING THE NUMBER OF EXPERTS ON K-MEANS-BASED ROUTING
We conduct an additional experiment varying the expert number for K-means-based routing (uncon- ditional routing is not used here). As shown in Table 16, K-means-base routing is highly sensitive to the expert number and fails scale effectively: increasing experts yields fluctuating performance with no clear gain, despite the increased parameter count. This contrast further highlights the scala- bility and robustness of ProMoR compared to clustering- routing.
20

---

Published as a conference paper at ICLR 2026
Table 16: Ablation of expert number in K-Means-hased Routing on base model size, trained with Rectified Flow on ImageNet (256x256) for 500K steps, evaluated with CFG-1.5. The setup uses activated expert and shared expert.
Number of Experts
3
5
13
21
FID 10K4
12.55
11.52
12.43
11,41
12.15
E MORE ABLATION STUDIES
E. ABLATION ON
Loss
As discussed in Sec. 4.3, the push-away term in our routing contrastive learning (RCL) serves role similar to load balancing. We verify this with an ablation in Table 17. Adding a conventional load-balancing loss on top of our method slightly degrades performance. We attribute this to RCL's explicit semantic guidance: it leverages token semantics to maintain diverse expert assignments, whereas load balancing loss only regularizes token counts and ignores assignment quality and se- mantics, thereby interfering with RCL. These results indicate that the semantic routing guidance from RCL more effective than traditional balancing losses.
Table 17: Ablation study using laa-bbllnncnng loss under Rectified Flow on ImageNet (566x266) after 500K training steps.
cfg=1.0
cfg=1.5
Model (500K)
FID50K |
FID50K J
1
w/ load-hal sead-halancing loss
24.98
59.04
6.53
151.37
w/o load- hal-balanniig loss
24.44
60.38
6.39
154.21
E.2 ABLATION ON Loss WEIGITT OF ROUTING CONTRASTIVE LEARNING
We vary the loss weight of routing contrastive learning (RCL) and report the results in Table 18. We observe that RCL is insensitive to the loss weight, as increasing it from to 10 yields only marginal gains. Therefore, which uses a weight of 10 based on this ablation study.
use default weight of for all experiments, except for
Table 18: Ablation study of ARCL in ProMoE-B-DDPM on ImageNet (256x256) after 500K train- ing steps.
cfg=1.0
cfg=1.5
ARCL (500K)
FIDSOK 4
IS f
FID50K
IS t
40.48
36.77
18.34
80.07
40.37
37.46
18.01
81.88
5
40.33
37.08
18.03
81.1
10
40.37
37.84
17.90
82.65
F USAGE OF LARGE LANGUAGE MODELS (LLMs)
In accordance with the ICLR 2026 policy, we report our use of a large language model (LLM) in preparing this manuscript. The L.M''s role was strictly confined to language polishing, such correcting grammar, refining wording, and improving readability. All scientific contributions, including the ideation, methodology, experimental design, and final conclusions, are entirely our own. The LLM was used solely writing-enhancemeet tool and did not contribute to the scientific aspects of the work. We have reviewed the manuscript and take full responsibility for its content.
21

---

Published as a conference paper at ICLR 2026
Algorithm ProMoE Layer (Training)
Input: X E Bx (input sequence), c E n (batch labels)
Variables: Ng (number of routed experts), K (number of activated routed experts), P E R (Learnable prototypes for routing), E (List of routed expert FFNs), E (Unconditional expert FFN), ES (Shared expert FFN), ÀRCL. (coef of Routing contrastive loss), т (temperature) Initialize final output
I: Initialize: O zeros zeros_like(X]
2: p*** Step 1. Functional Routing #**/
3: Mn (c =c= empty conditioning)
Get mask of unconditional image tokens + Get mask of conditional image tokens Get unconditional image tokens b Get conditional image tokens
4: Me f- -M
5: X= +- X(expand (M,.)
6: X + X(expand (M.)]
7: p*** Step 2. Unconditional Image Tokens Processing #**/
8: Ou < (X,)
9: O[M,] Ou
10: if (aan(MM)) then
11:
f*** Step 3. Prototypical Routing ***/
12:
X reshape (-1,
Flatten conditional image tokens D Get number of conditional image tokens Get mre-activation scores b Get token-expert affinity scores Get gating tensor and indices
13:
ne + [haphap[[O]
14:
Z E R
I.2 Normalize XX)) L.2 Normalize(P))
15:
S + Identity (Z)
16:
G R indices ZZeX (S, K)
17:
**** Step 4. Conditional Image Tokens Processing #**/
18:
Oc + zeros (irees))
19:
for i t to No
do m, < (indices == i) annnadimmmmm.... = 1) if (m;) then
Mask of tokens routed expert
21:
G, f- sum [[m,] (indices[m.) == i), dim =1) /m/m + Oc[mm G.unnquccze(() x E, [x,())))) end end for
Final gating scores > Update final output
23:
24:
25:
26:
Oc + reshape O X shape)
27:
[[M t [[M +
28: **** Step 5. Routing Contrastive Learning #*4/
29:
aux loss <- ARCL. x (xcu((X,, indices, P,
30: end if
31: /*** Step 6. Shared Expert Processing #**/
32:
O O E*
33: Return: O, aux loss
22

---

Published as a conference paper at ICLR 2026
Ang(liator dlsss distanos) (inrintra cluss distance)
Aug(inter claes distance)
Auglinter class distonse) Aug(intru class disturce)
23001
21.743
Avy (mtrru Elaes distance)
Layer Number: 7
Layer Number:
Layer Number:
Avg(iater class discance)
(vgee clags distance)
Au ginter class distance)) Angtintru class distunce)
(intrn class distance)
Argg(intu cluss distence)
Layer Number: 13
Layer Number: 15
Layer Number: 17
Figure 11: More t-SNE visualization results of Llama-3 8B on different layers.
Auglinter class distence) Angtintra plass distence)
Aug(inter class distance)
(inter class distunce)
0.733
Arg(mntr class distunce
0,748
Avg(intra class distence)
0.731
Layer Number: (Shallow Layer) Diffusion Timestep: 50 (Last Step)
Layer Number: 12 (Middle Layer) Diffusion Timestep: 50 (Last Step)
Layer Number: 24 (Deep Layer) Diffusion Timestep: 50 (Last Step)
Av Aiglinter- class distance)
Aug(inter class distunce)
Aug(inter class distance)
Aug (mtrra casss distance
Angtnntra class distance))
Auglintra- class distance
Layer Number: 12 Diffusion Timestep: 10 (Early)
Layer Number: 12 Diffusion Timestep: 25 (Middle)
Layer Number: 12 Diffusion Timestep: 40 (Late)
Figure 12: More t-SNE visualization results of DIT-XL/2 on different layers and diffusion timesteps.
23

---

Published as a conference paper at ICLR 2026
Figure 13: More samples generated by ProMoo-XXL-Flow after 2M iterations with cfg=4.0.
24

---

Published as a conference paper at ICLR 2026
photo a donut below cat
photo of two teddy bears
A photo of an elegant young woman
photo of dog
of
a photo of laptop
photo four donuts
photo of sheep
photo of blue train
of book
left of cow
photo of a pizza
photo of
photo of cake
photo of child
photo of bottle right
a photo of cake
photo of
A photo the Eiffel Tower on the moon's surface
photo a bench and sports elephant
photo a
Figure 14: Samples generated by ProMoE on the Text-to-Imanee task.
25


Reviews:
[{"review_id": "Oyce5uxnhm", "reviewer": "Reviewer_4DXQ", "summary": "This paper introduce a new expert routing method for Diffusion Transformers, by treating conditional and unconditional tokens independently. Routing guidance and contrastive learning are further introduced to enhance the performance.", "questions": "NA", "limitations": "", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The rationale for separating conditional and unconditional tokens is clear and well-founded.  \n2. The investigation into routing guidance and load balancing is insightful and valuable.", "weaknesses": "1. It would be beneficial to include ablation studies on dense models with conditional routing to determine whether the performance gain stems solely from conditional routing itself or requires combination with routing enhancements.\n\n2. Since one key advantage of MoE models is improved computational efficiency, the authors are encouraged to report training and inference times, as well as FLOPs, in comparison to both dense models and other MoE variants.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "TApOnZ1qIg", "reviewer": "Reviewer_1uA8", "summary": "The paper proposes ProMoE, a novel MoE framework for Diffusion Transformers that addresses the failure of prior MoE designs in vision via a two-step router with explicit routing guidance. It introduces conditional routing to separate functional roles and prototypical routing with learnable prototypes, enhanced by a routing contrastive loss.", "questions": "1.Could the authors provide quantitative statistics of expert utilization (e.g., token-per-expert ratios or activation entropy) to substantiate the claimed specialization and balance?\n2.Could the authors compare ProMoE with other unsupervised clustering methods that support top-K routing, such as GMM or deep clustering?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "1. The paper effectively addresses the core challenges of visual token redundancy and functional heterogeneity in Diffusion Transformers, introducing mechanisms that enable true expert specialization within the Mixture-of-Experts framework.\n2. The proposed method demonstrates strong and consistent scaling behavior across multiple model sizes, validating its robustness and efficiency under both Rectified Flow and DDPM training paradigms.", "weaknesses": "1. The experiments are conducted solely on ImageNet-1K for class-conditional generation, without evaluations on other datasets or modalities, which limits the evidence of generalization.\n2. The paper does not report quantitative expert utilization, such as the proportion of tokens or capacity per expert, making it hard to assess balance and specialization.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "UlA1kykaPg", "reviewer": "Reviewer_NvNP", "summary": "This paper introduces ProMoE, a framework that successfully applies Mixture-of-Experts (MoE) to Diffusion Transformers (DiTs), addressing why previous attempts failed. The authors argue that visual tokens, unlike language tokens, have high redundancy and functional differences, which hinders expert specialization. ProMoE solves this with a novel two-step router that first separates tokens by function and then assigns them to experts based on semantic content using learnable prototypes. This guided approach, enhanced by a new contrastive loss, enables strong expert specialization and achieves state-of-the-art results on ImageNet.", "questions": "See Weaknesses.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- This paper clearly diagnoses the problem of vision MoE and proposes an innovative ProMoE to solve it.\n\n- The ProMoE achieves validated, state-of-the-art results on the ImageNet benchmark.\n\n- The presentation is clear and easy to understand.", "weaknesses": "- What is the fundamental difference between prototypical routing and conventional MoE routing mechanisms, such as one using a standard linear layer? The paper introduces \"learnable prototypes\", but this seems functionally very similar to using the learnable weights of a linear layer to calculate token-expert affinities. Could you clarify what makes this prototypical approach a genuine innovation, rather than just a conceptual re-framing of a standard linear gating mechanism?\n\n- The routing mechanism in ProMoE appears to rely on pre-defined structures tailored for specific categories, unlike the autonomous expert specialization seen in LLMs. This raises questions about its generalizability—how would ProMoE handle open-ended conditional inputs, such as a natural-language prompt, rather than predefined categories? This design appears less flexible and general.\n\n- How is the number of experts determined, and what is the rationale for that specific choice? Are there more detailed ablation studies on the impact of varying the number of experts?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "NRWso52MR7", "reviewer": "Reviewer_vTUm", "summary": "The authors present ProMoE, a two-step routing framework for Diffusion-MoE models. In the first step, tokens are routed according to functional heterogeneity, sending unconditional tokens to dedicated unconditional experts. The remaining conditional tokens are then routed using a cosine similarity router to experts, with an additional routing contrastive loss designed to enhance intra-expert semantic similarity and inter-expert semantic differences.", "questions": "I'd strongly recommend the authors to include analysis of the kind seen in Table 4 but at training step=2M and cfg=4.0 across the MoE baselines and the dense baseline. This would provide a comprehensive analysis of the empirical benefits of ProMoE at convergence. Just choosing one size, ideally the largest, would be sufficient. If the authors can demonstrate the strong empirical results hold up at higher training step and cfg settings, I would consider raising my score, but for now it seems possible that the reported gains are too far from convergence to be meaningful. \n\nIf the authors could include an additional dataset that would also help boost the empirical contribution.", "limitations": "", "rating": 2, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "1. Experimental results on ImageNet are strong, with uniform improvements in generation quality and diversity\n\n2. The authors address scalability comprehensively, with experiments covering four model sizes. \n\n3. The paper is clear and easy to follow", "weaknesses": "**Limited technical novelty**. The authors claim the prototypical routing mechanism is novel, but I struggle to see the novelty here. The learnable prototypes appear to just be standard learnable expert embeddings and semantic similarity is computed using a cosine similarity between inputs and expert embeddings, which is fairly conventional in MoE [1,2, 3]. The authors do use an identity activation function which is non-standard, but this alone does not really qualify the entire routing mechanism as novel in my view. \n\n**Missing experimental results**. The reported results in Fig 3 show that ProMoE indeed offers substantive improvements over dense and MoE baselines up to 500k training steps, but the performance differences do appear to be converging on one another very quickly, with the differences starting to look more marginal towards 500K and with increasing cfg to 1.5. Given that the samples generated in figure 4 required 2 million training steps and a cfg of 4.0, it raises the question of whether the performance gains shown in Fig 3 and Table 4 are meaningful, as it seems the model is unlikely to be near convergence at 500K samples. Indeed, the loss curves seem to suggest that even at 1.2M steps the model is far from convergence. Given the trend in performance visible in Fig 3, it looks possible that at 2M and cfg=4.0 the improvement of ProMoE may no longer substantive, but the authors haven't included these important results. \n\n**Single dataset for experimental validation**. The authors present their experimental results on just ImageNet-1K. Though the authors do a good job of validating at multiple model sizes, the empirical contribution would be much more persuasive if the findings could be validated across multiple datasets. \n\n\n[1] On the representation collapse of spare moe [Chi et al, NeurIPS 2022]\n[2] Statistical advantages of perturbing cosine router in moe [Nguyen et al, ICLR 2025]\n[3] Sparse moe are domain generalizable learners [Li et al, ICLR 2023]", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)