Variational Flow Maps: Make Some Noise for One-Step Conditional Generation

Abbas Mammadov, So Takao, Bohan Chen, Ricardo Baptista, Morteza Mardani, Yee Whye Teh, Julius Berner

TL;DR: We propose a method for one-step conditional generation by tilting the noise space.

Abstract: 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.

Lay Summary: In contrast to iterative diffusion models, one/few-step generative models such as flow maps cannot naturally solve inverse problem (e.g. image de-blurring, inpainting, etc...) due to the lack of a sampling trajectory, which one can use to steer towards high-likelihood regions. To address this problem, we shifted the perspective of conditioning from "guiding a sampling path" to "finding the initial noise" that leads to high-likelihood samples. This helps to solve a wide variety of inverse problems, as well as more general reward alignment problems, while retaining the blazing fast inference speed of one/few-step flow maps.

Link To Code: https://github.com/abbasmammadov/VFM

Primary Area: Deep Learning->Generative Models and Autoencoders

Keywords: Generative Models, Flow Maps, Inverse Problems, Conditional Generation

Submission Number: 17728