Improved Distribution Matching Distillation for Fast Image Synthesis

要約

最近のアプローチでは、拡散モデルを効率的なワンステップジェネレーターに抽出できる可能性が示されています。
その中で、Distribution Matching Distillation (DMD) は、教師のサンプリング軌跡と 1 対 1 の対応を強制することなく、教師の分布に一致するワンステップ ジェネレーターを生成します。
ただし、安定したトレーニングを保証するために、DMD では、教師が決定論的サンプラーの多くのステップで生成したノイズ画像ペアの大規模なセットを使用して計算される追加の回帰損失が必要です。
これは、大規模なテキストから画像への合成にはコストがかかり、生徒の質を制限し、教師の元のサンプリング パスにあまりにも密接に結びついてしまいます。
この制限を解消し、DMD トレーニングを改善する一連のテクニックである DMD2 を紹介します。
まず、回帰損失を排除し、高価なデータセット構築の必要性を排除します。
結果として生じる不安定性は、偽の批評家が生成されたサンプルの分布を正確に推定していないことが原因であることを示し、改善策として 2 つのタイムスケールの更新ルールを提案します。
次に、GAN 損失を蒸留手順に統合し、生成されたサンプルと実際の画像を区別します。
これにより、実際のデータで生徒モデルをトレーニングできるようになり、教師モデルからの不完全な実際のスコア推定が軽減され、品質が向上します。
最後に、トレーニング手順を変更して、マルチステップ サンプリングを有効にします。
この設定では、トレーニング中に推論時間ジェネレーターのサンプルをシミュレートすることで、トレーニングと推論の入力不一致の問題を特定して対処します。
総合すると、私たちの改善により、ワンステップ画像生成の新たなベンチマークが設定され、ImageNet-64×64 では FID スコアが 1.28、ゼロショット COCO 2014 では 8.35 となり、推論コストが 500 分の 1 削減されたにもかかわらず、元の教師を上回りました。
さらに、私たちのアプローチがSDXLを蒸留することによってメガピクセル画像を生成できることを示し、数ステップの方法の中で優れた視覚品質を実証します。

要約(オリジナル)

Recent approaches have shown promises distilling diffusion models into efficient one-step generators. Among them, Distribution Matching Distillation (DMD) produces one-step generators that match their teacher in distribution, without enforcing a one-to-one correspondence with the sampling trajectories of their teachers. However, to ensure stable training, DMD requires an additional regression loss computed using a large set of noise-image pairs generated by the teacher with many steps of a deterministic sampler. This is costly for large-scale text-to-image synthesis and limits the student’s quality, tying it too closely to the teacher’s original sampling paths. We introduce DMD2, a set of techniques that lift this limitation and improve DMD training. First, we eliminate the regression loss and the need for expensive dataset construction. We show that the resulting instability is due to the fake critic not estimating the distribution of generated samples accurately and propose a two time-scale update rule as a remedy. Second, we integrate a GAN loss into the distillation procedure, discriminating between generated samples and real images. This lets us train the student model on real data, mitigating the imperfect real score estimation from the teacher model, and enhancing quality. Lastly, we modify the training procedure to enable multi-step sampling. We identify and address the training-inference input mismatch problem in this setting, by simulating inference-time generator samples during training time. Taken together, our improvements set new benchmarks in one-step image generation, with FID scores of 1.28 on ImageNet-64×64 and 8.35 on zero-shot COCO 2014, surpassing the original teacher despite a 500X reduction in inference cost. Further, we show our approach can generate megapixel images by distilling SDXL, demonstrating exceptional visual quality among few-step methods.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google