Source-linked AI summary
Fast Training of Diffusion Models with Masked Transformers
Hongkai Zheng, Weili Nie, Arash Vahdat, Anima Anandkumar
TL;DR
Diffusion models are expensive to train, and masked training is not straightforward with conventional U-Net backbones. MaskDiT masks image patches, combines an asymmetric transformer with reconstruction and score-matching objectives, and achieves competitive generation with roughly 30% of DiT’s training time.
Problem
Training large diffusion models is costly, while masked training cannot be directly applied to conventional U-Net backbones because their convolutions operate on dense grids.
Method
MaskDiT randomly masks input patches, processes visible patches with a transformer encoder and full patches with a lightweight decoder, and combines DSM with masked-patch reconstruction.
Results
31% of DiT’s training time produces FID 2.28 with guidance on ImageNet-256×256, while MaskDiT outperforms DiT-XL/2 on ImageNet-512×512 using about 29% of its training cost.
Takeaways & Limitations
MaskDiT provides a more training-efficient transformer-based diffusion approach while maintaining comparable or better generative performance.
Takeaways & Limitations
MaskDiT requires a few unmasking-tuning steps to match state-of-the-art guided FID and does not produce a good unconditional diffusion model.
Abstract
from arXiv · showhide
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.
1 Introduction
Diffusion-model training remains costly, and conventional U-Net backbones make masked training difficult. MaskDiT adapts masked transformer training with an asymmetric architecture and auxiliary reconstruction, reducing training cost while retaining competitive generation quality.
- Diffusion models require large-scale training for realistic image generation, but improving their training efficiency remains an open question.
- Masked training improves efficiency in representation learning, but U-Net convolutions operate on dense grids, making random patch masking difficult for diffusion models.
- Its objective combines denoising score matching on unmasked patches with mean-square-error reconstruction of masked patches, followed by unmasked finetuning costing less than 6% of total training.
- 31% of DiT’s training time yields FID 2.28 with guidance on ImageNet-256×256, while MaskDiT also outperforms DiT-XL/2 at 512×512 using about 29% of its training cost.
- MaskDiT randomly masks 50% of input patches and uses an asymmetric encoder-decoder architecture in which the encoder processes unmasked patches and a lightweight decoder processes full patches.
2 Related work
Diffusion research has progressed from convolutional U-Nets toward scalable transformer backbones, while efficiency work has largely targeted sampling. Masked training extends transformer-based generative modeling but remains an emerging direction for diffusion models.
- Backbone architectures of diffusion models: U-Net became the standard diffusion backbone, while newer transformer architectures exploit transformers’ expressivity and flexibility for image generation.
- Backbone architectures of diffusion models: DiT is a pure transformer diffusion architecture that achieves better class-conditional generation than U-Net counterparts with improved scalability.
- Efficiency in diffusion models: Diffusion models have high training and inference costs, but much prior efficiency research has focused on improving sampling through numerical solvers and related strategies.
- Masked training with transformers: Masked training is broadly used in representation learning and generative modeling across language and vision, including masked language modeling and masked image-token prediction.
- Masked training of diffusion models: DiffMAE applies masked diffusion mainly to discriminative and inpainting tasks, whereas MDT adds a masked objective to the original DiT loss for generative modeling.
3 Method
MaskDiT combines masked image patches with an asymmetric encoder-decoder and a two-part objective for efficient diffusion training. The method estimates scores on unmasked patches, reconstructs masked patches, and uses classifier-free guidance during generation.
- Classifier-free guidance: Classifier-free guidance modifies the class-conditional denoising function by combining conditional and unconditional predictions with guidance scale w ≥ 1.The method separately notes that masked training makes learning the unconditional score harder than learning the conditional score.
- Image masking: MaskDiT patches diffused images and randomly removes a fixed fraction of patches before passing the remaining tokens to the diffusion model.For an H × W image with patch size p × p, the number of patches is N = (HW)/p^2, and ⌊rN⌋ patches are masked.
- Asymmetric encoder-decoder backbone: The asymmetric backbone runs the transformer encoder only on unmasked patches and uses a lightweight decoder on the complete token set.The decoder adds shared learnable mask tokens and positional embeddings; its MAE decoder has fewer than 9% of DiT-XL/2's parameters.
- Training objective: The training objective estimates denoising scores on unmasked tokens instead of applying score matching to full tokens.The design avoids requiring the model to predict masked-patch scores solely from visible patches.
- Training objective: An auxiliary MSE loss reconstructs the diffused pixel values of masked patches to encourage global understanding and reduce local overfitting.The reconstruction target is the diffused image x0 + n, evaluated on masked patches.
- Training objective: The total objective combines denoising score matching and masked reconstruction as L = LDSM + λLMAE, with λ balancing the two terms.The paper notes that λ cannot be too large because it could move training away from the standard DSM update.
4 Experiments
Experiments evaluate MaskDiT’s setup, efficiency, generation quality, and ablations on ImageNet at two resolutions. MaskDiT substantially reduces computational and memory costs while maintaining competitive generation performance, with effectiveness depending on masking, reconstruction, score-matching, and unmasking-tuning choices.
- Experimental setup: MaskDiT is evaluated in a latent diffusion framework using a pretrained VAE, with default 50% masking, λ = 0.1, and batch size 1024.The evaluation uses FID as the primary metric, with sFID, Inception Score, and Precision/Recall as secondary metrics.
- Training efficiency: MaskDiT uses only 54.0% of DiT’s GFLOPs and 31.7% of MDT’s while providing higher training speed and lower memory consumption.These comparisons use similarly sized models evaluated on 8× A100 GPUs.
- Training efficiency: With batch size 1024 on ImageNet-256×256, MaskDiT trains 3.5× faster than DiT and 6.5× faster than MDT while using 45.0% and 19.2% of their memory, respectively.At 512×512, MaskDiT reaches 4.6× DiT’s training speed while using 45.7% of DiT’s GPU memory.
- Training efficiency: On ImageNet-256×256 with batch size 1024, MaskDiT reaches FID 10 within 40 hours, whereas DiT and MDT require more than 160 hours.The reported wall-time comparison implies a 4× speedup.
- Comparison with state-of-the-art: MaskDiT achieves FID 5.69 without guidance and FID 2.28 with guidance on ImageNet-256×256, reaching the guided result in 31% of DiT-XL/2’s training time.MaskDiT-G obtains FID 2.28 in 273 hours, compared with 868 hours for DiT-XL/2.
- Comparison with state-of-the-art: On ImageNet-512×512, MaskDiT achieves FID 10.79 without guidance and FID 2.50 with guidance, using about 29% of DiT’s training cost.The guided result outperforms the cited ADM and DiT baselines.
- Ablation studies: A 50% masking ratio benefits the asymmetric architecture, whereas 75% masking fails to preserve the same performance.At 50% masking, asymmetric-architecture FID changes from 13.71 to 12.31, while original DiT changes from 14.70 to 24.58.
- Ablation studies: Adding MAE reconstruction improves FID from 12.42 to 7.19 at 800k steps, while λ = 1.0 eventually worsens FID and λ = 0 saturates early.The authors attribute the benefit to improving global understanding beyond the unmasked-token subset.
5 Broader impacts and limitations
The work reduces diffusion-model training costs algorithmically without sacrificing generative performance, but guided generation still requires unmasking tuning and unconditional performance remains limited.
- MaskDiT reduces diffusion-model training costs without sacrificing generative performance.
- Matching state-of-the-art FID with guidance still requires a few steps of unmasking tuning.
- Masked training does not produce a good unconditional diffusion model, which remains future work.
6 Conclusions
MaskDiT trains diffusion models by masking image patches and combining an asymmetric encoder-decoder with auxiliary reconstruction. On class-conditional ImageNet benchmarks, it improves training efficiency while maintaining competitive generative performance.
- MaskDiT randomly masks image patches to reduce training overhead per iteration.
- Its asymmetric backbone uses a DiT encoder for visible tokens and a lightweight decoder for full tokens after masked-token injection.
- An auxiliary reconstruction loss predicts masked-token inputs while the DSM objective predicts scores for unmasked tokens.
- On class-conditional ImageNet-256×256 and ImageNet-512×512, MaskDiT provides better training efficiency with competitive generative performance.
A More experimental settings
The experiments use the EDM formulation with Heun sampling, cached latent representations, and mixed-precision training choices that differ between masked training and unmasking tuning.
- The experiments use EDM preconditioning with a σ-dependent skip connection instead of DiT’s ADM formulation.
- Heun’s second-order ODE solver uses N = 40, ρ = 7, tmax = 80, and tmin = 0.002 for sampling.
- Heun sampling reaches the same FID as DiT’s 250 DDPM steps using 79 rather than 250 function evaluations.
- A pretrained VAE encodes the dataset into cached latent representations, while masked training uses automatic mixed precision and unmasking tuning uses TensorFloat32.
B Benchmarking DiT and MDT
For fair benchmarking, the study uses official DiT and MDT implementations, a shared cached latent dataset, and their highest-capacity XL/2 models under default settings.
- The comparison uses official implementations of DiT and MDT.
- All three methods use the same VAE encoder and cached latent dataset to reduce redundant encoding.
- The benchmark evaluates the highest-capacity DiT-XL/2 and MDT-XL/2 models with their papers’ default settings.
C Compute used for the experiments
The experiments ran on 8×A100 GPUs, with substantial time devoted to the main experiment, ablations, efficiency studies, and preliminary or failed runs.
- The main MaskDiT experiment took 273 hours on 8×A100 GPUs, including 257 hours of training and 16 hours of unmasking tuning.
- Ablation studies took around 1250 hours in total.
- Training-efficiency experiments comparing MaskDiT, DiT-XL/2, and MDT-XL/2 took around 1000 hours.
- Preliminary or failed experiments not reported in the paper took around 800 hours.
D FID vs. training steps
MaskDiT learns at a similar or faster step-based rate than DiT-XL/2, with a larger advantage at bigger batch sizes and substantially lower cost per step.
- MaskDiT has a similar or faster learning speed than DiT-XL/2 in terms of training steps, with a larger advantage at larger batch sizes.
- 45%: every MaskDiT training step is around 45% cheaper than a DiT-XL/2 step.
- Figure 7 reports FID versus training steps for batch sizes 256 and 1024.For batch size 256, MDT FID is reported only before 220k steps because the official code encountered gradient explosion afterward under its default setting.
E Discussions on masking ratio schedule in unmasking tuning
Unmasking-tuning outcomes depend on the masking-ratio schedule, with zero-ratio and cosine-ratio schedules improving different FID settings.
- 4.54 to 2.28: disabling masking during finetuning improves FID with classifier-free guidance.This is the zero-ratio schedule.
- 5.95 to 5.69: a cosine-ratio schedule improves FID without guidance within 37.5k tuning steps.The masking ratio decreases from 50% to 0% according to r = 0.5 ∗cos4(π/2 ∗i/ntot).
- Different unmasking-tuning strategies lead to different results, and the paper leaves broader schedule exploration for future work.
F More generated samples
The paper provides additional class-conditional MaskDiT samples at 256×256 and 512×512 resolutions using 40-step deterministic EDM sampling.
- Figures 8, 9, and 10 show three diverse sets of class-conditional 256×256 images sampled with guidance 2.5 and 40 EDM steps.Figures 8–10 state that every four images in a row share the same class.
- Figures 11 and 12 show class-conditional 512×512 images sampled with guidance 1.5 and 40 deterministic EDM steps.