Diff-MoE: Diffusion Transformer with Time-Aware and Space-Adaptive Experts
Authors: Kun Cheng
Xiao He
Lei Yu
Zhijun Tu
Mingrui Zhu
Nannan Wang
Xinbo Gao
Jie Hu
Venue: ICML 2025 poster (2025)
Keywords: Diffusion Transformer
Mixture of Experts
Image Generation
Primary area: applications->computer_vision

Abstract: Diffusion models have transformed generative modeling but suffer from scalability limitations due to computational overhead and inflexible architectures that process all generative stages and tokens uniformly. In this work, we introduce Diff-MoE, a novel framework that combines Diffusion Transformers with Mixture-of-Experts to exploit both temporarily adaptability and spatial flexibility. Our design incorporates expert-specific timestep conditioning, allowing each expert to process different spatial tokens while adapting to the generative stage, to dynamically allocate resources based on both the temporal and spatial characteristics of the generative task. Additionally, we propose a globally-aware feature recalibration mechanism that amplifies the representational capacity of expert modules by dynamically adjusting feature contributions based on input relevance. Extensive experiments on image generation benchmarks demonstrate that Diff-MoE significantly outperforms state-of-the-art methods. Our work demonstrates the potential of integrating diffusion models with expert-based designs, offering a scalable and effective framework for advanced generative modeling.

Diff-MoE:
Diffusion Transformer with Time-Awar and Space-Adaptive Experts
Kun Cheng" 1
Xiao He
Lei Yu? Zhijun Tu
Mingrui Zhu
Nannan Wang'
Xinbo Gao
Jie Hu"
Abstract
DFSS 2712.
Switch-0
Ous
40
Swithho
Diffusion models have transformed generative modeling but suffer from scalability limitations due computational overhead and inflexible ar- chitectures that process all generative stages and tokens uniformly. In this work, introduce Diff- MoE, novel framework that combines Diffusion Transformers with Mixtureoo--Experts to exploit both temporarily adaptability and spatial flexi- bility. Our design incorporates expert-specific timestep conditioning, allowing each expert to process different spatial tokens while adapting to the generative stage, to dynamically allocate resources based on both the temporal and spa- tial characteristics of the generative task. Ad- ditionally, we propose a globally-aware feature recalibration mechanism that amplifies the rep- resentational capacity of expert modules by dy- namically adjusting feature contributions based input relevance. Extensive experiments on image generation benchmarks demonstrate that Diff-MoE significantly outperforms state-of-the- art methods. Our work demonstrates the potential of integrating diffusion models with expert- based designs, offering scalable and effective frame- work for advanced generative modeling. The code is available at https://// http://ggttt..//ee/eeeeee com/ waapcchhhnnnnff--feeeee
Swisch-
0us
ST30 21AXL
0un at.
FID-50K4 (400K iter.
10a Number of Parameters
10"
Figure 1. Comparison of FID performance and Parameters different sizes. We compare the class of DiT models trained for 400K iterations on ImageNet 256 256 generation with cfg 1.5. Diff-MoE achieves better FID across different sizes.
2023; Qi et al., 2023). However, their scalability remains limited by significant computational costs. Despite recent advancements, such as Diffusion Transformers (DiT) (Pee- bles Xie, 2023), the representation capacity and scala- bility of these models is constrained by two inherent limi- tations. First, the densely activated nature of DiT requires every parameter to be utilized across all stages of gener- ation, regardless of their relevance to the specific task or input, resulting in excessive computational overhead. Sec- ond, DiT employs a monolithic design that processes all diffusion timesteps uniformly, assuming that the same net- work components can handle both the coarse-graie gaarse-rrained global structures of early diffusion stages and the fine-grained local refinements in later stages. This lack of timestep-speciffi adaptation prevents the model from specializing in the dis- tinct requirements of different generative phases, ultimately hindering its representational capacity.
1. Introduction
Diffusion models (Ho et al., 2020; Song et al., 2020) have revolutionized generative modeling, achieving state-of-the- art results across various tasks, including image genera- tion (Dhariwal & Nichol, 2021; Rombach et al., 2022; Chen al., 2024; Esser et al., 2024), video generation (Blattmann et al., 2023; Yang al., 2024b) and editing (Brooks et al.,
Mixture-of-Experts (MoB) architectures have emerged as a scalable paradigm for deep learning (Shazeer al., 2017; Dai al., 2024), offering computational efficiency through dynamic parameter activation. By assigning distinet "ex- perts" to process different parts of the input, MoE models unlock unprecedented capacity without proportionally in- creasing computational cost. While MoE models excel in language tasks by routing inputs to specialized sub-netuwee networks ("experts"), their potential in diffusion- generative
'Equal contribution 'State Key Laboratory of Integrated Ser- vices Networks, Xidian University HHuawei Noah's Ark Lab, Cor- respondence to: Nannan Wang <nnwang@ xidian.edu.cn>.
Proceedings of the nd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

Diff-MoE: Diffusion Transformer with Time-Aware and Spacc-Adaptive Experts
frameworks remains underexplored. Prior works either par- tition experts across denoising stages (temporal) (Balaji et al., 2022; Feng et al., 2023) or spatial tokens (Fei et al., 2024; Sun et al., 2024), but none efficiently integrate both dimensions, leaving untapped performance gains, as illus- trated in Figure 2.
Expert
Gate
Gote
Gate
Takon
Token
! Timeltpp
Tolscn
Timadiw:
(4) Sporial MoFF
(b) Temporal
(e) Churs
In this paper, we propose Diff-Mof: diffusion framework that integrates timestep-aware and space-adaptive experts within Diffusion Transformers. Unlike naive multi-expert fusion methods (Xue et al., 2024) or rigid DiT-MoE hy- brids (Fei et al., 2024), our approach introduces expert- specific timestep conditioning: each spatial expert dynami- cally adapts to the current denoising stage via lightweight embeddings, enabling specialized processing of tokens based on both spatial content (e.g., object edges vs. tex- tures) and temporal context (early-stage global structures vs. late-stage refinements).
Figure 2. The comparison from mainstream diffusion-
MoE methods to the proposed Diff-MoE. (a) Standard spatial MoE: Experts are assigned based solely on spatial locations. (b) Temporal MoE: Experts are selected using a Top-K mechanism based on timestep information, allowing the model to focus on the dynamics of the diffusion process different stages. (c) The proposed Diff-MoE: Combines both spatial content and timestop information, modulating the input token to enable specialized processing of tokens based on both spatial content and timestep.
Furthermore, existing MoE designs often lack global input awareness, largely due to their application on autoregressive generation paradigms commonly used in language tasks. We propose an adaptive feature recalibration mechanism tailored specifically for image generation tasks, enhane- ing the representational capacity of expert modules. By incorporating globally contextual embeddings into the high- dimensional spaces within experts, this mechanism dynami- cally adjusts the contribution of different feature channels based on their relevance to the input distribution. This re- calibration strengthens the channel mixers' ability to model complex dependencies, emphasize informative features, and suppress noise, ultimately improving generation quality. To further mitigate parameter overhead, we project the expert- specific timestep conditioning and global context extraction processes into a compact, low-rank subspace, ensuring effi- cient scaling without sacrificing performance.
technique that enables compact and efficient modula- tion without sacrificing expressiveness, ensuring seal-
ability for large-scale models. Extensive experiments show that Diff-MoE significantly outperforms both dense DiT and state-of-the-ar expert- -aseed methods at the same parameter scale on conditional image gen- eration tasks.
2. Related Work
2.1. Diffuusion Transformer
Diffusion models (Ho et al., 2020; Rombach et al., 2022) have surpassed Generative Adversarial Networks (GANs) (Goodfellow al., 2020) the dominant paradigm for generative modeling, owing to their ability to capture complex data distributions and stable training dynamics. These models achieve state-of-the-art results across diverse domains, including image (Rombach et al., 2022; Chen et al., 2024), video (Mei & Patel, 2023), and 3D generation (Poole et al., 2022; Wang et al., 2024). Early architectures pre- dominantly employed U-Net backbones (Ronneberger et al., 2015; Ho et al., 2020; Nichol & Dhariwal, 2021), but recent advances like Diffusion Transformers (DiT) (Peebles & Xie, 2023) have demonstrated superior scalability by integrating transformer designs. Building on DIT, Stable Diffusion (SD3) (Esser et al., 2024) introduced a Multi-modal Dif- fusion Transformer (MM-DiT), scaling to 8B parameters for text-to-image synthesis. FLUX further expanded this to 12B parameters (Labs, 2023), showcasing DìT's poten- tial for high-resolution generation. However, scaling these dense architectures sample and per timestep (Hatamizadeh et al., 2024; Liu et al., 2024b; Karras et al., 2022)00iinee prohibitive com- putational costs, limiting practical deployment.
Our framework achieves significant improvements over both dense and expert-based sparse DiTs. Extensive experiments on standard benchmarks validate the superiority of Diff- MoE, achieving state-of-ther ofrthe-art results in quantitative met- rics. Our contributions can be summarized as follows:
A unified MoE-diffusion architecture optimizing tem- poral adaptation and spatial specialization through expert-specific timestep conditioning. This strategy en- ables dynamic expert adaptation across different stages of the diffusion process while keeping ooken-specifie flexiblity.
We introduce a global feature recalibration mechanism, dynamically adjusting the contributions of feature chan- nels based on their input relevance, empowering chan- nel mixers to model complex dependencies.
all parameters are activated per
To reduce parameter overhead, we employ a low-rank
2

---

Dift-MoE: Diffusion Transformer with Time-Aware and Space-Adaptive Experts
Diff-MoE Block
Mixture of Experts
Noise
Σ
Mixture Experts
Expert
Expert
...
Expert
Shared Expert
Linear and Reshape
Canditioning
Gloeal Contrntt
Assigned Hodden
AdaLN
Low-Rtan Condition Mapping
Context Extraction
MaE Router
Layer Norm
DWConv
Condition
Input Hidden
Diff-MoE Block x N
Expert
Multi-Head Self-Attention
Condition
Linzz
Patchify
Embed
Expert
Adais
Assigned Hidden
Linon
inesr
Output Hidden
AdaLN
Noised
Timestep
MLP
Latent
Global
Label
Conditioning
Input Tokens
Context
The Diff-MoF architecture extends the Diffusion Transformer (DiT) frame- work (Peebles & Xie, 2023). Inputs are first patchified into latent tokens, then processed through series of transformer blocks. Right: We replace the dense MLP in original DiT with carefully designed mixture of MLPs, enabling temporal adaptability (experts adjust diffusion timesteps) and spatial flexibility (token-speciff expert routing).
Figure 3. Overview of the Diff-MoF architecture. Left:
2.2. Mixture of Experts
3. Preliminaries
Mixture-of-Expeee (MoE) architectures (Shazeer et al., 2017; Dai et al., 2024) enhance model capacity by dynami- cally activating specialized sub-modules for distinct imputs, balancing expressiveness and efficiency. In the natural lan- guage processing (NLP) field, several studies have enhanced core component of MoE, including the optimization of routing mechanisms (Zhou et al., 2022: Huang et al., 2024) and the architectural design of expert models (Dai et al., 2024; Zoph et al., 2022). These improvements have suc- cessfully enabled large language models (LLM) (Dai et al., 2024; Yang et al., 2024a) to expand their capacity while min- imizing computational overhead. Inspired by the successful application of MoE NLP, the computer vision commu- nity has begun adopting MoE enhance the capacity of generative models, particularly diffusion models, which show significant potential. Time-centric approaches like DTR (Park et al., 2023) treat denoising as multi-task learn- ing through channel masking, while Switch-DiT (Park et al., 2024) employs sparse temporal gating, and MEME (Lee al., 2024) assigns experts timestep ranges. While these partially enhance capacity, their exclusive focus on tempo- ral partitioning ignores spatial complexity. Spatial MoE variants such DiT-MoE (Fei et al., 2024) (16.5B param- eters) and EC-DiT (Sun et al., 2024) (97B parameters via expert-choice routing) address teken-level diversity but fail to adapt to temporal dynamics of the denoising process. Our approach bridges this divide through joint spatiotemporal expert coordination, demonstrating efficient scalability.
3.1. Diffusion Models
Inspired by non-equilibrium thermodynamics, diffusion models learn to generate data by gradually reversing stochastic noise process. In the forward process, data sam- ples Do from the target distribution p(æo) are progressively corrupted by Gaussian noise through a series of timesteps t. The resulting noisy samples æ, are computed as follows:
q (m - (ats voraoo (1 o4.) I),
where ag defines the noise schedule. After a sufficient number timesteps, the data distribution approaches a standard Gaussian.
The denoising process aims to reverse the forward diffusion learning parameterized noise predictor eg (x, t). This predictor estimates the noise added at each step, enabling the recovery of æo from æ;, The reverse process is given as:
Pe (#t-1 N (mo (I,, yo t) (xx,, t)),
(2)
where (It, f) is parameterized by the noise predictor en (æ,, () and Σ (xt, () constant dependent on az.
3.2. Mixture of Experts
Mixture of Experts is modular neural network paradigm designed to enhance scalability and efficiency by dynami- cally activating a subset of specialized sub-models, called experts, during both training and inference. A standard MoE
3

---

Diff-MoE: Diffusion Transformer with Time-Aware and Spacc-Adaptive Experts
Table Configuraaions Diff-Mol architecture with different model sizes. align the Diff-MoE architecture configurations with DiT (Peebles& & Xie, 2023) in S. B and L. The St, B+ and XL configurations correspond to DiT-MoE (Fei et al., 2024). Specifically, "SEIAIS" indicates that a total of 8 sparse experts are used, with expert activated for each token, and expert shared by all tokens.
#Params
#Experts
#Blocks L
Hidden Dim. D
#Head n
ML.P Ratio
Dif--MoE-S
36M / 107M
8E1A1S
12
384
6
5/3
Dift-MoE-S+
77M / 211M
8E2A2S
12
384
Diff-MoE-B
137M / 402M
8EIAIS
12
768
12
5/3
Diff-MoE-B+
297M / 821M
8E2A2S
12
768
12
Diff-MoE-L
476M 1.4B
8EIAIS
24
1024
16
5/3
Dif--MoF-XI.
1.5B / 4.3B
8E2A2S
28
1152
16
4
system consists of N experts and a trainable gating mech- anism. The gate determines the importance score for each expert based on the input æ, enabling sparse activation:
into queries and keys before self-attention - technique proven to strengthen spatial dependency modeling both language (Touvron et al., 2023; Liu et al., 2024a) and diffu- sion models (Chu et al., 2025; Tian al., 2024; Esser et al., 2024). We also replace the standard MLP with gated linear units (GLU) (Shazeer, 2020) as feedforward network (FFN), introducing non-linearities that allow the model to learn more expressive feature representations. Inspired by pre- vious studies (Guo et al., 2022; Li et al., 2023; Tian et al., 2024), we introduce depthwise convolution layer prior to the MoE module, enhancing local feature extraction before channel mixing.
P(x) x softmax(W,, w))
where W are the gating weights. During each forward pass, only the top experts with the highest scores are activated, and their weighted sum is computed to produce the final output. By activating only small fraction of experts (k << N), MoE can achieve efficiently scaling without incurring the computational cost of dense architectures.
In addition to these architectural components, we introduce spatial Vixture-of-Expert (Dai et al., 2024), where each expert FFN shares an identical architecture. The spatial MoE enables different spatial tokens to be processed by N specialized experts, allowing the model to dynamically assign resources to different parts of the input data. This modular approach significantly improves the scalability of the model by reducing redundant computations, while still maintaining high representational capacity. The experts are sparsely activated, ensuring that computational resources are efficiently allocated according to the needs of each token. Inspired by (Dai et al., 2024; Fei et al., 2024), we incorporate additional N, experts as shared experts, where each token will be deterministically assigned to these shared experts, avoiding the inefficient utilization of parameters caused by knowledge sharing among experts. The transformer block of Diff-MoE can be formulated as:
4. Methodology
In this section, we propose Diff-MoE, a novel framework that integrates Diffusion Transformers with Mixture-of- Experts, exploiting both temporarily adaptability and spatial flexiblity. We begin by enhancing the original DiT with well-established advanced architecture designs and integrate spatial MoE, replacing dense MLPs with token-routed ex- perts to establish a nigh-performance baseline in Sec. 4.1. Next, we introduce Sec. 4.2, embedding lightweight temporal signals into indi- vidual experts to dynamically adapt their behavior across denoising stages- temporal adaptability with min- imal parameter overhead compared to stacking separate time/space experts. Furthermore, we propose a global fea- ture recalibration mechanism that injects contextual aware- ness into MoE modules in Sec. 4.3, enhaneing their ability to model spatially coherent structures in image generation tasks. Finally, to address the parameter burden introduced by these two modules, we apply low-rank decomposition technique in Sec. 4.4, ensuring the model remains efficient while incorporating the enhancements.
timestep conditioning in
x MHSA (AHSAAAdaNN(,,,o)) c)) + xx x
(4)
ac
Through the combination of these advanced architectural components and the MoE framework, Diff-MoB establishes a nigh-performance bascline that serves as the foundation for further innovations, discussed in later sections.
4.1. Basic Architecture Design
The primary objective of Diff-MoE, illustrated in Fig. 3, is to harness the powerful generative capabilities of DiT while introducing a more scalable and efficient architecture. We integrate rotary position embedding (RoPE) (Su al., 2024)
4.2. Expert-Specifie Timestep Conditioning
Conventional dense diffusion models employ static archi- tectures that uniformly process all timesteps despite the
4

---

DifT-MoE: Diffusion Transformer with Time-Aware and Spacc-Adaptive Experts
evolving demands of the denoising process-early stages requiring coarse structural modeling and later stages fine de- tail refinement. This uniform approach prevents the model from dynamically adapting to the differing needs of each timestep, limiting its representational power and efficiency.
generation tasks, where global coherence and local details
must be balanced for high-quality outputs.
To address this, we propose a global feature recalibration mechanism specifically designed to enhance the representa- tional capacity of expert modules in Diff-MoE. This mecha- nism allows each expert to incorporate global context into its processing, thereby improving the model's ability to capture complex dependencies between spatial tokens and enhancing its generative performance. The key idea behind this recalibration is the introduction of globally contextual embeddings Ig. which provide a high-level overview of the input data that is integrated into the expert's channel mixing process. This embedding enables each expert to dynami- cally adjust its contribution to the final output based on the broader context of the entire input sequence.
To address this issue, we introduce expert-specific timestep conditioning, building upon our spatial MoE baseline as described in Sec. 4.1. The core idea is to equip each ex- pert with dedicated timestep conditioning that allows it to specialize its behavior based on the current stage of the diffu- sion process. This enables the model to allocate processing resources more efficiently, focusing on relevant features at different stages of generation. The MoE block with expert- specific timestep conditioning can be expressed as:
NV
MoE(a;)= =
Σ Di.j (dN)
In practice, we first apply average pooling over the input tokens æ E Lxd along the spatial dimension, reducing the token sequence to global representation that captures the overall structure of the input data:
(5)
Pajj, Pi, E Topk < < N), K),
9i.j
0, otherwise,
Fg SSV.UU W giobal): )
(6)
where W global E Rd learnable weight matrices, SiLU represents the non-linear activation function and N is the expansion ratio of expert latent space. This step produces the globally contextual embedding T. E TR IxNd which encodes essential global features of the input. Next, we incorporate xy into each expert's forward pass, recalibrating the feature in high-dimensional space of GLU:
where Ij represents the j-th token of input x, is the timestep condition, and P., denotes the importance score of the j-th token corresponding to the i-th expert FFN, This expert-specific conditioning mechanism allows each expert to specialize in processing tokens at particular timesteps, where the nature of the information to be modeled differs significantly.
In standard DiT models, time-dependent AdaLN layers cre- ate entangled routing decisions, where temporal and spatial factors compete for optimization priority. This coupling forces routers to jointly consider denoising-stage dynamics and token complexity, often resulting in suboptimal equilib- ria and imbalanced expert utilization. The introduction of expert-specific timestep conditioning resolves this conflict by decoupling temporal adaptation from spatial specializa- tion. Each expert independently processes timestep embed- dings through dedicated conditioning pathways, enabling routers to focus exclusively on spatial token characteristics. This architectural refinement yields two synergistic bene- fits: routers develop sharper discriminative capabilities for token classification while experts achieve more fine-gradd -ne-grained timestep awareness.
(x,, = (SiLU(x, gate)) Ir W Wi
up Woown TFj (8)
where xy is the j-th token of input x. and W Wi E RdxNNd W Wiown ces. The channel-wise multiplication between the expert's token processing and the global context effectively recal- ibrates the expert's s output, enabling the model to adjust the contribution of different feature channels based on their relevance to the input distribution. By leveraging this mech- anism, Diff-MoE can model complex interactions more effectively, enhancing the overall quality of generation.
gate RdxNd
down E Ndxd are GLU weight matri-
up
4.4. Low-Rank Decomposition
To mitigate the parameter burden introduced by the pro- posed two modules, we apply low-rank decomposition to both components. For the eppert-specific time conditioning, we factorize the temporal projection matrix W, E dx of expert-specific AdaLN parameters (a, 3, 7) into low- rank components W.- = A, x B,, where A, E RR and A, E x3d with rank r < d. Similarly, the global feature recalibration weights W Wotobut& E d Nd are decomposed as W Wolobal = A x Bylobat, retaining only the dominant singular directions.
4.3. Global Feature Recalibration
While MoE architectures excel capturing local token de- pendencies, they often struggle with maintaining global in- put awareness. In typical autoregressive generative models, the absence of global context leads to suboptimal feature mixing and weak modeling of long-rat dependencies. This limitation becomes particularly prominent in image
5

---

Diff-MoE: Diffusion Transformer with Time-Aware and Spacc-Adaptive Experts
Table Quantitative Comparison with Dense DiTs on Ima-
Table 3. Quantitative Comparison with Time-aware Sparse
geNet x56x225 256 dataset with cfg = 1.0.
DIT's on ImageNet 256x dataset with cfg - 1.5. Specifi- cally, "woo RF" indicates training without rectified flow.
ImageNet 256x256, 400K, cfg=1.0
Model
#Params
FIDI
IS
ImageNet 2 256, 400K, cfg=1.5
Model
#Activated Params
FID
ISt
DIT-S/2
33M
66.59
20.68
SIT-S/2
33M
58.15
24.72
DTR-S/2
33M
37.43
38.97
SIT-LLaMA-S/2
33M
53.90
26.74
Switch-DiT-S/2
36M
33,99
42.99
Diff-MoF-S/2
36M/ 107M
44.27
34.27
Diff-MoF-S/2 w/ RF
36M
31.18
50.55
D/T-B/2
131M
42.84
33.66
DTR-B/2
131M
16.58
87.94
SiT-B/2
131M
35.54
42.33
Switch-DiT-B/2
144M
16.21
88.14
SiT-LLaMA-B/2
131
29.53
50.13
DifT-MoE-Bi BB2 w/o RF
137M
11.97
104.66
Dift- MoB--M
137M 402M
24.88
59.57
DiT-L/2
458
23.27
59.63
S/T-L/2
458
19.34
70.47
tivating 2 experts per token with 2 share experts (8E2A2S).
SiT-LLaMA-L/2
458
14.32
86.85
M/E-L/2
476M 1400M
13.98
90.72
Diff-MoE-B4o BBt for instance, scales to 821M total parame- ters (297M activated), preserving computational tractability while matching the scale of prior MoE-based DiT architec- tures (Fei et al., 2024). This hierarchical scaling ensures adaptability across hardware constraints, with the XL con- figuration (4.3B total parameters, 1.5B activated) targeting high-performance generation.
STR-S
DTR-3
SoRtE-OITS Ours K-
Switch-Dm Duri w/ RF
Implementation Details. We utilized pre-trained vari- ational autoencoder (VAE) model from Stable Diffu- sion (Rombach et al., 2022) to encode the image and decode the latent codes. By inputting images of 256 x 256 x 3, we obtain a latent representation with dimensions of 32 x 32 x 4. We conduct experiments on class-conditional generation tasks using the ImageNet dataset (Deng et al., 2009), which contains 1,281,167 training images across 1,000 distinct classes. We train all sizes of Diff-MoE for 400k iterations using the Adam optimizer with a learning rate of Ic-4. All models are trained with a batch size of 256. Following prior work (Park et al., 2023; Peebles & Xie, 2023), we apply exponential moving average (EMA) to the model pa- rameters during training. with decay factor of 0.9999, to enhance stability. Rectified flow and expert load balance loss are used by default, with further details provided in the supplementary material.
ning Staps
Taining Steps
Figure 4. Convergence comparison across different sizes. We compare the FID-10K at different training steps for expert-based DIT models on ImageNet 256 generation with cf9 = 1.5. For fair comparision, we report our results trained without rectified flow.
5. Experiments
5.1. Experimental Settings
Model Configurations. Diff-MoB is scaled across six configurations (S/B/L/XL and their Plus variants) to align with established Diffusion Transformer (DiT) benchmarks. As shown in Table 1, the base configurations (S/B/L) re- 12 transformer
We evaluate image generation qual- ity using Fréchet Inception Distance (FID) (Heusel et al., 2017), the standard metric for assessing sample quality and diversity in diffusion models. Adhering to the evaluation protocol from (Dhariwal & Nichol, 2021), FID scores calculated over 50K generated samples with 250 DDPM sampling steps. To complement this primary assessment, we additionally report Inception Score (IS) (Salimans et al., 2016) as a secondary diagnostic of sample fidelity.
Evaluation Metrics.
tain DIT's core hyperparameters--nnnddee blocks for S/B and 24 blocks for but replace dense MLPs with sparse MoE layers. Following previous MoE designs (Dai et al., 2024; Fei et al., 2024), we employ 8 task-specific experts by default and introduce shared ex- perts to mitigate knowledge redundancy. To maintain parity with DiT's activated parameter count, only 1 task-specific expert and shared expert are activated per token, while reducing the MLP ratio from 4 to 5/3. For example, Diff- MoE-S activates 36M parameters (107M total) under the 8EIAIS configuration (8 experts, activated per token, shared expert), achieving a 3x reduction active param- eters versus dense counterparts. The Plus variants expand capacity by reinstating DiT's original MLP ratio (4) and ac-
5.2. Comparison with State-of-the-Arts
Comparison with Dense DiTs. We perform a comprehen- sive comparison between Diff-MoE and dense DiT architec- tures across multiple model scales, including DiT (Fei et al.,
6

---

Dift-MoE: Diffusion Transformer with Time-Aware and Spacc-Adaptive Experts
Figure 5. Samples generated Diff-MoE 400K iterations. We show selected class-conditionl generated samples from Diff-MoE- XL/2 trained on ImageNet at 256 resolution.
2024), SIT (Ma et al., 2024), and SiT-LLaMA (Chu et al., 2024). DiT pioneered the integration of diffusion models with transformers, while SiT enhanced performance through improved prediction mechanisms and sampling strategies. SiT-LLaMA extends this progression by adopting a LLaMA- inspired architecture address diverse image generation tasks. All models were trained for 400K iterations and em- ployed classifier-free guidance (CFG) with a scale of 1.0 during the sampling phase. As shown in Table 2, Diff-MoE consistently outperforms dense counterparts at equivalent activation parameter counts. Notably, Diff- Diff-MoE-S reduces FID by 17.8% (53.90 44.27) and increases by 28.2% (26.74 + 34.27) compared state-of-the-art SiT-LLaMA- S. Scaling base scale models, Diff-MoE-B achieves 15.7% IS improvement (50.13 13 + 59.57) and 15.7% FID reduction (29.53 - 24.88) over SiT-LLaMA-BB At large scale, Diff-MoE further lowers FID from 14.32 to 13.98 (+3.0% improvement) and elevates IS from 86.85 to 90.72 (+4.5% gain). These results demonstrate -oE'' ability enhance both image quality (FID) and diversity (IS) while maintaining parameter efficiency.
Comparison with Expert- DiTs. We extend our eval- uation beyond dense DiT models to include expert- based architectures that employ conditional computation strate- gies, as illustrated in Fig. Current approaches specialize experts through either temporal decomposition (denoising stage selection) (Park et al., 2024; 2023) or spatial adapta- tion (token-wise routing) (Fei al., 2024). DiT-MoE (Fei et al., 2024) implements spatial specialization using Top-2 routing across 10 experts (2 shared + 8 task-specific), while DTR (Park et al., 2023) applies temporal windowing to acti- vate channel-specife experts. Switch-DiT (Park et al., 2024) explicitly assigns different denoising steps dedicated ex- pert modules. Unlike these dimension-specifie approaches, our proposed Diff-MoE jointly optimizes expert selection across both temporal and spatial axes. To ensure fair com- parison, we construct two model families with matched activation parameters: one aligned with temporal MoE base- lines and another with spatial MoE counterparts. All models follow standardized training protocols (Park et al., 2024) (400K iterations) and sampling parameters (CFG scale 1.5). Table demonstrates Diff-MoE-S" superiority over tempo-

---

Diff-MoE: Diffusion Transformer with Time-Aware and Spacc-Adaptive Experts
Table Quantitative Comparison with Space-adaptive Sparse
Table 6. Ablation Study MoB Design on ImageNet 256x256
DiT' on ImageNet x56 dataset with efg = 1.5.
dataset with cfg 1.5.
ImageNet 256x 256, 400K, cfg=1.5
ImageNet 256x256, 400K, cfg=1.5
Model
#Activated Params
FID
1ST
Model
#Acti. Params
FID
ISt
DiT-MoE-S/2
71M
17.92
80.50
MoE Baseline
38.37M
32.74
47.19
Difr-MoE-Se S/2+
77M
9.27
126.22
+ Expert- Spec. Time Cond.
49. 49.02M
28.71
54.57
+ Global Feat. Recalibration
53.74M
26.33
59.54
DiT-MoF-BAA
286M
7.19
147.67
Low Rank
41.84M
27.59
55.62
Diff-MoE- -B2+
297M
4.00
195.76
DiT-MoE-Xo X1.22
1.5B
3.42
214.73
Diff-MoE-XL/2
1.5B
2.69
262.22
(+4.1% 1% gain), while depthwise convolutions achieve FID
35.85 (+6.5% improvement) with minimal parameter over- head (+0.8%). The complementary integration of spatial MoR yields a final FID of 32.74 (+9.3% reduction).
Table 5. Ablation Study of Basic Architecture Design on Ima- geNet 2 256 dataset with efg = 1.5.
ImageNet 256x256, 400K, cfg=1.5
Model
#Params
FID4
ISt
MoE Design. We systematically evaluate the core com- ponents of Diff-MoE's spatiotemporal expert architecture through controlled ablations. As evidenced by Tab, 6, inte- grating expert-specific timestep conditioning into the spa- tial MoE framework substantially reduces FID from 32.74 to 28.71, confirming the necessity of modeling denoising- stage dynamics in expert design. Further performance gains emerge from our global feature calibration module, which lowers FID to 26.33 by injecting global context into locally processed tokens demonstrannns that global guidance en- hances local feature coherence despite expert specialization. To ensure parameter efficiency comparable to baseline DiT models, we implement low-rank decomposition techniques (r=64) that reduce total parameters by 28%. This optimiza- tion incurs a moderate performance trade-off (FID 27.59), yet still maintains superior results over standard spatial MoE implementations. The residual performance gap highlights our architecture's effectiveness in balancing capacity and efficiency.
DIT-S/2
32.96M
43.02
35.80
RoPE 2D
32.96M
39.80
39.10
+ GLU
32.94M
38.20
41.43
Depth Wise Conv
33.06M
35.85
42.85
Spatial MoE
38.37M
32.74
47.19
ral MoE baselines, achieving 12.3% FID reduction versus DTR-S and 8.3% improvement over Switch-DiT-S. This performance gap amplifies with model scale, confirming our method's 's enhanced scalability. Complementary com- parisons with spatial MoE architectures in Table reveal consistent superiority over DiT-MoE variants. Most notably, Diff-MoE-XL/2 achieves a 27% FID reduction (3.42
2.69) compared to DiT-MoE- XL/2, while maintaining sim- ilar number of parameters. The accelerated convergence rate of Diff-MoE, evidenced by training dynamics in Fig. 4, underscores the efficiency of our spatiotemporal paradigm. These consistent improvements across architectural scales and optimization timelines highlight the effectiveness of joint spatial-temporal expert coordination.
6. Conclusion
In this paper, we presented DifT-MoE, a novel framework that integrates Diffusson Transformers with Mixture-of- Experts (MoE) address the scalability and representa- tional limitations of existing diffusion models. Our archi- tecture introduces timestep-aware experts that dynamically adapt to denoising stages- stages-speciallzinn specializing in coarse semantic
5.3. Ablation Study
Basic Architecture Design. Building upon the founda- tion of DIT, we have made the following improvements to the network architecture: 1. Rotary Position Embed- ding (RoPE): Integrated into self-attention queries/keys to strengthen spatial awareness; 2. Gated Linear Unit (GLU): Replaces standard FFN MLP layers to enhance nonlinear feature learning: 3. Depthwise Convolution layers before the expert model: Boosting local feature extraction. We quanti- tatively validated the effectiveness of these modifications the conditional generation task using the ImageNet dataset. Quantitative evaluations on ImageNet conditional genera- tion (Table 5) demonstrate progressive improvements: RoPE integration reduces FID from 43.02 to 39.80 (+7.5% im- provement), GLU substitution further lowers FID to 38.20
shaping during early diffusion steps and fine-grained de- tail refinement in later stages-paired with adaptieee experts that perform localized feature processing through token-level routing. A global feature recalibration mech- anism further enhances representational power of expert modules by adapting feature contributions based on their relevance to the input. Our framework outperforms both dense DiT and existing MoE-based methods, offering a robust solution for nigh-performance generative modeling.
Limitation. While our experiments empirically validate
8

---

Dift-MoE: Diffusion Transformer with Time-Aware and Spacc-Adaptive Experts
Diff-MoE's efficacy 400K training iterations, acknowl- edge that computational resource constraints preclude eval- uation of large-scale variants (e.g., Diff-MoE-XL and DiT alignment protocols at 7M steps). This limitation leaves open the full exploration of our method's potential at ex- treme scales, which is left as our future work.
Dai, D., Deng, C., Zhao, C., Xu, R., Gao, H., Chen, D., Li, J., Zeng, W., Yu, X., Wu, Y., et al. Deepseekmoe: Towards ultimate expert specialization in mixture-of-experts lan- guage models. arXiv preprint arXiv:2401.06000 060666, 2024.
Deng. J., Dong, W., Socher, R., Li, I..JJ,, Li, K., and Fei-Fei, L. Imagenet: A large-scale hierarchical image database. In 2009 IEEE conference on computer vision and pattern recognition, pp. 248-244222 248-255, Icee, 2009.
Impact Statement
This paper presents work whose goal is advance the field of image generation. Similar to other image generation methods, our approach must be used cautiously to prevent potential misuse.
Dhariwal, P. and Nichol, A. Diffusion models beat gans on image synthesis. Advances in neural information processing systems, 34:8780-8708 8794, 2021.
Esser, P., Kulal, S., Blattmann, A., Entezari, R., Müller, J.,
Acknowledgments
Saini, H., Levi, Y., Lorenz, D., Sauer, A., Boesel, F,, et al. Scaling rectified flow transformers for high-resolution image synthesis. In Forty-first International Conference on Machine Learning. 2024.
This work supported in part by the National Natural Science Foundation of China under Grants U22A2096 and 62036007, in part by Scientific and Technological Innova- tion Teams in Shaanxi Province under grant 2025RS-CXTD- 011, in part by the Shaanxi Province Core Technology Research and Development Project under grant 2024QY2-
Fedus, W. Zoph, B., and Shazeer, N. Switch transformers: Scaling to trillion parameter models with simple and ef- ficient sparsity. Journal of Machine Learning Research, 23(120): 2022.
GJHX-1 11, in part by the Fundamental Research Funds for the Central Universities under CranttQT/X230422, in part by the Innovation Fund of Xidian University.
Fei, Z., Fan, M., Yu, C., Li, D., and Huang, J. Scaling diffusion transformers to 16 billion parameters. arXiv 2024.
References
preprint arXiv:2407.
Balaji, Y., Nah, S., Huang, X., Vahdat, A., Song, J., Zhang, Q., Kreis, K., Aittala, M., Aila, T., Laine, S., et al. ediff- i: Text-to-image diffusion models with an ensemble of expert denoisers. arXiv preprint arXiv:22/1.01322 2022.
Feng, Z., Zhang, Z., Yu, X., Fang, Y., Li, L., Chen, X., Lu, Y., Liu, J., Yin, W., Feng, S., al. Ernie-vilg 2.0: Improving text-to-image diffusion model with knowledge- enhanced
In Proceedings of the IEEE/CVF Conference Computer Vision and Pattern Recognition, pp. 10135-10145, 2023.
Blattmann, A., Dockhorn, T., Kulal, S., Mendelevitch, D., Kilian, M., Lorenz, D., Levi, Y., English, Z., Voleti, V., Letts, A., et al. Stable video diffusion: Scaling latent video diffusion models to large datasets. arXiv preprint arXiv:2311.15127, 2023.
Goodfellow, I., Pouget-Abadie, J., Mirza, M., Xu, B., Warde-Farley, D., Ozair, S., Courville, A., and Bengio, Y. Generative adversarial networks. Communications of the ACM, 63(11):139-144 2020.
Brooks, T., Holynski, A., and Efros, A. A. Instructpix2pix: Learning to follow image editing instructions. In Proceed- ings of the IEFF/CVF Conference on Computer Vision and Pattern Recognition, pp. 18392-18402, 2023.
Guo, J., Han, K., Wu, H.,, Tang, Y., Chen, X., Wang, Y., and Xu, C. Cmt: Convolutional neural networks meet vision transformers. In Proceedings of the IEEE/CVF conference computer vision and pattern recognition, pp. 12175- 12185, 2022.
Chen, J., Gc, C., Xie, E., Wu, Y., Yao, L., Ren, X., Wang. Z., Luo, P., Lu, H., and Li, Z. Pixart- Pixart-\sigmaa Asigma: Weak-to-strone training of diffusion transformer for 4k text-to-image generation. arXiv preprint arXiv:2403. 04692, 2024.
Hatamizadeh, A., Song, J., Liu, G., Kautz, J., and Vahdat, A. Diffit: Diffusion vision transformers for image generation. In European Conference on Computer Vision, pp. 37-55. Springer, 2024.
Chu, X., Su, J., Zhang, B., and Shen, C. Visionllama: A unified laaaa backbone for vision tasks. In European Conference Computer Vision, 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, 30, 2017.
Chu, X., Su, J., Zhang, B., and Shen, C. Visionllama: A unified lamaa backbone for vision tasks. In European Con- ference on Computer Vision, pp. 1-18. Springer, 2025.

---

Dift-MoE: Diffusion Transformer with Time-Aware and Spacc-Adaptive Experts
Ho, J., Jain, A., and Abbeel, P. Denoising diffusion proba- bilistic models. Advances in neural information process- ing systems, 33:6840-6851, 2020.
Park, B., Go, H., Kim, J.-Y., Woo, S., Ham, S., and Kim, C. Switch diffusion transformer: Synergizing denois- ing tasks with sparse mixture-of-expert, arXiv preprint arXiv:2403.09176 2024.
Huang, Q., An, Z., Zhuang, N., Tao, M., Zhang, C., Jin, Y., Xu, K., Chen, L.,, Huang, S., and Feng, Y. Harder tasks need more experts: Dynamic routing in moe models. arXiv preprint rrXiv:2403.076553 2024.
Peebles, W. and Xie, S. Scalable diffusion models with transformers. In Proceedings of the IEEE/CVF Interna- tional Conference on Computer Vision, pp. 4195 4205, 2023.
Karras, T., Aittala, M., Aila, T., and Laine, S. Elucidating the design space of diffusion-based generative models. Advances in Neural Information Processing Systems, 35: 2022.
Poole, B., Jain, A., Barron, J. T., and Mildenhall, B. Dream- fusion: Text-to-3d using 2d diffusion. arXiv preprint arXiv:2209. 2022.
Labs, B. F. Flux. flack-forest-- Cettoorsstlaab///e labs ffux, 2023.
Qi, C., Cun, X., Zhang, Y., Lei, C., Wang, X., Shan, Y., and Chen, Q. Fatezero: Fusing attentions for zero-shot text- based video editing. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pp. 15932- 15942, 2023.
Lec, Y. Kim, J., Go, H., Jeong, M., Oh, S., and Choi, S. Multi-architeccure multi-expert diffusion models. In Proceedings of the AAAI Conference on Artificial Intelli- gence, volume 38, pp. 13427-13436, 2024.
Rombach. R., Blattmann, A., Lorenz, D., Esser, P., and Ommer, B. High-resolution image synthesis with latent diffusion models. In Proceedings of the IEEE/CVF con- ference on computer vision and pattern recognition, pp. 10684-10695, 2022.
Li, Y., Zhang, K., Cao. J., Timofte, R., and Van Gool, L. Localvit: bringing locality to vision transformers (2021). arXiv preprint arXiv:2 1rXiv:2104.0500, 2023.
Liu, A., Feng, B., Xue, B., Wang, B., Wu, B., Lu, C., Zhao, C., Deng, C., Zhang, C., Ruan, C., Deepseek-
Ronneberger, O., Fischer, P., and Brox, T. U-net: Con- volutional networks for biomedical image segmenta- tion. In Medical image computing and computer-assisted intervention-MICCAl 2015: 18th international confer- ence, Munich, Germany, October 5-9, 2015, proceedings, part III 18, pp. 234 241. Springer, 2015.
v3 technical report. arXiv preprint arXiv:2412.19437, 2024a.
Liu, Q.. Zeng, Hc, J.. Yu, Q., Shen, X., and Chen, L.- C. Alleviating distortion in image generation via multi- resolution diffusion models and time-dependent layer nor- malization. Advances in Neural information Processing Systems, 37:133879 2024b.
Salimans, T., Goodfellow, I., Zaremba, W., Cheung, V., Radford, A., and Chen, X. Improved techniques for training gans. Advances in neural information processing systems, 29, 2016.
Liu, X., Gong, C., and Liu, Q. Flow straight and fast: Learning to generate and transfer data with rectified flow. arXiv preprint ar arXhv:2209.03003 2022.
Shazeer, N. Glu variants improve transformer. arXiv preprint arXiv:2002.05200, 2020.
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. arXiv preprint arXiv:2401.087740, 2024.
Shazeer, N., Mirhoseini, A., Maziarz, K., Davis, A., Le, Q., Hinton, G., and Dean, The sparsely- mixture- of-experts layer. Outrageously large neural networks, 2017.
Mei, K. and Patel, V. Vidm: Video implicit diffusion mod- els. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 37, pp. 9117-9125, 2023.
Song, J., Meng, C., and Ermon, S. Denoising diffusion implicit models. arXiv preprint arXiv:2010.02502, 2020.
Su, J., Ahmed, M., Lu, Y., Pan, S., Bo, W., and Liu, Y. Roformer: Enhanced transformer with rotary position embedding. Neurocomputing, 568:127063, 2024.
Nichol, A. Q. and Dhariwal, P. Improved denoising diffusion probabilistic models. In International Conference Machine Learning, pp. 8162-8171. PMLR, 2021.
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-choier routing, arXiv preprint arXiv:24 1.l:24100.000 10.00088,, 2024.
Park, B., Woo, S., Go, H., Kim, J.-Y., and Kim, C. De- noising task routing for diffusion models. arXiv preprint arXiv:2310.07. 2023.
10

---

Dift-MoE: Diffusion Transformer with Time-Aware and Space-Adaptive Experts
Tian, Y., Tu, Z.., Chen, H., Hu, J., Xu, C., and Wang. Y. U-dits: Downsample tokens in u-shaped diffusion trans- formers, arXiv preprint arXiv:2405.02730, 2024.
Touvron, H., Lavril, T., Izacard, G., Martinet, X., Lachaux, M.-A,, Lacroix, T., Rozière, B., Goyal, N., Hambro, E., Azhar, F, et al. Llama: open and efficient foundation language models. arxiv, arXiv preprint 2302:2230..13971, 2023.
Wang, Z., Lu, C., Wang, Y., Bao, F., Li, C., Su, and Zhu, Prolificdreamer: High-fidelity and diverse text-to-3d generation with variational score distillation. Advances in Neural Information Processing Systems, 36, 2024.
Xue, Z., Song, G., Guo, Q.. Liu, B., Zong, Z., Liu, Y. and Luo, Raphael: Text-to-image generation via large mix- ture of diffusion paths. Advances in Neural Information Processing Systems, 36, 2024.
Yang, A., Yang. B., Zhang, B., Hui, B., Zheng, B., Yu, B., Li, C., Liu, D., Huang, F., Wei, H., al. Qwen2.5 technical report. arXiv preprint 24
2024a.
Yang. Z., Teng, J., Zheng, W., Ding, M., Huang, S., Xu, J.. Yang. Y., Hong, W. Zhang, X., Feng, G., et al. Cogvideox: Text-to-video diffusion models with an expert transformer. arXiv preprint arXiv:2408.0600 2024b.
Zhou, Y., Lei, T., Liu, H., Du, N., Huang. Y., Zhao, V., Dai, A. M., Le. Q. Laudon, J., et al. Mixture-of- experts with expert choice routing. Advances in Neural Information Processing Systems, 35:7103-7114 2022.
Zoph, B., Bello, I., Kumar, S., Du, N., Huang. Y., Dean, J., Shazeer, N., and Fedus, W. Designing effective sparse expert models. arXiv preprint Mi0::2202.089006 2(3): 17, 2022.
Zuo, S., Liu, X., Jiao, J., Kim, J., Hassan, H., Zhang, R. Zhao, T., and Gao, Taming sparsely activated transformer with stochastic experts. arXiv preprint arX1v2211004226,, 2021.

---

Diff-MoE: Diffusion Transformer with Time-Aware and Spacc-Adaptive Experts
A. Diff-MoE Performance Summary
Table Performance Dift-MoE architecture with different model sizes.
ImageNet 256 256, 400K, cfg=1.5
#Params
FID
sFID
ISt
Precision
Recall
Diff-MoB-SS
36M / 107M
17.94
83.67
6.56
0.6472
0.5401
Diff-MoR-S/2+
77M / 21
9.27
126.22
5.72
0.7310
0.5285
Dift-MoE-B/2
137M / 402M
6.45
155.81
5.26
0.7629
0,5373
Diff-MoE-B/22
297M 821M
4.00
195.76
4.88
0.7990
0.5447
Diff-MoB-L/2
476M 1.4B
3.12
223.51
4.79
0,8162
0,5473
Diff-MoE-XI/2
1.5B 4.3B
2.69
262.22
4.54
0.8364
0.5579
B. Training
B.1. Rectified Flow
Rectified Flow (RF) (Liu et al., 2022) reformulates generation as deterministic ordinary differential equation (ODE).
learns a straight trajectory in probability space, directly mapping noise 2 ~ (O, 1) to data vo via a velocity field uz:
dæ; uo (a (), E [0, 1]:
dt
where x; (1 t)2 tæy enforces linear interpolation. The training loss minimizes trajectory curvature:
L 2,, dDg ((1 t)z + tæo, t) - (xo 2)1²
(10)
We present convergence analysis comparing the training dynamics of DDPM-based and Rectified Flow-based imple- mentations of our Diff-MoE- architecture, both trained for 400,000 iterations ImageNet 256x256 generation tasks under classifier-free guidance (CFG scale=1.5). As shown in Fig.6, the rectified flow-based model converges faster than the DDPM-based model. For context, Rectified Flow is used in SiT, SiT-Llama, and DiT-MoE, while DDPM is applied to other methods. To ensure a fair comparison, we report results without rectified flow in Tab.3 and Fig. 4.
Difi' McE 5 w/o RF
Traning Steps
Figure 6. Convergence comparision of DDPM and Rectified Flow-based DifT-MoE-S. We compare the DDPM-based Diff-MoE-S and Rectified Flow-based Diff-MoE-S trained for 400K iterations ImageNet 256 generation with efg 1.5.
B.2. Load Balance Loss
Building insights from prior work (Zuo et al., 2021) demonstrating that uniform expert utilization in MoE layers correlates with enhanced model performance, we incorporate a differentiable had-bblanning loss (Fedus et al., 2022) to enforce equitable workload distribution. For a MoB layer with N experts processing L input tokens, the load-balance loss can be calculated as:
N Σ
Lioad -N-N.
f P..
(11)
12

---

Dift-MoE: Diffusion Transformer with Time-Aware and Spacc-Adaptive Experts
where a controls the loss magnitude, and the expert-specific terms computed
fi arg max p(z) i}}
(12)
L P; (c)
(13)
Here P(x) denotes the routing probability distribution from Eq. 3, with fo representing the empirical fraction of tokens assigned to expert i, and P, the corresponding mean routing probability. This formulation penalizes discrepancies between actual token assignments and the router's probability mass allocation.
B.3. Final Loss
Our model diffusion model that utilizes either diffusion loss (Ho et al., 2020) or velocity field loss in Eq. as the generative training objective. The final loss function combines the generative loss with the load-balance loss, follows:
Ladipm + Lload; DDPM,
Dfimat
(14)
Crg + Lload, Rectified Flow.
C. Expert Selection Analysis
We analyze the routing dynamics of Diff-MoE-S/2 through comparative visualizations of its 12 MoE layers (Fig. 7), contrasting configurations with and without Expert-Speciff Time Conditioning, While this mechanism directly modulates expert parameters rather than router weights, its indirect influence on routing behavior emerges as a critical finding.
In MoE baseline (first variant of Tab. 6), time-dependent AdaLN layers in the main DiT path create entangled routing decisions where temporal and spatial factors compete for optimization priority. This coupling forces routers to jointly account for denoising-stee stage dynamics and token complexity, often converging to suboptimal equilibria characterized by imbalanced expert utilization.
The introduction of Expert-Specific Time Conditioning resolves this conflict by decoupling temporal adaptation from spatial specialization. Each expert independently processes timestep embeddings through dedicated conditioning pathways, enabling routers focus exclusively on spatial token characteristics. This architectural refinement yields two synergistic benefits: routers develop sharper discriminative capabilities for token classification while experts achieve more balancel workload distributions. Visual evidence in Fig. 7 confirms these improvements, showing reduced expert activation variance and clearer spatial feature boundaries compared to the unconditioned baseline.
13

---

Dift-MoE: Diffusion Transformer with Time-Aware and Spacc-Adaptive Experts
MoE Layer 0
MoE Layer
MoE Layer 2
MoE Layer
MoE Layer 4
MoE Layer
Inference step
100
200
MoE Layer
MoE Layer
Layer 8
MoE Layer
MoE Layer 10
Layer
Inference step
100
150
200
5
(a) w/o Expert Expert-Specifie Specitic Time Conditioning
MoE Layer 0
MoE Layer 1
MoE Layer 2
MoE Layer
MoE Layer 4
MoE Layer 5
Inference step
100
150
200
MoE Layer 6
MoE Layer 7
MoE Layer 8
MoE Layer
MoE Layer 10
MoE Layer 11
Inference step
100
200
0
5
(b) Expert- SSpeert-Specific Time Conditioning
Figure 7. Frequency for selected experts per timestep. We visualize the 12 MoE layers of the Dift-MoE-S/2 variant, both with and without Sxpert-Specific Time Conditioning. The x-axis represents the 8 experts in each layer, while the y-axis corresponds to the 250 DDPM steps used to sample the synthesized image. For each pair (expert e, inference step i), we show the average routing frequency across all timesteps assigned that specific expert e.
D. More Visualization Results
14

---

Dift-MoE: Diffusion Transformer with Time-Aware and Spacc-Adaptive Experts
Figure 8, More Samples generated by Diff-MoEE MoL-XL/2 trained on lamgeNet 256 at 400K iterations.
15


Reviews:
[{"review_id": "TN87uQ0NI0", "reviewer": "Reviewer_ap3M", "summary": "This paper proposes Diff-MoE, a method that captures both timestep and spatial contexts for expert routing. The approach consists of 1) Expert-Specific Timestep Conditioning – Unlike previous spatial MoE approaches, this enables each expert to adapt its operations based on the timestep, improving adaptability to different noise levels and 2) Feature Recalibration with Global Contexts – This enhances feature representations by incorporating global spatial information, leading to better expert specialization. These two techniques improve the model’s expert capabilities and global context awareness. Additionally, a parameter reduction technique using low-rank decomposition is employed to improve efficiency. Experimental results demonstrate that Diff-MoE outperforms both timestep-dependent routing methods and previous spatial MoE approaches.", "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": "Thanks for the thorough reviews. Below we try to solve  issues one-by-one.\n\nQ1:  Code release. \n\nWe sincerely appreciate the insightful feedback and recognition of this work. As astutely noted, training stability remains a critical challenge in MoE architectures, where conventional load-balancing losses offer partial mitigation. Our proposed expert-specific timestep conditioning mechanism further alleviates this issue by disentangling temporal adaptation from spatial routing, as analyzed in Supplementary Material Section C. To ensure reproducibility and foster community progress, we commit to open-sourcing all training frameworks, inference pipelines, and architectural implementations upon publication, to facilitate future research in scalable diffusion models.\n\nQ2: Including results on scaling over XL-size will be beneficial.\n\nWe fully agree that scaling up MoE to larger sizes and training more iterations can further validate the framework’s efficacy. However, current hardware limitations bound our exploration. Training configurations exceeding the XL scale (4.5B parameters) on our 8-GPU node infrastructure risk exceeding GPU memory capacity or incurring impractical training durations. We acknowledge this limitation in fully characterizing Diff-MoE’s scaling laws and commit to future large-scale evaluations upon securing expanded computational resources. Moreover, we will open-source the code to enable researchers with sufficient computing resources in the community to explore and build upon our work.\n\nQ3: Including the references about why the timstep awareness of networks is necessary will be beneficial.\n\nThanks for this valuable suggestion. We will discuss following papers in the revision:\n\n[1] Hatamizadeh, Ali, et al. \"Diffit: Diffusion vision transformers for image generation.\" ECCV, 2024.\n\n[2] Liu, Qihao, et al. \"Alleviating Distortion in Image Generation via Multi-Resolution Diffusion Models and Time-Dependent Layer Normalization.\" NeurIPS, 2024.\n\n[3] Karras, Tero, et al. \"Elucidating the design space of diffusion-based generative models.\" NeurIPS, 2022.\n\nQ4: Typos.\n\nThanks for pointing that out. We will fix all the typos in the revised version.", "claims_and_evidence": "I find most of their claims reasonable. Their primary claim is that their method combines previous temporal and spatial MoE approaches. By demonstrating that their approach outperforms both prior temporal and spatial MoE methods, they effectively validate this claim.", "methods_and_evaluation_criteria": "They follow the standard evaluation criteria used in the DiT-series evaluations.", "theoretical_claims": "N/A", "experimental_designs_or_analyses": "The ablations are well-conducted. However, one minor drawback is that the method does not achieve state-of-the-art (SOTA) performance. That said, I don’t find this to be a critical issue.", "relation_to_broader_scientific_literature": "### MoE in Diffusion Models \n\nTo enable efficient scaling of diffusion models, several works have explored MoE-based approaches. I agree with this paper's categorization, which classifies MoE methods into timestep-based and spatial MoEs. This work effectively combines both strategies to enhance the MoE architecture.\n\nAdditionally, training MoE models in diffusion frameworks is notoriously unstable, yet this paper seems a stable training process across various scale of base experts, which is a significant strength. If the code is made publicly available, it would greatly contribute to the field by providing insights into stabilizing MoE training in diffusion models.\n\n### Timestep-Aware Designs \n\nThis work aligns well with prior research advocating timestep-aware network operations in diffusion models. However, previous works have already demonstrated why timestep-aware design is necessary, and citing them would strengthen the paper’s argument. Including such references would reinforce the motivation behind their approach.\n\n### Diffusion Model Architectures \n\nSince DiT (Diffusion Transformers), most diffusion models have followed a transformer-based architecture, which is also prevalent in large-scale video diffusion models. Although this work does not achieve state-of-the-art (SOTA) performance, it explores an efficient scaling mechanism for MoE in DiTs.\n\nIt would be interesting to see how well this method scales to even larger models, though this is not a critical issue.\n\n### Overall Contribution \n\nIf the code is released, I believe this work would make a substantial contribution to MoE-based diffusion models, particularly by stabilizing training and providing an efficient scaling mechanism.", "essential_references_not_discussed": "Well discussed.", "other_strengths_and_weaknesses": "### Strength\n- This paper is well-written.\n- Proposed MoE architecture seems effective.\n\n### Weakness\n- Including the references about why the timstep awareness of networks is necessary will be beneficial.\n- Including results on scaling over XL-size will be beneficial.", "other_comments_or_suggestions": "Line 50: Please correct strange texts.", "questions_for_authors": "N/A", "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": "jdxShRbJ3x", "reviewer": "Reviewer_VKCP", "summary": "This paper introduces Diff-MoE, whihc is a novel framework combining Diffusion Transformers with Mixture-of-Experts to enhance scalability and flexibility in generative modeling. It achieves better FID scores across different model sizes compared to standard DiT models.", "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": "Thanks for the constructive comments and the recognition of novelty.\n\nQ1: Limited evaluation at extreme scales due to computational constraints.\n\nWe acknowledge the limitation in fully characterizing Diff-MoE's scaling laws and commit to conducting large-scale evaluations once additional computational resources are secured. Our contribution lies in synthesizing existing MoE methods in diffusion models, combining the benefits of timestep-based and spatial MoEs. This perspective offers a novel approach for scaling diffusion models, and extensive experiments validate its effectiveness. Furthermore, we will open-source the code to enable researchers with sufficient computational resources to explore and build upon our work.\n\nQ2: Comparison with more SOTA models.\n\nThank you for this constructive comment. State-of-the-art AR/diffusion methods can achieve an FID below 1.5 on ImageNet 256 generation. As discussed in Q1, current hardware limitations constrain our ability to explore larger scales or more iterations, which impacts our ability to achieve state-of-the-art performance. For reference, DiT-MoE-XL[1] achieves 1.72 FID after training for 7M iterations and 3.42 FID after 400K iterations. In comparison, Diff-MoE-XL achieves 2.69 FID after training for 400K iterations, which can serve as a benchmark. We will open-source all implementations to facilitate community-driven scaling efforts and will pursue large-scale training once expanded infrastructure becomes available. These steps aim to bridge the gap between methodological innovation and SOTA performance benchmarks in future work.\n\nQ3: Training and Inference computational costs.\n\nWe thank the reviewer for raising this critical aspect of MoE system design. Below we detail computational comparisons between Diff-MoE and DiT-MoE baseline under identical hardware (V100 GPU) and framework conditions.\n\nInferency Efficiency: Compared with expert-base method, Diff-MoE-S+ incurs only 6% FLOPs increase (16.05G → 17.01G) but 19% throughput reduction (278 → 225 samples/sec) compared to DiT-MoE-S, which may be due to the fact that the GPU optimizes different operators differently. When compared to dense DiT models, despite having similar parameter counts and theoretical FLOPs, our current implementation exhibits slower throughput due to the sequential computation of experts in a for-loop—an inherent challenge in MoE architectures. Nevertheless, Diff-MoE-S/2 achieves competitive FID (44.27 vs. DiT-B/2's 42.84) with 63% fewer activated parameters (36M vs. 131M), demonstrating superior memory efficiency critical for large-scale deployment.\n\nTraining Optimization: Building on fast-DiT[2] optimizations, we implement mixed precision training and pre-extracted VAE features. These adjustments yield 1.2 iterations/sec for Diff-MoE-S+ (211M params) vs. DiT-MoE-S’s 1.0 iter/sec (199M params), despite our model’s increased capacity.\n\nWhile our current implementation prioritizes architectural innovation over low-level optimizations, the sequential computation of experts in a for-loop exponentially exacerbates our speed disadvantage compared to sparse and dense models.  We recognize the need for low-level optimizations and identify clear pathways to mitigate throughput costs like parallel expert execution (DeepSeek-MoE[3] and DeepSpeedMoE[4]).  Following these advanced strategies, we believe there is significant room for improvement in inference speed. We will prioritize these optimizations after the code is open-sourced to continuously improve inference efficiency to bridge the efficiency gap while retaining the architectural advantages of Diff-MoE. Thanks again for this valuable question.\n\nQ4: Training Stability.\n\nDiff-MoE exhibits robust convergence behavior throughout the training process. As shown in Fig. 4, the FID consistently decreases as training advances. We also repeated the training for different model sizes multiple times. For Diff-MoE-S, the FID fluctuates by no more than 0.5 (44 ± 0.5). For Diff-MoE-XL, the FID fluctuates by no more than 0.05 (2.69 ± 0.05). Load balancing is an inherent problem in MoE training. While load balancing remains a persistent challenge in MoE architectures—partially addressed by conventional auxiliary losses—our expert-specific timestep conditioning mechanism introduces a critical refinement. By decoupling temporal adaptation (denoising stage dynamics) from spatial routing (token-level feature complexity), the framework redistributes computational loads more effectively, as discussed in Supplementary Material Section C.\n\n[1] Fei Z, Fan M, Yu C, et al. Scaling diffusion transformers to 16 billion parameters. arXiv:2407.11633, 2024.\n\n[2] https://github.com/chuanyangjin/fast-DiT\n\n[3] Liu, Aixin, et al. \"Deepseek-v3 technical report.\" arXiv:2412.19437 (2024).\n\n[4] Rajbhandari, Samyam, et al. \"Deepspeed-moe: Advancing mixture-of-experts inference and training to power next-generation ai scale.\" International conference on machine learning. PMLR, 2022", "claims_and_evidence": "Yes, most of the claims made have  convincing evidence.", "methods_and_evaluation_criteria": "Yes, proposed methods and evaluation criteria make sense.", "theoretical_claims": "Does not contain an explicit formal mathematical proof that needs to be verified for correctness.", "experimental_designs_or_analyses": "Limited evaluation at extreme scales (7M steps) due to computational constraints", "relation_to_broader_scientific_literature": "Diff-MoE integrates the Mixture-of-Experts paradigm, which has been effectively used to scale large models in Natural Language Processing, into diffusion models to achieve computational efficiency through dynamic parameter activation.\n\nDiff-MoE distinguishes itself from prior work on MoE in diffusion models by jointly optimizing for temporal adaptation and spatial specialization. Previous approaches often focused on either temporal partitioning of experts across denoising stages (e.g., DTR, Switch-DiT, MEME) or spatial routing of tokens to experts (e.g., DiT-MoE, EC-DiT). Diff-MoE proposes a more unified approach.", "essential_references_not_discussed": "This paper only compares DiT-based models and lacks a comparison with state-of-the-art (SOTA) image generation models.", "other_strengths_and_weaknesses": "Strengths:\nThe paper introduces several novel ideas, including joint spatiotemporal expert coordination, expert-specific timestep conditioning, and a globally-aware feature recalibration mechanism. The use of low-rank decomposition to reduce parameter overhead is also a notable contribution.\n\nWeaknesses:\nThe paper acknowledges that due to computational constraints, a full evaluation of the largest models (e.g., Diff-MoE-XL) at extended training durations was not performed.", "other_comments_or_suggestions": "Please refer to Strengths And Weaknesses.", "questions_for_authors": "While the paper extensively discusses parameter efficiency, it doesn't clearly address the training and inference computational costs (FLOPs/throughput) of Diff-MoE compared to baseline models. Could you provide quantitative comparisons of computational overhead and wall-clock time for training and inference? This would help evaluate the practical applicability of the approach beyond parameter efficiency.\n\nWere there any specific failure cases or training instabilities (e.g., dead experts, high variance in activation patterns) encountered during training?", "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": "cTQkw51SnU", "reviewer": "Reviewer_fsdw", "summary": "Diff-MoE introduces a novel integration of temporal and spatial adaptation in MoE for diffusion models. The module proposed in Diff-MoE takes into account the characteristics of the diffusion model. The experimental results are also impressive, proving the effectiveness and scalability of Diff-MoE.", "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": "Thanks for the thorough reviews. Below we try to solve  issues one-by-one.\n\nQ1: Typo on Tab.5.\n\nWe apologize for the typo in Table 5. The correct FID and IS scores for \"+GLU\" should be 38.20 and 41.43, respectively. The ability of GLU to enhance model capacity compared to a simple MLP has been discussed in works such as Llama [1] and StarNet [2]. We will fix this typo in the revised version.\n[1] Touvron, Hugo, et al. \"Llama: Open and efficient foundation language models.\" arXiv preprint arXiv:2302.13971 (2023).\n[2] Ma, Xu, et al. \"Rewrite the stars.\" Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. 2024.\n\nQ2:  Why the CFG value of Tab. 2 is 1.0, while others are 1.5?\n\nThe variation in classifier-free guidance (CFG) scales across Table 2 stems from our commitment to fair, methodologically consistent comparisons. For SiT-LLaMA, we directly adopt the reported CFG=1.0 results from its original paper, as no CFG=1.5 benchmarks were available. Conversely, other baselines (spatial/temporal MoE) are evaluated at CFG=1.5 following their established protocols.\n\nQ3: Code.\n\nWe sincerely appreciate the recognition of this work. To ensure reproducibility and foster community progress, we commit to open-sourcing all training frameworks, inference pipelines, and architectural implementations upon publication, to facilitate future research in scalable diffusion models.", "claims_and_evidence": "The quantitative and qualitative evaluations are overall adequate and support the claims of the paper. Strong baselines from prior work are considered for comparison. Authors trained models of different sizes, and the proposed method is clearly and consistently superior to the baselines across various sizes.", "methods_and_evaluation_criteria": "The paper is technically sound. All the proposed components are well motivated and serve a clear purpose: MoE for spatial dynamic computation, expert-specific timestep conditioning for temporal adaptation. \nThe evaluation is in line with what is used in previous work: standard ImageNet benchmark and FID metrics.", "theoretical_claims": "This paper verifies the effectiveness of architecture design through qualitative and quantitative experiments.", "experimental_designs_or_analyses": "The training strategy is aligned with existing methods. Moreover, several models of different sizes are designed to compare with existing methods in the case of the number of alignment parameters.\nThe comparison baselines are comprehensive, including dense and sparse (temporal or spatial) DiTs. The proposed method consistently outperforms different kinds of comparison methods across different scales.", "relation_to_broader_scientific_literature": "The current MoE-based DiT methods are still in its early days. Most methods only consider space or time, this work is the first to consider both time and space.", "essential_references_not_discussed": "To the best of my knowledge the paper offers a good coverage of related works. The papers listed are all relevant, are properly organized and accurately discussed.", "other_strengths_and_weaknesses": "Strengths:\n1.\tThe experimental setup is solid and the results are impressive. The performance of Diff-MoE is far superior to the existing methods, and the convergence speed is fast.\n2.\tThe MoE is a good way to scale up diffusion models to larger sizes, but this direction is still in the early stages of development. The motivation of this paper is clear, considering both spatial and temporal adaptability. The proposed module are specifically designed for diffusion models, rather than naively introducing MoE design from LLM.\n3.\tThis writing is good and well organized. The design motivation for each module is clear and easy to follow.\nWeaknesses: See questions.\nI recommend the paper for accept owing to the merits of the work and clear motivation behind each of the modules used. My final rating depends on how the authors address the concerns in the Question section.", "other_comments_or_suggestions": "See questions.", "questions_for_authors": "1.\tThere is a typo of Tab. 5. The FID of of +GLU should be 38.20, which is described in L433.\n2.\tWhy the CFG value of Tab. 2 is 1.0, while others are 1.5?\n3.\tDo the authors consider open source code? Open source makes a lot of sense to the community because the training process of MoE is usually accompanied by some tricks.", "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": "ldiztPjcha", "reviewer": "Reviewer_awrJ", "summary": "The paper introduces Diff-MoE, a framework integrating DiT with MoE to enhance scalability and performance in generative modeling. The proposed modules in Diff-MoE are specially designed for diffusion model, including the spatial-temporal adaptive experts and global feature recalibration. Extensive experiments show Diff-MoE significantly outperforms existing dense DiTs and prior MoE-based methods.\n\n## update after rebuttal\nMy concerns have been well addressed and would like to recommend accept.", "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": "Q1: Typo on Tab.5.\n\nWe apologize for the typo in Table 5. The correct FID and IS scores for \"+GLU\" should be 38.20 and 41.43, respectively. The ability of GLU to enhance model capacity compared to a simple MLP has been discussed in works such as Llama [1] and StarNet [2]. We will fix this typo in the revised version.\n\nQ2: Motivation of Depthwise Convolution.\n\nWe apologize for the unclear discussion of the motivation behind this design choice. The integration of depthwise convolution prior to the MoE module draws inspiration from hybrid vision architectures like CMT [3] and LocalViT [4], which strategically combine the spatial locality of CNNs with the global receptive fields of transformers. Our design addresses two critical requirements:\n\n● Spatial Locality Preservation: Depthwise convolution injects inductive biases for local feature extraction (edges, textures).\n\n● Parameter Efficiency: With computational complexity reduced from $k^2C^2$ (standard convolution) to $k^2C$, depthwise operations contribute only 0.1% additional parameters in Diff-MoE-S while maintaining spatial fidelity.\nIn our ablation studies (Table 5), we observed that removing depthwise convolutions led to a 6.5% increase in FID (35.85 → 38.20), highlighting their importance.\n\nQ3: Why fix the number of experts to 8?\n\nThe selection of 8 experts balances empirical performance gains against computational and architectural constraints, following DiT-MoE [5]. We report the results of the ablation experiments for the number of experts in the following table:\n| CFG=1.5, Imagenet256 | Params     | FID↓  | IS↑   |\n|----------------------|------------|-------|-------|\n| Diff-MoE-S-4E1A      | 36M / 66M  | 33.45 | 46.95 |\n| Diff-MoE-S-8E1A      | 36M / 107M | 31.18 | 50.55 |\n| Diff-MoE-S-16E1A     | 36M / 187M | 30.08 | 51.76 |\n\nWhile larger expert pools may benefit extreme-scale models, our focus on parameter-efficient diffusion training prioritizes balanced specialization over brute-force scaling. Open-sourced implementations will support flexible expert configurations for future hardware-augmented explorations.\n\n[1] Touvron, Hugo, et al. \"Llama: Open and efficient foundation language models.\" arXiv preprint arXiv:2302.13971 (2023).\n\n[2] Ma, Xu, et al. \"Rewrite the stars.\" Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. 2024.\n\n[3] Guo, Jianyuan, et al. \"Cmt: Convolutional neural networks meet vision transformers.\" Proceedings of the IEEE/CVF conference on computer vision and pattern recognition. 2022.\n\n[4] Li, Yawei, et al. \"Localvit: Analyzing locality in vision transformers.\" 2023 IEEE/RSJ International Conference on Intelligent Robots and Systems (IROS). IEEE, 2023.\n\n[5] Fei, Zhengcong, et al. \"Scaling diffusion transformers to 16 billion parameters.\" arXiv preprint arXiv:2407.11633 (2024).", "claims_and_evidence": "The claims are supported by extensive experiments. Diff-MoE consistently outperforms dense DiTs and MoE variants (Tables 2-4) with similar number of parameters.", "methods_and_evaluation_criteria": "Combining MoE with DiT for spatiotemporal adaptation is novel and well-motivated. The global recalibration mechanism addresses MoE’s local bias.\nClass-conditional image generation on ImageNet dataset is a widely used benchmark. The architecture design of generative model is usually validated on this dataset, and the metrics reported by the manuscript are persuasive.", "theoretical_claims": "The paper focuses on empirical contributions.", "experimental_designs_or_analyses": "Comprehensive ablations isolate contributions of each component, including basic architecture design, spatial-temporal adaptive experts and global feature recalibration.\nThe comparison with existing methods is also fair, as the authors have targeted the design of models of different sizes.", "relation_to_broader_scientific_literature": "This paper builds on Diffusion Transformers (Peebles & Xie, 2023) and MoE architectures (Shazeer et al., 2017). It advances prior diffusion-based MoE works (e.g., DiT-MoE, DTR) by unifying temporal and spatial adaptation.", "essential_references_not_discussed": "The related work discussed in this paper is more comprehensive.", "other_strengths_and_weaknesses": "Strengths:\n1.\tI completely agree that the current diffusion-based methods do not consider temporal and spatial flexibility simultaneously. The integration of temporal and spatial MoE for diffusion model is well-motivated. And the integration method is somewhat clever compared to the conventional MoE method.\n2.\tThe introduction of low-rank decomposition to reduce the number of parameters is natural under the design of this paper. The combined use of LoRA and AdaLN has no significant performance loss in this method.\n3.\tThe experiments in this paper are very comprehensive, including comparisons with dense and expert-based diffusion models. Compared with the existing methods across different scales, the performance of the proposed method is greatly improved.\nWeaknesses:\n1.\tThe FID value of +GLU in Tab.5 is not consistent with the text in Sec. 5.3, which could be a typo.\n2.\tThe design motivation of depthwise convolution in Basic Architecture Design section is not clear. Why add convolution to the pure transformer architecture? Based on Sec 5.3, this improvement resulted in a significant decrease in FID for reasons that need to be explained in further detail.", "other_comments_or_suggestions": "I'm positive about the paper. There are several contributions, which are well motivated, executed and ablated. The evaluation - especially the quantitative results - is convincing.", "questions_for_authors": "Why fix the number of experts to 8? Was this choice empirically validated against other configurations?", "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 paper introduces Diff-MoE, a novel framework combining Diffusion Transformers (DiT) with Mixture-of-Experts (MoE) to improve scalability and efficiency in generative modeling. The key contributions include expert-specific timestep conditioning, which enables dynamic adaptation to different diffusion stages, and a globally-aware feature recalibration mechanism that enhances the model's representational capacity. Extensive experiments demonstrate that Diff-MoE outperforms existing methods.

The strengths of this work lie in its innovative combination of spatial and temporal MoEs, which enhances both model adaptability and efficiency. The paper provides a solid experimental foundation with comprehensive ablation studies. The method is well-motivated, addressing key challenges in scaling generative models.
Decision: Accept (poster)