High Quality Diffusion Distillation on a Single GPU with Relative and Absolute Position Matching

TL;DR

This work presents RAPM, a diffusion distillation method that enables high-quality, text-to-image generation on a single GPU by aligning a student model with a teacher via relative and absolute diffusion position matching. By precomputing teacher trajectories using $M$-step DDIM over coarse time-slots and employing two discriminators for relative and absolute position losses, RAPM achieves strong distillation performance with batchsize-1 and as few as $4$ timesteps. The method demonstrates competitive FID scores on SD V1.5 and SDXL, outperforming several multi-GPU baselines while substantially reducing resource requirements. The approach introduces carefully designed losses and detach operations to maintain stability and effective trajectory alignment, advancing practical diffusion distillation for researchers with limited computational resources.

Abstract

We introduce relative and absolute position matching (RAPM), a diffusion distillation method resulting in high quality generation that can be trained efficiently on a single GPU. Recent diffusion distillation research has achieved excellent results for high-resolution text-to-image generation with methods such as phased consistency models (PCM) and improved distribution matching distillation (DMD2). However, these methods generally require many GPUs (e.g.~8-64) and significant batchsizes (e.g.~128-2048) during training, resulting in memory and compute requirements that are beyond the resources of some researchers. RAPM provides effective single-GPU diffusion distillation training with a batchsize of 1. The new method attempts to mimic the sampling trajectories of the teacher model by matching the relative and absolute positions. The design of relative positions is inspired by PCM. Two discriminators are introduced accordingly in RAPM, one for matching relative positions and the other for absolute positions. Experimental results on StableDiffusion (SD) V1.5 and SDXL indicate that RAPM with 4 timesteps produces comparable FID scores as the best method with 1 timestep under very limited computational resources.

Figures (5)

• Figure 1: 1024$\times$1024 samples produced by our 4-step generator distilled from SDXL. See Table \ref{['tab:prompts_RAPM_images']} for the text prompts.
• Figure 2: Demonstration of the functional losses being constructed in RAPM. Two sources of losses are introduced in RAPM, one for matching the relative diffusion position $\hat{\boldsymbol{z}}_n$ and the other one for matching the absolute diffusion position $\tilde{\boldsymbol{z}}_n$ .
• Figure 3: Qualitive comparison of RAPM with different loss functions over certain iterations when performing distillation over SDXL. The images on the top row are obtained by only matching the relative diffusion positions via the loss function (\ref{['equ:RAPM_S6']}). On the other hand, the images on the bottom are generated by matching both the relative and absolute positions via (\ref{['equ:loss_all']}). The number of timesteps was set to 4. See Table \ref{['tab:prompts_IPCM_loss']} for the associated text prompts.
• Figure 4: A simplified schematic diagram of our method from Fig. \ref{['fig:RAPM_demo']}. The detach operation is performed to obtain $\boldsymbol{z}_{n+1}$ before computing $\boldsymbol{z}_n^{\boldsymbol{\varphi}}$ to be able to make the position matching process decoupled across different time-slots.
• Figure 5: FID comparison of PCM and RAPM for distilation of SD V1.5 with 4 timesteps.