Source-linked AI summary
Align Your Flow: Scaling Continuous-Time Flow Map Distillation
Amirmojtaba Sabour, Sanja Fidler, Karsten Kreis
TL;DR
Diffusion and flow-based models are costly because they require many sampling steps, while consistency models degrade when used for multi-step sampling. The paper introduces Align Your Flow, a continuous-time flow-map distillation method with new objectives, autoguidance, and adversarial finetuning. AYF achieves state-of-the-art few-step results on ImageNet and text-to-image generation with small, efficient networks.
Problem
Diffusion and flow-based models require many sampling steps, while consistency models inevitably accumulate error as the number of sampling steps increases.
Method
AYF trains flow maps with two continuous-time objectives and adds stabilization techniques, autoguided teacher distillation, and optional adversarial finetuning.
Results
AYF achieves state-of-the-art few-step generation on ImageNet 64x64 and 512x512 and outperforms existing non-adversarial few-step generators in text-to-image synthesis.
Takeaways & Limitations
Flow maps remain effective across denoising step counts, and AYF combines efficient few-step sampling with high-quality image and text-to-image generation.
Takeaways & Limitations
The paper identifies discrete-time consistency training as sensitive to timestep choices and numerical ODE discretization, and leaves applications beyond image generation for future work.
Abstract
from arXiv · showhide
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.
1 Introduction
The paper identifies a multi-step weakness in consistency models and proposes Align Your Flow, a continuous-time flow-map distillation approach with new objectives and training techniques. AYF achieves strong few-step results across image and text-to-image generation tasks.
- Motivation: Consistency models can distill diffusion and flow-based models into few-step generators, but their performance degrades as sampling steps increase.The paper attributes this degradation to error accumulation from strictly predicting clean outputs.
- Approach: Flow maps connect any two noise levels in one step, supporting both efficient few-step and flexible multi-step sampling.AYF frames this mapping as aligning the teacher flow into a few-step sampler.
- Approach: AYF introduces two continuous-time training objectives that generalize existing consistency and flow-matching objectives.The objectives correspond to consistency conditions at the beginning or end of denoising trajectories.
- Training techniques: Autoguidance distills a low-quality guidance model together with the regular teacher, while adversarial finetuning further improves quality with minimal diversity loss.The paper presents these techniques as part of scaling AYF to high performance.
- Results: AYF attains state-of-the-art few-step performance on ImageNet 64x64 and 512x512 using small, efficient neural networks.Its 4-step ImageNet sampling is reported as fast as or faster than previous single-step generation.
- Results: Text-to-image AYF models distilled from FLUX.1 significantly outperform existing non-adversarially trained few-step generators.The text-to-image experiments use an efficient LoRA framework to avoid the overhead of many previous distillation approaches.
2 Background
The background presents diffusion and flow matching as continuous generative modeling frameworks, then describes consistency models and their discrete-time training limitations. Continuous-time consistency training avoids numerical ODE-solver discretization errors and provides different supervision signals.
- Diffusion Models and Flow Matching: Diffusion models add noise through a forward process and generate data by simulating a time-reversed process from Gaussian noise.Flow matching generalizes this setup by learning continuous flows between arbitrary distributions.
- Diffusion Models and Flow Matching: Flow matching learns a continuous flow between any two distributions, without requiring the noise distribution to be Gaussian.The paper uses this framework for the remainder of its development.
- Diffusion Models and Flow Matching: In flow matching, sampling starts from noise at t = 1 and simulates the probability flow ODE backward to t = 0.The neural network predicts the velocity field used by the ODE.
- Consistency Models: Consistency models map noisy inputs directly to clean samples and enforce consistent outputs between adjacent timesteps.Their parameterization also imposes the boundary condition fθ(x, 0) = x.
- Consistency Models: Discrete-time consistency models depend on the timestep gap, require manually designed annealing schedules, and can accumulate numerical discretization errors.The preceding noisy sample is often obtained by numerically solving the probability flow ODE.
- Continuous-time Consistency Models: Continuous-time consistency models avoid numerical ODE solvers and their discretization errors, while providing better supervision signals during training.The passage also reports improved results after stabilizing and scaling continuous-time consistency models.
3 Continuous-Time Flow Map Distillation
The paper presents flow maps as a unified alternative to consistency models for multi-step generation, then develops continuous-time distillation objectives and training techniques for Align Your Flow. It analyzes consistency-model error accumulation, introduces AYF-EMD and AYF-LMD, and uses autoguidance and stabilization methods to improve flow-map training.
- Consistency Models are Flawed Multi-Step Generators: Consistency-model multi-step sampling accumulates error because each step predicts clean data before reintroducing noise, causing performance to worsen beyond a certain step count.For a Gaussian data distribution, increasing steps beyond some integer N increases Wasserstein-2 distance to the ground-truth distribution; at c = 0.5, two steps produce the best samples in the cited experiment.
- Flow Map Formulation: Flow maps map noisy inputs directly between arbitrary PF-ODE times, reducing to consistency models at s = 0 and to Euler-integrated flow or diffusion sampling over many small steps.They are designed to remain effective across step counts and can also map images to noise for inversion and editing.
- Learning Flow Maps: AYF-EMD aligns flow-map outputs as starting points move along the teacher PF-ODE and generalizes continuous-time consistency and standard flow-matching objectives.It reduces to the continuous-time consistency objective when s = 0 and to flow matching in the limit as s →t.
- Learning Flow Maps: AYF-LMD instead aligns the trajectory ending at a fixed input point, but it is more stable in 2D experiments while producing overly smooth, lower-quality image samples.The paper therefore reports a stability-quality trade-off between the two objectives on different experimental settings.
- Sharpening the Distribution with Autoguidance: Autoguidance distills a teacher enhanced by a weaker model, sharpening the student, while stabilization techniques address continuous-time training instability.The paper stabilizes time embeddings and uses additional techniques; c = 0.1 is reported as critical for stable training.
4 Related Work
Related work spans consistency models, flow-map formulations, and diffusion-distillation methods. The paper positions AYF as a high-performance continuous-time flow-map distillation approach that uses autoguidance rather than relying exclusively on adversarial objectives.
- Consistency Models: Consistency-model research improved training stability and performance through new objectives, weighting functions, variance reduction, and truncated or continuous-time formulations.These methods generally target efficient few-step generation from flow-based models.
- Flow Map Models: Prior flow-map methods include CTMs, trajectory consistency distillation, bidirectional consistency models, multistep and phased consistency models, flow-map matching, and shortcut models.The cited methods connect timestep pairs or distributions, while the paper characterizes AYF as the first high-performance continuous-time flow-map model.
- Accelerating Diffusion Models: Diffusion acceleration includes knowledge, progressive, adversarial, variational score, and operator-learning distillation, alongside other few-step techniques.The paper notes that many methods use adversarial losses, which introduce GAN-like training complexity and can limit diversity.
5 Experiments
Experiments show AYF delivers efficient few-step image and text-to-image generation, with adversarial finetuning improving quality while preserving diversity.
- AYF achieves the best sample quality among non-adversarial few-step methods on ImageNet 64×64 with two sampling steps.This comes at the cost of slightly weaker one-step quality.
- 18% of the larger models’ compute is sufficient for AYF to match the best sCD model at two steps on ImageNet 512×512.Four-step AYF further improves quality while remaining over twice as fast as the large one-step sCM model.
- Adversarial finetuning significantly boosts AYF performance, especially for one-step generation, with minimal impact on recall-based sample diversity.The finetuned model achieves state-of-the-art few-step generation on ImageNet 64×64.
- AYF distills FLUX.1 into a four-step text-to-image generator using LoRA, receiving clear user preference over LCM and TCD.The comparison uses FLUX.1 [dev] for AYF and SDXL for LCM and TCD.
6 Conclusions
AYF is a continuous-time flow-map distillation method that generalizes flow-matching and consistency-based models. It maintains effectiveness across denoising step counts and achieves strong image and text-to-image results with efficient models.
- AYF trains flow maps with a continuous-time distillation method that generalizes flow-matching and consistency-based models.
- Autoguidance improves the teacher and distilled student, while adversarial finetuning adds quality with minimal impact on sample diversity.
- AYF reaches state-of-the-art non-adversarial distillation performance on ImageNet64, ImageNet512, and text-to-image generation.
- AYF’s small neural networks make two-step or four-step sampling as fast or faster than previous single-step generators.
- Future work could apply AYF to video distillation and domains such as efficient molecule or protein modeling.
B Limitations
AYF stabilizes multi-step performance but has trade-offs: slightly weaker one-step quality, a remaining teacher gap, and dependence on a pretrained flow-based teacher.
- AYF’s multi-step stability comes with slightly degraded one-step performance relative to methods focused solely on one-step generation.
- Adversarial finetuning can improve AYF across sampling counts while causing minimal loss in diversity.
- A small performance gap remains between AYF and its multi-step teacher regardless of the sampling-step count.The paper attributes this to AYF’s similarly sized model solving a more challenging task while preserving the same noise-to-data mapping.
- The method assumes access to a pretrained flow-based teacher, although prior work suggests direct consistency training can outperform distillation.
- On ImageNet 512×512, AYF maintains low FID across step counts, while consistency models deteriorate after only a few steps.
C.1 Theorem: Consistency Models are Flawed Multi-step Generators
The theorem analyzes consistency-model multistep sampling in a Gaussian setting and shows that repeated transitions accumulate error, eventually worsening distributional accuracy.
- In the Gaussian setting, the optimal consistency model and flow-matching velocity admit closed-form solutions used in the analysis.
- Beyond some integer N, adding sampling steps increases the Wasserstein-2 distance between generated samples and the ground-truth distribution.
- A consistency-model transition first predicts clean data and then re-noises that estimate to the next timestep.
- The resulting variance follows a recurrence across a uniform n-step schedule, enabling analysis of multistep sampling error.
- As n →∞, the sampled variance Var(0) diverges, so multistep consistency sampling accumulates errors beyond a certain point.
- The intuition is that imperfect denoising drifts off the data manifold, and re-noising compounds the resulting mismatch over successive steps.
C.3 Derivation: Tangent Warmup as Linearity Regularization
Equation (5) establishes that tangent warmup is equivalent to adding a flow-map regularizer that encourages linearity.
- Equation (5) shows equivalence between tangent warmup and a linearity-promoting regularization term on flow maps.
- The derived gradient contains the flow map, teacher velocity, and stop-gradient flow-map terms.
- The derivation uses three simplifications: dropping the nonnegative interval factor, identifying Fθ with Fθ−, and stopping gradients through vϕ.
D Connections to Existing Methods
AYF connects continuous-time flow-map distillation to flow matching, consistency models, trajectory distillation, related flow-map losses, and MeanFlow. Its stop-gradient design improves scalability and training efficiency, while AYF-EMD was the variant that scaled effectively to large image datasets.
- Overview: AYF extends prior approaches with new objectives, autoguided distillation, adversarial finetuning, theoretical analysis, and scaled text-to-image and ImageNet models.
- Flow Matching: AYF-EMD reduces to standard flow matching as s →t, up to a constant factor.
- Consistency Models: AYF-EMD reduces to the continuous-time consistency-model objective when s = 0, up to a constant factor.
- Consistency Trajectory Models: The TCD objective is a discrete approximation of AYF-EMD because both use the same discrete consistency loss, with TCD fixing the discretization schedule.
- Flow Map Matching: Removing stop-gradient recovers Boffi et al.’s EMD and LMD losses, whereas retaining it avoids backpropagation through the JVP and can speed training.
- Training Dynamics: Only AYF-EMD scaled effectively to large-scale image datasets and produced high-quality outputs in the experiments.
- MeanFlow: Under an Euler flow-map parameterization, MeanFlow’s objective matches AYF-EMD up to a constant, while MeanFlow trains from scratch and AYF focuses on distillation.
F.1 ImageNet Experiments
The ImageNet experiments evaluate AYF flow-map distillation with autoguided teachers, interval scheduling, stabilization techniques, and optional adversarial finetuning. AYF achieves strong few-step performance with small networks, including performance beyond existing baselines using four steps.
- Teacher Models: AYF distills publicly available EDM2 checkpoints after flow-matching alignment, using S and XS teacher models trained on 2147 million and 134 million images.
- Training Setup: Training uses tangent normalization, tangent warmup, w(t, s) = 1/|t−s|^2, learning rate 10^-4, batch size 2048, and 50,000 total steps.
- Timestep Scheduling: The interval schedule samples |t −s| from a transformed normal distribution, prioritizing medium-length intervals for stability and training only t > s pairs.
- Timestep Scheduling: For ImageNet-512 and ImageNet-64, effective schedule parameters are (Pmean, Pstd) = (−0.8, 1.0) and (−0.6, 1.6), respectively.
- Results: AYF models outperform all listed ImageNet-512 methods using only 4 sampling steps.
- Adversarial Finetuning: Adversarial finetuning combines AYF-EMD with a weighted RpGAN objective and generates negative samples through one-step flow-map sampling.
- Text-to-Image: Text-to-image AYF-LoRA distills FLUX.1-[dev] on 100K generated images selected from the text-to-image-2M dataset.
G Ablation Studies
Ablations show that AYF’s design choices matter across sampling regimes: AYF-EMD is stronger on images, autoguidance outperforms CFG, and AYF maintains an advantage over consistency and shortcut baselines.
- Objective Choice: AYF-EMD significantly improves image-generation quality over AYF-LMD, reversing the preference observed on the 2D toy distribution.
- Guidance: Replacing autoguidance with CFG consistently degrades performance across all sampling steps.
- Comparison with sCD: AYF stabilizes after 8 steps and remains better than sCD at 4 steps, while short adversarial finetuning boosts performance across all sampling steps.
- Comparison with sCD: With CFG instead of autoguidance, AYF still outperforms sCD for every setting except single-step generation.
- Comparison with Shortcut: AYF consistently outperforms Shortcut models, which improve with increasing NFEs but have substantially worse few-step performance.
H Additional Samples
Additional samples show AYF generating sharp, detailed text-to-image outputs and maintaining nearly unchanged quality as sampling steps increase. ImageNet examples cover one- and two-step generation across diverse classes and resolutions.
- AYF produces sharper, more detailed text-to-image samples with better prompt adherence than prior LoRA-based consistency models.
- Four-step AYF samples are qualitatively compared with LCM and TCD text-to-image models.
- Increasing sampling steps produces only tiny visible quality differences for the ImageNet-512 AYF model.
- The appendix includes selected one- and two-step ImageNet-512 samples across classes including goldfish, bald eagle, axolotl, macaw, wolves, landscapes, birds, and flowers.
- Additional uncurated one- and two-step samples are shown for the ImageNet-64 AYF-S model with randomly chosen class labels.