EDT: An Efficient Diffusion Transformer Framework Inspired by Human-like Sketching
Authors: Xinwang Chen
Ning Liu
Yichen Zhu
Feifei Feng
Jian Tang
Venue: NeurIPS 2024 poster (2024)
Keywords: Efficient AI; Lightweight architecture; Optimization; Diffusion
Primary area: optimization_for_deep_networks

Abstract: 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. This 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. Our code is available at https://github.com/xinwangChen/EDT.

EDT: An Efficient Diffusion Transformer Framework Inspired by Human-like Sketching
Xinwang Chen "l Ning Liu"l, Yichen Zhu', Feifei Feng', Jian Tang Midea Group, Beijing Innovation Center of Humanoid Robotics chen_xinwange ustb. edu. cn, (zhuyc25, feifei feng) @midmamaaa.... com, jian. tang@x cww.oxx-ummoiidd com
com
Abstract
Transformer-based Diffusion Probabilistic Models (DPMs) have shown more po- tential 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 lightweight-design diffusion model architecture, and training- freee Attention Modulation Matrix and its alternation arrangement in EDT inspired by juman-like sketching. Additionally, we propose a relation-enhanced augment its token relation learning capability. Our extensive experiments demon- strate the efficacy of EDT. The EDT framework reduces training and inference costs and surpasses existing transformer -baoodde----ddd diffusion models in image synthesis performance, thereby achieving significant overall enhancement. With lower FID, EDT-S, EDT-B, 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 the corresponding sizes of MDTv2. Our code available at here
token relation enhanced masking training strategy tailored explicitly for EDT to
Introduction
A person
Design framework
Review tree the whole picture
Review tree whole picture
Finish sketching
sketching
Improve the tree
Improve the tree
Global attention
Local attention
Global attention
Local attention
Global
attention
Figure 1: Illustration of the alternation process of local and global attention during sketching.
and practical applications 6 have validated the effectiveness of Diffusion Probabilistic Models (DPMs), establishing them as a mainstream method in image generation. In past years, predominant works [1] 2 10 have advanced diffusion models by incorporating a convolutional UNet-like [12 architecture as their backbone. On the other hand, transformers [ have achieved significant milestones in both natural language processing [4
Numerous studies
"Joint first authorship. Either author can be cited first. "Correponning author.
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

15 and computer vision [6 17 18 19,, prompting recent attempts to integrate these powerful transformer-aseee architectures into diffusion models with considerable success, For instance, U- VIT 20), an early work in diffusion leveraging ViT- transformers, surpassed the contemporary CNN-based U-Net DPMs in class-conditional image generation on ImageNet, demonstrating their potential. Similarly, diffusion transformer (DiT) [21], which employs transformers as its backbone instead of the traditional U-Net backbone in latent diffusion models (LDM) [[]] has shown excellent scalability. Further, masked diffusion transformer (MDT) [22] observes that DPMs often struggle to learn the relations among object parts in an image. To solve this, MDT introduces a masking training scheme to enhance the DPMs' ability to relation learning among object semantic parts in an image. MDT established a SOTA of class-condition image synthesis on the ImageNet.
While transformer- based DPMs offer scalability and higher performance ceiling than their CNN counterparts, they also require more computational resources. For instance, in each inference step, DiT-XL-2 consumes 1 118 GFLOPs and U-ViT-H requires 133 GFLOPs. This computational demand escalates with increasing time steps or token length, limiting their practical application. Despite their computational inefficiency, few studies have explored enhancing the efficiency of transformer-basee DPMs. Therefore, the trade-off between computation and performance underscores the importance of designing a lightweight model architecture that maintains excellent performance.
To improve the computational efficiency of transformers in DPMs, we introduced a comprehensive optimization framework named Efficient Diffusion Transformer (EDT). Specifically, we developed a lightweight diffusion transformer architecture based on a comprehensive computation analysis. Moreover, we devised the Attention Modulation Matrix (AMM) and its alternation arrangement in EDT inspired by human-like sketching. AMM, functioning as a plug-in, can be seamlessly integrated into diffusion transformers to enhance image synthesis performance significantly without requiring additional training. Additionally, we introduced novel token relation-enhanced masking training strategy tailored for EDT to enhance its relation learning capability.
Lightweight-ddsiinn diffusion transformer Based on the empirical analysis of the number of tokens, token dimensions, and the FLOPs, we propose two principles to design the lightweight diffusion transformer, and redesign and incorporate the down- nown-sampling, up-sampling. and long skip connection modules into diffusion transformers. The utilization of down-sampling module can reduce FLOPs, but harms performance, since the token merging operation in sampling and long skip connections modules leads to the loss of token information. To mitigate this loss, we enhance the key features by introducing token information enhancement and positional encoding supplement.
Attention Modulation Matrix The mind stores visual structures as top-down hierarchy passing from general shape to the relationships between parts down the detailed features of individual parts 23 24]. Based on this storage structure in the mind, humans tend to follow a coarse-to- fine drawing strategy 225 The logical structure of sketching of humans tends to first form a general framework (using global attention), then gradually refine local details (using local attention) driven by the global perspective (using global attention) shown in Figure [ Inspired by the sketching process, we integrate the alternation process of local and global attention to EDT, and propose Attention Modulation Matrix (AMM) to modulate from the default global attention self-attention mechanisms to local attention. AMM, functioning as a plug-in, which can be seamlessly integrated into diffusion transformers, enhancing image synthesis performance without necessitating additional training.
Token relation-enhanced masking training strategy The token compression in down-sampling modules may cause token information loss. Learning the relations among tokens can help token down-sampling modules compress tokens effectively. And it has been confirmed that masking training can enhance the DPMs' ability learn relations among object parts in images (22). We propose a novel masking training strategy to enhance the relation learning among tokens. Specifically, the full tokens are fed into EDT and the tokens masking is executed in nown-sampling modules. This forces models to learn token relations before some of the tokens are masked. We compare our masking training method to the counterpart in MDT, both implemented on EDT. Our masking training method achieves better generation.
We summarize the contributions of our work: 1. We develop an Efficient Diffusion Transformer (EDT) framework and design a lightweight diffusion transformer architecture based on a comprehensive computation analysis. 2. Inspired by human sketching, we design EDT with an alternation process between global attention and location attention. Moreover, to the best of our knowledge. we introduce Attention Modulation Matrix for the first time, which improves the detail of generated images of
2

---

pre- pre-trained diffusion transformers without any extra training cost. 3. We propose a novel token masking training strategy to enhance the token relation learning ability of EDT. 4. EDT has reached new SOTA and achieves faster training and inference speed compared to existing representative works DiT and MDTv2. We conduct series of exploratory experiments and ablation studies to analyze and summarize the key factors affecting the performance of EDT.
2 Method
2.1 Preliminaries
We briefly review several fundamental concepts necessary to understand classifier-free guidance class-condition diffusion models []]. The primary objective of diffusion models is to learn diffusion process that constructs a probability distribution for specific dataset, subsequently enabling the sampling of new images. Given a classifier-free guidance class-condition diffusion model Eq (Iz,c)., the model can generate images of specific class e from Gaussian noise over multiple denoising time steps. The model operates through two main processes: the forward and reverse processes. The forward process simulates training data 2t to be denoised at time step f. by adding Gaussian noise 4 ~ N IO to the original data Zy. This process is mathematically described by ((2))2)) = ' (Irs v Vā (1 ā²) I), where ãg denotes a hyperparameter. The reverse process samples noise-reduced data - based noise data 3y and class-condition The reverse process is represented as (0 Ft, c) = .N (124-1114 (I),c), ,,,,),,,, Eg c)), where lsg and Eg are the statisties of po- By optimizing the variational lower-bound of the log- lokeliiel (26) po (20) and reparameterizing 14g as a noise prediction network Eg. the model can be trained using simple mean- squared error between the predicted noise Eg (2t, and the ground truth e, sampled Gaussian noise:
(x G, 6x) = ||eg (It; c) - =elli. Additionally, eg (are, c) is a standard class-condition model; when c= 0, it functions an unconditional model. To allow the controllability of class-condition guidance, the prediction of models is further derived as ég (xte e) = 6g (Te, 0) + w (Eo (Zz, c) 6g (2r, 0)) where w 2 1 is class-condition guidance intensity.
In this work, we employ a classifier-free guidance class-condition diffusion transformer architecture operating on latent space. The pre-trained variational autoencoder (VAE) model (26 from LDM [I] remains frozen and is used to encode/decode the image/latent tokens.
2.2 Lightweight-desipe diffusion transformer
Transformer-based diffusion probabilistic models (DPMs) have demonstrated greater scalability and superior performance compared CNN-based DPMs 20 21 22 However, these models also entail significant computational overhead during both the training and inference phases. In response, we design a lightweight diffusion transformer architecture in this section. We undertake a computational complexity analysis of the cansform-based diffusion model. Based on the empirical analysis of the number of tokens, token dimensions, FLOPs, and the number of parameters, we establish two design principles: (I) reducing the number of tokens to decrease the FLOPs in the self- self-attention module through the down-son down-samplin samp-inm module; (2) ensuring that the FLOPs of each EDT stage post a down- lown-sampling module are significantly reduced compared to the stages prior to the down-sampling module, to effectively lower the overall FLOPs.
Building on the aforementioned design principles, we have redesigned and incorporated the down- sampling, up-sampling, and long skip connection modules into the transformer- diffusion model, successfully achieving a reduction in FI.OPs and increased inference speed. For instance, in comparison to DiT-S (21). our smaller version model EDT-S achieves an inference speed of 5.5 steps per second, versus 2.7 steps per second for DiT-S, effectively doubling the speed. Figure 2 illustrates the architecture of our lightweight-deiinned diffusion transformer. The model includes three EDT stages in the down-sampling phase, viewed as an encoding process where tokens are progressively compressed, and two EDT stages in the up-sampling phase, viewed as a decoding process where tokens are gradually reconstructed. These five EDT stages are interconnected through down-sampling, up-sampling, and long skip connection modules. Note that each EDT stage comprises several consecutive transformer blocks. For more details on the computational complexity analysis and model design, please refer to Appendix A.2 is important to note that the down-sampling and up-sampling phases can be viewed encoding and decoding processes, respectively. aligning with
3

---

Down sampling phase (see encoding process) enhanced by masking training strategy discussed in Section 3.4
Input model
Ousput model
EDT
EDT Stage
Down
Down
Up-sampling
N
module
VAE
di-//d <
Down-sampling module
Long
X X.X.K.
Lons
Long skip connection
N x N x 26
x
EDT
Stage
Stago
Veristional autoencoder
Image
Up-sampling phase decoding process) is
Predicted noise
Attention Modulation Matrix discussed Section 3.3
Figure 2: The architecture of lightweight-design diffusion transformer.
Long, Shp
EDT Stage
Long
tokens1
sampaing module
tokons
Long skip connection
(c) Up saappliiingggg module
Input/Dutput of module
xdd
Colenas
N x
Positional encoding supplement
Linear
Token Masking K.P
Reshape
Class-condition Information
AdaLN
AdeLN
N
N 2d.
for token information enhancerent
Linear
Reshape
Linear
4
N N x
masking
tokens]
6,4
takerssi
tokensz
G,
tokens2
operation
Figure 3: The design of down-sampling, long skip connection and up-sampling modules.
the conceptualizatio of drawing pictures in human sketching. These phases are crucial and will be
further discussed in the following Section 2.3 and Section
While we have successfully reduced the FLOPs of the model, the token merging operation in the nown-sampling and long skip connection modules inevitably leads to a loss of token information, including the positional encoding and contextual data essential for class-condition generation of images. To mitigate this loss of token information, we propose two improvements, as illustrated in Figure 3 token information enhancement and positional encoding supplement.
Token information enhancement We enhance the contextual information required for class-condition generation by employing Adaptive Layer Normalization (AdaLN) before the token merging process. AdaLN adjusts the output by learning scaling factors y and bias coefficients 3, which can scale, negate, or shut off the features (27). By utilizing AdaLN, we modulate the tokens based class conditions and time steps before merging, thereby preserving more contextual information and minimizing the loss of token information. As depicted in Figure class-condition information integrated by the 7 and 3 of AdaLN in both the down-sampling and long skip connection modules.
Positional encoding supplement We restore the absolute positional encoding of tokens following the token merging process. As illustrated in Figure after merging the tokens, we add the absolute positional encoding to the merged tokens at the end of both the down-sampling and long skip connection modules.
2.3
Making EDT "sketch" like human
The lightweight design of EDT might compromise the quality of image synthesis. To enhance the detail fidelity in generated images, have refined the decoding process (up-sampling phase) EDT by imitating the process of human sketching. We begin by examining how attention shifts during the act of sketching by humans. Human cognition stores visual structures as a top-down hierarchy passing from general shape to the relationships between parts down to the detailed features of individual parts 23 24 This hierarchical structuring visual information in the brain makes humans tend to follow a coarse-to-fine strategy in sketching (25). As shown in Figure the process of human

---

Input/Output of module
Scale
Lineer
AMM
Attention Modulation Matrix
Polmalsa
Dlement-wise multipiication
Feedlorward
Up-ampknng
AdaLN
softmax
Attnn
EDI Block with AMM
Multi-Head Self Attention
EDT Block without AMM
AdaLN
Linear
Lingar
skip connection
Lukens
An EDT stage lp--sampiing phase
An EDT block with AMM
Self-attention with AMM
Figure The position of Attention Modulation Matrix (local attention) in an EDT stage in the up-sampling phase.
sketching tends to first form general framework (using global attention), then gradually refine local details [24]22] (using local attention) hinted by the global perspective (using global attention). Even local detail, humans do not become completely detached from the overall framework. Therefore, humans periodically shift attention back to global view to scrutinize the local detail and further fine-tune it [29 30). This process reflects the alternation of global and local attention in the human brain when sketching.
when concentrating
Inspired by the sketching process, aim to integrate the alternation process of local attention and global attention to EDT. In the current series diffusion transformers [20] 22 21 only default global attention mechanisms are employed, which may lead to poor generation of local details. Therefore, we introduce the Attention Modulation Matrix (AMM) to enhance focus on local details. Moreover, mimic the alternation process in the EDT, we alternately incorporate the AMM into the lightweight-design transformer diffusion architecture.
2.3.1 Integrating local attention into the up-sampling phase of EDT
To imitate the alternation between global and local attention like the act of humans drawing, integrate local attention into the up-sampling phase of EDT by introducing Attention Modulation Matrix (AMM). In this section, we concentrate imitating the alternation process of attention. A detailed discussion of AMM is deferred to Section 2.3.2 We align the decoding process of EDT with the humans drawing pictures. Consequently, decoding process (up-sampling phase) of EDT. FigureHI iillustrates the placement of AMM (local attention) an EDT stage. As depicted in Figure aa), we alternately configure EDT blocks with and without the AMM,, thereby mimicking the alternation between global and local attention observed in drawing activities. The EDT block with AMM is shown in Figure b). As shown in Figure 4 c), the AMM is integrated into the self-attention module. The AMM and the global attention score matrix combined via Hadamard product modulate global attention into local attention.
incorporate local attention (AMM) into the
2.3.2 Attention modulation matrix
We develop the Attention Modulation Matrix (AMM) to modulate the default global attention in self-attention mechanisms into local attention, which imitates the local attention of humans during the act of drawing. Humans typically concentrate on either the actively engaged parts or the most salient aspects of visual scene [3]]. When drawing specific local region of image, region of interest tend to exhibit stronger contextual relations and thus warrant increased attention. Conversely, areas further from the region of interest generally show weaker contextual relations and can be allocated less attention. Thus, we articulate the principle: for local region, the strength of atlention on contextual relations within specific region inversely related to the distance between the local region and the specific region. In the self-attention mechanism, we regard the attention score between tokens as an indicator the strength of attention on contextual relations between regions. Similarly, aim modulate the strength of attention based on the distance among tokens
closer the
5

---

on the image. Based on this concept, we have developed the Attention Modulation Matrix (AMM), functioning as a plug-in, which can be seamlessly integrated into diffusion transformers, significantly enhancing image synthesis performance without necessitating additional training.
Formally, given a sequence of N x N tokens and its corresponding attention score matrix A we take two arbitrary tokens as an example to illustrate the formulation. The two arbitrary
R RN*x x tokens are denoted as 7 Token; Token, and their attention score dir, where dir E A. The coordinates of T Token; and Token, correspond to the (æi, Vi) and( Ur) in the original N N tokens grid, where i = Nx; + y/ and r = Nx, + Vr. The distance between these two tokens can be calculated by
Euclidean distance dir v (z: - x,)² + (y y-)² and we can derive the token distance matrix D E N We aim to modulate the global attention into local attention by multiplying the attention score matrix to the AMM, which is generated based on the token distance matrix D. The modulation matrix generation function is designed with adherence to two principles: (1) the generation function should be monotonically decreasing within the interval [0, d ensuring that inversely correlated with distance, where dmaz (N 1) v
the modulation matrix elements is the furthest distance, which is the distance between two diagonal opposite tokens; (2) the output range of this function should be limited to avoid significantly altering the original distribution of the attention score matrix. Based on the two principles, we utilize the monotonically decreasing interval in cosine function, cos (fd,r). dar [0, dmaz), where the monotonically decreasing interval can be flexibly adjusted by adjusting its period T or frequency f. According to the dmazz we set T 4ddmaz and = 27.. Further, we employ cos (fdrr) as the exponent of the Euler's number e, thereby smoothing the values of the modulation matrix elements. We obtain the final modulation matrix generation function F (dzr) = (e which can flcxibly scale the function value to [k, ke] by scaling factor k. We empirically set k 0.5 and the output range within [ 51. which allows the modulation matrix elements to appropriately adjust the attention scores. In addition, we define effective radius R for local attention to exclude the interactions of tokens that occur tokens with far distances. For each token pair, we only modulate the attention scores with distance
dir < R. where R is the effective radius for local attention. We set R = V(N - 1)2 + 4 based on experiments regarding the hyper-parameters of the AMM, detailed in Appendix A.3.3 And those dir > R, their attention scores are set to zero, indicating that tokens far from the region of interest exert less influence. Thus, we have the Attention Modulation Matrix M E NxN2 where Mar E M is defined as:
Mir {
F (ddr), 0,
der S R de R
(I)
The modulated attention score element is d' = Oir * Mirr where d'p- E A', and A' is the modulated attention score matrix. Further details about the entire process and an illustration of AMM can be found in Appendix A.3.T
2.4 Token relation-enhancel masking training strategy
The ability to learn relations among object parts in images is crucial for image generation, as highlighted in MDT (22]. However, the down-sampling process in EDT inevitably leads to the loss of token information. Establishing relations among tokens can alleviate performance degradation caused by the loss of token information. To enhance the relation-learning ability in EDT, we introduce a relation-enhanced masking training strategy. Before detailing the proposed masking training strategy, we first explore the integration of MDT into EDT. Figure 5 (a) shows the masking training strategy of MDT. In MDT, the training loss I. contains two parts as shown in Eqn. [
L= Lfull + Iimasked Lo (Ft, C, ()) + (mask: C,
Lfult is the loss when the input consists of the full token input, the Lmaskee is the loss when the input consists of the remained tokens after masking, and mask is matrix to mask tokens randomly.
However, our analysis reveals that the masking training method used in MDT excessively focuses on masked region reconstruction at the expense of diffusion training. potentially leading to a degradation in image generation performance. Additionally, our evaluation of MDT is observed conflict between the training objectives of Lfull and Lmaskedd Specifically, as Lfal decreases, I-masked increases, and vice versa, demonstrating the conflicting nature of these training objectives. To mitigate this conflict and allow the model to focus on the diffusion generation task, as shown in Figure 5 (b), we

---

(8) Masking Latent Inputs, model focuses too much on masked regions reconstruction.
Modules
MDT
Masking
Random masking
MDT
Latent or Tokens
Nolsed latent
Remained tokens
Denoised latent
[b] Postponing the operation of masking, model learns relation among takens before some tokens heing masked.
EDT(our)
Down sampling phase
Random Masking
Up-sompling phose
EDT Stage
EDT Stage
Remained tokens
Denoised latent
Figure 5: Token relation-enhanced masking training strategy. MDT is fed the remained tokens after token masking into models. EDT is fed full tokens into shallow EDT blocks, and the operation of token masking is performed in down-sampling modules.
design the token relation-enhanced masking training strategy. which feeds full tokens into shallow blocks and postpones the token masking operation to occur within the down-sampling modules. This strategy is designed to facilitate learning relationships among tokens and reduce the loss of token information without the issues arising from conflicting training objectives. When training, the masked tokens are unseen to the EDT blocks following the operation of masking, which forces the EDT blocks before the operation of masking to learn the relations among tokens. As the relations among tokens learned, the key information in each token is dispersed and stored across various tokens. This avoids reliance on certain tokens and reduces the loss token information from compression in down wwwnn-sllling.... Figure 3a) shows, with red arrow in the down- nown-sampling module, the specific point at which token masking is performed. After passing through the down-sampling modules, the EDT stages in up-sampling phase generate images solely relying on the remained tokens. The loss function of EDT with token relation-enhanced masking training strategy is shown in Eqn. where the token masking operation is executed in the down-sampling modules.
= full + T-masked = Lo( c, (.) + Lo(xe, C, mask, 6t)
(3)
We discover that EDT is particularly well-suited for masking training due to its up-sampling modules, which inherently are used for the reconstruction of tokens. Unlike MDT, which requires an additional interpolator module for reconstructing masked tokens, EDT eliminates the need for such module, thereby reducing unnecessary training overhead associated with an interpolator compared to MDT. For more detailed analysis, please refer Appendix(A.44 A.4.
3 Experiment
3.1 Implementation Details
Models: We develop three different sizes of EDT including small (EDT-S), base (EDT-B) and extra large (EDT-XL), each using a patch size of two. Details regarding token dimensions, head numbers, and parameter counts provided in Table of Appendix A.2.2 Training and evaluation: The training dataset is ImageNet [32 with 256x256 256 and 512x512 resolution. For a fair comparison,
follow the training settings of MDTv2 3 EDT uses the Adan 34 optimizer with global batch size 256 and without weight decay. The learning rate linearly decreases from 1e-3 to 5e-5 over 400k iterations. Masking training strategy: We set the mask ratio 0.4 0.5 in the first down-sampling module, and 0.1 ~ 0.2 in the second. The investigation of mask ratio refers to Appendix A.4.2 GPUs: Training is conducted on eight L40 48GB GPUs, while the speed test for inference is performed on single L40 48GB GPU. Evaluation metrics: Common metrics such Fre's he'chet Inception Distance (FID) 35 sFID (3). Inception Score (IS) [37], Precision, and Recall 38) are used to assess the model performance. The training speed is evaluated by iterations per second, and inference speed is assessed by steps per second using a batch size of 256 in FP32, For fair comparison, we follow 21 33 and employ the TensorFlow evaluation suite from ADM (4) reporting FID-50K results with 250 DDIM [0] sampling steps. These metrics are reported by default without the classifier-fres guidance.

---

Table 1: The comparison with existing SOTA methods on class-conditional image generation without
classifier-free guidance on ImageNet 256x256. We report the training speed (T-speed), inference speed (I-speed), and memory consumption (Mem.) of inference. The EDT* denotes the EDT without our proposed token relation-enhance masking training strategy.
Cost.
Params. (M)
T-speed (iter/s)
I-Speed (step/s)
Mem. (MB)
Model
GFLOPs]
FID
(IterXBS)
DIT-S
400Kx256
32.90
12.50
6.06
2.70
4296
68.40
SD-DiT-S 39
400Kx256
32.90
48.39
EDT-5
400Kx2.56
38.30
13.20
2.66
5,50
4268
38.73
MDT
400Kx256
33.10
2.25
6.07
2.40
4902
39.50
EDT-S (our)
400Kx256
38.30
8.86
2.66
5.50
4268
34.27
DiT-B 21
400Kx256
130.30
4.30
23.01
1.11
8978
43.47
SD-DiT-B
400Kx256
130.30
28.62
(ourB*(our)
400Kx256
152.00
5.80
10.20
2.20
8584
23.19
MDTv2-B
400Kx256
130.80
1.42
23.02
0.96
9212
19.55
MDT
1600Kx256
130.80
1,42
23.02
0.96
9212
13.60
EDT-B(our)
400Kx256
152.00
4.03
10.20
2.20
8584
19.18
EDT-B(our)
1000Kx256
152.00
4.03
10.20
2.20
8584
13.58
ADM
1980kx256
554.00
1120.00
10.94
LDM-4
178kx 120
400.00
104.00
DiT-XL
400Kx256
674.80
0.93
118.64
0.25
17538
19.47
SD-DIT-XI
39
1300Kx256
740.60
9.01
(DTr
400Kx256
698.40
1.49
51.83
0,51
14486
10.48
MDTv2-XL(33]
400Kx256
675.80
0.51
118.69
23436
7.70
EDT-XL(our)
400Kx256
698.40
0.98
51.83
0,51
14486
7.52
Table 2:
The comparison with existing transformer based models on class-conditional image generation without classifier-free guidance on ImageNet 512x512.
T-speed (iter/s)
Model
GFLOPs]
FID.
1St
sFID
DiT-S
2.26
31.42
85.21
23.68
13.53
MDTv2-S
0.53
31.46
51.16
29.94
8.57
EDT-S(our)
1.63
13.25
51.84
29.92
7.86
3.2 Comparison with SOTA transformer- diffusion methods
To validate the enhancements in speed and generation performance of EDT, we conducted compar- isons with both classical methods [4] [ 21 and recent advancements B3 39 40).
Experiment on ImageNet 256x256 The result of image generation without classifier-frs free guidance shown in Table Our comparisons across three different sizes demonstrate that EDT consistently achieves the best FID scores: EDTSS scored an FID of 34.2, EDT-B scored 19.1, and EDT-XL scored 7.5. Notably, EDT also showed significant reductions in GFLOPs compared to the second-best MDTv2 across all sizes (2.66 GFLOPs 6.07 GFLOPs, 10.2 GFLOPs vs. 23.02 GFLOPs. 51.83 GFLOPs vs. 118.69 GFLOPs). Moreover, EDT exhibited the lowest memory consumption during inference across all three sizes, underscoring the efficiency of our lightweight design. We further investigated the training speed of EDT. Given that both EDT and MDTv2 incorporate additional training strategies, we specifically compared the training speeds these models. Additionally, we assessed the training speed of EDT without the masking training strategy (denoted as EDT*) against other methods. In both scenarios, EDT trained faster than the baseline models. For example, EDT-XL* achieved training speed of 1.49 iter/s, compared to 0.93 iter/s for DIT-XL. In comparison MDTv2-XL. which trained 0.51 iter/s. EDT-XL was nearly twice fast at 0.98 iter/s. We further perform the experiment on the image generation with classifier-free guidance. The result is shown score
Table 13 Appendix I Under the same training cost, EDT-S-G achieves the lowest compared to MDTv2-S-C (9.89 vs. 15.62). Overall, these findings confirm that EDT significantly enhances both the speed and performance of image synthesis. Additionally, the training cost of EDT is efficient. We include training cost analysis of EDT in Appendix A.2.3

---

Experiment on ImageNet 512x512 As shown in Table 2 we train DiT-S, MDTv2-S and EDT-S on ImageNet 512x512 for 60 epochs. Under the same training cost, both MDTv2-S and EDT-S achieve better FID scores than DiT-S (51.16 vs. 85.21, 51.84 vs. 85.21). 1). In terms of training speed and inference overhead, EDT-S is 3.07 and 2.37 times faster than MDTvv--S respectively. This indicates that EDT achieves competitive image generation performance with lower resource overhead.
Table 3: Results on various models with (w) AMM and without (w/o) AMM. These models are trained for 400k iterations by default. We evaluate models using FID scores.
Model
W/o AMM
AMM
Model
Wo AMM
W AMM
EDT-S"
50.90
38.73
EDT-S
46.90
34.27
EDT-B"
33.19
23.19
EDT-B
26.30
19.18
EDT-XL"
14.92
10.48
EDT-XL
12.80
7.52
DIT-S
67.16
63.11
MDTv2-S
39.02
31.89
DiT-XL
18.48
14.73
DiT-XL-7000K
9.62
3.75
3.3
Ablation Study
3.3.1
Attention Modulation Matrix
Generated images using EDT-XL without AMM
Generated images using EDT-XL with AMM
A
Figure 6: EDT-XL with AMM achieves more realistic visual effects. Area There are some blue stains on the panda's 's arm, Area B: An unreasonable gray area. Area C: Black smoke in the red fog. Area Unrealistic eyes the fox. Area E: Fish with an odd shape. The parrot image generated by EDT-XL without AMM is realistic. And the parrot image generated by EDT-XL with AMM remains equally realistic. The add of AMM does not negatively affect the original quality.
Quantitative analysis We demonstrate the effectiveness and broad applicability of AMM across various models by comparing the FID scores between models with AMM and without AMM in Table Extensive results show that the pre-trained models with AMM consistently outperform models without AMM, thereby verifying the generality and effectiveness of AMM. For instance, MDTv2-S with AMM achieves better FIDD score than MDTvv-S without AMM (31.89 vs. 39.02). Using AMM enhances the FID of DiT-XI. from 18.48 to 14.73. EDT-XL also has a lower FID score of 7.52 compared to EDT-XI. without AMM of an FID score of 12.8.
Qualitative analysis We demonstrate the effectiveness of AMM by comparing the synthesis images from EDT-XL and DiT-XL with and without the AMM. As shown in Figure the red boxes highlight the unrealistic regions in the images generated by EDT-XL without AMM. In the corresponding regions 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. This visual analysis demonstrates the effectiveness of the AMM plugin. Please refer to A.3, and A.5.2 for more analysis about AMM. While AMM is effective, there potential for improvement. Please refer to A.6 details regarding its limitations.
3.3.2 Lightweight-desiin diffusion transformer
We investigate the effectiveness of the key components in our proposed diffusion transformer architec- ture. We denote the token information enhancement TIE and the positional encoding supplement

---

Table 4: The ablation study of the key components of the lightweight-desige and masking training strategy of EDT. The experiment conducted on the small-size EDT model (Wo AMM).
masking training strategy of MDT
masking training
Model
TIE
PES
FID
IS1
strategy EDT
Baseline
x
53.90
29.29
A
x
x
X
52.76
29.38
B
/
x
X
X
52.13
30.60
C
/
/
X
X
50.90
31.02
D
/
X 49.60
33.11
E
/
X
V
46.90
35.40
as PES in Table 4 Model A incorporates TIE without PES, with an FID score 52.8. Model integrates PES without TIE, with FID score of 52.1. Model C represents EDT-S*, and utilizes both components, with an FID score 50.9. Upon comparing Models A and C, we observed that the usage of token information enhancement leads to an improved FID score, decreasing from 52.8 to 50.9. Similarly, the comparison between Models B and C demonstrates that the addition of a positional encoding supplement also results in better FID score, reducing from 52.1 to 50.9. The experimental results confirm the effectiveness of both components in enhancing model performance.
3.3.3 Token relation-enhanced masking training strategy
We investigate the effectiveness of the token relation-enhanced masking training strategy and compare with the training strategy used MDT in Table 4 Model C does not employ any masking training strategy, with FID score of 50.9 and an IS score of 31.0. Model D, small-size EDT trained using the masking strategy of MDT, with an FID score of 49.6 and an IS score of 33.1. Model D shows only a slight improvement compared Model C. Model E is small-size EDT trained with the masking strategy of EDT, with an FID score of 46.9 and an IS score of 35.4. Model achieved the best performance in terms of both FID) and IS. This result suggests that the masking training strategy of EDT successfully improves performance by enhancing the learning ability of token relations.
Conclusions
In this work, we propose the Efficient Diffusion Transformer (EDT) framework, which includes lightweight-desige of diffusion transformer, a training-free Attention Modulation Matrix (AMM) inspired by human-like sketching, and the token relation-enhanced masking training strategy. Our lightweight-design reduces the number of tokens through damn-sampling to lower computational costs. We redesigned down-sampling module and masking training strategy to address token information loss caused by the reduction of tokens. During inference, we introduce local attention through AMM, further enhancing image generation performance. Extensive experiments demonstrate that the EDT surpasses existing SOTA methods in both inference speed and image synthesis performance.
5 Related Work
Diffusion Probabilistic Models: Denoising diffusion probabilistic models (DDPM) [. have marked significant advancement in generative models. DDPM improves image generation by progressively reducing noise. ADM 4 innovates further by introducing a classifier-guided approach to refine the balance between image diversity and fidelity. Subsequent developments include classifier-free method [9]. which increases the flexibility of diffusion models by eliminating classifier constraints. DIT (21 replacing U-Net with transformer in LDM (. achieving superior scalability. However, transformers- based models are computationally intensive. Efficient Diffusion: Various methods have been developed to enhance the efficiency of diffusion models. DDIM [10] redefines the diffusion process as non Mon-Markovian, speeding up sampling by removing dependencies on sequential time steps DDPM [ LDM [I]] reduces computational demands by transforming high-resolution images into a latent space for diffusion, thus balancing complexity with image detail. Current research in lightweight diffusion transformers is limited but offers potential for further efficiency improvements in diffusion model technologies. For more related work, please refer to Appendix A.]
10

---

References
[1] Jonathan Ho, Ajay Jain, and Picter Abbeel. Denoising diffusion probabilistic models. Advances in neural information processing systems, 33:6840-6851, 2020.
[2] Yang Song and Stefano Ermon. Generative modeling by estimating gradients of the data distribution. Advances in neural information processing systems, 32, 2019.
[3] Yang Song. Jascha Sohl-Dickstein, Diederik P Kingma. Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-basee generative modeling through stochastic differential equations. arXiv preprint arXiv:2011.13456. 2020.
[4] Prafulla Dhariwal and Alexander Nichol. Diffusion models beat gans on image synthesis. Advances in neural information processing systems, 34:8780- 14:8780-8794, 2021.
[5] Minjie Zhu, Yichen Zhu, Jinming Li, Junjie Wen, Zhiyuan Xu, Ning Liu, Ran Cheng, Chaomin Shen, Yaxin Peng. Feifei Feng, et al. Scaling diffusion policy in transformer to billion parameters for robotic manipulation. arXiv preprint arXiv:2409, 14411, 2024.
[6] Alex Nichol, Prafulla Dhariwal, Aditya Ramesh, Pranav Shyam, Pamela Mishkin, Bob McGrew, Ilya Sutskever, and Mark Chen. Glide: Towards photorealistic image generation and editing with text-guided diffusion models. arXiv preprint arXiv:2112. 2021.
[7] Chitwan Saharia, William Chan, Saurabh Saxena, Lala Li, Jay Whang, Emily I. Denton, Kamyar Ghasemipour, Raphael Gontijo Lopes, Burcu Karagol Ayan, Tim Salimans, et al. Photorealistic text-to- image difffusion models with deep language understanding, Advances in neural information processing systems, 36494, 2022.
[8] Aditya Ramesh, Prafulla Dhariwal, Alex Nichol, Casey Chu, and Mark Chen. Hierarchical text-conditional image generation with clip latents. arXiv preprint arXiv:2204.06120 1(2):, 2022.
[9] Jonathan Ho and Tim Salimans. Classifier-free diffusion guidance. arXiv preprint arXiv:2007112598 2022.
[10] Jiaming Song, Chenlin Meng, and Stefano Ermon. Denoising diffusion implicit models. arXiv preprint arXiv:2010. 020000, 2020.
[11] Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Bjöm Ommer. High-resolution image synthesis with latent diffusion models. In Proceedings of the IEFF/CVF conference on computer vision and pattern recognition, pages 10684-10695, 2022.
[12] Olaf Ronneberger. Philipp Fischer, and Thomas Brox. U-net: Convolutional networks for biomedical image segmentation. Medical image computing and computer-assisted intervention -MCCAI 2015: 18th international conference, Munich, Germany, October 5-9, 2015, proceedings, part III 18, pages 234-241. Springer, 2015.
[13] Ashish Vaswani. Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention all you need. Advances in neural information processing systems, 30, 2017.
[14] Jacob Devlin, Ming-Wos Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirec- tional transformers for language understanding. arXiv preprint 1810:181000400 2018.
[15] Tom Brown, Benjamin Mann. Nick Ryder. Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Language models are few-shot learners. Advances in neural information processing systems, 33: 197119999 2020.
[16] Nicolas Carion, Francisco Massa, Gabriel Synnaeve. Nicolas Usunier, Alexander Kirillov, and Sergey Zagoruyko. End-to-end object detection with transformers. In European conference on computer vision, pages 213-229. Springer, 2020.
[17] Sixiao Zheng, Jiachen Lu, Hengshuang Zhao, Xiatian Zhu, Zekun Luo, Yabiao Wang, Yanwei Fu, Jianfeng Feng, Tao Xiang, Philip HS Torr, Rethinking semantic segmentation from sequence-te peruuence-t-sseeee so-ueueeeee perspective with transformers. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pages 6881-6890, 2021.
[18] Alec Radford, Jong Wook Kim, Chris Hallacy, Aditya Ramesh, Gabriel Goh, Sandhini Agarwal, Girish Sastry, Amanda Askell, Pamela Mishkin, Jack Clark, et al. Learning transferable visual models from natural language supervision. In International conference on machine learning, pages 8748 8763. PMLR, 2021.

---

[19] Alexander Kirillov, Eric Mintun, Nikhila Ravi, Hanzi Mao, Chloe Rolland, Laura Gustafson, Tete Xiao, Spencer Whitchead, Alexander C Berg. Yen Lo, et al. Segment anything. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pages 4015-4026, 2023.
[20] Fan Bao, Shen Nic, Kaiwen Xue, Yue Cao, Chongxuan Li, Hang Su, and Jun Zhu. All are worth words: A vit backbone for diffusion models. In Proceedings of the IEEE//VE Conference on Computer Vision and Pattern Recognition. pages 22669-22679 2023.
[21] William Peebles and Saining Xie. Scalable diffusion models with transformers. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pages -205,, 2023.
[22] Shanghua Gao, Pan Zhou, Ming-Ming Cheng, and Shuicheng Yan. Masked diffusion transformer is strong image synthesizer. In Proceedings of the IEEE/CVF International Conference Computer Vision, pages 23164-23173, 23172, 2023.
[23] D Man and A Vision. A computational investigation into the human representation and processing of visual information. WH San Francisco: Freeman and Compony, San Francisco, 1:1, 1982.
[24] Jonathan Fish and Stephen Scrivener. Amplifying the mind's eye: sketching and visual cognition. Leonardo. 23(1):117-126, 1990.
[25] Mathias Eitz, James Hays, and Mare Alexa. How do humans sketch objects? ACM Transactions on graphics (TOG), (10)) 2012.
[26] Diederik Kingma and Max Welling, Auto-encoding variational bayes. arXiv preprint arXiv:1312.6111 2013.
[27] Ethan Perez, Florian Struh, Harm Vries, Vincent Dumoulin, and Aaron Courville. Film: Visual reason- ing with a general conditioning layer. In Proceedings of the AAAI Conference on Artificial Intelligence. 2018.
[28] Aaron Hertzmann. Toward modeling creative processes for algorithmic painting. arXiv preprint 1oXiv:2255001600 2022.
[29] Marion Botella, Franck Zenasni, and Todd Lubart. What are the stages of the creative process? what visual art students are saying. Frontiers in psychology. 9:389647, 2018.
[30] Kees Dorst and Nigel Cross. Creativity in the design process: co-cvolution of problem-solution. Design studies, 22(5):425-433 2001.
[31] Geoffrey M Boyaton. Attention and visual perception. Current opinion in neurobiology. 15(4):665 2005.
[32] Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A large-scale hierarchical image datahase. In 2009 IEEE conference on computer vision and pattern recognition, pages 248-255. leee, 2009.
[33] Shanghua Gao, Pan Zhou, Ming- Ming-Ming Cheng, and Shuicheng Yan. Mdtv2: Masked diffusion transformer is a strong image synthesizer. arXiv preprint arXiv:2303. 14389, 2023.
[34] Xingyu Xie. Pan Zhou, Huan Li, Zhouchen Lin, and Shuicheng Yan. Adan: Adaptive nesterov momentum algorithm for faster optimizing deep models. arXiv preprint urXiv:2208.06677 2022.
[35] Martin Heusel, Hubert Ramsauer, Thomas Unterthiner, Bernhard Nessler, and Sepp Hochreiter. Gans trained by two time-scale update rule converge to a local nash equilibrium. Advances in neural information processing systems, 30, 2017.
[36] Charlic Nash, Jacob Menick, Sander Dieleman, and Peter W Battaglia. Generating images with sparse representations. arXiv preprint arXiv:2103.03841 2021.
[37] Tim Salimans, Ian Goodfellow, Wojciech Zaremba, Vicki Cheung, Alec Radford, and Xi Chen. Improved techniques for training gans. Advances in neural information processing systems, 29, 2016.
[38] Tuomas Kynkäänniemi, Tero Karras, Samuli Laine, Jaakko Lehtinen, and Timo Aila. Improved precision and recall metric for assessing generative models. Advances in neural information processing systems, 32, 2019.
[39] Rui Zhu, Yingwei Pan, Yehao Li. Ting Yao, Zhenglong Sun, Tao Mei, and Chang Wen Chen. Sd- dit: Unleashing the power of self- discrimination in diffusion transformer. arXiv preprint arXiv:240311004, 2024.
12

---

[40] Katherine Crowson, Stefan Andreas Baumann, Alex Birch, Tanishq Mathew Abraham, Daniel Z Kaplan. and Enrico Shippole. Scalable high-resolution pixel-space image synthesis with hourglass diffusion transformers. arXiv preprint arXiv:240.11605, 2024.
[41] Alexey Dosovitskiy. Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai Thomas Unterthiner. Mostafa Dehghani, Matthias Minderer, Georg Heigold. Sylvain Gelly, et al. image is worth 16x16 words: Transformers for image recognition at scale. arXiv preprint urXiv:2010.11929, 2020.
[42] Tim Salimans and Jonathan Ho. Progressive distillation fast sampling diffusion models. arXiv preprint arXiv: 2022.
[43] Chenlin Meng. Robin Rombach, Ruiqi Gao, Diederik Kingma, Stefano Ermon, Jonathan Ho, Tim Salimans. On distillation of guided diffusion models. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 14297-14306, 2023.
[44] Hongkai Zheng, Weili Nie, Arash Vahdat, and Anima Anandkumar. Fast training of diffusion models with masked transformers. arXiv preprint arXiv:2306. 09305, 2023.
[45] Eslam Mohamed Bakr, Liangbing Zhao, Vincent Tao Hu, Matthieu Cord, Patrick Perez, and Mohamed Elho- sciny. Toddlerdiffusion: Flash interpretable controllable diffusion model. arXiv preprint arXiv:2311.1454,, 2023.
[46] Ze Liu, Yutong Lin, Yue Cao, Han Hu, Yixuan Wei, Zheng Zhang, Stephen Lin, and Baining Guo. Swin transformer: Hicrarchical vision transformer using shifted windows. In Proceedings of the IEEE/CVF international conference computer vision, pages 10012-10022, 2021.
[47] Syed Waqas Zamir, Aditya Arora, Salman Khan, Munawar Hayat, Fahad Shahbaz Khan, and Ming-Hsuan Hssun Yang. Restormer: Efficient transformer IFFF/CVF conference on computer vision and pattern recognition, pages 5728-5739, 2022.
nigh-resolution image restoration. Proceedings of
[48] Hu Cao, Yueyue Wang. Joy Chen, Dongsheng Jiang, Xiaopeng Zhang, Qi Tian, and Manning Wang. Snin-unet: Unet-like pure transformer for medical image segmentation. In European conference on computer vision, pages 205-218. Springer, 2022.
13

---

A Appendix
The appendix is structured as follows. In Appendix A.2
we provide computational complexity analysis of DiT and the lightweight architecture design of EDT. In Appendix AA.3] we show the computational process of the AMM and explore the usage methods of AMM, ppper-parameters... In Appendix A.4 we analyze the loss changes of different masking training strategies and explore the selection of mask ratio. In Appendix AA.5 we report additional experimental results and visualization
of generated images by EDTT-XL-200KK. Appendix A the limitations.
A.1 Related Work
Diffusion Probabilistic Models In the field of generative models, diffusion models have achieved significant success. The Denoising diffusion probabilistic model (DDPM) [T] is the most classic denoising diffusion model, which generates new images by gradually removing noise from the noising data. NCSN [2] guides the data from low to high probability density areas by predicting the gradient of the data distribution. SDE [3] considers the noise addition and removal process as continuous and uses stochastic differential equations to guide the data generation process, theoretically unifying DDPM and NCSN and providing new perspective for generative models. ADM [] introduces Massifier-guided mode, using classifier gradients to guide the diffusion model's generation process, balancing the diversity and fidelity of generated images. Building on ADM, [9] proposes a classifier- free guidance method, removing the constraints of the classifier and significantly enhancing the flexibility of conditional diffusion generation. Compared to earlier works based on the U-Net backbone, models based on the Transformers backbone achieved better results. U-ViT [20] adds U-Net's skip connections to VIT [1] for diffusion generation in the pixel space, showcasing the broad prospects for diffusion transformers. DiT [2] replaces the U-Net backbone with VIT based on LDM [1], surpassing previous U-Net-based performances and showing good scalability. However, due to the intensive computational nature of transformers, integrating them effectively into diffusion models still presents significant challenges.
Efficient Diffusion Several methods have been proposed previously to improve the inference effi- ciency of diffusion models. DDIM [10] defines the diffusion process non- Markov process, thus eliminating the dependency on adjacent time steps in the reverse process of DDPM(]. achieving faster sampling speeds. LDM [I]] compresses images from high-resolution pixel space to latent space and performs diffusion generation in latent space, achieving a balance between computational complexity and image detail. 42 introduces a method of gradual distillation for unconditional and classifier- guided diffusion models, which optimizes higher iteration counts into lower ones. Building on the work of [421. [3 proposes two-stage distillation method for classifier-free guided diffusion models. Studies [22] 44] improves the training speed of DiT [21] by masking inputs. HDiT [40] introduces hourglass-srructured diffusion transformer for generating high-resolution images, which requires 70% fewer inference FLOPPs in pixel space compared to DiT (1). but since HDiT operates in pixel space with high input and output resolution, still demands substantial computational resources. ToddlerDiffusion (45) proposes a novel approach that extends the diffusion framework into modality space, decomposing the complex task of RGB image generation into simpler, interpretable stages, Currently, research on lightweight diffusion transformers is relatively sparse, but it is orthogonal to the methods in [T] 42 43 Building on existing research, lightweight diffusion transformers will further enhance the efficiency of diffusion models, which is one of the core contributions of this paper.
dlassifierguidee
A.2 Computational analysis and lightweight design
We preset two lightweight design rules: (1) to reduce the FLOPs in the self-attention module, we decrease the number of tokens by token down-sampling; (2) to guarantee the total FLOPs significant reduction, the FLOPs of the EDT blocks after down-sampling should be significantly reduced compared to the EDT blocks before down-sampling Utilizing the down-sampling module is key to achieving these two design rules.
We first analyze the applicable scenarios of conventional down-sampling modules in Appendix A.2. In Appendix (2.222 we improve the down-sampling module to adapt to our scenario of latent difrusion models. Then, in Appendix A.2.3 the training costs of EDT, DiT, and MDT are reported.
14

---

Table 5: The model details of EDT across three different sizes.
Number of blocks
Blocks in each stage
Dimensions
Heads in each stage
Model
Params.
in each stage
[312,416, 520,416,312]
[6,8,
EDT-S/2
32,2M+6. IM
12
[2,2,2,3,3]
10,8,6]
[624,832,
[12,16,
EDT-B/2
128M+24.
12
[2,2,2,3,3]
1040,832.624)
20,16,12]
[936,1248,
[18.24.
EDT-XL/2
644M+54.4M
28
[6,4,4.7.7]
1560,1248,936)
30,24,18]
A.2.1 Applicable scenarios of the conventional down-sampling module
Firstly, we review the conventional down-sampling module. We have token sequence of shape N x N, where n is the number of tokens (n = N²). The dimension of token is denoted by d. The number of heads in multi-head attention is h. and D denotes the dimension of a head. Conventional 48) reduces the number of tokens by factor of k22 and increase the token dimensions by factor of k at the same time, where is is down-sampling factor (k 2 2). Large down- wown-sampling factor & will cause too many tokens to be merged, harming performance. Therefore, is generally equal of tokens n by down-sampling adjacent 2 tokens into one. Then we obtain 2 x N token sequence with fewer tokens and the token dimensions increase 2d from d. Table [ shows FLOPs analysis of DiT block. The total FLOPs F' is 2ndd + 12nd + 6ef and the number of parameters P is 18dl in a DiT block.
down-sampling module 46 47
That means, a N N token sequence, we reduce the number
Nv
Now we analyze how much can down-sampling module reduce FLOPs, and its applicable scenarios. To facilitate derivation and calculation, we set j = d. where j is the proportional coefficient between the number of tokens and the dimension of token. We explore the relationship between proportional coefficient and FLOPs drop ratio p after using conventional down-sampling.
In Figure 7 before down-sampling, the total FLOPPs of a DiT block F is 2j + 12jd + 6d2, where 2 and j = 7- After feeding tokens into conventional down-sampling, the number of tokens n' is and token dimensions d' is 2d. Then feeding these tokens into a DiT block, the total FLOPs of the DiT block becomes F' = 2nn + 12n'd + 6d ndd + 12nd + 24d and the number of parameters in the DiT block is p' p 72d
52d +12jd + 24d
Comparing DiT blocks after and before the down-sampling module, the FLOPs drop is VF F F- d 18d2. We calculate the relationship between the FLOPs drop ratio 7 and the proportional j:efficient
FF F
p
18d
23 + 12jd3 + 6d
7j2d 72
8j2d + 48jd +
(4)
7j
8j2d + 48jd
7j
8j + 48
7
+
The Eqn shows that 8/+-18 the upper limit of the FLOPs drop ratio p and proportional to j. Significant reductions in FLOPs can be achieved through down-sampling, only when the number of tokens n is larger than the token dimensions d, namely j > 1. For instance, when
15

---

Table 6: FLOPs in DiT block
Module
Operator
Input Shape
Params
Output Shape
FLOPs
AdaLN
d
6d
6d
6d
Att-kqy
1 R X d K: h
x 3d
1 X X 3d
3nd
D Q: X h X D x
1 n n
Attention
K@Q
dn?
(d= hD)
V
Att: h n X
1 (d = hD)
n
dn²
V: x h X D
1 n d
dxd
1 X n
ndf
fel
X n X d
d 4d
1 1 x 4d 1 X
4nd
FFN
fc2
1 n x 4d
4d x
4nd
Total
1 n x d
18d2
1xnxd
2n + 12nd + 6d
Input shape NxNxd
Input shape N NxNxd
Input shape NxNx d
Input shape N 2
Input shape
x 2d
x 2d
Down Sampling
block A
block A
block B
block
FLOPs F
FLOPs F
FLOPs F'
FLOPs F'
= dd + 12jd + 6d
F' f
N2 ) 12jd 24d (j
Figure 7: Input shape and FL.OPs of DiT block before and after the conventional down-sampling module.
j (or 6d), the FLOPs drop ratio p about 50%; and when j 1 (or n d), the FLOPs drop ratio is about 12.5%.
However, in our application scenarios of latent diffusion transformer, where 3 = < 1, the FLOPs drop ratio is smaller than 12.5%. Only using conventional down-sampling in our scenarios of latent hardly reduce the computational complexity of the DiT blocks. So it does not meet the design rule (2): the FLOPs of the blocks after down-sampling should be significantly reduced compared the blocks before down-sampling, to guarantee the total FLOPs reduction.
diffusion transformers
A.2.2 Redesign the down-sampling module
Now, we redesign the down-sampling to make the architecture meet rule (2).
According to Appendix A.2.1 before the down-sampling module, the total FLOPs of a DiT block is F = 2n²d + 12nd + 6d = 2 + 12jd + 6d². And after the down down-sampling module, the total FLOPs of DiT block after down- dawn-sampling is F = =² ndd of F and F' the second term 12nd predominates due to d by the down-sampling process. This results from the conventional down-sampling module, which reduces the number of tokens by factor of k22 and increases the token dimensions by factor of & at the same time, namely 12nd 12 x F X (kd)²
+ 12nd 24d. Among the three items However, this term isn't affected
To reduce the second item 12nd? we should redesign the process of down- wown-sampling module: we reduce the number of tokens by down-sampling factor à 2 and increase the token dimensions by a factor of r at the same time, where r is token dimension expansion cocfficient and < r < This design makes the second item reduce, namely 12
(rd)² 12nd? and the total FLOPs ri 3jr + 6r dl in the blocks after our down-sampling
F' rn?d 2d + 2nd 2r
16

---

module. The FLOPs drop ratio p is:
VF
F F F
p
F
rn'd + 3'rr2 dr
2ndd + 12nd + 6d
rn² 24nddr2 + 48dr
16n² 96nd + 48d
rj 24r + 48r²
16j + 96 + 48
rJ + 24,
16j + 96
(5)
j + 24r
16j + 96
j + 48
16j 96
16
672
3+418
2 r
16
672
14-48
0.43r
16j+96 shows that the FLOPs drop ratio is inversely proportional to the token dimension expansion coefficient and - 0.43r the lower bound of p. We can adjust the FLOPs of the DiT block after down-sampling by adjusting the token dimension expansion coefficient r. When is within the range [1, k] and 2, the FLOPs drop ratio p falls within the interval [0.57, 0.14]. To make the number of parameters in EDT network approximate to that other works (DiT and MDT), set 1.25. Namely, after down-sampling, we reduce the number of tokens by factor of 4 and increase the token dimensions by factor of 1.25 at the same time. This leads to about 47% FLOPs drop ratio of the block after down-sampling compared to that before down-sampling, which meets the requirement of rule (2).
In Equ p
+)+21
Table Ba shows three diffferent sizes of EDT. The number of blocks is consistent with that in the DiT network of the corresponding size. The number of parameters in the EDT network is also approximate to that in DIT. Our parameters consist of two parts, one is the EDT blocks parameter, and the other is the sampling module and long-skip connection module. In the table, the number of parameters is written the sum of these two parts.
Figure illustrates the architecture our lightweight-designed diffusion transformer. The model includes three EDT stages in the down-sampling phase, viewed as encoding process where tokens are progressively compressed, and two EDT Stages in the up-sampling phase, viewed as a decoding process where tokens gradually reconstructed. These five EDT stages are interconnected through down-sampling, up-sampling, and long skip connection modules. The 'blocks each stage" in 5 displays how many blocks there are in the corresponding EDT stage.
A.2.3 Training costs
On ImageNet, we estimated the training cost of EDT, MDTv2, and DiT on 48GB-L40 GPU in Table using batch size of 256 and FP32 precision. EDT achieves the best performance with a low training cost. GPU days refer the days required for training the models on single L40 GPU.
A.3 Exploring Attention Modulation Matrix
During the sketching process, humans alternately use global attention and local attention. Therefore, design Attention Modulation Matrix (AMM), to introduce local attention into the self-attention
17

---

Table 7: Training cost of EDT, MDTv2, and DiT on ImageNet
Model
Resolution
Epochs
Cost (Images)
GPU days
GFLOPs
FID
EDT-S*
256x256
80
102M
2.75
2.66
38.73
DiT-S
256x256
102M
2.96
6.06
68.40
MDTv2-S
256x256
80
102M
16.47
6.07
39.50
EDT-B
256x256
80
102M
9.19
10.20
19.18
DiT-B
256x256
80
102M
8.62
23.01
43.47
MDTv2-B
256x256
80
102M
26.17
23.02
19.55
EDT-XI
256x256
80
102M
37.79
51.83
7.52
DiT-XL
256x256
80
102M
39.82
118.64
19.47
MDTv2-XL
256x256
80
102M
72.62
118.69
7.70
EDTS
512x512
60
77M
19.02
13.25
51.84
DiT-S
512x512
60
77M
12.26
31.42
85.21
MDTv2-S
512x512
60
77M
51.96
31.46
51.16
Number of tokens
dy
(2t 2x,
(sk M6)
Token distance matrix
Calculate the Euclidean distance between tokens
N x N
di
Mar
0,
der > R
Tokens B x
Attention Modulation Matrix
N x N
Expand
Attention matrix 8 H x 2 x R
Attention Modulation Matrix N x N
Hadamard product
8: batch size number heads head dimensions number tokens row or column
i or r:: the index of token (i, [1,2, N?))
N2))
(Xp yı) coordinate of token-i in N x N grid
Modulated attention score matrix B x N x N
d distance between token-i and token-r
R: effective radius of local attention
Figure 8: The process of modulating the attention score matrix and the changes in tensor shape.
module which uses global attention by default. Appendix A.3.1 shows the process of modulating the attention score matrix and the changes in tensor shape. In Appendix [A.32] we explore the usage and arrangement of AMM in models. In Appendix A.3.3 explore the settings of the hyper-parameters of AMM. The computational cost of AMM is discussed in Appendix A.3.4
A.3.1 The process of modulating the attention
The process of modulating the attention score matrix and the changes in tensor shape are shown in Figure Image can be split into 2 patches and each token is the feature of a patch. Fach token (patch) corresponds
rectangular area of the image and has corresponding 2-D coordinate (x, y) in the image grid. We calculate an Euclidean distance value d for each pair of tokens, resulting in distance matrix D, which is an 2 2 tensor. Based on the distance matrix, generate modulation values m via the modulation matrix generation function F(d), which assigns lower modulation values tokens that are farther apart. These modulation values form an Attention Modulation Matrix (AMM), another N x 2 tensor. Importantly, we integrate the AMM into the pre-trained EDT without any additional training. The attention modulation matrix calculated when the model is instantiated. During inference, the modulated attention score matrix is obtained by performing Hadamard product between the attention modulation matrix and the attention score matrix.
18

---

Table 8: Comparison of adding AMM into EDT-S* during training versus inference on ImageNet 256 256.
adding AMM
adding AMM when inference
Model
FID
1ST
when training
A
50.9
31.0
B
V
52.1
30.3
X
38.7
36.4
Table 9: Performance of EDT-S* with varying insertion points of AMM on ImageNet 256 x 256.
AMM in encoder
AMM in decoder
Alternately inserting
Model
FID
1St
A
X
x
50.9
31.0
EC
B
V
/
60.4
24.9
X
X
44.4
35.1
D
X
38.7
36.4
A.3.2 When and where to use AMM in EDT
This section empirically discusses and demonstrates how to use AMM through experiments.
Just adding AMM into the pre-trained model when inference Table 8 shows when to use AMM in EDT. Model A is a pre-trained EDT-S* without AMM, getting FID score of 50.8. Model B, which is added AMM from initialization and then trained, achieves an FID score of 52.1. Model C gains FID score of 38.7, which is the model that added AMM to model A. The results show that using AMM during training will lead to poorer performance, and just adding AMM to the pre-trained model can greatly improve performance.
Alternately inserting AMM into the decoder of EDT We view the process of drawing pictures as decoding process. So we think the decoder of EDT should be aligned with the act of humans drawing pictures. This inspires us where to place AMM in EDT. Namely, the alternation of attention range between global and local in drawing AMM should be alternately or discontinuously inserted into the blocks of decoder.
of humans should be imitated in the decoder of EDT. The
To demonstrate the inspiration, in Figure 9 we try to compare Model A: pre-trained EDT-S* without AMM; Model B: alternately inserting AMM into encoder and decoder pre-trained EDT-S*; Model C: consecutively inserting AMM into decoder of pre- mre-trained EDT-S* and Model D: alternately inserting AMM into the decoder pre-trained EDT-S* As shown in Table 9 Model D which is alternately inserted AMM into decoder of pre-trained EDT-S*, gets the lowest FID
When integrating AMM into a pre-trained model, the best arrangement AMM in blocks varies across different models. Identifying the optimal placement and configuration of AMM requires testing and adjusting to realize its full potential.
A.3.3 The determination of hhper--paameen parameter
Table 10 shows the determination of hyper-parannn parameter about the effective radius local attention R E (0, dmaz], where dmaa = V2(N - 1) is the farthest distance x token grid. In the table, when R = /NV 1)2 4, *DT-S* achieves the lowest FID and highest Prec. So we set R = (NN 1)2 + 4 in all sizes of model.
19

---

Encoder
block
block
block
block
block
A
FID
block
block
block
block
block
50.9
Decoder
Encoder
block with AMM
block with AMM
block
block
block
B
FID
block with AMM
block with AMM
block
block
block
60.4
Decoder
Encoder
block
block
block
block
block
C
FID
block with AMM
block with
block
block
block
44.4
AMM
Decoder
Encoder
block
block
block
block
block
D
FID 38.7
block with AMM
block with AMM
block
block
block
Decoder
Figure 9: Different arrangement forms of AMM in EDT and their corresponding FID scores.
Table 10: Exploring the value of the effective radius of local attention in EDT-S* * for 256 256 resolution.
R
FIDSOK!
ISt
sFID
Prec.
Recallf
no AMM
50.9
31.0
13.3
0.427
0.604
2(
1)
38.9
36.6
9.3
0.472
0.626
(N
1)2 +
38.7
36.4
9.3
0.473
0.623
(N 1)2 +1
39.0
36.5
9.3
0.474
0.621
N-1
39.1
36.2
9.2
0.472
0.628
3N/4
41.0
35.4
9.6
0.483
0.604
N/2
47.4
32.1
11.7
0.495
0.583
N/4
72.0
23.9
23.6
0.476
0.516
20

---

A.3,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 no additional training costs. Secondly, the increased computational cost of AMM during inference is negligible. For instance, the last block EDT-XL, the attention score matrix and the Attention Modulation Matrix are both 18x256x256 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 merely 0.04% of the total FLOPs, making the computational cost of AMM negligible. In our experiments, we added AMM modules to pre-trained DiT-XL model, and the FID score decreased from 18 to 14.
A.4
Discussion about token relation-enhance masking training strategy
A.4.1 Analysis of the loss of EDT and MDT
MDT
EDT Less cent Meaking transig Strategy
Vakke Denatimm Masking braning Stratagy
2.-
6.-
300000
300000
300000
300000
304000
30009
30000
301006
30000
300000
304000
3040
Love Valuze over Berations Massiing Trainingg Strategy MOT
Lasn Velue over Denation Meaking training Strategy
20-
JM02000
300000
300000
300036
30/00C
300000
403000
Red line is the Lmaskedi Green line is the Lyutt-
Max-Min normalization preprocessing is performed on the loss values.
Figure 10: Comparing the loss changes of different masking training strategies.
In Figure I0 we separately applied the masking training strategies MDT and EDT to train EDT-S and extracted Lmaskeed and L f values at the 300k ~ 305k and 300k ~ 400k training iterations. The left-top of the figure depicts the loss changes when using MDT's masking training strategy. As I.full decreases, T-masked increases, and vice versa, illustrating the conflict between these two losses. This conflict arises because Tymasked in MDT causes the model to focus on masked token reconstruction while ignoring diffusion training. As shown in the bottom-left of the figure, both the Lfull and Lmookeed hardly decreased during the 300k to 400k training iterations. The right side the figure shows the loss changes when using EDT's masking training strategy. The Lmookeed and Lyalt exhibit synchronized changes, and the loss values continuously decrease during the 300k to 400k training iterations.
A.4.2 The determination of the mask ratio
We explore the mask ratio of our masking training strategy. There are two down-sampling modules in EDTs. So we have two positions to implement token masking. As shown in Table [1] we determine the masking ratio in the first down-sampling module by training EDT-S. The best masking ratio in the first down-sampling module is 0.4 0.5.
21

---

Table 11: Mask Ratio in the first down-sampling module.
Mask Ratio
FIDSOK
Ist
sFID
Prec.
Recall
0.1 0.2
51.3
31.1
13.7
0.434
0.612
0.2 ~ 0.3
48.6
33.5
13.2
0,441
0.631
0.3 ~ 0.4
47.3
34.9
13.2
0,442
0.622
0.4 ~ 0.5
46.4
34.1
11.7
0.449
0.621
0.5 0.6
48.8
33.4
13.5
0.431
0.623
Table 12: Mask Ratio in the second down-sampling module. (Based on the 0.4 0.5 mask ratio in the first down-sampling module)
Mask Ratio
FID5OK
ISt
sFID
Prec.
Recallf
0.1 ~ 0.2
45.4
34.8
13.2
0,440
0.621
0.2 0.3
46.2
34.6
13.5
0.441
0.619
0.3 ~ 0.4
46.5
34.7
13.2
0,438
0.620
0.4 0.5
47.2
34.4
13.4
0.432
0.620
Table 13: The comparison with existing SOTA methods on class-conditional image generation with classifier-free guidance on ImageNet 256x256 (CFG=2 in EDT; according to DiT and MDTvv, their optimal CFG settings are 1.5 and 3.8, respectively).
Cost (IterBBS)
Model
GFLOPs]
FID
DiT-S-G
400Kx256
6.06
21.03
MDTv2-S-G
400Kx256
6.07
15.62
EDT-S-G(our)
400Kx256
2.66
9.89
4DM-G
1980kx256
1120.00
4.59
LDM-4-G
178kx1200
104.00
3.60
D-T-XL-G
400Kx256
118.64
5.50
DIT-XL-GF21
7000Kx256
118.64
2.27
MDTV2-XL--G
33
4600Kx256
118.69
1.58
EDT-XL-G(our)
400Kx256
51.83
4.65
EDT-XL-G(our)
1000Kx256
51.83
4.30
EDT-XL-G(our) 2000Kx256
51.83
3.54
At the base of 0.4 0.5 mask ratio in the first down-sampling module, we then determine the mask ratio in the second down-sampling module shown in Table 12] According the results, the best masking ratio in the second down-sampling module is 0.1 ~ 0.2, at the base of 0.4 ~ 0.5 mask ratio in the first down-sampling module.
A.5
Additional Results
A.5.1 Image generation with classifier-free guidance ImageNet 256x256
The result of image generation with classifier-free guidance on ImageNet 256x256 is shown in Table [3 Under the same training cost, EDT-S-G achieves the lowest FID score compared MDTv2- S-G (9.89 vs. 15.62). EDT-XL-G achieves a good balance in training cost, inference GFLOPs, and image generation performance.
A.5.2
Comprehensive evaluation of AMM
Using AMM under different iterations We report the FID of EDT with and without AMM under different iterations in Table [4] The results show that all EDTs with AMM obtain lower FID.
22

---

Table 14: FID of EDTs* under different iterations on Imagenet 256 x 256.
EDT-S*
EDT-B
EDT-XL*
Iterations
no AMM
AMM
no AMM
AMM
no
AMM
AMM
50k
96.4
80.8
79.5
64.8
50.1
45.8
100k
64.5
58.5
46.0
40.4
23.7
21.3
150k
58.2
51.0
39.7
32.6
17.9
15.3
200k
55.8
47.0
37.3
29.0
15.6
12.6
250k
53.3
44.0
35.5
26.9
14.6
11.4
300k
52.4
42.0
34.2
25.1
14.1
10.7
350k
51.2
40.1
33.5
24.2
14.5
10.6
400k
50.9
38.7
33.2
23.2
14.9
10.5
Table 15: Results on various models across different sizes with (w) AMM and without (w/o) AMM on ImageNet.
Model
Resolution
FID
ISt
sFID
Prec. 1
Recall
(W/-)*(w/o)
256 x 256
50.9
31.0
13.3
0.427
0.604
EDT-S*
256 x 256
38.7
36.4
9.2
0.474
0.620
EDT-S(w/o)
256 x 256
46.9
35,4
13.5
0.442
0.624
EDT-S
256 x 256
34.3
42.6
13.1
0.501
0.612
EDT-S(w/o)
512 x 55.7
28.9
12.8
0.513
0.586
EDTS
512 x 512
51.8
29.9
7.9
0.563
0.572
EDT-B*((//)
256 256
33.2
50.0
10.5
0.512
0.648
EDT-B*
256 x 256
23.2
62.3
8.9
0.573
0.624
EDT-B(w/o)
256 x 256
26.3
64.5
10.3
0.544
0,659
EDTB
256 x 256
19.2
74.4
9.9
0.586
0.639
EDT-XL*((/O)
256 x 256
14.9
96.5
8.0
0.617
0.667
EDT-XL**
256 x 256
10.5
117.8
9.9
0.663
0.637
EDT-XL(W/o)
256 x 256
12.8
111.7
8.2
0.627
0.685
EDT-XII
256 x 256
7.5
142.4
7.4
0.684
0.648
DiT-XL/2(W/o)
256 x 256
18.5
71.3
6.1
0.641
0.632
DiT-XL/2
256 x 256
14.7
83.9
10.5
0.720
0.511
More types of evaluation indicators on EDT with AMM The Table 15 shows more types of evaluation indicators comparing EDTs without AMM and with A.M.. EDTs means the EDT models that don't 't use masking training strategy. In the table, when using AMM, the FID, and Precision are all improved. Furthermore, AMM is versatile and can be adapted to various diffusion transformers. For instance, the performance of DiT-XL is improved by AMM (18.5 FID vs. 14.7 FID).
Generated images using DiT-X without AMM
Generated images using DiT-XL with AMM
Figure 11: DÝT-XL-400k with AMM achieves more realistic visual effects.
23

---

Qualitative analysis on DiT-XL with AMM We conduct qualitative analysis on DiT with and without AMM in Figure [T, which demonstrates AMM can also improve the local details of the synthesis images of DiT. The red boxes highlight the unrealistic areas in the images generated by DiT-XI. without AMM. In the corresponding areas of the images generated by DĪT-XI. with AMM, the results appear more realistic. Area A: The otter lacks mouth. Area B: The red panda's mouth slightly distorted. Area C: The parrot's beak is not sharp enough. Area D: There are black stains the right side of the hot air balloon. Area E: The fox's 's eyes are white. The steam image generated by DiT-XL without AMM realistic. And the steam image generated by DiT-XL with AMM remains equally realistic. The addition of AMM does not negatively affect the original quality. The effectiveness of AMM on DiT-XL further demonstrates the universal applicability of AMM.
Figure 12: Visualization of images generated by the EDT-XL-2000K.

---

A.5.3 Visualization
More visualized examples of EDT-XXL-200K generated images are shown in Figure 12 step 250.
The sampling
A.6
Limitations
In this work, we propose the Efficient Diffusion Transformer (EDT) framework, which includes lightweight-design of diffusion transformer, a training Attention Modulation Matrix (AMM), and its alternation arrangement in EDT inspired by buman-like sketching and the token relation-enhancer masking training strategy. EDT effectively improves the training and inference speed of diffusion transformers. Notably, AMM is a new approach to improving diffusion models, with many aspects still worth exploring. When integrating AMM into pre-trained model, the insertion and arrangement of AMM in blocks differ across various models. Thus, identifying the optimal placement and configuration AMM requires testing to unlock its full potential. Moreover, the generation function of AMM still has room for improvement and deserves further exploration. For example, in addition to directly scaling attention scores by AMM, we can convert global attention into local attention during inference through methods like: (1) using attention window to exclude the interactions of tokens with far distance; (2) or scaling the attention based on token distance before the operation softmax in attention modules. It is encouraged to further explore the idea of AMM, which incorporates local attention into transformers during inference, to improve other transformer based models.
25

---

NeurIPS Paper Checklist
The checklist is designed encourage best practices for responsible machine learning research, addressing issues of reproducibility, transparency, research ethics, and societal impact. Do not remove the checklist: The papers not including the checklist will be desk rejected. The checklist should follow the references and precede the (optional) supplemental material. The checklist does NOT count towards the page limit.
Please read the checklist guidelines carefully for information on how to answer these questions. For each question in the checklist:
You should answer [Yes], [No] or [NA].
[NA] means either that the question is Not Applicable for that particular paper or the
relevant information is Not Available.
Please provide short (1-2 sentence) justificatin right after your answer (even for NA).
The checklist answers are an integral part of your paper submission. They are visible to the reviewers, area chairs, senior area chairs, and ethics reviewers. You will be asked to also include (after eventual revisions) with the final version of your paper, and its final version will be published with the paper.
The reviewers of your paper will be asked to use the checklist as one of the factors in their evaluation.
While "[Yes] " is generally preferable to "[No]" ", is perfectly acceptable to answer "[No] "provided proper justification is given (c.g., "error bars are not reported because it would be too computationally expensive" or "we were unable to find the license for the dataset we used"). ). general, answering "[No] or "[NA] is not grounds for rejection. While the questions are phrased in binary way, we acknowledge that the true answer often more nuanced, so please just use your best judgment and write justification to elaborate. All supporting evidence can appear either in the main paper or the supplemental material, provided in appendix. If you answer [Yes] to a question, in the justification please point to the section(s) where related material for the question can be found.
IMPORTANT, please:
Delete this instruction block, but keep the section heading "NeurIPS paper checklist",
Keep the checklist subsection headings, questions/answers and guidelines below.
Do not modify the questions and only the provided macros for your answers.
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: Please refer to the abstract and Sections T
We introduce the problem our work solved and the contributions made in both the abstract and final paragraph of the introduction.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made 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
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 be expected to generalize other settings.
is fine include aspirational goals as motivation long is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors? Answer: [Yes]
26

---

Justification:
We discuss the limitations of this work in Appendix A.6
Guidelines:
The answer NA means that the paper has limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model sell-specifiitiinn 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 few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or speech-to-texi system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and faimesss.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?
Answer: [Yes]
Justification: Please refer Section 3.3 Appendix A.3 and Appendix A.4.I
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that affects the main claims and/or conclusions of the paper (regardless whether the code and data are provided or not)?
Answer: [Yes]
Justification: Please refer Section 3.1 for the training details. Please refer Section 3.2 Section and and Appendix[A.3 for experimental results. And we will release the code and
checkpoints.
Guidelines:
27

---

The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing of a model checkpoint, other means that appropriate to the research performed.
also be provided via detailed
While NeurlPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution is primarily a new algorithm, the paper should make clear how to reproduce that algorithm. (b) If the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome deseribe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
Justification: All the data used is from public available dataset (ImageNet). Please refer to Section 3.1 And will release code and checkpoints.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://nips. cc/
for more details.
While we encourage the release of code and data, we understand that this might not be possible, "No" an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https
/ //nis.cc/public//uust http://////iiiee//edeuuubiiiiiii
for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
28

---

At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code is permitted.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: Please refer to Appendix A.3.3 and Appendix A.4.2 for hyperparameter selection. Please refer Section 3.1 for the training and details.
Guidelines:
The answer NA means that the paper does include experiments.
The experimental setting should be presented the core the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, supplemental material.
7. Experiment Statistical Significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [NA]
Justification: The training procedure for the proposed and baseline models takes consider- able amount of time. It is infeasible to perform enough repeated experiments calculate the statistical significance. Therefore, error bars are not reported.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.
The factors of variability that the error bars capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall
run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
should be clear whether the error bar is the standard deviation the standard error of the mean.
is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).
If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification: Please refer Section 3.1 and Appendix A.2.3
29

---

Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make into the paper).
9. Code Of Ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics
Answer: [Yes]
Justification: The research conducted in the paper conforms, in every respect, with the NeurlPS Code of Ethics.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special cireumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g.. if there is special consid- eration due to laws or regulations their jurisdiction).
Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [Yes]
Justification: In Appendix A.6 we discuss the positive impacts of our efficient transformer- based diffusion model design and Attention Modulation Matrix, and the limitations of Attention Modulation Matrix.
Guidelines:
The answer NA means that there no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no 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 (c.g., deploymeen of technologics 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 tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional unintentional) misuse of the technology.
there are negative societal impacts, the authors could also discuss possible mitigation strategies (c.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
30

---

Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification:
The paper poses no such risks.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow controlled use of the model, for example by requiring that users adhere usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Question: Are the creators or original owners of assets (c.g., code, data, models), used in the paper. properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification: The creators or original owners assets (e.g., code, data, models), used in the paper. properly credited and are the license and terms of use explicitly mentioned and properly respected.
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website). the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswi thcode. com/ datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that re- re-packaged, both the original license and the license of the derived asset (if has changed) should be provided.
If this information is not available online, the authors are encouraged to reach out to the asset's creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification:
The paper does not release new assets.
Guidelines:
The answer NA means that the paper does release assets.
Researchers should communicate the details the dataset/code/moodll as part 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.
31

---

At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification:
The paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Question:
Answer: [NA]
Justification: The paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does involve crowdsourcing nor research with human subjects.
Depending the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurlPS 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.
32


Reviews:
[{"review_id": "7Kjt7Tu7vi", "reviewer": "Reviewer_FhNv", "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.\nSecondly, 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.", "questions": "* Please show the effectiveness of your method on more datasets, as suggested earlier.\n* Integrate the AMM module into different methods to show its effectiveness as claimed.", "limitations": "* The limitations are not discussed.", "rating": 5, "confidence": 5, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "* Alternating between the local and global details inspired by human brains is novel.\n* Tackling an important application.\n* 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.\n* [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.\n* [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.\n* [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.\n* [Analysis] The analysis regarding the MVD-v2 masking mechanism is shallow. In addition, presenting the loss trends using screenshots of the terminal is shocking.\n* [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.\n* [Visualization] All the Figures must be at the top of the pages to enhance readability.\n* [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.\n\n[1] Bakr, Eslam Mohamed, et al. \"ToddlerDiffusion: Flash Interpretable Controllable Diffusion Model.\" arXiv preprint arXiv:2311.14542 (2023).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We appreciate the professional and insightful comments. We address each comment as follows:\n\n---\n\n***1. Performance improvement by AMM.*** \n\n\n- 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.\n\n---\n\n***2. More models incorporated with AMM.***\n- 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. \n\n|**On ImageNet 256×256 (class-conditional image synthesis)**||||\n|---|---|---|---|\n|Model|Training images|W/o AMM|W AMM|\n|EDT-S|400k * 256|42.60|**34.20**|\n|DiT-S|400k * 256|67.16|**63.11**|\n|MDTv2-S|400k * 256|39.02|**31.89**|\n|EDT-XL|400k * 256|12.80|**7.52**|\n|DiT-XL|400k * 256|18.48|**14.73**|\n|DiT-XL|7000k * 256|9.62|**3.75**|\n|**On CelebA-HQ 256 × 256 (unconditional image synthesis)**||||\n|Model|Training images|W/o AMM |W AMM|\n|EDT-S|100k * 240|17.01|**16.60**|\n|DiT-S|100k * 240|19.12|**18.41**|\n|MDTv2-S|100k * 240|17.34|**17.11**|\n\n**Rebuttal-Table 1**: Evaluation of the performance of pre-trained EDT, DiT, and MDT without and with AMM on ImageNet and CelebA-HQ.\n\n---\n\n***3. Experiments on extra dataset CelebA-HQ.*** \n- 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.** \n\n|Model|Training images|FID|\n|---|---|---|\n|EDT-S|100k * 240|**16.60**|\n|DiT-S|100k * 240|19.12|\n|MDTv2-S|100k * 240|17.34|\n\n**Rebuttal-Table 2**: Evaluation of unconditional image synthesis on CelebA-HQ.\n \n---\n\n***4. Clarification of the speed-ups in training.*** \n- 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.\n\n---\n\n***5. Analysis of masking training mechanism.*** \n- 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.\nIn 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.\nHowever, 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.**\nIn 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.\n\n---\n\n***6. Ablations.*** \n- We will rearrange the work and place the ablations table in the main paper.\n\n---\n\n***7. Visualization.*** \n- We will make all figures be placed at the top of the pages to enhance readability.\n\n---\n\n***8. Related work*** \n- 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.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "95OzM2M90v", "reviewer": "Reviewer_QU5R", "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.\n 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.", "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?\n\n2. 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.\n\n3. 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.", "limitations": "The evaluation and application of the proposed method is limited.", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "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.\n\nSignificance: 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.\n\nMethodology: 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.\n\nResults: 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.\n2. 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.\n3. 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.\n4. The computational process of the AMM would benefit from additional figures or diagrams to enhance understanding.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks for your insightful comments and suggestions. We will address your concerns in the following answers:\n\n---\n\n***1. Experiments on extra dataset CelebA-HQ.*** \n\n* 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.** \n\n|Model|Training images|FID|\n|---|---|---|\n|EDT-S|100k * 240|**16.60**|\n|DiT-S|100k * 240|19.12|\n|MDTv2-S|100k * 240|17.34|\n\n**Rebuttal-Table 2**: Evaluation of unconditional image synthesis on CelebA-HQ. \n\n---\n\n***2.Motivation for lightweight design.*** \n\n* **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.  \n\n---\n\n***3. The computational process of the AMM.***\n\n* **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.\n\n---\n\n***4. The computational cost of AMM.*** \n\n* **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.\n\n---\n\n***5. The loss of token information in long skip connection.*** \n\n* 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**.\n\n---\n\n***6. Training cost of EDT.*** \n\n* 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.)\n\n| Model|Epochs|Training images|GPU days| FID|\n|---|---|---|---|---|\n|**EDT-S** |80|102M|**2.75**|**38.73**|\n|DiT-S|80|102M|2.96|68.40|\n|MDTv2-S|80|102M|16.47|39.50|\n|**EDT-B**|80|102M|9.19|**19.18**|\n|DiT-B|80|102M|**8.62**|43.47|\n|MDTv2-B|80|102M|26.17|19.55|\n|**EDT-XL**|80|102M|**37.79**|**7.52**|\n|DiT-XL|80|102M|39.82|19.47|\n|MDTv2-XL|80|102M|72.62|7.70|\n\n**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.\n\n---\n\n***7. Adding AMM to Text-to-Image Models.*** \n* Due to time constraints, incorporating AMM into text-to-image models will be addressed in our future work.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "YjasL2Kf3i", "reviewer": "Reviewer_8pvq", "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.", "questions": "My main concerns are listed in the weaknesses of its evaluation of the CFG setting and other datasets.", "limitations": "Yes, they discussed it.", "rating": 5, "confidence": 5, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "+ The proposed network can reduce the computation of diffusion is a plus point although it requires more parameters compared to existing methods.\n\n+ Some experimental results on ImageNEt 256x256 show promising looks.", "weaknesses": "**[Update]: after rebuttal, I raised score from 4 to 5**\n\nThe proposed method shows some good merits in the architectural design, however, it remains several major concerns:\n\n+ 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).\n\n+ 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.\n\n[1] Relay Diffusion: Unifying diffusion process across resolutions for image synthesis, ICLR 2024", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the valuable feedback and suggestions. We will answer your questions in the following: \n\n---\n\n***1. The experiment under optimal CFG settings.*** \n\n* 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**.\n\n| Models|Training images|FID|\n|---|---|---|\n|EDT-S CFG=2|400k*256|**9.89**|\n|MDTv2-S CFG=3.8|400k*256|15.62|\n|DiT-S CFG=1.5|400k*256|21.03|\n|DiT-XL CFG=1.5|400k*256|5.50|\n|EDT-XL CFG=2|400k*256|**4.65**|\n\n**Rebuttal-Table 4**: Comparison and evaluation under optimal Cfg on ImageNet 256×256.\n\n---\n\n***2. Experiments on extra dataset CelebA-HQ.*** \n\n* 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.** \nDue to the limited time, we did not finish the experiments with the 512x512 resolution training. We will include it in the revision.\n\n|Model|Training images|FID|\n|---|---|---|\n|EDT-S|100k * 240|**16.60**|\n|DiT-S|100k * 240|19.12|\n|MDTv2-S|100k * 240|17.34|\n\n**Rebuttal-Table 2**: Evaluation of unconditional image synthesis on CelebA-HQ.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "tUslBePYq2", "reviewer": "Reviewer_MD5e", "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.", "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.\n2. 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.\n3. Please add potential limitations of the work.", "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.", "rating": 6, "confidence": 1, "soundness": 3, "presentation": 3, "contribution": 3, "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.\n2. 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.\n2. 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks for your positive feedback and insightful comments. We will answer your questions in the following:\n\n---\n\n***1. The usage of pre-trained VAE.*** \n* **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.\n\n---\n\n***2. Training cost of EDT.*** \n\n* 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.)\n\n| Model|Epochs|Training images|GPU days| FID|\n|---|---|---|---|---|\n|**EDT-S** |80|102M|**2.75**|**38.73**|\n|DiT-S|80|102M|2.96|68.40|\n|MDTv2-S|80|102M|16.47|39.50|\n|**EDT-B**|80|102M|9.19|**19.18**|\n|DiT-B|80|102M|**8.62**|43.47|\n|MDTv2-B|80|102M|26.17|19.55|\n|**EDT-XL**|80|102M|**37.79**|**7.52**|\n|DiT-XL|80|102M|39.82|19.47|\n|MDTv2-XL|80|102M|72.62|7.70|\n\n**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.\n\n---\n\n***3. Why use VAE.*** \n* 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.\n\n---\n\n***4. Experiments on extra dataset celebA-HQ.*** \n* 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.** \n\n|Model|Training images|FID|\n|---|---|---|\n|EDT-S|100k * 240|**16.60**|\n|DiT-S|100k * 240|19.12|\n|MDTv2-S|100k * 240|17.34|\n\n**Rebuttal-Table 2**: Evaluation of unconditional image synthesis on CelebA-HQ. \n\n---\n\n***5. Potential limitations.*** \n* 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. \n* 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. \n* Although significant progress has been made with AMM, the generation function of AMM still has room for improvement and warrants further exploration.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
We thank all the reviewers for their constructive comments and insightful suggestions. We carefully add experiments and figures according to the comments of all the reviewers. 

---

We are encouraged that the reviewers pointed out our work 
*"is novel"*, *"tackling an important application"* (**R FhNv**); 
*"well formulated and clearly explained"*, *"unique contributions"*,  and *"is noteworthy"* (**R QU5R**); 
*"reduce the computation of diffusion"* and *"experimental results show promising looks"* (**R 8pvq**); 
*"a novel way"*, *"pipeline is robust and lightweight"*, and *"provides specific quantitative results"* (**R MD5e**). 
We address the reviews below and will incorporate all changes in the revision.

---

***Summary and Responses to the Reviewers' Most Concerned Issues:***

---

***1. The performance improvements by 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. 

|**On ImageNet 256×256 (class-conditional image synthesis)**||||
|---|---|---|---|
|Model|Training images|W/o AMM|W AMM|
|EDT-S|400k * 256|42.60|**34.20**|
|DiT-S|400k * 256|67.16|**63.11**|
|MDTv2-S|400k * 256|39.02|**31.89**|
|EDT-XL|400k * 256|12.80|**7.52**|
|DiT-XL|400k * 256|18.48|**14.73**|
|DiT-XL|7000k * 256|9.62|**3.75**|
|**On CelebA-HQ 256 × 256 (unconditional image synthesis)**||||
|Model|Training images|W/o AMM |W AMM|
|EDT-S|100k * 240|17.01|**16.60**|
|DiT-S|100k * 240|19.12|**18.41**|
|MDTv2-S|100k * 240|17.34|**17.11**|

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

---

***2. Experiments on extra dataset CelebA-HQ 256×256.***

- CelebA-HQ 256×256 is used for the unconditional image synthesis task. We conducted a new experiment, training EDT-S, DiT-S, and MDTv2-S on CelebA-HQ 256×256, **as shown in Rebuttal-Table 2 in the global rebuttal. EDT-S achieved the lowest FID, demonstrating that EDT is also effective in unconditional image synthesis tasks.** 

|Model|Training images|FID|
|---|---|---|
|EDT-S|100k * 240|**16.60**|
|DiT-S|100k * 240|19.12|
|MDTv2-S|100k * 240|17.34|

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

---

***3. The training cost of EDT.***

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

| Model|Epochs|Training images|GPU days| FID|
|---|---|---|---|---|
|**EDT-S** |80|102M|**2.75**|**38.73**|
|DiT-S|80|102M|2.96|68.40|
|MDTv2-S|80|102M|16.47|39.50|
|**EDT-B**|80|102M|9.19|**19.18**|
|DiT-B|80|102M|**8.62**|43.47|
|MDTv2-B|80|102M|26.17|19.55|
|**EDT-XL**|80|102M|**37.79**|**7.52**|
|DiT-XL|80|102M|39.82|19.47|
|MDTv2-XL|80|102M|72.62|7.70|

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

---

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

---

***5.  The computational process of the AMM in a block.***

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

Meta-review: Overall, the reviews for this paper are generally positive, with three borderline accepts and one weak accept. The reviewers commend the work for addressing a challenging issue in diffusion, introducing novel and robust techniques, and providing substantial experimental evidence. The rebuttal successfully clarified several points of concern, further convincing the reviewers to some extent. The AC agrees with the reviewers’ assessment and recommends accepting this paper.

However, the reviewers have raised some concerns that should be addressed in the camera-ready version:

1. Provide more qualitative analysis within the main paper rather than in the appendix, including high-resolution comparisons with other methods.
2. Include the experiment of EDT_XL for full convergence and its discussion.
3. Incorporate the additional clarifications and evidence provided in the rebuttal into the final version.

We look forward to seeing this paper presented at NeurIPS.
Decision: Accept (poster)