D-Flow: Diferenciação através de Fluxos para Geração 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

Resumo

Domar o resultado da geração de modelos de última geração de Difusão e Flow-Matching (FM) sem a necessidade de re-treinar um modelo específico para a tarefa libera uma ferramenta poderosa para resolver problemas inversos, geração condicional e geração controlada em geral. Neste trabalho, introduzimos o D-Flow, uma estrutura simples para controlar o processo de geração ao diferenciar através do fluxo, otimizando o ponto de origem (ruído). Motivamos essa estrutura por nossa observação chave, que afirma que, para modelos de Difusão/FM treinados com caminhos de probabilidade gaussiana, diferenciar através do processo de geração projeta o gradiente na variedade de dados, injetando implicitamente o prior no processo de otimização. Validamos nossa estrutura em problemas de geração controlada linear e não linear, incluindo: problemas inversos de imagem e áudio e geração condicional de moléculas, alcançando desempenho de última geração em todos.

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.