拡散モデルのフローマップを用いたテスト時スケーリング

要旨

拡散モデルをテスト時に改善し、ユーザー指定の報酬に対してサンプルが高く評価されるようにする一般的な手法は、拡散過程のダイナミクスに報酬の勾配を導入することである。しかし、ユーザー指定の報酬は通常、生成の最終段階におけるデータ分布上でのみ適切に定義されるため、この手順はしばしば不適切な問題設定となる。この問題に対する一般的な回避策は、デノイザを使用してサンプルが生成終了時点でどのような状態であったかを推定することであるが、本研究では、フローマップを直接扱うという単純な解決法を提案する。フローマップと瞬間的な輸送を支配する速度場の関係を利用することで、Flow Map Trajectory Tilting (FMTT) というアルゴリズムを構築する。このアルゴリズムは、報酬の勾配を含む標準的なテスト時手法よりも、理論的に報酬に対するより良い上昇を行うことを証明する。このアプローチは、重要度重み付けによる正確なサンプリング、または報酬によって傾けられた分布の局所最大化子を特定する原理的な探索のいずれかに利用できる。我々は、他の先読み技術と比較して本手法の有効性を実証し、フローマップが複雑な報酬関数との連携を可能にすることで、例えば視覚言語モデルとのインターフェースを通じて、新しい形式の画像編集を実現する方法を示す。

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-time scaling of diffusions with flow maps

要旨

Support