Stable Consistency Tuning: Understanding and Improving Consistency Models

TL;DR

This paper introduces a new theoretical framework for consistency models using Markov Decision Processes and Temporal Difference Learning, leading to a novel Stable Consistency Tuning method that improves performance and achieves state-of-the-art results.

Contribution

It models consistency training as value estimation in an MDP, analyzes current limitations, and proposes SCT with variance reduction for better generative performance.

Findings

SCT significantly improves benchmark scores.

Achieves new state-of-the-art FID on ImageNet-64.

Provides a theoretical understanding of consistency models.

Abstract

Diffusion models achieve superior generation quality but suffer from slow generation speed due to the iterative nature of denoising. In contrast, consistency models, a new generative family, achieve competitive performance with significantly faster sampling. These models are trained either through consistency distillation, which leverages pretrained diffusion models, or consistency training/tuning directly from raw data. In this work, we propose a novel framework for understanding consistency models by modeling the denoising process of the diffusion model as a Markov Decision Process (MDP) and framing consistency model training as the value estimation through Temporal Difference~(TD) Learning. More importantly, this framework allows us to analyze the limitations of current consistency training/tuning strategies. Built upon Easy Consistency Tuning (ECT), we propose Stable Consistency Tuning (SCT), which incorporates variance-reduced learning using the score identity. SCT leads to significant performance improvements on benchmarks such as CIFAR-10 and ImageNet-64. On ImageNet-64, SCT achieves 1-step FID 2.42 and 2-step FID 1.55, a new SoTA for consistency models.