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

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.

Strengths

* The paper is structured well, making it easy to understand and follow. * The proposed mixing basis strategy is interesting as it achieves better performance with fewer parameters compared to existing multi-expert methods. * 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. * **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

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

Rating

6

Confidence

5

Soundness

3

Presentation

3

Contribution

2

Limitations

N/A

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.

Strengths

- The paper is well-motivated and represents a valuable step towards integrating the multi-expert denoising framework into standard diffusion models. - The main idea of the paper (using global mixers to compute the final experts) is novel and interesting to me in this context. - The method is simple and effective, making it more suitable for practical use cases. - The experiments are well-designed, and the ablations clearly illustrate the impact of various aspects of Remix-DiT. - 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) - 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

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.

Rating

7

Confidence

4

Soundness

3

Presentation

3

Contribution

3

Limitations

The authors have mentioned this in the 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.

Strengths

1. 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? 2. The 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. 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

Please refer to the weaknesses.

Rating

7

Confidence

4

Soundness

4

Presentation

3

Contribution

4

Limitations

This paper discusses limitations such as sparse gradients and the training difficulty associated with a large number of experts.

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.

Strengths

Novelty: Model mixers for efficient multi-expert diffusion model training is innovative and unique. Significance: Addressing the challenge of efficient training of multi-expert diffusion transformers is significant in the field of diffusion models. Methodology: The proposed algorithm is well-formulated and clearly explained. Results: 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. 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

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.

Rating

5

Confidence

4

Soundness

3

Presentation

3

Contribution

2

Limitations

Please refer to the weakness and question part.