17日前

拡散モデルは本質的にワンステップ生成器である

Bowen Zheng, Tianming Yang
拡散モデルは本質的にワンステップ生成器である
要約

拡散モデル(Diffusion Models, DMs)は、画像生成をはじめとする分野で著しい成果を上げている。事前に学習されたスコアモデルに基づき、SDE/ODEソルバーによって定義される軌道に沿って精密なサンプリングを行うことで、DMsは優れた高品質な出力を生成できる。しかし、このような高精度なサンプリングは多数のステップを要し、計算コストが非常に高くなる傾向にある。この課題に対処するために、より単純な学生モデル(student model)がより複雑な教師モデル(teacher model)の振る舞いを模倣することで、1ステップで生成可能なモデルをDMから蒸留(distillation)するインスタンスベースの蒸留手法が提案されてきた。しかしながら、本研究では、こうした手法に内在する制限を明らかにした。すなわち、教師モデルはより多くのステップとパラメータ数を持つため、学生モデルとは異なる局所最適解(local minima)に収束する傾向があり、その結果、学生モデルが教師モデルを再現しようとしても、最適な性能が得られないという問題が生じる。この問題を回避するため、我々は独創的な分布蒸留法(distributional distillation method)を提案する。この手法は特有の分布損失(distributional loss)を用いることで、従来のSOTA(State-of-the-Art)を上回る性能を達成しつつ、大幅に少ない訓練画像で学習が可能である。さらに、DMの各層が異なる時間ステップで非均一に活性化されることを示し、1ステップでの画像生成という本質的な能力を有していることを明らかにした。分布蒸留の過程で、DMの大部分の畳み込み層を固定(freeze)することで、この内在的な能力が発揮され、さらなる性能向上が実現された。本手法は、CIFAR-10(FID 1.54)、AFHQv2 64×64(FID 1.23)、FFHQ 64×64(FID 0.85)、ImageNet 64×64(FID 1.16)においてSOTAの結果を達成しており、極めて高い効率性を示している。これらの多くは、8台のA100 GPUを用いて6時間以内に、わずか500万枚の訓練画像で達成された。

拡散モデルは本質的にワンステップ生成器である | 最新論文 | HyperAI超神経