Mise à l'échelle au moment du test des diffusions avec des applications de flot

papers.abstract

Une méthode courante pour améliorer les modèles de diffusion lors des tests, afin que les échantillons obtiennent un score élevé selon une récompense spécifiée par l'utilisateur, consiste à introduire le gradient de cette récompense dans la dynamique de la diffusion elle-même. Cette procédure est souvent mal posée, car les récompenses spécifiées par l'utilisateur ne sont généralement bien définies que sur la distribution de données en fin de génération. Alors que les solutions habituelles à ce problème utilisent un débruiteur pour estimer ce qu'un échantillon aurait été en fin de génération, nous proposons une solution simple en travaillant directement avec une application de flot. En exploitant une relation entre l'application de flot et le champ de vitesse gouvernant le transport instantané, nous construisons un algorithme, Flow Map Trajectory Tilting (FMTT), qui réalise théoriquement une meilleure ascension de la récompense que les méthodes standards de test impliquant le gradient de la récompense. L'approche peut être utilisée pour effectuer un échantillonnage exact par pondération d'importance ou une recherche fondée qui identifie les maximiseurs locaux de la distribution pondérée par la récompense. Nous démontrons l'efficacité de notre approche par rapport à d'autres techniques de prévision et montrons comment l'application de flot permet d'utiliser des fonctions de récompense complexes, rendant possibles de nouvelles formes de retouche d'image, par exemple en interagissant avec des modèles de vision et langage.

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.

Mise à l'échelle au moment du test des diffusions avec des applications de flot

Test-time scaling of diffusions with flow maps

papers.abstract

Support