A General Framework for Inference-time Scaling and Steering of Diffusion Models

要約

拡散モデルは、画像やビデオからタンパク質の設計やテキストに至るまで、さまざまなモダリティで印象的な結果を生み出します。
ただし、ユーザー指定のプロパティを持つサンプルを生成することは依然として課題です。
最近の研究では、望ましい特性を捕捉して報酬を最大化するためにモデルを微調整することが提案されていますが、これらの方法は高価なトレーニングが必要であり、モード崩壊が起こりやすいです。
この研究では、報酬関数を備えた拡散モデルをステアリングするための推論時間フレームワークである、ファインマン Kac (FK) ステアリングを提案します。
FK ステアリングは、粒子と呼ばれる複数の相互作用する拡散プロセスのシステムをサンプリングし、ポテンシャルと呼ばれる関数を使用して計算されたスコアに基づいて中間ステップで粒子を再サンプリングすることによって機能します。
ポテンシャルは中間状態の報酬を使用して定義され、高い値が粒子が高い報酬のサンプルを生成することを示すように選択されます。
可能性、中間報酬、サンプラーのさまざまな選択肢を検討します。
テキストから画像へのモデルとテキスト拡散モデルで FK ステアリングを評価します。
人間の好みの報酬を使用してテキストから画像へのモデルをステアリングする場合、0.8B パラメーター モデルをステアリングする FK は、より高速なサンプリングとトレーニングなしで、プロンプト忠実度において 2.6B パラメーターで微調整されたモデルよりも優れていることがわかります。
テキストの品質と特定のテキスト属性に対する報酬を備えたテキスト拡散モデルのステアリングの場合、FK ステアリングは混乱度が低く、言語的に許容しやすい出力を生成し、毒性などの属性を勾配なしで制御できることがわかりました。
私たちの結果は、拡散モデルの推論時間のスケーリングとステアリングが、既製の報酬であっても、サンプル品質の大幅な向上と制御性の利点を提供できることを示しています。
コードは https://github.com/zacharyhorvitz/Fk-Diffusion-Steering で入手できます。

要約(オリジナル)

Diffusion models produce impressive results in modalities ranging from images and video to protein design and text. However, generating samples with user-specified properties remains a challenge. Recent research proposes fine-tuning models to maximize rewards that capture desired properties, but these methods require expensive training and are prone to mode collapse. In this work, we propose Feynman Kac (FK) steering, an inference-time framework for steering diffusion models with reward functions. FK steering works by sampling a system of multiple interacting diffusion processes, called particles, and resampling particles at intermediate steps based on scores computed using functions called potentials. Potentials are defined using rewards for intermediate states and are selected such that a high value indicates that the particle will yield a high-reward sample. We explore various choices of potentials, intermediate rewards, and samplers. We evaluate FK steering on text-to-image and text diffusion models. For steering text-to-image models with a human preference reward, we find that FK steering a 0.8B parameter model outperforms a 2.6B parameter fine-tuned model on prompt fidelity, with faster sampling and no training. For steering text diffusion models with rewards for text quality and specific text attributes, we find that FK steering generates lower perplexity, more linguistically acceptable outputs and enables gradient-free control of attributes like toxicity. Our results demonstrate that inference-time scaling and steering of diffusion models, even with off-the-shelf rewards, can provide significant sample quality gains and controllability benefits. Code is available at https://github.com/zacharyhorvitz/Fk-Diffusion-Steering .

arxiv情報

提供元, 利用サービス

arxiv.jp, Google