Remix-DiT: Mixing Diffusion Transformers for Multi-Expert Denoising

Authors: gongfan  fang
xinyin  ma
xinchao  wang
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-02 17:02:16

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.

Main paper: 
 introduction: 
 Transformer-based Diffusion models [39, 8, 29, 2] 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, 18, 28, 16, 27] , which employs multiple specialized diffusion models, each designed for distinct time intervals within the denoising process.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 [13, 32, 18] , requiring the network to make predictions at various noise levels. Previous work has shown that the tasks of diffusion models vary across different timesteps [13, 37, 27] . For instance, at higher noise levels, the model focuses more on low-frequency features, while at lower noise levels, the model emphasizes generating high-frequency details [37] . This variability inherently leads to a multi-task problem [16, 18] . However, due to the limited capacity of a single diffusion model, it is challenging to craft a 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 a specific subset of timesteps. Although this design introduces multiple models, only one model is activated at each timestep during the inference process, and therefore the overall computational complexity does not increase significantly.In a 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 an even-split strategy, which distributes the denoising process equally among a 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 [1] , 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, learnable strategy to craft multiple experts adaptively for denoising.Consider a denoising process of T steps. The objective of this paper is to construct N expert models, each handling an interval of length T 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 to 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 < N ) basis models, which can be used to craft N expert models efficiently for inference. To be exact, we construct K transformer-based diffusion models and introduce learnable mixing coefficients to fuse their parameters, which can be easily implemented as 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.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 a 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 in 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. 
related works: 
 Multiple Experts in Diffusion Models The multi-step denoising process inherent to the diffusion model [14, 33, 34] can be viewed as a multitasking problem [1, 18, 28, 16] . 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 [1] introduces a simple binary strategy to learn various experts and during inference, take the ensemble of experts for prediction. OMS-DPM [18] proposes a model scheduler, dynamically taking denoiser of different sizes from a model zoo for inference. MEME [16] 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 a novel multi-tasking strategy to learn a single model with scheduled task masks, activating specific sub-networks for different timesteps.Transformer-based Diffusion Models Transformer-based diffusion models have achieved impressive results on image generation [2, 29, 4, 10, 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 the U-Net [31, 22, 30] architecture for improved training efficiency and quality. Furthermore, recent works have expanded transformer-based diffusion models to video generation [20, 21, 24, 5, 3] , audio [19, 15] and 3D [25] which shows the power of transformers in modeling complicated data. However, training diffusion models remains inefficient, typically requiring millions of steps to produce satisfactory results. A series of works in the literature focus on the efficiency of diffusion-based transformers [40, 40, 23, 9] . MaskDiT [40] and MDTv2 [10] utilize masked transformers and an additional masked autoencoder task [12] to streamline training. UViT [2] 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 transformersQuestion: Does the paper discuss the limitations of the work performed by the authors? Answer: [Yes] Justification: A limitation section is included. Guidelines:• The answer NA means that the paper has no 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 a 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 the implications would be. • The authors should reflect on 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, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a 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 fairness. • While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a 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 important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations. • The authors should make sure to preserve anonymity (e.g., if there is a special consideration 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 societal impacts in the main paper. Guidelines:• The answer NA means that there is no societal impact of the work performed.• If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact. • Examples of negative societal impacts include potential malicious or unintended uses (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 a 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 disinformation. 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.Guidelines:• The answer NA means that the paper does not release new assets.• Researchers should communicate the details of the dataset/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 was obtained from people whose asset is used. • 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. 
preliminary: 
 Diffusion Probabilistic Models Diffusion Probabilistic Models (DPMs) train a denoiser network to revert a process of adding gradually increased noises [13] . Given an input x 0 , a T -step forward process is deployed to transform x 0 into latent x 1 , . . . x T , where the transformation is defined as q(x t |x t-1 ) = N (x t ; √ 1 -β t x t-1 , β t I) under a increased variance schedule β 1:T . To revert this process, a network q(x t-1 |x t ) is trained to predict x t-1 given x t . By defining α t = 1 -β t and ᾱt = t s=1 α s , the training objective is formalized asEQUATIONwhere ϵ θ is a trainable neural network and ϵ refer to the noises at timestep t. 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, 1, 37, 27] . Due to the limited capacity, a single model may struggle to fit all steps. Therefore, a more suitable yet natural approach is to use multiple models specifically learned for different timestep intervals [1] ."Multi-Expert Denoising To facilitate denoising with N expert models, the timesteps are typically divided into N intervals. Each model, denoted as ϵ θi , contains independent parameters θ i . For each model θ i , we minimize the training objective as formalized in Equation 1 on the corresponding EQUATION)Although the simple form of multi-expert denoising has been verified effective in previous works [1, 16] , the above objective still presents two major challenges: 1) The optimal number of experts and the best interval partition is unknown; 2) Training multiple independent experts can be expensive. In this work, we introduce a new approach, Remix-DiT to address the above issues. 
method: 
 Uniform Partition of Denoising Process Let's consider a multi-expert denoising problem with N experts for T timesteps. For simplicity, we assume that each expert's parameters θ i is a column vector with P trainable elements and the parameter matrix of all experts can be formalized asΘ N ×P = [θ 1 , θ 2 , • • • , θ N ] ⊤. This is natural since the parameters in a model can be flattened and concatenated into a vector. We consider a simple partition of timesteps, which equally divides the denoising process into N intervals, i.e., [0,T /N ] . . . [(N -1) • T /N, N • T /N ].The motivation behind the oracle partition is that adjacent steps share similar noise levels and training objectives. Following Equation 2, this leads to the learning objective to optimize each expert on their associated time intervals.EQUATIONThe 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 To address the problem of uniform partition, we leverage a 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β K×P = [β 1 , β 2 , . . . β K ] ⊤ .And the expert models can be crafted by mixing the basis parameters with certain coefficients. For each expert θ i , we associate it with a coefficients α i = [α i1 , α i2 , . . . α iK ] ⊤ and compute the mixed expert parameter through a weighted averaging θ i = k α ik β k . This leads to the coefficient matrix denoted as α N ×K = [α 1 , α 2 , . . . , α N ] ⊤ . Then, the expert's parameter matrix can be easily obtained through a simple matrix multiplication:EQUATIONWith Equation 4, we can freely craft different experts using different mixing coefficients. Now the problem lies in that how to learn a good coefficients α and basis models β.Algorithm 1 Mixed Linear (PyTorch-like Pseudo Code)1: function MIXEDLINEAR(x, w, coeffs, K) 2: // x: linear inputs 3:// w: linear parameters 4:// coeffs: mixing coefficients 5:// K: number of basis 6: w = w.view(K, w.shape [1] //K, -1) 7:w_ = (coeffs * w).sum(0) 8:return torch.nn.functional.linear(x, w_) 9: end function Architecture of Remix-DiT Let's further delve into the details of K basis models, denoted as β, alongside a coefficients matrix α. In this study, we adopt the DiT architecture proposed in [29] , which primarily consists of linear and normalization layers. Rather than constructing K distinct models, we propose a simple trick to construct a singular and extended DiT model to encapsulate all K basis models [36] . This is achieved by increasing the width of the linear layers by a factor of K, with a modified forward process to support the mixing of basis. An example of the extended linear layer is depicted in Algorithm 1. A notable distinction between the standard DiT and the proposed Remix-DiT is the integration of a weighted averaging step before the true computation. Despite the expansion in width by a 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.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, we employ a hierarchical sampling strategy whereby the expert θ i is selected first, followed by a random selection of timestep t within the interval[i × ⌊ T N ⌋, (i + 1) × ⌊ T N ⌋]. 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-DiT are delineated in Algorithm 2. The training protocol for Remix-DiT closely mirrors that of a standard DiT, with the exception that the sampled timesteps t 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 ∇ α i,k :EQUATION)where ⊙ is an element-wise product of two matrices. First, the update of basis parameter θi is 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 θi almost unchanged. In this case, the gradient to the coefficients will be constant. To address this issue, we introduce a prior regularization to force each basis θi to learn different objectives. Inspired by the oracle partition in multiple expert denoising, we craft prior coefficients α * 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 a cross-entropy between the prior coefficients and the learnable coefficients.EQUATIONAlgorithm 2 Remix-DiT Compute the mixing coefficients αi = Softmax(πi) 7:Obtain the i-th mixed expert θi = αiβ 8:Randomly sample the timestep t where γ is a hyperparameter to control the strength of regularization. In practice, we adopt an annealing strategy to linearly remove the regularization with γ → 0.∈ [i × ⌊ T N ⌋, (i + 1) × ⌊ T N ⌋] 9: Compute the loss Lt,i = ∥ϵ -ϵ θ i ( √ ᾱtx0 + √ 1 -ᾱtϵ, t)∥ 
experimental settings: 
 Network Architecture. 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-DiT-S is constructed by quickly combining multiple DiT-S models to form an expert model.Training Details. Since our mixed model is entirely consistent with the standard DiT, this method can be applied to pre-trained models. We first trained a standard DiT model on ImageNet and then initialized the basis model with the same pre-trained weights. We introduce prior as discussed in Equation 6 to the mixing coefficients to accelerate the learning of different basis models. In our experiments, we conducted 100 K fine-tuning on DiT-S/B/L models [29] , pre-trained for 2M/1M/1M steps correspondingly.In Table 1 , we present the results of fine-tuning standard DiT models using the Remix-DiT 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 successfully 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 baselines, where experts are trained independently. It is noteworthy that for Multi-expert baselines with eight experts, the total training budget is uniformly allocated to 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, RemixDiT 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 To further investigate the behavior of the experts learned by the algorithm, we visualized the mixing coefficients. As shown in Figure 4a , we conducted experiments on DiT-S. During the training of 20 expert 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-frequency detail features. In contrast, at late timesteps resembling random noise, the algorithm enabled and mixed multiple models to generate content rather than using only one expert. As will be illustrated in Figure 5 , mixing multiple models can improve the shape quality. Moreover, 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 8 expert models throughout the denoising process. Each mixed expert has a lower loss function within its respective timestep intervals. The farther from an expert model's current interval, the higher its loss. Notably, at step 0, we found that the ensemble-constructed model effectively reduces prediction loss.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 we explored three mixing methods: 1) Oracle Mixing: Basis models are manually assigned to different intervals as expert models, equivalent to 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 a real-valued coefficient without any constraints. 3) Softmax Mixer: This builds on the raw mixer by applying a softmax operation, ensuring that the mixed model is always a weighted average of the basis models. Table 2 shows the effectiveness of each mixer. We found that while oracle mixing improves model performance, it often fails to achieve optimal results due to manually designed partitions. The raw mixer, with its learnable coefficients, achieved better performance but included some impractical solutions, such as 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 to as the local mixer. The first strategy is simpler to train, as it requires only a small K × N coefficient matrix regardless of the network'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 to hundreds of millions of parameters. We validated both strategies in Table 2 . However, we found that a simple global mixer yielded slightly better performance compared to the local mixer. For simplicity, we use a global mixer in our experiments.In this part, we also explored the impact of different numbers of basis and expert models on the algorithm'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, it 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 be sparse. With a large coefficient table, the sparsity level of gradients will be higher, which introduces difficulty in optimization. However, compared to training the same K expert models independently, our algorithm's advantage lies in its ability to learnably allocate the model's capacity through mixing coefficients. This allocation is done through a soft partition based on weighted loss, rather than a hard partition on individual models. This allows us 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, as shown in Figure 4a , we can find highly similar coefficients between the timestep intervals [900, 950] and [950, 1000], which means we don't need to allocate different experts in these timesteps. In our experiments, we craft 20 experts with 4 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 to 2.29 due to the mixing of four basis models. Additionally, our method requires extra GPU memory to store the basis models. During inference, we 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 can achieve efficiency comparable to a standard DiT.Visualization of Generated Images. Table 5 compares the generated images of the proposed RemixDiT-B and 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 more model capacity to the early and intermediate stages of denoising, as illustrated in Figure 4 , which mainly contributes to the image contents rather than details. 
limitations and broader impact: 
 In this work, we introduce a learnable coefficient, implemented as a 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 to sparse gradients in the embedding layers, where coefficients without gradients can only be updated with momentum rather than accurate gradients. One solution to alleviate this issue is distributed training, where processes craft different expert 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 a similar model. In addition, this work will not introduce negative societal impact. 
conclusion: 
 In this work, we introduce a multi-expert method to enhance the quality of transformer-based diffusion models while maintaining an acceptable inference overhead. The core contribution lies in the ability to craft a large number of experts from a few basis models, thereby significantly reducing the training effort. Besides, with our method, we don'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: MOE-T2EP20122-0006). 
 
 Appendix:  


Decision: 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.

Reviews: review: paper_summary: paper_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.
strength_weakness: strength_weakness: 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.

2. 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.

3. 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.
questions: 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.

2. 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: limitations: Please refer to the weakness and question part.
paper_summary: paper_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.
strength_weakness: strength_weakness: 1.	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?
2.	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?
3.	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.
questions: questions: Please refer to the weaknesses.
limitations: limitations: This paper discusses limitations such as sparse gradients and the training difficulty associated with a large number of experts.
paper_summary: paper_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.
strength_weakness: strength_weakness: - 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)

- 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.


**Minor comments:**

Please fix the following issues in terms of writing in your draft:
- L114 "refer to" -> "refers to"
- L144 -> citation is missing
- L215 -> I assume 100M steps should be 100K steps
- L290 -> it seems that it should be written as N experts because K is the number of base models
- L295 -> "can found" should be "can find"

Please also cite GigaGAN [1] as the mixing part of the paper is related to their method of mixing different convolution kernels during training.

[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).
questions: 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?

2. 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)?

3. Can you also visualize some examples generated by DiT and Remix-DiT? While the metrics are valuable, a qualitative evaluation is interesting as well.

[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: limitations: The authors have mentioned this in the paper.
paper_summary: paper_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.
strength_weakness: strength_weakness: * **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.

* **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.

[1] Park et al., Denoising Task Routing for Diffusion Models, ICLR 2024.

[2] Park et al., Switch Diffusion Transformer: Synergizing Denoising Tasks with Sparse Mixture-of-Experts, ECCV 2024.
questions: questions: *  Is the Exponential Moving Average (EMA) model used to further train a pre-trained diffusion model?

* 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: limitations: N/A

Comments: title: Response to the rebuttal
comment: Thank you very much for your response. It well addressed my concerns. I will increase my rating to weak accept.
Thank you so much for the positive feedback! We will polish our draft with the above experiments.
I would like to thank the authors for taking the time to answer my questions in detail. Since I believe that the role of time step in diffusion networks has been less explored and my concerns have been resolved by the rebuttal, I would like to increase my score from 6 to 7.