拡散モデルの逐次誤差を学習ベースの補正器で抑える

拡散モデルは、ランダムノイズから画像を段階的に復元する。学習時、デノイザーが受け取るのは通常、実画像にノイズを加えた入力である。しかし推論時には、次のステップの入力が前のステップでモデル自身が出した予測に依存する。つまり、学習時にはほとんど見ていない種類の入力を処理することになる。初期段階の小さな誤りが後続ステップへ伝われば、最終画像の品質低下につながる可能性がある。

「Correction Training for Diffusion Models」は、この学習と推論の差を補うため、基盤デノイザーに補正器を追加する。手法の流れは次の通りだ。

• 基盤デノイザーが、現在の時刻におけるクリーンサンプルを予測する。
• 補正器がその予測を入力として受け取る。
• 基盤モデルと同じ正解のクリーンサンプルを目標に、補正器を学習する。
• サンプリング時には中間予測を補正し、その処理を補正器自身の出力に対して繰り返すこともできる。

この構成のポイントは、予測をどう修正するかを独立した学習課題として扱うことにある。単にサンプラーの設定を変えたり、外部のガイダンスを強めたりするのではなく、基盤デノイザーが実際に出す予測を材料にして、目標との差を縮める。中間段階で結果を修正できれば、早い段階で生じたアーティファクトが最後まで残ることも抑えられると考えられる。

論文ではDiT、SiT、JiTを用いて評価し、無誘導生成のFIDについて、従来報告された結果を上回る改善を示したと報告している。特に、分類器フリーガイダンスを使わない無条件生成では、高品質な画像を得ることが難しい。ガイダンスは品質向上に役立つ一方、生成の多様性や彩度に影響を与える場合がある。提案手法は、ガイダンスを強くする代わりにモデル自身の予測を修正し、品質を高めながら多様性を大きく損なわないことを狙う。

補正器を別モジュールとして追加する点も実用上の特徴だ。基盤デノイザーを全面的に置き換えず、反復サンプリングで現れる誤差を補正器に担当させる。この発想は、より大きなモデルや強いガイダンスだけに頼らず、モデルが自分の中間出力を評価し、修正する方向性を示している。

提示された素材だけでは、すべてのデータセット、解像度、サンプリング回数での挙動までは判断できない。今後は適用範囲や追加計算量を含めた検証が必要だが、反復生成における誤差蓄積を直接扱う、比較的シンプルな改善策として注目される。

出典：Hugging Face Daily Papers