Test-tijd schaling van diffusies met stroomafbeeldingen

Samenvatting

Een gangbare methode om diffusiemodellen tijdens het testen te verbeteren, zodat gegenereerde voorbeelden hoog scoren op een door de gebruiker gespecificeerde beloning, is het introduceren van de gradiënt van deze beloning in de dynamica van het diffusieproces zelf. Deze procedure is vaak slecht gesteld, omdat gebruikersspecifieke beloningen meestal alleen goed gedefinieerd zijn op de dataverdeling aan het einde van de generatie. Terwijl gebruikelijke oplossingen voor dit probleem een denoiser gebruiken om in te schatten wat een voorbeeld aan het einde van de generatie zou zijn geweest, stellen wij een eenvoudige oplossing voor door rechtstreeks met een stroomafbeelding te werken. Door een relatie te benutten tussen de stroomafbeelding en het snelheidsveld dat het instantane transport bepaalt, construeren we een algoritme, Flow Map Trajectory Tilting (FMTT), dat aantoonbaar een betere stijging op de beloning bereikt dan standaard testtijdmethoden die de gradiënt van de beloning gebruiken. De aanpak kan worden gebruikt om ofwel exacte steekproeven te nemen via importance weighting, ofwel voor een principieel zoekproces dat lokale maximale waarden van de beloning-gekantelde verdeling identificeert. We demonstreren de doeltreffendheid van onze aanpak tegenover andere look-ahead technieken en tonen aan hoe de stroomafbeelding interactie met complexe beloningsfuncties mogelijk maakt, wat nieuwe vormen van beeldbewerking mogelijk maakt, bijvoorbeeld door koppeling met vision-language modellen.

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.

Test-tijd schaling van diffusies met stroomafbeeldingen

Test-time scaling of diffusions with flow maps

Samenvatting

Support