Transformer-based Diffusion Probabilistic Models (DPMs) have shown more potential than CNN-based DPMs, yet their extensive computational requirements hinder widespread practical applications. To reduce the computation budget of transformer-based DPMs, this work proposes the Efficient Diffusion Transformer (EDT) framework. The framework includes a lightweight-design diffusion model architecture, and a training-free Attention Modulation Matrix and its alternation arrangement in EDT inspired by human-like sketching. Additionally, we propose a token relation-enhanced masking training strategy tailored explicitly for EDT to augment its token relation learning capability. Our extensive experiments demonstrate the efficacy of EDT. The EDT framework reduces training and inference costs and surpasses existing transformer-based diffusion models in image synthesis performance, thereby achieving a significant overall enhancement. With lower FID, EDT-S, EDT-B, and EDT-XL attained speed-ups of 3.93x, 2.84x, and 1.92x respectively in the training phase, and 2.29x, 2.29x, and 2.22x respectively in inference, compared to the corresponding sizes of MDTv2. The source code is released at https://github.com/xinwangChen/EDT.
Paper
Similar papers
Peer review
Summary
Transformer-based Diffusion Probabilistic Models (DPMs) have shown great potential in image generation tasks but are often hindered by extensive computational requirements. This paper introduces the Efficient Diffusion Transformer (EDT) framework to address these computational challenges. The EDT framework features a lightweight diffusion model architecture and incorporates a classifier-free Attention Modulation Matrix inspired by human-like sketching. Additionally, the authors propose a token relation-enhanced masking training strategy tailored explicitly for EDT to improve token relation learning. Extensive experiments demonstrate that the EDT framework significantly reduces both training and inference costs while surpassing the performance of existing transformer-based diffusion models in image synthesis. The paper highlights the effectiveness of the EDT framework in achieving a balance between computational efficiency and high performance in image generation tasks. The proposed methods and strategies demonstrate potential for broader applications and future research in the optimization of transformer-based models.
Strengths
1. The paper introduces a novel way of approaching Transformer-based Diffusion Probabilistic Models, especially a classifier-free approach. The paper also attempts to understand the relationship between tokens, thereby making the pipeline efficient. 2. The pipeline is robust and lightweight. The paper introduces Attention Modulation matrix, that provides specific quantitative results for where to focus and how much to focus on an image. The authors have provided visual results on how this method has improved the quality of the images.
Weaknesses
1. The pre-trained VAE might be trained on a classifier task, that might have some impact on categorization on certain classes. 2. The paper does not mention about the number of epochs or an estimate of time required by EDT to be able to produce this result.
Questions
1. Is there any specific reason to use a VAE? There are other lightweight image feature extractor models available, that might make the pipeline even more efficient. 2. The experimental validation, while thorough, might benefit from a broader range of datasets to ensure the robustness and generalizability of the EDT framework across different image generation tasks. 3. Please add potential limitations of the work.
Rating
6
Confidence
1
Soundness
3
Presentation
3
Contribution
3
Limitations
The paper is an extremely interesting read. The paper talks about the trade-off between computational resources and performance. However, there is no specific section dedicated, where the potential limitations have been discussed.
Summary
This paper introduces the Efficient Diffusion Transformer (EDT) framework to address the high computational requirements of transformer-based Diffusion Probabilistic Models (DPMs). The EDT framework features a lightweight diffusion model architecture, a training-free Attention Modulation Matrix inspired by human sketching, and a token relation-enhanced masking training strategy. Extensive experiments show that EDT reduces training and inference costs while surpassing existing transformer-based diffusion models in image synthesis performance. For example, EDT-S achieves a lower FID score of 34.27 and offers significantly faster training and inference speeds compared to MDTv2-S.
Strengths
+ The proposed network can reduce the computation of diffusion is a plus point although it requires more parameters compared to existing methods. + Some experimental results on ImageNEt 256x256 show promising looks.
Weaknesses
**[Update]: after rebuttal, I raised score from 4 to 5** The proposed method shows some good merits in the architectural design, however, it remains several major concerns: + Most results are reported with not very optimal settings. Specifically, classifier-free guidance is the standard practice, but it is omitted, and cannot show the advantage of the proposed method. This results in the not-very convincing evaluation that the proposed method is superior to the existing approaches with SOTAs (DiT, MDT, RDM [1], etc). + More datasets (or resolutions, such as 512) need to be conducted to have a more conclusive finding for the effectiveness of the proposed method as the proposed method is lightweight, it should not be a big deal with a higher resolution training. [1] Relay Diffusion: Unifying diffusion process across resolutions for image synthesis, ICLR 2024
Questions
My main concerns are listed in the weaknesses of its evaluation of the CFG setting and other datasets.
Rating
5
Confidence
5
Soundness
3
Presentation
2
Contribution
2
Limitations
Yes, they discussed it.
Thank you for your helpful suggestions and for providing a method for submitting qualitative images. We have submitted two sets of qualitative images at [https://postimg.cc/gallery/ZS1FXcb](https://postimg.cc/gallery/ZS1FXcb). --- ***Figure 1 is a qualitative analysis of the AMM plugin.*** AMM is a train-free plugin that can be inserted into pre-trained models to improve image generation performance. We compared images generated by EDT-XL with and without AMM to qualitatively analyze the effect of AMM on image generation. In Figure 1, the red boxes highlight the unrealistic areas in the images generated by EDT-XL without AMM. **In the corresponding areas of the images generated by EDT-XL with AMM, the results appear more realistic.** Moreover, the parrot image generated by EDT-XL without AMM is realistic and the parrot image generated by EDT-XL with AMM still remains equally realistic. Therefore, **adding AMM does not negatively affect the original quality.** This visual analysis demonstrates the effectiveness of the AMM plugin. --- ***Figure 2 is a qualitative results comparison of reconstructing masked images using models trained by the masking training strategy of MDT and EDT.*** We conducted inpainting experiments using EDT-S under different masked ratios. We utilized two versions of the pre-trained EDT-S models, each employing the masking training strategy of MDT and EDT, respectively. In Figure 2, We applied various mask ratios to the images and used these two models to reconstruct the masked areas. **When the mask ratio reached 50%, the EDT-S model using MDT's masking strategy struggled to reconstruct the images, while the EDT-S model using EDT's masking strategy was still able to do so. EDT trained by EDT's masking training strategy demonstrates better image reconstruction capabilities.** This visual analysis demonstrates the effectiveness of the EDT's masking training strategy.
Please note that we have highlighted the difference between the samples, as suggested by reviewer FhNv, and edited the comment above.
Summary
The Efficient Diffusion Transformer (EDT) framework is developed, featuring a lightweight architecture designed based on thorough computational analysis. Inspired by human sketching, EDT alternates between global attention and location attention. Additionally, the Attention Modulation Matrix enhances the detail of generated images in pre-trained diffusion transformers without requiring extra training. A novel token masking training strategy is proposed to improve the token relation learning capability of EDT. EDT achieves a new state-of-the-art performance and faster training and inference speeds compared to existing models like DiT and MDTv2. A series of exploratory experiments and ablation studies were conducted to analyze and identify the key factors affecting EDT's performance.
Strengths
Novelty: The three innovative techniques introduced for designing efficient diffusion transformers represent a significant advancement in the field. These techniques, which include the development of a lightweight architecture, the introduction of the Attention Modulation Matrix, and a novel token masking training strategy, are unique contributions that set this work apart from existing approaches. Significance: The challenge of designing efficient diffusion transformers for training and inference is a critical issue in the development of diffusion models. The work’s focus on enhancing efficiency without compromising performance is particularly relevant in the context of growing model sizes and the demand for faster computational methods. Methodology: The proposed algorithm is well formulated and clearly explained. The approach includes a comprehensive computational analysis to design a lightweight diffusion transformer architecture, inspired by human sketching with an alternation process between global attention and location attention. The introduction of the Attention Modulation Matrix is particularly noteworthy, as it improves image detail in pre-trained diffusion transformers without additional training costs. The novel token masking training strategy enhances the learning ability of the model, demonstrating a sophisticated understanding of token relationships. Results: The experimental results show improvements over existing methods such as DiT and SD-DiT. A series of exploratory experiments and ablation studies further validate the robustness and effectiveness of the proposed techniques, providing a detailed analysis of the key factors influencing EDT's performance.
Weaknesses
1. The training dataset is limited to ImageNet, which may not fully represent the broader applicability of the method. 2. The Attention Modulation Matrix (AMM) could be tested on text-to-image models, such as the Pixart Series, to demonstrate the wider applicability of the proposed method. 3. The motivation behind the lightweight design is not entirely clear. It seems to balance between UNet and transformer architectures but appears more similar to the UNet. 4. The computational process of the AMM would benefit from additional figures or diagrams to enhance understanding.
Questions
1. Why do long skip connection modules inevitably lead to a loss of token information? From my perspective, long skip connections represent a fundamental difference between UNet and transformer architectures, potentially enhancing the training process by facilitating information flow across layers. If these connections are perceived as a drawback, wouldn't it be more effective to adopt the vanilla transformer architecture, which typically avoids such connectivity patterns? 2. What is the additional computational cost incurred by integrating the Attention Modulation Matrix into pre-trained DiT models? Understanding the impact on computational resources is crucial for assessing the feasibility and scalability of implementing this enhancement across different model configurations and training scenarios. 3. Can you provide a detailed comparison of the training costs associated with your proposed method versus other relevant approaches? Analyzing the computational overhead, training time, and resource requirements relative to existing methods will provide valuable insights into the practical advantages and trade-offs of adopting your approach for training and deploying diffusion transformers.
Rating
5
Confidence
4
Soundness
3
Presentation
3
Contribution
3
Limitations
The evaluation and application of the proposed method is limited.
Additional qualitative results for both AMM and masking training strategy
As reviewers FhNv and 8pvq suggested, we added qualitative results to visualize our results. --- We highlight the different areas by the red box on the samples. Please refer to the new Link at [https://postimg.cc/gallery/ZS1FXcb](https://postimg.cc/gallery/ZS1FXcb). --- ***1.The difference between the provided samples generated by the models with/without AMM.*** In Figure 1, the red boxes highlight the unrealistic areas in the images generated by EDT-XL without AMM. **In the corresponding areas of the images generated by EDT-XL with AMM, the results appear more realistic.** Moreover, the parrot image generated by EDT-XL without AMM is realistic and the parrot image generated by EDT-XL with AMM still remains equally realistic. Therefore, **adding AMM does not negatively affect the original quality.** This visual analysis demonstrates the effectiveness of the AMM plugin. --- ***2.The difference between the provided samples generated by the models with the proposed masking training strategy and MDT's masking strategy.*** In Figure 2, We applied various mask ratios to the images and used these two models to reconstruct the masked areas during inference. **When the mask ratio reached 50%, the EDT-S model using MDT's masking strategy struggled to reconstruct the images, while the EDT-S model using EDT's masking strategy was still able to do so.** EDT trained by EDT's masking training strategy demonstrated better image reconstruction capabilities. This visual analysis demonstrates the effectiveness of the EDT's masking training strategy. --- Thank you for your valuable and insightful suggestions. We are looking forward to hearing your feedback on EDT. If you have any questions, we are more than willing to provide further clarification and address any issues promptly.
Summary
The paper introduces a new efficient diffusion-based model, namely EDT. First, they revisit the masking strategy proposed in MDT and provide some insights regarding the discrepancy in the training objective. To this end, EDT uses a more efficient masking mechanism that focuses on the main generation task instead of paying more attention to recovering the masked regions. Secondly, the main design is inspired by the human brain, which is nice, where they implement an alternating mechanism to alternate between local and global attention.
Strengths
* Alternating between the local and global details inspired by human brains is novel. * Tackling an important application. * The paper's writing is good, which makes the paper easy to follow in general.
Weaknesses
* [Methodolgy] The papers tackle an important application. However, I am concerned about the execution way. Designing architecture inspired by human brains is very important; however, the gain seems limited. * [Expirements] It is mentioned in lines 43, 62, and 193 that the AMM could be easily integrated into any existing method. However, there are no experiments to support this claim. This is crucial to show the effectiveness of the proposed module. * [Expirements] Conducting the experiments on only one dataset is insufficient. It is recommended to show the gain of the proposed method on Celeb-HQ and Lsun-Churches datasets. I suggested these two as their scale is considerably small thus conducting these experiments should not be challenging. * [Results] The claims in the abstract related to the efficiency by saying your method achieves 4 x speed up in training is misleading. Despite your clarification, this is only valid for the small variant, still I was expecting significant improvements on the other variants as well, which is not the case. * [Analysis] The analysis regarding the MVD-v2 masking mechanism is shallow. In addition, presenting the loss trends using screenshots of the terminal is shocking. * [Ablations] Ablation studies are very important so differing them to the appendix could be problematic. I recommend saving some space and including the ablations in the main paper. * [Visualization] All the Figures must be at the top of the pages to enhance readability. * [Experiments] The related work is too short and not enough. At least, I was expecting a longer version in the appendix. In addition, some important work related to this is missing. I would suggest the author include Toddler [1] as its motivation, designing an efficient diffusion-based generative model inspired by human brains, which is very similar. [1] Bakr, Eslam Mohamed, et al. "ToddlerDiffusion: Flash Interpretable Controllable Diffusion Model." arXiv preprint arXiv:2311.14542 (2023).
Questions
* Please show the effectiveness of your method on more datasets, as suggested earlier. * Integrate the AMM module into different methods to show its effectiveness as claimed.
Rating
5
Confidence
5
Soundness
3
Presentation
2
Contribution
3
Limitations
* The limitations are not discussed.
Thank you for your professional and insightful comments. These comments enhance the quality of our work. We appreciate the recognition of our efforts to improve the score. Below are our responses to the follow-up questions. --- We highlight the different areas by the red box on the samples. Please refer to the new Link at [https://postimg.cc/gallery/ZS1FXcb](https://postimg.cc/gallery/ZS1FXcb). --- ***1.The difference between the provided samples generated by the models with/without AMM.*** In Figure 1, the red boxes highlight the unrealistic areas in the images generated by EDT-XL without AMM. **In the corresponding areas of the images generated by EDT-XL with AMM, the results appear more realistic.** Moreover, the parrot image generated by EDT-XL without AMM is realistic and the parrot image generated by EDT-XL with AMM still remains equally realistic. Therefore, **adding AMM does not negatively affect the original quality.** This visual analysis demonstrates the effectiveness of the AMM plugin. --- ***2.The difference between the provided samples generated by the models with the proposed masking training strategy and MDT's masking strategy.*** In Figure 2, We applied various mask ratios to the images and used these two models to reconstruct the masked areas during inference. **When the mask ratio reached 50%, the EDT-S model using MDT's masking strategy struggled to reconstruct the images, while the EDT-S model using EDT's masking strategy was still able to do so. EDT trained by EDT's masking training strategy demonstrates better image reconstruction capabilities.** This visual analysis demonstrates the effectiveness of the EDT's masking training strategy. --- ***3.The question for the reported FID on CelebA-HQ.*** Due to the limited time and resources, we train 100k steps for both baselines and our proposed EDT for a fair comparison, using the standard training settings. We agree that the FID is not very promising for the optimal result, but the result still can show better performance than the other baseline methods. We will perform a more steps (400k) training experiment for future revision. --- ***4.High resolution result.*** Indeed, the high-resolution result makes the work more convincing. Currently, we are performing the experiment on the Imagenet 512x512, and we will provide the result in the revision.
Thank you for the insightful discussion!
Thanks for the rebuttal. It partly resolved my concerns, however, still remained. First, putting some numeric numbers is not very convincing. In the submission (including the appendix), the qualitative results (generated samples) are absent and it is very important to see the generation tasks to support their quantitative numbers. In addition to the missing comparison with other methods on 512 resolution, the author heavily referred to and compared with MDTv2 but they only focused on the earlier convergence of training steps. In MDTv2 they compared with baseline DiT on both earlier (e.g. 400k steps) and full convergence (4600k steps) and MDTv2 reached their best FID 1.58 on ImageNet 256x256. The authors claim the efficiency of the method with much lower computation cost, which raises the question of why they ignored their optimal and final results as their models are trained until fully converged. Can it beat the best FID of MDTv2? Qualitative results comparison to support their numerical number should also be carefully prepared for both their optimal generated images and the comparison of reconstructed images using the masking strategy of MDT and their masking strategy (see MDTv1 ICCV version for the referred image reconstruction with masking). I expect that the reconstructed output of the proposed method with their masking is better than MDT.
We sincerely thank you for your valuable suggestions regarding the comparative experiments and qualitative analysis of the masking training strategy of EDT. These additions make the results more understandable and convincing than relying solely on numerical data. --- ***1. Qualitative results comparison of inpainting images using models trained by the masking training strategy of MDT and EDT.*** After following your advice regarding the qualitative results of masking training and referring to MDTv1, we conducted inpainting experiments using EDT-S under different masked ratios. We utilized two versions of the pre-trained EDT-S models, each employing the masking training strategy of MDT and EDT, respectively. We applied various mask ratios to the images and used these two models to inpaint the masked areas. **When the mask ratio reached 50%, the EDT-S model using MDT's masking strategy struggled to reconstruct the images, while the EDT-S model using EDT's masking strategy was still able to do so.** We plan to include this visual comparison in the appendix. --- ***2. The experiment of EDT-XL for full convergence.*** Our primary focus has been on the efficiency of diffusion transformers. Our proposed EDT offers the community a more efficient diffusion model architecture, which achieves speed-ups of 2.29x, 2.29x, and 2.22x in inference speed on small, base, and xlarge sizes respectively. Unfortunately, training an EDT-XL for 4600k iterations would require approximately two months for one round of training, and we could not complete this due to resource limitations. Nonetheless, we have validated that EDT maintains acceptable accuracy while significantly improving speed within our capacity. **Training diffusion model is exceptionally costly, and we are currently training an EDT-XL for 2000k iterations, with plans to include the relevant results in the updated version of the paper.** --- ***3. Experiment on different datasets to evaluate the effectiveness of the EDT.*** Unfortunately, because of time and resource limitations, we were unable to conduct experiments at 512 resolution. Instead, we performed experiments on a different dataset (CelebA-HQ), where EDT-S showed competitive results. --- In addition, please note that the masking training strategy in our work is just one part of what we have contributed. It aims to mitigate the performance loss of the efficient diffusion transformer architecture. **Our contributions also encompass the model architecture design and the AMM plugin. These modules are not proposed by the current DiT works, and we believe they can offer new insights for developing this field.**
Thank you for the response; it seems to address my concerns. However, I'm curious why the comparison images weren't provided now but waited to be put into the Appendix. Could the authors upload the qualitative images to a free platform like https://postimg.cc/ and share the link here, ensuring anonymity? I've seen this link used in conferences like ICLR and ICML for reviews and rebuttals. From what I've quickly read, NeurIPS policy seems to allow external links. See [1] for reference to sharing images. [1] Relay Diffusion: Unifying diffusion process across resolutions for image synthesis, ICLR 2024.
Thanks to Authors
I appreciate the author's rebuttal and efforts. After carefully reading the other reviewers' rebuttals and comments, I am more inclined to raise my score to 5 instead of 4. As I have some concerns, as follows: - The qualitative results in response to the reviewer "8pvq" do not show the superiority of the proposed methods. I see no difference between the provided samples. - The quantitative results show some gains but are still not so convincing. For instance, the reported FID on CelebHQ is too high. The diffusion-based models, even the small ones, can easily get around 7 FID on CelebHQ. - I also agree with other reviewers about the importance of showing the effectiveness of the proposed method on higher resolution, e.g., 512. However, I also understand the hardware challenges; thus, I will not put any weight on this.
Decision
Accept (poster)