Wavelet Diffusion: 最高速の拡散モデル
3つの要点
✔️ 拡散モデルは高い出力画質と高い多様性を持ち、GANを上回っているが、推論速度が非常に遅いため、リアルタイムでの使用が難しい
✔️ 先行研究のDiffusionGANはGANの仕組みと組み合わせることで、推定速度を大幅に向上できた
✔️ DiffusionGANをベースにして、低周波数成分と高周波数成分に変換することで、入力を４倍圧縮して、高画質維持しながら、拡散モデルの中で最高速を達成した
Wavelet Diffusion Models are fast and scalable Image Generators
written by Hao Phung, Quan Dao, Anh Tran
(Submitted on 29 Nov 2022 (v1), last revised 22 Mar 2023 (this version, v2))
Comments: Accepted to CVPR 2023
Subjects: Computer Vision and Pattern Recognition (cs.CV); Image and Video Processing (eess.IV)
code：
本記事で使用している画像は論文中のもの、紹介スライドのもの、またはそれを参考に作成したものを使用しております。
はじめに
拡散モデルは最近登場し、急速に成長し、多くの研究者の関心を引いています。これらのモデルは、ランダムなノイズ入力から高品質な画像を生成します。特に画像生成タスクにおいては、最先端の生成モデルであるGANよりも優れた成果を上げています。拡散モデルはさまざまな条件付き入力を柔軟に扱えるため、テキストから画像の生成、画像から画像への変換、画像の修復など、さまざまなアプリケーションが可能です。これにより、AIベースのデジタルアートや他の領域への応用が期待されています。
拡散モデルは大きなポテンシャルを持っていますが、推定速度が非常に遅いため、GANのように広く採用されることが妨げられています。基本的な拡散モデルでは、望ましい出力品質を得るために数分かかります。多くの研究が推定時間を短縮するために提案されていますが、最速のアルゴリズムでも、32×32の画像を生成するのに数秒かかります。DiffusionGANは、拡散モデルとGANを組み合わせることで推定時間を劇的に向上させましたが、それでもまだ大規模またはリアルタイムのアプリケーションには適していません。
リアルタイムのアプリケーションでの利用のために、今回の解説論文では、WaveletDiffusionという新しい拡散を提案します。離散ウェーブレット変換を使用し、低周波数成分と高周波数成分に変換することで、入力を４倍圧縮して推論時間を大幅に削減しました。また、ウェーブレットの特徴をより効率に活用し、出力品質を維持するために、ウェーブレットに特化した生成器も提案します。実験結果により、WaveletDiffusionは高画質維持しながら、拡散モデルの中で最高速を達成していると確認されました。
便宜上、後の節ではDiffusionGANをDDGANと略します。
提案手法
ウェーブレットベースの拡散スキーム
本論文では、入力画像を4つのウェーブレットサブバンドに分解し、それらを単一の対象として拡散プロセスに連結します（図１に示す）。このようなモデルは、元の画像空間ではなくウェーブレットスペクトル上で動作します。その結果、モデルは生成された画像の詳細をさらに増やすために高周波情報を活用できます。一方、ウェーブレットサブバンドのサイズは元の画像の4倍小さいため、サンプリングプロセスの計算量が大幅に低減されます。
本論文の手法はDDGANモデルに基づいており、入力はウェーブレット変換の4つのウェーブレットサブバンドです。入力画像x ∈ R 3×H×Wが与えられると、それを低および高のサブバンドに分解し、さらに連結して行列y ∈ R 12× H 2 × W 2を形成します。この入力は、最初の線形層を介してベースチャネルDに投影され、DDGANと比較してネットワークの幅は変更されません。したがって、ほとんどのネットワークは空間次元が4つ減少したことにより、その計算が大幅に削減されます。
学習の損失関数
・敵対的損失
DDGANと同様、敵対的損失を通して、生成器と識別器を最適化します：
・再構成損失と総合の損失関数
上の敵対的損失に加えて、周波数情報の損失を防ぎ、ウェーブレットサブバンドの一貫性も保持するために、再構築項も追加します。これは、生成された画像とそのグラウンストゥルースとの間のL1損失として定式化されます。
生成器の全体的な目的は、敵対的損失と再構成損失の次の線形結合となります：
ここで、λは重み付けハイパーパラメータです。定義された数回のサンプリングステップの後、推定されたノイズ除去されたサブバンドy'0が取得されます。最終的な画像は、ウェーブレット逆変換を使用して復元することができます: x'0 = IWT(y'0)。
ウェーブレットを組み込んだ生成器
図2は、提案されたウェーブレット埋め込み生成器の構造を示しています。提案された生成器は、M個のダウンサンプリングブロックとM個のアップサンプリングブロックを持つUNet構造に従います。また、同じ解像度のブロック間にスキップ接続があります。ただし、通常のダウンサンプリングとアップサンプリング演算子の代わりに、周波数に敏感なブロックを使用します。最低解像度では、低周波数と高周波数成分に対するより良い注意を払うために周波数ボトルネックブロックを採用しています。
最後に、オリジナルの信号Yをエンコーダーの異なる特徴ピラミッドに組み込むために、ウェーブレットダウンサンプルレイヤーを使用した周波数残差接続を導入しています。ここで、Yは入力画像を示し、FiはYのi番目の中間特徴マップを示します。
周波数を考慮したダウンサンプリングとアップサンプリングブロック
従来のアプローチでは、エイリアシングアーティファクトを軽減するために、ダウンサンプリングとアップサンプリングプロセスにぼかしカーネルを使用していました。代わりに、本論文はウェーブレット変換の固有の特性を利用して、より良いアップサンプリングとダウンサンプリングを行っています（図3に示されています）。
これにより、これらの操作での高周波情報の認識が強化されます。特に、ダウンサンプリングブロックは、入力特徴Fi、潜在的なz、および時間埋め込みtのタプルを受け取り、それらを一連の層を通して処理してダウンサンプリングされた特徴と高周波サブバンドを返します。これらの返されたサブバンドは、アップサンプリングブロック内で周波数の手がかりに基づいて特徴をアップサンプリングするための追加の入力として機能します。
実験
データセット
32×32のCIFAR-10、64×64のSTL-10、および256×256のCelebA-HQとLSUN-Churchデータセットで実験を行いました。また、高解像度で、天安手法の効果を検証するために、CelebA-HQ（512＆1024）の高解像度画像での実験も行いました。
評価指標
画像の画質をFrechet Inception Distance（FID）で、サンプルの多様性をRecallで測定します。DDGANと同様に、FIDとRecallは、50,000個の生成されたサンプルに対して計算されます。推定速度は、バッチサイズ100で、300回の試行にわたる平均推論時間で測定されます。また、CelebA-HQ 512×512のような高解像度画像の推論時間は、25サンプルのバッチから計算されます。
実験結果
各データセットで、VAE、GAN、拡散モデルなどの代表的な生成モデルとの比較結果は表１、２と３に示されています。VAEのSOTA、つまり最強モデルと比較すると、提案手法はすべての評価指標で大きく上回りました。特に、画質を表すFIDはVAEの４倍以上でした。
GANのSOTAと比較すると、推定速度はほぼ同じレベルで、画質も上位です。多様性に関して、どの場合でもGANよりも１０％以上高いです。
拡散モデル、DDGANと比較すると、提案手法は拡散モデルの中で最高の推定速度を達成しました。特に、DiffusionのSOTAよりも５００倍以上早くなっています。画質と多様性も上位で、場合によって、拡散モデルのSOTAより1~2ポイントほど高いです。また、すべての評価指標で、先行研究であるDDGANを上回りました。
ウェーブレットを組み込んだ生成器の効果
提案された生成器の各個々のコンポーネントの効果がCelebA-HQ 256×256で検証されました。ここで、フルモデルには残差接続、アップサンプリング、ダウンサンプリングブロック、およびボトルネックブロックが含まれています。表４に示すように、各コンポーネントがモデルのパフォーマンスに良い影響を与えています。提案された3つのコンポーネントをすべて適用することで、5.94で最高のパフォーマンスを達成します。ただし、性能の向上には推定速度というわずかな代償が伴います。
1枚の画像を生成する際の実行時間
さらに、実際のアプリケーションでは期待されるように、単一の画像における提案手法の優れた速度を示します。表5では、その時間と主要なパラメータを示しています。提案手法は、1024×1024までの画像をわずか0.1秒で生成することができ、これは拡散モデルとして初めて、ほぼリアルタイムの性能を達成したことになります。
結論
今回、画質とサンプリング速度の両方で優れた性能を示すWavelet Diffusionという新しい拡散モデルを紹介しました。画像と特徴空間の両方にウェーブレット変換を組み込むことで、提案方法は拡散モデルにおける最先端の実行速度を実現し、GANのSOTAとの差を縮めると同時に、StyleGAN2や他の拡散モデルと比較してほぼ同等の画像生成品質を得ることができます。さらに、提案方法はベースラインのDDGANよりも速い収束を提供し、提案されたフレームワークの効率を確認できます。
この記事に関するカテゴリー
