D-Flow: Дифференцирование через потоки для управляемой генерации
D-Flow: Differentiating through Flows for Controlled Generation
February 21, 2024
Авторы: Heli Ben-Hamu, Omri Puny, Itai Gat, Brian Karrer, Uriel Singer, Yaron Lipman
cs.AI
Аннотация
Управление результатом генерации современных моделей диффузии и согласования потоков (Flow-Matching, FM) без необходимости повторного обучения модели для конкретной задачи открывает мощный инструмент для решения обратных задач, условной генерации и контролируемой генерации в целом. В данной работе мы представляем D-Flow — простую структуру для управления процессом генерации путем дифференцирования через поток с оптимизацией исходной (шумовой) точки. Мы обосновываем этот подход нашим ключевым наблюдением, которое гласит, что для моделей диффузии/FM, обученных с гауссовыми вероятностными траекториями, дифференцирование через процесс генерации проецирует градиент на многообразие данных, неявно внося априорную информацию в процесс оптимизации. Мы проверяем нашу структуру на линейных и нелинейных задачах контролируемой генерации, включая обратные задачи для изображений и аудио, а также условную генерацию молекул, достигая наилучших результатов во всех случаях.
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.