D-Flow: Differenziazione attraverso Flussi per la Generazione Controllata

D-Flow: Differentiating through Flows for Controlled Generation

February 21, 2024

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

cs.AI

Abstract

Domare i risultati generativi dei modelli all'avanguardia di Diffusion e Flow-Matching (FM) senza dover riaddestrare un modello specifico per il compito sblocca uno strumento potente per risolvere problemi inversi, generazione condizionata e generazione controllata in generale. In questo lavoro introduciamo D-Flow, un framework semplice per controllare il processo di generazione differenziando attraverso il flusso, ottimizzando il punto sorgente (rumore). Motiviamo questo framework con la nostra osservazione chiave che afferma che per i modelli Diffusion/FM addestrati con percorsi di probabilità gaussiana, differenziare attraverso il processo di generazione proietta il gradiente sulla varietà dei dati, iniettando implicitamente il prior nel processo di ottimizzazione. Validiamo il nostro framework su problemi di generazione controllata lineari e non lineari, inclusi: problemi inversi di immagini e audio e generazione condizionata di molecole, raggiungendo prestazioni all'avanguardia in tutti i casi.

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.