Post-Hoc Guidance for Consistency Models by Joint Flow Distribution Learning

TL;DR

This paper introduces JFDL, a lightweight method that enables effective post-hoc guidance for pre-trained Consistency Models, matching the guidance flexibility of Diffusion Models without requiring a separate teacher model.

Contribution

JFDL allows guidance in pre-trained CMs without distillation, unlocking guided generation and reducing FID scores on CIFAR-10 and ImageNet datasets.

Findings

JFDL provides an adjustable guidance knob for CMs.

Guided images from JFDL have similar characteristics to CFG.

JFDL reduces FID scores on CIFAR-10 and ImageNet 64x64 datasets.

Abstract

Classifier-free Guidance (CFG) lets practitioners trade-off fidelity against diversity in Diffusion Models (DMs). The practicality of CFG is however hindered by DMs sampling cost. On the other hand, Consistency Models (CMs) generate images in one or a few steps, but existing guidance methods require knowledge distillation from a separate DM teacher, limiting CFG to Consistency Distillation (CD) methods. We propose Joint Flow Distribution Learning (JFDL), a lightweight alignment method enabling guidance in a pre-trained CM. With a pre-trained CM as an ordinary differential equation (ODE) solver, we verify with normality tests that the variance-exploding noise implied by the velocity fields from unconditional and conditional distributions is Gaussian. In practice, JFDL equips CMs with the familiar adjustable guidance knob, yielding guided images with similar characteristics to CFG. Applied to an original Consistency Trained (CT) CM that could only do conditional sampling, JFDL unlocks guided generation and reduces FID on both CIFAR-10 and ImageNet 64x64 datasets. This is the first time that CMs are able to receive effective guidance post-hoc without a DM teacher, thus, bridging a key gap in current methods for CMs.