見出し画像

"Generative Modeling by Estimating Gradients of the Data Distribution"を読んだ（後半）

岩隈啓悟

はじめに

株式会社GA technologies、AI Strategy Center（以下AISC）の岩隈です。最近スプラトゥーン３でウデマエがS＋50になりました。Xマッチの実装が待ち遠しいですね！

今回の内容

先日自分が社内勉強会で取り上げた、Stanford UniversityのSongらによる”Generative Modeling by Estimating Gradients of the Data Distribution”（NeurIPS 2019）の内容を前半・後半の2回に分けてブログにしようと思います。
後半の今回は、score-based generative modelingに関する課題と論文での提案をまとめようと思います（使用する式番号は前半に続く形になっています）。
前半の内容（論文に関する予備知識とscore-based generative modelingに関する定義まで）に関しては、前回の記事をご確認下さい！

論文の内容に関する概要を知りたい方は、本編の「論文の概要」の節があるのでそちらをご確認下さい！

取り上げた理由（再掲）

最近、深層生成モデルの新しいフレームワークとして拡散モデル（diffusion models）が注目を集めています。その火付け役である、Denoising Diffusion Probabilistic Models（以下DDPM）を紹介しようと思ったのですが、DDPMとの関係がよく言及されるNoise Conditional Score Networksの方が時系列的に先だったのでこちらを優先することにしました。
DDPMについてもこのブログで次回以降取り上げようと思います。

注意（再掲）

このブログには筆者（岩隈）の理解による間違いが含まれている可能性があります。お気づきの際はご指摘頂けるとありがたいです。
参考にした内容については参考文献を記載しているので詳細はそちらを参照して下さい。
本編で使用している画像については特に断りがない限り、今回取り上げた論文である参考文献[3]から引用しております。

本編

論文の概要（再掲）

この論文では、データ分布のscore（＝確率変数に関する確率分布の勾配）をモデリングの対象とし、Score Matchingによる学習とLangevin dynamicsによるサンプリングからなる生成モデルのフレームワーク、score-based generative modelingを定義しています。また、適用時の課題を分析しそれを回避するためのモデルの定式化（Noise Conditional Score Network）とデータのサンプリング手法（annealed Langevin dynamics）を新しく提案しています。最後に、実験としてCIFAR-10でのimage generationタスクを行い、Inception scoreで当時のSOTAを達成しています。

従来の課題

従来のscore-based generative modelingの課題としてここでは大きく2つ取り上げます。

1つ目は多様体仮説に関してです。経験的に、実在するデータは高次元な空間内の低次元な多様体に存在するとされています。そのため、１）実データのみではscoreを定義できない空間が存在する、２）Score Matching（以下SM）ではデータの分布が空間全体で定義されている場合のみ一貫したscoreの推定が行える、という点で実用上の妨げになっています。
Fig.1は、CIFAR10とResNetを用いてSMによりscore関数を学習させたときの損失の推移のグラフです。左ではノイズを加えず、右ではノイズを加えています。画像のピクセル値を$${[0, 1]}$$に正規化されており、それに対してノイズは$${\mathcal{N}(0, 0.0001)}$$と極めて小さいにも関わらず学習の安定に寄与していることが分かります。個人的な理解ですが、実データのみではサポート外のscoreの推定に何もペナルティがかからないので一貫性が保てず学習が安定しないのだと思っています。
（補足：この実験では、SMにDenoising Score Matching（以下DSM）ではなくこの論文の著者らが以前に提案しているSliced Score Matching（以下SSM）[7]を使用しているため、ノイズを加えていないデータの確率分布に対するscore関数の学習が可能です。）

Figure 1: Left: Sliced score matching (SSM) loss w.r.t iterations. No noise is added to data. Right: Same but data are perturbed with N(0, 0.0001).

2つ目はデータの低密度領域に関してです。このような領域に対してはscoreの推定とscoreからのサンプリング両方の面で課題があります。

まずデータの低密度領域のscoreの推定に関しては学習に十分なサンプルが足りず不正確になる傾向があります。
Fig.2は、混合ガウス分布$${p_\textrm{data} = \frac{1}{5}\mathcal{N}((-5, -5), I) + \frac{4}{5} \mathcal{N}((5, 5), I) }$$を用いたscoreの真値$${\nabla_\mathbf{x} \log p_{\textrm{data}}(\mathbf{x})}$$と推定値$${\mathbf{s}_\textbf{\textit{\char"03b8}}(\mathbf{x})}$$の可視化結果です。両者を比較すると、$${\mathbf{s}_\textbf{\textit{\char"03b8}}(\mathbf{x}) \approx \nabla_{\mathbf{x}} \log p(\mathbf{x})}$$とみなせるのは赤い波線で囲まれた領域のみとなります。

Figure 2: Left: ∇​x log p(x); Right: s(x). The data density p(x) is encoded using orange colormap: darker color implies higher density. Red rectangles highlight regions where ∇​x log p(x) ≈ s(x).

次にデータの低密度領域のscoreからのサンプリングに関しては、データ分布が低密度領域を跨いでいくつかのモードで構成されている場合、モードごとの重みを適切に考慮したサンプリングに時間がかかる、もしくは出来ないといったことが起こります（直感的には、低密度領域では更新式における乱数項が支配的になるため）。
Fig.3は先ほどと同じ混合ガウス分布を用いて、（a）実際にサンプリングしたもの、（b）scoreの真値からLangevin dynamicsを用いてサンプリングしたもの、（c）scoreの真値から後述するannealed Langevin dynamicsを用いてサンプリングしたもの、をそれぞれ表しています。Fig.3（b）から、低密度領域を跨ぐモードを持つデータ分布に関するLangevin dynamicsではモードの重みを考慮できていないことが分かります。

Figure 3: Samples from a mixture of Gaussian with different methods. (a) Exact sampling. (b) Sampling using Langevin dynamics with the exact scores. (c) Sampling using annealed Langevin dynamics with the exact scores. Clearly Langevin dynamics estimate the relative weights between the two modes incorrectly, while annealed Langevin dynamics recover the relative weights faithfully.

論文での提案

前節で挙げた課題に対する解決策の1つは、データ分布にノイズを加えることです。
データ分布にノイズを加えることでモデリングの対象は真のデータ分布ではなくなってしまいますが、高次元な空間内で低次元な多様体に限定されていた分布のサポートを空間全体に広げることができ、さらに分布の低密度領域を補完する効果があります。
特に後者では、加えるノイズの大きさによってモデリングの対象となる分布（＝ノイズを加えたデータ分布）の学習のしやすさと真のデータ分布との近さのトレードオフが存在します。加えるノイズが大きければモデリングの対象となる分布の低密度領域が減らすことができるためより多くの学習のための信号を獲得できるようになりますが、モデリングの対象自体は本来求めたい真のデータ分布からより離れてしまう、といったトレードオフです。

これらの観察からこの論文では、１）様々なレベルのノイズを加えた複数のデータ分布を用意する、２）単一のネットワークの学習でその全てのノイズのレベルでscoreを推定する、の2点によってscore-based generative modelingを改善することを提案しています。論文同様、提案の詳細についてネットワーク、学習、サンプリングの順に分けて紹介します。

Noise Conditional Score Networks

ノイズのレベルとして、$${\frac{\textit{\char"03c3}_1}{\textit{\char"03c3}_2} = \cdots = \frac{\textit{\char"03c3}_{L-1}}{\textit{\char"03c3}_{L}} > 1 }$$を満たす正の等比数列$${\{\textit{\char"03c3}_i\}^L_{i=1}}$$を考えます。ノイズを加えたデータ分布は$${q_\textit{\char"03c3}(\mathbf{x}) \triangleq \int p_\textrm{data}(\mathbf{t}) \mathcal{N}(\mathbf{x} | \mathbf{t}, \textit{\char"03c3}^2 \mathbf{I}) \mathrm{d}\mathbf{t}}$$と表記することにします。
ここで$${\sigma_1}$$は従来の課題として挙げた要因を緩和できるほど大きく、$${\sigma_L}$$は実データへの影響をほとんどないほど小さくなるよう恣意的に選ばれています（このハイパーパラメータの適切な選び方に関しては参考文献[5]の第14回講義資料や参考文献[6]で検討されていますがここでは割愛します）。
様々なノイズを加えたデータ分布からscoreを推定するために、$${\forall\sigma \in \{\sigma_i\}^L_{i=1}: \mathbf{s}_\textbf{\textit{\char"03b8}}(\mathbf{x}, \textit{\char"03c3}) \approx \nabla_{\mathbf{x}} \log q_\textit{\char"03c3}(\mathbf{x})}$$となるような条件付きのニューラルネットワークを学習します。論文でこのネットワークはNoise Conditional Score Network（以下NCSN）と呼ばれており、アーキテクチャにはU-Net構造が採用されています。

SMによるNCSNの学習

NCSNの学習にはDSMを利用します。
各ノイズのレベルでのDSMの目的関数は以下のように書けます。

$$
\begin{align}
\tag{10}
\ell(\textbf{\textit{\char"03b8}};\textit{\char"03c3}) \triangleq
\frac{1}{2}\mathbb{E}_{p_\textrm{data}(\mathbf{x})}
\mathbb{E}_{\~\mathbf{x} \sim \mathcal{N}(\mathbf{x}, \textit{\char"03c3}^2\mathbf{I})}
\Big[ \Big\|
\mathbf{s}_\textbf{\textit{\char"03b8}} (\~\mathbf{x}, \textit{\char"03c3}) - 
\frac{\mathbf{x} - \~{\mathbf{x}}}{\textit{\char"03c3}^2}
\Big\|^2_2 \Big]
\end{align}
$$

全てのノイズのレベル$${\textit{\char"03c3} \in \{\textit{\char"03c3}_i\}^L_{i=1}}$$で統合した目的関数は以下になります。

$$
\begin{align}
\tag{11}
\mathcal{L}(\textbf{\textit{\char"03b8}}; \{\textit{\char"03c3}_i\}^L_{i=1})
\triangleq \frac{1}{L} \sum_{i=1}^L 
\lambda(\textit{\char"03c3}_i)
\ell(\textbf{\textit{\char"03b8}};\textit{\char"03c3}_i)
\end{align}
$$

$${\lambda(\sigma_i)}$$はノイズのレベル$${\sigma_i}$$に依存した係数関数です。全てのノイズのレベル$${\textit{\char"03c3} \in \{\textit{\char"03c3}_i\}^L_{i=1}}$$で$${\lambda(\textit{\char"03c3}_i)\ell(\textbf{\textit{\char"03b8}};\textit{\char"03c3}_i)}$$がおよそ等しいオーダーになるように、経験的に$${\lambda(\sigma) = \sigma^2}$$が用いられています。

Annealed Langevin dynamics

先に、学習したNCSNからサンプリングを行う方法として、この論文で提案されたannealed Langevin dynamicsの擬似コードをAlg.1に示します。

Algorithm 1: Annealed Langevin dynamics.

annealed Langevin dynamicsでは、最初に一様分布のような何らかのサンプリングが容易な事前分布から初期値$${\~\mathbf{x}_0}$$を用意します。
次に、その初期値とステップサイズ$${\alpha_1}$$を用いて、$${q_{\sigma_1}(\mathbf{x})}$$からのLangevin dynamicsによるサンプリングを行います。その後、$${q_{\sigma_1}(\mathbf{x})}$$からのサンプリング値を初期値とし、小さくしたステップサイズ$${\alpha_2}$$を用いて、$${q_{\sigma_2}(\mathbf{x})}$$からのLangevin dynamicsによるサンプリングを行います。同様に、ノイズが十分小さいとみなせる$${q_{\sigma_L}(\mathbf{x})}$$からのLangevin dynamicsによるサンプリングが終わるまで繰り返し、最終的なサンプリング値$${\~\mathbf{x}_T}$$を得ます。

このサンプリング方法では、十分大きいノイズのレベルを適用した$${q_{\sigma_1}(\mathbf{x})}$$からのLangevin dynamicsによるサンプリングから始めることで初期の多様体仮説やデータの低密度領域に関する課題を回避しています。また、連続するノイズのレベルを適用した分布どうしには大きな違いがなく前のサンプリング値が後の初期値として良い点を与える（＝scoreの推定が不正確な低密度領域に陥りにくい）ことを利用することでその後の課題も回避しています。

実験による評価

この論文では、画像全体を生成するimage generationのタスクと画像の覆った一部を生成するimage impaintingのタスク、2つの実験を行なっていますがここではimage generationタスクの結果のみを紹介します。

学習データセットにはMNIST, CelebA, CIFAR-10の3つを使用し、 CelebAに対しては解像度を$${32 \times 32}$$に落としています。また、ピクセル値は$${[0, 1]}$$に正規化されています。各ハイパーパラメータは次のように設定されています。

  • ノイズのレベル数：$${L = 10}$$

  • ノイズの初期レベル：$${\sigma_1 = 1}$$

  • ノイズの最終レベル：$${\sigma_L = 0.01}$$

  • 各ノイズのレベルでのLangevin dynamicsのステップ数：$${T = 100}$$

  • 最終的なステップサイズ：$${\epsilon = 2 \times 10^{-5}}$$

Fig.4にimage generationの結果を示します。近年のDiffusion Modelと比べると見劣りますが、多様で識別できる程度の品質を達成していることが確認できます。

Figure 4: Uncurated samples on MNIST, CelebA, and CIFAR-10 datasets.

次にTab.1にCIFAR-10を用いたimage generationにおける定量的な評価を示します。InceptionやFIDといった生成された画像の品質を表すスコアで当時のSOTAやそれに匹敵する性能を達成していることが確認できます。

Table 1: Inception and FID scores for CIFAR-10

おわりに

今回はscore-based generative modelingが抱える課題と論文での提案をまとめました。拡散モデルとは違う視点でノイズを導入している点が個人的には面白くて勉強になったなと思っています。その一方で、この論文の時点で、実験結果から分かるように生成画像の解像度が低く、また画像の生成に多くの更新が必要になるため時間がかかるといった課題があります。今後、これらの課題を解決する提案を行なった論文の紹介もしていけたらと思います。

ブログのテーマや内容に関しても引き続き改善していけるように（そして継続できるように…）頑張ります！！

前半同様、最後に参考文献を共有します。

参考文献

  1. How to Train Your Energy-Based Models [PDF]
    Song, Y. and Kingma, D.P., 2021. arXiv preprint arXiv:2101.03288.

    1. Energy-Based Modelのチュートリアル論文です。EBMの理解の参考にしました。

  2. A connection between score matching and denoising autoencoders [PDF]
    Vincent, P., 2011. Neural computation, Vol 23(7), pp. 1661--1674. MIT Press.

    1. Denoising Score Matchingを提案した論文です。

  3. Generative Modeling by Estimating Gradients of the Data Distribution  [PDF]
    Song, Y. and Ermon, S., 2019. Advances in Neural Information Processing Systems, pp. 11895--11907.

    1. 今回取り上げたNoise Conditional Score Networkを提案した論文です。

  4. Yong Song, "Generative Modeling by Estimating Gradients of the Data Distribution", 5 May. 2021,  https://yang-song.net/blog/2021/score/

    1. NCSNの主著者がNCSNから後続研究までをカジュアルに解説したブログです。

  5. Stanford CS236, "Deep Generative Models", 2021, https://deepgenerativemodels.github.io/

    1. NCSNの著者らによる講義スライドです。深層生成モデル全般の理解の参考にしました。

  6. Improved Techniques for Training Score-Based Generative Models  [PDF]
    Song, Y. and Ermon, S., 2020. Advances in Neural Information Processing Systems 33: Annual Conference on Neural Information Processing Systems 2020, NeurIPS 2020, December 6-12, 2020, virtual.

    1. NCSN v2として、NCSNに関する様々な改善を行った論文です。

  7. Sliced score matching: A scalable approach to density and score estimation  [PDF]
    Song, Y., Garg, S., Shi, J. and Ermon, S., 2020. Uncertainty in Artificial Intelligence, pp. 574--584.

    1. Sliced Score Matchingを提案した論文です。

この記事が気に入ったらサポートをしてみませんか？