Expert Race: A Flexible Routing Strategy for Scaling Diffusion Transformer with Mixture of Experts
Authors: Yike Yuan
Ziyu Wang
Zihao Huang
Defa Zhu
Xun Zhou
Jingyi Yu
Qiyang Min
Venue: ICML 2025 poster (2025)
Keywords: Mixture of Experts
Diffusion Models
Primary area: deep_learning->foundation_models
TL;DR: MoE with flexible routing strategy designed for diffusion models

Abstract: Diffusion models have emerged as mainstream framework in visual generation. Building upon this success, the integration of Mixture of Experts (MoE) methods has shown promise in enhancing model scalability and performance. In this paper, we introduce Race-DiT, a novel MoE model for diffusion transformers with a flexible routing strategy, Expert Race.  By allowing tokens and experts to compete together and select the top candidates, the model learns to dynamically assign experts to critical tokens. Additionally, we propose per-layer regularization to address challenges in shallow layer learning, and router similarity loss to prevent mode collapse, ensuring better expert utilization. Extensive experiments on ImageNet validate the effectiveness of our approach, showcasing significant performance gains while promising scaling properties.

Expert Race:
A Flexible Routing Strategy for Scaling Diffusion Transformer with Mixture of Experts
Yike Yuan Ziyu Wang" 12
Zihao Huang'
Defa Zhu'
Xun Zhou'
Jingyi Yut?
Qiyang Min
Abstract
20.0
Model Size
Diffusion models have emerged as mainstream framework in visual generation. Building upon this success, the integration of Mixture of Experts (MoE) methods has shown promise in enhancing model scalability and performance. per, we introduce Race-DiT, a novel MoE model for diffusion transformers with flexible routing strategy, Expert Race. By allowing tokens and experts to compete together and select the top candidates, the model learns to dynamically as- sign experts to critical tokens. Additionally, we propose per-layer regularization to address chal- lenges shallow layer learning, and router sim- ilarity loss prevent mode collapse, ensuring better expert utilization. Extensive experiments ImageNet validate the effectiveness of our ap- proach, showcasing significant performance gains while promising scaling properties.
17.5
B-cense
0.18
0,38
FID-50K ImageNet 256x556
15.0
1,08
2.7B
12.5
this pa-
10.0
B-moe-Kin32
L-dense
KL-dense
5.0
Meoo-4in32
6mme-4in322
42n32
2.5
0.1
0.2
0.3
0.4
0.5
0.6
0.7
0.8
Activation Parameters (B)
Figure I. Comparison between Dense and our MoE models. Our model significantly outperforms the Dense model with fewer acti- vation parameters.
ties but also an equally successful pursuit of larger models.
In the quest for larger models, the Mixture of Experts (MoE) approach, proven effective in scaling large language models (LLMs) (Jiang al., 2024), exhibits promising potential when incorporated into diffusion models. Essentially, MoE utilizes routing module to assign tokens among experts (typically, set of Feed-Forward Networks (FFN)) based on respective scores. This router module, pivotal to MoE's functionalityy, employs common strategies such as token- choice and expert-choice.
1. Introduction
Recent years have seen diffusion models earning consider- able recognition within the realm of visual generation. They have exhibited outstanding performance in multiple facets such as image generation (Ramesh et al., 2022; Nichol et al., 2021; Saharia et al., 2022; Rombach et al., 2022b; Esser et al., 2024), video generation (Gupta et al., 2024; Brooks et al., 2024), and 3D generation (Zhang et al., 2024; Ben- sadoun al., 2024). Thus, diffusion models have solidified their position pivotal milestone in the field of visual generation studies. Mimicking the triumph of transformer- based large language models (LLMs), diffusion models have effectively transitioned from U-Net to DiT and its variants. This transition yielded not only comparable scaling proper-
Meanwhile, observe that the visual signals processed by diffusion models exhibit two distinct characteristics pared to those in LLMs. First, visual information tends to have high spatial redundancy. For instance, signif- icant disparity in information density exists between the background and foreground regions, with the latter typically containing more critical details. Second, denoising task complexity exhibits temporal variation across different timesteps (Go et al., 2023). Predicting noise at the begin- ning of the denoising process is substantially simpler than predicting noise towards the end, as later stages require finer detail reconstruction. These unique characteristics necessi- tate specialized routing strategies for diffusion models.
'Equal contribution 'Seed-LLM-Model Team, ByteDance "ShanghaiTech University. Correspondence to: Qiyang Min <minqiyang wawiingiyangbbeedaancccoom www.ebyceeecdmm....... Jingyi Yu <yu- @ wwwgggaithhhhhccu.....
Proceedings of the ndnd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).
Consider these characteristies under the MoF, the presence

---

Expert Race: A Flexible Routing Strategy for Scaling Diffusion Transformer with Mixture of Experts
of routing module can adaptively allocate computational resources. By assigning more experts to challenging to- kens and fewer to simpler ones, we can enhance model utilization efficiency. Previous strategies like expert-choice anticipated this, but their routing design limit the assign- ment flexibility to image spatial regions without considering temporal denoising timestep complexity.
MoE Block
Scale
Expert
Expert
Expert 3
MoE Block
Token Assignment
Scale Shift
Top-K Selection
LayerNorm
In this paper, we introduce Race-DiT, novel family of MoE models equipped with enhanced routing strategies, Expert Race. We find that simply increasing strategy flexi- bility greatly boost the model's performance. Specifically, we conduct a "race" among tokens from different samples, timesteps, and experts, and select the top-k tokens from all. This method effectively filters redundant tokens and optimizes computational resource deployment by the MoE.
Flatten
Scale
Gating Logits
E1
E2
Mutt-Head Sell-Attetion
E3
T3
Scale Shift
Router
LayerNorm
Input Tokens
Expert Race introduces high degree of flexibility in token allocation within the MoE framework. However, there are several challenges when extending DiT to larger parame- ter scales using MoE. First, we observe that routing in the shallow layers of MoE struggles to learn the assignment, es- pecially with high-noise inputs. We believe this is due to the weakening of the shallow components in the identity branch of the DiT framework. To address this, we propose aux- iliary loss with layer-wise regularization to aid learning. Second, given the substantial expansion of the candidate space, to prevent the collapse of the allocation strategy, we extend the commonly used balance loss from single experts to expert combinations. This extension is complemented by router similarity loss, which ensures effective expert utilization by regulating pairwise expert selection patterns.
Figure 2. The Race-DiT Architecture. We replace the Multi-Layer Perceptron (MIP) with the MoE block, which consists of Router and multiple Experts. In Race-DiT, the token assignment is done once for all. Each token can be assigned to any number of experts, and cach expert can process any number of tokens (including zero).
2. Related Work
2.1. Mixture of Experts
Mixture of Experts (MoE) improves computational effi- ciency by activating only subset of parameters at a time and forcing the other neurons to be zero. Typically, MoE is used to significantly scale up models beyond their cur- rent size leveraging the natural sparsity of activations. This technique has been widely applied in LLMs first (Lepikhin et al., 2021; Fedus et al., 2022) and then extended to the vision domain (Riquelme et al., 2021). The most commonly used routing strategy in MoE is token-choice, in which each token selects a subset of experts according to router scores. For its variants, THOR (Zuo et al., 2022) employs a random strategy, BASELayer (Lewis et al., 2021) addresses the lin- ear assignment problem, HASHLayer (Roller et al., 2021) uses a hashing function, and MoNE (Jain et al., 2024) uses greedy top-k. All of these methods allocate fixed number of experts to each token. DYNMoE (Guo et al., 2024) and ReMoE (Wang et al., 2024c) activates different number of experts for each token by replacing TopK with threshold and using additional regularization terms to control the total bud- get. Also, some auxiliary regularization terms are applied to constrain the model to activate experts uniformly (Zoph et al., 2022; Dai et al., 2024; Wang et al., 2024a). Expert choice (Zhou et al., 2022) has been proposed to avoid load imbalance without additional regularizations and enhance routing dynamics, but due to conflicts with mainstream causal attention, it is less commonly applicd in large lan- guage models (LLMMs).
To validate the proposed method, we conducted experiments on ImageNet (Deng et al., 2009), performing detailed abla- tions on the proposed modules and investigating the scaling behaviors of multiple factors. Results show that our ap- proach achieves significant improvements across multiple metrics compared baseline methods.
In summary, our main contributions include
Expert Race, novel MoE routing strategy for diffu- sion transformers that supports high routing allocation flexibility in both spatial image regions and temporal denoising steps.
Router similarity loss, a new objective that optimizes expert collaboration through router logits similarity, effectively maintaining workload equilibrium and di- versifying expert combinations without compromising generation fidelity.
Per-layer Regularization that ensures effective learning in the shallow layers of MoE models.
Detailed MoE scaling analysis in terms of hidden split and expert expansion provides insights for extending this MoE model to diverse diffusion tasks.
2

---

Expert Race: A Flexible Routing Strategy for Scaling Diffusion Transformer with Mixture of Experts
2.2. Multiple Experts in Diffusion
3.2. Mixture of Experts
Diffusion follows multi-task learning framework that shares the same model across different timesteps. Con- sequently, many studies have explored whether perfor- mance can be enhanced by disentangling tasks according to timesteps inside the model. Ernie (Feng et al., 2023) and e- diff (Balaji et al., 2022) manually separate the denoising pro- cess into multiple stages and train different models to handle each stage. MEME (Lee et al., 2024) uses heterogeneous models and DTR (Park et al., 2024) heuristically partitions along the channel dimension. DyDiT (Zhao et al., 2024) introduce nested MLPs and channel masks to fit varing com- plexities across time and spatial dimensions. DiT-MoE (Fei et al., 2024), EC-DiT (Sun et al., 2024), and Raphael (Xue et al., 2024) have applied MoE architectures, learning to assign experts to tokens in an end-to-end manner. Compared with previous works, our methods proposed build on the MoE but further enhancing its flexibility on dynamically allocate experts on all dimensions to unleash its potential.
Mixture-of-Experts (MoE) is neural network layer com- prising router R and set {E,} of NE experts, each spe- cializing in subset of the input space and implemented as FFN. The router maps the input X E BxLx into token- expert affinity scores S E R BxLXX, trailed by a gating function G:
S = G((()))
The input will be assigned to subset of experts with top-k highest scores for computation and its output is the weighted sum of these experts' output. A unified expression is as follows:
S, if E TopK (S)
G =
(4)
0, Otherwise
MoB(X)= = E G(())+E(M)
(5)
iENE
3. Preliminaries
Before introducing our MoE design, we briefly review some preliminaries of diffusion models and mixture of experts.
where G E R BXLEE is the final gating tensor and Topk(;; KC) is an operation that builds a set with K. largest value in tensor.
3.1. Diffusion Models
To maintain a constant number of activated parameters while increasing the top-k expert selection, the MoE model often splits the inner hidden dimension of each expert based on the top-k value, named fine-grained expert segmentation (Dai et al., 2024). In the subsequent discussions, an "x-in-y" MoE means there are y candidate experts, with the top- x experts activated, and the hidden dimension of expert's intermediate layer will be divided by x.
Diffusion models (Ho et al., 2020) are a class of generative models that sample from noise distribution and learn a gradual denoising process to generate clean data. It can be seen as an interpolation process between the data sample To and the noise A typical Gaussian diffusion models formulates the forward diffusion process as
Dt Vorro + v1 0456
where E (O, I) is the Gaussian noise and 0g is mono- tonically decreasing function from to O. Diffusion mod- els use the neural networks to estimate the reverse denois- ing process po (2+) = NV (z)) Ex(z))) They are trained by minimizing the following objectives:
3.3. The Rationality of Using MoE in Diffusion Models
Diffusion models possess several distinctive characteristics.
Multi-task in nature, tasks at different timesteps pre- dicting the target are not identical. Prior works like e-diff (Balaji et al., 2022) validate this dissimilarity.
min E.....e [lly- - Fo c, 1)
(2)
Redundancy of image tokens. The information den- sity varies across different regions, leading to unequal difficulties in generation.
where t is the timestep which uniformly distributed be- tween 0 to T, e is the condition information. e.g. class labels, image or text-prompt. The training target y can be a Gaussian noise e, the original data sample 20 or velocity vāro.
Given these traits, MoE presents a suitable architecture for diffusion models. Its routing module can flexibly allocate and combine tokens and experts based on the predicted diffi- distribution culties. We consider the allocation process of computational resources. More challenging timesteps and complex image patches should be allocated to more experts. Achieving this requires a routing strategy with sufficient flexibility to distribute resources with broader degrees of freedom. Our method is designed following this principle.
e VI- āe
Early diffusion models used U-net (Dhariwal & Nichol, 2021; Rombach et al., 2022a) as their backbone. Recently, Transformer-based diffusion models (Peebles & Xie, 2023) with adaptive layer normalization (AdaLN) (Perez et al., 2018) have become mainstream, showing significant advan- tages in scaling up.
3

---

Expert Race: A Flexible Routing Strategy for Scaling Diffusion Transformer with Mixture of Experts
4. Taming Diffusion Models with Expert Race
4.1. General Routing Formulation
B
B
B
For computational tractability,, we decompose the original score tensor into two operational dimensions through per- mutation and reshaping, obtaining matrix S' E RDDxDD where
E
E
(a) Token Choice
(bb Expert Choice
(c) Expart Race
DR: Size of the expert candidate pool;
Figure 3. Top-KC Selection Flexibility. B: batch size; L: sequence
DD:: Number of parallel selection operations.
length; F: the number of experts. (a) Token Choice selects top-K. experts along the expert dimension for each token. (b) Expert Choice selects top-X tokens along the sequence dimension for each expert. (c) Expert Race selects top- across the entire set.
This dimensional reorganization enables independent top- selection within each row while preserving cross-row independence.
Following the sparse gating paradigm in (Zhou et al., 2022; Lepikhin et al., 2021), we control the MoE layer sparsity through parameter k, which specifies the expected number of activated experts per token. To satisfy system capacity constraints, the effective selection size per candidate pool is defined as:
Table 1. Specific design choices of different routing strategies.
Method
DA
DB
K
Token-choice Expert-choice Expert-Race
B x L
E
k
B E
L
&+L/E
B+L+E
B=Lak
k NE
K
DB-
(6)
4.2. Expert Race
The routing objective, aligning with the optimization frame- work in (Lewis et al., 2021), formalizes as the maximization of aggregated gating scores:
To address these limitations, we propose Expert-Race, which performs global top-K selection across all gating scores in single routing pass. The "Race" mechanism pro- vides optimal solution to Equation (7) by setting A 1, ensuring the selected K. elements are globally maximal. This design maximizes router flexibility to learn adaptive alloca- tion patterns, enabling arbitrary expert-to tote-tenttteknottt assignments and dynamic allocation based on computational demands.
DA
ΣΣ
max
-E
where To denotes the set of indices corresponding to the top-X values in the i-th row of S'
However, applying Expert-Race presents two challenges.
Suboptimal in Conventional Strategies As shown in Table 1, Figure 3, the unified framework generalizes exist- ing routing methods through top-X. selection in S'. How- ever, standard row-wise approaches like
Gating Function Conflict. While softmax over the expert dimension is standard for score normalization in existing routing strategies, it disrupts toksn score ordering in Race. Additionally, applying softmax across the full se- quence incurs high computational costs and risks numerical underflow as sequence length grows. We therefore explore alternative activation functions, finding through Table 2 that the i identity ((x) yields improved results.
and Expert-Choice exhibit inherent sub-optimality. These se- lection methods struggle to achieve optimal allocation in practice, as the required uniform distribution of top IC D elements across rows, which is necessary for attaining the theoretical optimum in Equation (7), rarely holds with real- world data distributions.
Training-Inference Mismatch. Batch-wise candidate ag- gregation creates a fundamental mismatch between training and inference. During training, samples influence each other's routing selection and timesteps are randomly sam- pled per batch, whereas inference operates on independent samples with consistent timesteps. Since timesteps directly control noise mixing levels, this inconsistency degrades gen- eration quality and can lead to model failure. At the same time, the mutual influence between samples during rout- ing selection causes unstable inference. To mitigate these effects, we propose a learnable threshold 7 that estimates the KC-th largest value through exponential moving average
In practical scenarios like image diffusion model training, generation complexity varies across two key dimensions: denoising timesteps (B) and spatial image regions (L). To address this computational heterogeneity, the routing mod- ule must dynamically allocate more experts to tokens with greater generation demands. However, the token-choice strategy, since D is constituted by dimensions B&T, both dimensions will receive an identical amount of activation experts. Expert-Choice mitigates this issue but remains constrained by its L-dimensional top-IC selection, limiting optimal allocation potential.

---

Expert Race: A Flexible Routing Strategy for Scaling Diffusion Transformer with Mixture of Experts
Algorithm Pytorch-style Pseudocode of Expert-Race
Token
Token
Token
Token
Gating logils
mi momentum
tau: ema updated threshold
Expert
input of shape (B, L, D)
Export 2
experts: list of FFN
Expert 3
Compute router logits for each token
logits router (x) $ (B, L, E)
Expert 4
score logits. flatten()
gates identity() (logits) activation
Balance but collapse
expect_k B
Token
Token
Token
Token
Gating logils
Get kthvalue and update threshold if training:
Expert
kth val kthvalue (score, k=expect_k)
mask = score >= vth
Expert
tau tau (1. - m) vth
clse:
Expert
mask score
Expert
Aggregate the token for each expert indices dispatch (mask) gather_inpon ix, indices)
Balance and diverse
gates Secher_gateed (gates, indices)
Figure 4. Toy examples of token assignment. Both of the show perfect load halance that each expert process two tokens. But in case above, experts and 2 are assigned the same token, as are experts and 4, where the 2-in-4 MoE collapse into 1-in-2. The example below shows more diverse assignment, making use of the expert specialization.
Process tokens by each expert and combine outs = [experts[i] (x[i]) for i in range (E)
result combine (outs, gates, indices)
(EMA) updates during training.
per expert, this approach fails prevent collapse between experts with similar selection rules. As shown in Figure if multiple experts follow the same rules for selecting to- kens, they are downgraded to one wider expert. Although such configurations satisfy balance loss constraints, they undermine the specialization benefits of ane-ggrined ex- pert design (Dai et al., 2024), ultimately degrading overall performance.
+ (1 m).
Σ i,C
(8)
DA
where S', s., represents the 1C-th largest element in the i-th row of S'. This adaptive threshold directly applied during inference, ensuring sample independence and consistent performance.
Pseudocode. We provide core pseudocode in PyTorch style in Algorithm illustrating how the expert selects the k-th largest logits and updates the threshold. Our algorithm is easy implement, requiring only minor modifications to the existing MoE framework.
Router Similarity Loss. To tackle this issue, propose maximizing expert specialization by promoting pairwise diversity among experts. Specifically, inspired by (Zbontar et al., 2021), we compute cross-correlation matrices and minimize their off-diagonal elements to encourage expert differentiation. Given the router logits E ((BxL)x we apply softmax along the expert dimension obtain normalized probabilities P, and compute two correlation matrices
5. Load Balancing via Router Similarity Loss
In MoE systems, balanced token allocation across experts remains a critical engineering challenge. For our pro- posed Race strategy, the increased policy flexibility imposes greater demands on routing balancing.
M' M P' PT=PTP
where M is the indicator matrix that Mi, 1 if expert selects the i-th token and 0 otherwise.
Mode Collapse in Balancing Loss. The conventional bal- ancing loss (Shazeer et al., 2017; Fedus al., 2022), origi-
nally designed for token-choice, promotes load balance by enforcing uniform token distribution across experts, thereby preventing dominance by small subset experts. How- ever, by only constraining the marginal distribution of scores
Then, we define the router similarity loss:
Crim 1
Σ W(i Pl. [.EELE
(10)
5

---

Expert Race: A Flexible Routing Strategy for Scaling Diffusion Transformer with Mixture of Experts
where W(1,)) is the weighting function defined as:
To mitigate this issue, we introduce a per-layer regulariza- tion that enhances gradients in a supervised manner with- out altering the core network structure. Specifically, given the hidden output h, from the [-th layer, we add a projec- tion layer H TR - RLxd to predict the final target y E R Riidd where and L' represent the number of patches before and after the patchify operation. The per-layer loss is defined as:
M
E.
i i
((G
Eisy Mi
(1I)
Mr
- Mi
In more detail, the off-diagonal elements denote the similar- ity between each pair of experts based on token selection patterns in the current batch. From a probabilistic perspec- tive, Pi, captures the joint probability of a token being routed to both expert and j. This formulation regular- izes consistent selectectin patterns across experts while promoting diverse expert combinations. Regarding the diag- onal elements, Pi,, represents geometric mean version of the balance loss, effectively encouraging individual expert utilization.
N Σ N
bm H(h,)) |
PPL = T Reo
(12)
where N is the total number of patches, and n is the patch index. In our implementation, the projection layer is inte- grated into the MLP router as an additional prediction head. By supervising the projection layer's predictions against final targets, we enhance shallow layer contributions during training, improving overall MoE performance.
6. Per-Layer Regularization for Efficient Training of Shallow Layers
7. Experiments
Implementation Details We follow training configura- tion in (Peebles & Xie, 2023) and conduct experiments on ImageNet (Deng et al., 2009). The metrics used include FID (Heusel et al., 2017), CMMD (Jayasumana et al.,, 2024), and CLIP Score (Radford al., 2021). We present se- ries of MoE size configurations, denoted as k-in-E where represents the total number of experts and k indicates the number of average activated experts. Additionally, we set the inner hidden dimension of each expert to be 1/k of its dense counterpart to ensure that the number of activated parameters remains the same. For all experiments, we train DiT-B and its 2-in-8 MoE variant for 500K iterations unless specified otherwise.
Individual Output Nomm
MoE
175
MoE-PLR
150
125
100
75
50
25
1
2
3
5
7
8
9
10
11
12
Layer Index
Figure 5. The norm of each block's output before added the shortcuts. The output norm increases rapidly in deep layers, result- ing in the weakening of shallow-layer components. This issue is alleviated with our proposed per-layer regularization.
Routing Strategy Expert-race enables extensive ex- ploration within the logit space during training, en- abling complexity-aware expert allocation across diffusion timesteps. As shown in Figure 6, our method dynamically assigns more experts to tokens at timesteps requiring higher image detail (lower timestep indices). In contrast, both token-choice and expert-choice strategies maintain fixed average expert allocations per timestep, lacking the tempo- ral dynamic allocation capability. Within the framework proposed in Section 4.1, we conducted ablation studies on routing strategies for combinations of different dimensions. See Section for more results.
DiT employs adaptive layer normalization (adaL.) when introducing conditions. In the pre-normalization (pre-norm) architecture, we observe that adaLN progressively ampli- fies the outputs of deeper layers. This causes the output magnitudes of shallow layers be relatively diminished, illustrated in Figure 5. This imbalance results in the learning speed of shallow layers lagging behind that of deeper lay- ers, which is detrimental to the MoE training process. This imbalance has both advantages and disadvantages. On one hand, the outputs from deeper layers are more accurate, and their larger magnitudes make them less susceptible to the substantial noise present in shallow layers, facilitating more precise regression results. On the other hand, due to the presence of the normalization module in the final layer, the component of shallow layers in the residuals is diminished, posing a risk of gradient vanishing and resulting in the learn- ing speed of shallow layers lagging behind that of deeper layers, which detrimental to the MoE training process.
Gating Function Table 2 shows that identity gating out- performs both softmax and sigmoid variants. In this experi- ment, we isolate other components (learnable threshold and regularizations) to verify the impact of the gating function on performance. We found that identity gating outperforms softmax and sigmoid under expert-race, and enhances both token-choice and expert-choice compared to softmax.
6

---

Expert Race: A Flexible Routing Strategy for Scaling Diffusion Transformer with Mixture of Experts
Experts assigned each token
Experts
Token
Choice
Expert
Choice
Expert
Race
Gemerate
Average # experts assigned to each token
Token Choice
Expert Choice
Expert Race
Sample
timestep
32
104
156
209
261
313
365
417
463
526
578
630
682
734
780
839
891
943
909
Figure 6. Average token allocation different time steps. Expert-Race assigns experts the complex denoising time steps, which occur at lower timestep indices that handle finer-grain image details.
Table 2. Routing strategy and gating function.
Table 4. Load balance for 4 4-in-32 MoE.
Routing
Gating
FID
CMMD
CLIP
Setting
FID..
MaxVio,
Comb!
Token Choice
17.28
.7304
21.87
No Constraint
11.38
6.383
18.98
Expert Choice
softmax
16.71
.7267
21.95
Balance Loss
11.67
2.052
72.36
Expert Race
16.47
.7104
21.97
Router Similarity
10.77
0.850
83.10
Token Choice
15.25
.6956
22.09
Expert Choice
sigmoid
15.73
.6821
22.06
Expert Race
13.85
.6471
22.23
100
Constraint Balance Loss
Token Choice
15.98
.6338
22.01
Combination Usage([%)
Router Similarity
Expert Choice
identity
15.70
.6963
22.04
80
Expert Race
13.66
.6317
22.25
60
Table Ablation study of core components.
40
Setting
FID.
CMMD,,
CLIP
Expert Racing (softmax)
16.47
.7104
21.97
20
Identity Gating
13.66
.6317
22.25
Leaamable Threshold
11.56
.5863
22.56
Per-Layer Reg.
8.95
.4847
22.94
16
32
64
128
Router Similarity
8.03
.4587
23.09
Number of Experts
Figure 7. Combination usage between different number of experts.
Load Balance Table 4 compares our proposed router sim- ilarity loss with conventional load balancing loss (Dai et al., 2024). This setting similar to the gating function ablation above, but here the MoE is 4-in-32 to further observe the impact of load balancing. The MaxVio (Wang et al., 2024b) metric measures how much the most violated expert exceeds its capacity limit. Combination Usage (Comb) measures the selection frequency of each expert pair across all tokens. It first sorts these frequencies descending order, then counts expert pairs that appear when the cumulative frequency reaches 95%, and finally ratios this count the total number of possible combinations. This metric thus estimates the utilization of all expert pairs (higher is better).
Core Components Table demonstrates the improve- ments brought by each component. Starting from the base- line of expert racing with softmax gating, incrementally added identity gating, learnable thresholds, per-layer regu- larization, and router similarity loss.
Scaling Law We first scale the base model sizes in the full pipeline, as detailed in Table 5. The comparison of our 4-in-32 MoE models with their Dense counterparts are across four sizes (B/M/L/XL). The number of activation parameter of Dense and MoE models are nearly identical. Results in Figure 1 demonstrate that model significantly outperforms the corresponding Dense models given same activation parameter count. Furthermore, our MoE-4in32 model surpasses the XL-Dense model with less than half the total number of parameters, further showcasing the effi- ciency of our model design.
Figure 7 further demonstrates the evaluation of MoE con- figurations across multiple scales 14.32-16,3,,64, 64,,2,6,,18), high- lighting our approach's 's capability to diversify expert activa- tion pattern compared to existing methods.

---

Expert Race: A Flexible Routing Strategy for Scaling Diffusion Transformer with Mixture of Experts
Table Model specifications and evaluation results of the comparison between MoE and Dense models.
Model Config.
Total Param.
Activate Param.
# Layers
Hidden
Heads
FIDJ
CMMD1
CLIP1
B/2-Denss
0.127B
0.127B
12
768
12
18.03
.7532
21.83
M/2-Dense
0.265BB
0.265BB
16
960
16
11.18
.5775
22.56
L/2-Dense
0.453B
0.453B
24
1024
16
7.88
.4842
23.00
XL/2-Dense
0.669B
0.669B
1152
.4388
23.27
B/2-MoE-44nn32
0.531B
0.135B
12
768
12
7.35
.4660
23.15
M/2-MoFE-4in32
1.106B
0.281B
16
960
16
5.16
.3507
23.50
L/2-MoFF-4in32
1,889B
0.479B
24
1024
16
4.04
.2755
24.12
XL/2-MoE-4in33
2.788B
0.707B
1152
3.31
.1784
24.68
1in4
Model Size
Table 6. Evaluation results of different MoE configurations with
FID-50K- ImageNet 256x256
9.5
0.36
same number activation parameters.
9.0
MoE
Hidden
Total
FID
CMMD2
CLIP
1in16
Config.
Split
Param.
8.5
1-in-4
0.297B
9.70
.5200
22.82
8.0
1-in-8
0.524B
9.05
.497
22.91
2in16
ain16
7.5
2in32
Ain32
Bin32
2-in-8
1-in-16
0,978B
0.297B
8.03
8.65
.4887
.0199
22.92
23.09
7.0
Aln64
Bino4
2-in-16
2
0.524B
7.78
.4607
23,06
2-in-32
0.977B
7.57
.4833
23.12
6.5
8in
4-in-16
0.297B
7.78
.6288
23.09
4-in-32
4
0.524B
7.35
.4600
23.15
6.0
4-in-64
0.977B
6.91
.4244
23.21
1
8-in-32
0.2977B
7.56
.5166
23.11
MoE Hidden Split Ratio
8-in-64
8
0.524B
6.87
4263
23.24
8-in-128
0.977B
6.28
.0155
23.35
Figure Scaling results of DīT-B in different MoF. configurations. Our method demonstrates linear performance improvement when increasing expert split ratios and expanding candidate expert pools.
ing selection freedom generally outperform conventional fxxe-dimension top-k selection approaches (such as token- choice/expert-choice)) Notably, Expert Race achieves the best performance across all evaluated routing strategies.
We further expand the model's size while maintaining the same number of active parameters under the same model configurations (DiT-B). The model expansion achieved by varying the hidden split ratios of experts and increasing the number of candidate experts. As shown in Figure 8 and Table 6, increasing both the number of candidate experts and splitting the hidden dimensions of MoE leads to im- provement in performance, highlighting the potential of our MoE architecture for scaling up.
Comparison with Leading Methods We further provide comparison with leading approaches shown in Table 8, where Samples is reported as training iterations x batch size. The classifier-free guidance scale is 1.5 for evaluation.
We train vanilla DiT (marked with *) and MoE model with Expert Race following the training protocol from the original DiT paper (Peebles & Xie, 2023), except that use larger batch size 1024 to improve memory utilization and train for only 1.75M steps. The learning rate remains unchanged at 10 Our MoE model achieves better perfor- mance with a similar activated parameter amount,
Extended Routing Strategies We extend the token- choice and expert-choice routing strategies by introducing varying degrees of routing selection flexibility, aiming to investigate how training-stage router freedom impacts final model performance, as shown Figure 9. All compared methods are trained for 500K iterations under identical configurations: full pipeline with DYT-B/2-MoE-2-11 architecture. To ensure experimental consistency, all ap- proaches employ learnable thresholds for inference queries. Specifically, we develop three new strategies: BL-Choice, OE-Choice, and LE-Choice, obtained through pairwise com- binations of selection dimensions Batch (B), Sequence Length (L), and Expert Count (E). Experimental results in Table demonstrate that strategies with higher train-
Computational Overhead During training, Expert-Race maintains the same total number of tokens assigned ex- perts other routing methods, differing only in their dis- tribution. The load imbalance can be mitigated through router similarity regularization, thus the training efficiency of Expert-Race is comparable to other routing methods. However, during inference, samples within a batch share identical timesteps. Variations expert activations across different timesteps cause fluctuations in computational cost.
8

---

Expert Race: A Flexible Routing Strategy for Scaling Diffusion Transformer with Mixture of Experts
B
B
B
1.
(a) Token Choice
(b) Expert Choice
(c) BL Choice
(d) BE Choice
(e) LE Choice
(1) Expert Race
Figure 9. Top- selection flexibility in extended routing strategies.
Table 7. Design Choices and Evaluation Results of Different Routing Strategies
Token Choice
Expert Choice
BL Choice BE Choice
LE Choice
Expert Race
DA
B E
B
L
B
DB
L
B L B E
L E
B L E
X
+ L/E
B k/EE
+
FIDL
9.50
10.13
9.08
8.28
8.89
8.03
CMMD4
.5202
.5639
.5145
.4636
.4871
.4887
CLIP
22.81
22.73
22.87
23.05
22.99
23.09
Table 8. Comparison with other methods on ImageNet 256x256.
Model Config.
Total
Activated
Samples
FIDq
1S1
Precisiont
Recall
ADM-G (Dhariwal & Nichol, 2021)
0.608B
0.608B
2.0M 256
4.59
186.70
0.82
0.52
LDM-8-G (Rombach et al., 2022a)
0.506B
0.506B
4.8M x 64
7.76
209.52
0.84
0.35
MDT (Gao et al., 2023)
0.675B
0.675B
2.5M x 256
2.15
249.27
0.82
0.58
MDT (Gao et al., 2023)
0,675B
0,675B
6.5M x 256
1.79
283.01
0.81
0.61
DIT-XL/2 (Peebles & Xie, 2023)
0.669B
0.669B
7.0M x 256
2.27
278.24
0.83
0.57
SiT-XL (Ma et al., 2024)
0.669B
0.669B
7,0M x 256
2.06
277.50
0.83
0.59
MaskDiT (Zheng et al., 2023)
0.737B
0.737B
2.0M x 1024
2.50
256.27
0.83
0.56
DiT-MoF-XI122 (Fei et al., 2024)
4. 4.105B
1,530B
7.0M x 1024
1.72
315.73
0.83
0.64
DiT-XL/2*
0.669B
0.669B
1.7M X 1024
3.02
261.49
0.81
0.51
2.788B
0.707B
1.7M 1024
2.06
318.64
0.83
0.60
8. Conclusion
Table Computational Overhead. Inference evaluation of the number of activated experts token on RaceDiT-XL/2 The first column refers to batch size and denotes the batch is composed of samples from different timesteps.
This paper proposes Expert Race, novel Mixture-of- Experts (MoE) routing strategy that enables stable and effi- cient scaling of diffusion transformers. Compared to previ- ous methods with fixed degrees freedom in expert-token assignments, our strategy achieves higher routing flexibility by enabling top-k selection across the full routing space spanning batch, sequence, and expert dimensions. This expanded selection capability provides greater optimiza- tion freedom, significantly improving performance when scaling diffusion transformers. To address challenges from increased flexibility, we propose an EMA-based threshold adaptation mechanism that mitigates timestep-induces dis- tribution shifts between training (randomized per-sample timesteps) and inference (uniform timesteps), ensuring gen- eration consistency. Additionally, per-layer regularization improves training stability, while router similarity loss pro- motes diverse expert combinations and better load balancing, as shown on x56x256 ImageNet generation tasks. As a gen- eral routing strategy, future work will extend Expert Race to broader diffusion- based visual tasks.
Mean
Max
Batch Size
t=0
t=100
t=0
t=1000
5.35
3.79
5.81
4.56
32
5.40
3.78
5.48
4.01
32*
4.03
4.27
As shown in Table 9, the peak computational cost reaches 5.81 experts per token, potentially limiting inference effi- ciency. To address this issue, one solution to construct batches containing samples from multiple timesteps. This be achieved through asynchronously interleaving ples from different timesteps within one batch, similar to pipeline parallelism. Such an approach can reduce the peak cost to 4.27 experts per token and the average cost to 4.03, significantly mitigating dynamic routing overhead.
9

---

Expert Race: A Flexible Routing Strategy for Scaling Diffusion Transformer with Mixture of Experts
Impact Statement
Improving text-to-image diffusion model with knowledge- enhanced mixture-of-denoisingg ference Computer Vision and Pattern Recognition (CVPR), pp. 10135-10145, 2023.
In IEEE Con-
We consistently adhere to ethical standards and responsi- ble AI development principles. Our models are trained on the widely-used, open-souree ImageNet dataset, ensuring diverse and representative outputs while minimizing the risk of biased or unsuitable content. We are committed to rig- orously reviewing the content generated by our models to prevent any harmful material.
Gao, S., Zhou, P., Cheng, M.-M., and Yan, S. Masked diffu- Sion transformer is a strong image synthesizer. In IEEE Conference on Computer Vision and Pattern Recognition (CVPR), pp. 23164-23173. 2023.
References
Go, H., Lee, Y., Lee, S., Oh, S., Moon, H., and Choi, S. Ad- dressing negative transfer in diffusion models. Advances in Neural Information Processing Systems (NeurIPS), 36: 27199-27231 2023.
Balaji, Y., Nah, S., Huang, X., Vahdat, A., Song, J., Zhang, Q.. Kreis, K., Aittala, M., Aila, T., Laine, S., et al. ediff- Text-to-image diffusion models with an ensemble of
expert denoisers. arXiv preprint arXiv:221.001324 2022.
Guo, Y., Cheng, Z., Tang, X., Tu, Z., and Lin, T. Dynamic mixture of experts: An auto-tuning approach for efficient transformer models. arXiv preprint arXhv:2405.14297 2024.
Bensadoun, R., Monnier, T., Kleiman, Y., Kokkinos, F,, Siddiqui, Y., Kariya, M., Harosh, O., Shapovalov, R., Graham, B., Garreau, E., et al. Meta 3d gen. arXiv preprint arXiv:2407.02500 2024.
Gupta, A., Yu, L., Sohn, K., Gu, X., Hahn, M., Li, F.- F., Essa, I., Jiang, L., and Lezama, J. Photorealistic video generation with diffusion models. In European Conference on Computer Vision (ECCV), pp. 393 411. Springer, 2024.
Brooks, T., Peebles, B., Holmes, C., DePue, W., Guo, Y., Jing, L., Schnurr, D., Taylor, J., Luhman, T., Luhman, E., Ng, C., Wang, R., and Ramesh, A. Video generation models as world simulators. 2024.
Dai, D., Deng, C., Zhao, C., Xu, R., Gao, H., Chen, D., Li, Zeng, W., Yu, X., Wu, Y., et al. Deepseekmoe: Towards ultimate expert specialization in mixture-of-experts lan- guage models. arXiv preprint urXiv:2401.06066 2024.
Heusel, M., Ramsauer, H., Unterthiner, T., Nessler, B., and Hochreiter, S. Gans trained by a two time-scale update rule converge to a local nash equilibrium. Advances in Neural Information Processing Systems (NeurIPS), 30, 2017.
Deng, J., Dong, W., Socher, R., Li, L.-J., Li, K., and Fei-Fei, L. Imagenet: A large-scale hierarchical image database. In IEEE Conference on Computer Vision and Pattern Recognition (CVPR). leee, 2009.
Ho, J., Jain, A., and Abbeel, P. Denoising diffusion proba- bilistic models. Advances in Neural Information Process- ing Systems (NeurIPS), 33:6840-6851. 2020.
Dhariwal, P. and Nichol, A. Diffusson models beat gans on image synthesis. Advances in Neural Information Processing Systems (NeurIPS), 34:8780-8794, 2021.
Jain, G., Hegde, N., Kusupati, A., Nagrani, A., Buch, S., Jain, P., Arnab, A., and Paul, S. Mixture of nested experts: Adaptive processing of visual tokens. arXiv preprint arXiv:2407.1 19985, 2024.
Esser, P., Kulal, S., Blattmann, A., Entezari, R., Müller, J., Saini, H., Levi, Y., Lorenz, D., Sauer, A., Boesel, F., al. Scaling rectified flow transformers for high-resolution image synthesis. In International Conference on Machine Learning (ICML), 2024.
Jayasumana, S., Ramalingam, S., Veit, A., Glasner, D., Chakrabarti, A., and Kumar, Rethinking fid: Towards better evaluation metric for image generation. In IEEE Conference on Computer Vision and Pattern Recognition (CVPR), pp. 9307-9315. 2024.
Fedus, W., Zoph, B., and Shazeer, N. Switch transformers: Scaling trillion parameter models with simple and ef- ficient sparsity. Journal of Machine Learning Research (JMLR), 23(12)):1-33 2022.
Jiang, A. Q., Sablayrolles, A., Roux, A., Mensch, A., Savary, B., Bamford, C., Chaplot, D. S., Casas, D. d. 1., Hanna, E. B., Bressand, F., et al. Mixtral of experts. arXiv preprint arXiv:2401.04083 2024.
Fei, Z., Fan, M., Yu, C., Li, D.,, and Huang, J. Scaling diffusion transformers 16 billion parameters. arXiv preprint ar A00iv:2407.1 1633, 2024.
Lee, Y., Kim, J., Go. H., Jeong, M., Oh, S., and Choi, S. Multi-architecture multi-expert diffusion models. AnA/A/ Conference on Artificial Intelligence (AAAI), volume 38, pp. 13427-13436, 2024.
Feng, Z. Zhang, Z. Yu, X., Fang, Y., Li, L., Chen, X.,
Lu, Y. Liu, J., Yin, W., Feng, S., et al. Ernie-ill 2.0:
10

---

Expert Race: A Flexible Routing Strategy for Scaling Diffusion Transformer with Mixture of Experts
Lepikhin, D., Lee, H., Xu, Y., Chen, D., Firat, O., Huang, Y., Krikun, M., Shazeer, N., and Chen, Z.. (GSS)ard: Scaling giant models with conditional computation and automatic sharding. In International Conference on Learning Rep- resentations (ICIR), 2021.
diffusion models. In IEEE Conference on Computer Vi- sion and Pattern Recognition (CVPR), pp. 10684-10695, 2022a.
Rombach, R., Blattmann, A., Lorenz, D., Esser, P., and
Ommer, B. High-resolution image synthesis with latent diffusion models. In IEEE Conference on Computer Vi- sion and Pattern Recognition (CVPR), pp. 10684-10695, 2022b.
Lewis, M., Bhosale, S., Dettmers, T., Goyal, N., and Zettle- moyer, L. Base layers: Simplifying training of large, sparse models. In International Conference on Machine Learning (ICML), pp. 6265-6274, 2021.
Saharia, C.,, Chan, W., Saxena, S., Li, L.., Whang. J., Denton, E. L., Ghasemipour, K., Gontijo Lopes, R., Karagol Ayan, B., Salimans, T., et al. Photorealistic text-to-image diffu- sion models with deep language understanding. Advances in Neural Information Processing Systems (NeurIPS), pp. 36479-36494, 2022.
Ma, N., Goldstein, M., Albergo, M. S., Boffi, N. M., Vanden- Eijnden, E., and Xie, S. Sit: Exploring flow and diffusion- based generative models with scalable interpolant trans- formers. In European Conference on Computer Vision (ECCV), pp. 23-40, 2024.
Nichol, A., Dhariwal, P., Ramesh, A., Shyam, P., Mishkin, P. McGrew, B., Sutskever, I., and Chen, M. Glide: Towards photorealistic image generation and editing with text-guided diffusion models. arXiv preprint arXiv:21 2021.
Shazeer, N., Mirhoseini, A., Maziarz, K., Davis, A., Le, Q., Hinton, G., and Dean, J. Outrageously large neural networks: The sparsely-gated mixture-of-experts layer. In International Conference on Learning Representations (ICLR), 2017.
Sun, H., Lei, T., Zhang, B., Li, Y., Huang, H., Pang, R., Dai, B., and Du, N. Ec-dit: Scaling diffusion transform- ers with adaptive expert-choie routing. arXiv preprint 24Xiv:2410.0222 2024.
Park, B., Woo, S., Go, H., Kim, J.-Y., and Kim, C. Denois- ing task routing for diffusion models. In International Conference on Learning Representations (ICLR), 2024.
Peebles, W. and Xie, S. Scalable diffusion models with transformers. In IEEE International Conference on Com- puter Vision (ICCV), pp. 1195-4205 2023.
Wang, L.,, Gao, H., Zhao, C., Sun, X., and Dai, D. Auxiliary- loss-free load balancing strategy for mixture-of-expper arXiv preprint arXiv: 15664, 2024a.
Perez, E., Strub, F., De Vries, H., Dumoulin, V., and Courville, A. Film: Visual reasoning with general conditioning layer. In AAAI Conference Artificial Intelligence (AAAI), volume 32, 2018.
Wang, L.,, Gao, H., Zhao, C., Sun, X., and Dai, D. Auxiliary- loss-free load balancing strategy for mixture-of-experss arXiv preprint arXiv:2408. / 15664, 2024b.
Wang. Z., Chen, J., and Zhu, J. Remoe: Fully differen- tiable mixture-of-experts with relu routing, arXiv preprint arXiv: 1471/, 2024c.
Radford, A., Kim, J. W., Hallacy, C., Ramesh, A., Goh, G., Agarwal, S., Sastry, G., Askell, A., Mishkin, P., Clark, J., et al. Learning transferable visual models from natural language supervision. In International Conference Machine Learning (ICML), pp. 8748-8763 8763. PMLR, 2021.
Xue, Z., Song, G., Guo, Q., Liu, B., Zong, Z., Liu, Y., and Luo, P. Raphael: Text-to-image generation via large mix- ture of diffusion paths. Advances in Neural Information Processing Systems (NeurIPS), 36, 2024.
Ramesh, A.. Dhariwal, P., Nichol, A., Chu, C., and Chen, M. Hierarchical text-conditional image generation with clip latents. arXiv preprint arXiv:2204.06122 2022.
Zbontar, J., Jing, I.., Misra, I., LeCun, Y., and Deny, S. Bar- low twins: Self-supervised learning via redundancy reduc- tion. In International Conference on Machine Learning (ICML), pp. 12310-12320. PMLR, 2021.
Riquelme, C.,, Puigcerver, J., Mustafa, B., Neumann, M., Je- natton, R., Susano Pinto, A., Keysers, D., and Houlsby, N. Scaling vision with sparse mixture of experts. Advances in Neural Information Processing Systems (NeurlPS), 34: 8583 8595, 2021.
Zhang, L., Wang, Z., Zhang, Q., Qiu, Q., Pang, A., Jiang, H., Yang, W., Xu, I., and Yu, J. Clay: A controllable large-scale generative model for creating high-quality 3d assets. ACM Transactions on Graphics (TOG), 43(4): 1-20, 2024.
Roller, S., Sukhbaatar, S., Weston, J., et al. Hash layers for large sparse models. Advances in Neural Information Processing Systems (NeurIPS), 34:17555-1756. 2021.
Zhao, W., Han, Y., Tang, J., Wang, K., Song, Y., Huang, G., Wang, F., and You, Y. Dynamic diffusion transformer. arXiv preprint 24.iv:4400.03454 2024.
Rombach, R., Blattmann, A., Lorenz, D., Esser, P., and
Ommer, B. High-resolution image synthesis with latent

---

Expert Race: A Flexible Routing Strategy for Scaling Diffusion Transformer with Mixture of Experts
Zheng. H., Nie, W., Vahdat, A., and Anandkumar, A. Fast training of diffusion models with masked transformers. arXiv preprint arXiv:2306.09) 09305, 2023.
Zhou, Y., Lei, T., Liu, H., Du, N., Huang, Y., Zhao, V., Dai, A. M., Le, Q. V., Laudon, J., et al. Mixture-of-expers with expert choice routing. Advances in Neural Infor- mation Processing Systems (NeurIPS), 35:7103-7114, 2022.
Zoph, B., Bello, Kumar, S., Du, N., Huang, Y., Dean, J., Shazeer, N., and Fedus, W. St-moe: Designing stable and transferable sparse expert models. arXiv preprint arXiv:2202.08940 2022.
Zuo, S., Liu, X., Jiao, Kim, Y. J., Hassan, H., Zhang, R., Gao, J., and Zhao, T. Taming sparsely activated transformer with stochastic experts. In International Con- ference Learning Representations (ICLR), 2022.
12


Reviews:
[{"review_id": "g6WGUozUaF", "reviewer": "Reviewer_qFyg", "summary": "This paper introduces Race-DiT, a novel approach for applying Mixture of Experts (MoE) to diffusion transformers with a flexible routing strategy called Expert Race. The key innovation is allowing tokens and experts to compete together in a global selection process, enabling dynamic allocation of computational resources based on token complexity. The authors also propose per-layer regularization to address shallow layer learning challenges and router similarity loss to prevent mode collapse among experts. Experiments on ImageNet demonstrate significant performance gains over dense models and other MoE routing strategies. The method shows promising scaling properties with both increased expert numbers and hidden dimension splits. Overall, Race-DiT achieves better FID scores with fewer activated parameters compared to traditional dense models, showing the potential of flexible MoE routing in diffusion models.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Per-layer Regularization lacks strong theoretical evidence. While the authors identify that shallow layers contribute less to DiT's pre-norm architecture, the fundamental necessity for layer balance is inadequately justified. The ablation study shows empirical improvements but doesn't conclusively demonstrate why balancing shallow and deep layers is optimal versus simply allowing deeper layers to dominate. The visual evidence in Figure 5 shows the phenomenon clearly, but the causal relationship between layer balance and overall performance could use more rigorous analysis. Additionally, the interaction between the MoE architecture and layer dynamics needs more theoretical grounding beyond empirical results.", "methods_and_evaluation_criteria": "The evaluation methodology generally makes sense for diffusion models.\n\nHowever, the computational efficiency comparison between different MoE strategies (Expert Race vs. Token/Expert Choice) is notably absent. Given that MoE's primary advantage is computational efficiency, understanding the overhead introduced by global token selection would be crucial. The paper focuses heavily on quality metrics (FID, CMMD, CLIP) but lacks throughput, latency, or memory usage comparisons. Additionally, while they show strong performance against dense models, evaluation against other contemporary diffusion MoE approaches like EC-DiT or Raphael is limited to the appendix rather than being central to the main evaluation.", "theoretical_claims": "I checked the correctness of the Analysis of Router Similarity Loss section. The derivation showing how router similarity loss extends traditional balance loss by considering pairwise expert interactions is mathematically sound.", "experimental_designs_or_analyses": "The experimental design is generally valid, but comparisons with other MoE-based diffusion research are somewhat limited. The internal comparisons between different MoE strategies (Expert Race vs. Token/Expert Choice) are comprehensive.\n\nHowever, the paper fails to clearly quantify the additional computational complexity and actual computational cost of Expert Race compared to Expert Choice or Token Choice. Since Expert Race requires flattening and processing the entire score tensor, it likely has poorer cache efficiency and higher computational overhead. This trade-off between quality improvement and computational cost is not explicitly analyzed, which is particularly concerning for an optimization technique like MoE where efficiency is a primary consideration.", "relation_to_broader_scientific_literature": "The paper makes a significant contribution by expanding MoE routing to consider all dimensions (B, L, E) with equal weight. Particularly valuable is connecting this approach to diffusion models' unique temporal characteristics across different denoising timesteps. By recognizing that diffusion complexity varies both spatially (across image regions) and temporally (across timesteps), the authors provide insights applicable beyond this specific implementation. This connection between model architecture design and diffusion's inherent multi-task nature advances the field's understanding of how to build more efficient generation models.", "essential_references_not_discussed": "The paper lacks a reference to \"Addressing negative transfer in diffusion models\" (Go et al., NeurIPS 2023), which provides explicit evidence for diffusion's multi-task nature through demonstrated negative transfer effects across timesteps. Race-DiT repeatedly emphasizes diffusion's varying temporal complexity as motivation for flexible routing, and Go et al.'s work would significantly strengthen this foundational claim with empirical support.", "other_strengths_and_weaknesses": "# Strengths\n\n- The model achieves impressively lower FID scores compared to dense models even when considering total parameter count, not just activated parameters, demonstrating true parameter efficiency.\n- The proposed Expert Race mechanism effectively addresses temporal complexity variation in diffusion models, showing marked improvement in early denoising timesteps where detail preservation is critical.\n- The authors perform extensive ablation studies that clearly isolate the contribution of each proposed component.\n- Router Similarity Loss shows strong theoretical grounding with clear connections to existing balance loss formulations.\n\n# Weaknesses\n\n- The dramatic change in layer dynamics after applying Per-Layer Regularization, particularly in layers 1-2, lacks detailed analysis of why these specific layers are most affected.\n- The paper provides limited analysis of expert specialization patterns - what specific features or timestep characteristics different experts learn to handle.\n- Testing is limited to ImageNet at 256×256 resolution without validation on more diverse datasets or higher resolutions.\n- The throughput-vs-quality tradeoff analysis is insufficient for practical deployment considerations.", "other_comments_or_suggestions": "The paper would benefit from a brief discussion of inference-time optimizations for Expert Race, since the global top-k selection might introduce latency challenges in production environments. Consider adding visualization of what different experts specialize in, perhaps showing attention maps or feature visualizations from different experts. Figure 6 currently shows allocation patterns but doesn't reveal what features drive those allocations. A minor typo appears in equation (7) where the summation indices could be clarified.", "questions_for_authors": "- How does Expert Race compare to Token/Expert Choice in terms of computational overhead during both training and inference? Could you provide wall-clock time comparisons on equivalent hardware?\n- Have you explored whether the benefits of Per-Layer Regularization are specific to MoE architectures, or would dense DiT models also benefit from this approach?", "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": "wboWkL7tJW", "reviewer": "Reviewer_h4on", "summary": "- This paper proposes an innovative MoE block integrated with a simple routing strategy and multiple regularization techniques within diffusion transformer frameworks, to concurrently optimize performance and computational efficiency. In particular, the authors broaden the routing strategy's exploration space by simultaneously considering token, batch, and expert dimensions, effectively addressing limitations inherent in existing methods that primarily emphasize either expert or token dimensions alone. Moreover, grounded in thorough empirical observations and logical justification, the authors substitute the conventional softmax routing mechanism with an identity function, complemented by a learnable threshold designed to alleviate inconsistencies arising from noise magnitude discrepancies between training and inference phases. Drawing inspiration from the Barlow Twins framework, the authors further propose an analogous loss formulation to combat mode collapse. Additionally, an auxiliary regularization branch is introduced to maintain stable output magnitudes across layers, thereby preventing gradient vanishing issues during model training. Empirical evaluations demonstrate that the proposed approach yields substantial performance gains on image generation benchmarks using the ImageNet dataset, thereby validating its effectiveness and advancing the state-of-the-art in this domain.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "- In general, the claims of this paper are supported by clear and convincing evidence.", "methods_and_evaluation_criteria": "- Evaluating the proposed methods solely on ImageNet for image generation may not be sufficient. Even when focusing exclusively on image generation (without considering other modalities), additional datasets such as COCO could offer valuable insights. Furthermore, FID is no longer considered the most suitable metric for open-domain image generation; alternative metrics like ImageReward or HPS might provide more meaningful evaluations.", "theoretical_claims": "- I have reviewed the theoretical claims presented in Section C of the supplementary material and found them to be correct.", "experimental_designs_or_analyses": "- Yes, I have assessed the soundness and validity of the experimental designs and analyses presented in the paper. Regarding the experimental designs, comparisons with similar approaches are limited; however, this limitation arises primarily due to the scarcity of existing methods related to diffusion models, particularly those employing DiT. Also, the analysis is adequate considering the contents in the supplementary material.", "relation_to_broader_scientific_literature": "- From my perspective, the key contribution of this paper is introducing a simple yet efficient MoE to the DiT architecture, accompanied by theoretical analysis. Given that very few related approaches exist, this paper makes a meaningful contribution to the broader scientific literature.", "essential_references_not_discussed": "- In general, this paper discusses related approaches adequately.", "other_strengths_and_weaknesses": "- This paper is well-crafted and effectively articulates both the proposed methodology and the corresponding experimental outcomes.\n- The implementation of the approach is methodical and straightforward, facilitating its practical applicability.\n- Comprehensive implementation details significantly enhance the reproducibility of the research.\n- However, certain aspects lack sufficient discussion, such as the influence of timestep and effects at different resolutions.", "other_comments_or_suggestions": "- The figures are somewhat low in quality; I suggest that the authors use vector graphics for better clarity.", "questions_for_authors": "- (1) From my perspective, the influence of timestep selection on the MoE, as well as the impact of varying input complexities (e.g., performance differences across categories), is particularly intriguing. Although the authors provide a basic illustration and a brief subsection discussing these aspects, incorporating theoretical analyses would further enhance the contribution of this paper.\n- (2) I'm curious about the consistency and generalizability of the proposed method across different resolutions or datasets, such as COCO.\n- (3) Additionally, what performance does the method achieve when evaluated using alternative metrics like ImageReward or HPS?", "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": "UUUMHnjICx", "reviewer": "Reviewer_ERH6", "summary": "The authors introduce a novel Mixture-of-Expert model for Diffusion Transformers. The main novelty aspect lies in its routing strategy, Expert Race, which allows to tune experts assignment not only by token, but also by batch element (ie, by time-step in the diffusion process). As a consequence, the model gains the additional flexibility of employing a different number of experts depending on the stage of the diffusion process it finds itself: arguably, this is beneficial, as at the beginning of the denoising procedure (when the image is still mostly noise) noise identification is simpler than towards the end (where the denoising procedure is almost complete, and more fine-grained details must be generated). While granting additional flexibility, this strategy comes with its challenges to ensure effective training. To address these, the paper proposes a number of adjustments to the training procedure, mainly the Router Similarity loss, to ensure experts specialisation, and a per-layer regularisation, to boost gradient in shallower model layers. Additional modifications ensure stability (dropping softmax as gating function) and consistency in predictions between training and inference (this is necessary due to the different distributions of time-steps in a batch seen by the model during these two phases).\nThe effectiveness of the model is tested by training on ImageNet and reporting FID, CMMD, and CLIP scores of the distribution of generated images, comparing them against alternative routing strategies available in the literature. Ablation studies confirm the advantages provided by the various adaptations introduced to the training procedure. Scaling laws confirm that the improvements are sustained also for large model sizes (up to ~O(1B)).", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Given the results provided, I’m reasonably convinced of the validity of the method.", "methods_and_evaluation_criteria": "Evaluation criteria are based on commonly used metrics for image quality evaluation (FID, CMMD, CLIP), and training is done on the ImageNet dataset. Both these choices make perfect sense to me.", "theoretical_claims": "The only theoretical claim casts the proposed Router Similarity loss to a generalised load balance loss. I haven't checked the derivation in detail, but seems legit at a glance.", "experimental_designs_or_analyses": "Code wasn’t made available, so I couldn’t skim through their implementation. From what I could read, the experiments design seems solid.", "relation_to_broader_scientific_literature": "The method proposed in the paper is grounded in a direct generalisation of MoE routing strategy for diffusion models, expanding both (Zoph et al., 2022, Zhou et al., 2022)", "essential_references_not_discussed": "The references provided are in my opinion sufficient to properly contextualise this work.", "other_strengths_and_weaknesses": "The paper is -for the most part- explained well, with the main idea illustrated clearly, and its element of novelty wrt previous work properly defined. The structure is solid, and the story reads well. Results are reasonably detailed and convincing.", "other_comments_or_suggestions": "- L63,C1 “flexibilities greatly boost” -> “flexibility greatly boosts”\n- L111,C1 “share” -> “shares”\n- L134,C2 “build a set with \\mathcal{K} largest value in tensor” -> “builds a set with the \\mathcal{K} largest value in a tensor”\n- L281,C2: “patchily operation ()” -> reference missing?", "questions_for_authors": "- Q1 On training vs inference - I’m honestly having a hard time understanding how your method works at inference/generation. If I got it correctly, at training time the routing strategy is deciding how to allocate the various experts while looking at the *whole* batch, which includes various samples at different time-steps / noise levels. This way, the model can learn to adaptively allocate more experts to less noisy tokens, if needed. Now, at generation, the time-steps distribution within the batch (if you’re doing batch generation at all!) is clearly different: first we de-noise samples that are all at time $t=T$, then all at $t=T-1$, and so on until $t=0$. I reckon this is the mismatch you hint at in the Training-Inference Mismatch paragraph. What I don’t understand is why the choice $\\tau$ in Alg1 should help in this regard? I was expecting that some sort of explicit (learnable) dependency on $t$ was gonna be baked into the threshold directly? \n- Q2 Considerations on workload during generation - From my understanding (correct me if I’m wrong), both Token-choice and Expert-choice guarantee constant workload per time-step at generation, in the sense that, for each $t$, there is a certain fixed number of $E_i(X_j)$ functions that must be evaluated (although the distributions of tokens per expert / expert per token might vary). In your case, this is not valid anymore (Fig6 highlights this). I was wondering what repercussions this had, if any, on the total wall-clock time / memory resources necessary for generation: it’s hard for the gauge this from the top of my head, as I reckon it depends on code implementation and parallelisability considerations. To help clarify what I mean, consider the extreme case where only 1 expert is invoked for each time-step between $T$ and 1, and all remaining $k-1$ expert invocations are used at the very last diffusion step $t=0$. If all experts invocations are perfectly parallel and we have infinite resources, at the end of the day the overall time for generation doesn’t change, but the last step would generally require a spike in computational demands, which might bottleneck operations and slow down generation. Does this make sense? If so, could you comment of this?", "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:

Meta-review: This work proposes a new technique to implement MoE that allows expert selection by considering both spatial image regions and temporal denoising steps. This allows for more flexibility and cost reduction. Moreover, the authors propose a mode collapse balancing loss and per-layer regularization to prevent vanishing gradients.

The experimental results seem to convince the reviewers in general.

The rebuttal from the authors was not delivered to the reviewers since it was submitted as AC private message slightly after the deadline. Therefore, there has been no modification to the original scores during the rebuttal period.
Decision: Accept (poster)