Align Your Tangent: Training Better Consistency Models via
  Manifold-Aligned Tangents
With diffusion and flow matching models achieving state-of-the-art generating
performance, the interest of the community now turned to reducing the inference
time without sacrificing sample quality. Consistency Models (CMs), which are
trained to be consistent on diffusion or probability flow ordinary differential
equation (PF-ODE) trajectories, enable one or two-step flow or diffusion
sampling. However, CMs typically require prolonged training with large batch
sizes to obtain competitive sample quality. In this paper, we examine the
training dynamics of CMs near convergence and discover that CM tangents -- CM
output update directions -- are quite oscillatory, in the sense that they move
parallel to the data manifold, not towards the manifold. To mitigate
oscillatory tangents, we propose a new loss function, called the manifold
feature distance (MFD), which provides manifold-aligned tangents that point
toward the data manifold. Consequently, our method -- dubbed Align Your Tangent
(AYT) -- can accelerate CM training by orders of magnitude and even out-perform
the learned perceptual image patch similarity metric (LPIPS). Furthermore, we
find that our loss enables training with extremely small batch sizes without
compromising sample quality. Code: https://github.com/1202kbs/AYT
68e3605973e20ab5778420cf
Beomsu Kim
68e3605973e20ab5778420d0
Byunghee Cha
68e3605973e20ab5778420d1
Jong Chul Ye
Align Your Tangent: Training Better Consistency Models via
  Manifold-Aligned Tangents
With diffusion and flow matching models achieving state-of-the-art generating
performance, the interest of the community now turned to reducing the inference
time without sacrificing sample quality. Consistency Models (CMs), which are
trained to be consistent on diffusion or probability flow ordinary differential
equation (PF-ODE) trajectories, enable one or two-step flow or diffusion
sampling. However, CMs typically require prolonged training with large batch
sizes to obtain competitive sample quality. In this paper, we examine the
training dynamics of CMs near convergence and discover that CM tangents -- CM
output update directions -- are quite oscillatory, in the sense that they move
parallel to the data manifold, not towards the manifold. To mitigate
oscillatory tangents, we propose a new loss function, called the manifold
feature distance (MFD), which provides manifold-aligned tangents that point
toward the data manifold. Consequently, our method -- dubbed Align Your Tangent
(AYT) -- can accelerate CM training by orders of magnitude and even out-perform
the learned perceptual image patch similarity metric (LPIPS). Furthermore, we
find that our loss enables training with extremely small batch sizes without
compromising sample quality. Code: https://github.com/1202kbs/AYT
Align Your Tangent (AYT) improves Consistency Model training by reducing oscillatory tangents and enabling faster convergence with small batch sizes.
diffusion models flow matching models Consistency Models CMs diffusion sampling probability flow ordinary differential equation PF-ODE training dynamics CM tangents manifold feature distance MFD Align Your Tangent AYT learned perceptual image patch similarity metric LPIPS