"Generative Modeling by Estimating Gradients of the Data Distribution"を読んだ（前半）
はじめに
株式会社GA technologies、AI Strategy Center（以下AISC）の岩隈です。
この度、AISCでブログを始めることになりました！！
AISC所属社員が自分が好きなテーマについて好きに投稿していく予定なのでぜひウォッチして頂けると嬉しいです！
またAISCではウェブサイトも公開しているので、こちらもぜひ見に来て頂けると嬉しいです！
今回の内容
AISCでは２週間に１回程度、所属社員持ち回りで気になった論文を選び部署内で共有する勉強会があります。ブログの初回となる今回は、先日自分が勉強会で取り上げた、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を達成しています。
Energy-Based Models （EBM）
生成モデル（generative models）とは、観測されたデータ$${\{\mathbf{x}_i\}_{i=0}^{N}}$$（$${\mathbf{x} \in \mathbb{R}^d}$$）から、その背後にあるデータの確率分布$${p_\textrm{data}(\mathbf{x})}$$を、（データ$${\mathbf{x}}$$を入力とする）$${\textbf{\textit{\char"03b8}}}$$でパラメータ化された関数$${p_\textbf{\textit{\char"03b8}}(\mathbf{x})}$$で表現しようとしたときの$${p_\textbf{\textit{\char"03b8}}(\mathbf{x})}$$のことを言います。
この$${p_\textbf{\textit{\char"03b8}}(\mathbf{x})}$$をどのように表すかという方法の一つにEnergy-Based Models（以下EBM）というものがあります。
EBMでは$${p_\textbf{\textit{\char"03b8}}(\mathbf{x})}$$を以下の式$${(1)}$$のように表します。
$$
\begin{align}
\tag{1}
p_\textbf{\textit{\char"03b8}}(\mathbf{x})
= \frac{\exp(-E_\textbf{\textit{\char"03b8}}(\mathbf{x}))}{Z_\textbf{\textit{\char"03b8}}}
\end{align}
$$
確率分布の要求事項の一つに定義域で非負であることがありますが、EBMでは指数関数で使うことでそれを満たすことができています。
この指数部分にある$${E_\textbf{\textit{\char"03b8}}}$$をエネルギー関数（energy function）と呼び、EBMではこれをパラメータ系列$${\textbf{\textit{\char"03b8}}}$$で表現することを目指します。
また、分母の$${Z_\textbf{\textit{\char"03b8}}}$$は、分子の$${\exp(-E_\textbf{\textit{\char"03b8}}(\mathbf{x}))}$$を確率変数$${\mathbf{x}}$$で積分したもので、分配関数（partition function, normalizing constant）と呼ばれます。
これは確率分布のもう一つの要求事項である$${\int p_\textbf{\textit{\char"03b8}}(\mathbf{x}) \mathrm{d}\mathbf{x} = 1}$$を満たすための正規化定数となっています。
一般にこの分配関数は解析的に計算できないため、この計算に工夫が必要となります。
$$
\begin{align}
\tag{2}
Z_\textbf{\textit{\char"03b8}} = \int \exp(-E_\textbf{\textit{\char"03b8}}(\mathbf{x})) \mathrm{d}\mathbf{x}
\end{align}
$$
Score Matching (SM)
次に、このEBMが持つパラメータ系列$${\textbf{\textit{\char"03b8}}}$$をどのように学習するかを考えます。今回、EBMにはニューラルネットワークを使う想定で、パラメータ系列$${\textbf{\textit{\char"03b8}}}$$の学習には勾配法による逐次的な更新を使用します。従って、パラメータ系列の勾配を求めるための目的関数が必要です。
EBMのパラメータ系列$${\textbf{\textit{\char"03b8}}}$$を学習するための代表的な目的関数の一つにScore Matching （以下SM）を利用する方法があります。SMでは式$${(3)}$$で表される、２つの分布間の乖離度を測る指標（Fisher divergence）を目的関数として使用します。
$$
\begin{align}
\tag{3}
D_F(
p_\textrm{data}(\mathbf{x}) \parallel
p_{\textbf{\textit{\char"03b8}}}(\mathbf{x})
) =
\mathbb{E}_{p_\textrm{data}(\mathbf{x})}
\Big[ \frac{1}{2} \|
\nabla_\mathbf{x} \log p_\textrm{data}(\mathbf{x}) -
\nabla_\mathbf{x} \log p_{\textbf{\textit{\char"03b8}}}(\mathbf{x})
\|^2 \Big]
\end{align}
$$
確率分布の自然対数を確率変数で微分して得られる勾配を（確率変数の）scoreと呼びます。
$$
\begin{align*}
\underbrace{\nabla_\mathbf{x} \log p(\mathbf{x})}_{\textit{score}} \\
\end{align*}
$$
SMの目的は、データの真の確率分布$${p_\textrm{data}(\mathbf{x})}$$とそのモデルによる確率分布$${p_{\textbf{\textit{\char"03b8}}}(\mathbf{x})}$$のscoreを近づけること、つまり、式$${(3)}$$を最小化するパラメータ系列$${\textbf{\textit{\char"03b8}}}$$を求めることになります。
$$
\begin{align}
\tag{4}
\min\limits_\textbf{\textit{\char"03b8}}
\mathbb{E}_{p_\textrm{data}(\mathbf{x})}
\Big[ \frac{1}{2} \|
\nabla_\mathbf{x} \log p_\textrm{data}(\mathbf{x}) -
\nabla_\mathbf{x} \log p_{\textbf{\textit{\char"03b8}}}(\mathbf{x})
\|^2 \Big]
\end{align}
$$
EBMをモデルに用いた場合にscoreを使うメリットとして、分配関数$${Z_\textbf{\textit{\char"03b8}}}$$が確率変数$${\mathbf{x}}$$に寄らないため$${Z_\textbf{\textit{\char"03b8}}}$$の計算を避けることができます（以下の式参照）。
$$
\begin{align*}
\nabla_\mathbf{x} \log p_\textbf{\textit{\char"03b8}}(\mathbf{x}) =
- \nabla_\mathbf{x} E_\textbf{\textit{\char"03b8}} (\mathbf{x})
- \underbrace{\nabla_\mathbf{x} \log Z_\textbf{\textit{\char"03b8}}}_{=0} =
- \nabla_\mathbf{x} E_\textbf{\textit{\char"03b8}} (\mathbf{x})
\end{align*}
$$
Denoising Score Matching (DSM)
この時点で、SMの目的である式$${(4)}$$には計算できない未知の項として$${\nabla_\mathbf{x} \log p_\textrm{data}(\mathbf{x})}$$があります。この未知の項を取り除く工夫は色々提案されていますが、今回はDenoising Score Matching（以下DSM）を紹介します。
まずはじめに、予め決められたノイズ分布$${q(\~{\mathbf{x}}|\mathbf{x})}$$によってノイズを与えられたデータ$${\~{\mathbf{x}}}$$を考えます。DSMとは、このノイズを与えられたデータ$${\~{\mathbf{x}}}$$に対する周辺確率分布$${q(\~{\mathbf{x}})=\int q(\~{\mathbf{x}}|\mathbf{x})p_\textrm{data}(\mathbf{x})\mathrm{d}\mathbf{x}}$$とそれをモデル化した確率分布$${p_{\textbf{\textit{\char"03b8}}}(\~{\mathbf{x}})}$$におけるSMを行う、というものです。
ここで、目的関数である$${q(\~{x})}$$と$${p_\textbf{\textit{\char"03b8}} (\~{\mathbf{x}})}$$のFisher divergenceは以下の式$${(5)}$$のように書くことができます（詳しい導出は参考文献[2]のAppendixまたは参考文献[5]の第13回講義資料を参照して下さい）。
$$
\begin{align*}
& &
D_F(
q(\~{\mathbf{x}}) \parallel
p_\textbf{\textit{\char"03b8}} (\~{\mathbf{x}})
)
&=
\mathbb{E}_{q(\~{\mathbf{x}})}
\Big[ \frac{1}{2} \|
\nabla_{\~{\mathbf{x}}} \log q(\~{\mathbf{x}}) -
\nabla_{\~{\mathbf{x}}} \log p_\textbf{\textit{\char"03b8}} (\~{\mathbf{x}})
\|^2 \Big]
\\
& &
&=
\mathbb{E}_{
\underbrace{q(\~{\mathbf{x}}|\mathbf{x})p_\textrm{data}(\mathbf{x})}_{q(\~{\mathbf{x}},\mathbf{x})}
}
\Big[ \frac{1}{2} \|
\nabla_{\~{\mathbf{x}}} \log q(\~{\mathbf{x}}|\mathbf{x}) -
\nabla_{\~{\mathbf{x}}} \log p_\textbf{\textit{\char"03b8}} (\~{\mathbf{x}})
\|^2 \Big] + \textrm{constant}
& & (5)
\end{align*}
$$
ここで、より具体的に、共分散行列が$${\sigma^2 \mathbf{I}}$$のガウシアンノイズを与えるノイズとして考えます。従って、ノイズを与えられたデータ$${\~{\mathbf{x}}}$$は以下のようになります。
$$
\begin{align*}
\~{\mathbf{x}} = \mathbf{x} + \textbf{\textit{\char"03b5}}, \quad
\textbf{\textit{\char"03b5}} \sim \mathcal{N}(\mathbf{0}, \sigma^2 \mathbf{I})
\end{align*}
$$
上記の式より、$${\mathbf{x}}$$が与えられたときの$${\~{\mathbf{x}}}$$の事後確率分布（＝ノイズ分布）は、中心を$${\mathbf{x}}$$、共分散行列を$${\sigma^2 \mathbf{I}}$$に持つガウス分布となることがわかるので、ノイズ分布はパラメータ$${\sigma}$$を用いて以下のように書けます。
$$
\begin{align}
\tag{6}
q_\sigma (\~{\mathbf{x}}|\mathbf{x}) =
\mathcal{N}(\~{\mathbf{x}}; \mathbf{x}, \sigma^2 \mathbf{I}) =
\frac{1}{(2\pi)^{d/2}\sigma^d}
\exp (- \frac{1}{2\sigma^2} \| \~{\mathbf{x}} - \mathbf{x} \|^2 )
\end{align}
$$
$$
\begin{align}
\tag{7}
\therefore \quad
\nabla_{\~{\mathbf{x}}} \log q_\sigma(\~{\mathbf{x}} | \mathbf{x}) =
\frac{\mathbf{x} - \~{\mathbf{x}}}{\sigma^2}
\end{align}
$$
式$${(6), (7)}$$より、DSMの目的は以下の式で表すことができます。
$$
\begin{align}
\tag{8}
\min\limits_\textbf{\textit{\char"03b8}}
\mathbb{E}_{p_\textrm{data}(\mathbf{x})q(\~{\mathbf{x}}|\mathbf{x})}
\Big[ \frac{1}{2} \Big\|
\underbrace{\nabla_{\~{\mathbf{x}}} \log p_\textbf{\textit{\char"03b8}} (\~{\mathbf{x}})}_{\textrm{score}} -
\frac{\mathbf{x} - \~{\mathbf{x}}}{\sigma^2}
\Big\|^2 \Big]
\end{align}
$$
$${\frac{1}{\sigma^2}(\mathbf{x}-\~{\mathbf{x}})}$$で表される方向は、ノイズを加えられたデータ$${\~{\mathbf{x}}}$$から元のデータ$${\mathbf{x}}$$に戻す方向であり、式$${(8)}$$は「ノイズを除去する方向としてscoreを学習」していると捉えることができます。
ここで挙げたノイズ分布$${q_\sigma (\~{\mathbf{x}}|\mathbf{x})}$$はカーネル密度推定などで使われるガウシアンカーネルの形をしています。直感的には、観測データに対して局所的にカーネル密度推定を行い計算可能な教師信号を用意していると捉えることもできます。
ただし、前述の通り、DSMではあくまでもノイズを与えられたデータ$${\~{\mathbf{x}}}$$に対する周辺確率分布$${q_\sigma(\~{\mathbf{x}})=\int q_\sigma(\~{\mathbf{x}}|\mathbf{x})p_\textrm{data}(\mathbf{x})\mathrm{d}\mathbf{x}}$$とそれをモデル化した確率分布$${p_{\textbf{\textit{\char"03b8}}}(\~{\mathbf{x}})}$$でSMを行うため、$${q_\sigma({\mathbf{x}}) \approx p_\textrm{data}(\mathbf{x})}$$とみなせる場合のみ、本来意図した学習が行われることに注意が必要です。
Langevin dynamics
続いて、EBMとしてモデル化した確率分布$${p_{\textbf{\textit{\char"03b8}}}(\mathbf{x})}$$からデータをサンプリングする方法を考えます。
ここでは、scoreのみを用いて確率分布$${p(\mathbf{x})}$$からサンプリングすることができる、Langevin dynamicsと呼ばれる方法を紹介します。
Langevin dynamicsでは、固定されたステップサイズ$${\epsilon > 0}$$とサンプリングが容易な事前分布$${\pi}$$（一様分布など）からサンプリングされた初期値$${\~{\mathbf{x}} \sim \pi(\mathbf{x})}$$を用いて、以下の更新式を$${T}$$回繰り返し適用することでサンプリングを行います。
$$
\begin{align}
\tag{9}
\~\mathbf{x}_t = \~\mathbf{x}_{t-1} +
\frac{\epsilon}{2} \nabla_{\~\mathbf{x}} \log p(\~\mathbf{x}_{t-1}) +
\sqrt{\epsilon} \mathbf{z}_t
\end{align}
$$
上記の式で$${\mathbf{z}_t \sim \mathcal{N}(0, \mathbf{I})}$$であり、$${\epsilon \rightarrow 0}$$かつ$${T \rightarrow \infty}$$のとき$${\~\mathbf{x}_T}$$の分布は$${p(\mathbf{x})}$$と等しくなり、いくつかの制限下で$${\~\mathbf{x}_T}$$は$${p(\mathbf{x})}$$からの正確なサンプルとみなせます（詳しくは僕から説明することができないので参考文献[3]のreference[62]を参照して下さい）。実装上は、ステップサイズ$${\epsilon}$$を小さく、更新回数$${T}$$を大きくすることでこの条件を代替します。
Score-based Generative Modeling
SMを用いた学習（＝式$${(4)}$$）とLangevin dynamicsを用いたサンプリング（＝式$${(9)}$$）では、どちらもモデリングしようとしている確率分布$${p(\mathbf{x})}$$はscoreの形$${\nabla_{\mathbf{x}} \log p(\mathbf{x})}$$で使用されるため、あらためてscoreをモデリングしようという、$${\mathbf{s}_\textbf{\textit{\char"03b8}}(\mathbf{x}) \approx \nabla_{\mathbf{x}} \log p(\mathbf{x})}$$のようなパラメータ化で置き換えることができます。
このような生成モデルのフレームワークをこの論文ではscore-based generative modelingと定義しています。
おわりに
今回は論文を理解するための予備知識とscore-based generative modelingの定義までをブログにまとめました。後半では、score-based generative modelingが抱える課題と論文での提案をまとめます。
ブログにしては冗長な表現が多く図も少なくて読みにくいかもしれないですが今後改善していけるように（そして継続できるように…）頑張ります！！
著者ら発信の講義資料など含め理解の助けになる素晴らしいリソースがネット上には多く存在するので一部を参考文献で共有します。
参考文献
How to Train Your Energy-Based Models [PDF]
Song, Y. and Kingma, D.P., 2021. arXiv preprint arXiv:2101.03288.Energy-Based Modelのチュートリアル論文です。EBMの理解の参考にしました。
A connection between score matching and denoising autoencoders [PDF]
Vincent, P., 2011. Neural computation, Vol 23(7), pp. 1661--1674. MIT Press.Denoising Score Matchingを提案した論文です。
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.今回取り上げたNoise Conditional Score Networkを提案した論文です。
Yong Song, "Generative Modeling by Estimating Gradients of the Data Distribution", 5 May. 2021, https://yang-song.net/blog/2021/score/
NCSNの主著者がNCSNから後続研究までをカジュアルに解説したブログです。
Stanford CS236, "Deep Generative Models", 2021, https://deepgenerativemodels.github.io/
NCSNの著者らによる講義スライドです。深層生成モデル全般の理解の参考にしました。
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.NCSN v2として、NCSNに関する様々な改善を行った論文です。
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.Sliced Score Matchingを提案した論文です。
この記事が気に入ったらサポートをしてみませんか？