Variational Flow Maps: Make Some Noise for One-Step Conditional Generation
Flow maps enable high-quality image generation in a single forward pass. However, unlike iterative diffusion models, their lack of an explicit sampling trajectory impedes incorporating external constraints for conditional generation and solving inverse problems. We put forth Variational Flow Maps, a framework for conditional sampling that shifts the perspective of conditioning from "guiding a sampling path", to that of "learning the proper initial noise". Specifically, given an observation, we seek to learn a noise adapter model that outputs a noise distribution, so that after mapping to the data space via flow map, the samples respect the observation and data prior. To this end, we develop a principled variational objective that jointly trains the noise adapter and the flow map, improving noise-data alignment, such that sampling from complex data posterior is achieved with a simple adapter. Experiments on various inverse problems show that VFMs produce well-calibrated conditional samples in a single (or few) steps. For ImageNet, VFM attains competitive fidelity while accelerating the sampling by orders of magnitude compared to alternative iterative diffusion/flow models. Code is available at https://github.com/abbasmammadov/VFM
69b00e6d9dbab2a62ec280b6
Abbas Mammadov
69b00e6d9dbab2a62ec280b7
So Takao
69b00e6d9dbab2a62ec280b8
Bohan Chen
69b00e6d9dbab2a62ec280b9
Ricardo Baptista
69b00e6d9dbab2a62ec280ba
Morteza Mardani
69b00e6d9dbab2a62ec280bb
Yee Whye Teh
69b00e6d9dbab2a62ec280bc
Julius Berner
Variational Flow Maps: Make Some Noise for One-Step Conditional Generation
Flow maps enable high-quality image generation in a single forward pass. However, unlike iterative diffusion models, their lack of an explicit sampling trajectory impedes incorporating external constraints for conditional generation and solving inverse problems. We put forth Variational Flow Maps, a framework for conditional sampling that shifts the perspective of conditioning from "guiding a sampling path", to that of "learning the proper initial noise". Specifically, given an observation, we seek to learn a noise adapter model that outputs a noise distribution, so that after mapping to the data space via flow map, the samples respect the observation and data prior. To this end, we develop a principled variational objective that jointly trains the noise adapter and the flow map, improving noise-data alignment, such that sampling from complex data posterior is achieved with a simple adapter. Experiments on various inverse problems show that VFMs produce well-calibrated conditional samples in a single (or few) steps. For ImageNet, VFM attains competitive fidelity while accelerating the sampling by orders of magnitude compared to alternative iterative diffusion/flow models. Code is available at https://github.com/abbasmammadov/VFM
Variational Flow Maps introduce a novel framework for conditional image generation that learns optimal initial noise distributions to achieve fast, calibrated sampling while maintaining high fidelity.
flow maps conditional sampling variational objective noise adapter model data prior inverse problems diffusion models image generation sampling trajectory noise-data alignment