Test-Time-Skalierung von Diffusionen mit Flussabbildungen
Test-time scaling of diffusions with flow maps
November 27, 2025
papers.authors: Amirmojtaba Sabour, Michael S. Albergo, Carles Domingo-Enrich, Nicholas M. Boffi, Sanja Fidler, Karsten Kreis, Eric Vanden-Eijnden
cs.AI
papers.abstract
Ein gängiges Verfahren zur Verbesserung von Diffusionsmodellen zur Testzeit, damit Stichproben hohe Werte bezüglich einer benutzerdefinierten Belohnungsfunktion erzielen, besteht darin, den Gradienten der Belohnung in die Dynamik der Diffusion selbst einzuführen. Dieses Verfahren ist oft schlecht gestellt, da benutzerdefinierte Belohnungsfunktionen üblicherweise nur auf der Datenverteilung am Ende der Generierung wohldefiniert sind. Während gängige Lösungsansätze für dieses Problem einen Denoiser verwenden, um abzuschätzen, wie eine Stichprobe am Ende der Generierung ausgesehen hätte, schlagen wir eine einfache Lösung vor, indem wir direkt mit einer Flussabbildung arbeiten. Durch Ausnutzung einer Beziehung zwischen der Flussabbildung und dem Geschwindigkeitsfeld, das den momentanen Transport beschreibt, konstruieren wir einen Algorithmus – Flow Map Trajectory Tilting (FMTT) –, der nachweislich einen besseren Aufstieg bezüglich der Belohnung erreicht als Standard-Testzeit-Verfahren, die den Gradienten der Belohnung einbeziehen. Der Ansatz kann verwendet werden, um entweder exakte Stichprobenziehung mittels Importance-Weighting oder eine prinzipielle Suche durchzuführen, die lokale Maximierer der belohnungsgekippten Verteilung identifiziert. Wir demonstrieren die Wirksamkeit unseres Ansatzes im Vergleich zu anderen Look-Ahead-Techniken und zeigen, wie die Flussabbildung die Nutzung komplexer Belohnungsfunktionen ermöglicht, die neue Formen der Bildbearbeitung realisierbar machen, beispielsweise durch die Schnittstelle zu Vision-Language-Models.
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.