EDT: An Efficient Diffusion Transformer Framework Inspired by Human-like Sketching

Thanks for your positive feedback and insightful comments. We will answer your questions in the following:

---

1. The usage of pre-trained VAE.

• In the field of Latent Diffusion Models, pre-trained VAE is a commonly used model. When training diffusion models or generating images, the pre-trained VAE remains frozen. The VAE encoder is used to encode images into latent representations and the VAE decoder is used to decode latent representations back into images.

---

2. Training cost of EDT.

• We reported the training cost and speed in Table 1 in the main paper. For ImageNet 256×256, we estimated the training cost for EDT, MDTv2, and DiT on a 48GB-L40 GPU in Rebuttal-Table 3 in the global rebuttal, using a batch size of 256 and FP32 precision. EDT achieves the best performance with a low training cost. (GPU days refer to the number of days required for training on a single L40 GPU.)

Rebuttal-Table 3: The training cost and FID of EDT, DiT, and MDTv2 on ImageNet 256×256 with batch size of 256 and FP32 precision.

---

3. Why use VAE.

• VAE is used to reduce the computational cost of diffusion models. The diffusion model generates images through multiple iterations, whose input in each iteration is its output in the previous iteration. To reduce the computational cost of each iteration, we make the diffusion model train and predict in latent space, since the feature size in latent space is smaller than that in the pixel space of images. During inference, the diffusion model is used to generate the latent representations of images, and the VAE decoder is used to decode latent representations back into images.

---

4. Experiments on extra dataset celebA-HQ.

• We conducted a new experiment on CelebA-HQ 256×256 for the unconditional image synthesis task. We train EDT-S, DiT-S, and MDTv2-S on CelebA-HQ 256×256 under the same training epochs with the default settings in their papers, respectively. As shown in Rebuttal Table 2, EDT-S achieved the lowest FID, demonstrating that EDT is also effective in unconditional image synthesis tasks.

Rebuttal-Table 2: Evaluation of unconditional image synthesis on CelebA-HQ.

---

5. Potential limitations.

• The experiments were conducted on class-conditional image synthesis tasks and unconditional image synthesis. Future work will explore other image synthesis tasks such as text-to-image tasks.
• When integrating AMM into a pre-trained model, the insert and arrangement of AMM blocks varies across different models. Identifying the optimal placement and configuration of AMM requires testing to fully realize its potential.
• Although significant progress has been made with AMM, the generation function of AMM still has room for improvement and warrants further exploration.

We thank the reviewer for the valuable feedback and suggestions. We will answer your questions in the following:

---

1. The experiment under optimal CFG settings.

• According to DiT and MDTv2, their optimal CFG settings are 1.5 and 3.8, respectively. Based on our experimental exploration, the optimal Cfg for EDT is 2. As shown in Rebuttal-Table 4, when compared under their optimal CFG settings with EDT and at the same training cost, EDT-S achieves the best performance with the lowest FID.

Rebuttal-Table 4: Comparison and evaluation under optimal Cfg on ImageNet 256×256.

---

2. Experiments on extra dataset CelebA-HQ.

• We conducted a new experiment on CelebA-HQ 256×256 for the unconditional image synthesis task. We train EDT-S, DiT-S, and MDTv2-S on CelebA-HQ 256×256 under the same training epochs with the default settings in their papers, respectively. As shown in Rebuttal Table 2, EDT-S achieved the lowest FID, demonstrating that EDT is also effective in unconditional image synthesis tasks. Due to the limited time, we did not finish the experiments with the 512x512 resolution training. We will include it in the revision.

Rebuttal-Table 2: Evaluation of unconditional image synthesis on CelebA-HQ.

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.

---

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.

Thanks for your insightful comments and suggestions. We will address your concerns in the following answers:

---

1. Experiments on extra dataset CelebA-HQ.

• We conducted a new experiment on CelebA-HQ 256×256 for the unconditional image synthesis task. We train EDT-S, DiT-S, and MDTv2-S on CelebA-HQ 256×256 under the same training epochs with the default settings in their papers, respectively. As shown in Rebuttal Table 2, EDT-S achieved the lowest FID, demonstrating that EDT is also effective in unconditional image synthesis tasks.

Rebuttal-Table 2: Evaluation of unconditional image synthesis on CelebA-HQ.

---

2.Motivation for lightweight design.

• Our motivation for a lightweight design is to decrease the computational overhead by reducing the number of tokens in the intermediate blocks. We achieve this by using down-sampling modules to compress the tokens. This idea is similar to U-Net, but we implemented it to the transformer-based diffusion models. However, while compressing tokens reduces the computational overhead of EDT, it also compromises token information. To address this issue, we have enhanced the down-sampling modules and long skip connection modules, incorporating techniques such as ''token information enhancement'' and ''positional encoding supplement''. As shown in Table 4 of the Appendix, the effectiveness of these improvements is demonstrated through ablation experiments.

---

3. The computational process of the AMM.

• In Rebuttal-Figure 2 in the PDF of global rebuttal, we illustrate the process of attention modulation. Image can be split into N patches and each token is the feature of a patch. Each token (patch) corresponds to a rectangular area of the image and has a corresponding 2-D coordinate (x, y) in the image. We calculate an Euclidean distance value $d$ for each pair of tokens, resulting in a distance matrix, which is an N×N tensor. Based on the distance matrix, we generate modulation values $m$ via the modulation matrix generation function $F(d)$, which assigns lower modulation values to tokens that are farther apart. These modulation values form an Attention Modulation Matrix (AMM), another N×N tensor. Importantly, we integrate the AMM into the pre-trained EDT without any additional training. The attention modulation matrix is calculated when the model is instantiated. During inference, the modulated attention score matrix is obtained by performing a Hadamard product between the attention modulation matrix and the attention score matrix.

---

4. The computational cost of AMM.

• The addition of AMM introduces minimal computational costs. Firstly, AMM can be incorporated into a pre-trained model without requiring additional fine-tuning, resulting in no additional training costs. Secondly, the increased computational cost of AMM during inference is negligible. For instance, in the last block of EDT-XL, the attention score matrix and the Attention Modulation Matrix are both 18×256×256 tensors. The computational cost of the Hadamard product between the attention score matrix and AMM is only 1.18M FLOPs for multiplication calculations, out of a total of 2819.3M FLOPs for the block. This amounts to merely 0.04% of the total FLOPs, making the computational cost of AMM negligible. In our experiments, we added 5 AMM modules to a pre-trained DiT-XL model, and the FID score decreased from 18 to 14.

---

5. The loss of token information in long skip connection.

• Long skip connection modules are crucial for the performance of diffusion models, but the token merging operations within these modules can result in the potential loss of token information. Before the long skip connection, we have two sets of tokens, each represented as an N × D tensor. During the long skip connection, these two sets of tokens are concatenated into an N × 2D tensor and then merged into an N × D tensor through a linear layer. This dimensionality reduction from 2D to D leads to a loss of token information and disrupts positional information. To address this issue, we introduced Token Information Enhancement and Positional Encoding Supplement within the long skip connections, as demonstrated in Table 4 of the Appendix.

---

6. Training cost of EDT.

• We reported the training cost and speed in Table 1 in the main paper. For ImageNet 256×256, we estimated the training cost for EDT, MDTv2, and DiT on a 48GB-L40 GPU in Rebuttal-Table 3 in the global rebuttal, using a batch size of 256 and FP32 precision. EDT achieves the best performance with a low training cost. (GPU days refer to the number of days required for training on a single L40 GPU.)

Rebuttal-Table 3: The training cost and FID of EDT, DiT, and MDTv2 on ImageNet 256×256 with batch size of 256 and FP32 precision.

---

7. Adding AMM to Text-to-Image Models.

• Due to time constraints, incorporating AMM into text-to-image models will be addressed in our future work.

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.

---

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.

We appreciate the professional and insightful comments. We address each comment as follows:

---

1. Performance improvement by AMM.

• We demonstrate that AMM is both effective and efficient through several experiments detailed in Tables 9, 10, and 11 in the Appendix. Tables 9 and 10 compare model performance with and without AMM across different sizes of EDTs. Models incorporating AMM achieve lower FID scores and higher IS scores. For example, Table 9 shows that EDT-S with AMM achieves a lower FID score of 34.2 compared to 46.9 for EDT-S without AMM. Notably, AMM is integrated into pre-trained EDT-S without additional training. Table 11 illustrates AMM's effectiveness across different iterations of EDT. Furthermore, AMM and its arrangement in blocks are designed by mimicking the brain's logical process—alternating between global and local attention as observed in human sketching. We will include a detailed ablation study of AMM in the main paper.

---

2. More models incorporated with AMM.

• We integrated AMM into pre-trained EDT, DiT, and MDT, which are transformer-based models. We conducted experiments on ImageNet and Celeba-HQ in Rebuttal-Table 1. The models with AMM obtain lower FID compared to the models without AMM, which further demonstrates the effectiveness of AMM in different models and datasets.

Rebuttal-Table 1: Evaluation of the performance of pre-trained EDT, DiT, and MDT without and with AMM on ImageNet and CelebA-HQ.

---

3. Experiments on extra dataset CelebA-HQ.

• We conducted a new experiment on CelebA-HQ 256×256 for the unconditional image synthesis task. We train EDT-S, DiT-S, and MDTv2-S on Celeb-HQ 256×256 under the same training epochs with the default settings in their papers, respectively. As shown in Rebuttal Table 2, EDT-S achieved the lowest FID, demonstrating that EDT is also effective in unconditional image synthesis tasks.

Rebuttal-Table 2: Evaluation of unconditional image synthesis on CelebA-HQ.

---

4. Clarification of the speed-ups in training.

• As shown in Table 1 of the main paper, EDT-S, EDT-B, and EDT-XL attain 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, when compared to the corresponding sizes of MDTV2. We will provide a more explicit description regarding the speed-up details.

---

5. Analysis of masking training mechanism.

• By observing the loss changes in Rebuttal-Figure 1 in the PDF of the global rebuttal, we identified a conflict between $L_{masked}$ (loss when the input consists of the remaining tokens after masking) and $L_{full}$ (loss when the input consists of the full token input) in MDTv2. Although MDTv2 also observed similar issues, they did not solve them effectively. In MDTv2, it was found that using only $L_{masked}$ caused the model to overly focus on reconstructing the masked tokens, thereby neglecting diffusion training. To address this, both the full token input and the masked token input were fed to the diffusion model, resulting in the inclusion of both $L_{masked}$ and $L_{full}$ in MDTv2's loss function. However, our work discovered a conflict between $L_{masked}$ and $L_{full}$. We separately applied the masking training strategies of MDTv2 and EDT to train diffusion models and extracted $L_{masked}$ and $L_{full}$ values at the 300k~305k training iterations. As shown in Rebuttal-Figure 1 in the PDF of the global rebuttal, we visualized the changes of $L_{masked}$ and $L_{full}$. The left-top of Rebuttal-Figure 1 depicts the loss changes when using MDTv2's masking training strategy. As $L_{full}$ decreases, $L_{masked}$ increases, and vice versa, illustrating the conflict between these two losses. This conflict arises because $L_{masked}$ causes the model to focus on masked token reconstruction while ignoring diffusion training. As shown in the bottom-left of Rebuttal-Figure 1, both the $L_{full}$ and $L_{masked}$ hardly converged during the 300k to 400k training iterations. In our work, to resolve this conflict, we performed masking in the intermediate layer instead of before input. This method forces the model learning to establish relations between tokens before they are masked. The right side of Rebuttal-Figure 1 shows the loss changes when using EDT's masking training strategy. With EDT's strategy, $L_{masked}$ and $L_{full}$ exhibit synchronized changes, and the loss values continuously decrease during the 300k to 400k training iterations. Due to the page limit, the aforementioned discussion will be included in the supplementary material for the camera-ready version.

---

6. Ablations.

• We will rearrange the work and place the ablations table in the main paper.

---

7. Visualization.

• We will make all figures be placed at the top of the pages to enhance readability.

---

8. Related work

• We will improve our related work and provide more comprehensive writing in the appendix, and we will discuss important work such as Toddler in the related work.

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.

---

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.