論文概要
タイトル：High-Resolution Image Synthesis with Latent Diffusion Models
著者：Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, Björn Ommer
学会：CVPR 2022
URL：https://arxiv.org/abs/2112.10752
内容：Stable Diffusionの元論文。画像生成において、拡散モデルがSOTAになりつつある。しかし、従来の拡散モデルはピクセル空間 (RGBなど) でサンプリングを行うので、計算量が大きい。その品質や精度を保ったまま、拡散モデルを限られた計算資源で訓練できるように、潜在空間上で拡散モデルを適用するLatent Diffusion Model (LDM) を提案している。テキストからの画像生成や画像生成 (条件の有無に関わらず) 、超解像などのタスクで、従来の拡散モデルに匹敵するスコアを達成した。
はじめに
（この論文を読む前に、「Denoising Diffusion Probablistic Models」の論文を読んだり、内容を調べたりすると、理解が容易になります。）
画像生成は、相当大きい計算量を必要とするタスクである。
このタスクにおいて、ここ数年拡散モデルはGANのようなモード崩壊や学習の不安定性もなく、自己回帰モデルのように数十億のパラメータを必要とせずに、画像の非常に複雑な分布をモデル化することができる。
高解像度画像合成の民主化
従来の拡散モデルは、学習と推論どちらもRGB画像の高次元空間での評価と勾配計算の繰り返しが必要になるため、計算量が非常に大きいという欠点がある。実際最もパワフルな拡散モデルは、学習にV100で数百日、推論でもA100で5日ほどかかる。そこで、学習と推論の両方において性能を損なわずに計算量を減らすことは拡散モデルを多くの人に扱えるようにするための鍵となる。
潜在空間への旅立ち
下図は、訓練済みモデルの歪み率とレートのトレードオフを表した図である。図中の横軸のbits/dimは、知覚できないほどの細部を表す。
拡散モデルの学習は2つの段階に分けられる。
- 知覚的圧縮（Perceptual Compression）：高周波の詳細なノイズを除去するが、意味的な変化はほとんど学習されない
- 意味的圧縮（Semantic Compression）：実際の生成モデルがデータの意味的な構成を学習する。
デジタル画像のほとんどのビットは、知覚できない細部に対応している。拡散モデルを用いると、損失項を最小化することでこの無意味な情報を抑制できるが、勾配（学習時）とニューラルネットワークのバックボーン（学習と推論）は依然としてすべてのピクセルで評価する必要があり、余計な計算につながる。
そこで、効果的な生成モデルとして潜在拡散モデル（LDM）を提案し、知覚できない細部のみを除去するマイルドな圧縮段階を別に設ける。この場合、LDMが意味的圧縮を行い、マイルドな圧縮段階が知覚的圧縮を行う。
著者らはトレーニングを2つの異なるフェーズに分けた。まず、データ空間と知覚的に等価な低次元の表現空間を提供するオートエンコーダを訓練する。空間次元に対してより優れたスケーリング特性を示す学習済み潜在空間で拡散モデルを訓練するため、過度の空間圧縮に頼る必要がない。また、複雑さが軽減されたことで、1回のネットワークパスで潜在空間から効率的に画像を生成することができる。この結果得られたモデルを潜在拡散モデル（LDM）と呼ぶ。
このアプローチの利点は、普遍的なオートエンコーダを1回しか訓練する必要がないため、複数の拡散モデルの訓練や、全く異なるタスクに再利用できることである。これにより、様々な画像から画像、テキストから画像へのタスクに対して、多数の拡散モデルを効率的に探索することができる。
提案手法
提案されたモデルのアーキテクチャを以下の図で示す。大きく2つの部分に別れており、ピンクの部分が画像を潜在空間にエンコード、もしくはその逆を行う知覚的画像圧縮であり、緑の部分が潜在空間において拡散プロセスを行う潜在拡散モデルである。
知覚的画像圧縮
知覚圧縮モデルとして、知覚損失とパッチベースの目的関数を組み合わせて訓練したオートエンコーダを用いている。これによって、再構成したデータが画像に限定されることを保証し、L2やL1ノルムなどのピクセル空間の損失のみに依存することで生じるぼやけを回避している。
潜在空間が過度に高分散化するのを避けるため、2つの異なる正則化を試した。KL-regはVAEと同様に、学習した潜在空間上の標準正規形に対するわずかなKLペナルティを課す。一方VQ-regはオートエンコーダのデコーダ内でのベクトル量子化による正則化項である。いずれにせよ、エンコードされた潜在変数がガウシアン正規分布（平均0、分散が単位行列）に近づくようにする。
これらから、画像の局所的な詳細の特徴量を失うことなくエンコードすることができるようになる。
潜在拡散モデル
知覚圧縮モデルの学習により、高周波で知覚できない細部が抽象化された効率的な低次元の潜在空間が利用できるようになりました。この空間は、高次元の画素空間と比較して、尤度ベースの生成モデルに適しており、（i）データの重要な意味的ビットに焦点を当て、（ii）低次元で計算効率の高い空間で訓練することができる。U-Netアーキテクチャから継承された拡散モデルの帰納的バイアスを利用すると、空間構造を持つデータに対して特に有効であるため、従来のアプローチで要求された積極的で品質を下げる圧縮レベルの必要性を緩和できた。
潜在拡散モデル内の行っていることはDDPMとほとんど変わらない。違いは潜在空間にエンコード・デコードしているという点のみ。
条件付けの機構
このモデルでは、text-to-image・超解像度化・inpainting といったマルチモーダルな画像生成を可能にするために、テキストデータ、もしくはセマンティクスセグメンテーション画像などを 別のエンコーダでエンコードした特徴ベクトルを、T個のU-Netの各エンコーダの中間層とデコーダの中間層それぞれにクロスアテンションと呼ばれる層で入力し、条件付け生成を行えるようにしている。クロスアテンションとは、Transformerで提案されているattention構造の1種であり、Queryに、U-Netの中間表現の線形変換、KeyとValueに条件付け入力を 別のエンコーダでエンコードした特徴ベクトルの線形変換を用いている。これにより、画像やテキスト全体の大域的な特徴量を失うことなく後段のネットワークに伝搬できるようになる。
条件付けの機構を含めた潜在拡散モデルの目的関数は以下のようになる。ただし、$x$は実データ、$y$は条件付け用の入力（テキストデータや、低解像度画像など）、$\varepsilon(x)$は知覚圧縮モデルのエンコーダ、$\epsilon_\theta()$は誤差を予測するモデル（U-Net）、$z_t$はタイムステップtでの潜在変数、$\tau_\theta(y)$は、条件付け入力用のエンコーダを表す。
$$L_{LDM} = \mathbf{E}_{\varepsilon(x),y,\epsilon \sim N(0,1),t} \left[||\epsilon - {\epsilon} _\theta (z_t,t,{\tau} _\theta (y)) ||_2^2 \right]$$
$\epsilon_\theta$と$\tau_\theta$を上の目的関数を用いて同時に学習させる。
実験
このモデルを用いて、画像合成に関するいくつかのタスクについての実験を行った。
知覚的圧縮のトレードオフについて
クラス条件付き画像生成モデルを、画像圧縮モデルのダウンサンプリング係数 $f$（何倍に圧縮するか）を変えて行った結果が以下の図になる。FIDは低ければ良い値、Inception Scoreは高ければ良い値となる。LDM-fはダウンサンプリング係数を表す。また、下2つの図の横軸のスループットは高いほど時間効率が良いことを示し、同じ$f$での異なる点は、左に行くほど拡散モデルのタイムステップ数が大きくなる。
$f$が小さいと、学習の進行が遅くなることがわかる。
逆に$f$が大きすぎても、トレーニングステップ数が小さい段階でFIDが停滞する。これは、圧縮が強すぎて情報損失が生じているためと考えられる。
また、生成時の生成画質とスループットを両立させるには、解像度を下げたほうが良く、$f$が4または8が最良と言える。
ベースラインとの比較
複数のデータセットにおいて無条件画像生成と、テキスト条件付き画像生成を行った結果をベースラインと比較したものを以下の図に示す。
無条件画像生成では、FIDにおいて、StyleGAN～StyleGAN2と同程度、そしてテキスト条件付き画像生成では、FIDにおいてGLIDEと同程度を、軽量化して達成した。
条件付き画像生成
text-to-image
text-to-image用のLDMでは、条件つけ生成のためのEncoderとしてTransformerを用いている。またTokenizerとしては、大規模自然言語モデルであるBERTのBERT-tokenizerを用いている。
学習用データセットとしては、テキストと画像の４億枚のペアデータセットであるLAION-400Mを使用している。
（Stable Diffusion v1はtext-to-image用のLDMとの違いとしては、学習用データセットとしてLAION-5Bと呼ばれる50億枚のデータセットを使用し、Transformerで実装される条件付け入力のためのエンコーダとして、事前学習済み CLIP ViT-L/14モデルを使用したものが Stable Diffusion v1である。）
超解像
条件付け入力に低解像度の画像を用いることで、超解像を行うことが可能になる。また、この場合、条件付けの入力をエンコードした後に、クロスアテンションではなく、U-Netの入力層に結合する。
下図の結果より、LDMはDiffusion Modelベースの超解像の強力なモデルであるSR3と比較して優れた品質を達成している。
Inpainting
一部分が抜け落ちている画像を入力とし、その部分を補完した画像を出力するInpaintingというタスクにおいても、従来のGANベースのLaMaと比較して良い結果となった。
まとめ
画像生成のための拡散モデルの欠点であった計算量の大きさを、潜在空間を用いることで、大きく軽減させた。
潜在空間を用いたり、条件付けをU-Netに追加するなどの点はすでにGANなどで行われていたが、それを拡散モデルに適用させ、また大規模なデータセットで学習させたStable DiffusionをOSSとして公開したのが新規性・貢献だと推測。