Масштабирование диффузий во время тестирования с помощью карт потока

Test-time scaling of diffusions with flow maps

November 27, 2025

Авторы: Amirmojtaba Sabour, Michael S. Albergo, Carles Domingo-Enrich, Nicholas M. Boffi, Sanja Fidler, Karsten Kreis, Eric Vanden-Eijnden

cs.AI

Аннотация

Распространенным методом улучшения диффузионных моделей во время тестирования для получения сэмплов с высокими показателями заданной пользователем функции вознаграждения является введение градиента этой функции в динамику самой диффузии. Однако данная процедура часто оказывается некорректно поставленной, поскольку пользовательские функции вознаграждения обычно определены лишь на распределении данных в конце генерации. В то время как стандартные обходные пути этой проблемы используют денойзер для оценки того, каким был бы сэмпл в конце генерации, мы предлагаем простое решение путем непосредственной работы с потоковым отображением. Используя взаимосвязь между потоковым отображением и полем скорости, управляющим мгновенным переносом, мы конструируем алгоритм FMTT (Flow Map Trajectory Tilting – Наклон Траектории Потокового Отображения), который теоретически обеспечивает лучший подъем по функции вознаграждения, чем стандартные методы тестирования, использующие градиент вознаграждения. Данный подход может быть использован для выполнения точной выборки посредством важностного взвешивания или для принципиального поиска, идентифицирующего локальные максимизаторы распределения, наклоненного по вознаграждению. Мы демонстрируем эффективность нашего подхода в сравнении с другими методами упреждения и показываем, как потоковое отображение позволяет работать со сложными функциями вознаграждения, что открывает новые возможности редактирования изображений, например, путем взаимодействия с языковыми моделями компьютерного зрения.

English

A common recipe to improve diffusion models at test-time so that samples score highly against a user-specified reward is to introduce the gradient of the reward into the dynamics of the diffusion itself. This procedure is often ill posed, as user-specified rewards are usually only well defined on the data distribution at the end of generation. While common workarounds to this problem are to use a denoiser to estimate what a sample would have been at the end of generation, we propose a simple solution to this problem by working directly with a flow map. By exploiting a relationship between the flow map and velocity field governing the instantaneous transport, we construct an algorithm, Flow Map Trajectory Tilting (FMTT), which provably performs better ascent on the reward than standard test-time methods involving the gradient of the reward. The approach can be used to either perform exact sampling via importance weighting or principled search that identifies local maximizers of the reward-tilted distribution. We demonstrate the efficacy of our approach against other look-ahead techniques, and show how the flow map enables engagement with complicated reward functions that make possible new forms of image editing, e.g. by interfacing with vision language models.