Align Your Flow: Scaling Continuous-Time Flow Map Distillation
Authors: name: Amirmojtaba Sabour, name: Sanja Fidler, name: Karsten Kreis

Diffusion- and flow-based models have emerged as state-of-the-art generative
modeling approaches, but they require many sampling steps. Consistency models
can distill these models into efficient one-step generators; however, unlike
flow- and diffusion-based methods, their performance inevitably degrades when
increasing the number of steps, which we show both analytically and
empirically. Flow maps generalize these approaches by connecting any two noise
levels in a single step and remain effective across all step counts. In this
paper, we introduce two new continuous-time objectives for training flow maps,
along with additional novel training techniques, generalizing existing
consistency and flow matching objectives. We further demonstrate that
autoguidance can improve performance, using a low-quality model for guidance
during distillation, and an additional boost can be achieved by adversarial
finetuning, with minimal loss in sample diversity. We extensively validate our
flow map models, called Align Your Flow, on challenging image generation
benchmarks and achieve state-of-the-art few-step generation performance on both
ImageNet 64x64 and 512x512, using small and efficient neural networks. Finally,
we show text-to-image flow map models that outperform all existing
non-adversarially trained few-step samplers in text-conditioned synthesis.