Adaptive Discretization for Consistency Models

TL;DR

This paper introduces ADCMs, an adaptive discretization framework for consistency models that optimizes the discretization step during training, leading to improved efficiency and performance in generative tasks on CIFAR-10 and ImageNet.

Contribution

It presents a unified, automatic discretization method for consistency models using an optimization framework and Gauss-Newton method, enhancing training efficiency and adaptability.

Findings

Significantly improves training efficiency of CMs

Achieves superior generative performance on CIFAR-10 and ImageNet

Demonstrates strong adaptability to advanced diffusion model variants

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.