Unified Continuous Generative Models

TL;DR

This work introduces UCGM, a unified framework that bridges diffusion, flow-matching, and consistency models through a single training objective and a universal sampler. By parameterizing a consistency ratio $\lambda$ and employing time transformations, enhanced target scores, self-boosting, and extrapolative sampling, UCGM achieves state-of-the-art or competitive results across both few-step and multi-step generation on ImageNet and CIFAR-10, while reducing sampling costs. Empirical results demonstrate robust performance gains when applying UCGM-T and UCGM-S to pre-trained models and during joint training, with significant reductions in NFEs and improved FID scores at high resolutions. The approach promises practical impact by enabling efficient, cross-paradigm, high-fidelity generative modeling suitable for latent-space diffusion transformers and related architectures.

Abstract

Recent advances in continuous generative models, including multi-step approaches like diffusion and flow-matching (typically requiring 8-1000 sampling steps) and few-step methods such as consistency models (typically 1-8 steps), have demonstrated impressive generative performance. However, existing work often treats these approaches as distinct paradigms, resulting in separate training and sampling methodologies. We introduce a unified framework for training, sampling, and analyzing these models. Our implementation, the Unified Continuous Generative Models Trainer and Sampler (UCGM-{T,S}), achieves state-of-the-art (SOTA) performance. For example, on ImageNet 256x256 using a 675M diffusion transformer, UCGM-T trains a multi-step model achieving 1.30 FID in 20 steps and a few-step model reaching 1.42 FID in just 2 steps. Additionally, applying UCGM-S to a pre-trained model (previously 1.26 FID at 250 steps) improves performance to 1.06 FID in only 40 steps. Code is available at: https://github.com/LINs-lab/UCGM.

Figures (8)

• Figure 1: Generated samples from two $675\mathrm{M}$ diffusion transformers trained with our UCGM on ImageNet-1K $512\!\times\!512$. The figure showcases generated samples illustrating the flexibility of Number of Function Evaluation (NFE) and superior performance achieved by our UCGM. The left subfigure presents results with NFE $=40$ (multi-step), while the right subfigure shows results with NFE $=2$ (few-step). Note that the samples are sampled without classifier-free guidance (CFG) or other guidance techniques.
• Figure 2: Ablation studies of UCGM on ImageNet-1K $256\!\times\!256$. These studies evaluate key factors of the proposed UCGM. Ablations presented in (a) and (c) utilize XL/1 models with the VA-VAE autoencoder. For the results shown in (b), B/2 models with the SD-VAE autoencoder are used to facilitate more efficient training.
• Figure 3: Case studies of UCGM on three synthetic datasets. These intuitive studies evaluate the ability of our UCGM to capture the latent data structure for both few-step generation ($\lambda=1$) and multi-step generation ($\lambda=0$) tasks.
• Figure 4: Intermediate images generated during $60$-step sampling from UCGM-S. Columns display intermediate images $\hat{\mathbf{x}}_t$ produced at different timesteps $t$ during a single sampling trajectory, ordered from left to right by decreasing $t$. Rows correspond to models trained with $\lambda \in \{0.0, 0.5, 1.0\}$, ordered from top to bottom. Note that the initial noise for generating these images is the same.
• Figure 5: Visualization of generated images ($512\times512$) from pre-trained EDM2-S karras2024analyzing.

Theorems & Definitions (65)

• Lemma 1: Gradient of a Squared Norm
• proof
• Lemma 2: Stop-Gradient Simplification
• proof
• Lemma 3: Finite-Difference Definition
• proof
• Theorem 1: Gradient Approximation via Finite Difference
• proof
• Lemma 4
• proof