EC-DIT: Scaling Diffusion Transformers with Adaptive Expert-Choice Routing
Authors: Haotian Sun
Tao Lei
Bowen Zhang
Yanghao Li
Haoshuo Huang
Ruoming Pang
Bo Dai
Nan Du
Venue: ICLR 2025 Poster (2025)
Keywords: Diffusion transformer
text-to-image synthesis
Mixture-of-Experts
Primary area: foundation or frontier models, including LLMs

Abstract: Diffusion transformers have been widely adopted for text-to-image synthesis. While scaling these models up to billions of parameters shows promise, the effectiveness of scaling beyond current sizes remains underexplored and challenging. By explicitly exploiting the computational heterogeneity of image generations, we develop a new family of Mixture-of-Experts (MoE) models (EC-DIT) for diffusion transformers with expert-choice routing. EC-DIT learns to adaptively optimize the compute allocated to understand the input texts and generate the respective image patches, enabling heterogeneous computation aligned with varying text-image complexities. This heterogeneity provides an efficient way of scaling EC-DIT up to 97 billion parameters and achieving significant improvements in training convergence, text-to-image alignment, and overall generation quality over dense models and conventional MoE models. Through extensive ablations, we show that EC-DIT demonstrates superior scalability and adaptive compute allocation by recognizing varying textual importance through end-to-end training. Notably, in text-to-image alignment evaluation, our largest models achieve a state-of-the-art GenEval score of 71.68% and still maintain competitive inference speed with intuitive interpretability.

Published as a conference paper at ICLR 2025
EC-DIT: SCALING DIFFUSION TRANSFORMERS WITH ADAPTIVE EXPERT-CHOICE ROUTING
Haotian Sun 1:2+t Tao Lei', Bowen Zhang',
Yanghao Li',
Haoshuo Huang',
Ruoming Pang', Bo Dai?,
Nan Du
AAplle AI/ML Georgia Institute of Technology
ft _ei2, nowen_zhang,
yanghao_li, haoshuoh, dunan)eapple.com com,
haotian, sun@gatech. edu, bodai bodaiècc.on 8cc. grecc.atech.ede
ABSTRACT
Diffusion transformers have been widely adopted for text-to-image synthesis. While scaling these models up billions of parameters shows promise, the effec- tiveness of scaling beyond current sizes remains underexplored and challenging. By explicitly exploiting the computational heterogeneity of image generations. we develop a new family of diffusion transformers with expert-choice routing. EC-DIT learns to adaptively optimize the compute allocated to understand the input texts and generate the re- spective image patches, enabling heterogeneous computation aligned with varying text-image complexities. This heterogeneity provides an efficient way of scal- ing EC-DIT up to 97 billion parameters and achieving significant improvements in training convergence, text-to-image alignment, and overall generation quality over dense models and conventional MoE models. Through extensive ablations, show that EC-DIT demonstrates superior scalability and adaptive compute al- location by recognizing varying textual importance through end-to-end training. Notably, in text-to-image alignment evaluation, our largest models achieve state- of-the-art GenEval score of 1.66% and still maintain competitive inference speed with intuitive interpretability.
(MoE) models (EC-DiT) for
INTRODUCTION
An sirplane flying
Algnmeet Score (96
73.
BATCH MORM
NORM.
BATCH MORM
A sweatshirt with 'Batch Normi
Dense
32 Experts
64 Experts
Model Configurations
(14
[8] Figure 1: Performance of EC-DIT. (a) Across four model configurations, EC-DIT consistently
(b)
demonstrates superior performance in text-to-image alignment compared to the baseline models with similar activated parameters per prediction. (b) Scaling up with EC-DIT improves text-to-image alignment and visual detail rendering. The alignment score (a) is the average of GenEval (Ghosh et al., 2023) and DSG scores (Cho et al., 2024).
Diffusion models (Ho et al., 2020; Rombach et al., 2022; Podell et al., 2023) have demonstrated remarkable success in generation across various modalities. There has been growing interest in scaling these models to billions of parameters to improve performance and control computation
Work done while interning at Apple.
1 "Corrsppnndigg authors.

---

Published as a conference paper at ICLR 2025
over the generation process. One of the primary focuses of scaling diffusion models is to enhance scalability for rext-conditione image generation. As a promising approach for this synthesis task, diffusion transformers (DiT) (Peebles & Xie, 2023) combine the strengths of diffusion models and transformers and exhibit great potential in scalability. Current DiT-based models, such as Stable Diffusion 3 (SD), reach up billion parameters (Esser et al., 2024) and achieve promising perfor- mance in text-to-image generation. However, the effectiveness and optimal approaches for scaling DiTs beyond this size remain underexplored, primarily due to the significant impacts on training and inference efficiency associated with larger models.
The sparse Mixture-of -f-eettrts (MoE) technique (Zoph al., 2022; Lepikhin et al., 2020; Du al., 2022: Zhou et al., 2022) has proven to be an efficient method for scaling diffusion models like DiT. Sparse MoEs effectively scale up model capacity while maintaining relatively low computa- tional overhead. In MoE architectures, the router selects subset of parameters (experts) to process each group of token embeddings and then combines the final outputs from each expert. Current works (Yatharth Gupta, 2024; Fei et al., 2024; Xue et al., 2023) on sparse DiT often adopt token- choice routing strategy, which initially proposed for scaling language models (Zoph et al., 2022; Lepikhin et al., 2020; Fedus al., 2021). Despite technical differences, these methods, in princi- ple, rely on routing each image token to the top-ranked experts. However, this routing strategy is suboptimal with diffusion models, which requires more adaptive routing approach. Specifically, unlike autoregressive generation, DiT processes the entire image sequence at a time. This suggests that the global information of generated samples accessible throughout the denoising procedure. Additionally, different image areas often contain varying levels of detail. Thus, an optimal router should adaptively allocate heterogeneous computation for different patterns.
Driven by these motivations, we propose scaling DiT with adaptive expert-choice routing (EC-DiT), which is of independent interest, and tailor for cet-to-immae generation. This approach leverages global information from each image to optimize computational resource allocation that is aligned with varying image complexities, thus naturally aligning with the diffusion procedure. We scale EC-DIT up to 97 billion parameters with 64 experts. Figure showcases the promising scalability of EC-DIT. Despite a less than 30% increase computational overhead, our approach significantly improves training convergence, text-to-image alignment, and overall generation quality compared to dense and sparse variants with conventional token-choice MoEs (Yatharth Gupta, 2024; Fei et al., 2024). Notably, in text-to-image alignment evaluation, EC-DIT achieves new GenEval score of 71.68% while maintaining competitive inference speed compared to other strong baselines. Our main contribution is summarized as follows:
Scaling DiT with adaptive computation. We introduce EC-DIT, sparsely scaled DiT incor- porating expert-choice routing for text-to-image synthesis. This novel approach leverages global context information each generation step achieve heterogeneous compute allocation adaptive to different patterns within the generated images.
Promising performance at scale. We examine EC-DIT's scalability and effectiveness by scaling up to 97 billion parameters. larger model scales, EC-DIT demonstrates faster loss convergence, improves image quality, and reaches state-of sate-of-the-ase thett-rrt text- text-image alignment. Moreover, scaling with more experts consistently improves performance across various dimensions without significantly increasing computational overhead.
Comprehensive experiments and analysis. We conduct extensive experiments validate EC- DIT's superior scalability compared to dense and token-choice variants. Visualizations further demonstrate the model effectively learns adaptive computation allocation strongly aligned with textual significance.
2 METHODOLOGY
Starting with preliminaries of rectified flow model and MoE, this section presents the main design and properties of our EC-DIT models.
2.1 PRELIMINARIES
Image generation via rectified flow. DiT processes images as sequences of patches in latent space. An input image is first equally divided into grid of patches, then flattened and projected
2

---

Published as a conference paper at ICLR 2025
Rour
N' Blocks
More compute
The intemalonnal
FFN
FFN n
Adaptively alfocaee more compute desalled area
Space Starion flyng fron of the moon
Lece compute
EC-DIT Router W;
Overal Allccatoon
Encoder
Multi Head Cross Attention
EDNTER
Each expert so lelest tap takens (cenctes as across the entie socuence. (C- 6 this case.)
Intermediatl
Self-Atenion
Figure 2: EC-DIT architecture. The router leverages information from the entire sequence to adaptively route the most suitable tokens to each expert. Through this heterogeneous routing, more computation allocated to detailed image areas, such as the space station and moon, while less computation is used for rendering the background.
into a sequence of embeddings. We denote E R as the image sequence of length S in latent space, with dy representing the model's hidden dimension. Diffusion models progressively perform time-dependent denoising procedures on the image sequence Given the observation from data distribution and noise X1 ~ = W (0, I). rectified flow (Liu et al., 2022; Lipman al., 2022: Esser al., 2024) transfers between two distributions Po and Pı with straight-line path X, [x1 + (1 - ())x- This construction induces a probability flow ordinary differential equation (ODE) dx,
po(x)
(x, t)dt, where (x,, 1) represents the velocity direction of the path, parame- terized by the DiT backbone. Since uo (x(,,, t) dx.
Xg. optimize the following L2
objective (Liu et al., 2022; Esser et al., 2024):
(1) where T--) refers to the timestep density during training, for which we select the logit-normal sam- pling strategy (Atchison Shen, 1980): = log - with N(4,, ot). The hyperparameters ft and ar control the frequency for intermediate timesteps being sampled during the training process.
C (0) = E -N1001) [[](x] - xb) us (tx, + (1 ()=t)xx. t)|?
Mixture of Experts. Shazeer et al. (2017) first introduced the MoE layer that consists of a set {Ex())) E experts (each which is often FeedForward network (FFN)) and a learnable router with the weight W.. For given token representation x. the router selects the highest top-k experts based on the gating value of x W,. The output of the MoE layer will then be the weighted combination of the selected expert's computation where the expert weights are the normalized gating values via the softmax distribution. Since the token chooses its best set of experts, this top-k routing also termed the token-choice routing.
2.2 TEXT CONDITIONING WITH
MODULES
In the text-to-image synthesis, the model generates image samples x conditioned on the text prompt y. We utilize text encoder Enc (y) E Lxd, to extract the text embeddings of where L is the number text tokens, and d, is the text's hidden dimension. Following the approach Chen et al. (2023a), we incorporate multi-head cross-attention layer (Vaswani et al., 2023) between the self-attention and FFN layers in each DiT block. The cross artss-attention C₁ for head h is written as:
Q2 = = Wqqeery
K Enc (y) whey Q.-KT Vdd
V Enc (y) value
(2)
C softmax
where Qn is the query matrix for x, and K,, V, are the key and value matrices for respectively;
quueery E qddxxx= and Whey Waalue E Rd xd, are the projection matrices for the image and text representations. Then, the outputs from all H heads W to obtain the final output, f.e., MHCA(x) = concat ((C,att((C........ Cnl) Wrott This modification effectively injects textual information into the image generation.
concatenated and linearly transformed by
3

---

Published as a conference paper at ICLR 2025
2.3 SCALING DIT WITH EXPERT-CHOICE ROUTING
Because the majority of computation takes place in the dense FeedForward (FFN) layers (Peebles & Xie, 2023) of DiT, replacing the FFN layer with an MoE layer is one efficient method of scaling. Introducing the MoE layer decouples the effective computation from the model capacity. That is, we can effectively scale up the model capacity without significantly sacrificing the inference speed since only a selected subset experts will be activated conditioned on the input representation. While this approach seems plausible and effective, it can be further improved for the following reasons:
Uniform computation. The computation assigned to each input token is the same regardless of its representation even though the selected set of experts can be different. However, different image patches naturally demand varying amount of compute, e.g., foreground needs more attention while background and low-frequency areas may be disregarded.
Local context. Token-choice routing cannot access to the global information of the sequence due to the constraint of auto-regressive decoding. However, this global information can be better exploited in the diffusion process.
Load balancing. To improve the load balance of each expert, some sorts of auxiliary load balance loss (Shazeer et al., 2017; Lepikhin et al., 2020; Du et al., 2022; Zoph et al., 2022) are often required for token-choice routing to work properly.
Therefore, we propose to scale up the diffusion models by better exploiting the structure of the diffusion process, Specifically, following Pecbles & Xie (2023); Chen et al. (2023a), we embed the timestep using two-layer MLP with dimensionality equal to the hidden dimension d2, and feed this embedding into each AdaLN layer. Additionally, as shown in Figure 2, let x, be the output of the self-attention module, and x' be the input the router, i.e.,
x' x5 + MHC)(x,) Sxd,
Therefore, the input sequence x/ contains timestep information and integrates different modalities through equation 3.
Each expert is a two-layer FeedForward component, where the i-th expert represented as (x) GeLU(x Wi) Wi. Here, Wi E RRd, xd, and W E Rdd, xdz are the weight matrices for the i-th expert. For each MoE layer, the router is parameterized by the expert embedding W, E Rd, x where E is the total number of experts of the layer. Given the input x', the router first produces a token-expert affinity score tensor A SxS via so softmax along the expert dimension, i.e.
exp (xx W Wr)...)
A.,i
(5)
Σ exp (x'' W.)..))
This affinity score tensor assesses the relevance between each pair of expert and input token. Unlike the token-choice routing where each token x' selects the top-k experts from As,s, EC-DIT works from the expert-choice (Zhou al., 2022) view where each expert independently selects the top-C tokens in the descending order from A,, and C = S x fe/ E represents the average capacity of each expert, where fe denotes the capacity factor and reflects the average number of experts assigned to process each token.
Comparing to the token-choice routing, which assigns each token independently, BC-DIT selects the most relevant C' tokens from the entire sequence. To achieve this, we compute the gating tensor G S S follows:
G.,- = AA As, top-k ((A.,) < $)
0, otherwise,
where G.,, is the weighting score representing expert i's preference over the s-th token.
With equation 3 and equation 6, EC-DIT directs the computational focus of the experts towards tokens with significant textual information and visual patterns, while also being adaptive to varying denoising timesteps. Specifically, such allocation is not one-to-one. For instance, Figure 2, tokens corresponding to the moon and space station contain substantial detail and therefore assigned to multiple experts for finer processing. In contrast, tokens representing plain background may not be assigned to any experts at all. This makes the computation more adaptive and efficient, concentrating resources where they are most needed.
shown in
4

---

Published as a conference paper at ICLR 2025
Algorithm Pseudocode of EC-DIT's Routing Layer
batch size, 8: sequence length, d: hidden dimension
# E: number of experts, C: expert cspacity
experts: List of length consaining expert FFNs
def pe_dit_routing IX_P, K_r, experts]:
1, compute token-expert affinicy scores logita = eineum ('bad, de->bse", X_3, M affinity soltmax
shape (B,
din- 1)
shape: (B,
affinity einsum ('bse->bes', affinity)
shape: (B,
2. select the top-k tokens for each expert gating, index top_k fat_kiattnntyy dispatch hot (index, num_classes=@1
dim=-1)
shape: (R,
shape: (B, E, C, 8)
3. Process the tokena by each expert combing dispatch, X_P)
X_in einaun ( beca, bed-bbbed
shape: (B,
x_= [ep (xini:; +l) for in range (E)
stack dimell
shape: (8,
__out (insum/"beess bec, baod->had', dispetch, gating, x_0) K_out
shape: (B,
We define set of indexing vector {I expert, iee,
to filter the input tokens allocated to each
I {s
G..²
0, s < S}
The output Xout E Sxds of the sparse layer is then obtained by combining the results from each expert using the gating tensor G, i.e.
E Xout Σ (GI,a) & (xI..).
(8)
Note that equation inherently ensures full utilization of each expert and perfect load balance by design. Therefore, EC-DIT is trained directly via optimizing equation 1, without the need for an additional balan-llnning loss (Zoph et al., 2022). The pseudocode of EC-DIT's routing strategy presented in Algorithm 1.
The proposed EC-DIT shares several key benefits, making particularly well-suited for DiT-based diffusion models in text-to-image generation tasks.
Remark:
Adaptive computation. EC-DrT is implicitly aware of the textual context and denoising stages. As will demonstrate in Section 3, EC-DIT effectively enables context-aware routing and dy- namically adjusts its computation across both model depth and denoising timesteps.
Global context information. Diffusion models, unlike autoregressive models, process the tire sequence at each denoising step. As described equation 6, EC-DIT leverages this global context, combining textual and image information optimize token-to-expert allocation.
Load balancing. While token-choice routing often requires auxiliary load-balancing loss, EC- DIT 's routing mechanism inherently ensures balanced utilization of all experts. This prevents bottlenecks and enhances overall efficiency by allocating computations more effectively.
3 EVALUATIONS
In this section, empirically evaluate EC-DIT's performance across various metrics for text- image alignment and image quality and validate the training and inference efficiency. We then visualize the heterogeneous compute allocation within the proposed EC-DIT. Finally, present comparative results against dense and token-choice baselines. Throughout the paper, we use the naming convention DENSE-<Configg for dense models and EC-DrT-<con's ig>><>>E, where <Config> is the model configuration name; <n> is the number of experts.
3.1 EXPERIMENT SETUP
Following Chen et al. (2023a), we adopt a modified DiT architecture with additional cross-attention modules for text-to-image generation. For all models evaluated, we use 670M clip-baac based text encoder with the T5 tokenizer (Raffel et al., 2023) for the text input and a 34M variational autoencoder (VAE) with 8 channels et al., 2022). Both encoders remain frozen during DiT training. The transformer component is configured with four model sizes: XI, XXL, 3XL, and M, as detailed in Table 1. Starting from
Model architecture.
for the image input (Hessel
5

---

Published as a conference paper at ICLR 2025
Table 1: Model size and specifications across four backbone configurations. For Total Params., nE denotes EC-DIT with n experts in a single MoE layer. Activated Params. refer to the average number activated parameters per token in the forward pass. The detailed calculation depicted in Appendix D. In Model Arch., #Head is the number of query heads, and #KV represents the number of key/value heads for grouped-query attention (GQA). 'Note that a trade-off was made in the model depth for the M configuration, where DENSE-M has 46 layers, while EC-DIT-M has 38 layers in order to be able to fit into the HBM of training accelerators.
Total Params.
Activated Params.
Model Arch.
Config.
DENSE
8F
16F
32F
64F
EC-DIT
#layers
Hidden dim.
#Head
#KV
XI.
1.47B
2.51B
3.70B
6.088
1.62B
28
1,152
XXL
2.35B
4.87B
7.73B
13.47B
2.71B
38
1,536
24
6
3XL
4,50B
10.74B
17.87B
32.15B
5.18B
42
2,304
36
6
M
8.03B
97.21B
8.27B
38/46"
3,072
48
12
the second layer, scale each dense model by interleavingly replacing all even-numbered dense layers with sparse MoE layers, where each expert is an FFN with dimensions identical to the dense counterpart. We set the capacity factor fc = 2.0 throughout the training and inference stages for all sparse models. We scale the XL-3XL dense models with 8, 16, and 32 experts. To further test the model's scalability, we scale the M configuration with 64 experts. To explore the trade-off between model depth and width during scaling, we reduce the number of layers in EC-DrT-M-64B to 38, while maintaining 46 layers in the DENSE. With these settings, our EC-DIT-3XL-32E reaches 32.15 billion parameters, and the largest EC-DIT-M-64E attains 97.21 billion parameters beyond all the current sparse diffusion models (Fei al., 2024; Lin et al., 2024; Balaji et al., 2022).
Training details. We collect and utilize approximately 1.2 billion text-image pairs from the In- ternet (McKinzie et al., 2024; Lai et al., 2023). The model resolution is set to 256 256, with a patch size of 2 x 2. To accelerate training, we employ the masking technique proposed in (Zheng et al., 2023) with a masking ratio of 0.5. This results in the input sequence length of 128 per image. Model training is conducted on v4 and v5p TPUs with batch size 4096. We use the RMSProp with momentum optimizer (Hinton, 2012) with a learning rate 1e-4 and 20K warmup steps. All models trained with Distributed Data Parallelism (DDP) or Fully Sharded Data Parallel (FSDP) for 800K steps. The expert dimension of EC-DIT is fully shared across the TPU mesh. To enhance training efficiency, we regroup the first two dimensions of the input sequence (batch_s ze, sequence_length
into (outer batch, rum-group, group-size). We maintain a group_s ze of 1024 for training and 512 for inference, which allows EC-DIT process two images in sequence at inference time.
Evaluation metries. To assess the model's text-to-image capabilities, we primarily use GenEval (Ghosh et al., 2023) and Davidsonian Scene Graph (DSG) scores (Cho et al., 2024). GenEval segments and detects objects in the generated samples evaluate their alignment with the conditioned texts. For DSG, we use an internal version of Gemini to perform Visual Question Answering (VQA) to assess the text-image alignment. To evaluate the image quality of the gen- erated images, we measure zero-shot Fréchet Inception Distance (FID) (Heusel et al., 2017) along with CLIP Score (Hessel al., 2022) on the MS-COCO 256 256 dataset using 30K samples (Lin et al., 2015). We also provide generated samples from subset of Partiprompts (Yu et al., 2022) in Appendix E.
3.2 TEXT-TO-IMAGE ALIGNMENT
We present the comparative evaluation results on GenEval and DSG, both of which are commonly used to measure text-to-image alignment. Table 2 shows the comprehensive GenEval results of EC-DIT compared to current state-of-the-art models. Specifically, our largest sparse model, EC- DIT-M-64E (resolution 256 x 256), reaches GenEval of 71.68% and surpasses SD3 (resolution 512x in the overall score (68.00%) and most sub-tasks. Notably, the solely pretrained EC-DIT- M-64E without fine-tuning stage even outperforms the SD3 model fine-tuned with DPO. Notably, both DPO fine-tuning and training with higher resolution positively impact the GenEval score, as

---

Published as a conference paper at ICLR 2025
Table 2: GenEval comparison. Our largest scaled model (EC-DIT-M with 64 experts) surpasses all existing models in overall performance and most sub-tasks of the GenEval Metrics (Ghosh al., 2023), including the aPO-finetuned version the current state-of-the-arg SDr-Large (Esser et al., 2024) (depth=38, resolution 512 x 512). The scores for models preceding ours are sourced from Esser et al. (2024). The best and second-best entries are highlighted.
Model ()) / Score (-+)
Overall Single obj.
Two obj.
Counting Colors Position Color attr.
SD v1.5 (Romhach et al., 2022)
43.00
97.00
38.00
35.00
76.00
4.00
6.00
PixArt-a (Chen et al., 2023a)
48.00
98.00
50.00
44.00
80.00
8.00
7,00
SD v2.1 (Rombach et al., 2022)
50.00
98.00
51.00
44.00
85.00
7.00
17.00
DALL-E (Ramesh
2022)
$2.00
94.00
66.00
49.00
77.00
10.00
19.00
SDXL (Podell ct al., 2023)
55.00
98.00
74.00
39.00
85.00
15.00
23.00
SDXL Turbo (Podell et al., 2023)
55.00
100.00
72.00
49.00
80.00
10.00
18.00
TF-XI (Saharia et al., 2022)
61.00
97.00
74.00
66.00
81.00
13.00
35.00
DALL-E (Shi et al., 2020)
67.00
96.00
87.00
47.00
83.00
43,00
45.00
SD3-Large (Esser et al., 2024)
68.00
98,00
84,00
66.00
74.00
40.00
43.00
SD3-Large (Esser et al., 2024) w/ DPO
71.00
89,00
73.00
83.00
34.00
47.00
DENSE-3XL
68.92
99.69
86.00
69.41
81.67
20.58
$6.19
EC-DIT-3XL-32E
70.91
99.64
87.88
72.53
83.84
21.19
60.40
EC-DTT-M-64E
71.68
99.84
88.67
73.69
85.77
21.33
60.80
SD3 generally yields higher GenEval scores equipped with these techniques (Esser et al., 2024). This further underscores the potential benefits of scaling dense model with EC-DIT. Furthermore, EC-DIT-3XL-322 with 5.18B activated parameters, achieves a nearly equivalent GenEval score of 70.91% % compared to the 71 1.00% score from SD3-Large with DPO, with only 64% activation size of the latter. The EC-DTT-3XL-32F3 inference time only 33.8% of the 8B DENSE-M. These results demonstrate that our scaling approach can achieve performance comparable to significantly larger dense models, even without additional fine-tuning.
We further validate the effectiveness of EC-DrT through DSG scores (Cho et al., 2024), as depicted in Figure 3. EC-DIT consistently improves DSG performance over the dense variants across different model configurations. Ad- ditionally, for each model size, increasing the number of experts further enhances text-image alignment. In both GenEval and DSG evaluations, increasing the model's ca- pacity with the EC-DIT directly enhances its text-following ability and leads to more accurately aligned image genera- tion. A detailed DSG score breakdown is provided in Ap- pendix
Podal Typa
Benso
322
15.4
15.
DSG Score
76.6
JXLL
Figure 3: DSG comparison. Scal- ing up with EC-DrT elevates perfor- mance.
3.3 INFERENCE- EFFICIENCY
Figure illustrates the inference efficiency of EC-DIT compared to dense models. Across all base model configurations,FC-DD consistently delivers improved performance in text-image alignment. The actual inference overhead introduced by EC-DIT ranges from 20% to 28% across different configurations. This practical overhead is greater than the theoretical one of up to 15%, which is calculated in Table 1. For EC-DIT-M, although the theoretical overhead around 3%, the ac- tual overhead is measured 23%. This difference might be attributed the varying efficiency in inference- sierence-time parallelism: EC-DIT-M uses model parallelism to fit on 8xH100 GPUs, whereas the dense model utilizes FSDP. Despite these factors, scaling with EC-DIT reliably enhances text- to-image alignment up to the state-of-the-art level with less than 30% additional overhead.
3.4 HETEROGENEOLS COMPUTE ALLOCATION
In Figure 6, we visualize the heterogeneity of compute allocation through heatmaps that show the number of experts assigned each image token. We observe that EC-DIT allocates more computa- tion to areas with clear textual significance, such as major objects or detailed patterns. For example, in Figure 6(a), the main object (such the moon) and the rendered text receive the most compu-

---

Published as a conference paper at ICLR 2025
72
32E-XXL
32E-3XL
64E-M
20
16E-3XL
00:00
05-145 161162
GenEval (%6
70
32E-XL
16E-XXL
Dense-M
23c
E0
522
8E-3XL
Training Ix
Dense-3XL
0.286
16E-XL
BE-XXL
0.2284
8E-XL
Dense-XXL
Activated Pararrs
18
0.2282
Denus 60-126
Dense-XL
05-322
200
400
50
100
200
400
Training Steps Ix
Inference Time 500 samples)
Figure 4: Inference-time efficiency. The circle size is pro- portional the total activated parameters. Inference time rep- resents the time elapsed to generate 500 samples on 8xH100 GPUs. EC-DIT shows superior performance compared dense models, with less than 30% additional overhead.
Figure 5:
Comparison with
token-choice baselines on FID and CLIP Score of EC-DIT-XXL (EC) and GShard (GS) with top-2
token-choice routing.
Can'
Esparta kosaled to each takon
samples
Allocation per sparse layer at the denoising step
40
Exports allocated each
No.
10
12
18
20
22
24
28
28
32
36
98
38
40
42
44
48
(b) Allocation per denosing timestep t
Figure 6: Visualization compute allocation for EC-DIT-XXL-33 XXL-32E (38 layers, with even layers being sparse). Each point in the heatmap represents the number of experts selecting the correspond- ing token in the latent space. Darker colors indicate areas where more compute is allocated. The heatmaps reveal that the router tends to assign more computations to significant objects, such as the moon in the first sample, and to detailed areas, such as text in the second sample.
lation. In contrast, the background (composed of nearly monotone colors) receives much lighter computation. This allocation heterogeneity tends be more pronounced in later layers. Figure 6(b) also demonstrates that later layers exhibit heterogeneity in fewer denoising steps. We hypothesize that the early layers handle most "lowwffeeuuecc"" areas, such as the background tokens, and provide context for the later stages. Meanwhile, the routers from later layers more quickly converge to the most textually representative tokens. This coincides with previous observations pixel space (Lei al., 2023). Overall, EC-DIT 's compute allocation exhibits high heterogeneity, where a single im- age token can be processed by up to 44% of the total computation to capture intricate details, while only few experts process non-detailed areas even skipped in certain layers. Notably, EC-DIT learns to achieve this adaptive and heterogeneous compute allocation through end-to-end training- We hypothesize that the input the router, derived from the cross-attention module, likely contains essential cross-modal information from both text and image, which can be effectively leveraged to achieve this adaptive routing.
8

---

Published as a conference paper at ICLR 2025
Dense
Dense EC-DIT-BE DIT-16E EC-DIT-32E
Dense
EC-NT--EE
EC-DIT-BE
ECIDIT-16E
EC EC-DIT
EC-DIT-32E
EC-DIT-32E
22
0.279
0,281 0.283 CLIP Score
0.285
0.279
0.281
0.283 CLIP Score
0.285
0.279
0.281 0.283
0.285
CLIP Score
Figure 7: FID vs. CLIP Score curve during training. The closer a point is to the upper right corner (7²), ) the better the model performs in terms of higher image quality (lower FID) and text- image alignment (higher CLIP Score). For each model, eight points are plotted from left to right, representing the model's 's performance at 100k training step intervals.
3.5 COMPARISON WITH TOKEN-CHOICE ROUTING
Figure 5 compares EC-DIT with models scaled using token-choice MoE (Yatharth Gupta, 2024; Xue et al., 2023; Fei et al., 2024). For fair comparison, we maintained identical experimental settings, except for replacing the routing strategy with token-choice and incorporating an auxiliary load-balancing loss in the training objective. The token-choice bascline uses top-2 routing, which matches the activation size of EC-DrT with capacity factor fc 2.0. Our method consistently demonstrates superior training convergence and performance throughout the entire training period. Notably, EC-DIT with experts rivals the token-choice baseline with 16 experts in both generation quality and text-image alignment, while EC-DIT with more experts significantly outperforms this token-choice baseline. One reason for the slower training convergence in token-choice routing is the equal compute allocation to different image areas, which makes the compute less adaptive. Detailed image areas may not receive sufficient computation to render complex patterns, while simpler areas waste the computation they receive. Moreover, token-choice routing processes each token inde- pendently without considering sequence information, leading to a lack of perception of the overall relationships between local tokens. Additionally, the inherent load imbalance token-choice rout- ing results in some experts being overloaded with an excessive number of tokens and creating a training bottleneck (Zhou al., 2022). Appendix C provides comprehensive comparison of EC- DIT and the token-choice baseline regarding GenEval performance. As shown in Table 5, EC-DIT outperforms the token- choice model across all tested configurations, which further validates the advantages of EC-DIT in text-image alignment.
3.6 SCALING EC- MODELS WITH MORE EXPERTS
From previous results, scaling a dense model with more experts consistently brings performance gains both convergence and generation quality. In Figure increasing the number of experts leads to slightly improved loss convergence. Figure 3 and Figure 4 collectively demonstrate that adding more experts to the same dense model consistently enhances the model's text-image align- ment capability. Additionally, Figure 4 indicates that increasing the number of experts effectively results in a more powerful model without increasing the activation size nor the inference time, as the expert capacity C is dynamically adjusted based on the number of experts. In Figure 7, we further examine the trend FID vs. CLIP Score during training. Despite fluctuations the curve, the general trend shows that adding more experts shifts the curve towards improved image gener- ation quality and text-image alignment. Therefore, we suggest increasing the number of experts the computational budget allows, this consistently yields performance gains without additional inference overhead.
3.7 TRAINING DYNAMICS.
We also analyze the training dynamics and convergence of dense models and EC-DIT with varying model configurations and numbers of experts, as shown in Figure 8. Across all model settings, we observe a significant improvement in loss reduction between the dense models and sparse EC-DiT throughout the training period. EC-DIT demonstrates more effective ability to learn denoising for rectified flow generation and leads to better loss convergence. For example, EC-DIT-XXI. models trained with 200K steps achieve training loss comparable to DENSE-XXI models that require

---

Published as a conference paper at ICLR 2025
Models
XXL Models
3XL Models
0.28
Dense ECIDIT-BE
Dense
Dence
ECIDIT-BE
EC-DIT-BE
EC-DIT-1GE
EC-DIT-16E
EC-DIT-1GE
EC-DIT-32E
EC-DIT-32E
EC-DIT-32E
Lo5
0.26
EC-DIT-64F
200 400
600
800
0
200
400
600
800
0
200
400
600
800
Training Steps (x10P)
Training Steps (x10F)
Training Steps (x1DF))
Figure 8: Training loss of the EC-DIT. The results show noticeable loss gap between the dense model and the sparse EC-DIT. The training loss and convergence speed is further improved by increasing the number of experts.
around 400K steps. Additionally, the sparse EC-DIT-XXL models eventually stabilize at a lower loss than their dense counterparts.
4 RELATED WORK
Experts. Many MoE models utilize a token-choice routing strategy, where each token independently selects the top-k experts for processing. Sparsely Gated MoE (Shazeer et al., 2017) first introduced top-k gating in LSTM models. Such routing strategy was later extended to trans- formers for language modeling in GShard (Lepikhin et al., 2020), Switch Transformer (Fedus et al., 2021), GLaM (Du et al., 2022). DeepSpeed-MoF: (deepspeed), and ST-MoF (Zoph et al., 2022). To address potential load imbalances among experts, these token-choice approaches typically in- corporate an auxiliary load-balaneing loss (Shazeer al., 2017). Alternatively, Zhou al. (2022) proposed an expert-choice routing in text domain that selects the most suitable text tokens for each expert. Inspired by this, we adapt and scale up the DiT model using expert-choice routing tailored specifically for diffusion-basse generative tasks. With this approach, each image token from differ- ent image areas receives heterogeneous computation allocation. Additionally, the method inherently achieves perfect load balance, as each expert processes an equal number of image tokens, which optimizes both efficiency and performance for the diffusion process.
Scaling diffusion models with MoE. Recent advancements have leveraged MoEs to scale up dense diffusion models for image generation tasks. Several studie have focused on class-conditioned image generation. DTR (Park et al., 2023) and Switch-DiT (Park et al., 2024) approach the routing problem as multitask learning and utilize token-choice MoE for different denoising stages. Simi- larly, MEME (Lee et al., 2023) employs an ensemble of denoising experts specialized in distinct timestep ranges. Furthermore, some models have extended this approach for text-to-image gener- ation. ERNIE-ViLC 2.0 (Feng et al., 2022) and eDiff-I (Balaji et al., 2022) employ ensembles of specialized experts for different denoising stages. Using U-Net as their backbone, these methods scale up to 24 billion parameters. However, these models activate only one expert at each timestep, which reduces efficiency at large scales. Recently, some works have shifted towards sparse MoEs. SegMoE (Puigcerver et al., 2023) employs top-k routing to scale up SDXL (Podell et al., 2023) with up to four experts. RAPHAEL (Xue et al., 2023) combines the concepts of distinct timestep experts and token-choice MoBS. DiT-MoE (Fei et al., 2024) applies token-choice routing to produce a sparse DiT with up to 16 billion parameters. Nevertheless, these methods often suffer from load imbalance issues inherent to the token-choice scheme, requiring an additional auxiliary balancing loss to mitigate this problem. The most similar work to ours is MoMa (Lin et al., 2024), which also utilizes expert-choice routing. However, MoMa explores model scaling only on a relatively small scale, with its largest model reaching 7 billion parameters. In contrast, we leveraged the proposed adaptive routing to scale EC-DIT up to 97 billion parameters and demonstrated the DiT scaling potential on much larger scope.
5 CONCLUSION
We propose EC-DIT, a novel scaling approach for DiT using expert-choice routing. EC-DIT lever- ages global image information and allocates computation to each image patch adaptive to its pattern complexity. Our model effectively scales up to 97 billion parameters and significantly improves training efficiency, text-to-image alignment, and overall generation quality compared to dense mod- els and sparse DiT's with token-choice routing.
10

---

Published as a conference paper at ICLR 2025
ETHICS STATEMENT
In this paper, we propose new scaling strategy for text-to-image generation, improving both the quality of generated images and their alignment with given textual inputs. However, this general approach potentially carries ethical concern.. For example, the generative model could be misused generate harmful or malicious content, such as imagery based on violent or discriminatory text prompts. We acknowledge these concerns and are committed to addressing them by developing more advanced dataset filtering mechanisms and model calibration techniques. Our future work will prioritize increasing the safety and ethical use of text-to-image generation models.
REPRODUCIBILITY STATEMENT
In Section 2, we present method with detailed formulations, examples, and visual demonstra- tions to clarify the model structure and mechanism (such as Figure 2). In Section 3, we introduce the training dataset and describe the most effective model parameters and components. We also provide specific details regarding the hardware used for training and inference like Figure 4. Additionally, we report the training dynamics of the model, depicted in Figure 8. Through various tables and figures, we present the performance of the proposed method from multiple perspectives, such as Table 2, Figure 3, and Figure 7. We compared our approach with existing methods targeting similar domains overlapping methodologies, including both theoretical analysis and experimental eval- uations (e.g., Figure 5). Notably, all evaluation experiments were conducted on public datasets or benchmarks, such MSCOCO, GenEval, FID, and CLIP Score. To further ensure reproducibility, we plan to release the model weights contingent on the acceptance of this work.
ACKNOWLEDGMENTS
We would like to thank John Peebles, Xianzhi Du, Chen Chen, Wenze Hu, Riley Roberts, Wei Liu, Zhengfeng Lai, Vasileios Saveris, Peter Grasch, and Yinfei Yang for constructive feedback and discussions.
REFERENCES
J. Atchison and S.M. Shen. Logistic-normal distributions:Smne properties and uses. Biometrika, 67(2):261-272, 08 1980. ISSN 0006-3444, doi: 10.1093/biomeet667..222 URL https:
//doi. wrg.10.gg11000 10t.:/1099bbmeeeeee 2..
Yogesh Balaji, Seungjun Nah, Xun Huang, Arash Vahdat, Jiaming Song, Qinsheng Zhang, Karsten Kreis, Miika Aittala, Timo Aila, Samuli Laine, Bryan Catanzaro, Tero Karras, and Ming- Yu Liu. eDiff-I: Text-to-lmage Diffusion Models with an Ensemble of Expert Denoisers. arXiv, 2022. doi: 10.48550/arxi.2221..oo
Black Forest Labs. Announcing black forest labs, 2024. URL htRphhttppss/blaccforeetttt a bamnouncingg-blaa--ffeeee
Junsong Chen, Jincheng Yu, Chongjian Ge, Lewei Yao, Enze Xie, Yue Wu, Zhongdao Wang. James Kwok, Ping Luo, Huchuan Lu, and Zhenguo Li. Pixart-a: Fast training of diffusion transformer for photorealistic text-to-image synthesis, 2023a. URL pettss:/arree www.xxii.....aaaoo.mom wbwg//bb////ee 2310 .00426.
2230
Weifeng Chen, Jie Wu, Pan Xie, Hefeng Wu, Jiashi Li, Xin Xia, Xuefeng Xiao, and Liang Lin. Control-a-vide:: Controllable text-to-video generation with diffusion models, 05 2023b.
Jaemin Cho, Yushi Hu, Roopal Garg, Peter Anderson, Ranjay Krishna, Jason Baldridge, Mohit Bansal, Jordi Pont-Tuset. and Su Wang. Davidsonian scene graph: Improving reliability in fine- grained evaluation for text-to-image generation, 2024. URL https: /ta::///xxiiie /aaxiaiiii..////aa wowggabs/// 10. 82335.
Nan Du, Yanping Huang, Andrew Dai. Simon Tong, Dmitry Lepikhin, Yuanzhong Xu, Maxim Krikun, Yanqi Zhou, Adams Wei Yu, Orhan Firat, Barret Zoph, Liam Fedus, Maarten Bosma,

---

Published as a conference paper at ICLR 2025
Zongwei Zhou, Tao Wang, Yu Emma Wang, Kellie Webster, Marie Pellat, Kevin Robinson, Kath- leen Meier- Hellstern, Toju Duke, Lucas Dixon, Kun Zhang, Quoc Le, Yonghui Wu, Zhifeng Chen, and Claire Cui. Glam: Efficient scaling of language models with mixture-offexperrs Pro- ceedings of the 39 th International Conference on Machine Learning, Baltimore, Maryland, USA, PMLR 162, 2022., 2022. UtLhttps //tp////avvvv//// www.wwwwwwwwww..a...ooo..ooooooe wogg//dee....oo. /tfff/2110 /2/22 06905.
Patrick Esser, Sumith Kulal, Andreas Blattmann, Rahim Entezari, Jonas Müller, Harry Saini, Yam Levi, Dominik Lorenz, Axel Sauer, Frederic Boesel, Dustin Podell, Tim Dockhorn, Zion English, Kyle Lacey, Alex Goodwin, Yannik Marek, and Robin Rombach. Scaling rectified flow trans- formers for high-resolution image synthesis, 2024. URL perttts://arxiiee org/ ha/g//bb//2450 .03206.
William Fedus, Barret Zoph, and Noam Shazeer. Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity. arXiv, 2021. doi: (0.48550/amxiv22001.00
Zhengcong Fei, Mingyuan Fan, Changqian Yu, Debang Li, and Junshi Huang. Scaling Diffusion Transformers to 16 Billion Parameters. arXiv, 2024.
Zhida Feng, Zhenyu Zhang, Xintong Yu, Yewei Fang, Lanxin Li, Xuyi Chen, Yuxiang Lu, Ji- axiang Liu, Weichong Yin, Shikun Feng, Yu Sun, Li Chen, Hao Tian, Hua Wu, and Haifeng Wang. ERNIE-ViLG 2.0: Improving Text-to-Image Diffusion Model with nhowledge-Enhancee Sewiroood-Deooiinn-eeeeeeee arXiv, 2022. doi: 10.48550/arxvv22210.0
Oran Gafni, Adam Polyak, Oron Ashual, Shelly Sheynin, Devi Parikh, and Yaniv Taigman. Make- a-scene: Scene-based text-to-image generation with human priors, 2022. URL httph:t/p////at w.org//bs//2 /2/aaabs//22... 2203.2223.2 13131.
Dhruba Ghosh, Hanna Hajishirzi, and Ludwig Schmidt. Geneval: An object- framework for evaluating text-to-image alignment, 2023. URL
2210.11513
Yuwei Guo, Ceyuan Yang, Anyi Rao, Maneesh Agrawala, Dahua Lin, and Bo Dai. Sparsectrl: Adding sparse controls to text-to-video diffusion models, 2023. URL https:// /ttp:///a/xxvo///// abs/2311 6933.
Jack Hessel, Ari Holtzman, Maxwell Forbes, Ronan Le Bras, and Yejin Choi. Clipscore: A reference-free evaluation metric for image captioning, 2022. URL https: /ttpsv/////vv..///aeeaaeea org/ab ab s/2104. 08718.
Martin Heusel, Hubert Ramsauer, Thomas Unterthiner, Bernhard Nessler, and Sepp Hochreiter. Gans trained by a two time-scale update rule converge local nash equilibrium. In I. Guyon, U. Von Luxburg, S. Bengio, H. Wallach, R. Fergus, S. Vishwanathan, and R. Garett (eds.), Advances in Neural Information Processing Systems, volume 30. Curran Associates, Inc., 2017. URL https: httt:////ococeeeeeeeeeeeeee neurips cc/paper. fiwefiiee//eeeee httpeeeeffelaeppeeeeeeeeeeeeeeee /ess/apper/0////leeee 211899701eb0eefe6877116679999999
Paper....
Geoffrey Hinton. Lecture 6.5 .ms-roppop:: Divide the gradient by a running average of its recent
magnitude. coursera: Neural networks for machine learning, 4, 26-31., 2012.
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models, 2020. URL
https: : / /aaxivv.... org/ abs/2005 11239.
Zhengfeng Lai, Haotian Zhang, Bowen Zhang, Wentao Wu, Haoping Bai, Aleksei Timofeev, Xi- anzhi Du, Zhe Gan, Jiulong Shan, Chen-Nee Chuah, Yinfei Yang, and Meng Cao. Veclip: Im- proving clip training via visual-enriched captions, 2023. URL https: http://axxvvvvvo org/ org/abs////// /bs/ 2310.0766 07699.
Yunsung Lee. Jin- Young Kim, Hyojun Go, Myeongho Jeong, Shinhyeok Oh, and Seungtack Choi. Multi-Architecture Multi-Expert Diffusion Models. arXiv, 2023. doi: 10.48550/arxiv22306000 0.
Tao Lei, Junwen Bai, Siddhartha Brahma, Joshua Ainslie, Kenton Lee, Yanqi Zhou, Nan Du, Vincent Y. Zhao, Yuexin Wu, Bo Li, Yu Zhang, and Ming-Wei Chang. Conditional adapters: Parameter-efficien transfer learning with fast inference. NeurIPS 2023, 2023. URL https: //arxiv.org/aat 200.33444443433300 04947.
12

---

Published as a conference paper at ICLR 2025
Dmitry Lepikhin, HyoukJoong Lee, Yuanzhong Xu, Dehao Chen, Orhan Firat, Yanping Huang, Maxim Krikun, Noam Shazeer, and Zhifeng Chen. GShard: Scaling Giant Models with Condi- tional Computation and Automatic Sharding. arXiv, 2020. doi: 20tp55500rrxxx220..6666000
Tsung- Yi Lin, Michael Maire, Serge Belongie, Lubomir Bourdev, Ross Girshick, James Hays, Pietro Perona, Deva Ramanan, C. Lawrence Zimicck and Piotr Dollár. Microsoft coco: Common objects in context, 2015. URL www.hhttt///cc
org/ abs /1405. 0312.
Xi Victoria Lin, Akshat Shrivastava, Liang Luo, Srinivasan Iyer, Mike Lewis, Gargi Gosh, Luke Zettlemoyer, and Armen Aghajanyan. MoMa: Efficient Farly-Fusion Pre-training with Mixture of Modality-Aware Experts. arXiv, 2024.
Yaron Lipman, Ricky T Q Chen, Heli Ben-Hamu, Maximilian Nickel, and Matt Le. Flow Matching for Generative Modeling. arXiv, 2022. doi: 20.5500/rxiv.2210.02222
Xingchao Liu, Chengyue Gong, and Qiang Liu. Flow Straight and Fast: Learning to Generate and Transfer Data with Rectified Flow. arXiv, 2022. doi: 1..48550/arxiv.2209 220.003000000.0
Brandon McKinzie, Zhe Gan, Jean-Philippe Fauconnier, Sam Dodge, Bowen Zhang, Philipp Dufter, Dhruti Shah, Xianzhi Du, Futang Peng. Floris Weers, Anton Belyi, Haotian Zhang, Karanjeet Singh, Doug Kang, Ankur Jain, Hongyu Hè, Max Schwarzer, Tom Gunter, Xiang Kong, Aonan Zhang, Jianyu Wang, Chong Wang, Nan Du, Tao Lei, Sam Wiseman, Guoli Yin, Mark Lee, Zirui Wang, Ruoming Pang, Peter Grasch, Alexander Toshev, and Yinfei Yang. Mml: Methods, analysis & insights from multimodal Ilm pre- pre-training, 2024 European Conference on Computer Vision, 2024. URL https: / org/ 220022201..0 05596.
Byeongjun Park, Sangmin Woo, Hyojun Go, Jin- Young Kim, and Changick Kim. Denoising Task Routing for Diffusion Models. arXiv, 2023. doi: (0-48500/axxiv2210.0075.
Byeongjun Park, Hyojun Go, Jin-Young Kim, Sangmin Woo, Seokil Ham, and Changick Kim. Switch Diffusion Transformer: Synergizing Denoising Tasks with Sparse Mixture-of-Eaie of-Exttss. arXiv, 2024. doi: 11445550/rxxv.2443..00
William Peebles and Saining Xie. Scalable diffusion models with transformers, 2023. URL https:
Dustin Podell, Zion English, Kyle Lacey, Andreas Blattmann, Tim Dockhorn, Jonas Müller, Joe Penna, and Robin Rombach. Sdxl: Improving latent diffusion models for high-resolution image synthesis, 2023. URL httpttttp/// /arxiv. wwtpiv.or/..v..ooooooo..ooooorrr abs/
Joan Puigcerver, Carlos Riquelme, Basil Mustafa, and Neil Houlsby. From Sparse Soft Mixtures of Experts. arXiv, 2023. doi:
Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J. Liu. Exploring the limits of transfer learning with a unified text-to-text transformer, 2023. URL httph:tt//// /atxiviivvv////////////a/////ava org att//bb///99 10683.
Aditya Ramesh, Prafulla Dhariwal, Alex Nichol, Casey Chu, and Mark Chen. Hierarchical text- conditional image generation with clip latents, 2022. URL https://aase http://awwww..//aaaa///// wwtp/bs/r////////e/oo/o/oo/ss 22 04. 06125.
Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Bjöm Ommer: High- resolution image synthesis with latent diffusion models, 2022. URL wwtph::s//aaii..oooo..
abs/2112. 10752.
Chitwan Saharia, William Chan, Saurabh Saxena, Lala Li, Jay Whang, Emily Denton, Seyed Kam- yar Seyed Ghasemipour, Bureu Karagol Ayan, S. Sara Mahdavi, Rapha Gontijo Lopes, Tim Sal- imans, Jonathan Ho, David J Fleet, and Mohammad Norouzi. Photorealistic text-to-image diffu- sion models with deep language understanding, 2022. URL https: /at:://axxxa /arxiv.vivv./////////a////aaatva org/ abs/ 22 05. 11477.
Noam Shazeer, Azalia Mirhoseini, Krzysztof Maziarz, Andy Davis, Quoc Le, Geoffrey Hinton, and Jeff Dean. Outrageously large neural networks: The sparsely-gate mixture-of-expert layer, 2017. URL https:// www.wxwwww..w.....eeooeeeeeeeeee 1701 06538.
13

---

Published as a conference paper at ICLR 2025
Zhan Shi, Xu Zhou, Xipeng Qiu, and Xiaodan Zhu. Improving image captioning with better of http///axvvvv///////////a wwtp:/ooorr/aaaaoaeeee 2006 6. 11807.
captions, 2020. URL htwps:ttppse..t /
Chameleon Team. Chameleon: Mixed-modal early-fusion foundation models, 2024. URL https:
09818.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need, 2023. URL https: //ar
xiv.org/ abs 1106/700.0000 03762.
Zeyue Xue, Guanglu Song, Qiushan Guo, Boxiao Liu, Zhuofan Zong, Yu Liu, and Ping Luo. RAPHAEL: Imaee Generation via Large Mixture of Diffusion Paths. arXiv, 2023. doi: 10. (0.445550/axxi2233518225
Harish Prabhala Yatharth Gupta, Vishnu V Jaddipal. Segmoc. https::/hugginggfett
egmind/ SegMoE- 4x2 v0, 2024.
Jiahui Yu, Yuanzhong Xu, Jing Yu Koh, Thang Luong, Gunjan Baid, Zirui Wang. Vijay Vasudevan, Alexander Ku, Yinfei Yang. Burcu Karagol Ayan, Ben Hutchinson, Wei Han, Zarana Parekh, Xin Li, Han Zhang, Jason Baldridge, and Yonghui Wu. Scaling autoregressive models for content-rich text-to-image generation. Trans. Mach. Learn. Res., 2022, 2022. URL https: ://s::/appppp. sema nticscholar. Sareg/CorpussD.. 2449228466
Hongkai Zheng, Weili Nie, Arash Vahdat, and Anima Anandkumar. Fast Training of Diffusion Models with Masked Transformers. arXiv, 2023. doi: 10448550arri..23660993333
Chunting Zhou, Lili Yu, Arun Babu, Kushal Tirumala, Michihiro Yasunaga, Leonid Shamis, Jacob Kahn, Xuezhe Ma, Luke Zettlemoyer, and Omer Levy. Transfusion: Predict the next token and diffuse images with one multi-modal model, 2024. URL htw.t:ttp//// //tps::///xxivv wwwp/br/////aee.e.o .11039.
Yanqi Zhou, Tao Lei, Hanxiao Liu, Nan Du, Yanping Huang, Vincent Zhao, Andrew Dai, Zhifeng Chen, Quoc Le, and James Laudon. of with expert choice routing, 2022. URL https: // a rxiv. org/ abbs/////// 2202..222222 09368.
Barret Zoph, Irwan Bello, Sameer Kumar, Nan Du, Yanping Huang, Jeff Dean, Noam Shazeer, and William Fedus. ST-MoE: Designing Stable and Transferable Sparse Expert Models. arXiv,, 2022.
doi: 10.48550/arxiv 2t..00/arxx.222.8890666
14

---

Published as a conference paper at ICLR 2025
A FUTURE WORK
In this work, our approach primarily focuses on adaptive computation based on image information, such as pattern complexity. There may also be additional factors, such object semantics compo- sitional relationships. that could further improve generation quality, and we plan to investigate how these additional sources of information can be leveraged to further boost EC-DIT's performance.
Additionally, the adaptive routing mechanism of EC-DIT can seamlessly integrate into any sequence- chunk-wis encoding and generation process. For instance, adaptive routing can be directly applied to the tokenization (Gafni et al., 2022) commonly used in mixed-modal, early- fusion language models (Team, 2024). Similarly, in multi-modal models unifying autoregressive and diffusion paradigms (Zhou et al., 2024), expert-choice routing can be naturally integrated into the diffusion-based image generation component. Furthermore, EC-DIT could potentially enhance video generation methods (Chen et al., 2023b; Guo et al., 2023), where the global context within each video chunk could be effectively utilized. We will also explore these directions in our future work.
B DETAILED DSG SCORES
Table 3: DSG Score breakdown. An internal version of Gemini was used to assess each sample's quality through VQA, as proposed in Cho et al. (2024).
Model (U) Score (-) Overall Counting Real Text Paragraph
Poses
TIFA-160
Relation Defying
DENSE-XL
72.50
71.50
50.00
60.70
86.70
74.80
83.70
80.40
73.90
EC-DIT-XL-SE
74.03
72.63
52.13
65.90
87.10
75.63
85.83
78.70
76.20
EC-DIT-XL-I6E
74.03
74,20
52.17
64.57
87.00
73.50
86.13
79.60
76.63
EC-DIT-XL32E
74.58
73.53
52.20
66.33
86.58
74.60
86.85
80.48
79.15
DENSE-XXL
74.80
74.40
51.60
66.20
87.30
76.10
86.30
82.10
77.80
EC-DIT-XXL-8E
75.03
74.17
53.55
67.95
86.47
73.58
87.15
81.35
78.60
EC-DIT-XXL-IGE
75.63
76.33
53.60
68.52
87.40
72.80
88.18
82.40
78.60
EC-DIT-XXIL-22F
75.93
76.55
53.63
71.57
88.18
73.63
87.80
80.85
78.08
DENSE-3XL
75.03
74.47
52.83
70.33
87.73
69.40
86.77
83,00
78.10
EC-DIT-BXL-SE
75.60
76.43
52.23
72.60
87.80
72.00
88,38
82.13
76.55
EC-DIT-3XL-I6E
76.10
75.63
54.10
73.10
88.07
73.70
88,00
81.57
77.90
EC-DIT-SXI-32E
76.20
78.27
54.47
71.33
87.63
73.17
87.90
82.00
78.13
EC-DIT-M-64E
76.40
77.33
54.67
72.50
88.23
74.33
87.27
82.40
77.97
Table 4: DSG comparison with token-choice baselines. GShard (GS) are implemented with two base model configurations (XL and XXL). GShard uses top-2 token-choice routing.
Model (L) Score (-) Overall Counting Real user
Text Paragraph Puses 16FA-160
Relation Defying
GS-XL GS-XL-8E
73.47
73.27
51.43
63.73
86.00
73.50
86.70
78.67
76.37
G6-
73.90
74.30
52.80
62.80
86.90
74.00
85,90
81.00
74.30
GS-XT--32E
74.07
74.13
52.17
64.53
86.50
75.23
86.60
78.47
77.23
GS-XXL-BE
75.10
76.10
53.80
68.60
87.40
72.70
87.70
80.50
76.00
-6-XXI--
75.30
74.50
52.83
66.57
87.27
77.67
87.73
80.63
78.33
GS-XXL-32E
75.40
76.40
52.70
67.20
87.10
78.90
88.40
79.80
76.30
Table 3 presents the detailed DSG scores for DENSE and EC-DIT across various model config- urations. Scaling with EC-DIT consistently improves performance several categories, such as Counting and Text, and leads to better overall performance. Table 4 shows that EC-DIT demonstrates superior performance over the token-choice baselines across most evaluation DSG categories.
15

---

Published as a conference paper at ICLR 2025
C GENEVAL COMPARISON WITH TOKEN-CHOICE BASELINE
Table 5: GenEval comparison with token-choice baselines. EC-DIT (EC) and GShard (GS) are implemented with two base model configurations (XL and XXL). GShard uses top-2 token-choice routing.
Model (L) Score (%) (-)) Overall Single obj.
Two obj.
Counting Colors Position Color aftr.
DENSE-XI,
67.08
99.80
82.47
64.28
82.46
19.61
53.86
GS-XL-8E
67.82
99.69
81.19
68.85
80.95
20.83
55.44
G6-
68.00
99.71
81.64
66.39
82.85
20.56
56.84
GS-XL-32E
68.52
99.92
82.88
69.18
83.23
20.30
55.63
EC-DIT-XL-8E
68.62
99.69
83.71
68.44
81.85
19.78
58.28
EC-DrT-XL-I6E
68.76
99.69
83.21
66.45
82.36
22.05
58.83
EC-DrT-XI--32F
69.38
99.79
84.14
68.61
83.68
21.02
59.03
DENSE-XXL
67.82
99.67
82.28
69.77
81.25
19.09
54.88
GS-XXL-8E
68.30
99.61
83,96
68.83
80.19
19.81
57.38
GS-XXL-I6E
68.49
99.59
84.88
69.04
80.37
19.97
57.09
GS-XXL-32E
69,96
99.38
86.17
70.00
85.51
20.44
58.25
EC-DIT-XXL-SE
69.11
99.84
85.43
68.59
81.85
20.31
58.61
1E-DIT-XXL-IGE
69.43
99.34
83.46
71.62
83.83
20.94
57.42
XC-DIT-XXL-32E
70.86
99.53
86.47
72.34
84.39
21.00
61.39
Table 5 further compares the GenEval performance EC-DIT with the token-choice base- lines (Lepikhin al., 2020). By leveraging global image information and adaptive computation, EC-DIT demonstrates superior performance over the token-choice baseline across most evaluation categories.
D CALCULATING ACTIVATED PARAMETERS INCREMENT
We present the pseudocode of calculating the activated parameters increment in Algorithm 2. Here, the parameters hi dden_dim, num.sparse.... www.sppareellaerr....m num num_heads, and num_kv for each model configuration specified in Table 1. We set ffn_factor=4.0 0 and capacity-factor-2 globally.
Algorithm Pseudocode for Activated Parameters Increment Calculation
Dense PEN Size
dense_ dense_fer_aize fér__ire 2 hidden_dim hidden_dim. ffc_factor
Router Sise
couter_size mun_experts hidden_dim
Atcention size
attr_size hidder_dim (num + nun_ky 2 num_heads) attn_key_din
Avated Dense Total
selivated_dense_LLoe (atin_eize Lense_fin_siee nun_spatee_
Activated Sparse
act lvated__ppar sparse_tttat
alta_size
hiddan_dim hidden_dim ffn_factor cspacity_factor
router_size nun_apazze_layerz
Acc Increment
aet inceeeeen aotivaaed__parae__too aetivaaedd_eetse__cdd
16

---

Published as a conference paper at ICLR 2025
E MORE GENERATED SAMPLES AND VISUAL COMPARISONS
Dense-3XL
EC-DIT-3XL-8E
EC-DIT-3XL-32E
EC-DIT-M-64E
Two cups of coffee, one with latte art the Eiffel tower. The other has latle art Statue Liberty.
EC DIT
EC DIT E
store front that has the word 'EC-
DIT written
EC-DIT
SCONT
shiny wearing race
suill and black visor slands proudly is setting on oityccape in the backiground. comic book illustration.
DON + PANiC!
DON'T PANIC!
DON'T T P PANIC!
book with the words Don't Panic!'
DON'T PANIC!
Figure 9: Visual comparison of three model configurations on the Partiprompts subset. For each prompt, the model generates four candidates, and the one with the best visual satisfaction is selected.
We provide an additional comparison on the PartiPrompt subset (Yu al., 2022) across different model configurations and also benchmark against two external text-to-image generative models: SD3-Large (Esser al., 2024) and FLUX. (Black Forest Labs, 2024).
Both SD3-Large and FLUX.1 [dev] have comparable higher activation sizes than EC-DIT-M- 64E. Despite this, EC-DIT demonstrates superior text-to-image alignment compared to these mod- els. For example, in Prompt #3, FL.U.1 fails to align with the textual elements hol liddd ding cane and holding abstract cubism smb In contrast, EC-DIT effectively captures both the described actions and the specified artistic style. Similarly, Prompt #5, SD3-Large and FLUX. [[dev] fail include the background detail of Dois Irmãos in Rio de Janeiro, while EC-DiT correctly incorpo- rates the described view.
garbage bag simultaneously. nor adhering to the requested style of
Furthermore, EC-DIT generates images with lower hallucination. For instance, in Prompt #2, SD3- Large hallucinates a third claw on the hamster, and FLUX. [dev] adds an unnecessary extra "!" mark. In contrast, EC-DrT generation appears visually accurate and reasonable.
17

---

Published as a conference paper at ICLR 2025
EC-DIT-M-644 M4EE (88)
SD3-Large (88)
FLUX.1 1 [dev] (128)
Prompt
This dreamilke digita captures vibrant, laleidecoopic bird lush rainforest
Prompt f2
A high contrast portrait photo of flufly hamster wearing orange beanle sunglasses sign that says "Ec-dit"
Ec-dit
Ec-dit
Ec-ditl
Prompt
A raccoon wearing tormal clomees, wearing top hat and halding a cane. The raccoon holding garbage bag. painting style abstract cublsm
Prompt 44
a portrait statue the Egyptian god Anubis wearing gogules, white tshirt and jacket Angeles the background.
Promet
A leddy bear wearing motorcycle helmet and cape riding motorcycle
Dols Immãos the background
Promet as
cozy corgi on the above rouch and round coffee table in front of couch and vase of flowers coffee table
painting
Prompt.87
drawing happy pig butcher selingg meat in is shop
Figure 10: Visual comparison of EC-DIT-M-64E and two external text-to-image models the Partiprompts subset. Numbers in parentheses indicate the activated parameter size for each model. All models leverage rectified flow for generation. SD3-Large inherently generates 1024 x 1024, while FLUX.1[det generates at 512 x 512.
18

---

Published as a conference paper at ICLR 2025
soliary figuro atrouded in mists poors from the cobbiastona streat ar the imposing dark gathic building surounding old-fashioned shines nearby, Oil painting
cheeseurger with juicy beef patlies and meted cheese sits on top of a foler that looks lhk a throne and stands in the middlo the royal chamber
A bowi soup that looks Mik a moneter made out of plasticnne
A gundam tall with skyscrapers
blue standing next red yellow standing in a lield with white flowers. inpressionielic painting
Three-quärters front view a otue 1977 Corvette coming mountain looking over a groen valley on cloudy day
raised. city the distance, with a mountain and ocean in the background. A dark moon le in the sky. Realatic nigh-contrast antme ilustration
in
A rusty speceship bissts of the foreground A cily with fall skyscrapera the distance, with mountain and background sky. high contrast ilusttation
small office made out car parts
A mischievous ferret with playlol grin squeezes itsel! into a large glase jar. surrounded by coloriul candy. The jar wooden table cazy kitchen. and sunlight nearby window
through
Figure 11: More generated samples with EC-DrT-M-64F on the Partiprompts subset. For each prompt, the model generates four candidates, and the one with the best visual satisfaction is selected.
19


Reviews:
[{"review_id": "gy3kaFhjD4", "reviewer": "Reviewer_gySw", "summary": "This paper proposes a Mixture-of-Experts (MoE) model, called EC-DIT, for diffusion transformers. By applying the prior work of expert-choice routing to Diffusion Transformer (DiT), EC-DIT allocates a fixed amount of tokens to each expert. As the authors suggest, the method has the advantage that tokens on the more complex parts of images can be assigned to more experts, resulting in higher computation allocated to more important tokens and thus higher efficiency.", "questions": "- We know that DiT performs the diffusion in the latent space instead of the actual image pixel space. Does it undermine the motivation of this paper to allocate more computations to image regions of higher complexity? Please explain how working in latent space affects their approach to computational allocation.\n- From Section 2.1, it seems that the authors use the Rectified Flow method on the DiT architecture with MoE. Can the authors explain why they chose to use Rectified Flow and how it integrates with their overall approach? This would help clarify the role of Rectified Flow in their method.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The proposed method is well-motivated. An image typically contains regions of different complexity, which require different computation costs during generation. It would be beneficial to use the MoE method to allocate computation to different image regions adaptively. In Figure 6, the visualization of the number of experts for each token shows that more experts are used for the image tokens corresponding to the objects (compared to the background) on the images. The figure provides good support for the motivation of this work.\n- The quantitative scores of EC-DIT show its effectiveness (Table 2). The evaluation includes multiple metrics, including GenEval, Davidsonian Scene Graph (DSG) scores, FID, and the CLIP Score.\n- The paper is well-written and the method is clearly explained with a pseudocode provided.", "weaknesses": "- For Table 2, can the authors include the FID metric to compare the generated images' visual quality across the models? \n- Can the authors add a column for inference time or throughput (e.g. images/second) in Table 2? It would help the readers better understand the efficiency of different models.\n- This paper has a good contribution to the community by showing the benefits of applying expert-choice routing in DiT, although the main approaches are mainly based on existing works, which makes the proposed method slightly less novel. \n- Can the authors make the notations in Section 2.1 regarding the rectified flow clearer? Please explain $p_0$, $x_0$, $x_t$, and $v_\\theta$ in more details.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "S54fGDvBQN", "reviewer": "Reviewer_MJns", "summary": "This paper focuses on scaling diffusion transformers for text-to-image synthesis. It proposes EC-DIT, a new family of Mixture-of-Experts (MoE) models with adaptive expert-choice routing. By leveraging the computational heterogeneity of image generation, EC-DIT adaptively allocates compute to understand input texts and generate image patches, enabling heterogeneous computation aligned with text-image complexity. It scales up to 97 billion parameters and shows significant improvements in training convergence, text-to-image alignment, and overall generation quality compared to dense and conventional MoE models. Its contributions include introducing EC-DIT for adaptive computation in text-to-image synthesis, achieving promising performance at scale with faster loss convergence and better quality, and conducting comprehensive experiments to validate its superiority and demonstrate effective adaptive compute allocation.", "questions": "see weakness", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. EC-DIT can scale up to 97 billion parameters, achieving a state-of-the-art GenEval score of 71.68% in text-to-image synthesis.\n2. The adaptive expert-choice routing is a novel design. It enables efficient compute allocation by having experts select relevant tokens, ensuring load balance without an extra load-balancing loss. This approach better utilizes computational resources, adapting to the complexity of different image regions and text, and improving overall performance and efficiency in handling inputs.\n3. By embedding timestep and using cross-attention modules, EC-DIT effectively integrates global information, which allows the model to understand the context of the image and its relationship with the text, guiding compute allocation at different stages. It results in hierarchical processing, enhancing image quality and text-to-image alignment, and making the generation process more context-aware.", "weaknesses": "1. The best score for SD3 (Esser et al., 2024) w/ DPO is 74% in the original paper, which is actually higher than that of EC-DIT (71.68%). Even with the value 71% listed in Table 2,  EC-DIT is only gaining marginally. Given its additional overhead, it is hard to see the value of the gain. \n2. The difference between the theoretical and actual inference overhead (e.g., for EC-DIT-M, the theoretical is around 3% but the actual is 23%) indicates that there may be room for optimization in the inference process. The authors may want to explore the factors contributing to this overhead or propose specific strategies for mitigating it, which could further enhance the practical applicability of the model, especially for applications where real-time or fast inference is crucial.\n3. The generalizability of the adaptive expert-choice routing is unclear. The performance of EC-DIT is mainly demonstrated in the context of text-to-image synthesis. How would it fit into the task of pre-training mixed-modal, early-fusion language models? It would be interesting to explore its effectiveness and efficiency in other related tasks. \n4. This paper could benefit from more detailed ablation studies on specific components of the proposed model. For example, understanding the individual contributions of the timestep embedding, cross-attention modules, and the expert-choice routing mechanism in more detail. This lack of in-depth analysis makes it harder to understand which parts are truly driving its performance.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "S3hUoCUbAP", "reviewer": "Reviewer_J1KC", "summary": "This paper proposes a MoE based diffusion transformer framework to scale diffusion models. The basic idea is to combine DIT and MoE. The experimental results demontrate the effectiveness.", "questions": "While efforts have been made to enhance the scalability of diffusion models for improved performance and efficiency, several issues remain:\ni) The novelty of this approach requires clarification, especially in comparison to existing diffusion models incorporating MoE, as detailed in the related work section (lines 057-063).\nii) Certain claims should be refined or revised. For instance, the statement: “our approach significantly improves training convergence, text-to-image alignment, and overall generation quality compared to dense and sparse variants with conventional token-choice MoEs (Lepikhin et al., 2020b; Du et al.,2022; Zoph et al., 2022).” First, I do not see the comparison with these methods in experiments; Second, the performance is only comparable with SOAT, thereby questioning the usage of 'significantly.'; Third, the reason of improvement might be mainly attributed to increased model parameters, e.g., in experience, more parameters make the training converge faster. Thus, deeper analysis could be given.\n“EC-DIT leverages this global context, combining multi-modal information to optimize token-to-expert allocation.” multi-modal?\n“The model resolution is set to 256 × 256, with a patch size of 2.” “This results in the input sequence length of 128 per image.” Patch size is 2*2?\niii) Those existing diffusion models with MoE should be compared in experiments, if they are available.\niv) Fig. 3 only shows the performance on DSG score. Additional evaluation on other metrics should be showed, as the performance across various metrics might exhibit different trends. \nv) Fig. 4 should show more comparisons with other methods (e.g., in Table 2), rather than solely focusing on the proposed methods themselves. \nvi) How to set C in Top-k? There lacks the comparison with the different C values.\nvii) The performance is only comparable with SOAT methods, such as SD3. In fact, the parameter quantities are in the same magnitude, SD3 (8B) vs EC-DIT-M-64E(>8B, observed in Fig. 4). Thus, this raises questions about the core contributions of this work. Is the primary focus acceleration? If so, there lacks many comparisons to current methods.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 4, "contribution": 2, "strengths": "i) The technique line is reasonable.\nii) The writing is easy to follow.", "weaknesses": "i) This work should be incremental in essence, as the basic idea is to combine DIT and MoE.\nii) Certain statements exhibit overclaims that need refinement. (see Questions in detail)\niii) Crucial experiment comparisons with diffusion methods utilizing MoE are notably absent. (see Questions in detail)\n\nThe rebuttal addresses my concerns, so I raise the score.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "hq7xlMazTV", "reviewer": "Reviewer_AMiW", "summary": "This paper explores how to scale diffusion transformers (DiT) to larger model capacities. The authors developed EC-DiT, a MoE-based (Mixture of Experts) DiT model leveraging expert-choice routing, expanding the DiT model parameters to 97 billion. EC-DiT enables context-aware routing and dynamically adjusts computation, inherently ensuring load balancing across experts. Extensive experiments were conducted to verify and analyze the effectiveness of EC-DiT.", "questions": "Please refer to the weaknesses part.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- EC-DiT successfully scales the DiT model to 97 billion parameters, demonstrating exceptional performance.\n- Expert-choice routing outperforms traditional token-choice routing strategies while maintaining competitive inference speed.\n- The paper presents extensive experiments to showcase EC-DiT’s capabilities.", "weaknesses": "- While I recognize the authors’ contribution in scaling up DiT and support the paper’s acceptance, I believe it does not demonstrate sufficient technical innovation beyond prior work, appearing more as a practical integration of existing techniques\n- A clearer and more precise definition of ‘Activated Params.’ should be provided. In Algorithm 1, it appears that all parameters are activated. Does this refer to the average number of parameters activated per token?\n- I would like the authors to include a comparison of memory usage during training and inference.\n- It would be beneficial if the authors could open-source the relevant scripts and model weights upon acceptance.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "AxMCoBm0RC", "reviewer": "Reviewer_U4nK", "summary": "The paper presents EC-DIT, a new Mixture-of-Experts (MoE) model specifically designed for diffusion transformers in text-to-image synthesis. EC-DIT leverages computational heterogeneity through expert-choice routing, which enables the model to allocate computational resources adaptively based on the complexity of the input text and the corresponding image patches. This approach supports scaling up to 97 billion parameters, leading to improvements in text-to-image alignment, and generation quality.", "questions": "Please refer to the weaknesses.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "1. Investigating methods to scale up diffusion transformers is a crucial research direction, as it directly impacts the effectiveness and potential of large-scale text-to-image synthesis models. \n\n2. The proposed EC-DIT model's effectiveness is verified by quantitative metrics.", "weaknesses": "1. As shown in Figures 1 and 9, the generated images are relatively simple and lack sufficient complexity. For more intricate images, the aesthetic quality appears subpar, and the benefits of scaling up are not convincingly perceptible. Offering 64-128 more generated images with more complex scenes could better demonstrate the value of the proposed method.\n\n2. Models like SD3 and Flux, which have significantly fewer parameters than EC-DIT-64E, manage to produce images with notably superior visual quality. Please clarify why this discrepancy exists and address what specific advantages scaling up provides if the visual quality is still lacking. This raises questions about the true value that scaling up offers in terms of image quality.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)