D-Flow: Diferenciación a través de flujos para generación controlada

D-Flow: Differentiating through Flows for Controlled Generation

February 21, 2024

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

cs.AI

Resumen

Domar los resultados de generación de los modelos de última generación de Difusión y Emparejamiento de Flujos (FM) sin necesidad de volver a entrenar un modelo específico para cada tarea desbloquea una herramienta poderosa para resolver problemas inversos, generación condicional y generación controlada en general. En este trabajo presentamos D-Flow, un marco simple para controlar el proceso de generación diferenciando a través del flujo, optimizando el punto de origen (ruido). Motivamos este marco con nuestra observación clave que establece que, para los modelos de Difusión/FM entrenados con trayectorias de probabilidad gaussianas, diferenciar a través del proceso de generación proyecta el gradiente en la variedad de datos, inyectando implícitamente el previo en el proceso de optimización. Validamos nuestro marco en problemas de generación controlada lineales y no lineales, incluyendo: problemas inversos de imágenes y audio, y generación condicional de moléculas, alcanzando un rendimiento de última generación en todos ellos.

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.