MDTv2: Masked Diffusion Transformer is a Strong Image Synthesizer

Shanghua Gao; Pan Zhou; Ming-Ming Cheng; Shuicheng Yan

Source-linked AI summary

MDTv2: Masked Diffusion Transformer is a Strong Image Synthesizer

Shanghua Gao, Pan Zhou, Ming-Ming Cheng, Shuicheng Yan

arXiv:2303.14389v2Mar 25, 2023cs.CV

TL;DR

Diffusion probabilistic models can learn object parts independently rather than their contextual relations, slowing training. MDT introduces latent token masking and an asymmetric transformer to reconstruct masked information, while MDTv2 adds macro-architecture and training improvements; it achieves higher image-synthesis performance and faster learning, including a 5× speedup over MDT.

Problem
Diffusion probabilistic models often fail to learn associated relations among semantic object parts because per-pixel prediction overlooks those relations, slowing training.
Method
MDT masks latent image tokens and uses an asymmetric diffusion transformer to predict masked tokens from unmasked contextual inputs; MDTv2 further improves the macro network structure and training strategy.
Results
MDTv2 achieves higher image-synthesis performance and learns approximately five times faster than MDT, while MDT achieves FID 6.23 compared to DiT’s 9.62.
Takeaways & Limitations
Masked latent modeling improves contextual representation and relation learning among image semantics while substantially improving image-synthesis performance and training progress.

Abstract

from arXiv · show

Despite its success in image synthesis, we observe that diffusion probabilistic models (DPMs) often lack contextual reasoning ability to learn the relations among object parts in an image, leading to a slow learning process. To solve this issue, we propose a Masked Diffusion Transformer (MDT) that introduces a mask latent modeling scheme to explicitly enhance the DPMs' ability to contextual relation learning among object semantic parts in an image. During training, MDT operates in the latent space to mask certain tokens. Then, an asymmetric diffusion transformer is designed to predict masked tokens from unmasked ones while maintaining the diffusion generation process. Our MDT can reconstruct the full information of an image from its incomplete contextual input, thus enabling it to learn the associated relations among image tokens. We further improve MDT with a more efficient macro network structure and training strategy, named MDTv2. Experimental results show that MDTv2 achieves superior image synthesis performance, e.g., a new SOTA FID score of 1.58 on the ImageNet dataset, and has more than 10x faster learning speed than the previous SOTA DiT. The source code is released at https://github.com/sail-sg/MDT.

1 INTRODUCTION

Diffusion probabilistic models learn relations among semantic object parts slowly, motivating masked latent modeling and an asymmetric transformer that reconstructs masked tokens from contextual inputs. MDTv2 further improves this approach with a more efficient network and training strategy, yielding faster learning and strong ImageNet synthesis results.

Motivation: DPMs often learn semantic object parts independently, slowing convergence because per-pixel prediction loss ignores relations among those parts.Examples show incomplete or inaccurately positioned dog features persisting across training steps.
Results: 3× faster learning than DiT is reported for MDT, while MDT achieves superior FID scores.The comparison is measured with respect to training steps and time on 8×A100 GPUs.
Method: MDT masks image tokens in latent space and uses an asymmetric diffusion transformer to predict masked tokens from unmasked ones.The design preserves the diffusion generation process while encouraging contextual relation learning.
Method: MDT reconstructs full image information from incomplete contextual input and learns associated relations among semantic parts.The paper contrasts MDT’s coordinated synthesis of both dog eyes with DiT’s difficulty modeling correct part relations.
Results: MDT achieves a new SOTA result for class-conditional image synthesis on ImageNet and synthesizes images with less training time than SOTA methods.The supplied passages identify the ImageNet result and reduced training time but do not provide the final FID value here.
Results: 5× faster training than MDT is reported for MDTv2 after introducing a more efficient macro network structure and enhanced training strategies.The improvements include long encoder shortcuts, dense decoder input shortcuts, Adan, adapted loss weights, and expanded masking ratios.

2 RELATED WORKS

Diffusion models use iterative stochastic processes for image synthesis, while related work has pursued architectural, masking, and training-based improvements. Mask modeling has shown benefits for representation quality, scalability, convergence, and generative modeling across language and vision.

Diffusion Probabilistic Models: DPMs add Gaussian noise through a forward SDE and reverse the process with a time-inverted SDE that progressively maps noise into samples.Their iterative sampling procedure contributes to high training and inference costs.
Diffusion Probabilistic Models: Diffusion research has reduced inference costs through efficient sampling, latent-space diffusion, and multi-resolution cascaded generation.Other training schemes address approximate maximum likelihood training and loss weighting.
Diffusion Architectures: Diffusion architectures evolved from UNet-like networks with attention and normalization toward transformer-based designs motivated by transformers’ broad applicability.The passage identifies spatial self-attention, group normalization, additional attention heads, residual blocks, and adaptive normalization as architectural developments.
Mask Modeling: MDT’s position-aware designs indicate that mask latent modeling benefits from a stronger diffusion network.The authors identify further exploration of these networks as future work.
Mask Modeling: Mask modeling supports representation pretraining and language generation in NLP, and has also proved feasible for vision recognition and generation.In vision recognition, mask-based pretraining is associated with representation quality, scalability, and faster convergence.

3 MASKED DIFFUSION TRANSFORMER

MDT improves diffusion training by masking latent tokens and reconstructing them from contextual information with an asymmetric transformer. MDTv2 adds shortcut-rich architecture and training strategies to accelerate convergence, while power-cosine guidance balances diversity and fidelity across sampling steps.

3.1 Overview: MDT performs diffusion learning in latent space, where noisy image tokens are partially masked and reconstructed from unmasked contextual tokens.A pre-trained VAE provides the latent representation, reducing computational costs compared with pixel-space diffusion.
3.1 Overview: The masking objective addresses slow learning of relations among semantic parts by forcing reconstruction of complete image information from incomplete contextual input.The approach is intended to make relations such as those among an image’s object parts available during training.
3.3 Asymmetric Diffusion Transformer: MDT’s asymmetric diffusion transformer uses an encoder, side-interpolater, and decoder to predict masked tokens while retaining the diffusion generation process.The encoder processes unmasked tokens during training; the side-interpolater restores a full token sequence for the decoder and is removed during inference.
3.3 Asymmetric Diffusion Transformer: Local relative positional bias supplies token-position relations that facilitate masked latent modeling.The bias is selected from relative positional differences and updated during training.
3.4 Masked Diffusion Transformer v2: MDTv2 adds U-Net-style long shortcuts in the encoder and dense input-shortcuts in the decoder to improve diffusion and masking-reconstruction learning.Dense input-shortcuts connect the encoder input to every decoder block, preserving noisy input information unavailable through masked encoder patches.
3.4 Masked Diffusion Transformer v2: The shortcut-enhanced macro-structure improves image-generation quality while reducing training time.MDTv2 also uses a faster Adan optimizer, timestep-adapted loss weights, and an expanded masking ratio.
3.5 Training and Inference: Power-cosine classifier-free guidance keeps guidance low early in sampling and increases it quickly late, supporting diversity early and quality late.The schedule uses s to control the increase rate; the paper sets s to 4 and maximum guidance scale w to 3.8.

4 EXPERIMENTS

Experiments show that MDT and MDTv2 improve ImageNet image synthesis over DiT while requiring fewer training steps, with ablations supporting their asymmetric masking design and macro-structure.

Comparison Results: MDT achieves higher FID scores than DiT across model scales with fewer training costs while maintaining similar parameter and inference costs.The additional MDT modules have negligible inference overhead.
Comparison Results: 57.01 versus 68.40 FID: MDT-S/2 trained for 300k steps outperforms DiT-S/2 trained for 400k steps.MDT-S/2 trained for 2,000k steps achieves similar performance to the larger DiT-B/2 under a similar computational budget.
Comparison Results: 1.92 FID improvement: MDTv2-XL/2 trained for 400k steps surpasses DiT-XL/2 trained for 7,000k steps, indicating more than 18× training acceleration.MDTv2-XL/2 also reaches FID 6.11 versus 6.23 after 700 versus 6,500 training steps for original MDT.
Comparison Results: 1.58 FID: MDTv2 sets a new SOTA for class-conditional ImageNet 256×256 generation with fewer training iterations.MDTv2 lowers MDT’s FID from 6.23 to 5.06 with under a third of MDT’s training iterations before the reported SOTA comparison.
Convergence Speed: 5× faster learning: MDTv2 improves training speed over MDT, while MDT is about 3× faster than DiT in both training steps and training time.The comparison is reported using experiments on 8×A100 GPUs.
Ablation on MDT: The asymmetric architecture, side-interpolater, masked shortcut, all-token loss, and positional awareness each improve reported FID in MDT ablations.The asymmetric design reduces training cost and focuses learning on contextual information; positional embeddings and relative positional modeling improve FID from 51.58 to 50.26 and from 53.56 to 50.26, respectively.
Ablation on MDTv2: MDTv2’s long encoder shortcuts and dense decoder input shortcuts reduce FID, while access to noised masked patches improves noise prediction.Expanding decoder blocks from 2 to 6 also produces a stable performance improvement.
Ablation on MDTv2: Middle placement of the side-interpolater gives the best FID, whereas early or final placement harms performance; decoder depth effects depend on training duration.With 600k iterations, six decoder blocks outperform four, but after 1,500k iterations their performance is comparable.

5 CONCLUSION

The paper proposes masked latent modeling with an asymmetric diffusion transformer to improve contextual representation and semantic relation learning in DPMs. Experiments report higher image-synthesis performance and faster training, including a new ImageNet SOTA.

5 CONCLUSION: MDT introduces masked latent modeling and an asymmetric diffusion transformer to improve contextual representation and relation learning among image semantics.The method is designed for diffusion probabilistic models.
5 CONCLUSION: Experiments report higher image-synthesis performance and substantially improved learning progress during training, achieving new SOTA performance on ImageNet.The conclusion summarizes the method’s reported performance and training-progress gains.