Align Your Flow: Scaling Continuous-Time Flow Map Distillation

TL;DR

This paper introduces flow map models that generalize diffusion and consistency models, enabling efficient one- or few-step image generation with improved performance and flexibility, validated on challenging benchmarks.

Contribution

The paper proposes new continuous-time flow map training objectives and techniques, enhancing efficiency and quality in generative modeling beyond existing methods.

Findings

Achieves state-of-the-art few-step generation on ImageNet benchmarks.

Demonstrates effective text-to-image synthesis outperforming existing non-adversarial methods.

Shows that flow maps maintain effectiveness across different step counts.

Abstract

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.