Align Your Flow: Scaling Continuous-Time Flow Map Distillation (2506.14603v1)

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.

• The paper introduces the AYF framework with Eulerian and Lagrangian Map Distillation to enhance multi-step generative sampling.
• The study applies autoguidance and stabilization techniques like tangent normalization and adaptive weighting to improve sample quality.
• The paper validates AYF on ImageNet and text-to-image tasks, showcasing its versatility and superior performance over traditional consistency models.

Overview of "Align Your Flow: Scaling Continuous-Time Flow Map Distillation"

The paper "Align Your Flow: Scaling Continuous-Time Flow Map Distillation" explores innovative techniques to improve the efficiency and quality of generative models, specifically focusing on diffusion and flow-based methodologies. The authors propose the Align Your Flow (AYF) framework to leverage continuous-time flow maps, thereby addressing the limitations observed in traditional consistency models, especially regarding multi-step sampling.

Key Contributions

1. Analytical Insights into Consistency Models: The authors provide analytical evidence showing that while consistency models are efficient for few-step generation, their efficacy deteriorates significantly during multi-step processes due to error accumulation. This insight challenges the prevailing assumption that consistency models are generally robust across various step counts.
2. Introduction of AYF Methodology: AYF introduces two novel objectives for training flow maps — Eulerian Map Distillation (AYF-EMD) and Lagrangian Map Distillation (AYF-LMD). These objectives generalize existing training paradigms and offer a unified framework that bridges consistency models and flow maps.
3. Autoguidance in Distillation: The paper highlights the application of autoguidance, an emerging technique used to improve model performance without adversarial training complexities. By incorporating autoguidance, AYF achieves sharper sample quality, enhancing the generative process of distilled flow maps.
4. Implementation Details and Stability Enhancements: To ensure stability and scalability, the authors introduce several training techniques, including tangent normalization, adaptive weighting, and regularized tangent warmup, making AYF a practically viable method for large-scale generative modeling.

Numerical Results and Implications

The AYF framework demonstrates significant improvements in few-step generation performance. When applied to ImageNet benchmarks, AYF achieves state-of-the-art results — a testament to its potential in practical applications. Notably, the model maintains high output quality across varying step counts, contrasting the degradation observed in traditional consistency models. This robustness positions AYF as a compelling option for scenarios where computational efficiency and sample quality are both critical, such as real-time image synthesis.

Moreover, the application of AYF in text-to-image generation, using an efficient LoRA framework, showcases its versatility in generative tasks beyond image synthesis. This demonstrates the broader applicability of AYF across different modalities, with promising implications for future developments in AI-driven content generation.

Future Directions

The paper opens several avenues for future research. The methodological advancements and empirical successes of AYF suggest potential expansions into video generation and other modalities requiring efficient generative processes. Furthermore, the insights into the limitations of consistency models inspire further exploration into refining these models to enhance their multi-step generation capabilities.

AYF's approach could also be significant in fields such as drug discovery, where rapid generation of molecular structures is essential. By enabling efficient few-step sampling, AYF may facilitate faster data-driven innovation cycles, thereby impacting areas where speed and quality of generative outputs are paramount.

In conclusion, "Align Your Flow" proposes a robust framework for generative modeling that challenges existing paradigms and offers innovative solutions to enhance the efficiency and quality of the generative process. Through theoretical insights, novel objectives, and stabilization techniques, AYF positions itself as a leading approach in the continuous-time flow map domain, with broad implications for both practical applications and theoretical advancements.