【LDDGAN】最高速の推論が可能な拡散モデル

Diffusion Model 2024年09月29日

3つの要点
✔️ 拡散モデルは画質や多様性、学習安定性でGANを上回っているが、推論速度が非常に遅いためリアルタイムでの使用が難しい
✔️ 先行研究のDiffusionGANとWDDGANは推論速度を大幅に向上できたが、GANと比較してまだ遅く、生成画質が低下するという課題がある
✔️ LDDGANは低次元潜在空間でGANの敵対的学習を活用することで、高画質や多様性を維持し、拡散モデルの中で最も高速な処理を実現した

Latent Denoising Diffusion GAN: Faster sampling, Higher image quality
written by Luan Thanh Trinh, Tomoki Hamagami
(Submitted on 17 Jun 2024)
Comments: Published on arxiv.
Subjects: Computer Vision and Pattern Recognition (cs.CV)

code：

本記事で使用している画像は論文中のもの、紹介スライドのもの、またはそれを参考に作成したものを使用しております。

はじめに

拡散モデルは、高品質で多様な画像を生成する強力な手法として注目されており、多くの場合、GANを上回る性能を示しています。しかし、推論速度が遅いため、リアルタイムでの応用が難しいという課題があります。この問題に対処するため、DiffusionGANは条件付きGANを活用して、ノイズ除去のステップ数を大幅に減らし、推論速度を向上させました。その改良版であるWavelet Diffusionは、データをウェーブレット空間に変換することで、さらに処理を高速化しました。しかし、これらのモデルは依然としてGANと比較してまだ遅く、画質も低下しました。

これらのギャップを埋めるために、今回の解説論文はLatent Denoising Diffusion GAN（LDDGAN）を提案しました。このモデルは、事前学習済みのオートエンコーダーを使用して画像をコンパクトな潜在空間に圧縮し、推論速度と画質を大幅に向上させます。さらに、多様性と画質を高めるためのWeighted Learningの訓練戦略も提案しました。

CIFAR-10、CelebA-HQ、LSUN Churchデータセットでの実験結果は、LDDGANが拡散モデルの中で最高速の実行速度を達成したことを示しています。DiffusionGAN（DDGAN）やWavelet Diffusion（WDDGAN）といった先行研究と比較して、すべての評価指標において顕著な改善を示しました。特に、GANと同等の推論速度と画質を保ちながら、生成画像の多様性がGANを大きく上回りました。

提案手法

全体の概要

LDDGANの概要は図 1 のようになり、次の４ステップがあります。（i）事前学習済みの Encoder で、入力画像を低次元潜在変数に変換します。ここで、WDDGANのように 4 倍圧縮の制限がなく、8 倍、16 倍も可能となります。（ii）拡散過程を行って、ガウス分布に限らず、マルチモーダル分布も許可します。従来 Diffusion のように、サンプリング回数 T は数百、数千回ではなく、T ≤ 4 と設定されます。（iii）生成器は識別器のフィードバックに基づいて、逆変換のマルチモーダル分布の予測を学習して逆拡散過程を行います。(iv) 事前学習済みの Decoder で、潜在変数からピクセル空間に変換して、元画像を復元します。

LDDGANの期待できる効果としては、まず、入力画像をできるだけ多く圧縮するので、拡散モデルの学習による計算コストを大幅に削減できることと、先行研究よりも高速の推論という２点を挙げられます。また、低次元潜在空間は尤度ベースの生成モデルである Diffusion に最適し、出力画質・多様性の向上も可能となります。

Autoencoderの学習

LDDGANのオートエンコーダーの構造は、Esserらが提案したVQGANを基にしています。その特徴的な点は、デコーダー内に量子化層を組み込んでいることです。従来の手法では、オートエンコーダーの損失関数にカルバック・ライブラー（KL）ダイバージェンス・ペナルティを用いることが一般的でした。このアプローチは、学習された潜在空間を正規分布に近似させることを促し、モデルの学習方針がガウス分布に大きく依存する場合に効果的であることが確認されています。しかし、LDDGANでは正規分布に限らず、複雑でマルチモーダルな分布も許容します。そのため、このKLペナルティは使用せず、オートエンコーダーが潜在空間を自由に利用できるようにしています。これにより、画像の圧縮・復元能力を優先させることが可能となります。

表1の結果は、この仮説の正当性を示しています。適切な潜在空間を自由に探索できるようにしたことで、ほとんどの場合において結果が大幅に改善されました。特筆すべきは、CELEBA-HQデータセットでの結果です。ここでは、KLペナルティを用いたオートエンコーダよりも再構成FIDが劣るオートエンコーダを使用したにもかかわらず、メインモデルはより優れたFIDとRecallを達成しました。

学習損失とWeighted Learning

LDDGANの生成器と識別器の敵対的損失は次の式のようになります。

敵対的損失だけで学習する場合、実データにそっくりの画像を生成できますが、識別器を通じて間接的に学習するため、収束に時間がかかります。そこで、生成器の学習の収束を促進するために、下の式に示すように、元の画像と生成画像の違いを表す再構成損失も導入しました。

複数の損失関数を扱う場合、従来の手法では固定のパラメータを用いた線形結合で最終的な損失を合成していました。これは、再構成損失の重要度が一定に保たれることを意味します。しかしながら、再構成損失は入力データと同じデータを異なるノイズで生成するよう制約を課すため、生成されるサンプルの多様性が低下する可能性があります。そのため、LDDGANでは次の式のように、Weighted Learningが提案されました。図２はWeighted Learningの１つの例です。

まず学習の初期段階では、再構成損失の重要度をほぼ1に設定し、収束を促進します。その後、学習が進むにつれてこの重要度を徐々に下げ、敵対的損失を優先することでサンプルの多様性を高めていきます。学習の終盤に近づくと、再構成損失の減少率を緩やかにし、全体の安定性を優先します。この手法により、学習の収束を高速化しながら、画質、多様性、そして学習の安定性を維持するという効果が期待できます。

表2の実験結果は、この仮説を裏付けるものとなりました。再構成損失を採用することで、両データセットにおいて、敵対的損失のみに依存する場合と比較して、より優れた画質（FID）が得られました。しかし、その一方で多様性（Recall）は低下しました。これに対し、Weighted Learningを導入した場合、画質と多様性の両方が改善される結果となりました。

実験

データセットと評価指標

LDDGANの効果を検証するため、低解像度のCifar10（32x32サイズ）、高解像度のCeleba-HQおよびLSUN（256x256サイズ）のデータセットを用いて実験が行われました。評価指標として、推論時間、画質を表すFID（Fréchet Inception Distance）、そして多様性を表すRecallが使用されました。推論時間は、100枚の画像バッチを生成する処理を300回試行し、その平均時間を測定しています。

先行研究との比較

表3、４，5の結果は、LDDGANが拡散モデルの弱点をさらに改善し、高い画質と多様性を維持しながら、拡散モデルの中で最先端の実行速度を達成していることを示しています。

SDE ScoreやDDPMなど、LDDGANよりも優れたFIDを達成する拡散モデルも存在します。しかし、LDDGANはScore SDEの5000倍、DDPMの1000倍のサンプリング速度を実現しており、速度面で圧倒的な優位性を示しています。

特に注目すべきは、先行研究であるDDGANやWDDGANとの比較結果です。提案手法は、全ての評価指標においてこれらの手法を上回っています。

さらに、GANのSOTAとされるStyleGANと比較した場合、LDDGANは同等の画質と推論速度を達成しつつ、多様性の面で大きく上回っていることが確認されました。

図3は定性的なと比較です。LDDGANは明らかにより良いサンプル品質を達成しています。CelebA-HQデータセットでは、DDGANとWDDGANの両方が、はっきりとした完全な人間の顔を生成することに苦戦しており、しばしば歪んだ特徴を生み出しています。同様に、LSUN Churchデータセットでも、これらのモデルは建築物の直線的で水平な細部を正確に描写することが難しいです。それに対してに、LDDGANは一貫して現実的で鮮明な画像を生成しています。