D-Flow: Differenzierung durch Flüsse für kontrollierte Generierung

D-Flow: Differentiating through Flows for Controlled Generation

February 21, 2024

papers.authors: Heli Ben-Hamu, Omri Puny, Itai Gat, Brian Karrer, Uriel Singer, Yaron Lipman

cs.AI

papers.abstract

Die Steuerung der Generierungsergebnisse modernster Diffusions- und Flow-Matching (FM)-Modelle, ohne ein aufgabenspezifisches Modell neu trainieren zu müssen, eröffnet ein leistungsstarkes Werkzeug zur Lösung inverser Probleme, bedingter Generierung und kontrollierter Generierung im Allgemeinen. In dieser Arbeit stellen wir D-Flow vor, ein einfaches Framework zur Steuerung des Generierungsprozesses durch Differenzierung entlang des Flusses, wobei der Ausgangspunkt (Rauschen) optimiert wird. Wir begründen diesen Ansatz mit unserer zentralen Beobachtung, dass bei Diffusion/FM-Modellen, die mit Gaußschen Wahrscheinlichkeitspfaden trainiert wurden, die Differenzierung durch den Generierungsprozess den Gradienten auf die Datenmannigfaltigkeit projiziert und somit implizit das Prior in den Optimierungsprozess einbringt. Wir validieren unser Framework an linearen und nichtlinearen Problemen der kontrollierten Generierung, einschließlich inverser Probleme bei Bildern und Audio sowie der bedingten Molekülgenerierung, und erreichen dabei state-of-the-art Leistungen in allen Bereichen.

English

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.