Ridimensionamento al tempo di test delle diffusioni con mappe di flusso

Test-time scaling of diffusions with flow maps

November 27, 2025

Autori: Amirmojtaba Sabour, Michael S. Albergo, Carles Domingo-Enrich, Nicholas M. Boffi, Sanja Fidler, Karsten Kreis, Eric Vanden-Eijnden

cs.AI

Abstract

Una ricetta comune per migliorare i modelli di diffusione durante il test, in modo che i campioni ottengano punteggi elevati rispetto a una ricompensa specificata dall'utente, consiste nell'introdurre il gradiente della ricompensa nelle dinamiche della diffusione stessa. Questa procedura è spesso mal posta, poiché le ricompense specificate dall'utente sono solitamente ben definite solo sulla distribuzione dei dati alla fine della generazione. Mentre le soluzioni alternative comuni a questo problema prevedono l'uso di un denoiser per stimare quale sarebbe stato un campione alla fine della generazione, noi proponiamo una soluzione semplice a questo problema lavorando direttamente con una mappa di flusso. Sfruttando una relazione tra la mappa di flusso e il campo di velocità che governa il trasporto istantaneo, costruiamo un algoritmo, Flow Map Trajectory Tilting (FMTT), che dimostrabilmente esegue una risalita sulla ricompensa migliore rispetto ai metodi standard di test che coinvolgono il gradiente della ricompensa. L'approccio può essere utilizzato per eseguire un campionamento esatto tramite pesatura d'importanza o una ricerca principiata che identifica i massimizzatori locali della distribuzione distort dalla ricompensa. Dimostriamo l'efficacia del nostro approccio rispetto ad altre tecniche di look-ahead e mostriamo come la mappa di flusso consenta di interagire con funzioni di ricompensa complesse che rendono possibili nuove forme di editing di immagini, ad esempio interfacciandosi con modelli linguistici visivi.

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.