Diffusion Transformers (DiTs) From Scratch

Introduction

A clean PyTorch re-implementation of Diffusion Transformers (DiTs), introduced in “Scalable Diffusion Models with Transformers” (2023) [1]. Since the DiT paper builds on “Improved Denoising Diffusion Probabilistic Models” [2] for training details, such as the loss functions, noise scheduler, and diffusion process, this implementation also incorporates those components following that work.

• The codebase is designed to be easy to train, modify, and extend.
• Everything is fully configurable, including DiT model architecture, noise scheduler, diffusion process, sampling methods for inference, optimization settings, EMA, data pipeline components, training hyperparameters, etc.
• All components, including attention modules, the model architecture, diffusion process, noise scheduler, sampling methods, and EMA, have been implemented from scratch and are thoroughly documented.

Quick Installation

1. Create a virtual environment

make venv

2. Activate it

source .venv/bin/activate

3. Install Dependencies and CUDA-Enabled PyTorch

make install-gpu
make install-dev

If you want to change the version of cuda enabled torch (currently CUDA 12.8), you can modify install-gpu section in Makefile.

4. Verify GPU support

make check-gpu

How To Use

Training

After installing the project, you can start training using the dit train command. The training CLI is fully configuration-driven and allows you to override any field in the default experiment configs located in src/diffusion_transformer/config/ via the --set section.field=value syntax. Multiple --set flags can be provided to customize the run without editing code.

As a minimal example, the following overrides some of the default values defined in src/diffusion_transformer/config/ before training a DiT model:

dit train \
--set model.num_attention_heads=6 \
--set noise_scheduler.name="cosine_scheduler" \
--set sampling.sampling_method="ddpm" \
--set dataset.name="cifar10" \
--set optimizer.weight_decay=0.0 \
--set lr_scheduler.scheduler_name='cosine' \
--set loss.loss_name="mse_kl" \
--set training.gradient_accumulation_steps=1

For all available fields that you can override, please see src/diffusion_transformer/config/.

Also, stopped training can be resumed, just pass the path to saved checkpoint, for example:

dit train \
  --set training.resume_path="./checkpoints/2026-02-07_21-53-15/last.pth"

Moreover, for practical examples, please refer to the scripts/ folder.

To monitor training metrics, such as loss curves, generated samples, and learning rate, launch TensorBoard using:

tensorboard --logdir=runs

Inference

After training, you can generate samples from a saved checkpoint using the dit infer command. Similar to training, the inference CLI is configuration-driven: it loads the experiment configuration directly from the checkpoint and allows overriding any field via the --set section.field=value syntax.

At minimum, you must provide a checkpoint path and an output directory for generated images. As a minimal example:

dit infer \
--checkpoint ./checkpoints_path/checkpoint_name.pt \
--save-dir ./samples \
--num-samples 64

You can control how class labels are used during generation:

All classes (default):

--class-mode all

Specific classes:

dit infer \
--checkpoint ./checkpoints_path/checkpoint_name.pt \
--save-dir ./samples \
--num-samples 40 \
--class-mode specific \
--class-labels 1 3 7

Unconditional generation:

--class-mode unconditional

Overriding Sampling Settings You can override inference-specific parameters without retraining:

dit infer \
--checkpoint ./checkpoints/best.pt \
--save-dir ./samples \
--num-samples 64 \
--set sampling.sampling_method="ddim" \
--set sampling.num_inference_steps=100 \
--set sampling.guidance_scale=4.0

Notes for Inference:

• The model configuration is automatically restored from the checkpoint.
• For inference, CLI should be used to overrides (--set) feild thats affect inference-time behavior.
• Lower num_inference_steps speeds up sampling but may reduce quality.
• Sampling parameters such as sampler (ddim, ddpm, etc), number of inference steps, guidance scale, etc are controlled via the sampling section of the experiment config and can be overridden using --set sampling.*. For available field names, see the config files in src/diffusion_transformer/config/sampling/.

Results

I trained Diffusion Transformer (DiT) models from scratch on CIFAR-100, CIFAR-10, German Traffic Sign Recognition Dataset (GTSRB), and MNIST. To do the same, you can run follwoing files: scripts/training-cifar-100.sh, scripts/training-cifar-10.sh, scripts/training-gtsrb.sh, and scripts/training-mnist.sh.

In the following sections, I report the FID score along with visualizations of model outputs, losses, learning rates, and related metrics.

Results are obtained without careful hyperparameter selection. All configurations and hyperparameters can be modified via the CLI. For better performance, you can adjust these values as needed.

German Traffic Sign Recognition Benchmark (GTSRB)

Generated samples

Train and Validation simple loss + VLB loss

FID and Learning Rate

CIFAR-10

Generated samples

Train and Validation simple loss + VLB loss

FID and Learning Rate

CIFAR-100

Generated samples

Train and Validation simple loss + VLB loss

FID and Learning Rate

MNIST

Generated samples

Train and Validation simple loss + VLB loss

FID and Learning Rate

License

Released under the MIT License. For used datasets, please check their respective licenses.

References

- Peebles, William, and Saining Xie. "Scalable diffusion models with transformers." In Proceedings of the IEEE/CVF international conference on computer vision, pp. 4195-4205. 2023.
- Nichol, A.Q. and Dhariwal, P., 2021, July. Improved denoising diffusion probabilistic models. In International conference on machine learning (pp. 8162-8171). PMLR.