Remix-DiT: Mixing Diffusion Transformers for Multi-Expert Denoising
Authors: Gongfan Fang
Xinyin Ma
Xinchao Wang
Venue: NeurIPS 2024 poster (2024)
Keywords: Diffusion Models
Multiple Experts
Primary area: diffusion_based_models

Abstract: Transformer-based diffusion models have achieved significant advancements across a variety of generative tasks. However, producing high-quality outputs typically necessitates large transformer models, which result in substantial training and inference overhead. In this work, we investigate an alternative approach involving multiple experts for denoising, and introduce RemixDiT, a novel method designed to enhance output quality at a low cost. The goal of RemixDiT is to craft N diffusion experts for different denoising timesteps, yet without the need for expensive training of N independent models. To achieve this, RemixDiT employs K basis models (where K < N) and utilizes learnable mixing coefficients to adaptively craft expert models. This design offers two significant advantages: first, although the total model size is increased, the model produced by the mixing operation shares the same architecture as a plain model, making the overall model as efficient as a standard diffusion transformer. Second, the learnable mixing adaptively allocates model capacity across timesteps, thereby effectively improving generation quality. Experiments conducted on the ImageNet dataset demonstrate that RemixDiT achieves promising results compared to standard diffusion transformers and other multiple-expert methods.

Remix-DiT: Mixing Diffusion Transformers for Multi-Expert Denoising
Gongfan Fang Xinyin Ma Xinchao Wang"
National University of Singapore
{gongfan an mww.mmxiinnnn)... nus. edu, zinchao@mus edu
Abstract
Transformer-basee diffusion models have achieved significant advancements across variety of generative tasks. However, producing high-quality outputs typically necessitates large transformer models, which result in substantial training and inference overhead. In this work, we investigate alternative approach involving multiple experts for denoising, and introduce Remix-DiT, a novel method designed enhance output quality low cost. The goal of Remix-DiT is craft N diffu- sion experts for different denoising timesteps, yet without the need for expensive training of N independent models. To achieve this, Remix- employs К basis models (where K < N) and utilizes learnable mixing coefficients to adaptively craft expert models. This design offers two significant advantages: first, although the total model size is increased, the model produced by the mixing operation shares the same architecture as plain model, making the overall model efficient as standard diffusion transformer. Second, the learnable mixing adaptively allo- cates model capacity across timesteps, thereby effectively improving generation quality. Experiments conducted on the ImageNet dataser demonstrate that Remix- DiT achieves promising results compared standard diffusion transformers and other multiple-expert methods.
Introduction
have shown significant potential in generating high-quality images and videos [24, 20, 3 However, achieving such quality often requires large transformer architectures (35), which incur considerable training and inference costs. To alleviate these computational burdens, multi-expert denoising has emerged as a promising approach [1 [8] which employs multiple specialized diffusion models, each designed for distinct time intervals within the denoising process.
Transformer-basee Diffusion models 39 8 29
28 16 27
The goal of multi-expert denoising is to increase the overall capacity of diffusion models while keeping an acceptable overhead. The denoising process of diffusion models involves multiple different timesteps [3] 32, 18), requiring the network make predictions various noise levels. Previous work has shown that the tasks of diffusion models vary across different timesteps 37 27
For instance, higher noise levels, the model focuses more on low-free low-frequency features, while lower noise levels, the model emphasizes generating high-frequency details [37]. This variability inherently leads multi-task problem [16] 18). However, due to the limited capacity of single diffusion model, it is challenging to craft comprehensive and balanced model that performs well across the entire denoising process. To alleviate this issue, multiple-expert denoising deploys specialized expert models at different timesteps. To this end, each model only needs to learn the denoising task for specific subset of timesteps. Although this design introduces multiple models, only one model
*Corresponding Author
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

Lurable Milig Cofficiat (lo-))
g
20x
D 200
08000
Dizb
4X
Dagge
Haleon Soo
Halln Sou 2x11
4x
4X
Faacazing
(a) Training larger DIT
(b) 4x Experts over timesteps
(c) 20x RemixDiT with Basis DITs (Ours)
Figure 1: (a) Re-training larger DiT incurs significant training and inference costs. (b) Multi- expert denoising trains multiple expert models to improve generation quality while maintaining low inference overhead. However, training multiple experts still results in significant training costs. (c) This work introduces RemixDiT, a learnable method to craft any number of experts by mixing basis models.
activated at each timestep during the inference process, and therefore the overall computational complexity does not increase significantly.
In multi-step denoising process, such as a 1000-step denoising chain, determining the optimal number of intervals and their partitioning remains an open problem. Existing approaches typically employ even-split strategy, which distributes the denoising process equally among predefined number of expert models. This uniform allocation often leads to suboptimal utilization of model capacity. In addition, the training of multiple expert models introduces substantial computational costs. Due to the inherent task similarity between adjacent intervals in the denoising process, training separate experts without accounting for this similarity can result in redundant and inefficient training. Although some recent methodologies, such as tree-structured training [. have proposed hierarchical strategies to mitigate these issues, the process of training and capacity allocation remains largely manual and heuristic. To address these challenges, this paper introduces a novel, Iearnable strategy to craft multiple experts adaptively for denoising.
Consider denoising process of T steps. The objective of this paper is to construct N expert models, each handling an interval of length N as shown in Figure 1 (b). A straightforward approach to improve performance is to increase the number of experts, as more experts provide greater total capacity, leading enhanced performance. However, this also results in a linear increase in training costs. The core idea of the proposed method lies in training only K K(K < N) basis models, which can be used to craft N expert models efficiently for inference. To be exact, we construct K transformer-basee diffusion models and introduce learnable mixing coefficients to fuse their parameters, which can be casily implemented an embedding layer. During training, the coefficients will learn to allocate the model capacity across different timesteps. During inference, we precompute those mixed models to do a chained multi-expert denoising.
We validated the effectiveness of our method on ImageNet-256x256 6 Owning to the learnable nature of Remix-DiT, our approach adaptively allocates model capacity across timesteps, thereby achieving superior generation results compared to independently trained multiple expert models. Our analysis of the learned mixing coefficients revealed several insightful findings: Firstly, the coefficients exhibit a natural similarity between adjacent timesteps, indicating that tasks at neighboring steps are relatively similar. Conversely, the coefficients for timesteps that are further apart show distinct differences, supporting the multi-expert denoising hypothesis that the learning target at different steps can be quite diverged. Additionally, our method's learnable coefficients tend to allocate more capacity to early timesteps, suggesting that the algorithm identifies predictions at lower noise levels as more challenging and thus requires greater capacity to generate finer details. In contrast, at higher noise levels (i.e., timesteps approaching T), the algorithm integrates multiple basis models for higher utilization of model capacity. Furthermore, we observed that experts created through hybridization exhibit specialization, achieving lower loss at specific timesteps while incurring higher loss at others. Through collaborative denoising, such specialization reduces overall prediction errors, highlighting the effectiveness of our method.

---

The contribution of this work lies in a novel method for multi-expert denoising, which creates N experts by combining K (K<N) basis models. This approach effectively reduces training costs and improves performance.
2 Related Works
Multiple Experts in Diffusion Models The multi-step denoising process inherent to the diffusion
can be viewed multitasking problem I 18 28 T6]. At each step, the model receives inputs with varying levels of noise and makes predictions accordingly. Given the disparate learning objectives for each denoising step, the utilization of a singular model across all stages is, to a certain extent, inefficient and challenging. Several prior works suggested employing multiple models, each dedicated to handling partial timestep intervals. For example, E-Diff [ introduces simple binary strategy to learn various experts and during inference, take the ensemble of experts for prediction. OMS-DPM [8] proposes a model scheduler, dynamically taking denoiser of different sizes from model 200 for inference. MEME [66] deploys a similar paradigm, yet training lightweight models for different steps. The ensemble of these lightweight denoisers can achieve superior performance compared to their larger counterparts while maintaining efficiency during inference. Another exploration to enhance efficiency using multiple experts is T-Stich [27], which stitches the denoising trajectories of different pre-trained models, enabling dynamic inference cost without additional training. However, it's noteworthy that multiple diffusion models may incur additional storage and context switch costs. Thus, DTR (28) introduces novel multi-tasking strategy to learn a single model with scheduled task masks, activating specific sub-networks for different timesteps.
model 14
33 34
Transformer-based Diffusion Models Transformer-based diffusion models have achieved im- 20
17]. DiT [29] explores the scalability of transformers for image generation, achieving competitive performance compared to CNN-based diffusion models. UVIT [2] independently explores transformer design for generation, incorporating
pressive results image generation 2 29
the U-Net [31 22] 30 architecture for improved training efficiency and quality. Furthermore, recent
works have expanded transformer- bassomdeer-bbeddddde diffusion models to video generation 22
21 24 3),
(25] which shows the power of transformers in modeling complicated data. However, training diffusion models remains inefficient, typically requiring millions of steps pro- duce satisfactory results. A series of works in the literature focus on the efficiency of diffusion- MaskDiT 44 and MDTv2 [10 utilize masked transformers and an additional masked autoencoder task [12 to streamline training. UVIT ( also demonstrates the benefits of skip connections for accelerating convergence. In this work, we introduce a new strategy to enhance the training efficiency of plain diffusion transformers
audio [19 [5| and 3D
transformers 40 40
3 Preliminary
Diffusion Probabilistic Models Diffusion Probabilistic Models (DPMs) train a denoiser network revert a process of adding gradually increased noises [13] Given an input zo. T'step forward process is deployed to transform zo into latent zi.... zT, where the transformation is defined as q (1) Nxx; vI -1, Br1) under a increased variance schedule BI... To revert this process, a network q( - (x)) is trained to predict rz-1 given 2t. By defining ā. = II Ot. the training objective is formalized as
1 - B, and
( 10r. [||c -o-oooVVaa + v1 ā.c, t)
(1)
where Eg is trainable neural network and refer to the noises at timestep A common practice in the literature is to train a single neural network for all timesteps t. However, it has been observed that the denoising behavior usually varies across timesteps 13 37 27). Due to the limited capacity, single model may struggle to fit all steps. Therefore, more suitable yet natural approach is to use multiple models specifically learned for different timestep intervals
Multi-Exper Denoising To facilitate denoising with N expert models, the timesteps are typically
divided into N intervals. Each model, denoted as 60,, contains independent parameters 0.. For each model 0,, minimize the training objective formalized in Equation I on the corresponding
3

---

Coefficients
Mixed Linear Experts ab
Linear Layer Basis θ
Inference 5T/6]
Mix
[0. T/6]
4xLinear
Figure 2: An example of mixing linear layers basis into 6 expert layers. Each expert linear layer is a weighted averaging of the basis layers. At each denoising interval, only one expert is activated for inference training. To increase the number of experts, increase the number of coefficients a. which is more efficient than independently training new experts.
timestep interval E ( x J. (i++)) x [ as follows:
((O,) : E, [ex & (++1)xx | (0..0-//(1111 [l|c [ rec-ce.(varr + a ā.e,t)]
t)]
(2)
Although the simple form of multi-expert denoising has been verified effective in previous works [T 16), the above objective still presents the best interval partition is unknown; 2) Training multiple independent experts can be expensive. In this work, we introduce
major challenges: The optimal number of experts and
approach, Remix-DiT to address the above issues.
4 Method
Uniform Partition of Denoising Process Let's consider multi-expert denoising problem with N experts for T timesteps. For simplicity, we assume that each expert's parameters 0; a column vector with P trainable elements and the parameter matrix of all experts can be formalized as ONXP 0,, 12. ....... 0 0v] This is natural since the parameters in model can be flattened and concatenated into vector. We consider simple partition of timesteps, which equally divides the denoising process into intervals, i.e., NO.T/N]...[(N - 1) T/N, N T/NJ. The motivation behind the oracle partition is that adjacent steps share similar noise levels and training objectives. Following Equation this leads to the learning objective to optimize each expert on their associated time intervals.
Σ C(O.).
arg min
0,E0
The above process involves N independent optimization problem, which presents some issues. First, the optimal number of experts is unknown. It's difficult to increase the number of experts since this will introduce huge training costs. Besides, the optimal partition is also unclear, making the uniform partition inefficient. In this work, we investigate such a problem: Is it possible to learn any number of experts, while keeping an affordable overhead?
Crafting Experts by Mixing address the problem of uniform partition, we leverage set of basis models to avoid straightforward training on experts. The core idea lies in that, it is possible to fuse the parameters of two diffusion models to achieve better performance 38 Formally, instead of training N independent models directly, the core idea of Remix-DiT is to learn K (K < N) basis models with the same architecture as experts, parameterized by 3 3xxP B2... Bx]] And the expert models can be crafted by mixing the basis parameters with certain coefficients. For each expert 0, associate with coefficients a; [am,, 0:2,.... [ak] and compute the mixed expert parameter through a weighted averaging 0, = Σn Ou This leads to the coefficient matrix denoted as xN K = [a,, .2,..., xN]] Then, the expert's parameter matrix can be easily obtained through a simple matrix multiplication:
NxP aNXK
(4)

---

With Equation 4 we can freely craft different experts using different mixing coefficients. Now the problem lies in that how to learn a good coefficients a and basis models B3.
Architecture of Remix-DiT Let's fur- ther delve into the details of K basis mod- els, denoted as (3, alongside coefficients matrix a. In this study, we adopt the DiT architecture proposed in [29]. which pri- marily consists of linear and normalization layers. Rather than constructing K distinct models, we propose a simple trick to con- struct a singular and extended DiT model encapsulate all K basis models [36]. This achieved by increasing the width of the linear layers by factor of K, with a modi- fied forward process to support the mixing of basis. An example of the extended linear layer is depicted in Algorithm [ Remix-DiT is the integration of weighted averaging step before the true computation. Despite the expansion in width by factor of K, the effective weights engaged during forward propagation are equivalent to those of a normal DiT of the original width. The additional computational cost is attributed solely to the element-wise operation during mixing. Furthermore, during the backward propagation phase, the computational cost of the mixed network almost mirrors those of a standard DiT, except for the additional overhead caused by the mixing operation.
Algorithm Mixed Linear (PyTorch-lik) like Pseudo Code)
1:
function LINER,,, w, coeffs, K)
2:
x: linear inputs
3:
w: linear parameters
4:
coefi's: mixing coefficients
5:
K: number of basis
6:
W w.view(K, w.shape( 1J//K,
7:
w_ = (coeffs * w).sum(()
8:
return torch.an..uncctiiii.. www.aa...nnnnnaaiineneeo.oa W_)
9:
end function
A notable distinction between the standard DiT and the proposed
Network Training A notable challenge presented by Remix-DiT is the restriction of activating only one mixed expert during each forwarding and backwarding phase. To address this limitation, employ hierarchical sampling strategy whereby the expert 0, is selected first, followed by random selection of timestep within the interval Ii x 1 (i + 1) x L J]- Consequently, only one expert is activated per training step. Despite this constraint, a critical advantage of Remix-DiT is that regardless of which expert is activated, all basis models remain updatable. This feature distinctly sets our approach apart from methodologies that train experts independently. The training and inference processes of Remix-Dii DiT are delineated in Algorithm 2 The training protocol for Remix- DiT closely mirrors that of standard DiT, with the exception that the sampled timesteps must originate from a singular time interval during each step. During inference, the compact nature of the basis models allows for reduced GPU memory usage compared to maintaining N separate experts, providing a significant advantage. Additionally, it is still feasible to pre-compute the N experts for more efficient inference without runtime mixing.
Transform Pre-trained DiT to Remix-DiT with Prior Coefficients Training generative models from scratch is usually inefficient. Note that the proposed method Builds a DiT model with K-times width. This allows it to inherit the weights of pre-trained DiTs, by replicating the pre-trained weight K times. This leads to K basis models with the same parameters. However, this will be problematic if we inspect the gradient to each mixing coefficient, denoted as V V
VOLk Σ Vo. C
(5)
where 3 is an element-wing product of two matrices. First, the update of basis parameter 6 usually slow, this makes all mixed experts similar to each other at the beginning of training and also makes the gradient w.r.t the mixed expert Ó, almost unchanged. In this case, the gradient to the coefficients will be constant. To address this issue, we introduce prior regularization to force each basis Q to learn different objectives. Inspired by the oracle partition in multiple expert denoising, we craft prior coefficients a. with one-hot vectors. For example, the hand-coded prior coefficient [0, 1, 0, 0] directly assigns the second basis as the expert, this encourages a polarization for different bases. And the regularization can be easily implemented as cross-entropy between the prior coefficients and the learnable coefficients.
R 7 Σ ai log(acu))
(6)
5

---

Algorithm 2 Remix-DiT
1: // Training
2: Initialize K basis models 3 12..3 from random or pretrained models
3: Randomly initialize mixing logits [Ti, 2........
4: while Training terminated do
5:
Randomly sample expert index E N) N)
6: Compute the mixing coefficients a; = Softmax(T) (T.)
7:
Obtain the i-th mixed expert 0, = ox,
8:
Randomly sample the timestep li x 1 (i+ [
9:
Compute the loss Lr. = ||e ce, (vaizu + -1- ā,c, t)ll
10: Update and @ by back packppropngttion
end while
// Inference
13: Compute the mixing coefficients a (oftmma Softmax (Trt)).....S.Sfma(m
Prepcompute experts with a" 15: Inference with
where is hyperparameter to control the strength of regularization. In practice, we adopt an annealing strategy to linearly remove the regularization with 7 - O.
5 Experiments
5.1 Experimental Settings
This paper utilizes the DiT [29] models as the fundamental architecture. By reloading the computation process of the linear layers, additional mixing coefficients is introduced to create RemixDiT. For a fair comparison, we ensure that the dimensions of all mixed experts are completely consistent with the original DiT. For instance. Remix-DiTe is constructed by quickly combining multiple DIT-S models to form an expert model.
Network Architecture.
Training Details. Since our mixed model is entirely consistent with standard DiT, this method can be applied to pre-trained models. We first trained standard DiT model on ImageNet and then initialized the basis model with the same pre-trained weights. We introduce prior as discussed in Equation to the mixing coefficients to accelerate the learning of different basis models. In our experiments, we conducted 100 fine-tuning on DiT-S/B/L models (29, pre-trained for 2M/IM/IM steps correspondingly.
6
5.2 Transform Pretrained DiT Remix-DiT
In T present the results of fine-tuning standard DiT models using the Remix-Dis approach. Given the architectural consistency between the mixed model and the standard DIT, it is feasible to initialize K basis models within the standard DiT framework. This enables the construction of RemixDiT with minimal additional training steps since the pre-trained DiT has been comprehensively trained at all temporal steps. For instance, starting with a pre-trained DiT-B model, our method suc- cessfully crafts a Remix-B with only 100K training steps, achieving superior performance compared to both the original models and baselines such as continual training and multiple experts (1.
Moreover, within the same computational budget, our approach outperforms Multi-expert basclines, where experts are trained independently. It is noteworthy that for Multi-expert baselines with eight experts, the total training budget is uniformly allocated each expert. As the number of experts increases, the allocated training steps for each expert become more limited. In contrast, Remix- DiT trains shared basis models throughout the entire training process, allowing each model to be sufficiently updated.
As discussed earlier, RimixDiT allocates gradients from different time steps by employing mixing coefficients. Therefore, for K basis models initialized similarly, greater diversity in the initial mixing coefficients facilitates the rapid learning of different models. To achieve this, we sequentially distribute the entire denoising process among the basis models. During subsequent fine-tuning, the

---

Conditional Diffusion Transformers (DiT) ImageNet 256 256 efg=115
Method
#Effective Params
#Gflops per step
18
FID
Precision
Recall
DiT-L/2-IM
458. M
80.79
196.34
3.73
0.8187
0.5398
100k Cont. Training
458.10 M
80.79
200.16
3,57
0.8193
0.5356
100k Multi Experts XxL
8x458.
80.79
205.39
3.41
0.8149
0.5448
100k Remix-L-4-20
4x458.10
80.79
207.54
3.22
0.8175
0.5451
DiT-B/2-IM
130.51 M
23.04
119.74
10.11
0.7348
0.5487
100K Cont. Training
130.51 M
23.04
125.32
9.31
0,7418
0.5499
100k Multi Experts 8x
8x 130.51
23.04
126.58
9.28
0.7396
0.5504
100k Remix-B-4-20
4x 130.51
23.04
127.42
9.02
0.7453
0.5553
DIT-S/2-IM DIT-S/2-2M
32.96 M
6.07
48.46
32.55
0,5476
0,5677
32.96 M
6.07
59.39
26.51
0.5896
0,5613
100k Cont. Training
32.96 M
6.07
60.34
26.05
0.5941
0.5672
100k Multi Experts 8xS 1
3x M
6.07
63.40
25.00
0.6017
0.5622
100k Remix-S-4-20
4x M
6.07
69.23
22.84
0.5998
0.5705
Table 1:
Finetuning pre-trained DiT for 100k steps on ImageNet-256x255 During inference, all
expert share the same architecture as standard DiT.
variation in coefficients will enable refined allocation of model capacity. In Table compare the fine-tuned Remix-B to other diffusion models.
5.3 Visualization of Learned Experts
To further investigate the behavior of the experts learned by the algorithm, we visualized the mixing coefficients. shown in Figure Mal we conducted experiments on DiT-S. During the training of 20 ex- pert models, we observed that the algorithm assigned more one-hot coefficients to timesteps close to 0. At these steps, the denoising model focuses more on high-frequene detail features. In contrast, at late timesteps resembling random noise, the algorithm enabled and mixed multiple models to generate con- tent rather than using only one expert. illustrated in Figure mixing multiple models improve the shape quality. Morcover, it can be observed that adjacent timesteps exhibit similar mixing coefficients, whereas more distant timesteps show greater differences. This finding supports the core hypothesis of multi-expert reasoning. To validate this observation, we further trained a RemixDiT with 8 experts. Figure 4b also shows the distribution of mixing coefficients, revealing results similar to the previous experiment. Additionally, Figure 4c visualizes the training loss functions of the expert models throughout the denoising process. Each mixed expert has lower loss function within its respective timestep intervals. The farther from an expert model's 's current interval, the higher its loss. Notably, at step O, we found that the ensemble-constrrueee model effectively reduces prediction loss.
Method
#Params per Step
FID
ADM 7
554M
10.94
IDDPM 26
270M
12.26
VQ-Diffusion
11.89
DIT-B/2 IM 29
130 M
10.11
MЕME 16
103
13.19
130
9.02
Figure 3: Comparision to existing methods
will be
5.4 Analytical Experiments
In this section, we validate some key design aspects of RemixDiT. This includes the model's mixing method, the number of expert models and basis models, and whether the coefficients are independent or shared across layers.
Model Mixers. The core of the proposed method lies in the mixing of multiple models, and explored three mixing methods: 1) Oracle Mixing: Basis models are manually assigned to different intervals as expert models, equivalent training independent expert models, which makes the algorithm revert to the standard multi-expert training strategy regardless of N. 2) Raw Mixer: Each expert is assigned real-valued coefficient without any constraints. 3) Softmax Mixer: This builds on the raw mixer by applying softmax operation, ensuring that the mixed model is always weighted average of the basis models. Table shows the effectiveness of each mixer. We found that while oracle mixing improves model performance, it often fails to achieve optimal results due manually
7

---

Basis 1- 0.89
0.57 0.30 0.13
0.04
0.01
0.01 0.02 0.03 0.04 0.08 0.11 0.15 0.20 0.23 0.25 0.28 0.29 0.31 0.32
Basis
-0.06
0.07 0.18
0.32
0.46
0.57 0.64 0.67 0.63 0.57 0.47 0.37 0.27 0.21 0.19 0.20 0.22 0.23 0.23 0.22
0-0.02 0.11 0.17 0.20 0.23 0.25 0.28 0.31 0.40 0.43 0.43 0.40 0.27 0.21 0.18 0.17 0.17
Basis 4 40.02 0.25 0.35 0.35 0.29 0.19 0.09 0.04 0.03 0.04 0.06 0.09 0.14 0.19 0.24 0.28 0.29 0.29 0.30 0.30
10, 50)
150. 100)
1100, 150)
1250. 200)
1200. 250)
1250. 3n0)
[300, 350)
1350, 400)
[4n0. 450)
1450, 500)
[500, 550)
(550, soo)
1600, 650)
l6s0, 700)
1700. 750)
1750. and)
rano. 8s0)
leso, 900)
1000. 950)
1950. 1000)
(a) Learned mixing coefficients for RemixDiT-S-44
which crafts 20 experts by mixing 4 basis models.
14
Busis
0.66
0.28
0.04
0.02
0.08
0.19
0.25
0.28
Bosis
0.12
0.57
0.73
0.56
0.35
0.25
0.25
0.26
Bosis 3
0.11
0.07
0.14
0.75
0.32
0.25
0.16
0.10
Bodic
0.11
0.08
0.10
0.17
0.26
0.31
0.34
0.35
125
200)
375)
soa)
lex 7501
1858 875)
1075, 1008)
1125
iss
LPS
(b) Learned mixing coefficients for Remix Dimix--iT-SSe DiT-S-4-8 --8--8
(c) Training losses of 8 experts over timesteps
Treastee
Figure 4: (a) Learned Coefficients for a Remix-DiT-S-4-5 which mixes basis DiT-S to obtain 8 models, each associated with a 125-step interval. The x-axis shows the corresponding timestep intervals associated with each mixed model. The value in each grid refers to the coefficients that will be used to weight the corresponding basis models. At early timesteps T + 999, Remix-DiT tend to use an ensembled model for inference, which is averaged over all basis models. And at the late timesteps T > O, more specialized models, such the first basis model are picked for fine- mne-grained prediction.
designed partitions. The raw mixer, with its learnable coefficients, achieved better performance but included some impractical solutions, such negative coefficients, which increased learning difficulty. The softmax mixer yielded the best results.
Global or Local Mixer. A DiT model typically comprises multiple layers. When mixing models, we can use globally shared mixing coefficients, referred to as the global mixer, or employ a layer-wise independent mixing strategy, referred as the local mixer. The first strategy is simpler to train, as it requires only small K x N coefficient matrix regardless of the network's 's depth. The second strategy, however, offers a larger model space by introducing trainable parameters that grow linearly with the number of layers. Despite this, the additional parameter size for both methods is very small compared to the original DIT model, which usually has tens hundreds of millions of parameters. We validated both strategies in Table [ However, we found that a simple global mixer yielded slightly better performance compared to the local mixer. For simplicity, we use global mixer in our experiments.
The Number of Experts and Basis In this part, we also explored the impact of different numbers of basis and expert models on the algorithm's 's performance. It is evident that indefinitely increasing the number of expert models does not continuously enhance performance. There are three reasons for this: 1) First, the total capacity of the basis models is limited. With 4 basis models, is difficult to craft 1,000 varied experts for denoising: 2) As illustrated in Figure 4a some denoising steps share similar mixing coefficients, which limits the maximal number of effective experts in our method; 3) Due to the fact that only one expert is activated in forward and backward passes, the gradient w.r.t the mixing coefficients will sparse. With large coefficient table, the sparsity level gradients will be higher, which introduces difficulty in optimization. However, compared to training the same K expert models independently, our algorithm's 's advantage lies its ability to learnably allocate the model's 's capacity through mixing coefficients. This allocation is done through a soft partition based on
8

---

Ablation
#Params (M)
#GBlops per Step
IS-10K
FID-I0K
Precision-I0K
Recall-10K
Mixer Type
DiT-S/2 2M
32,96
6.07
53,99
32.54
0.5691
0.6150
Onchot Mixer
4x.32.96
6.07
56.56
32.00
0.5726
0.5967
Raw Mixer
4x32.96
6.07
$7.39
31.46
0.5688
0.6122
Softmax Mixer
4x32.96
6.07
57.93
31.11
0,5702
0.6071
Global Local Mixer
DiT-S/2 2M
4x.22.96
6.07
53.99
32.54
0.5691
0.61.50
Global Mixer
4x32.96
6.07
57.93
31.11
0.5702
0.6071
Local Mixer
4x:2.99
6.07
57.65
31.21
The Number of Basis and Expert Models
DIT-S/2
4x,2,96
6.07
53.99
32.54
0.5691
0.6150
Remix-S.444
0.5666
Remix-S- S-4-8
4x32.96
6.07
57.15
31.60
0.5743
0.6109
Remix-$-4-20
4x32.96
6.07
57.93
31.11
0.5702
0.6071
Remix- Remix-S-4-100 S-4-100
4x322.96
6.07
31.48
Remix S 5-4-1000
1000
4x32.96
6.07
49.66
35.60
0.5571
0.6039
Remix-S-2-20
2x32.96
6.07
56.78
31.67
0,5703
0,6080
8x32.96
6.07
57.12
32.05
0,5642
Table 2: Ablation with 10K fintuning over a pre-trained DiT-S/2. For efficiency, we compute the FID-1OK with 100 sampling steps.
Model
Steps per See
GPU Mem. (MiB)
Latency (Mixing)
Latency (pre-computed)
DiT-B/2 29
2.93
13.152 16,942
15.77
15.77 ms
Demiix-/2-440
2.23
18.65 ms
15.78 ms
BB2ix-B/224-88
16,926
15.82 ms
DIT-S/2 29
3.40
9,758
5.66 ms
5.66 ms
Remix-S/2-4-23
3.24
10.578
7.16
5.71
Remix- S/2-4-8 29
3.30
10,558
7.04 ms
5.69 ms
Table Training and inference efficiency Remix-DiT, For inference, we can craft expert models by runtime mixing or pre-computing all experts, which is usually more efficient.
weighted loss, rather than hard partition on individual models. This allows to roughly choose an appropriate N which is slightly larger than the optimal N and the learnable coefficients can "merge" some experts if necessary by generating the same coefficients. For example, shown Figure 4al we can find highly similar coefficients between the timestep intervals [900, 950] and [950, 1000), which means we don't t need allocate different experts in these timesteps. In our experiments, we craft 20 experts with basis models.
Inference and training Efficiency. In Table 3 we present the training and inference efficiency of the proposed Remix-DiT. We found that the training speed slightly decreased from 2.93 2.29 due to the mixing of four basis models. Additionally, our method requires extra GPU memory to store the basis models. During inference, estimated the latency of compact models that craft experts at runtime and a standard multi-expert denoising approach, where all experts are precomputed. Post-training, our method achieve efficiency comparable standard DIT.
compares the generated images of the proposed RemixDiT-B and DYT-B. can observed that the shape of objects can be improved using the proposed methods. This is as expected since we allocate more model capacity to the early and intermediate stages of denoising, as illustrated in Figure which mainly contributes the image contents rather than details.
Visualization of Generated Images. Table 5
6 Limitations and Broader Impact
In this work, we introduce a learnable coefficient, implemented as simple embedding layer for mixing. However, due to the limitations of the deep learning framework, we can only create one expert per forward and backward pass. This leads sparse gradients in the embedding layers, where coefficients without gradients can only updated with momentum rather than accurate gradients. One solution to alleviate this issue is distributed training, where processes craft different expert

---

RemixDiT-B
DIT-B
Figure 5: Visualization of generated samples from DiT-B and Remix-DiT-B.
models. Despite this, the challenge remains significant when training a Remix- DiT with a large number of experts, such as 1000. However, according to the visualization of learned coefficients, we find that 1000 experts may not be necessary since many adjacent timesteps share similar model. In addition, this work will not introduce negative societal impact.
Conclusion
In this work, introduce multi- method to enhance the quality of transformer-baie basedd diffusion models while maintaining an acceptable inference overhead. The core contribution lies the ability to craft a large number of experts from few basis models, thereby significantly reducing the training effort. Besides, with our method, we don't 't have to accurately estimate the optimal number of required experts, since the learnable coefficients will adaptive merge experts if necessary, which brings huge flexibility to the practice.
Acknowledgment
This project is supported by the Ministry of Education, Singapore, under its Academic Research Fund Tier 2 (Award Number:
10

---

References
[1] Yogesh Balaji, Seungjun Nah, Xun Huang, Arash Vahdat, Jiaming Song, Qinsheng Zhang, Karsten Kreis, Miika Aittala, Timo Aila, Samuli Laine, ediff-i: Text-to-image diffusion models with an ensemble of expert denoisers. arXiv preprint arXiv:2211.01222 2022.
[2] Fan Bao, Shen Nic, Kaiwen Xue, Yue Cao, Chongxuan Li, Hang Su, and Jun Zhu. All are worth words: A vit backbone for diffusion models. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 22669-22679, 2023.
[3] Tim Brooks, Bill Peebles, Connor Holmes, Will DePue, Yufei Guo, Jing, David Schnurr, Joe Taylor. Troy Luhman, Eric Luhman, Clarence Ng, Ricky Wang, and Aditya Ramesh, Vidco generation models as world simulators. 2024.
[4] He Cao, Jianan Wang, Tianhe Ren, Xianbiao Qi, Yihao Chen, Yuan Yao, and Lei Zhang. Exploring vision transformers diffussio learners. arXiv preprint 21X11:211..1377,, 2022.
Shoufa Chen, Mengmeng Xu, Jiawei Ren, Yuren Cong. Sen Hc, Yanping Xie, Animesh Sinha, Ping Luo, Tao Xiang, and Maan- Perez-Rua. Gentron: Delving deep into diffusion transformers for image and 2023.
vidco generation. arXiv preprint arXiv:
[6] Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A large-scale hicrarchical image database. In 2009 IEEE conference computer vision and pattern recognition, pages 248-255. lece, 2009.
[7] Prafulla Dhariwal and Alexander Nichol. Diffusion models beat gans on image synthesis. Advances in neural information processing systems, 34:8780- 8794,, 2021.
[8] Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner. Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al. An image is worth 16x16 words:
Transformers for image recognition scale. arXiv preprint arXiv:2010.11 11909, 2020.
[9] Gongfan Fang. Xinyin Ma, and Xinchao Wang. Structural pruning for diffusion models. Advances in Neural Information Processing Systems, 2023.
[10] Shanghua Gao, Pan Zhou, Ming-Ming Cheng, and Shuicheng Yan. Masked diffusion transformer is a strong image synthesizer. In Proceedings of the IEEE/CVF International Conference Computer Vision, pages 23 1316-233173 173, 2023.
[11] Shuyang Gu, Dong Chen, Jianmin Bao, Fang Wen, Bo Zhang, Dongdong Chen, Lu Yuan, and Baining Guo. Vector quantized diffusion model for text-to-image synthesis. Proceedings of the IEEF/CVF Conference on Computer Vision and Pattern Recognition, pages 10696-10706, 2022.
[12] Kaiming He, Xinlei Chen, Saining Xie, Yanghao Li, Piotr Dollár, and Ross Girshick. Masked autoencoders scalable vision learners. In Proceedings the IEEE/CVF conference computer vision and pottern recognition, pages 16000-16009, 2022.
[13] Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in neural information processing systems, 33:6840 6851, 2020.
[14] Jonathan Ho, Ajay Jain, and Pieter Abheel, Denoising diffusion probabilistic models. Advances in neural information procesingg systems, 23:6840-68511 2020.
[15] Xin Jing, Yi Chang, Zijiang Yang, Jiangjian Xie, Andreas Triantafyllopoulos, and Bjoern W Schuller. U-dit tts: U-diffusion vision transformer for text-to- text-to-speech. In Speech Communication; 15th ITG Conference, pages 56 60. VDE, 2023.
[16] Yunsung Lee, Jin Young Kim, Hyojun Go, Mycongho Jeong, Shinhyeok Oh, and Seungtaek Choi. Multi- architecture multi-expert diffusion models. In Proceedings of the AAAI Conference on Artificial Intelligence. volume 38, pages 13427-13436, 2024.
[17] Elad Levi, Eli Brosh, Mykola Mykhailych, and Meir Perez. Dit: Conditioned layout generation with joint discrete-comtinuous diffusion layout transformer. In Proceedings of the IEEE/CVE International Conference on Computer Vision, pages 2106-2115, 2023.
[18] Enshu Liu, Xuefei Ning, Zinan Lin, Huazhong Yang. and Yu Wang. Oms-dpm: Optimizing the model schedule for diffusion probabilistic models. arXiv preprint arXiv:2306. 08860, 2023.
11

---

[19] Huadai Liu, Rongjie Huang. Xuan Lin. Wenqiang Xu, Maozong Zheng. Hong Chen. Jinzheng He, and Zhou Zhao. Vit-Its: visual lext-to-speech with scalable diffusion transformer. arXiv preprini 2.Kv:23005.12066 2023.
[20] Haoyu Lu, Guoxing Yang, Nanyi Fei, Yuqi Huo, Zhiwu La, Ping Luo, and Mingyu Ding. Vdt: General- purpose video diffusion transformers via mask modeling. In The Twelfth International Conference on Learning Representations, 2023.
[21] Xin Ma, Yaohui Wang, Gengyun Jia, Xinyuan Chen, Ziwei Liu, Yuan-Fang Li, Cunjian Chen, and Yu Qiao. Latte: Latent diffusion transformer for video generation. arXiv preprint arXiv:2401.0338 100048, 2024.
[22] Xinyin Ma, Gongfan Fang, and Xinchao Wang. Deepeache: Accelerating diffusion models for free. arXiv preprint arXiv:23/2.008. 2023.
[23] Xinyin Ma, Gongfan Fang, and Xinchao Wang. Deepeache: Accelerating diffusion models for free. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 15762- 15772, 2024.
[24] Willi Menapace, Aliaksandr Siarohin, Ivan Skorokhodov, Ekaterina Deyneka, Tsai-Shien Chen, Anil Kag, Yuwei Fang. Aleksei Stoliar, Elisa Ricci, Jian Ren, et al. Snap video: Scaled spatiotemporal transformers for text-to-video synthesis. arXiv preprint 24K.:22402.1497,, 2024.
[25] Shentong Mo. Enze Xie, Ruihang Chu, Lanqing Hong, Matthias Niessner, and Zhenguo Li. Dit-3d: Exploring plain diffusion transformers for 3d shape generation. Advances in Neural Information Processing Systems, 36, 2024.
[26] Alexander Quinn Nichol and Prafulla Dhariwal. Improved denoising diffusson probabilistic models. In International conference machine learning, pages 8162-8171. PMLR, 2021.
[27] Zizheng Pan, Bohan Zhuang, De-An Huang, Weili Nie, Zhiding Yu, Chaowei Xiao, Jianfei Cai, and Anima Anandkumar. T-stitch: Accelerating sampling in pre-trained diffusion models with trajectory stitching. arXiv preprint arXiv:2402. 14167, 2024.
[28] Byeongjun Park, Sangmin Woo, Hyojun Go, YnnYoung Kim, and Changick Kim. Denoising task routing for diffusion models. arXiv preprint arXiv:2310.0713 2023.
[29] William Peebles and Saining Xie. Scalable diffusion models with transformers. In Proceedings of the JEEE/CVF International Conference on Computer Vision, pages 4 4195-4205, 2023.
[30] Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Bjöm Ommer. High-resolution image synthesis with latent diffusion models. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pages 10684 10695, 2022.
[31] Olaf Ronneberger, Philipp Fischer, and Thomas Brox. U-net: Convolutional networks for biomedical image segmentation. In Medicul image computing and computer-axxisted intervention-MICCAl 2015: 18th international conference, Munich, Germany, October 5-9, 2015, proceedings, part III 18. pages 234-241. Springer, 2015.
[32] Jiaming Song, Chenlin Meng, and Stefano Ermon. Denoising diffusion implicit models. arXiv preprint axX1v20000.02502, 2020.
[33] Yang Song and Stefano Ermon. Generative modeling by estimating gradients of the data distribution. Advances in neural information processing systems, 32, 2019.
[34] Yang Song, Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-hased generative modeling through stochastic differential equations. arXiv preprint 13345 2020.
[35] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention all you need. Advances neural information processing systems, 30, 2017.
[36] Xingyi Yang, Jingwen Yc, and Xinchao Wang. Factorizing knowledge in neural networks. In European Conference on Computer Vision, pages 73-91. Springer, 2022.
[37] Xingyi Yang. Daquan Zhou, Jiashi Feng, and Xinchao Wang. Diffusion probabilistic model made slim. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 22552- 22562, 2023.
12

---

[38] Xingyi Yang, Daquan Zhou, Songhua Liu, Jingwen Ye, and Xinchao Wang. Deep model reassembly. Advances in neural information processing systems, 35:25739-25753, 2022.
[39] Xiulong Yang, Sheng-Min Shih, Yinlin Fu, Xiaoting Zhao, and Shihao Ji. Your vit is secretly hybrid disciiminative-eenerttiee diffusion model. arXiv preprint arXiv:2208. S00.22208007911, 2022.
[40] Hongkai Zheng, Weili Nie, Arash Vahdat, and Anima Anandkumar. Fast training of diffusion models with masked transformers. arXiv preprint arXiv:2306.09305 2023.
13

---

1. Claims
Question:
Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification:
This submission introduced new method for multi-expert denoising.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.
2. Limitations
Question:
Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification:
limitation section is included.
Guidelines:
The answer NA means that the paper has limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g.. independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what implications would be.
The authors should reflect the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and faimesss
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?
Answer: [NA]
Justification:
This work does not involve theoretical results.
Guidelines:
14

---

The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should properly referenced.
4. Experimental Result Reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided not)?
Question:
Answer: [Yes]
Justification:
We provide hyper-paramene the appendix.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is specific model and empirical evaluation, may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution is primarily a new algorithm, the paper should make it clear how reproduce that algorithm. (b) If the contribution is primarily
model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome describe the particular way they provide for reproducibility, In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
Justification: Code is available in the supplemental material.
15

---

Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://nips.co cc/
for more details.
While encourage the release of code and data, we understand that this might not be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https: for more details.
/ //nips wwtpcccppppcccccccccccc.ccccccle wwwpollccgggddee//eeeeoeee wwwooddeumisssooonnoioo
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if' applicable).
Providing much information possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code permitted.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer:
[Yes]
Justification: Training details summarized in the paper and appendix. Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to a level of detail that is necessary appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, or as supplemental material.
7. Experiment Statistical Significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [No]
Justification: This submission does not include error bars.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call a library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
should be clear whether the error bar is the standard deviation or the standard error of the mean.
It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
16

---

For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out range (e.g. negative enror rates).
If error bars reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
The number of GPUs, training efficiency and memory are reported in the paper.
Justification:
Guidelines:
The answer NA means that the paper does include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary failed experiments that didn't make into the paper).
9. Code Of Ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics
Answer: [Yes]
Justification:
The research conducted with the NeurIPS Code of Ethics.
Guidelines:
The answer NA means that the authors have not reviewed the NeurlPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g.. if there special consid- eration due to laws or regulations in their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [Yes]
Justification:
We discuss the socictal impacts in the main paper.
Guidelines:
The answer NA means that there no societal impact of the work performed.
If the authors answer NA or No. they should explain why their work has no socictal impact why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious unintended
(e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinfornation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
17

---

The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional unintentional) misuse of the technology.
there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: the paper poses no such risks.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper. properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification: All papers and assets properly credited.
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (c.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website). the copyright and terms of service of that source should provided.
assets are released, the license, copyright information, and terms use in the package should be provided. For popular datasets, haperswwthcode cwm/ datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided. If this information is not available online, the authors are encouraged to reach out to the asset's 's creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification: the paper does not release new assets.
18

---

Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the datasct/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create anonymized URL or include an anonymized zip file.
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification: the paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: the paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
19


Reviews:
[{"review_id": "ae5rwXtJ07", "reviewer": "Reviewer_2dkC", "summary": "To improve the generation quality of diffusion transformers, Remix-DiT proposes to enhance output quality at a lower cost and aims to create N diffusion experts for different denoising timesteps without the need for expensive training of N independent models. Remix-DiT achieves this by employing K basis models (where K < N) and using learnable mixing coefficients to adaptively craft expert models. This approach offers two main advantages: although the total model size increases, the model produced by the mixing operation shares the same architecture as a plain model, maintaining efficiency comparable to a standard diffusion transformer. Additionally, the learnable mixing adaptively allocates model capacity across timesteps, effectively improving generation quality. Experiments on the ImageNet dataset show that Remix-DiT achieves promising results compared to standard diffusion transformers and other multiple-expert methods.", "questions": "1. Performance Comparison Between Multi-Expert and Single Larger Models: Is it possible for the multi-expert small models to outperform a single, larger model? To fully validate the potential of the multi-expert approach, it is crucial to provide a thorough performance comparison. This should include quantitative metrics and benchmarks that demonstrate the advantages, if any, of using multiple experts over a single larger model in terms of both output quality and computational efficiency.\n\n2. Scalability and Efficiency of Increasing the Number of Experts: If the number of experts is increased for the same basis models, how easily can the system be scaled, and does this lead to more efficient training? It would be important to discuss the scalability of the multi-expert framework, including any potential challenges or limitations in transferring the model to a larger number of experts. Additionally, insights into how the efficiency of training might be affected by increasing the number of experts would be valuable.", "limitations": "Please refer to the weakness and question part.", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "Novelty: Model mixers for efficient multi-expert diffusion model training is innovative and unique.\n\nSignificance: Addressing the challenge of efficient training of multi-expert diffusion transformers is significant in the field of diffusion models.\n\nMethodology: The proposed algorithm is well-formulated and clearly explained.\n\nResults: Experimental results demonstrate promising improvements over existing methods such as DiT.", "weaknesses": "1. Lack of Visualization Results: The paper does not include any visualization results. Providing visual examples of generated outputs is crucial for qualitatively evaluating the effectiveness of the proposed method.\n\n2. Insufficient Motivation for Multi-Expert Training: The rationale behind adopting a multi-expert training approach is not fully well-motivated, particularly in the context of quantitative comparisons. A more detailed explanation of why multi-expert training is beneficial and how it compares quantitatively to other methods would strengthen the argument. Clarifying the advantages and potential trade-offs in performance and efficiency would provide a more compelling case for this approach.\n\n3. High Training Cost: The training cost associated with the proposed method is substantial. It would be beneficial to provide a thorough analysis of the computational resources, time, and energy required for training compared to other existing methods. Discussing potential ways to mitigate these costs or offering insights into why the increased training cost is justified by the performance gains would add valuable context for evaluating the practicality of the method.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "> **Q1: Lack of Visualization Results: The paper does not include any visualization results. Providing visual examples of generated outputs is crucial for qualitatively evaluating the effectiveness of the proposed method.**\n\nThanks for the suggestion. We supplement visualization results in the attached PDF, which compares the proposed RemixDiT-B to a DiT-B. It can be observed that the shape of objects can be improved using the proposed methods. This is as expected since RemixDiT allocates more model capacity to the early and intermediate stages of denoising as illustrated in Figure 4 of the paper, which mainly contributes to the global shape rather than detailed patterns. We will incorporate visualization results into the experiment sections and polish the draft following the advice.\n\n> **Q2: Insufficient Motivation for Multi-Expert Training: The rationale behind adopting a multi-expert training approach is not fully well-motivated, particularly in the context of quantitative comparisons. A more detailed explanation of why multi-expert training is beneficial and how it compares quantitatively to other methods would strengthen the argument. Clarifying the advantages and potential trade-offs in performance and efficiency would provide a more compelling case for this approach.**\n\nThe key motivation for the multi-expert approach lies in the limitations of a single model's capacity for the 1000-step denoising tasks. Figure 4 (c) provides a quantitative perspective on this issue by illustrating the performance of experts in both their specialized and non-specialized regions. It can be observed that all experts achieve low losses within their designated time slots, but in the \"non-professional\" timesteps, an expert may yield relatively large losses. This disparity is due to the limited model capacity when faced with the extensive number of timesteps. To address this, we propose utilizing multi-expert strategies to enhance diffusion models without significantly increasing the number of parameters. We will incorporate the reviewer's advice to clarify the motivation in the revised version.\n\nBased on this, the main advantage of RemixDiT lies in its ability to provide a flexible trade-off between the number of experts and the total training costs through the mixing of a few basis parameters.  This allows training $N$ experts for different regions by optimizing $K (K<N)$ basis models.\n\n> **Q3: High Training Cost: The training cost associated with the proposed method is substantial. It would be beneficial to provide a thorough analysis of the computational resources, time, and energy required for training compared to other existing methods. Discussing potential ways to mitigate these costs or offering insights into why the increased training cost is justified by the performance gains would add valuable context for evaluating the practicality of the method.**\n\nThe training cost of RemixDiT can be found in Table 3, which measures the training throughput and memory usage. To obtain 20 experts for different regions, our method incurs 5\\% ~ 20\\% additional costs compared to a standard DiT. Notably, the training cost does not linearly increase with the number of experts or basis models, as the additional computation primarily arises from the lightweight mixing operation, which will not introduce too much forwarding or backwarding efforts during training. We will revise lines 298-303 of the paper to include a more comprehensive discussion on training and inference efficiency.\n\n> **Q4: Performance Comparison Between Multi-Expert and Single Larger Models: Is it possible for the multi-expert small models to outperform a single, larger model? To fully validate the potential of the multi-expert approach, it is crucial to provide a thorough performance comparison. This should include quantitative metrics and benchmarks that demonstrate the advantages, if any, of using multiple experts over a single larger model in terms of both output quality and computational efficiency.**\n\nBased on our results in the submission, it is challenging for RemixDiT-B with four basis models (FID=9.02) to outperform a DiT-L that is four times larger (FID=3.73), with around 1M training steps. This limitation arises because the number of effective parameters in RemixDiT-B remains equivalent to that of a single DiT-B. Currently, our proposed RemixDiT can ensure better performance under a comparable inference cost. We appreciate the suggestion to explore the upper bound of this method and are working on training DiT baselines of different sizes. \n\n> **Q5: Scalability and Efficiency of Increasing the Number of Experts: If the number of experts is increased for the same basis models, how easily can the system be scaled, and does this lead to more efficient training? It would be important to discuss the scalability of the multi-expert framework, including any potential challenges or limitations in transferring the model to a larger number of experts. Additionally, insights into how the efficiency of training might be affected by increasing the number of experts would be valuable.**\n\nWe agree that scalability is important for this method. As shown in Table 2, our key observation is that increasing the number of experts is not always beneficial. There are three reasons for this: 1) As illustrated in Figure 4 (a,b), some denoising steps share similar mixing coefficients, which limits the maximal number of effective experts in our method; 2) The total capacity is also bounded by the number of basis models; 3) As discussed in the limitations (Line 310), optimizing a large number of experts can lead to sparse gradients in the mixing coefficients, making optimization more difficult as the number of experts increases.\n\nTherefore, we chose to use 20 experts in our experiments, which we found to be a good balance between performance and optimization difficulty.", "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": "1tbPXEydjf", "reviewer": "Reviewer_cagX", "summary": "The paper proposes Remix-DiT, a model architecture designed to enhance the capacity of a standard DiT model without significantly increasing inference costs. This is accomplished by training mixing coefficients to adaptively fuse multiple DiT models and developing specialized experts for multi-expert denosing. A key advantage highlighted in this paper is that Remix-DiT achieves better generation quality while maintaining inference speed comparable to that of a standard DiT. Experimental results on ImageNet-256 demonstrate favorable outcomes compared to baseline methods.", "questions": "Please refer to the weaknesses.", "limitations": "This paper discusses limitations such as sparse gradients and the training difficulty associated with a large number of experts.", "rating": 7, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 4, "strengths": "1.\tThe visualization results in Figure 4 are very interesting. It seems that the model has a certain preference in allocating the capacity of basis models, with clear segmentation across the timesteps. Additionally, a high coefficient is observed at early timesteps, such as 0-150. Does this imply that those steps are more challenging for the diffusion model to learn?\n2.\tThe idea of mixing multiple basis models is clear and easy to implement. It does not requires the expensive training of independent experts for different steps.", "weaknesses": "1.\tUsing multiple base models may introduce more training costs. However, in Table 3, the GPU memory usage only slightly increases from 13G to 16G for DiT-B. Can the authors provide more details about the reason? Will Remix-DiT introduce a substantial backward and forward footprint?\n2.\tThis method utilizes the pre-trained model as the initialization. This might make the mixed experts always the same after mixing since they are working on the same basis model initially. Will this be a problem?\n3.\tWhy does the proposed method outperform naively training independent experts? In this method, the experts are crafted by mixing, which should theoretically be upper bounded by the naïve method mentioned above.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "> **Q1: The visualization results in Figure 4 are very interesting. It seems that the model has a certain preference in allocating the capacity of basis models, with clear segmentation across the timesteps. Additionally, a high coefficient is observed at early timesteps, such as 0-150. Does this imply that those steps are more challenging for the diffusion model to learn?**\n\nThank you for the question. Allocating a high score to a certain slot indicates that this step is distinct from others and might be challenging during training. In the diffusion process, the shape of objects emerges quickly, followed by the incorporation of finer details in the later stages. The algorithm adaptively allocates model capacity to these specific steps, ensuring that critical stages receive the necessary resources for accurate denoising.\n\n> **Q2: This method utilizes the pre-trained model as the initialization. This might make the mixed experts always the same after mixing since they are working on the same basis model initially. Will this be a problem?**\n\nThanks for the comment, this issue indeed exists at the early training steps. Therefore, we introduce the prior coefficients in Line 188 of the paper to force a prior allocation of model capacity. For example, we initialize the coefficients for the first timestep region with [1,0,0,0], this forces the first basis model to learn from early steps. This trick will lead to diverge basis models, which is important for our method.\n\n> **Q3: Using multiple base models may introduce more training costs. However, in Table 3, the GPU memory usage only slightly increases from 13G to 16G for DiT-B. Can the authors provide more details about the reason? Will Remix-DiT introduce a substantial backward and forward footprint?**\n\nThank you for the comment; this issue indeed exists in the early training steps. To address this, we introduce prior coefficients, as mentioned in Line 188 of the paper, to enforce a prior allocation of model capacity. For example, we initialize the coefficients for the first timestep region with [1,0,0,0], which forces the first basis model to learn from the early steps. This approach helps diverge the basis models, which is crucial for the effectiveness of our method.\n\n> **Q4: Why does the proposed method outperform naively training independent experts? In this method, the experts are crafted by mixing, which should theoretically be upper bounded by the naïve method mentioned above.**\n\nThanks for the invaluable question. Our method is more parameter efficient than training $N$ independent experts. First, our method can support training 20 experts simultaneously by optimizing 4 basis models. Under the same budget, our method can be fully optimized while the naive expert training is still underfitted. In addition, our method is able to adaptively allocate model capacity to different timesteps, which improves the utilization of network parameters.", "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": "EZld1XOId6", "reviewer": "Reviewer_7vev", "summary": "The paper introduces Remix-DiT, a modification to the diffusion transformer architecture that incorporates the multi-expert denoiser framework during both training and inference. Unlike traditional multi-expert methods that train $N$ separate individual experts independently for each time interval, Remix-DiT employs $K$ base models combined with $N$ mixing coefficients to dynamically compute time-specific experts. This approach enhances efficiency and leverages task similarities between adjacent intervals more effectively. Experiments on ImageNet demonstrate that Remix-DiT improves the performance of DiT across various model sizes.", "questions": "1. The EDM paper [1] suggests that for denoising, only the middle noise levels are important, while this paper suggests that the noise levels towards 0 are more crucial. Do you have an intuition on the difference between these two conclusions?\n\n2. Is the performance of Remix-DiT more sensitive to the number of sampling steps compared to a normal DiT? In other words, how do the experts perform when using a deterministic sampler with low NFEs (<50)?\n\n3. Can you also visualize some examples generated by DiT and Remix-DiT? While the metrics are valuable, a qualitative evaluation is interesting as well.\n\n[1] Karras T, Aittala M, Aila T, Laine S. Elucidating the design space of diffusion-based generative models. Advances in neural information processing systems. 2022 Dec 6;35:26565-77.", "limitations": "The authors have mentioned this in the paper.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The paper is well-motivated and represents a valuable step towards integrating the multi-expert denoising framework into standard diffusion models. \n\n- The main idea of the paper (using global mixers to compute the final experts) is novel and interesting to me in this context.\n\n- The method is simple and effective, making it more suitable for practical use cases. \n\n- The experiments are well-designed, and the ablations clearly illustrate the impact of various aspects of Remix-DiT. \n\n- The paper is well-written and generally easy to understand.", "weaknesses": "- While the authors show the benefits of Remix-DiT on finetuning a pretrained DiT model, it would be interesting to see its effect when training all components from scratch. If the compute budget allows, I suggest that the authors also add this experiment for better insights into what happens if one uses the remixing scheme from the beginning of training (perhaps after a small warmup)\n\n- The performance gain seems to diminish as the size of the base model increases. Hence, a more detailed discussion on this issue is needed for the final version. For example, the performance gain is almost 30% for DiT-S, while it drops to only 15% for DiT-L.\n\n\n**Minor comments:**\n\nPlease fix the following issues in terms of writing in your draft:\n- L114 \"refer to\" -> \"refers to\"\n- L144 -> citation is missing\n- L215 -> I assume 100M steps should be 100K steps\n- L290 -> it seems that it should be written as N experts because K is the number of base models\n- L295 -> \"can found\" should be \"can find\"\n\nPlease also cite GigaGAN [1] as the mixing part of the paper is related to their method of mixing different convolution kernels during training.\n\n[1] Kang M, Zhu JY, Zhang R, Park J, Shechtman E, Paris S, Park T. Scaling up gans for text-to-image synthesis. InProceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition 2023 (pp. 10124-10134).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "> **Q1: While the authors show the benefits of Remix-DiT on finetuning a pretrained DiT model, it would be interesting to see its effect when training all components from scratch. If the compute budget allows, I suggest that the authors also add this experiment for better insights into what happens if one uses the remixing scheme from the beginning of training (perhaps after a small warmup)**\n\nThanks for the invaluable suggestion. We use pre-trained models to avoid training basis models independently since they share some common ability across timesteps. We agree with the reviewer that it is valuable to explore scratch training with our method. Due to the time limits of this short rebuttal period, we can only provide 400K results for scratch training. It can observed that our method still achieves competitive FID in the setting of scratch training.\n| Model        | FID @ 400K |\n|--------------|------------|\n| DiT-S        | 43.88      |\n| RemixDiT-S   | 36.68      |\n\n> **Q2: The performance gain seems to diminish as the size of the base model increases. Hence, a more detailed discussion on this issue is needed for the final version. For example, the performance gain is almost 30\\% for DiT-S, while it drops to only 15% for DiT-L.**\n\nThank you for the insightful comment! The key motivation behind RemixDiT is that the capacity of a single model is insufficient for multi-step denoising tasks. Therefore, with a smaller model like DiT-S, our method can offer more significant benefits compared to relatively larger models like DiT-L. Additionally, training large models can be more challenging. We will include a more detailed discussion about this phenomenon in the revised version, following the reviewer's advice.\n\n> **Q3: Please fix the following issues in terms of writing in your draft.**\n\nWe will revise the draft accordingly.\n\n> **Q4: Please also cite GigaGAN [1] as the mixing part of the paper is related to their method of mixing different convolution kernels during training.**\n\nIt's indeed a highly related paper from the perspective of techniques. We will cite it properly.\n\n> **Q5: The EDM paper [1] suggests that for denoising, only the middle noise levels are important, while this paper suggests that the noise levels towards 0 are more crucial. Do you have an intuition on the difference between these two conclusions?**\n\nA large coefficient indicates that this slot is more unique and challenging, necessitating greater capacity to handle it. The segmented patterns in Figure 4(a) demonstrate the importance of both early and intermediate steps. Our method allocates more model capacity to the early stages (0-50) and the intermediate stages (150-500). So, the results is, to some extend, aligned with the observation in the EDM paper. \n\n> **Q6: Is the performance of Remix-DiT more sensitive to the number of sampling steps compared to a normal DiT? In other words, how do the experts perform when using a deterministic sampler with low NFEs (<50)?**\n\nIn Table 2 of the submission, we present results for 100-step sampling, where our method continues to show positive results compared to the baseline DiT. In response to your question, we further reduced the number of steps to 25. The following table demonstrates that even with this lower number of NFEs, our method, RemixDiT-S, still outperforms the baseline DiT-S.\n| Model        | FID-10K (Steps=25, cfg=1.5) |\n|--------------|-----------------------------|\n| DiT-S        | 49.82                       |\n| RemixDiT-S   | 44.75                       |\n\n> **Q7: Can you also visualize some examples generated by DiT and Remix-DiT? While the metrics are valuable, a qualitative evaluation is interesting as well.**\n\nThanks for the suggestion. We supplement visualization results within the attached PDF, which compares the proposed RemixDiT-B to a DiT-B. It can be observed that the shape of objects can be improved using the proposed methods. This is as expected since we allocate model capacity to the early and intermediate stages of denoising, which mainly contributes to the image contents rather than details. We will incorporate visualization results into the experiment sections.", "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": "uRlQWBcbUk", "reviewer": "Reviewer_KY1f", "summary": "This paper proposes Remix-DiT, which creates multiple experts by mixing fewer basis diffusion transformers, allowing each expert to specialize in the denoising task for corresponding timestep intervals. It achieves performance improvements by having each expert responsible for a larger number of timestep intervals with fewer total trainable parameters than previous multi-expert methods. Also, the paper analyzes the coefficients of how much each expert uses bases, demonstrating the denoising task similarity for adjacent timesteps, as well as the use of specialized bases for lower timesteps.", "questions": "*  Is the Exponential Moving Average (EMA) model used to further train a pre-trained diffusion model?\n\n* It would be better that the authors provide an affinity matrix between 20 timestep clusters based on the learned mixing coefficients. I think the affinity matrix could explain the similarity between denoising tasks.", "limitations": "N/A", "rating": 6, "confidence": 5, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "* The paper is structured well, making it easy to understand and follow.\n\n* The proposed mixing basis strategy is interesting as it achieves better performance with fewer parameters compared to existing multi-expert methods.\n\n* Ablation studies on mixing methods are comprehensive.", "weaknesses": "* **Lack of experiments.** The authors have to validate the performance of Remix-DiT by reporting comparisons with previous methodologies on the FFHQ or MS-COCO datasets. It would make the manuscript more solid if Remix-DiT achieves consistent performance improvements on multiple datasets.\n\n* **Lack of comparison.** There are two methods, DTR [1] and Switch-DiT [2], to address the multi-task learning aspect of diffusion training by designing distinct denoising paths for 1000 timesteps in a single model. These are more parameter-efficient methods where they use no additional parameters or 10%, respectively. The authors should analyze them with respect to Remix-DiT.\n\n[1] Park et al., Denoising Task Routing for Diffusion Models, ICLR 2024.\n\n[2] Park et al., Switch Diffusion Transformer: Synergizing Denoising Tasks with Sparse Mixture-of-Experts, ECCV 2024.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "> **Q1: Lack of experiments. The authors have to validate the performance of Remix-DiT by reporting comparisons with previous methodologies on the FFHQ or MS-COCO datasets. It would make the manuscript more solid if Remix-DiT achieves consistent performance improvements on multiple datasets.**\n\nThank you for the invaluable advice. We agree with the reviewer that incorporating more datasets is important to validate the robustness of RemixDiT and will add more datasets to polish this work. However, due to the limited rebuttal period, we are not able to present the results on large-scale datasets like COCO. We are currently working on FFHQ following DTR and should be able to update the results in the coming days once they are available.\n\n> **Q2: Lack of comparison. There are two methods, DTR [1] and Switch-DiT [2], to address the multi-task learning aspect of diffusion training by designing distinct denoising paths for 1000 timesteps in a single model. These are more parameter-efficient methods where they use no additional parameters or 10\\%, respectively. The authors should analyze them with respect to Remix-DiT.**\n\nThank you for the highly relevant pointers. We have supplemented the following results to compare RemixDiT to the mentioned baselines. We trained the proposed RemixDiT from scratch for 400K steps on ImageNet, which yielded a comparable FID to DTR. However, we also observed that the FID is higher than that of Switch-DiT-S, as training RemixDiT from scratch is less efficient, as discussed in Line 188 of our paper, where we have four basis models to build. Another reason for this might be the non-optimal hyperparameters for the scratch training. These baseline methods are very insightful and helpful. We will continue to improve this work to make our method more robust to different training configurations and will incorporate these results into the manuscript.\n\n| **Model**          | **Train Steps** | **FID** |\n|--------------------|-----------------|----------------|\n| DiT-S              | 400K            | 43.88          |\n| DTR-S [1]          | 400K            | 37.43          |\n| Switch-DiT-S [2]   | 400K            | 33.99          |\n| RemixDiT-S         | 400K            | 36.68          |\n\n[1] Park et al., Denoising Task Routing for Diffusion Models, ICLR 2024.  \n[2] Park et al., Switch Diffusion Transformer: Synergizing Denoising Tasks with Sparse Mixture-of-Experts, ECCV 2024.", "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:
We would like to extend our sincere gratitude to all the reviewers for their time, effort, and insightful feedback on our submission. In response to reviewers' questions, we included some visualization results in the attached PDF file to compare the RemixDiT-B to a standard DiT-B, where our method is able to improve the object shape by allocating more model capacity to early and intermediate timesteps.

Meta-review: The paper presents Remix-DiT, a multi-expert denoising method that enhances diffusion transformers by mixing basis models with learnable coefficients, achieving high-quality outputs without extensive training costs. Reviewers appreciated the novel approach and its potential to efficiently improve generation quality across different timesteps. 

While reviewers praised the paper’s practicality and clear presentation, they recommended additional evaluations on diverse datasets and comparisons with state-of-the-art methods like DTR and Switch-DiT. Concerns about high training costs and scalability were noted, along with suggestions to include more visualizations to better showcase the improvements. Addressing these points in the camera-ready would further solidify the paper’s contributions and impact in the field of diffusion models.
Decision: Accept (poster)