Energy-Based Model：ノイズを復元して生成モデルを学習する

本記事はロボットとAI技術の専門誌『日経Robotics』のデジタル版です

［画像のクリックで別ページへ］

　深層学習を使ってデータを生成する深層生成モデルとして、変分自己符号化器（VAE）、敵対的生成モデル（GAN）、自己回帰モデル、正規化フローなどが知られている。これらのモデルは従来のモデルには難しかった高次元（数万〜数百万次元）の自然画像や音声データ生成に成功しており、生成されたサンプルは本物と見分けがつかないほどに高忠実な生成が可能となっている。

著者の岡野原大輔氏

　これに対し、別の生成モデルとしてエネルギーベースモデル（Energy-Based Model、以下EBM）が近年注目されている。深層学習研究の中心人物の一人であり2018年にチューリング賞を受賞したYann LeCun教授もICLR 2020の招待講演^1）の中で次の重要なモデルはEBMであると話している。

　EBMは他の生成モデルにない優れた特性を持つ。確率モデルを扱う際に不可欠な分割関数（全てのデータを列挙し積分をとる操作）をサンプリング時に扱う必要がなく、複数の候補解を列挙でき、また、複数のモデルを組み合わせることもできる。しかし、EBMの学習は難しく、生成モデルの実用的な性能は他の生成モデルに劣っていた。

　しかし近年の改良の結果、他の生成モデルによる最先端の結果に匹敵する品質の画像生成や音声生成を実現できるようになってきている。

　本稿ではEBMおよび、近年の高次元データの学習を成功させた学習手法としてデノイジングスコアマッチングや強いつながりを持つ拡散確率モデルについて紹介していく。

エネルギーベースモデル（EBM）

　EBMは入力を受け取ると、スカラー値であるエネルギーを返すエネルギー関数$F(x)$を使ってデータをモデル化する。この$F(x)$はデータセット中のデータ上では小さな値をとり、そうでない場合は大きな値をとるように学習するのが目標である。EBMはエネルギー関数から導出されるギブス分布を使って確率モデルに変換することができる。

\[ p(x) = \frac{\exp(-\beta F(x))}{\int_{x'} \exp (-\beta F(x'))} \]

　この$\beta$は逆温度パラメータとよばれ（$1/\beta$が温度に対応）、$\beta \rightarrow 0$であれば一様分布、 $\beta \rightarrow \infty$であればエネルギーが一番低いところだけ大きな確率を持つような分布となる。この確率モデルを明示的に扱うことはなく、学習やサンプリングは分割関数（分母の正規化項）を扱わずに処理できるように工夫する。そのため、EBMのエネルギー関数は、非正規化密度関数の一種とみなすこともできる。

　このモデルからデータをサンプリングするには尤度比$p(x)/p(x')$を使ったマルコフ連鎖モンテカルロ法（MCMC） を使ことが一般的であった。尤度比は、正規化項がキャンセルされ、扱う必要がないことに注意してほしい。しかし高次元データで多峰性があるデータではMCMCによるサンプリングは品質が低いことが分かっている。そこで本稿で扱う手法では対数尤度の入力についての勾配であるスコア関数を使ってサンプリングしている。これについては後で述べる。

　次に、このエネルギー関数をパラメータ$\theta$を使ったニューラルネットワークなどでモデル化$F_{\theta}(x)$し、訓練データ$D= \{x_i\}_{i=1}^N$を使って学習することをみていこう。さきほどのギブス分布の対数尤度を最大化する最尤推定でパラメータを推定する場合、負の対数尤度$\mathcal{L}$のパラメータ$\theta$についての勾配が求まればよい。この勾配は

\[ \frac{\partial \mathcal{L}}{\partial \theta} = \mathbb{E}_{x \sim D} \left [ \frac{\partial F_{\theta}(x)}{\partial \theta} \right ] - \mathbb{E}_{x' \sim p_{\theta}(x)} \left [ \frac{\partial F_{\theta}(x')}{\partial \theta} \right ] \]

と求められる（$\beta$は省略）。つまり、勾配は、データ分布上の勾配の期待値と、モデル分布上の勾配の期待値の差として表せられる。

　対称性があり、きれいな形をしているが、第二項のモデル分布からのサンプリングが難しいため高次元データをこの勾配を計算して学習することは難しく、別の学習手法が提案されていた。本稿ではデノイジングスコアマッチングとよばれる手法について扱っていく。

デノイジングスコアマッチング

　対数尤度$\log p(x)$を入力について勾配をとった関数$\nabla_x \log p(x)$をスコア関数とよぶ。スコア関数は正規化項が含まれない（正規化項は入力依存項がないので、入力について勾配は$0$）。データ分布のスコア関数とモデル分布のスコア関数を一致させることできれば分布を学習できる。また、スコア関数が得られれば、ランジュバンサンプリングとよばれる方法を使って高次元でも効率的にサンプリングすることができる。これはエネルギー関数の勾配と少しのガウシアンノイズに従ってたどっていきサンプリングする手法である。

　観測データ$x$にガウシアンノイズを加えたデータ$q_{\sigma}(\tilde{x}|x) =\mathcal{N}(\tilde{x}|x, \sigma^2 I)$を元のデータにデノイジングする方向をノイズについて期待値をとった値はスコア関数と一致することが分かっている。これを使ってスコア関数を学習する手法をデノイジングスコアマッチング（DSM：Denoising Score Matching）とよぶ。具体的には次の目的関数を最小化するような関数$s_{\theta}(\tilde{x},\sigma)$を求めることで、スコア関数を学習する。

\[ \frac{1}{2}\mathbb{E}_{q_{\sigma}(\tilde{x}|x)p(x)p(\sigma)} \left [||\sigma s_{\theta}(\tilde{x}, \sigma) + \frac{\tilde{x} -x}{\sigma}||_2^2 \right ] \]

$\sigma$はノイズの大きさを示し、学習中にスケジューリングし調整する。様々な改良を加えることで、この手法を使って学習された生成モデルはGANなどに匹敵する性能を持つと報告されている^2-3）。

拡散確率モデル

　拡散確率モデルは、データにノイズを徐々に加えていき全ての情報が失われて完全なノイズになるマルコフ過程を逆向きに辿ることで生成モデルを学習する。各拡散ステップが小さい量のガウシアンノイズである場合、その逆向きの生成に対応する遷移も条件付けされたガウシアンからのサンプリングでモデル化できることが知られている。このため、生成モデルはノイズからデータ分布を生成できるような条件付けガウシアンを使って学習できる。

　拡散確率モデルの学習は対数尤度の変分下限を最大化することで実現されるが、その目的関数はさきほどのDSMとほとんど同じ式になることが分かっており、また高品質の画像を生成できると示された^4）。拡散確率モデルは表現力が高く、自己回帰モデルなどと違って並列にサンプリング操作ができ、推論時に計算コストと生成品質のトレードオフをとることができる特徴を持つ。スペクトログラムからの波形生成でも自己回帰モデルを使った手法に比べて高速に計算可能であり、リアルタイム生成を可能としている。

EBMの今後

　EBMの大きな特徴は複数の候補解を列挙できること（特に連続値の場合）、複数のEBMを組み合わせて（例えば$\alpha F(x) +\beta F'(x)$）別のEBMを構成できること、複数変数間の双方向的な推論ができること（現在の候補出力結果に応じて入力を補正するなど）などがある。また生成モデルだけでなく、スコア関数は入力の欠損値を復元したり、異常値を発見したりと別の有用性がある。

本記事はロボットとAI技術の専門誌『日経Robotics』のデジタル版です

［画像のクリックで別ページへ］

　また最近は複数の生成モデルを組み合わせてモデル化することも多くなってきた。EBMもその長所を生かした形で別の生成モデルと組み合わせて使うことも今後増えると考えられる。

岡野原 大輔（おかのはら・だいすけ）
Preferred Networks 代表取締役 最高執行責任者

2006年にPreferred Infrastructureを共同創業。2010年、東京大学大学院博士課程修了。博士（情報理工学）。未踏ソフト創造事業スーパークリエータ認定。東京大学総長賞。