Escalonamento no momento do teste de difusões com mapas de fluxo
Test-time scaling of diffusions with flow maps
November 27, 2025
Autores: Amirmojtaba Sabour, Michael S. Albergo, Carles Domingo-Enrich, Nicholas M. Boffi, Sanja Fidler, Karsten Kreis, Eric Vanden-Eijnden
cs.AI
Resumo
Uma receita comum para melhorar modelos de difusão durante o teste, de modo que as amostras obtenham pontuações elevadas em relação a uma recompensa especificada pelo usuário, é introduzir o gradiente da recompensa na dinâmica da própria difusão. Este procedimento é frequentemente mal formulado, uma vez que recompensas especificadas pelo usuário geralmente só são bem definidas na distribuição de dados no final da geração. Embora soluções alternativas comuns para este problema envolvam o uso de um desruidor para estimar como uma amostra teria sido no final da geração, propomos uma solução simples para este problema trabalhando diretamente com um mapa de fluxo. Explorando uma relação entre o mapa de fluxo e o campo de velocidade que governa o transporte instantâneo, construímos um algoritmo, *Flow Map Trajectory Tilting* (FMTT), que comprovadamente realiza uma ascensão melhor na recompensa do que os métodos padrão de teste que envolvem o gradiente da recompensa. A abordagem pode ser usada para realizar amostragem exata via ponderação por importância ou uma busca fundamentada que identifica maximizadores locais da distribuição inclinada pela recompensa. Demonstramos a eficácia da nossa abordagem contra outras técnicas de antecipação e mostramos como o mapa de fluxo permite o engajamento com funções de recompensa complexas que possibilitam novas formas de edição de imagem, por exemplo, através da interface com modelos de linguagem visual.
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.