Correcting Flows with Marginal Matching
Keywords: flow matching, rectified flow, diffusion, generative models
TL;DR: Marginal Matching improves flow matching models by steering generated samples towards the true data distribution, enhancing sample quality with minimal computational cost.
Abstract: Flow matching models, ODE-based generative models, generate samples by gradually morphing a simple source distribution into a target distribution.
In practice, these models still fall short of perfectly replicating the target distribution, mainly due to imperfections of the learned mapping. Previous work mainly focus on alleviating discretization error, which rises from sampling a continuous trajectory with a finite number of steps. In this work we focus on prediction error, an error that is inherent in the model.
Our main contribution is identifying a trajectory that complies with the imperfect flow model and leads exactly to the target distribution. Based on this finding, we propose Marginal Matching—a simple inference-time correction scheme to steer the generated samples in the direction of the data.
This scheme proves to reduce a bound on the distance between the data and the learned distribution, motivating two different implementations for the correction function. We show that our proposed method improves sample quality on CIFAR-10 and ImageNet-64, with minimal overhead in computation time, or non at all when applying approximated correction.
Primary Area: generative models
Code Of Ethics: I acknowledge that I and all co-authors of this work have read and commit to adhering to the ICLR Code of Ethics.
Submission Guidelines: I certify that this submission complies with the submission instructions as described on https://iclr.cc/Conferences/2025/AuthorGuide.
Anonymous Url: I certify that there is no URL (e.g., github page) that could be used to find authors’ identity.
No Acknowledgement Section: I certify that there is no acknowledgement section in this submission for double blind review.
Submission Number: 4926
Loading