[論文レビュー] Stable Target Field for Reduced Variance Score Estimation in Diffusion Models

Yilun Xu, Shangyuan Tong+1 more|arXiv (Cornell University)|Feb 1, 2023

Advanced Neuroimaging Techniques and Applications被引用数 5

ひとこと要約

本稿では、自己正規化重要度サンプリングを用いて参照バッチから重み付き条件付きスコアを計算することで、拡散モデルにおけるスコア推定の分散を低減するための安定的ターゲットフィールド（STF）目的関数を提案する。この手法により、特に中間ノイズ領域における訓練ターゲットの分散が低減され、CIFAR-10でEDMと組み合わせた場合、FIDが1.90に向上し、3.6倍高速化された訓練が達成された。

ABSTRACT

Diffusion models generate samples by reversing a fixed forward diffusion process. Despite already providing impressive empirical results, these diffusion models algorithms can be further improved by reducing the variance of the training targets in their denoising score-matching objective. We argue that the source of such variance lies in the handling of intermediate noise-variance scales, where multiple modes in the data affect the direction of reverse paths. We propose to remedy the problem by incorporating a reference batch which we use to calculate weighted conditional scores as more stable training targets. We show that the procedure indeed helps in the challenging intermediate regime by reducing (the trace of) the covariance of training targets. The new stable targets can be seen as trading bias for reduced variance, where the bias vanishes with increasing reference batch size. Empirically, we show that the new objective improves the image quality, stability, and training speed of various popular diffusion models across datasets with both general ODE and SDE solvers. When used in combination with EDM, our method yields a current SOTA FID of 1.90 with 35 network evaluations on the unconditional CIFAR-10 generation task. The code is available at https://github.com/Newbeeer/stf

研究の動機と目的

• 拡散モデルにおけるノイズ除去スコアマッチング（DSM）の訓練ターゲットに生じる高い分散を是正すること、特に中間ノイズ領域における分散に焦点を当てる。
• 複数のデータモードが同じノイズを含むサンプルに影響を与える場合に生じるデータソースの曖昧性が分散の原因であることを特定する。
• 参照バッチと重要度サンプリングを用いて訓練ターゲットを安定化させる一般化されたスコアマッチング目的関数を提案する。
• 参照バッチサイズが増加するにつれて漸近的に不偏性を保ちつつ、訓練ターゲットの共分散のトレースを低減することを示す。
• 複数の拡散モデルアーキテクチャおよびデータセットにおいて、訓練効率、モデル安定性、生成性能の向上を実現する。

提案手法

• 前方拡散プロセスを三段階に分解する：近接領域（低ノイズ）、中間領域（高分散）、遠方領域（低信号）。
• 安定的ターゲットフィールド（STF）目的関数を、大規模な参照バッチを用いて重み付き条件付きスコアを計算する一般化DSM目的関数として定義する。
• 自己正規化重要度サンプリングを適用し、参照サンプルからの寄与を集約することで、正確な正規化を必要とせずにターゲット分散を低減する。
• データポイントの参照バッチを用いて条件付きスコアフィールドを推定し、各参照点が摂動分布下での尤度に比例して寄与するようにする。
• STF目的関数は漸近的に不偏であり、訓練ターゲットの共分散のトレースを、参照バッチサイズに依存する要因で低減する。
• 本手法はODEおよびSDEソルバと両立可能であり、EDM、VE、VP SDEsなどの既存の拡散モデルに統合可能である。

Figure 5: FID scores in the training course with varying reference batch size.

実験結果

リサーチクエスチョン

• RQ1 拡散モデルにおけるノイズ除去スコアマッチングの訓練ターゲットに高い分散が生じる原因は何ですか？
• RQ2 前方拡散プロセスのどの領域でターゲット分散が最も顕著になりますか？
• RQ3 データポイントの参照バッチを重要度サンプリングを用いてスコア推定を安定化させるために利用できるでしょうか？
• RQ4 参照バッチサイズが増加するに従って、STF目的関数のバイアス・バリアンストレードオフはどのように変化しますか？
• RQ5 STFは、異なる拡散モデルアーキテクチャにおいて、訓練速度、FID、モデル安定性をどの程度向上させますか？

主な発見

• STF目的関数は、分散が最も高い中間ノイズ領域における訓練ターゲットの共分散のトレースを低減する。
• EDMと組み合わせた場合、CIFAR-10における非条件付き生成で、1.90という新たなSOTA FIDスコアを達成し、ネットワーク評価をたった35回に抑えた。
• VEおよびVP SDEsを含む複数の拡散モデルにおいて、FIDおよびインセプションスコアが向上し、CIFAR-10およびCelebA 64²で一貫した向上が得られた。
• 本手法は、異なるランダムシードにおいてVEモデルの画像品質の劣化やノイズの多い生成を低減し、訓練安定性を向上させた。
• CIFAR-10のVEモデルでは、FIDを維持または向上させながら、最大3.6倍の訓練速度向上を達成した。これは、拡散モデルの訓練を高速化する最初の手法である。
• 参照バッチサイズが増加するに従って、STFによるバイアスは漸近的に消滅し、本手法の一貫性が裏付けられた。

Figure 6: FID versus NFE using ODE samplers on CIFAR-10 (left) and CelebA $64\times 64$ (right)