Elucidating the Preconditioning in Consistency Distillation

TL;DR

This paper provides theoretical insights into preconditioning in consistency distillation for diffusion models, proposing an analytic method to optimize it, leading to faster training and better trajectory alignment.

Contribution

It introduces a theoretical framework for preconditioning in consistency distillation and proposes Analytic-Precond, an optimization method that improves training efficiency and model performance.

Findings

Achieves 2x to 3x training acceleration.

Enhances alignment of student and teacher trajectories.

Facilitates learning of trajectory jumpers.

Abstract

Consistency distillation is a prevalent way for accelerating diffusion models adopted in consistency (trajectory) models, in which a student model is trained to traverse backward on the probability flow (PF) ordinary differential equation (ODE) trajectory determined by the teacher model. Preconditioning is a vital technique for stabilizing consistency distillation, by linear combining the input data and the network output with pre-defined coefficients as the consistency function. It imposes the boundary condition of consistency functions without restricting the form and expressiveness of the neural network. However, previous preconditionings are hand-crafted and may be suboptimal choices. In this work, we offer the first theoretical insights into the preconditioning in consistency distillation, by elucidating its design criteria and the connection to the teacher ODE trajectory. Based on these analyses, we further propose a principled way dubbed \textit{Analytic-Precond} to analytically optimize the preconditioning according to the consistency gap (defined as the gap between the teacher denoiser and the optimal student denoiser) on a generalized teacher ODE. We demonstrate that Analytic-Precond can facilitate the learning of trajectory jumpers, enhance the alignment of the student trajectory with the teacher's, and achieve $2\times$ to $3\times$ training acceleration of consistency trajectory models in multi-step generation across various datasets.