arXiv preprint ·
On the Stability and Convergence of Flow Map Training
- Published
- Source
- arXiv preprint, posted before peer review
- Original
- Read it at arXiv
- Topics
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.
The authors’ abstract, as arXiv lists it. arXiv’s listings, abstracts included, are free for anyone to reuse.
Read it on ReadSearch
Each opens after you create an account, and brings you back to this paper.
More in Machine Learning (Statistics)
Subdata Selection for Measurement-Constrained Linear Regression under Model Misspecification
Efficient Agnostic Learning of Univariate Densities Under Local Differential Privacy
Large-Step Last-Iterate Convergence of Random-Reshuffling Stochastic Gradient Descent
Score Mismatch and Set-Constrained Sampling in Zero-Shot Conditional Discrete Diffusion
State-Matched Sequence Conformal Prediction
Get the few papers worth your time in Machine Learning (Statistics) every morning.
ReadSearch reads the new papers in your field each day and brings you the few worth reading, each with its key finding in a sentence.