Fast Training of Diffusion Models with Masked Transformers
Authors: Hongkai Zheng
Weili Nie
Arash Vahdat
Anima Anandkumar
Venue: Accepted by TMLR ()

Abstract: We propose an efficient approach to train large diffusion models with masked transformers.
While masked transformers have been extensively explored for representation learning, their application to generative learning is less explored in the vision domain. Our work is the first to exploit masked training to reduce the training cost of diffusion models significantly. Specifically, we randomly mask out a high proportion (e.g., 50\%) of patches in diffused input images during training. For masked training, we introduce an asymmetric encoder-decoder architecture consisting of a transformer encoder that operates only on unmasked patches and a lightweight transformer decoder on full patches. To promote a long-range understanding of full patches, we add an auxiliary task of reconstructing masked patches to the denoising score matching objective that learns the score of unmasked patches. Experiments on  ImageNet-256x256 and ImageNet-512x512 show that our approach achieves competitive and even better generative performance than the state-of-the-art Diffusion Transformer (DiT) model, using only around 30\% of its original training time. Thus, our method shows a promising way of efficiently training large transformer-based diffusion models without sacrificing the generative performance. Our code is available at https://github.com/Anima-Lab/MaskDiT.

Published in Transactions on Machine Learning Research Mlarch/2024/
Fast Training of Diffusion Models with Masked Transformers
Hongkai Zheng"
hazheng@oaltech. edu
Caltech
Weili Nie*
unieennridaa.
NVIDIA
Arash Vahdat
avakdat@nvidia..om com
NVIDIA
Anima Anandkumar
anima@caltech. edu
Caltech
Reviewed on OpenReview:
https: opw.reveee.e.eeeeeeem net/ forum? id- ddBBBBBtGio
Abstract
We propose an efficient approach train large diffusion models with masked transformers. While masked transformers have been extensively explored for representation learning, their application to generative learning is less explored in the vision domain. Our work is the first to exploit masked training to reduce the training cost of diffusion models significantly. Specifically, we randomly mask out high proportion (e.g., 50%) of patches in diffused input images during training. For masked training, we introduce an asymmetric encoder-decoder architecture consisting of a transformer encoder that operates only on unmasked patches and a lightweight transformer decoder on full patches. To promote a long-range understanding of full patches, we add an auxiliary task of reconstructing masked patches to the denoising score matching objective that learns the score of unmasked patches. Experiments on ImageNet- 256 256 and ImageNNe-5112x122 show that our approach achieves competitive and even better generative performance than the state-of-the-arl Diffusion Transformer (DiT) model, using only around 30% of its original training time. Thus, our method shows promising way efficiently training large transformer-based diffusion models without sacrificing the generative performance. Our code available at hnttps://githhn http//ggtthbbooo/////
cotplluub..om//iiiaaaaaaa Anima-Liimaaaaaaaaa MaskDiT
1 Introduction
Diffusion models (Schl-Dickstein et al. 2015 Ho et al., 2020 [onn al. 2021b) have become the most popular class of deep generative models, due to their superior image generation performance (Dhariwal
& Nichol 2021 particularly in syuthesizing high-quality and diverse images with text inputs (Ramesh
et al. 2022: Rombach et al., 2022 Saharia et al. 2022b Balaji et al. 2022). Their strong generation performance has enabled many applications, including image-to-imane
translation (Saharia et al.) 2022a Liu et al. 2023), personalized creation (Ruiz et al.) 2022 Gal et al.) 2022 and model robustness (Nie
et al.j 2022 Carlini et al. 2022). Large-scale training of these models is essential for their capabilities of gencrating realistic images and creative art. However, training these models requires a. large amount of computational resources and time, which remains major bottleneck for further scaling them up. For
example, the original stable diffusion (Rombach et al. 2022) was trained on 256 A100 GPUs for more than 24 days. While the training cost can be reduced to 13 days 256 A100 GPUs with improved infrastructure
"Equal contribution. Correspondence bezhengSealltehh.cch wnie@nvidia.com.

---

Published in Transactions on Machine Learning Research March/2022
DSM target
target targer on
200
LDM-8
Linear and Reshspe
Uneer Reshape
LDM-4
Layer Norm
DIT-XU/2
Ours
Decoder
MDT-XL/2
Leamatle sokeng
Diffusion models we guidance
Insert
17
Diameter
DIT Encoder
DIT Encoder
LDM-8
GRLOPs
Encode visitle tokens
13
mum
Remove masked takens
LDM-4
DIT-XL22
Time, label
Time, label
Ours
MDT-XL22
Di)
(b) MaskDIT
Diffusion guidance
Figure 1: A comparison of our MaskDiT architecture with DiT (Peebles & Xie 2022). During training, we randomly mask high proportion
Figure 2: Generative performance of
the state-of-the-art diffussion models on ImageNet-256x in two settings: with and without guidance. The area of each bubble indicates the FLOPs for single forward pass during training. Our method is more compute-efiiient with competitive performance.
input patches. The encoder operates on unmasked patches, and after adding learnable masked tokens (marked by gray), full patches processed by small decoder. The model is trained to predict the score unmasked patches (DSM target) and reconstruct the masked patches (MAE target). At inference, all the patches are processed for sampling.
and implementation (Cory Stephenson 2023), it is still inaccessible to most researchers and practitioners. Thus, improving the training efficiency of diffusion models is still an open question.
In the representation learning community,, masked training widely used to improve training efficiency domains such as natural language processing (Devlin et al. 2018), computer vision (He et al. (022). and vision-language understanding (Li et al.) 2022b). Masked training significantly reduces the overall training time and memory, especially in vision applications (He et al. 2022 Li et al. 2022b). In addition, masked training is an effective self-supervised learning technique that does not sacrifice the representation learning quality, as high redundancy in the visual appearance allows the model to learn from unmasked patches. Here, masked training relies heavily on the transformer architecture (Vaswani et al. 2017) since they operate on patches, and is natural to mask a subset them.
Masked training, however, cannot be directly applicd to diffusion models since current diffusion models mostly use U-Net (Ronneberger et al., 2015) the standard network backbone, with some modifications (Ho al. 2020 Song et al., 2021b Dhariwal & Nichol 2021). The convolutions in U-Net operate regular dense grids, making it challenging to incorporate masked tokens and train with a random subset. of the input patches. Thus, applying masked training techniques to U-Net models is not straightforward.
Recently, a few works (Peebles & Xie 2022 Bao et al. 2023 show that replacing the U-Net architecture with vision transformers (ViT) (Dosovitskiy et al.) (2220) as the backbone of diffusion models can achieve similar or even better performance across standard image generation benchmarks. In particular, Diffusion Transformer (DiT) (Peebles & Xie
2022) inherits the best practices of ViTs without relying the U-Net
2

---

Published in Transactions on Machine Learning Research (March/2024)
inductive bias, demonstrating the scalability of the DiT backbone for diffusion models. Similar to standard ViTs, training DiT's is computationally expensive, requiring millions of iterations. However, this opens up new opportunities for efficient training of large transformer-baseee diffusion models.
In this work, we leverage the transformer structure of DiTs to enable masked modeling for significantly faster and cheaper training. Our main hypothesis that images contain signíficant redundancy the pixel space. Thus, it is possible train diffusion models by minimizing the denoising score matching (DSM) loss only subset of pixels. Masked training has two main advantages: 1) It enables us to pass only small subset of image patches to the network, which largely reduces the computational cost per iteration. 2) It allows us to augment the training data with multiple views, which can improve the training performance, particularly in limited-data settings.
Driven by this hypothesis, we propose an efficient approach to train large cransormer-baaee diffusion models, termed Masked Diffusion Transformer (MaskDiT). We first adopt an asymmetric encoder-decoder architec- ture, where the transformer encoder only operates on the unmasked patches, and light transformer decoder operates on the complete set of patches. We then design new training objective, consisting of two parts: 1) predicting the score of unmasked patches via the DSM loss, and 2) reconstructing the masked patches of the input via the mean square error (MSE) loss. This provides a global knowledge of the complete image that prevents the diffusion model from overfitting to the unmasked patches. For simplicity, we apply the same masking ratio (e.g., 50%) across all the diffusion timesteps during training. After training, we finetune the model without masking to close the distribution gap between training and inference and this incurs only a minimal amount of overhead (less than 6% of the total training cost).
By randomly removing 50% patches of training images, our method reduces the per-iteration training cost by 2x. With our proposed new architecture and new training objective, MaskDiT can also achieve similar performance to the DiT counterpart with fewer training steps, leading to much smaller overall training costs. In the class-conditional ImageNet-256x2 256 generation benchmark, MaskDiT achieves an FID of 5.69 without guidance, which outperforms previous (non-cascaded) diffusion models, and achieves an FID of 2.28 with guidance, which is comparable to the state-of-the-art models. To achieve the result FID 2. 2.28, the total training cost of MaskDiT is 273 hours on 8x A100 GPUs, which is only 31% of the DiT's training time. For dlasscoonditional MmageNet-5122512 MaskDiT achieves FID of 10.79 without guidance and 2.50 with guidance, outperforming DiT-XL/2 and prior diffusion models. The total training cost is 209 A100 GPU days, which is about 29% of that of DIT- validating the scalability and effectiveness of MaskDiT. All these results show that our method provides a significantly better training efficiency in wall-clock time.
Our main contributions are summarized as follows:
We propose a new approach for fast training of diffusion models with masked transformers by randomly masking out a high proportion (50%) of input patches.
For masked training, we introduce an asymmetric encoder-decoder architecture and a new training objective that predicts the score of unmasked patches and reconstructs masked patches.
We show our method has laster training speed and less memory consumption while achieving com- parable generation performance, implying its scalability to large diffusion models.
2
Related work
Backbone architectures of diffusion models The success of diffusion models in image generation does not come true without more advanced backbone architectures. Since DDPM (Ho et al. 2020). the convolutional U-Net architecture (Ronneberger et al., 2015) has become the standard choice of diffusion backbone (Rombach et al.] 2022 Saharia et al.) 2022b). Built on the U-Net proposed in (Ho et al. 2020
several improvements have also been proposed (e.g., (Song et al. (2021b] Dhariwal & Nichol] 2021) ) to achieve better generation performance, including increasing depth versus width, applying residual blocks from BigGAN (Brock et al., 2019),, using more attention layers and attention heads, rescaling residual
connections with I_ etc. More recently, several works have proposed transformer based architectures for /
3

---

Published in Transactions on Machine Learning Research Mlarch/2024/
Figure 3: Generated samples from MaskDiT. Upper panel: without CFG. Lower panel: with CFG (5cale=1.5). Both panels use the same batch of input latent noises and class labels.
U-ViT (Bao et al. 2023 RIN (Jabri et al. 2022 and DiT (Peebles & Xie 2022), largely due to the wide expressivity and flexibility of transformers (Vaswani let al. 2017 Dosovitskiy et al. 2020). Among them, DiT a pure transformer architecture without the U-Net inductive bias and it obtains a better Ilaas-conditionall generation result than the U-Net counterparts with better scalability (Peebles & Xie (022)). Rather than developing a better transformer architecture for diffusion models, our main goal is to improve the training efficiency of diffusion models by applying the masked training to the diffusion transformers (c.g., DiT).
diffusion models, such as GenViT (ang et al. 2022
Efficiency in diffusion models Despite the superior performance of diffusion models, they suffer from high training and inference costs, especially large-scale, high-resolution image datasets. Most works have focused on improving the sampling efficiency of diffusion models, where some works (Song et al.) 2021aa] Lu et al. 2022 Karras et al. 2022) improve the sampling strategy with advanced numorical solvers while others (Salimans & Ho! 2022: Meng el al.j 2022 Zheng et al. 2022 Song el al., 2023 train surrogate network by distillation. However, they are not directly applicable reducing the training costs of diffusion models. For better training efficiency, Vahdat al. (2021); Rombach al. (0022) have proposed to train latent-space diffusion model by first mapping the high-resolution images to the oow-dimensional latent space. (Ho et al. 2 developed a cascaded diffusion model, consisting of a base diffusion model on low-resolution images and several lightweight super-resolution diffusion models. More similarly, prior work (Wang et al. 2023) proposed a patch-wise training framework where the score matching is performed at the patch level. Different from our work that uses a subset small patches as the transformer input, (Wang et al. 2023 uses single large patch to replace the original image and only evaluates their method on the U-Net architecture.
Masked training with transformers As transformers become the dominant architectures in natural language processing (Vaswani et al. 2017) and computer vision (Dosovitskiy et al. 2020), masked training
4

---

Published in Transactions on Machine Learning Research March/2024
has also been broadly applied to representation learning (Devlin et al. 2018 He et al.
2022 Li et al. 2022b
and generative modeling (Radford et. al. 2018 Chang et al. 2022: 2023) in these two domains. One of its notable applications is masked language modeling proposed by BERT (Devlin et al. 2018 In computer
vision, a line of works has adopted the idea of masked language modcling to predict masked discrete tokens by first tokenizing the images into discrete values via VQ-VAE (Van Den Oord et al.) 2017), including BEiT [Bao et al., 2022) for self-supervieed learning, MaskGiT (Chang et al., 2022) and MUSE (Chang al. 2023 for image synthesis, and MAGE (Li et al., 2022a) for unifying representation learning and generative modeling. On the other hand, MAE (He et al. 2022) applied masked image modeling to directly predict the continuous tokens (without discrete tokenizer). This design also takes advantage of masking to reduce training time and memory and has been applied to various other tasks (Feichtenhofer et al.) 2022 Geng et al. 2022). Our work is inspired by MAE, and it shows that masking can siguificantly improve the training efficiency of diffusion models when applying it to an asymmetric decoder architecture.
Masked training of diffusion models More recently, it starts to attract more atlention combine masked training and diffusion models. Wei et al. (2023) formulated diffusion models as masked autoencoders (DiffMAE) by denoising masked patches corrupted with Gaussian noises at different levels. But DiffMAE mainly focuses on the discriminative tasks, besides being able to solve the inpainting task. Gao et al. (2023) recently proposed MDT that adds the masked training objective to the original DiT loss for better generative performance. Our method has several differences with MDT: First, MDT still keeps the original DSM loss for training which takes full patches. In contrast, MaskDiT applics DSM only to unmasked tokens and proposes a reconstruction loss on masked patches, which leads to more efficient training. Second, MDT relies on extra modules like side-interpolater and relative positional bias while MaskDiT only needs minimal changes to the original DiT' architecture. Last but not least, our main goal is fast training of diffusion models with little loss of performance while MDT aims at better generative performance with additional network modules and more computational cost per training iteration.
3 Method
In this section, we first introduce the preliminaries of diffusion models that we have adopted during training and inference, and then focus on the key designs of our proposed MaskDiT.
3.1 Preliminaries
Diffusion models We follow (ong ct al.] (2021b). which introduces the forward diffusion and backward demoising processes of diffusion models in continuous time setting using differential equations. In the forward process, diffusion models diffuse the real data zn ~ (do towards a noise distribution ær (10 !) through the following stochastic differential equation (SDE):
dx f(z t)dt + g(f)dw,
where f is the (vector-valued) drift coefficient, g is the diffusion coefficient, w is a standard Wiener process, and the time & flows from to T. In the reverse process, sample generation can be done by the following SDE:
dx (f(x, - V( log (x))ddt + g({)dw
(2)
where w standard reverso-time Wiener process and dt is an infinitesimal negative timestep. The reverse SDE can be converted to probability flow ordinary differential equation (ODF) (Song al. 2021b):
dr = (f(xtt) 2 _ V le (z)]dt
(3)
which has the same marginals pe (z) as the forward SDE at all timesteps t. We closely follow the EDM formulation (Karras et al., 2022) by setting f(x t) = 0 and g(t) := V2t. Specifically, the forward SDE
reduces to x = Zo + where n ~ ((0,t1), and the probability flow ODE becomes
dx - log (x)))
(4)
5

---

Published in Transactions on Machine Learning Research (March/200 /2024)
To learn the score function sx,, V lo EDM parameterizes denoising function Do(æ, t) to minimize the denoising score matching loss:
E Peor-pann Vo..
(D Do + -n,) xoll
(5)
and then the estimate score 8(x, 1) (Do(x,1) -
Classifter-free guidance For class-conditional generation, classifier-Iree guidance (CFG) (Ho & Salimans, 2022) is widely used sampling method to improve the generation quality of diffusion models. In the EDM formulation, denote by (2 c) the class-conditional denoising function, CFG defines a modified denoising function: Da t, = De (z, t) + w(Do(z, t,c) Dz t)), where w 2 1 is the guidance scale. To get the unconditional model for the CFG sampling, we can simply use mull token Ø (e.g., an all-zero vector) replace the class label e, i.e., Do(z,t) Ds(x,t,0). During training, we randomly Ø with some fixed probability Puncond-
the null token
3.2 Key designs
Image masking Given a clean image zn and the diffusion timestep t, we first get the diffused image z; by adding the Gaussian noise n. We then divide (or "patchify") I; into a grid of N non-overlapping patches, each with a patch size of p x p. For an image of resolution H W, we have N = (HW)/p². With a fixed masking ratio r, we randomly remove [rN] patches and only pass the remaining N [rN] unmasked patches the diffusion model. For simplicity, we keep the same masking ratio for all diffusion timesteps. A high masking ratio largely improves the computation efficiency but may also reduce the learning signal of the estimation. Given the large redundancy of ze, the learning signal with masking may be compensated by the model's ability to extrapolate masked patches from neighboring patches. Thus, there may exist an optimal spot where we achieve both good performance and high training efficiency.
Asymmetric encoder-decoder backbone Our diffussio backbone is based DřT, standard ViT- based architecture for diffusion models, with few modifications. Similar to MAE (He et al.) 2022), we apply an asymmetric encoder-decoder architecture: 1) the encoder has the same architecture as the original DiT except without the final lincar projection layer, and only operates on the unmasked patches; 2) the decoder is another DiT architecture adapted from the lightweight. MAE decoder, and it takes the full tokens as the input. Similar to DiT, our encoder embeds patches by linear projection with standard ViT froquency-based positional embeddings added to all input tokens. Then masked tokens removed before being passed to the remaining encoder layers. The decoder takes both the encoded unmasked tokens in addition to new mask tokens as input, where each mask token is a shared, learnable vector. We then add the same positional embeddings to all tokens before passing them to the decoder. Due to the asymmetric design (e.g., the MAE decoder with <9% paramcters of DiT-XL/2), masking can dramatically reduce the computational cost per iteration.
Training objective Unlike the general training of diffusion models, we do not perform denoising score matching on full tokens. This is because it is challenging to predict the score of the masked patches by solely relying on the visible unmasked patches. Instead, we docompose the training objective into two subtasks: 1) the score estimation task on unmasked tokens and 2) the auxiliary reconstruction task on masked tokens, Formally, denote the binary masking label by E (0, 1J.. We uniformly sample [r:]] patches without replacement and mask them out. The denoising score matching loss unmasked tokens reads:
LDSM =
Exoerpaat E E Em (Do((zo + n) o 11-m).1) mo) @ (1-m)l²
m)
(6)
where @ denotes the element-wist multiplication along the token length dimension of the "patchified" image, and (0o + n) o (1-m) denotes the model D, only takes the unmasked tokens as input. Similar to MAE, the reconstruction task performed by the MSE loss on the masked tokens:
LMAE Eao 110 1m (Do((zo + +n) @ (1 m), (2o + n)) 3 ml|
(7)

---

Published in Transactions on Machine Learning Research MMarch/2024)
where the goal is to reconstruct the imput (i.e., the diffused image Zo + n) by predicting the pixel values for each masked patch. We hypothesize that adding this MAF reconstruction loss can promote the masked transformer to have a global understanding of the full (diffused) image and thus prevent the denoising score matching loss in (6) from overfitting to a local subset of visible patches.
Finally, our training objective is given by
= ALMAE
where the hyperparameter À controls the balance between the score prediction and MAE reconstruction losses, and it cannot be too large to drive the training away from the standard DSM update.
Unmasking tuning Although the model trained on masked images, it can be used to generate full images with high quality in the standard class-conditional sctting. However, we observe that the masked training always leads to a worse performance of the CFG sampling than the original DiT training (without masking). This is because dassifier-frre guidance relies on both conditional and unconditional score. Without label information, learning unconditional score from only 50% tokens is much harder than that of conditional score. Therefore, the unconditional model (by setting the class label e to the null token Ø) is not as good as the conditional model. To further close the gap wetween the training and inference, after the masked training, we finetune the model using a smaller batch size and learning rate at minimal extra cost. We also explore two different masking ratio schedules during unmasking tuning: 1) the zero-ratio schedule, where the masking ratio r = 0 throughout all the luming steps, and 2) the cosine-ralio schedule, where the masking ratio r = co5 (m/2 with and Root being the current steps and the total number of tuning steps, respectively. This unmasking tuning strategy trades slightly higher training cost for the better performance of CFG sampling.
4 Experiments
4.1 Experimental setup
Model settings Following DiT (Peebles & Xie 2022 we consider the latent diffusion model (LDM) framework, whose training consists of two stages: 1) training an autoencoder to map images to the latent space, and 2) training a diffusion model on the latent variables. For the autoencoder, we use the pre-trained VAE model from Stable Diffusion (Rombach
al. 2022). with a downsampling factor of 8. Thus, an image of resolution 256 x 256 3 results in 32 32 x 4 latent variable while an image of resolution 512 x 512 3 results in a 64 x 64 x 4 latent variable. For the diffusion model, we use DiT-XL/2 2 as encoder architecture, where "XL" denotes the largest model size used by DiT and "2" denotes the patch size of 2 for all input patches. Our decoder has the same architecture with MAE (He et al. 2022). except that we add the adaptive layer norm blocks for conditioning on the time and class embeddings Peebles & Xie 2022
Training details Most training details are kept the same with the DIT work: AdamW (oshchhlov & Hutter 2017) with a constant learning rate of 1c-4, no weight decay, and an exponential moving average (EMA) of model weights over training with a decay of 0.9999. Also, we use the same initialization stralegies with DIT. By default, we use a masking ratio of 50%, an MAE coefficient À A= 0.1, a probability of dropping class labels Puncond 0.1, and batch size of 1024. Because masking greatly reduces memory usage and FLOPs, we find it more efficient to larger batch size during training. Unlike DiT which uses horizontal flips, we do not use any data augmentation, as our early experiments showed it does not. bring better performance. For the unmasked tuning, we change the learning rate to 5e-5 and use full precision for better training stability. Unless otherwise noted, experiments on ImageNet 256 x 256 are conducted on 8x A100 GPUs, each with 80GB memory, whereas for ImageNet 512 x 512, we use 32x A100 GPUs.
Evaluation metrics Following DiT (Peebles & Xie 2022 we mainly use Fréchet Inception Distance (FID) (Heusel et al.) 2017) to guide the design choices by measuring the generation diversity and quality of our method, as it is the most commonly used metric for generative models. To ensure fair comparison with previous methods, we also use the ADM's TensorFlow evaluation suite (Dhariwal & Nichol 2021) to compute FID-50K with the same reference statistics. Similarly, we also use Inception Score (TS) (Salimans et al. 2016).
7

---

Published in Transactions on Machine Learning Research Mlarch/2024/
Training memory per GPU (GB)
DiT
250
DiT
Training speed (steps/sec)
5
MDT
MDT
Ours
tf32
200
Ours_ furs_tf32
Ours
Ours
150
3
100
2
Batch=256
Batch=1024
Batch=256
Batch=1024
256x 256 diffusion models
256x256 diffusion models
350
Training memory per GPU (GB
DIT
Training speed (steps/sec
Ours_ __33
300
Ours tf32
1.50
Ours
Ours
250
1.25
200
1.00
150
100
0.50
0.25
50
0.00
Batch 22tc
Batch=1024
Batch= 256
Batch=1024
512x512 diffusion models
512x512 diffusion models
Figure 4: Training speed (in steps/scc) and memory-per-GPU (in GB) for DiT, MDT and MaskDiT (Ours), measured on a single node with 8xA100 GPUs, each with 80 GB memory. (Top): Results on ImageNet- x56x256. (Bottom): Results ImageNet 512x512. Here and MDT where DiT applies the TensorFloat322 (TF32) precision and MDT applies the Float32 precision. Our MaskDiT, by default, applies Mixed Precision. We also add the TF32 precision of MaskDiT (termed "Ours ft:2")) for reference. If the GPU memory requirements surpass the GPU hardware of our hardware, we enable gradient accumulation and estimate the GPU memory usages accounting for typical overhcads in distributed training communication.. MDT paper has no results on 512x512 so we exclude from the comparison on 512x512.
the latest official implementation of DIT
sFID Nash 2021 and Trecision/Recall (Kynkäänniemi et al. 2019) as secondary metrics to compare with previous methods. Note that for FID and sFID, the lower is better while for IS and Precision Recalll the higher is better. Without stating explicitly, the FID values we report in experiments refer to the cases without classifier-free guidance.
4.2 Training efficiency
We compare the training elficiency of MaskDiT, DIT-XL/2 (Peebles & Xie 2022), and MDT-XL/2 (Gao et al. 2023) on 8x A100 GPUs, from three perspectives: GLOPs, training speed and memory consumption, and wall-time training convergence. All three models have a similar model size.
First, in Figure 2 the GFLOPs of MaskDiT, defined as FLOPs for a single forward pass during training, are much fewer than DiT and MDT. Specifically, our GFLOPs are only 54.0% DiT and 31.7% of MDT, As reference, LDM-8 has similar number of GFLOPs with MaskDiT, but FID much worse than MaskDiT. Second, in Figure 4 the training speed of MaskDiT much larger than the two previous methods while memory-per-GPP cost is much lower, matter whether Mixed Precision is applied. Also, the

---

Published in Transactions on Machine Learning Research Mlarch/2024/
DIT-XL/2
70
70
MDT-XL/2
MDT-XL/2
60
MaskDiT
60
MaskDIT
30
90
20
20
60 Training hours
40
80
Training hours
Figure 5: FID vs. training hours on ImageNet x56x256: Left: batch size 256. Right: batch size 1024. For DiT and MDT, we use their official implementations with default settings. For batch size 256, we only report the FID MDT before 220k steps because its gradient explodes after that.
Comparison of MaskDiT with various statc-of-the-art generative models on class-conditional Ima- Table 1:
geNet 256x256, where -G means the classifier-free guidance.
Method
FID4
sFID
IST
Prec.t
Rec.t
BigGAN-ddeep Brock et al. 2018)
6.95
7.36
171.40
0.87
0.28
StyleGAN-XL Sauer et al. 2022
2.30
4.02
265.12
0.78
0.53
MaskGIT Chang et al. 2022)
6.18
182.10
0.80
0.51
CDM Ho al. 2022
4.88
158.71
ADM Dhariwal & Nichol
2021
10.94
6.02
100.98
0.69
0.63
ADM-U
7.49
5.13
127.49
0.72
0.63
LDM-8
Rombach et al.
2022
15.51
79.03
0.65
0.63
LDM-4
10.56
209.52
0.84
0.35
U-VIT-H/2
Bao et al. 2023
6.58
DiT-XL/2
i
Peebles & 2022
9.62
6.85
121.50
0.67
0.67
MDT-XL/2
Gao et al.
2023
6.23
5.23
143.02
0.71
0.65
MaskDiT
5.69
10.34
177.99
0.74
0.60
ADM-G DDhariwal & Nichol
2021
4.59
5.25
186.70
0.82
0.52
ADM-G, ADM-U
3.94
6.14
215.84
0.83
0.53
LDM-8-G
(Rombach et al.
2022
7.76
103.49
0.71
0.62
LDM-4-G
3.60
247.67
0.87
0.48
U-VIT-H/2-G
Bao et 2023
2.29
5.68
263.88
0.82
0.57
DIT-XT./2-G
Peebles & Xie
2022
2.27
4.60
278.24
0.83
0.57
MDT-XL/2-G
Gao 2023
1.79
4.57
283.01
0.81
0.61
MaskDiT-G
2.28
5.67
276.56
0.80
0.61
improvement becomes higher use larger batch size. For resolution 256x256 with batch size of 1024, our training speed is 3.5x of DiT and 6.5x of MDT, while memory cost is only 45.0% of DiT's and 19.2% of MDT's. For resolution 512x512 with batch size of 1024, our training speed is 4.6x of DiT's, and our GPU memory cost remains as low as 45.7% of DiT's. Finally, in Figure 5) MaskDiT also consistently yields better wall-time training convergence than DiT and MDT with different batch sizes. Similarly, increasing the batch size further improves the efficiency of MaskDiT, but does not benefit DiT-XI,/2 and MDT-XL/2. For instance on ImageNet 256x256, when using batch size of 1024, MaskDiT can achieve an FID of within 40 hours, while the other two methods take more than 160 hours, implying 4x speedup.
9

---

Published in Transactions on Machine Learning Research Alarch/2024/
Table 2: Comparison of MaskDiT with various state-of -f-t--rrttt-- generative models on coass-onniitooall Ima- geNet 512x512, where -G mcans the classificr-free guidance.
Method
FID4
sFID
IST
Prec.
Rec.
BigGAN-deep Brock et al.
2018
8.43
8.13
177.90
0.88
0.29
StyleGAN-XT, (auee et al.
2022
2.41
4.06
267.75
0.77
0.52
ADM Dhariwal Nichol
2021
23.24
10.19
58.06
0.73
0.60
ADM-U
9.96
5.62
121.78
0.75
0.64
DiT-XT,/2 (Pebbles & Xie
2022
12.03
7.12
105.25
0.75
0.64
MaskDiT
10.79
13.41
145.08
0.74
0.56
ADM-G (Dhrriall Nichol
2021
7.72
6.57
172.71
0.87
0.42
ADM-G, ADM-U
3.85
5.86
221.72
0.84
0.53
DIT-XL/2-G (Peebles Xie
2022
3.04
5.02
240.82
0.84
0.54
MaskDiT-G
2.50
5.10
256.27
0.83
0.56
4.3 Comparison with state-of-the-an
We compare against state-of-the-ae class-conditional generative models in Table I and where our 256x256 results are obtained after 2M training steps and 512x512 results are obtained after 1M training steps.
ImageNeet2250 256. Without CFG, we further perform the ummasking tuming with the cosine-ratio schedule for 37.5k steps. MaskDiT achieves a better FID than all the non-cascaded diffusion models, and it even outperforms the MDT-XL/2, which is much more computationally expensive, by decreasing FID from 6.23 5.69. Although the cascaded diffusion model (CDM) has better FID than MaskDiT, it underperforms MaskDiT regarding IS (158.71 vs. 177.99).
When using CFG, we perform the unmasking tuning with the zero-ralio schedule for 75k steps. MaskDiT-G achieves very similar FID (2.28) with DiT-XI./2-G (2.27), and they both share similar values of TS and Precision/ Recal.. Notably, MaskDiT achieves an FID of 2.28 with only 273 hours (257 hours for training and hours for unmasking tuning) on 8x A100 GPUs while DiT-XL/2 (7M steps) takes 868 hours, implying 31% of the DiT-XL//2's training time. Compared with the MDT-XT1/2-G, MaskDiT-G performs worse terms of FID and 1S but has comparable Precision/Rese Recall scores.
ImageNett-512x512. As shown in Table 2 without CFG, MaskDiT achieves an FID of 10.79, outperform- ing DFT-XL/2 that has an FID of 12.03. With CFG, MaskDiT-G G achieves an FID of 2.50, outperforming all previous methods, inchuding ADM (FID: 3.85) and DiT (FID: 3.04). The total training cost including unmasked tuning is around 209 A100 GPU days, which is about 29% of the training cost of DiT (712 A100
GPU days). We provide 512x512 MaskDiT samples in Figure 11 and Figure 12
We also observe that increasing the training budgets in the unmasking tuning stage can consistently improve the FID, but it also harms the training efficiency of our method. Given that our main goal is not to achieve new stato-of-tho-art and our generation quality is already good as demonstrated by Figure and Appendix we do not explore further in this direction.
4.4 Ablation studies
We conduct comprehensive ablation studies to investigate the impact of various design choices on our method. Unless otherwise specified, we use the default settings for training and inference, and report FID without guidance the evaluation metric.
Masking ratio and diffusion backbone As shown in Table 3a our asymmotric encoder-decoder archi- tecture always better than the original DiT architecture. When FID 100, we consider these two backbones to behave equally poorly. The performance gain from the asymmetric backbone is the most significant with
10

---

Published in Transactions on Machine Learning Research MMarch/20244
Table 3: Ablating our method. Unless stated otherwise, we use the following setting: our asymmetric encoder-decoder architecture, 50% masking ratio, the MAE coefficient À 0, and applying DSM on un- masked tokens only. By default, we report FID after training for 400k steps with a batch size of 1024.
(a) Impact of masking ratio and diffusion backbone
mask 0%
mask 50%
mask 75%
DiT backbone
14.70
24.58
107.55
Our backbone 13.71
12.31
121.16
(b) Impact of MAE reconstruction
== 0
0.01
0.1
1.0
400k steps
12.31
13.76
8.87
9.76
800k steps
12.42
12.93
7.19
12.74
(c) Impact of DSM design
X = 0
0.1
DSM on unmaskod tokens
12.31
8.87
DSM on full tokens
10.59
9.72
proper masking ratio (i.e., 50%). Without masking, the performance gain (FID: 14.70 - 13.71) mainly comes from the larger model size brought by the newly added decoder layers, confirming the observed ex- pressivity of DiT-like architectures (Peebles & Xie, 2022p. With a masking ratio of 50%, we see a large performance decrease (FID: 14.70 - 24.58) in the original DiT architecture but a performance gain (FID: 13.71 + 12.31) in the asymmetric encoder-decoder architecture. It implies that the interplay of the image masking and the asymmetric diffusion backbone contributes to the success of our method. Finally, the failure of masking 75% suggests that cannot arbitrarily increase the masking ralio for higher efficiency.
MAE reconstruction As shown in Table 3b adding the MAE reconstruction auxiliary task can largely benefit the generation quality of our method (e.g., FID from 12.42 to 7.19 at 800k steps). Moreover, the MAE coefficient x = 0.1 works the best, where we also observe the consistent improvement training iterations. A carefully tuned the MAE coefficient is crucial for our method: On one hand, if the MAE coefficient is too large (i.e., X = 1.0), FID first decreases (FTD=9.76 at 400k steps) but then increases (F11D=12.74 800k steps) during training. We hypothesize that the training in the late stage has been gradually dominated by the MAE reconstruction task, and the MAE training alone does not produce photorealistic images (He et al. 2022). On the other hand, if the MAE coefficient is too small (i.e., x= 0), the FID improvement soems to be saturated at the carly training stage, evidenced by FID=12.31 at 400k steps versus F1D=12.42 at 800k steps. It confirms our intuition that without the MAE reconstruction task, the training easily overfits the local subset of unmasked tokens as it lacks a global understanding of the full image, making the gradient update less informative.
FID over
DSM design Since our default choice is add the DSM loss on the unmasked tokens only, we compare its performance with adding the DSM loss on the full tokens. As shown in Table 3d without the MAE reconstruction loss (A = 0), adding the DSM loss the unmasked tokens only performs worse than that the full tokens. It can be explained by the fact that the score prediction on full tokens closes the gap between training and inference with our masking strategy. On the contrary, with the MAE reconstruction loss (A = 0.1), adding the DSM loss on the unmasked tokens only performs better than that on the full lokens. It is because, in this setting, we jointly perform the score prediction and MAE reconstruction tasks on the masked, invisible tokens, which are contradictory to each other. Furthermore, we also observe that across these four settings in Table 3c| the best FID score is achieved by "DSM on unmasked tokens & X = 0.1". which highlights the important role of the interplay between our DSM design and MAE reconstruction.
11

---

Published in Transactions on Machine Learning Research (March /2024)
Ok
25k
75k
Guidance scale w
Figure 6: Impact of unmasking tuning FID with guidance where we vary both the guidance scale and the number of tuning steps.
Unmasking tuning To show the impact of unmasking tuning, we report FID with guidance of different scales in Figure [ where we compare three tuning steps: Ok (no tuning), 25k, and 75k. We observe that without unmasking tuning, the best FID=4.54 obtained by the guidance scale w = 1.2, and the score keeps increasing as we further increase w. As we perform the unmasking tuning, the best FID with guidance improves over iterations (e.g., from FID-2.42 at 25k sleps to FID=2.28 at 75k steps). Interestingly, the optimal guidance scale w" also shifts to larger value over iterations (e.g., from w" 1.4 at 25k steps w* = 1.5 75k steps). This is mainly because that masked training does not produce good unconditional model, as evidenced by the observation that unconditional FID is 47.10 when training with mask ratio of 50% for 500k steps and it becomes 32.19 without masking. However, unmasking tuning improves the unconditional model. Thus, we can use a large w for a better FID as the number of tuming steps increases.
5 Broader impacts and limitations
Our work can reduce the training costs of diffusion models from the algorithmic perspective without sacri- ficing the generative performance. Our algorithmic innovation is orthogonal to the improved infrastructure and implementation (Cory Stephenson 2023 A combination of these two improvements can further make the development of large diffusion models more accessible to the broad community. Although our work is directly tied to particular applications, MaskDiT shares with other image synthesis tools similar potential benefits and risks, such as generating Deepfakes for disinformation, which have been discussed extensively (e.g., (Vaccari & Chadwick (020))
One of our limitations is that still need few steps of unmasking tuning to match the state-of-the-art FID with guidance. It is directly related to another issue of masked training: it does not produce good unconditional diffusion model. We leave it as future work to explore how to improve the unconditional image generation performance of MaskDiT.
6 Conclusions
In this work, we propose MaskDiT, an efficient approach training diffusion models with masked transform- ers. We randomly mask out large portion of the image patches, which significantly reduces the training overhead per iteration. To better accommodate masked training for diffusion models, first introduce asymmetric encoder-decoder diffusion backbone, where the DiT encoder only takes visible tokens and the lightweight DiT decoder lakes full lokens after injecting masked tokens. We also incorporate an auxiliary reconstruction loss into the DSM objective, where we reconstruct the input values of masked tokens and predict the score of unmasked tokens. On the class-conditional ImageNet-256x2B5 and ImagaNet-512x511 generation benchmarks, we demonstrated a better training efficiency for our proposed MaskDiT with a competitive performance with the state-of-the-art generative models.
12

---

Published in Transactions on Machine Learning Research Mlarch/2024/
Acknowledgements
We thank Sangwoo Mo for sharing insights into the DiT training, and Zhiding Yu, Shiyi Lan and Boyi Li for helpful discussions. Anima Anandkumar is supported by Bren Chair professorship and Schmidt Sciences AI 2050 Senior fellow funding.
References
Yogesh Balaji, Seungjun Nah, Xun Huang, Arash Vahdat, Jiaming Song, Karsten Kreis, Mika Aittala, Timo Aila, Samuli Laine, Bryan Catanzaro, et al. ediffi: Text-to-image diffusion models with an ensemble of expert denoisers. arXiv preprint orXiv:2211.0199,, 2022.
Fan Bao, Shen Nie, Kaiwen Xuc, Yue Cao, Chongxuan Li, Hang Su, and Jun Zhu. All are worth words: A vit backbone for diffusson models. CVPR, 2023.
Hangbo Bao, Li Dong, Songhao Piao, and Furu Wei. Beit: Bert pre-training of image transformers. In International Conference on Learning Representations, 2022.
Andrew Brock, Jeff Donahue, and Karen Simonyan. Large scale gan training for high fidelity natural image synthesis. arXiv preprint arXiv:1809. 1096, 2018.
Andrew Brock, Jeff Donahue, and Karen Simonyan. Large scale gan training for high fidelity natural image synthesis. In International Conference on Learning Represenlations, 2019.
Nicholas Carlini, Florian Tramer, Zico Kolter, et al. (certified!!) adversarial robustness for free! arv preprinl ar 22v: 10550, 2022.
Huiwen Chang, Han Zhang, Lu Jiang, Ce Liu, and William Freeman. Maskgit: Masked generative image transformer. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 11315-11325, 2022.
Huiwen Chang, Han Zhang, Jarred Barber, AJ Maschinot, Jose Lezama, Lu Jiang, Ming- Hsuan Yang, Kevin Murphy, William T Freeman, Michael Rubinstein, al. Muse: Text-to-image generation via masked generative transformers. arXio preprint arXi::2301.0070/, 2023.
Landan Seguin Cory Stephenson. Training Stable Diffusion from Scratch Costs <160k. https://github. http//oosaacnl/iiffussssoo-o---o
2023.
Jacob Devlin, Ming- Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidiroctional transformers for language understanding. arXiv preprint arXiv:1810.04800 2018.
Prafulla Dhariwal and Alexander Nichol. Diffusion models beat gans on image synthesis. Advances in Neural Information Processing Systems, 34:8780-888 8794, 2021.
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Un- terthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al. An image is worth 16x16 words: Transformers for image recognition at scale. arXiv preprint arXiv:2010. 11929, 2020.
Christoph Feichtenhofer, Yanghao Li, Kaiming He, et al. Masked autoencoders spatiotemporal learners. Advances in neural information processing systems, 35:35946-35958, 2022.
Rinon Gal, Yuval Alaluf, Yuval Atzmon, Or Patashnik, Amit H. Bermano, Gal Chechik, and Daniel Cohen- Or. An image is worth one word: Personalizing text-to-image generation using textual inversion, 2022.
URL
Shanghua Gao, Pan Zhou, Ming-Ming Cheng, and Shuicheng Yan. Masked diffusion transformer is strong image synthesizer. arXiv preprint (TXiw:2303.14380 2023.
Xinyang Geng, Hao Liu, Lisa Lee, Dale Schuurans, Sergey Levine, and Pieter Abbeel. Multimodal masked autoencoders learn transferable representations. arXio preprint arXim:2205.1 2022.
13

---

Published in Transactions on Machine Learning Research MMarch/2024/
Kaiming He, Xinlei Chen, Saining Xie, Yanghao Li, Piotr Dollár, and Ross Girshick. Masked autoencoders are scalable vision learners. In Proceedings of the IERE/CVF Conference on Computer Vision and Pattern Recognition, pp. 16000-16000 2022.
Martin Heusel, Hubert Ramsauer, Thomas Unterthiner, Bernhard Nessler, and Sepp Hochreiter. Gans trained by two time-scale update rule converge to local nash equilibrium. Aduances in neural infor- malion processing systems, 30, 2017.
Jonathan Ho and Tim Salimans. Classífier-free diffusion guidance. arXiv preprint ar Xrx: 12598, 2022.
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in Newral Information Processing Systems, 33:6840 6851, 2020.
Jonathan Ho, Chitwan Saharia, William Chan, David J Fleet, Mohammad Norouzi, and Tim Salimans. Cascaded diffusion models for high fidelity image generation. J. Mach. Learn. Res., 23(47):1-33, 2022.
Allan Jabri, David Fleet, and Ting Chen. Scalable adaptive computation for iterative generation. arXiv preprinl 2iv: 22212.11991 11972, 2022.
Tero Karras, Miika Aittala, Timo Aila, and Samuli Laine. Flucidating the design space of diffusion-based generative models. In Proc. NeurIPS, 2022.
Tuomas Kynkäänniemi, Tero Karras, Samuli Laine, Jaakko Lehtinen, and Timo Aila. Improved precision and recall metric for assessing generative models. Advances in Newral Information Processing Systems, 32, 2019.
Tianhong Li, Huiwen Chang, Shlok Kumar Mishra, Han Zhang, Dina Katabi, and Dilip Krishnan. Mage: Masked generative encoder unify representation learning and image synthesis. arXio preprint arXiv:2211.09117 2022a.
Yanghao Li, Haoqi Fan, Ronghang Hu, Christoph Feichtenhofer, and Kaiming He. Scaling language-inage pre-training via masking. arXiv preprint arXiv:2212.00799, 2022b.
Guan-Horng Liu, Arash Vahdat, De-An Huang, Evangelos A Theodorou, Weili Nie, and Anima Anandkumar. Psb: Image-to-imago schrödinger bridge. arXiv preprint arXim:2302.0587 2023.
Tya Loshchilov and Frank Hutter. Decoupled weight decay regularization. arXie preprint arXiu:1711.051001 2017.
Cheng Lu, Yuhao Zhou, Fan Bao, Jianfei Chen, Chongxuan Li, and Jun Zhu. Dpm-solver: A fast ode solver for diffusion probabilistic model sampling in around 10 steps. arXiv preprint Xtv:2220 00927, 2022.
Chenlin Meng, Ruiqi Gao, Diederik P Kingma, Stefano Ermon, Jonathan Ho, and Tim Salimans. On distillation of guided diffusion models. arXiv preprint arXio:2210.08142, 2022.
Charlic Nash, Jacob Menick, Sander Dieleman, and Peter Battaglia. Generating images with sparse repre- sentations. In International Conference on Machine Learning, pp. 7958-7968. PMLR, 2021.
Weili Nie, Brandon Guo, Yujia Huang, Chaowei Xiao, Arash Vahdat, and Anima Anaudkumar. Diffusion models for adversarial purification. In International Conference on Machine Learning (ICML), 2022.
William Peebles and Saining Xie.
Scalable diffusion models with transformers. arXie preprint
arXiv:2212.099 09748, 2022.
Alec Radford, Karthik Narasimhan, Tim Salimans, Ilya Sutskever, et al. Improving language understanding by generative pre-training. OpenAl, 2018.
Aditya Ramesh, Prafulla Dhariwal, Alex Nichol, Casey Chu, and Mark Chen. Hicrarchical text-conditional image generation with clip latents. Xix preprint arXiv:2204.000 06600200061,555 2022.
14

---

Published in Transactions on Machine Learning Research (March/200 /2024)
Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Björn Ommer. High-resolution image synthesis with latent diffusion models. In Proceedings of the TEEE/CV Conference on Compuler Vision and Pattern Recognition, pp. 10684 10695, 2022.
Olaf Ronneberger, Philipp Fischer, and Thomas Brox. U-net: Convolutional networks for biomedical image segmentation. In Medical Image Compuling and Compuler-Assisleee Intervenlion-MiCCS MICCAI 2015: 18th International Conference, Munich, Germany, October 5-9, 2015, Proceedings, Part 1I1 18, pp. 2412 Springer, 2015.
Nataniel Ruiz, Yuanzhen Li, Varun Jampani, Yael Pritch, Michael Rubinstein, and Kfir Aberman. Dream- booth: Fine tuning text-to-image diffusion models for subject-driven generation. 2022.
Chitwan Saharia, William Chan, Huiwen Chang, Chris Lee, Jonathan Ho, Tim Salimans, David Fleet, and Mohammad Norouzi. Palette: Image-to-image diffusson models. In ACM SIGGRAPH 2022 Conference Proceedings, pp. 1-10, 2022a.
Chitwan Saharia, William Chan, Saurabh Saxena, Lala. Li, Jay Whang, Emily I. Denton, Kamyar Ghasemipour, Raphael Gontijo Lopes, Bureu Karagol Ayan, Tim Salimans, et al. Photorealistic text- to-image diffusion models with deep language understanding. Advances in Neural Information Processing Systems, 35:36479 033449996644,, 2022b.
Tim Salimans and Jonathan Ho. Progressive distillation for fast sampling of diffusion models. In International Conference on Learning Representations, 2022.
Tim Salimans, Ian Goodfellow, Wojeiech Zaremba, Vicki Cheung, Alec Radford, and Xi Chen. Improved techniques for training gans. Advances in neural information processing systems, 29, 2016.
Axel Sauer, Katja Schwarz, and Andreas Geiger. Stylegan-xl: Scaling stylegan to large diverse datasets. In ACM SIGGRAPH 2022 conference proceedings, pp. 10, 2022.
Jascha Sohl-Dickstein, Eric Weiss, Niru Maheswaranathann, and Surya Ganguli. Deep unsupervised learning using nonequilibrium thermodynamics, In International Conference on Machine Learning, pp. 2256 2265. PMLR, 2015.
Jiaming Song, Chenlin Meng, and Stefano Ermon. Denoising diffusson implicit models. In International Conference Learning Representations, 2021a.
Yang Song, Jascha Sohl-Dickstein, Diederík P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. In International Conference Learning Representations, 2021b.
Yang Song, Prafulla Dhariwal, Mark Chen, and Tly Sutskever. Consistency models. ar preprinl 3031022303.01169, 2023.
Cristian Vaccari and Andrew Chadwick. Deepfakes and disinformation: Exploring the impact of synthetic political video on deception, uncertainty, and trust in news. Social Media+ Society, 6(1):2056305) 209003400, 2020.
Arash Vahdat, Karsten Kreis, and Jan Kautz. Score-based generative modeling in latent space. Advances in Neural Information Processing Systems, 34:11287 11302, 2021.
Aaron Van Den Oord, Oriol Vinyals, et al. Neural discrete representation learning. Advances in neural information processing systems, 30, 2017.
Ashish Vaswani, Noam Shazoer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems, 30, 2017.
15

---

Published in Transactions on Machine Learning Research Mlarch/2024/
Zhendong Wang, Yifan Jiang, Huangjie Zheng, Peihao Wang, Pengeheng He, Zhangyang Wang, Weizhu Chen, and Mingyuan Zhou. Patch diffusion: Faster and more data-efficient training of diffusion models. arXiv preprint arXi::2304..100 12526, 2023.
Chen Wei, Karttikeya Mangalam, Po-Ya Yao Huang, Yanghao Li, Haoqi Fan, Hu Xu, Huiyu Wang, Cihang Xic, Alan Yuille, and Christoph Feichtenhofer. Diffusion models as masked autoencoders. arXiv preprint arXiv:2304.0390 2023.
Xiulong Yang. Sheng-Min Shih, Yinlin Fu, Xiaoting Zhao, and Shihao Ji. Your vit is secretly a hybrid diseriiinatvve-gneeaaiee diffusion model. arXiv preprint ar:::22008.077 2022.
Hongkai Zheng, Weili Nie, Arash Vahdat, Kamyar Arizzaeneshelli, and Anima Anandkumar, Fast sampling of diffusion models via operator learning. arxiv preprint orXiv:2211.13449 2022.
Appendix
A More experimental settings
Diffusion configurations. Different from DIT which is basod on the ADM formulation (Dhariwal Nichol 20210, we use the EDM formulation (Karras et al. 2022) for training simplicity and inference efficiency. Specifically,
the EDM preconditioning via o-dependent skip connection with the de- fault hyperparameters (see (Karras et al. 2022) for more details). Thus, do not need to learn ADM's parameterization of the noise covariance as in DIT. During inference, we use the default time schedule
4<N = (thax + N_ (tmin )), where = 40, 40, p = 7, tmax = 80 and tmin = 0.002, and the second-order Heun's s method as the ODE solver for sampling. Our early experiments with the original DîT architecture show that Heun's 's mothod reaches the same FID as 250 DDPM sampling steps used in DiT with a much lower number of function evaluations (79 vs. 250 evaluations), confirming the observations in (Karras el al. 2022).
Training details. We use the pre-trained VAE model ft-MSE with scaling factor 8 from the latent diffusion model (Rombach et al.) 2022). To avoid repeatedly encoding the same images during training, we use the VAE encoder to encode the dataset into latent space and store them on the disk. All the experiments are running on this latent dataset, including the runs of DiT and MDT. For all the masked training of MaskDiT, we enable automatic mixed precision. During the unmasking tuning, we instead use TensorFFoatt) for better numerical stability.
B Benchmarking DiT and MDT
To make fair comparison with DiT (Pebles Xie 2022) and MDT Gao el al. 2023), we use their official implementations. Since DiT, MDT and our MaskDiT use the same VAE encoder, we replace the original dataset with the latent dataset to reduce the number redundant encoding processes. We benchmark their highest-capacity models DiT-XL/2 and MDT-XL/2, with the default setting in their paper.
C Compute used for the experiments
We run all the experiments on 8x A 100 GPUs. The main experiment of MaskDiT takes 273 hours, of which 257 hours are spent on training and 16 hours on unmasking tuning. All the experiments in the ablation studies take around 1250 hours total. The experiments on the training efficiency MaskDiT, DIT-XL/2,, and MDT-XL/2 take around 1000 hours in total. The preliminary or failed experiments reported in the paper take around 800 hours in total.
16

---

Published in Transactions on Machine Learning Research Mlarch/2024/
D FID vs. training steps
Figure 7 plots the FID vs. training steps of the same experiments report in Figure 5. MaskDiT, DiT- XL/2, and MDT-XL/2 are trained for the same amount of time. As shown in the figure, MaskDiT has a similar or even faster learning speed as 2 in terms of steps. The advantage of MaskDiT over DiT- XL/2 becomes more significant when the batch size is larger. Additionally, every training step MaskDiT is around 45% cheaper than DiT-XL/2. Therefore MaskDiT is much more efficient than DiT-XL/2. MDT- X2/ has laster learning speed than DIT-XT./2 and MaskDiT, but its per-step cost is much higher. We do not report MDT's results after 220k steps in the setting of batch size 256 because encountered numerical issue while running MDT's code.
DIT-XL/2
DIT-XL/2
70
70
MDT-XL/2
MDT-XL/2
60
MaskDiT
60
MaskDiT
50
50
40
40
30
30-
20
20
10 Training steps
1.2 1e6
le6
Training steps
Figure 7: FID training steps. Left: batch size 256. Right: batch size 1024. For batch size 256, we only reported the FID of MDT before 220k steps, because it encountered gradient explosion issue afterwards when we ran the official MDT code with their default setting.
E Discussions on masking ratio schedule in unmasking tuning
In this paper, we explore two strategies for unmasking tuning. The first one is disable masking during finetuning completely (i.e., the zero-ratio schedule). This simple strategy effectively improves the FID with classifier-free guidance from 4.54 to 2.28. The second strategy adopts a cosine-ratio schedule where the mask ratio gradually decreases from 50% to 0%. More precisely, the masking ratio is calculated as = 0.5 * cos (T/2 i/Trot) with i and not.. being the current steps and the total number of tuning steps, respectively. We find this strategy improves the FID without guidance from 5.95 to 5.69 within 37.5k steps. Clearly, different unmasking tuning stratogics lead different results. We do not fully explore this direction in this paper as it is not the focus of this paper. However, it would be interesting future direction to investigate different tuning strategies for masked training of diffusion transformers.
F More generated samples
Figures S
g and 10 show three diverse sets of images conditionally sampled from our MaskDiT with guidance
2.5 using 40 EDM Karras et al. 2022) sampling steps.
17

---

Published in Transactions on Machine Learning Research (March/200)
"alp"
3
"black swan"
"daisy"
"Eskimo dog"
"hamster"
D nanas
"orange"
"promontory"
Figure 8: Class- conditional x56x256 image generation from MaskDiT with guidance 2.5 and 40 deterministic EDM
Karras et al. 2022) sampling steps. Every four images in row are from the same class.
18

---

Published in Transactions on Machine Learning Research (March /2024)
"beacon"
"brambling"
"lakeside"
"coral reef"
"tarantura"
"tusker"
"great white shark"
Figure 9: Class-conditional 256 image generation from MaskDiT with guidance 2.5 and 40 deterministic EDM (Karras al. 2022) sampling steps. Every four images in row from the same class.
19

---

Published in Transactions on Machine Learning Research (March /2024)
"minibus"
"wood rabbit"
"giant panda"
"white wolf"
"volcano"
"Siamese cat"
"pier"
Figure 10: Class-conditional 256x256 image generation from MaskDiT with guidance 2.5 and 40 determin- istic EDM (Karras et al. 2022 sampling steps. Every four images in a are from the same class.
20

---

Published in Transactions on Machine Learning Research (March
/2024)
Figure 11: Class- consicoondiionnl 512x512 image generation from MaskDiT with guidance 1.5 and 40 determin- istic EDM KKrras al. 2022) sampling steps.
21

---

Published in Transactions on Machine Learning Research (March /2024)
Figure 12: Class conditonn 512x512 image generation from MaskDiT with guidance 1.5 and 40 determin-
istic EDM Karras et al. 2022) sampling steps.
22


Reviews:
[{"review_id": "jIfkkT5lKp", "reviewer": "Reviewer_difn", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths: \n- The authors propose a simple technique to improve training efficiency of training diffusion transformer models. The approach is motivated by  the same intuitions as the Masked Auto Encoder paper, and benefits from the same simplicity. I like that the approach is simple to implement, and is quite effective. \n- I enjoyed reading the paper! It's very easy to follow, and the experiments section is fairly exhaustive. The ablations in the paper are well done and insightful. \n\nWeaknesses:\nGiven the scope of the paper (improving diffusion transformer model training efficiency while maintaining generation performance), the paper claims are well supported by the experiments. However, there is one additional experiment that I'd like to see: \n- I am curious to see if the proposed approach can be effective for higher resolution images. To complete the manuscript, it would be nice to show comparison on class-conditional ImageNet 512x512 with DiT.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "Thea authors propose a new faster training paradigm for transformer-based diffusion models which retains generation performance of the original model (Diffusion Transformer, DiT) while taking only 31% of its training time. The approach involves randomly masking out input patches, and only performing denoising score matching on unmasked tokens. This significantly reduces the memory and compute required to train on full tokens thus achieving faster training of these models. The intuition behind using masked tokens is that due to high visual redundancy in images, learning to denoise only a small set of unmasked patches in the image doesn't affect the generation performance. To prevent overfitting on unmasked patches and to understand global visual information, they add the MAE objective of reconstructing the masked patches via MSE loss as an auxiliary loss. The paper show experiments on class-conditional ImageNet-256x256 generation benchmark.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "See Weaknesses.", "broader_impact_concerns": "None", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "4Vyb8aOhWy", "reviewer": "Reviewer_eJnL", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "The claim of improving training speed is very relevant and well supported by the experiments.\nThe paper is well written except for some minor spelling and grammar mistakes (see list below of those I caught). \n\nMy major concern is that of using a loss function without a principled grounding in the marginal likelihood (such as a variational bound) and using it directly on a set of heuristics, like FID and IS. My personal experience with training such models is limited, so I have very little intuition on how often such models and their learning rules can escape such metrics and game them.\nFrom what I know, randomly leaving out a reconstruction terms can be seen as an unbiased estimate of that reconstruction. Leaving out inputs/conditions however can lead to bias in the context of variational inference. A theoretical analysis would have been a great plus.\n\nStill, I think that the paper and the experiments are executed well.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The authors propose an architectural change and a special loss function to speed up training for diffusion models. The method appears to work according to well-established metrics of generative performance of image data.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "Please fix the following mistakes.\n\n- \"our approache achieves ~~the~~ similar performance ...\" (Abstract)\n- \"Our main hypothesis is that ~~image~~ images contain* significant redundancy in the pixel space.\" (Introduction)\n- \"We propose a new approach for ~~the~~ fast training of ...\" (Introduction)\n- \"trades a slightly ~~more~~ higher training cost for ~~the~~ better performance of ~~the~~ CFG sampling.\" (3.2 Key designs)\n- Replace upper/bottom with left/right in the caption of fig 5\n\nAlso, you only have to cite Dosovitskiy once per section per mentioning of ViT.", "broader_impact_concerns": "none", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "qPjOFdikdq", "reviewer": "Reviewer_FZfH", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "**Strengths**\\\n        **a**. The proposed MaskDiT largely saves the number of parameters to train diffusion models using the asymmetric encoder-decoder masked DiT as the backbone of diffusion models and the MAE loss.\\\n        **b**. The presentation is in general clear and comprehensive. The whole paper is easy to follow.\\\n\n**Weaknesses**\\\n        **a**. The experiments are limited to ImageNet-256$\\times$256, where the commonly applied benchmark datasets for diffusion models with higher resolutions are not tested. Therefore, the main contribution of this work on the training efficiency should be further evaluated. The commonly applied high-resolution datasets for diffusion models include the CelebA-HQ dataset consisting of 30,000 images at $1024\\times 1024$ resolution.\\\n        **b**. A minor weakness is that the motivation should be more coherent. There have been a growing number of pretrained diffusion models which are publicly available. Given these pretrained models, the sampling efficiency is the key to the efficiency of the diffusion models. To enhance the sampling efficiency, there have been a number of effective methods (Song et al., 2021; Lu et al., 2022; and Song et al., 2023). Therefore, why is it still necessary to  propose MaskDiT for higher the higher training efficiency?\\\n        **c**. Fast training and sampling of diffusion models has been also studied in other existing work including what follows the theoretical insights from optimization or advanced numerical solvers (Wu et al., 2023). However, the authors claimed that they cannot be directly applied to reduce the training cost of diffusion models. It would enhance the comprehensiveness of this work to provide a discussion on the differences between the other existing method (Wu et al., 2023) and MaskDiT.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "In this work, the authors proposed MaskDiT to accelerate the training speed of diffusion models, following the asymmetric encoder-decoder architecture. For the experiments on ImageNet-256$\\times$256, MaskDiT largely reduces the computational cost without significant loss on performance compared with existing diffusion models.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "See the weaknesses.\\\nIn general, the experiments really should be extended and articulation of the related work also needs to be improved.", "broader_impact_concerns": "NA", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: As acknowledged by the reviewers, the claims are empirically supported and the paper is easy to follow. Moreover, the authors improved their submission using the reviewers' feedback.
Decision: Accept (poster)