D-Flow: Differentiating through Flows for Controlled Generation

Heli Ben-Hamu, Omri Puny, Itai Gat, Brian Karrer, Uriel Singer, Yaron Lipman
Proceedings of the 41st International Conference on Machine Learning, PMLR 235:3462-3483, 2024.

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.

Cite this Paper

BibTeX
@InProceedings{pmlr-v235-ben-hamu24a, title = {D-Flow: Differentiating through Flows for Controlled Generation}, author = {Ben-Hamu, Heli and Puny, Omri and Gat, Itai and Karrer, Brian and Singer, Uriel and Lipman, Yaron}, booktitle = {Proceedings of the 41st International Conference on Machine Learning}, pages = {3462--3483}, year = {2024}, editor = {Salakhutdinov, Ruslan and Kolter, Zico and Heller, Katherine and Weller, Adrian and Oliver, Nuria and Scarlett, Jonathan and Berkenkamp, Felix}, volume = {235}, series = {Proceedings of Machine Learning Research}, month = {21--27 Jul}, publisher = {PMLR}, pdf = {https://raw.githubusercontent.com/mlresearch/v235/main/assets/ben-hamu24a/ben-hamu24a.pdf}, url = {https://proceedings.mlr.press/v235/ben-hamu24a.html}, 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.} }
Endnote
%0 Conference Paper %T D-Flow: Differentiating through Flows for Controlled Generation %A Heli Ben-Hamu %A Omri Puny %A Itai Gat %A Brian Karrer %A Uriel Singer %A Yaron Lipman %B Proceedings of the 41st International Conference on Machine Learning %C Proceedings of Machine Learning Research %D 2024 %E Ruslan Salakhutdinov %E Zico Kolter %E Katherine Heller %E Adrian Weller %E Nuria Oliver %E Jonathan Scarlett %E Felix Berkenkamp %F pmlr-v235-ben-hamu24a %I PMLR %P 3462--3483 %U https://proceedings.mlr.press/v235/ben-hamu24a.html %V 235 %X 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.
APA
Ben-Hamu, H., Puny, O., Gat, I., Karrer, B., Singer, U. & Lipman, Y.. (2024). D-Flow: Differentiating through Flows for Controlled Generation. Proceedings of the 41st International Conference on Machine Learning, in Proceedings of Machine Learning Research 235:3462-3483 Available from https://proceedings.mlr.press/v235/ben-hamu24a.html.

Related Material