On the Stability and Convergence of Flow Map Training

Abstract:Flow maps are a recent method of probabilistic generation that amortize diffusion trajectories into one-step samplers, unlocking efficient generation. While progress has been made in improving their empirical performance, there is still a lack of understanding behind the mechanisms underlying their success. Specifically, many flow map learning methods employ stop gradient operators as a crucial part of their training target. Adaptive weighting also drastically improves the generation capabilities of flow maps. Although many interpretations have been proposed to support the inclusion of these heuristics, none have analyzed the impact of these assumptions on the training dynamics of flow maps themselves. In this work, we address this issue by providing grounded theory explaining why such tricks are required. Analyzing the functional gradient flow induced by the MeanFlow and Lagrangian losses, we show their convergence to the true flow map, settling the validity of these objectives. For finite updates, we highlight a severe roadblock to the usage of these losses in practice: the possible unboundedness of the update. Adaptive normalization not only provides bounded updates, but, combined with a specific weighting, additionally provides convergence guarantees towards a neighborhood of the solution. On synthetic data, we show that our functional gradient flow analysis exactly describes the true training dynamics for flow maps. Experiments on standard image datasets demonstrate that a proper reweighting scheme is essential for strong empirical performance, which is only achieved for a specific range of reweighting strengths.

Machine Learning (stat.ML); Machine Learning (cs.LG)

arXiv:2610.20674 [stat.ML]

(or arXiv:2610.20674v1 [stat.ML] for this version)

https://doi.org/10.48550/arXiv.2610.20674 Focus to learn more arXiv-issued DOI via DataCite

arXivLabs is a framework that allows collaborators to develop and share new arXiv features directly on our website.

Both individuals and organizations that work with arXivLabs have embraced and accepted our values of openness, community, excellence, and user data privacy. arXiv is committed to these values and only works with partners that adhere to them.

Have an idea for a project that will add value for arXiv's community? Learn more about arXivLabs.