D-Flow: Differentiëren door middel van stromen voor gecontroleerde generatie

D-Flow: Differentiating through Flows for Controlled Generation

February 21, 2024

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

cs.AI

Samenvatting

Het beheersen van het generatieresultaat van state-of-the-art Diffusion en Flow-Matching (FM) modellen zonder een taakspecifiek model opnieuw te hoeven trainen, biedt een krachtig hulpmiddel voor het oplossen van inverse problemen, conditionele generatie en gecontroleerde generatie in het algemeen. In dit werk introduceren we D-Flow, een eenvoudig raamwerk voor het beheersen van het generatieproces door differentiatie door de flow, waarbij het bronpunt (ruis) wordt geoptimaliseerd. We motiveren dit raamwerk door onze belangrijkste observatie, die stelt dat voor Diffusion/FM modellen die zijn getraind met Gaussische waarschijnlijkheidspaden, differentiatie door het generatieproces de gradiënt projecteert op de datamanifold, waardoor impliciet de prior in het optimalisatieproces wordt geïnjecteerd. We valideren ons raamwerk op lineaire en niet-lineaire gecontroleerde generatieproblemen, waaronder: inverse problemen bij beeld en audio en conditionele molecuulgeneratie, waarbij we state-of-the-art prestaties behalen op alle gebieden.

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.