MDT

Project Introduction: Masked Diffusion Transformer V2

Overview

Masked Diffusion Transformer V2 (MDTv2) is a cutting-edge image synthesis framework designed to improve upon the existing capabilities of diffusion probabilistic models (DPMs). Compared to previous models like DiT, MDTv2 stands out by delivering superior image generation performance along with a remarkably faster learning speed, surpassing its predecessors by over 10 times in training efficiency.

The MDTv2 Advantage

• State-of-the-Art Performance: MDTv2 has achieved a new state-of-the-art Frechet Inception Distance (FID) score of 1.58 on the ImageNet dataset, a benchmark for measuring the quality of generated images.
• Enhanced Learning Speed: It provides a fivefold acceleration over the original Masked Diffusion Transformer (MDTv1), demonstrating a significant improvement in processing time and efficiency.
• Contextual Reasoning: One of MDTv2's key enhancements is its ability to understand and learn from the contextual relationship among different parts of an image, which allows it to generate more coherent and detailed images.

How It Works

MDTv2 leverages a unique approach to image reconstruction using a mask latent modeling scheme. This method enhances DPMs by introducing masked tokens that the model predicts from the unmasked ones, maintaining the diffusion generation process intact. This mechanism empowers MDT to rebuild an image's full information from incomplete input, teaching it to learn relationships among different image segments effectively.

Experimental Success

The extensive tests conducted on MDTv2 reveal superior performance metrics:

Model Accessibility

The pre-trained model is available for download on Hugging Face, a popular platform for AI model hosting. Users can easily integrate it into their systems using:

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

Getting Started

Setup

To start using MDTv2, ensure PyTorch 2.0 or higher is installed in your environment. Additionally, downloading the code repository and setting up the Adan optimizer, known for its faster convergence than AdamW, are essential steps:

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

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

Data Preparation

MDTv2 supports ImageNet and CIFAR datasets directly. For custom datasets, ensure images are named in the format ClassID_ImgID.jpg to allow proper class recognition during training.

Training and Evaluation

MDTv2 provides comprehensive scripts for training and evaluating the model in different configurations, supporting single or multi-node setups. Detailed instructions are available within the project’s repository to facilitate these processes.

Visualization

To visualize results directly from MDTv2, the 'infer_mdt.py' script can be used to generate images from the trained models.

Final Thoughts

MDTv2 represents a significant leap forward in the field of image synthesis, offering both improved performance and faster training capabilities. It provides a robust platform for researchers and developers looking to explore advanced AI-driven image generation techniques.