OYS、拡散モデルのノイズ除去タイムステップを直接探索し、5ステップで50ステップ時の品質の89〜94%を維持

拡散モデルの推論コストは通常、ノイズ除去ネットワークのforward pass回数で測定される。実際のエンジニアリングではデフォルトのタイムステップをそのまま短縮することが多いが、少ステップ時に最高の品質を維持できるとは限らない。コーネル大学などの研究者はOptimize Your Sampling（OYS）を提案した。OYSはタイムステップやguidance strengthなどのサンプリングパラメータをブラックボックスとして扱い、目標指標を改善できる設定をベイズ最適化によって直接探索する。完全なサンプリングチェーンを通したbackpropagationは不要で、スコア関数が微分可能である必要もないため、HPS、LPIPS、PSNR、または特定タスク向けの非微分スコアを最適化できる。

これは、拡散SDEの離散化におけるKL上界を最適化するAlign Your Steps（AYS）の手法とは異なる。OYSは各イテレーションで実際に画像のバッチを生成し、surrogate modelとacquisition functionを用いて次のタイムステップ候補を選択する。コストは高いものの、最終的に実際に使用する品質指標を直接最適化できる。この探索コストが発生するのはモデルとタスクごとのチューニング時だけであり、その後は大量のオンライン推論によって償却できる。

COCO Captionsを用いた5ステップのテストでは、DeepFloyd、Stable Diffusion 1.5、SDXLにおけるOYSのHPSは、それぞれ0.252、0.238、0.245だった。同じステップ数のデフォルトスケジュールでは、それぞれ0.215、0.230、0.229だった。SDXLでは、OYSの画像がデフォルトスケジュールに対して70.70%のHPS勝率を記録した。もともと少ステップ生成向けに蒸留されたSDXL-Turboでも、3ステップのOYSによってHPSが0.287から0.298へ向上した。FLUX.1-devとQwen Imageの5ステップ時の勝率は、それぞれ82.69%と77.19%だった。58人が参加したSDXLの実験でも、忠実度とテキスト整合性の両面でOYSが好まれた。

研究では、5ステップのスケジュールによって50ステップ時の品質の89〜94%を維持できると要約している。これは、成熟したモデルで重みを再学習することなく、ネットワーク評価回数を約10分の1に削減できることを意味する。ただし、エンドツーエンドのサービス全体が10倍高速化するわけではない。テキストエンコーディング、デコーディング、バッチスケジューリング、データ転送には依然として固定コストがかかる。また、最適なスケジュールはモデル、プロンプト分布、sampler、評価指標に大きく依存する。たとえば、FLUX.1-devのチューニングには5,376枚の画像生成が必要だった。デプロイ担当者は今後、スケジュールがLoRA、量子化バージョン、実際の中国語プロンプトに対してどの程度汎化できるか、また著者が再現可能な探索コードとスケジュールファイルを公開するかを注視すべきだ。