Masked Diffusion Transformer V2

The official codebase for Masked Diffusion Transformer is a Strong Image Synthesizer.

MDTv2: Faster Convergeence & Stronger performance

MDTv2 achieves superior image synthesis performance, e.g., a new SOTA FID score of 1.58 on the ImageNet dataset, and has more than 10× faster learning speed than the previous SOTA DiT.

MDTv2 demonstrates a 5x acceleration compared to the original MDT.

MDTv1 code

Introduction

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 10× faster learning speed than the previous SOTA DiT.

Performance

Pretrained model download

Model is hosted on hugglingface, you can also download it with:

from huggingface_hub import snapshot_download
models_path = snapshot_download("shgao/MDT-XL2")
ckpt_model_path = os.path.join(models_path, "mdt_xl2_v1_ckpt.pt")

A hugglingface demo is on DEMO.

NEW SOTA on FID.

Setup

Prepare the Pytorch >=2.0 version. Download and install this repo.

git clone https://github.com/sail-sg/MDT
cd MDT
pip install -e .

Install Adan optimizer, Adan is a strong optimizer with faster convergence speed than AdamW. (paper)

python -m pip install git+https://github.com/sail-sg/Adan.git

DATA

• For standard datasets like ImageNet and CIFAR, please refer to 'dataset' for preparation.
• When using customized dataset, change the image file name to ClassID_ImgID.jpg, as the ADM's dataloder gets the class ID from the file name.

Training

Evaluation

The evaluation code is obtained from ADM's TensorFlow evaluation suite. Please follow the instructions in the evaluations folder to set up the evaluation environment.

Visualization

Run the infer_mdt.py to generate images.

Citation

@misc{gao2023masked,
      title={Masked Diffusion Transformer is a Strong Image Synthesizer}, 
      author={Shanghua Gao and Pan Zhou and Ming-Ming Cheng and Shuicheng Yan},
      year={2023},
      eprint={2303.14389},
      archivePrefix={arXiv},
      primaryClass={cs.CV}
}

Acknowledgement

This codebase is built based on the DiT and ADM. Thanks!