harahan/rtdmd

RTDMD is a research code release accompanying an academic paper on training fast text-to-image generation models. Standard diffusion models generate images through many small steps, often 50 or more. This work focuses on training a model that produces good images in only four steps by combining distribution-matching distillation with reward-guided reinforcement learning. The training process works in two stages. The first stage, called AC-DMD, trains a student model to match the output distribution of a larger teacher model by computing a consistency loss across intermediate generation steps. The second stage, called RTDMD, continues training using a reinforcement learning approach that rewards the student model for producing outputs that score well on metrics like image quality, alignment with the text prompt, and human preference ratings. The two objectives are applied jointly so the student model improves on both dimensions at once. The results reported in the paper show the four-step RTDMD model surpassing its teacher on most evaluated metrics for three model families: SD3-M, SD3.5-M, and FLUX.2 4B. The FLUX.2 4B model trained with RTDMD scores higher than the larger FLUX.2 9B teacher on seven of nine metrics when both are compared at their standard settings, including improvements in image reward, aesthetic score, PickScore, and GenEval. The repository provides training scripts organized into two command-line trainers, one for each stage, along with YAML configuration files for five model backbones. Inference and reward evaluation scripts are also included. Pretrained checkpoints are published on Hugging Face. The code requires Python 3.10 or later and PyTorch.