D-Flow: Differentiating through Flows for Controlled Generation
Authors: name: Heli Ben-Hamu; name: Omri Puny; name: Itai Gat; name: Brian Karrer; name: Uriel Singer; name: Yaron Lipman
Published: 2024-02-21 18:56:03

Abstract: Taming the generation outcome of state of the art Diffusion and Flow-Matching
(FM) models without having to re-train a task-specific model unlocks a powerful
tool for solving inverse problems, conditional generation, and controlled
generation in general. In this work we introduce D-Flow, a simple framework for
controlling the generation process by differentiating through the flow,
optimizing for the source (noise) point. We motivate this framework by our key
observation stating that for Diffusion/FM models trained with Gaussian
probability paths, differentiating through the generation process projects
gradient on the data manifold, implicitly injecting the prior into the
optimization process. We validate our framework on linear and non-linear
controlled generation problems including: image and audio inverse problems and
conditional molecule generation reaching state of the art performance across
all.