Progressive Seed Pruningとは？拡散モデルの推論時スケーリングを高効率化

• 多数のノイズシードをデノイズ初期段階で評価し、有望でない候補を段階的に枝刈りすることで、固定の計算予算のまま画像生成の品質を高める手法です。
• SDXLのGenEvalスコアはBest-of-Nの0.629に対しPSPが0.645を達成し、importance sampling系のFK-Steeringも上回りました。
• 249名による8,295枚の人手評価でもプロンプト整合性でBest-of-Nを上回り、既存の拡散モデルに追加学習なしで適用できる汎用性を示しました。

研究の背景

大規模言語モデル（Large Language Model、LLM）の世界では、推論時により多くの計算を投じるほど回答の質が上がる「テスト時計算スケーリング」が定着しています。ところが画像生成を担う拡散モデル（Diffusion Model）やFlow Matchingには、同じように計算を積み増して品質を伸ばす確立した方法がまだ乏しいのが実情でした。

拡散モデルの出力品質は、生成の出発点となる初期ノイズ、いわゆるシード（乱数の種）の選び方に大きく左右されます。そのため、複数のシードを試して良いものを選ぶ「シード探索」が有力な方向として注目されてきました。カリフォルニア工科大学のRogerio GuimaraesとPietro Peronaによる本研究は、このシード探索をより計算効率よく行う枠組みを提案します。

従来の代表的な手法には、N個のシードを最後まで生成して最良の1枚を選ぶBest-of-N（BoN）、粒子数を一定に保ちながら途中でリサンプリングするimportance sampling系のFK-Steering、そして木探索（tree-search）などがあります。これらはいずれも、推論中のメモリ使用量を一定に保つ制約のもとで動いていました。

提案手法PSPの仕組み

Progressive Seed Pruning（PSP、段階的シード枝刈り）の発想は、推論中のメモリを一定に保つという前提をあえて緩めることにあります。まず通常より多くのシードから生成を始め、デノイズが進む初期の段階で各候補の中間的な品質を評価します。見込みの薄い候補を早めに切り捨て、残った有望な軌道だけに残りの計算を集中させる流れです。

具体的なスケジュールは明快です。デノイズを64ステップで行うSD1.5やSDXLの場合、8個のシードから開始し、25%地点にあたる16ステップ目で候補を4個に、50%地点の32ステップ目で2個へと半減させ、最後まで残った2候補を完走させます。32ステップで動くSD3.5でも同じ割合で枝刈りします。この配分によって、他手法と同じ総ステップ数のまま2倍のシードを初期に評価できます。

枝刈りの判断には、途中経過の画像を採点する報酬モデルを使います。ImageRewardやHPS v2といった既存の評価モデルに、その時点のノイズから一段階で予測したクリーン画像（x̂0）を渡してスコアを得る仕組みです。このx̂0はサンプラーの内部で通常計算される値を再利用するため、追加コストがほとんどかかりません。論文では、ImageRewardの中間スコアが最終品質と比較的早い段階から高い相関を示すと報告されています。

計算予算は「計算倍率」N̄で表します。これは投じた総デノイズステップ数を1軌道分の長さで割った値で、N̄が1のとき標準の1回生成に相当します。N̄が4のとき、PSPは4軌道分の予算で2N̄すなわち8個の初期シードを扱う計算になります。この柔軟さが、限られた予算を有望な軌道へ寄せる余地を生みます。

GenEvalと人手評価

実験では、拡散モデルのSD1.5とSDXL、Flow MatchingベースのSD3.5という3つのバックボーンで評価しました。テキスト指示への忠実さを測る自動評価指標GenEvalでは、計算倍率N̄が4のもとでPSPが多くの設定でBoNとFK-Steeringを上回っています。

SD1.5とSDXLでは、標準生成からの伸びも枝刈りによる上積みもはっきり出ています。SD3.5はもともとの品質が高く、PSPとBoNが同点に並ぶ結果でした。生成モデルの実力が上がるほど改善の余地は縮む傾向がうかがえます。

人手評価はさらに大規模に実施されました。249名の評価者が8,295枚の画像を採点し、各画像を3名がプロンプトとの整合性で判定した結果（計24,885件、評価者間の一致率80.3%）、PSPはすべてのモデルでBoNとFK-Steeringを上回りました。たとえばSDXLではPSPが0.713、BoNが0.682、FK-Steeringが0.676という差がつきました。自動指標だけでなく、人の目で見ても改善が確認できた点は説得力があります。

計算倍率を増やすほどPSPの品質は伸び続け、同じ計算量ではBoNを一貫して上回りました。小さいモデルにN̄を大きく与えることで、より大きなモデルに近い品質へ届く場合があることも示されています。Mage-Flowのような高速なFlow Matching基盤モデルと組み合わせれば、軽量モデルでも大型モデル並みの結果を狙える余地が広がりそうです。

限界と今後の展望

PSPにも適用しにくい場面があります。枝刈りは単一のスコアで候補を順位付けする前提のため、複数の目的を同時に満たしたい場合や、そもそも数値化しにくい品質を扱う場合には表現しづらくなります。報酬設計そのものが手法の効き方を左右する点は留意が必要です。

効果はタスクにも依存します。生成の後半で決まる細かな描写、たとえば美的な仕上がりが品質を左右するような目的では、初期段階でのシード選択が効きにくく、改善幅は小さくなります。実際、SDXLでHPSを指標にした場合の伸びは比較的控えめでした。一方で、本番環境ではメモリを一定に保つ制約が必ずしも厳しくないため、候補数を柔軟に変える発想は今後の推論時スケーリングの設計指針として活きる可能性があります。

コードはGitHubで公開されており、既存の拡散モデルやFlow Matchingモデルに追加学習なしで組み込める点が実用面での強みです。同じ計算予算でより良い画像を得たい場面で、有力な選択肢になりそうです。

Inference-Time Scaling of Diffusion Models via Progressive Seed Pruning

Abstract page for arXiv paper 2607.21591: Inference-Time Scaling of Diffusion Models via Progressive Seed Pruning

arxiv.org