D-Flow : Différenciation à travers les flux pour une génération contrôlée

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

Maîtriser les résultats de génération des modèles de pointe en Diffusion et Flow-Matching (FM) sans avoir à réentraîner un modèle spécifique à une tâche ouvre la voie à un outil puissant pour résoudre des problèmes inverses, la génération conditionnelle et la génération contrôlée en général. Dans ce travail, nous introduisons D-Flow, un cadre simple pour contrôler le processus de génération en différenciant à travers le flux, en optimisant le point source (bruit). Nous justifions ce cadre par notre observation clé qui indique que pour les modèles de Diffusion/FM entraînés avec des chemins de probabilité gaussiens, la différenciation à travers le processus de génération projette le gradient sur la variété des données, injectant implicitement l'a priori dans le processus d'optimisation. Nous validons notre cadre sur des problèmes de génération contrôlée linéaires et non linéaires, incluant : des problèmes inverses en image et audio, ainsi que la génération conditionnelle de molécules, atteignant des performances de pointe dans tous ces domaines.

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.