Escalado en tiempo de prueba de difusiones con mapas de flujo

Resumen

Una receta común para mejorar los modelos de difusión en tiempo de prueba, de modo que las muestras obtengan puntuaciones altas según una recompensa especificada por el usuario, es introducir el gradiente de la recompensa en la dinámica de la propia difusión. Este procedimiento suele estar mal planteado, ya que las recompensas especificadas por el usuario generalmente solo están bien definidas en la distribución de datos al final de la generación. Si bien las soluciones habituales a este problema consisten en utilizar un desruidosor para estimar cómo habría sido una muestra al final de la generación, nosotros proponemos una solución simple trabajando directamente con un mapa de flujo. Al explotar una relación entre el mapa de flujo y el campo de velocidad que gobierna el transporte instantáneo, construimos un algoritmo, Inclinación de Trayectoria por Mapa de Flujo (FMTT, por sus siglas en inglés), que demuestrablemente realiza un ascenso mejor sobre la recompensa que los métodos estándar de tiempo de prueba que involucran el gradiente de la recompensa. El enfoque puede utilizarse para realizar un muestreo exacto mediante ponderación por importancia o una búsqueda fundamentada que identifique maximizadores locales de la distribución sesgada por la recompensa. Demostramos la eficacia de nuestro enfoque frente a otras técnicas de anticipación y mostramos cómo el mapa de flujo permite interactuar con funciones de recompensa complejas que posibilitan nuevas formas de edición de imágenes, por ejemplo, mediante la interfaz con modelos de lenguaje 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.

Escalado en tiempo de prueba de difusiones con mapas de flujo

Test-time scaling of diffusions with flow maps

Resumen

Support