Adaptive Discretization for Consistency Models

TL;DR

This work tackles the sensitivity of Consistency Models to discretization by proposing Adaptive Discretization for Consistency Models (ADCMs), a unified framework that optimizes the discretization step to balance trainability and stability. It formulates a constrained optimization with local consistency as the objective and global consistency as a constraint, relaxed via a Lagrange multiplier and solved analytically with the Gauss-Newton method to yield an adaptive Δt. The approach yields significant improvements in training efficiency and one-step generation quality on CIFAR-10 and ImageNet 64×64, and demonstrates robustness across DM variants such as Flow Matching. The method is practical, with adaptive weighting and distance metrics enhancing stability, and scales to high-resolution data, making ADCMs a versatile tool for efficient CM-based generation.

Abstract

Consistency Models (CMs) have shown promise for efficient one-step generation. However, most existing CMs rely on manually designed discretization schemes, which can cause repeated adjustments for different noise schedules and datasets. To address this, we propose a unified framework for the automatic and adaptive discretization of CMs, formulating it as an optimization problem with respect to the discretization step. Concretely, during the consistency training process, we propose using local consistency as the optimization objective to ensure trainability by avoiding excessive discretization, and taking global consistency as a constraint to ensure stability by controlling the denoising error in the training target. We establish the trade-off between local and global consistency with a Lagrange multiplier. Building on this framework, we achieve adaptive discretization for CMs using the Gauss-Newton method. We refer to our approach as ADCMs. Experiments demonstrate that ADCMs significantly improve the training efficiency of CMs, achieving superior generative performance with minimal training overhead on both CIFAR-10 and ImageNet. Moreover, ADCMs exhibit strong adaptability to more advanced DM variants. Code is available at https://github.com/rainstonee/ADCM.

Figures (7)

• Figure 1: ADCMs significantly improve the training efficiency on both (a) unconditional CIFAR-10 and (b) class-conditional ImageNet $64 \times 64$. ADCMs achieve superior generation quality with only a minimal amount of training data. $*$ indicates a smaller model.
• Figure 2: Discretization strategies of different models. CMs consider only local consistency during discretization, while DMs consider only global consistency. ADCMs balance local and global consistency and adaptively adjust the discretization step size based on the information from both.
• Figure 3: (a) Training time cost of ADCMs. (b) Training dynamics of different discretization methods. Compared to other CMs' approaches, ADCMs have a faster convergence rate and better final performance. (c) Training dynamics for different $\lambda$. A large $\lambda$ improves stability but hurts final performance, while a too-small $\lambda$ reduces stability and hinders convergence.
• Figure 4: Adaptive Discretization on (a) EDM (VE) and (b) Flow Matching.
• Figure 5: 1-step samples from unconditional CIFAR-10 trained with a budget of $76.8$M images.

Theorems & Definitions (2)

• Remark 3.1
• Remark 3.2