原論文から解き明かす生成AI

著:菊田 遥平

Amazon

楽天市場

ポチップ

どんな内容の本か？

ChatGPTをはじめとしたいろいろな生成AIが出てきている中、どのような仕組みで動いているか？をみなさんどの程度理解しているでしょうか？
私自身、生成AIはいろいろな要素が組み合わさっているため、1つ１つは知っているけど、全部を網羅的に理解するというのはかなり難しいと感じています。しかも、発展のスピードが速く関連論文が毎日のように数十本公開されたりしています。
このような現状を受けて、この本では重要な原論文をベースに生成AIの基礎を紹介し、かつ、論文を独力で読み解く技術を解説し、自分で論文を読む際の役に立つ技術を紹介している本になります。
生成AIといってもいろいろありますが、この本ではChatGPTをはじめとするテキストを中心に画像生成にも触れられている本になっています。
このため、テキストと画像に関する基礎部分は大体抑えられる印象です。

どんな人にお勧めか？

以下のような人にお勧めな本かと思っています。

1. 生成AIを勉強したいけど何から手を出したらいいかわかってない人
2. 論文の読み方について他の人がどうやっているかを知りたい人

各論文の解説記事はいろいろwebに公開されている一方、資料がバラバラなのと、生成AIの知識が浅い人にとってはどの論文が重要なのか調べるのも大変な印象です。このため、重要な論文だけに焦点を当ててくれているこの本をまず読むのは最初の手としては良いと感じました。
また、論文の読み方は研究室などで他の人のやり方をたまに聞く程度で、他の人がどのように読んでいるかを知る機会はそれほど多くないと思っています。
この本では最初に１章かけてこの論文の読み方を解説していて、私のように他の人がどう読んでいるかを知る機会がなかった人には良い参考情報になるのではないかと思っています。

個人的に良かった点

個人的には以下の点が良かったです。

1. 他の人の論文の読み方を知ることができた
2. 生成AIで抜けていた技術について知ることができた

1は先ほどお勧めの人で書いた通り、この本では論文の読み方を１章かけて紹介しており、他の人がどう読んでいるか知る機会がなかった自分にとっては非常に有用でした。
また、これに加えて、生成AIは要素が本当にいろいろあり、私自身、さらっとしか知らなかった部分がいくつかあるのですが、その部分についても深く知ることができてよい機会になりました。

終わりに

今回は献本でいただいた「原論文から解き明かす生成AI」の紹介記事でした。生成AIに関する本はいろいろありますが、深く理解したいと思っている人には非常に良い本だと感じました。
定期的にこのように本の紹介記事を書いてます。もし興味があれば他の本の記事も見ていただければ幸いです。

The post [書評] 原論文から解き明かす生成AI ー 話題のAIを深く理解したいと思ったら読む一冊 first appeared on まったり勉強ノート.

生成ＡＩのしくみ　〈流れ〉が画像・音声・動画をつくる (岩波科学ライブラリー)

岡野原 大輔

Amazon

楽天市場

Amazon Kindle

ポチップ

どんな内容の本か？

皆さん、生成AIのニュースが毎日のようにでている気がしますが画像や動画などのデータを生成する仕組みを簡単に説明することができるでしょうか？私は、自分の親にこれらの仕組みを説明できるほど分かりやすい説明はできないと思っていますが、この本はそんなAIについて詳しくない人向けに画像や動画などをどのようにAIが作っているか、技術の進化の歴史を交えて分かりやすく説明した本です。

この本の特徴としては普通の人が挫折せずに読めるように数式が一切出てこず、分かりにくい英語の専門用語も極力日本語訳をあてて、難しいところはイラストを使って分かりやすく説明されているのを感じました。
一方、これらの影響で私のような専門に近ければ近いほど、日本語訳に馴染みがなくて、これがどの論文の話をいっているのかぱっと分かりにくいところがあるうえに、数式が何もないこともあり、細かい部分が全くわからない本にもなっています。

このため、詳しく知りたい人は別の本、例えば同じ著者のこちらの本を読むほうが良いと思います。

どんな人にお勧めか？

おそらく以下のような人にお勧めかなと思っています。

1. 生成AIについて知りたいが数式が得意ではない人
2. 生成AIの発展の歴史を知りたい人

特にお勧めなのが1の人で、生成AIについて技術的な説明をしている本やすごい軽く説明している本はいくつかありますが、その中間くらいの数式は苦手だけどどのような仕組みで動いているのか知りたい人にはちょうどよい印象の本です。

また、これまで生成AIの取り組みはいろいろあったのですが、その発展の歴史を知りたいと思った人にもお勧めかと思います。
この本ではノーベル賞で話題になったホップフィールドネットワークの話題から生成AIへの流れを順番に説明されいて、ノーベル賞と生成AIがどうつながっているのかもわかるようになると思います。

個人的に良かった点

個人的に良かった点としては以下の通りです。

1. 生成AIの歴史について飛び飛びの部分を補完できた
2. 人に生成AIを説明するときどういう言い回しがいいのかの参考になった

特に1に関しては私も専門分野とはいえ、2017年くらいから深層学習の技術開発を始めた関係で、それよりも前の部分は人に説明できるような状態ではなかったのですが、そのあたりの歴史について知ることができたのは良かったです。

また、2に関して、生成AIは人に説明が難しいところが多いのですが、それをどういう風に説明すれば良いか少し分かった気がするので、その点も良かったです。

終わりに

今回は「生成ＡＩのしくみ　〈流れ〉が画像・音声・動画をつくる」という本を紹介しました。最初のほうにも書きましたが、この本は普通の人向けの本であり、技術的な詳細を知りたい方は例えば同じ著者の「拡散モデル データ生成技術の数理」や手を動かして勉強できる「ゼロから作るDeep Learning ❺ ―生成モデル編」が良いかと思っています。
「拡散モデル データ生成技術の数理」の方は紹介記事を書いたので参考にしてください。

定期的にこのような書評を書いているので、もし興味があれば他の本の記事も見ていただければ幸いです。

The post [書評] 生成ＡＩのしくみ ー 生成AIの進化の歴史が分かる一冊 first appeared on まったり勉強ノート.

今回の記事はスコアベースモデル (SBM)はすでに理解している前提で説明していきます。もしスコアベースモデルがよくわからないという方はこちらに簡単な解説を書いたので参考にしてください。

今回のコードは以下のところにあげてありますので、コード全体を見たい方はこちらをご覧ください。

https://github.com/shu65/diffusion-model-book/blob/main/diffusion_model_book_3_VE_SDE.ipynb

分散発散型確率微分方程式(VE-SDE)とは？

この本の3章の前半部分で、2章で紹介しているスコアベースモデル (SBM)とデノイジング拡散確率モデル (DDPM) を確率微分方程式 (SDE) とみなすことができるという説明をしています。このうち、SBMのほうをSDE表現してでてくるものが分散発散型確率微分方程式 (VE-SDE)です。

確率微分方程式(SDE)

確率微分方程式（SDE; Stochastic Differential Equations）は次の式で与えられます。

$$ \begin{align} \text{d}\boldsymbol{x} = \boldsymbol{f}(\boldsymbol{x}, t)\text{d}t + \boldsymbol{G}(\boldsymbol{x}, t)\text{d}\boldsymbol{w} \tag{3.1} \end{align} $$

この式において\(\text{d}\boldsymbol{x} \)は\(\boldsymbol{x}\)の変化量です。この変化量は決定的に変化量である\(\boldsymbol{f}(\boldsymbol{x}, t)\text{d}t\)とランダムに変化する量である\(\boldsymbol{G}(\boldsymbol{x}, t)\text{d}\boldsymbol{w}\)の和で構成されています。

ここで、\(\boldsymbol{w}\)は標準ウィーナー過程またはブラウン運動ともよばれ、\(\text{d}\boldsymbol{w}\)は微小時間間隔\(\tau\)において平均が0、分散が\(\tau\)の正規分布とみなすことができます。

この確率微分方程式において\(\boldsymbol{f}(\cdot, t)\)はドリフト係数、\(\boldsymbol{G}(\cdot, t)\)は拡散係数と呼びます。

ただし、一般に拡散モデルで扱う確率微分方程式以下のようにドリフト係数が時間のみに依存する関数\(\boldsymbol{f}(t)\)と入力\(\boldsymbol{x}\)の積、拡散係数は時間のみに依存してスカラ値を出力する\(g(t)\)を使った確率微分方程式が利用されます。

$$ \begin{align} \text{d}\boldsymbol{x} = f(t)\boldsymbol{x}\text{d}t +g(t)\text{d}\boldsymbol{w} \tag{3.2} \end{align} $$

スコアベースモデルの拡散過程をSDEで表現する

スコアベースモデル(SBM)の拡散過程は以下のようになっていました。

$$ \begin{align} q(\boldsymbol{x}_i | \boldsymbol{x}) = \mathcal{N}(\boldsymbol{x}, \sigma_i^2\boldsymbol{I}) \tag{3.3} \end{align} $$

ここで\(i = 0,…, N\)です。この場合の拡散過程の1ステップは次のようになります。

$$ \begin{align} q(\boldsymbol{x}_i | \boldsymbol{x}_{i-1}) = \mathcal{N}(\boldsymbol{x}_i;\boldsymbol{x}_{i-1}, (\sigma_i^2 – \sigma_{i-1}^2)\boldsymbol{I}) \tag{3.4}
\end{align} $$

式(3.3), (3.4)は2章のほうで説明されています。この拡散過程の1ステップは変数変換を使うと以下のようになります。

$$ \begin{align}
\boldsymbol{x}_i &= \boldsymbol{x}_{i-1} + \sqrt{\sigma_i^2 – \sigma_{i-1}^2}\boldsymbol{z}_{i-1} \tag{3.5} \\
\boldsymbol{z}_{i-1} &\sim \mathcal{N}(0, \boldsymbol{I}) \tag{3.6}
\end{align} $$

ここで簡略化のために\(\sigma_0 = 0\) として考えます。

ここから\(N \rightarrow \infty\) とした極限を考えていきます。この時、\(i\)の代わりに\(t\)を用いて、\({\boldsymbol{x}_i}_{i=1}^N\)を連続的な確率過程\({\boldsymbol{x}_t}_{t=0}^1\)とし、\(\sigma_i\)を関数\(\sigma(t)\)、\(\boldsymbol{z}_{i}\)は\(\boldsymbol{z}(t)\)とします。

また、\(\Delta t=1/N\)とし、\(t \in \left\{0, \frac{1}{N},…, \frac{N-1}{N} \right\}\)とします。

この時式(3.5)の式は以下のようになります。

$$ \begin{align} \boldsymbol{x}(t + \Delta t) = \boldsymbol{x}(t) + \sqrt{\sigma(t + \Delta t)^2 – \sigma(t)^2}\boldsymbol{z}_{i-1} \tag{3.7} \end{align} $$

ここで\(\sigma(t + \Delta t)^2 – \sigma(t)^2\)の部分で1次近似を利用して式変形します。1次近似は以下の近似を指します。

$$ \begin{align} f(x + \Delta x) \approx \frac{\text{d}f(x)}{\text{d}x} \Delta x + f(x) \tag{3.8} \end{align} $$

この1次近似の式において\(f(x)\)の部分を\(\sigma(t)^2\)として置き換えると以下のようになります。

$$ \begin{align} \sigma(t + \Delta t)^2 \approx \frac{\text{d}[\sigma(t)^2]}{\text{d}t} \Delta t + \sigma(t)^2 \tag{3.9} \end{align} $$

この式の両辺を\(\sigma(t)^2\)で引くと以下のようになります。

$$ \begin{align} \sigma(t + \Delta t)^2 – \sigma(t)^2 \approx \frac{\text{d}[\sigma(t)^2]}{\text{d}t} \Delta t \tag{3.10} \end{align} $$

この式(3.10)を式(3.7)に代入すると以下のようになります。

$$ \begin{align} \boldsymbol{x}(t + \Delta t) = \boldsymbol{x}(t) + \sqrt{\frac{\text{d}[\sigma(t)^2]}{\text{d}t} \Delta t}\boldsymbol{z}(t) \tag{3.11} \end{align} $$

このあとの説明のために以下のように少し式変形をします。

$$ \begin{align}
\boldsymbol{x}(t + \Delta t) &= \boldsymbol{x}(t) + \sqrt{\frac{\text{d}[\sigma(t)^2]}{\text{d}t} \Delta t}\boldsymbol{z}(t) \\
\boldsymbol{x}(t + \Delta t) – \boldsymbol{x}(t) &= \sqrt{\frac{\text{d}[\sigma(t)^2]}{\text{d}t} \Delta t}\boldsymbol{z}(t) \\
\boldsymbol{x}(t + \Delta t) – \boldsymbol{x}(t) &= \sqrt{\frac{\text{d}[\sigma(t)^2]}{\text{d}t}} \left(\sqrt{\Delta t}\boldsymbol{z}(t) \right) \tag{3.12}
\end{align} $$

さて、ここから\(\Delta t \to 0\) にしたときのことを考えます。式(3.12)の左辺のほうは以下のようになります。

$$ \begin{align}
\lim_{\Delta t \to 0} \boldsymbol{x}(t + \Delta t) – \boldsymbol{x}(t) = \text{d}\boldsymbol{x} \tag{3.13}
\end{align} $$

問題は右辺の\(\sqrt{\Delta t}\boldsymbol{z}(t) \)の部分です。これは結果的には以下のようになります。

$$ \begin{align}
\lim_{\Delta t \to 0} \sqrt{\Delta t}\boldsymbol{z}(t) = \text{d}\boldsymbol{w} \tag{3.14}
\end{align} $$

この部分ですが本も元論文[3] のほうにもこの式変形のところで言及がないのでわかりにくいので少し説明します。

まず、そもそも\(\text{d}\boldsymbol{w}\)は何であったか？ですが、これは最初に説明した通り標準ウィーナー過程またはブラウン運動ともよばれ、\(\text{d}\boldsymbol{w}\)は微小時間間隔\(\tau\)において平均が0、分散が\(\tau\)の正規分布とみなすことができます。このことから以下のように表すことができます。

$$ \begin{align}
\text{d}\boldsymbol{w} \sim \mathcal{N}(0, \tau \boldsymbol{I}) \tag{3.15}
\end{align} $$

ここで\(\boldsymbol{z}(t)\)は

$$ \begin{align}
\boldsymbol{z}(t) \sim \mathcal{N}(0, \boldsymbol{I}) \tag{3.16} \\
\end{align} $$

なので、\(\text{d}\boldsymbol{w}\)は以下のようになります。

$$ \begin{align}
\text{d}\boldsymbol{w} = \sqrt{\tau} \boldsymbol{z}(t) \tag{3.17}
\end{align} $$

\(\tau\)が微小時間間隔なので式(3.14)と式(3.17)を見比べるとなんとなく式(3.14)が成り立ちそうだなぁと思います。ただ、極限を素直に考えると以下のようになるのでは？とずっと思ってました。

$$ \begin{align}
\lim_{\Delta t \to 0} \sqrt{\Delta t}\boldsymbol{z}(t) = 0
\end{align} $$

この部分、私は気になってしょうがなかったので、少し調べました。結論からいうとこの部分の式変形に関してはウィーナー過程の条件から導出できそうだということがわかりました。詳しくは以下のサイトが分かりやすかったので、詳しく知りたい方はご覧ください。

http://takashiyoshino.random-walk.org/memo/keikaku_ensyu/node4.html

ここでは簡単に説明します。まずウィーナー過程 \(\boldsymbol{w}(t)\)を考えます。ウィナー過程の条件より以下が成り立ちます。

$$ \begin{align}
\boldsymbol{w}(t + \Delta t) – \boldsymbol{w}(t) \sim \mathcal{N}(0, \Delta t \boldsymbol{I}) \tag{3.18}
\end{align} $$

ここで式(3.18)を右辺を見ると平均０、分散\(\Delta t\)の正規分布です。このため、式(3.18)は左辺は以下のように表すこともできます。

$$ \begin{align}
\boldsymbol{w}(t + \Delta t) – \boldsymbol{w}(t) = \sqrt{\Delta t}\boldsymbol{z}(t) \tag{3.19} \\
\end{align} $$

この式(3.19)の右辺は式(3.14)の左辺の\(\lim_{\Delta t \to 0}\)の中と同じになります。また式(3.19)の左辺は\(\Delta t \to 0\)のとき以下のようになります。

$$ \begin{align}
\lim_{\Delta t \to 0} \left( \boldsymbol{w}(t + \Delta t) – \boldsymbol{w}(t) \right) &= \text{d}\boldsymbol{w} \tag{3.20}
\end{align} $$

よって式(3.14)は式(3.19)と(3.20)を使うと以下のようになります。

$$ \begin{align}
\lim_{\Delta t \to 0} \sqrt{\Delta t}\boldsymbol{z}(t) &= \lim_{\Delta t \to 0} \left( \boldsymbol{w}(t + \Delta t) – \boldsymbol{w}(t) \right) \\
&= \text{d}\boldsymbol{w} \tag{3.21}
\end{align} $$

この式変形なら個人的には納得できました。よって最終的に式(3.12)で\(\Delta t \to 0\) を考えると式(3.13)と式(3.21)より以下のようになります。

$$ \begin{align}
\text{d}\boldsymbol{x} &= \lim_{\Delta t \to 0} \boldsymbol{x}(t + \Delta t) – \boldsymbol{x}(t) \\
&= \lim_{\Delta t \to 0} \sqrt{\frac{\text{d}[\sigma(t)^2]}{\text{d}t}} \left(\sqrt{\Delta t}\boldsymbol{z}(t) \right) \\
&= \sqrt{\frac{\text{d}[\sigma(t)^2]}{\text{d}t}} \boldsymbol{w}(t) \tag{3.22}
\end{align} $$

この式(3.22)を見るとドリフト係数\(f(t)\) と拡散係数\(g(t)\)が以下のようなSDEであることが分かります。

$$ \begin{align*}
f(t) &= 0 \tag{3.23} \\
g(t) &= \sqrt{\frac{\text{d}[\sigma(t)^2]}{\text{d}t}} \tag{3.24} \\
\end{align*} $$

これでSBMをSDEで表現することができました。このSBMの式から導出したSDEを分散発散型確率微分方程式 (VE-SDE)と呼びます。

VE-SDEの学習

VE-SDEの各時刻\(t\)のスコアを学習するあために、次の条件付き確率（拡散カーネル）を知る必要があります。

$$ \begin{align*}
p_{0t}(\boldsymbol{x}(t)|\boldsymbol{x}(0)) \tag{3.25}
\end{align*} $$

ここで\(p_{0t}\)は\(\boldsymbol{x}(0)\)を条件付けしたときの\(\boldsymbol{x}(t)\)の確率を表しています。

ここでSDEが以下の形として考えていきます。

$$ \begin{align} \text{d}\boldsymbol{x} = f(t)\boldsymbol{x}\text{d}t + g(t)\text{d}\boldsymbol{w} \end{align} \tag{3.26}$$

この場合、式(3.26)の条件付き確率は以下のような正規分布で表すことができます[1, 2]。

$$ \begin{align}
p_{0t}(\boldsymbol{x}(t)|\boldsymbol{x}(0)) =& \mathcal{N}(s(t)\boldsymbol{x}(0), s(t)^2\sigma^{\prime}(t)^2\boldsymbol{I}) \tag{3.27} \\
s(t) =& \text{exp}\left(\int_0^tf(\xi)\text{d}\xi\right) \tag{3.28} \\
\sigma^{\prime}(t) =& \sqrt{\int_0^t \frac{g(\xi)^2}{s(\xi)^2}\text{d}\xi} \tag{3.29} \\
\end{align} $$

本のほうでは式(3.27)と式(3.29) の\(\sigma^{\prime}(t)\)の部分は\(\sigma(t)\)という表記になっています。ただ、VE-SDEのほうにも\(\sigma(t)\)があって区別ができないので、この記事では式(3.27)と(3.29)に登場する\(\sigma(t)\)を\(\sigma^{\prime}(t)\)として説明していきます。

VE-SDEの場合はこの式を使うと簡単に\(p_{0t}(\boldsymbol{x}(t)|\boldsymbol{x}(0))\)の形がわかるので、以下に示していきます。

まず、\(s(t)\)の部分ですが、VE-SDEの場合、式(3.23)から以下のようになります。

$$ \begin{align}
s(t) &= \text{exp}\left(\int_0^tf(\xi)\text{d}\xi\right) \\
&= \text{exp}\left(\int_0^t 0 \text{d}\xi\right) \\
&= \text{exp}\left(0 \right) \\
&= 1 \tag{3.30} \\
\end{align} $$

次に\(\sigma^{\prime}(t)\)に関してです。まず式(3.26)を使って式変形します。

$$ \begin{align}
\sigma^{\prime}(t) &= \sqrt{\int_0^t \frac{g(\xi)^2}{s(\xi)^2}\text{d}\xi} \\
&= \sqrt{\int_0^t \frac{g(\xi)^2}{1^2}\text{d}\xi} \\
&= \sqrt{\int_0^t g(\xi)^2\text{d}\xi} \tag{3.31}
\end{align} $$

ここでVE-SDEの\(g(t)\)は式(3.24)で分かっているのでこれを利用してさらに式変形します。

$$ \begin{align}
\sigma^{\prime}(t) &= \sqrt{\int_0^t g(\xi)^2\text{d}\xi} \\
&= \sqrt{\int_0^t \left( \sqrt{\frac{\text{d}[\sigma(\xi)^2]}{\text{d}\xi}} \right)^2\text{d}\xi } \\
&= \sqrt{\int_0^t \frac{\text{d}[\sigma(\xi)^2]}{\text{d}\xi} \text{d}\xi } \\
&= \sqrt{\sigma(t)^2 – \sigma(0)^2} \tag{3.32}
\end{align} $$

式変形した式(3.30)、(3.32)を式(3.27)に代入すると最終的には以下のようになります。

$$ \begin{align}
p_{0t}(\boldsymbol{x}(t)|\boldsymbol{x}(0)) &= \mathcal{N}(s(t)\boldsymbol{x}(0), s(t)^2\sigma^{\prime}(t)^2\boldsymbol{I}) \\
&= \mathcal{N}(\boldsymbol{x}(0), \left[\sigma(t)^2 – \sigma(0)^2\right]\boldsymbol{I}) \tag{3.33}
\end{align} $$

これによりVE-SDEの拡散過程の条件付き確率の式がわかりました。

本の説明では\(\sigma(t)\)が具体的にどのような式を使うのかまでは示してないため、式変形はここまでになっています。

一方、このブログではコードに落とすところまでをやるため、ここからさらに式変形していきます。ここから元論文の[3]を参考にして式変形していきます。

[3]の論文で使われている\(\sigma(t)\)と同じものを用いて説明していきます。[3]では以下のものが使われています。

$$ \begin{align}
\sigma(t) &= \sigma_{min}\left( \frac{\sigma_{max}}{\sigma_{min}} \right)^t, & \ t &\in (0, 1] \\
\sigma(0) &= 0, & \ t &= 0 \\
\tag{3.34}
\end{align} $$

ここで\(\sigma_{min}\)と\(\sigma_{max}\)はハイパーパラメータです。

これを使って式(3.24)の\(g(t)\)と式(3.33)の条件付き確率の式変形をしていきます。

まず、式(3.24)の\(g(t)\)に関してです。

$$ \begin{align*}
g(t) &= \sqrt{\frac{\text{d}[\sigma(t)^2]}{\text{d}t}} \\
&= \sqrt{\frac{\text{d}}{\text{d}t} \left( \sigma_{min}\left( \frac{\sigma_{max}}{\sigma_{min}} \right)^t \right)^2} \\
&= \sqrt{\frac{\text{d}}{\text{d}t} \sigma_{min}^2\left( \frac{\sigma_{max}}{\sigma_{min}} \right)^{2t} } \\
&= \sqrt{\sigma_{min}^2 \frac{\text{d}}{\text{d}t} \left( \frac{\sigma_{max}}{\sigma_{min}} \right)^{2t} } \tag{3.35}
\end{align*} $$

ここで\( \frac{\text{d}}{\text{d}t} \left( \frac{\sigma_{max}}{\sigma_{min}} \right)^{2t} \)の部分に注目します。以下のような指数関数の微分公式を利用します。

$$ \begin{align*}
\frac{\text{d}}{\text{d}x} a^x = a^x \log a \tag{3.36}
\end{align*} $$

(参考：https://manabitimes.jp/math/1112)

この公式を利用すると以下のようになります。

$$ \begin{align*}
\frac{\text{d}}{\text{d}t} \left(\frac{\sigma_{max}}{\sigma_{min}} \right)^{2t} &= \left( \frac{\sigma_{max}}{\sigma_{min}} \right)^{2t} \log \left(\frac{\sigma_{max}}{\sigma_{min}} \right)^2
\tag{3.37}
\end{align*} $$

この式(3.37)を式(3.35)に代入して式変形していくと以下のようになります。

$$ \begin{align*}
g(t) &= \sqrt{\sigma_{min}^2 \frac{\text{d}}{\text{d}t} \left( \frac{\sigma_{max}}{\sigma_{min}} \right)^{2t} } \\
&= \sqrt{\sigma_{min}^2 \left( \frac{\sigma_{max}}{\sigma_{min}} \right)^{2t} \log \left(\frac{\sigma_{max}}{\sigma_{min}} \right)^2 } \\
&= \sigma_{min} \left( \frac{\sigma_{max}}{\sigma_{min}} \right)^{t} \sqrt{\log \left(\frac{\sigma_{max}}{\sigma_{min}} \right)^2 } \\
&= \sigma_{min} \left( \frac{\sigma_{max}}{\sigma_{min}} \right)^{t} \sqrt{2 \log \left(\frac{\sigma_{max}}{\sigma_{min}} \right)} \tag{3.38}
\end{align*} $$

次に式(3.33)の条件付き確率のほうを式変形していきます。この式には分散のほうにだけ\(\sigma(t)\)が登場するので、この部分だけ注目します。この分散に式(3.34)の\(\sigma(t)\)を代入して式変形していくと以下のようになります。

$$ \begin{align}
\sigma(t)^2 – \sigma(0)^2 &= \left[\sigma_{min}\left( \frac{\sigma_{max}}{\sigma_{min}} \right)^t \right]^2 – 0^2 \\
&= \sigma_{min}^2\left(\frac{\sigma_{max}}{\sigma_{min}} \right)^{2t} \\ \tag{3.39}
\end{align} $$

よって式(3.33)の条件付き確率は以下のようになります。

$$ \begin{align}
p_{0t}(\boldsymbol{x}(t)|\boldsymbol{x}(0)) &= \mathcal{N}(\boldsymbol{x}(0), \left[\sigma(t)^2 – \sigma(0)^2\right]\boldsymbol{I}) \\
&= \mathcal{N}\left(\boldsymbol{x}(0), \sigma_{min}^2\left(\frac{\sigma_{max}}{\sigma_{min}} \right)^{2t}\boldsymbol{I}\right) \tag{3.40}
\end{align} $$

これらを用いてデノイジングスコアマッチングをロス関数としてスコア関数\(s_{\theta}\)を学習します。VE-SDEの場合のデノイジングスコアマッチングの関数はSBMのときと同じ形になります。具体的には以下のようになります。（変数はVE-SDEに合わせています。）

$$ \begin{align}
L(\theta) :=&
E_t \left[ \lambda(t) E_{\boldsymbol{x}(0) \sim p_{data}(\boldsymbol{x}),\boldsymbol{x}(t) \sim p_{0t}(\boldsymbol{x}(t)|\boldsymbol{x}(0))} \left\{ \right. \right. \\
& \quad \left. \left. \left| \nabla_{\boldsymbol{x}(t)} \log p_{0t}(\boldsymbol{x}(t)|\boldsymbol{x}(0)) – s_{\theta}(\boldsymbol{x}(t), t) \right|^2 \right\} \right] \tag{3.41}
\end{align} $$

ここで、\(\lambda(t)\)は各\(t\)における重みづけです。

これを実装するために、SBMのときと同じようにスコア \( \nabla_{\boldsymbol{x}(t)} \log p_{0t}(\boldsymbol{x}(t)|\boldsymbol{x}(0)) \)の部分を式変形します。これはSBMのときと同じなので本の２章と以前私が書いたSBMの解説の記事をご覧ください。

結果として以下のようになります。

$$ \begin{align}
\nabla_{\boldsymbol{x}(t)} \log p_{0t}(\boldsymbol{x}(t)|\boldsymbol{x}(0)) &= \frac{-\epsilon}{\sigma_{min}^2\left(\frac{\sigma_{max}}{\sigma_{min}} \right)^{2t}} \tag{3.42} \\
\epsilon &\sim \mathcal{N}\left(0, \sigma_{min}^2\left(\frac{\sigma_{max}}{\sigma_{min}} \right)^{2t}\boldsymbol{I}\right) \tag{3.43}
\end{align} $$

式(3.41)を式(3.42)、(3.43)を使って変形すると以下のようになります。（式が長すぎるので\(\boldsymbol{x}(0), \boldsymbol{x}(t), \epsilon\)の分布を省略してます。）

$$ \begin{align}
L(\theta) :=&
E_t \left[ \lambda(t) E_{\boldsymbol{x}(0),\boldsymbol{x}(t)} \left\{ \left| \nabla_{\boldsymbol{x}(t)} \log p_{0t}(\boldsymbol{x}(t)|\boldsymbol{x}(0)) – s_{\theta}(\boldsymbol{x}(t), t) \right|^2 \right\} \right] \\
=& E_t \left[ \lambda(t) E_{\boldsymbol{x}(0),\epsilon} \left\{ \left| \frac{-\epsilon}{\sigma_{min}^2\left(\frac{\sigma_{max}}{\sigma_{min}} \right)^{2t}} – s_{\theta}(\boldsymbol{x}(t), t) \right|^2 \right\} \right] \tag{3.44}
\end{align} $$

これをPyTorchを使ってコードにすると以下のようになります。

def sigma(t, sigma_min=sigma_min, sigma_max=sigma_max):
  return sigma_min * (sigma_max / sigma_min) ** t 

def ve_sde_marginal_prob_statistics(x, t, sigma_min, sigma_max):
  mean = x
  std = sigma(t=t, sigma_min=sigma_min, sigma_max=sigma_max)
  return mean, std

def ve_sde_drift(t, sigma_min, sigma_max):
  drift = torch.zeros_like(t)
  return drift

def ve_sde_diffusion(t, sigma_min, sigma_max):
  std = sigma(t=t, sigma_min=sigma_min, sigma_max=sigma_max)
  diffusion = std * torch.sqrt(2 * (torch.log(sigma_max) - torch.log(sigma_min))) # (30)
  return diffusion

def dsm_loss(score_model, samples, sigma_min, sigma_max):
  eps = 1.0e-8
  t = torch.distributions.uniform.Uniform(torch.tensor([eps], device=samples.device), torch.tensor([1], device=samples.device)).sample([samples.shape[0]]) 
  z = torch.randn_like(samples)
  mean, std = ve_sde_marginal_prob_statistics(x=samples, t=t, sigma_min=sigma_min, sigma_max=sigma_max)
  noise = z * std
  perturbed_samples = mean + z * std
  scores = score_model(perturbed_samples, t)
  target = - 1 / (std ** 2) * noise
  
  target = target.view(target.shape[0], -1)
  scores = scores.view(scores.shape[0], -1)
  g = ve_sde_diffusion(t=t, sigma_min=sigma_min, sigma_max=sigma_max)
  lmd = g ** 2
  loss = torch.sqrt(((scores - target) ** 2).sum(dim=-1)) * lmd
  return loss.mean()

ここで本によると\(\lambda(t)=g(t)^2\)のときにスコアマッチングの目的関数は負の対数尤度の上限となっていることが証明できるそうです。このため、上記のコードでは\(\lambda(t)=g(t)^2\)を利用しています。

VE-SDEのサンプリング

VE-SDEのサンプリングをするためには拡散過程を逆にたどる逆算過程を知る必要があります。

拡散過程のSDEは式(3.1)で与えらえるとするとこの逆算過程は以下のようになります。

$$ \begin{align} \text{d}\boldsymbol{x} =& \left\{f(\boldsymbol{x}, t) – \nabla \left[ \boldsymbol{G}(\boldsymbol{x}, t) \boldsymbol{G}(\boldsymbol{x}, t)^\text{T} \right] \right. \\
& \quad \left. – \left[ \boldsymbol{G}(\boldsymbol{x}, t) \boldsymbol{G}(\boldsymbol{x}, t)^\text{T} \right] \nabla_{\boldsymbol{x}} \log p_t(\boldsymbol{x})\right\} \text{d}t \\
& \quad+ \boldsymbol{G}(\boldsymbol{x}, t)\text{d}\bar{\boldsymbol{w}} \tag{3.45} \end{align} $$

ただし、\(\text{d}\bar{\boldsymbol{w}}\)は時刻Tから0まで客向きに辿ったときの標準ウィーナー過程です。

ただし、一般的に拡散もモデルで使われる確率微分方程式は式(3.2)の形だそうです。このため式(3.2)で使われている\(f(t), g(t)\)で式(3.45)を書き直すと以下のようになります。

$$ \begin{align} \text{d}\boldsymbol{x} =& \left[f(t) – g(t)^2\nabla \log p_t(\boldsymbol{x})\right] \text{d}t + g(t)\text{d}\bar{\boldsymbol{w}} \tag{3.46} \end{align} $$

式(3.45)と式(3.46)の式変形の説明も本当はやろうと思ったのですが、かなり長い式変形になるのと、本の付録のほうに詳しい説明があるのでこの記事では省略します。

この式(3.46)に基づいて拡散モデルのサンプリングをする方法としてオイラー・丸山先生によるサンプリングが本で紹介されています。疑似コードは以下の通りです。(「拡散モデル　データ生成技術の数理」Algorithm 3.1の引用)

1. \(\boldsymbol{x} \sim \mathcal{N}(0, \boldsymbol{I})\))
2. for \(i=T,…,1\) do
3. \(\quad \boldsymbol{z}_i \sim \mathcal{N}(0, \boldsymbol{I})\)
4. \(\quad \boldsymbol{x} := \boldsymbol{x} – \left[f(t_i) – g(t_i)^2 s_{\theta}(\boldsymbol{x}, t_i)\right] \Delta t_i + g(t)\sqrt{|\Delta t_i|} \boldsymbol{z}_i \)
5. end for
6. return \(\boldsymbol{x}\)

これをPyTorchで実装すると以下のようになります。

def euler_maruyama_sample(n_samples, score_model, device=device, n=1000):
  with torch.no_grad():
    x = torch.randn(n_samples, 2, device=device)
    dt = torch.tensor(1.0 / n, device=x.device)
    for t in range(n, 0, -1):
      t_tensor = torch.full((n_samples, 1), t/n, device=device)
      z = torch.randn(n_samples, 2)
      f = ve_sde_drift(t_tensor, score_model.sigma_min, score_model.sigma_max)
      g = ve_sde_diffusion(t_tensor, score_model.sigma_min, score_model.sigma_max)
      g2 = g ** 2
      score = score_model(x, t_tensor)
      x = x - (f*x - g2 * score) * dt + g * torch.sqrt(dt) * z
    return x

コードの実行例

ここでは先ほど紹介したロス関数とサンプリング関数を利用して実際にVE-SDEでスコア関数のパラメータを学習し、サンプリングした例を示します。

参考例として入力となる\(\boldsymbol{x}\)のサンプリングする分布の確率密度関数は以下のように平均が違うガウス分布二つの混合分布とし、サンプリングしたデータを正規化して使用します。

n_samples = int(1e6)
sigma = 0.01

dist0 = torch.distributions.MultivariateNormal(torch.tensor([-2, -2], dtype=torch.float).to(device), sigma*torch.eye(2, dtype=torch.float).to(device))
samples0 = dist0.sample(torch.Size([n_samples//2]))
    
dist1 = torch.distributions.MultivariateNormal(torch.tensor([2, 2], dtype=torch.float).to(device), sigma*torch.eye(2, dtype=torch.float).to(device))
samples1 = dist1.sample(torch.Size([n_samples//2]))
samples = torch.vstack((samples0, samples1))

mean = torch.mean(samples, dim=0)
std = torch.std(samples, dim=0)

normalized_samples = (samples - mean[None, :])/std[None, :]

使用する\(\boldsymbol{x}\)を2Dのヒストグラムで可視化すると以下のようになります。

次にスコア関数のモデルと学習コードです。基本的には先ほど紹介したロス関数を使ってモデルを学習形になります。

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F

class ScoreModel(nn.Module):
  def __init__(self, sigma_min, sigma_max, n_channels=2):
    super(ScoreModel, self).__init__()
    self.sigma_min = sigma_min
    self.sigma_max = sigma_max
    self.model = nn.Sequential(
        nn.Linear(n_channels, 2*n_channels),
        nn.ELU(),
        nn.Linear(2*n_channels, 16*n_channels),
        nn.ELU(),
        nn.Linear(16*n_channels, 2*n_channels),
        nn.ELU(),
        nn.Linear(2*n_channels, n_channels),
    )

  def forward(self, x, t):
    y = self.model(x)
    sigma_t = sigma(t=t, sigma_min=self.sigma_min, sigma_max=self.sigma_max)
    return y/sigma_t

batch_size = 512
n_steps = 100000

dataloader = torch.utils.data.DataLoader(dataset, batch_size=512, shuffle=True, num_workers=0)
dataloader_iter = iter(dataloader)

score_model = ScoreModel(sigma_min=sigma_min, sigma_max=sigma_max).to(device)

optimizer = torch.optim.Adam(score_model.parameters())
lr_scheduler = torch.optim.lr_scheduler.OneCycleLR(optimizer, max_lr=0.001, total_steps=n_steps)

for i in range(n_steps):
  try:
    x = next(dataloader_iter)[0]
  except StopIteration:
    dataloader_iter = iter(dataloader)
    x = next(dataloader_iter)[0]
  x = x.to(device)

  optimizer.zero_grad()
  loss = dsm_loss(score_model, x, sigma_min=sigma_min, sigma_max=sigma_max)
  loss.backward()
  optimizer.step()
  lr_scheduler.step()
  if (i % 1000) == 0:
    print(f"{i} steps loss:{loss}")

学習が終わったら最後に以下のようにサンプリングする関数を呼び出してサンプリングします。

samples_pred = euler_maruyama_sample(n_samples=100000, score_model=score_model)

サンプリングされたデータの2Dのヒストグラムは以下の通りです。

ほぼ元の分布と同じサンプリングが得られることが確認できました。

終わりに

今回は「拡散モデル　データ生成技術の数理」の中で紹介されている分散発散型確率微分方程式 (VE-SDE)の部分を紹介しました。コードは先月の中旬にはできていたのですが、今回紹介する部分の式変形でぱっと見てわからないところがいくつかあり、それを調べていたらだいぶ時間がかかりました。また、説明のために必要な式の打ち込みにもかなり時間がかかってしまいました。

ただ、頑張ったおかげでかなりVE-SDEの部分の理解が進んだので記事にまとめてよかったです。

今後に関してはVP-SDEに関してもやろうと思っていますが、先に最近流行りのChatGPT, LLM, LangChainあたりに関していろいろ調べてみようと思うのでそちらの記事をいくつか書いてからになると思います。

この記事が他の方の役に少しでもなれば幸いです。

参考文献

1. 確率微分方程式　入門から応用まで
2. Särkkä, S., & Solin, A. (2019). Applied Stochastic Differential Equations (Institute of Mathematical Statistics Textbooks). Cambridge: Cambridge University Press. doi:10.1017/9781108186735
3. Song, Y., Sohl-Dickstein, J.N., Kingma, D.P., Kumar, A., Ermon, S., & Poole, B. (2020). Score-Based Generative Modeling through Stochastic Differential Equations. ArXiv, abs/2011.13456.

The post [勉強ノート] 「拡散モデル　データ生成技術の数理」 3.1-3.5のVE-SDE部分について first appeared on まったり勉強ノート.

その第3弾として「2.3 デノイジング拡散確率モデル」で説明されているデノイジング拡散確率モデル（DDPM）の学習とそれを使ったサンプリングについてPython(深層学習部分はPytorch)でコードを書いて試したのでそのまとめになります。今回の記事ではDDPMの細かい数式を説明すると記事の量がすごいことになりそうなので、重要な部分だけ説明していきます。

また、この本を買うか迷っている方は私が読んだ感想をこちらの記事に書いてますので参考にしてみてください。

また、この記事で紹介したコードは以下にあげてありますので、コード全体を確認したい方はこちらをご覧ください。

https://github.com/shu65/diffusion-model-book/blob/main/diffusion_model_book_2_3_ddpm.ipynb

デノイジング拡散確率モデルとは

デノイジング拡散確率モデル（DDPM）はデータに対して徐々にノイズを加えていく拡散過程を逆向きに辿っていく逆拡散過程によってデータ生成を行います。図でみると分かりやすいと思うので拡散過程と逆拡散過程の関係図を以下に示します。

まずは上の図の上段の拡散過程についてです。ここでは\(\boldsymbol{x}_0\)がオリジナルのデータで、これに徐々にノイズを加えていき、\(\boldsymbol{x}_1, \boldsymbol{x}_2, \boldsymbol{x}_3, …, \boldsymbol{x}_T\)といいうデータを作っていきます。これを確率密度関数で表すと以下のようになります。

$$ \begin{align*}
q(\boldsymbol{x}_{1:T}|\boldsymbol{x}_0) :=& \prod_{t=1}^T q(\boldsymbol{x}_t|\boldsymbol{x}_{t-1}) \tag{2.1.1} \\
q(\boldsymbol{x}_{t}|\boldsymbol{x}_{t-1}) :=& \mathcal{N}(\boldsymbol{x}_t; \sqrt{\alpha_t}\boldsymbol{x}_{t-1}, \beta_t \boldsymbol{I}) \tag{2.1.2} \\
\end{align*} $$

ここで\(\beta_t\) は分散の大きさを制御するパラメータで、\(0<\beta_1<\beta_2<…<\beta_T<1\)です。また、\(\alpha_t := 1 – \beta_t\) で、\(\alpha_t, \beta_t\)を合わせてノイズスケジュールと呼びます。

ここで、\(\mathcal{N}(\boldsymbol{x}_t; \sqrt{\alpha_t}\boldsymbol{x}_{t-1}, \beta_t \boldsymbol{I})\) について詳しく見ていきます。この拡散過程を繰り返していくと、\(\beta_t\) は徐々に大きくなります。結果として、ノイズ成分は大きくなっていきます。

一方、\(\beta_t\)が大きくなるということは\(\alpha_t := 1 – \beta_t\) なので\(\boldsymbol{x}_{t-1}\)の係数の\(\sqrt{\alpha_t}\)はどんどん小さくなっていきます。結果として、拡散過程を繰り返していくと任意の\(\boldsymbol{x}_{0}\)に対して、以下のような近似ができるようになります。

$$ \begin{align*}
q(\boldsymbol{x}_{T}|\boldsymbol{x}_{0}) :=& \mathcal{N}(\boldsymbol{x}_T; 0, \boldsymbol{I}) \tag{2.1.3} \\
\end{align*} $$

次に、図の下段の逆拡散過程です。これは\(\mathcal{N}(\boldsymbol{x}_T; 0, \boldsymbol{I}) \)からスタートして拡散過程を逆向きに辿っていく処理になります。

この逆拡散過程の各ステップを正規分布で表し、この正規分布の平均と共分散行列を一つ前のステップの変数\(\boldsymbol{x}_{t}\)と時刻\(t\)を入力としてパラメータ\(\theta\)を使ったモデル(\(\mu_{\theta}(\boldsymbol{x}_{t}, t), \Sigma(\boldsymbol{x}_{t}, t))\))として表します。これを使うと逆拡散過程は以下のような式で表すことができます。

$$ \begin{align*}
p_{\theta}(\boldsymbol{x}_{0:T}) :=& p(\boldsymbol{x}_T)\prod_{t=1}^T p_{\theta}(\boldsymbol{x}_t-1|\boldsymbol{x}_{t}) \tag{2.1.4} \\
p_{\theta}(\boldsymbol{x}_{t-1}|\boldsymbol{x}_{t}) :=& \mathcal{N}(\boldsymbol{x}_{t-1}; \mu_{\theta}(\boldsymbol{x}_{t}, t), \Sigma_{\theta}(\boldsymbol{x}_{t}, t)) \tag{2.1.5} \\
p(\boldsymbol{x}_{T}) =& \mathcal{N}(\boldsymbol{x}_T; 0, \boldsymbol{I}) \tag{2.1.6}
\end{align*} $$

この\(\mu_{\theta}(\boldsymbol{x}_{t}, t), \Sigma_{\theta}(\boldsymbol{x}_{t}, t)\)は後ほど示す通りニューラルネットワークなどを使ってモデル化します。これによって逆拡散過程を実現しています。

では次から実際にこのモデルのパラメータ\(\theta\)の学習方法とそれを使ったサンプリングを見ていきます。

デノイジング拡散確率モデルの学習

デノイジング拡散確率モデルのモデルの学習方法の説明を本来はしたいのですが、この説明はすごく長いものになります。このため、詳しい説明は本を見ていただくとして、ここでは学習を回すうえで重要な変数と式に関する簡単な説明にとどめておきます。

まず、先ほど示した式(2.1.2)などを利用すると以下の式が導けます。

$$ \begin{align*}
q(\boldsymbol{x}_{t}|\boldsymbol{x}_{0}) :=& \mathcal{N}(\boldsymbol{x}_t; \sqrt{\bar{\alpha}_t}\boldsymbol{x}_{0}, \bar{\beta}_t \boldsymbol{I}) \tag{2.1.7} \\
\bar{\alpha}_t := \prod_{s=1}^t \alpha_s \tag{2.1.8} \\
\bar{\beta}_t := 1 – \bar{\alpha}_t \tag{2.1.9} \\
\end{align*} $$

式(2.1.7)から式(2.1.9)の導出の証明に関しては本の式(2.1)の下に証明がありますので詳しく知りたい方はそちらをご覧ください。

これにより、わざわざ式(2.1.2)に従って、\(\boldsymbol{x}_{0}\)から徐々に\(t\)を大きくして\(\boldsymbol{x}_{t}\)を生成しなくても、式(2.1.7)を使うことで正規分布のサンプリングを１度すれば任意の\(t\)の\(\boldsymbol{x}_{t}\)のデータを生成できることになります。学習ではこれを利用して高速にランダムな\(t\)のデータを生成して学習に利用します。

次にデノイジング拡散確率モデルの学習で用いるロス関数についてです。これに関しては前の記事でも紹介したスコアベースモデルと同じようにデノイジングスコアマッチングを利用します。また、ロス関数の導出は先ほど説明した通り、ちゃんと説明しようとするとすごく長いので本で示されている最終的な結果を以下に示します。

$$ \begin{align*}
L_{\gamma}(\theta) =& \sum_{t=1}^T w_t E_{\boldsymbol{x}_0, \epsilon} \left\{ \left\| \epsilon – \epsilon_{\theta}(\sqrt{\bar{\alpha}_t}\boldsymbol{x}_0 + \sqrt{\bar{\beta}_t}\epsilon, t) \right\|^2 \right\} \tag{2.1.10} \\
\gamma =& \left\{ w_1, w_2, …, w_T \right\} \tag{2.1.11}
\end{align*} $$

この式の\(w_t\)に関しては本によると\(w_t = 1\) がよくつかわれるとのことなので、このあとの実装のコードでも\(w_t = 1\)としています。

ここで、式の導出を省略して関係が分かりにくくなっていため、逆拡散過程の説明で出てきた式(2.1.5)の中の\(\mu_{\theta}(\boldsymbol{x}_{t}, t), \Sigma_{\theta}(\boldsymbol{x}_{t}, t)\)と式(2.1.10)の中で出てきている\(\epsilon_{\theta}(\sqrt{\bar{\alpha}_t}\boldsymbol{x}_0 + \sqrt{\bar{\beta}_t}\epsilon, t)\) との関係を説明しておきます。

逆拡散過程を行う上で\(\mu_{\theta}(\boldsymbol{x}_{t}, t), \Sigma_{\theta}(\boldsymbol{x}_{t}, t)\)のパラメータ\(\theta\)を学習する必要があります。ここで、本の説明によると\(\Sigma_{\theta}(\boldsymbol{x}_{t}, t)\)に関してはパラメータ\(\theta\)に依存しない固定の\(\Sigma_{\theta}(\boldsymbol{x}_{t}, t)) = \sigma_t^2 \boldsymbol{I}\)を使うことが多いそうです。先ほど示したロス関数も\(\Sigma_{\theta}(\boldsymbol{x}_{t}, t)) = \sigma_t^2 \boldsymbol{I}\)として式変形しています。

次に\(\mu_{\theta}(\boldsymbol{x}_{t}, t)\)の部分です。こちらはロス関数の導出の過程で結局は\(t\)の時点で加えられたノイズを予測できるモデルに置き換えることができます。このため、\(\mu_{\theta}(\boldsymbol{x}_{t}, t)\)ではなく、ノイズを予測する\(\epsilon_{\theta}(\boldsymbol{x}_{t}, t)\)がロス関数の中で登場しています。

また、\(\boldsymbol{x}_{t}\) は式(2.1.7)から以下のようになります。

$$ \begin{align*}
\boldsymbol{x}_{t} =& \sqrt{\bar{\alpha}_t}\boldsymbol{x}_0 + \sqrt{\bar{\beta}_t}\epsilon \\
\epsilon \sim& \mathcal{N}(0, \boldsymbol{I}) \\
\tag{2.1.12}
\end{align*} $$

本のほうには丁寧にこのロス関数の導出が書かれているので詳細を知りたい方はぜひ本を読んでください。

この式(2.1.10)をロス関数として利用したデノイジング拡散確率モデルの学習の疑似コードは以下の通りです。(「拡散モデル　データ生成技術の数理」Algorithm 2.2の引用)

1. repeat
2. \(\quad \boldsymbol{x}_0 \sim p_{\text{data}}(\boldsymbol{x}_0)\)
3. \(\quad t \sim \text{Uniform}({1, …, T})\)
4. \(\quad \epsilon \sim \mathcal{N}(0, \boldsymbol{I})\)
5. \(\quad g := \nabla_{\theta} w_t \left\| \epsilon – \epsilon_{\theta}(\sqrt{\bar{\alpha_t}} \boldsymbol{x}_0 + \sqrt{\bar{\beta_{t}}} \epsilon, t)\right\|^2 \)
6. \(\quad \theta := \theta – \alpha g \)
7. until converged

ここで本では特に説明されてないですが、式(2.1.10)の最初の\(\sum_{t=1}^T\)の部分は少し変更して、\(t=1\)から\(t=T\)ランダムに\(t\)を選んで使用するように変更されています。

このアルゴリズムでは、まず最初にデータ\(\boldsymbol{x}_0 \)と\(t\)を選び、それらを用いて式(2.1.10)のロス関数を計算します。その後、勾配を計算してパラメータのアップデートをするということを繰り返します。6行目の勾配を使ったパラメータのアップデートは深層学習の基本的なパラメータ更新の確率的勾配降下法を利用したコードになっています。この部分は確率的勾配降下法以外のもの、例えばAdamなどでも問題ありません。

この疑似コードを基にPyTorchで実装するとこのようになります。

batch_size = 512
n_steps = 100000

dataloader = torch.utils.data.DataLoader(dataset, batch_size=512, shuffle=True, num_workers=0)
dataloader_iter = iter(dataloader)

model = Model().to(device)

optimizer = torch.optim.Adam(model.parameters())
lr_scheduler = torch.optim.lr_scheduler.OneCycleLR(optimizer, max_lr=0.001, total_steps=n_steps)
loss_func = torch.nn.MSELoss(reduction="none")

for i in range(n_steps):
  try:
    x0 = next(dataloader_iter)[0]
  except StopIteration:
    dataloader_iter = iter(dataloader)
    x0 = next(dataloader_iter)[0]
  x0 = x0.to(device)

  optimizer.zero_grad()

  noise = torch.randn_like(x0)
  t = torch.randint(0, len(x0), (x0.shape[0],), device=device)
  x = torch.sqrt(alpha_bars[t])[:, None] * x0 + torch.sqrt(beta_bars[t])[:, None] * noise
  noise_pred = model(x, t)
  w = 1.0
  losses = w * loss_func(noise_pred, noise)
  loss = losses.mean()

  loss.backward()
  optimizer.step()
  lr_scheduler.step()

基本的には疑似コードとほぼ同じですが、ロス関数周りで違うところがあるので、簡単に説明します。まず、深層学習なのでミニバッチを使った学習に置き換えています。このため、\(\boldsymbol{x}_0 \)も一つだけでなく、バッチサイズ分ランダムに選んでいます。これに伴って、\(t\)もバッチサイズ分ランダムに選んで利用します。そして、最終的には各データ点のロス関数の値の平均を計算して勾配を計算するという形に置き換えています。

また、パラメータの最適化の部分は確率勾配降下法ではなくAdamを利用しています。

デノイジング拡散確率モデルを使ったサンプリング

ここから先ほど紹介した方法で学習したモデルを利用してどのようにサンプリングしていくか、について説明します。

基本的には式(2.1.5)に従って逆拡散過程のステップを繰り返すことで実現します。

ここで、説明を省略してしまいましたが、式(2.1.5)の中にでてくる\(\mu_{\theta}(\boldsymbol{x}_{t}, t)\)を学習したノイズを予測するモデル\(\epsilon_{\theta}(\boldsymbol{x}_{t}, t)\)を使って表すと以下のようになります。

$$ \begin{align*}
\mu_{\theta}(\boldsymbol{x}_{t}, t) =& \frac{1}{\sqrt{\bar{\alpha}}} \left( \boldsymbol{x}_{t} – \frac{\beta_t}{\sqrt{\bar{\beta_t}}} \epsilon_{\theta}(\boldsymbol{x}_{t}, t) \right) \tag{2.1.13}
\end{align*} $$

また、一方、\(\Sigma_{\theta}(\boldsymbol{x}_{t}, t))\)は先ほど説明した通り、\(\Sigma_{\theta}(\boldsymbol{x}_{t}, t)) = \sigma_t^2 \boldsymbol{I}\)です。

これらと式(2.1.5)に従ったサンプリングの疑似コードは以下の通りです。(「拡散モデル　データ生成技術の数理」Algorithm 2.3の引用)

1. \(\boldsymbol{x}_T \sim \mathcal{N}(0, \boldsymbol{I})\)
2. for \(t=T, …, 1\) do
3. \(\quad \boldsymbol{u}_t \sim \mathcal{N}(0, \boldsymbol{I})\)
4. \(\quad\) if \(t=1\) then \(\boldsymbol{u}_t := 0\)
5. \(\quad \boldsymbol{x}_{t-1} := \frac{1}{\sqrt{\bar{\alpha}}} \left\{ \boldsymbol{x}_{t} – \frac{\beta_t}{\sqrt{\bar{\beta_t}}} \epsilon_{\theta}(\boldsymbol{x}_{t}, t) \right\} + \sigma_t \boldsymbol{u}_t \)
6. end for
7. return \(\boldsymbol{x}_0\)

基本的には徐々にノイズを取り除くことで目的のデータをサンプリングするという流れです。

PyTorchのコードとしては以下のようになります。

def ddpm_sample(n_samples, model, alphas, betas, beta_bars):
    xt = torch.randn(n_samples, 2)
    T = len(alphas)
    for t in range(T -1, -1, -1):
      print(f"t:{t}")
      ut = torch.randn(n_samples, 2)
      if t == 0:
        ut[:, :] = 0.0
      with torch.no_grad():
        noise_pred = model(xt, t*torch.ones(n_samples, dtype=xt.dtype))
        sigma_t = torch.sqrt(betas[t])
        xt = 1 / torch.sqrt(alphas[t]) * (xt - betas[t]/torch.sqrt(beta_bars[t])*noise_pred) + sigma_t*ut
    return xt

ここで、n_samplesがサンプリングするサンプル数、modelが\(\epsilon_{\theta}(\boldsymbol{x}_{t}, t)\)、alphas、betas, beta_barsがそれぞれ\(\alpha_t, \beta_t, \bar{\beta}_t\) のリストです。

実行例

先ほど紹介したPythonコードを実際に動かした例も示しておきます。参考例として入力となる\(\boldsymbol{x}\)のサンプリングする分布の確率密度関数は以下のように平均が違うガウス分布二つの混合分布とし、サンプリングしたデータを正規化して使用します。

n_samples = int(1e6)
sigma = 0.01

dist0 = torch.distributions.MultivariateNormal(torch.tensor([-2, -2], dtype=torch.float).to(device), sigma*torch.eye(2, dtype=torch.float).to(device))
samples0 = dist0.sample(torch.Size([n_samples//2]))
    
dist1 = torch.distributions.MultivariateNormal(torch.tensor([2, 2], dtype=torch.float).to(device), sigma*torch.eye(2, dtype=torch.float).to(device))
samples1 = dist1.sample(torch.Size([n_samples//2]))
samples = torch.vstack((samples0, samples1))

mean = torch.mean(samples, dim=0)
std = torch.std(samples, dim=0)

normalized_samples = (samples - mean[None, :])/std[None, :]

使用する\(\boldsymbol{x}\)を2Dのヒストグラムで可視化すると以下のようになります。

このデータを再現できるようにデノイジング拡散確率モデルを学習します。コードとしては先ほど示した通りです。

学習が終わったら次は以下のようにサンプリングを行います。

samples_pred = ddpm_sample(n_samples=100000, model=model, alphas=alphas, betas=betas, beta_bars=beta_bars)

サンプリングされたデータの2Dのヒストグラムは以下の通りです。

可視化結果をみると元の分布の平均の近くにデータが集中しているので、うまくいっていると考えられます。

終わりに

今回は「拡散モデル　データ生成技術の数理」の2.3のデノイジング拡散確率モデルの簡単な説明とコードを書いたのでそのまとめの記事になります。先日スコアベースモデルのコードを用意したことで、そのコードを参考に今回のデノイジング拡散確率モデルをすぐに作ることができたのですが、説明はすごい大変でした。

スコアベースモデルのほうも気になるという方はこちらをご覧ください。

今後、3章で紹介されている連続時間化拡散モデルのVE-SDEのほうも紹介予定です。コードは昨日できました。ただ、思ったよりも説明が大変そうなので、記事を書くのに時間がかかると思います。

この記事が少しでもみなさんの理解の助けになれば幸いです。

参考文献

1. Ho, J., Jain, A., & Abbeel, P. (2020). Denoising Diffusion Probabilistic Models. ArXiv, abs/2006.11239.
2. https://github.com/hojonathanho/diffusion

The post [勉強ノート] 「拡散モデル　データ生成技術の数理」 2.3 デノイジング拡散確率モデル first appeared on まったり勉強ノート.

その第2弾として「2.2 スコアベースモデル」で説明されているスコアベースモデルの学習とそれを使ったサンプリングについてPython(深層学習部分はPytorch)でコードを書いて試したのでそのまとめになります。

また、この本を買うか迷っている方は私が読んだ感想をこちらの記事に書いてますので参考にしてみてください。

また、この記事で紹介したコードは以下にあげてありますので、コード全体を確認したい方はこちらをご覧ください。

https://github.com/shu65/diffusion-model-book/blob/main/diffusion_model_book_2_2_score_based_model.ipynb

スコアベースモデルとは

１章で紹介されているデノイジングスコアマッチングは以下の２つの問題点があると本では紹介されています。

1. デノイジングスコアマッチングで推定されたスコア関数はデータ分布の密度が小さい領域で不正確
2. データ分布が多峰性を持つ場合、あるモード（確率が大きい領域）から他のモードに移る際、確率が小さい領域を通過するために非常に多くのステップを必要とする

これらの問題を解決するためにスコアベースモデル（SBM）[1, 2] では複数の異なる強度のノイズによって攪乱した攪乱後分布を用意して、それらの攪乱後分布上のスコアを求めるようにしています。

スコアベースモデルの学習

スコア関数 \(s_{\theta}(\boldsymbol{x}, \sigma_t)\) を学習する際は以下のロス関数を使います。

$$ \begin{align*}
L_{\text{SBM}}(\theta) := \sum_{t=1}^T w_t E_{p_{\sigma_t}}(\tilde{\boldsymbol{x}}) \left\{ \left\| \nabla_{\tilde{\boldsymbol{x}}} \log p_{\sigma_t}(\tilde{\boldsymbol{x}}) – s_{\theta}(\tilde{\boldsymbol{x}}, \sigma_t) \right\|^2 \right\} \tag{2.2.1}
\end{align*} $$

ここで\(\sigma_t \) はノイズの強さを表す変数で\( \sigma_{min} = \sigma_1 < \sigma_2 <… < \sigma_T = \sigma_{max}\)の合計\(T\)個をスコアベースモデルでは利用します。そして、\(p_{\sigma_t}(\tilde{\boldsymbol{x}}) \) \(x\)は\(x\)の分布\(p(x)\)を\(\sigma_t\)の強さで攪乱したあとの分布を表しています。

この式(2.2.1)を本の1.5.5の「デノイジングスコアマッチング」で説明されている通り、デノイジングスコアマッチングを使って式を書き換えると以下のようになります。

$$ \begin{align*}
L_{\text{DSM-SBM}}(\theta) := \sum_{t=1}^T w_t E_{\boldsymbol{x} \sim p_{data}(\boldsymbol{x}),\tilde{\boldsymbol{x}} \sim \mathcal{N}(\boldsymbol{x}, \sigma_t^2\boldsymbol{I})} \left\{ \left\| \frac{\boldsymbol{x} – \tilde{\boldsymbol{x}}}{\sigma_t^2} – s_{\theta}(\tilde{\boldsymbol{x}}, \sigma_t) \right\|^2 \right\} \tag{2.2.2}
\end{align*} $$

詳細は本にわかりやすくかいてあるので本を参照してください。

ここで本の式(1.9)のデノイジングスコアマッチングの式において最初に\(1/2\)があるのに式(2.2.2)ではそれが省略されています。これに関して本にはちゃんと書いてない気がしますが、おそらくこれは\(w_t\)の中に\(1/2\)が含まれているから、もしくは\(1/2\)は定数であり、最適化の際にパラメータが移動する方向は\(1/2\)のありなしで変わらないということで省略しているのではないかと思っています。

ここから２章にはちゃんと書いてないですが、Pythonで実装するためにさらに式変形していきます。\(\tilde{\boldsymbol{x}} \sim \mathcal{N}(\boldsymbol{x}, \sigma_t^2\boldsymbol{I})\)なので、\(\tilde{\boldsymbol{x}}\)を\(\epsilon \sim \mathcal{N}(0, \sigma_t^2 \boldsymbol{I})\)を使って表すと以下のようになります。

$$ \begin{align*}
\tilde{\boldsymbol{x}} = \boldsymbol{x} + \epsilon \tag{2.2.3}
\end{align*} $$

この式(2.2.3)を使って式(2.2.2)を式変形すると以下の通りです。

$$ \begin{align*}
L_{\text{DSM-SBM}}(\theta) :=& \sum_{t=1}^T w_t E_{\boldsymbol{x} \sim p_{data}(\boldsymbol{x}),\tilde{\boldsymbol{x}} \sim \mathcal{N}(\boldsymbol{x}, \sigma_t^2\boldsymbol{I})} \left\{ \left\| \frac{\boldsymbol{x} – \tilde{\boldsymbol{x}}}{\sigma_t^2} – s_{\theta}(\tilde{\boldsymbol{x}}, \sigma_t) \right\|^2 \right\} \\
=& \sum_{t=1}^T w_t E_{\boldsymbol{x} \sim p_{data}(\boldsymbol{x}),\epsilon \sim \mathcal{N}(\boldsymbol{x}, \sigma_t \boldsymbol{I})} \left\{ \left\| \frac{-\epsilon}{\sigma_t^2} – s_{\theta}(\tilde{\boldsymbol{x}}, \sigma_t) \right\|^2 \right\} \tag{2.2.4}
\end{align*} $$

この式を見たときに\(t=1\)から\(t=T\)までの和をとっている部分、\(T\)のサイズによっては計算量がすごいことにならないか？ということを思いました。このため、何か実装するときに工夫があるのかも？ということで[2]著者実装である[3]を見にいきました。すると2023/03/03時点では\(t=1\)から\(t=T\)ランダムに\(t\)を選び、その平均をとるということをしていました。

Pythonのコードのほうが分かりやすいと思うので、以下にPythonのコードも示しておきます。

def dsm_loss(score_model, samples, sigmas):
  t = torch.randint(0, len(sigmas), (samples.shape[0],), device=sigmas.device)
  used_sigmas = sigmas[t].view(samples.shape[0], *([1] * len(samples.shape[1:])))
  noise = torch.randn_like(samples) * used_sigmas
  perturbed_samples = samples + noise
  target = - 1 / (used_sigmas ** 2) * noise
  scores = score_model(perturbed_samples, used_sigmas)
  target = target.view(target.shape[0], -1)
  scores = scores.view(scores.shape[0], -1)
  w = used_sigmas.squeeze(-1) ** 2
  loss = ((scores - target) ** 2).sum(dim=-1) * w
  return loss.mean()

ここでscore_modelがスコア関数 \(s_{\theta}(\boldsymbol{x}, \sigma_t )\) 、samplesが\(\boldsymbol{x}\)、sigmasが\(\{\sigma_1,…,\sigma_T\}\)の配列となっています。また、\(w_t\)は本にならって\(w_t=\sigma_t^2\)を使っています。

この関数では最初にランダムに\(t\)を選び、それに従ってノイズを生成し、\(\tilde{\boldsymbol{x}}\)を作ります。その後、スコア関数のscore_modelを使ってスコアを計算し、式(2.2.4)を使ってロス関数を計算します。

このロス関数を使ってスコア関数のパラメータを学習していきます。

ここで１つ、スコア関数のモデルに関して注意点があります。スコア関数は\(s_{\theta}(\boldsymbol{x}, \sigma_t) \)は\(\boldsymbol{x}\)だけでなく\(\sigma_t\)も引数にとります。このため、モデルの中でどうにかして\(\sigma_t \)と\(\boldsymbol{x}\)の入力を組み合わせる必要があります。これに関して今回のコードでは[3]の実装にならって、以下のようにして\(\boldsymbol{x}\)だけを入力として受け取るスコア関数\(s_{\theta}^{\prime}(\boldsymbol{x})\)の出力を\(\sigma_t\)で割るという形にしています。

$$ \begin{align*}
s_{\theta}(\boldsymbol{x}, \sigma_t) = s_{\theta}^{\prime}(\boldsymbol{x}) / \sigma_t \tag{2.2.5}
\end{align*} $$

また、後ほど示しますが、今回は２つのガウス分布の混合分布を入力とします。この分布はシンプルな分布なため、今回は簡単なMLPをスコア関数のモデル使用します。コードとては以下のようになります。

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F

class ScoreModel(nn.Module):
  def __init__(self, n_channels=2):
    super(ScoreModel, self).__init__()

    self.model = nn.Sequential(
        nn.Linear(n_channels, 2*n_channels),
        nn.ELU(),
        nn.Linear(2*n_channels, 16*n_channels),
        nn.ELU(),
        nn.Linear(16*n_channels, 2*n_channels),
        nn.ELU(),
        nn.Linear(2*n_channels, n_channels),
    )

  def forward(self, x, sigma):
    y = self.model(x)
    return y/sigma

スコアベースモデルを使ったサンプリング

ここから学習済みのスコア関数\(s_{\theta}(\boldsymbol{x}, \sigma_t) \)を使ったサンプリングについて説明していきます。

スコアベースモデルを使ったサンプリング１章で紹介されたランジュバン・モンテカルロ法をベースにしています。ランジュバン・モンテカルロ法の部分についてはこちらに解説しています。

詳細は上の記事にかいてありますが、ランジュバン・モンテカルロ法は最初、ランダムに\(\boldsymbol{x}_0\)を生成後、以下のランジュバン・モンテカルロ法の更新則を\(K\)回繰り返すことで\(p(\boldsymbol{x})\)からサンプリングしたようなデータを作ります。

$$ \begin{align*}
\boldsymbol{x}_k := \boldsymbol{x}_{k-1} + \alpha \nabla_\boldsymbol{x} \log p(\boldsymbol{x}_{k-1}) + \sqrt{2\alpha}\boldsymbol{u}_k \tag{2.2.6}
\end{align*} $$

スコアベースモデルのサンプリングでは更新則のスコア（\(\nabla_\boldsymbol{x} \log p(\boldsymbol{x}_{k-1})\)）を学習したスコア関数に置き換えた以下の更新則を利用します。

$$ \begin{align*}
\boldsymbol{x}_{t, k} := \boldsymbol{x}_{t, k-1} + \alpha_t s_{\theta}(\boldsymbol{x}_{t, k-1}, \sigma_t)+ \sqrt{2\alpha_t}\boldsymbol{u}_k \tag{2.2.7}
\end{align*} $$

この更新則を用いたスコアベースモデルのサンプリングの疑似コードは以下の通りです。(「拡散モデル　データ生成技術の数理」Algorithm 2.1の引用)

1. \(\boldsymbol{x}_0\)を初期化(\(\boldsymbol{x}_0 \sim \mathcal{N}(0, \sigma_T^2 \boldsymbol{I})\))
2. for \(t=1,…,T\) do
3. \(\quad \alpha_t := \alpha \sigma_t^2\/\sigma_T^2)\
4. \(\quad\) for \(k=1,…,K\) do
5. \(\qquad \boldsymbol{u}_k \sim \mathcal{N}(0, \boldsymbol{I})\)
6. \(\qquad\) if \(t=1\) and \(k=K\) then \(\boldsymbol{u}_k := 0\)
7. \(\qquad \boldsymbol{x}_{t, k} := \boldsymbol{x}_{t, k-1} + \alpha_t s_{\theta}(\boldsymbol{x}_{t, k-1}, \sigma_t)+ \sqrt{2\alpha_t}\boldsymbol{u}_k \)
8. \(\quad\) end for
9. \(\quad \boldsymbol{x}_{t-1, 0} := \boldsymbol{x}_{t, K}\)
10. end for
11. return \(\boldsymbol{x}_{0, 0}\)

ここで\(\alpha\)はステップ幅のスケール、\(K\)はステップ回数です。アルゴリズムを見て分かる通り、ノイズの強度を変えながらランジュバン・モンテカルロ法を使って少しずつ\(\boldsymbol{x}_{t, k}\)を変化させています。また、７行目にある通り、各ノイズの強度の最後のステップではデノイジングのみを行うことでサンプリングの品質を向上させています。

この疑似コードをPythonのコードにするとこのようになります。

def sbm_sample(n_samples, score_model, sigmas, alpha=0.1):
    sigma_T = sigmas[-1]
    x_0 = torch.randn(n_samples, 2)*sigma_T
    x_tk = x_0
    K = 200
    for t in range(len(sigmas) -1, -1, -1):
      sigma_t = sigmas[t]
      alpha_t = alpha*(sigma_t**2)/(sigma_T**2)
      print(f"t:{t}, sigma_t:{sigma_t}, alpha_t:{alpha_t}")
      for k in range(K+1):
        u_k = torch.randn(n_samples, 2)
        if (k == K) and t == 0:
          u_k[:, :] = 0.0
        with torch.no_grad():
          score = score_model(x_tk, sigma_t)
          x_tk = x_tk + alpha_t * score + np.sqrt(2 * alpha_t) * u_k
    return x_tk

n_samplesが生成するサンプル数、score_modelがスコア関数、sigmasがノイズ強度の配列、alphaがステップ幅のスケールになっています。

実行例

先ほど紹介したPythonコードを実際に動かした例も示しておきます。参考例として入力となる\(\boldsymbol{x}\)のサンプリングする分布の確率密度関数は以下のように平均が違うガウス分布二つの混合分布とし、サンプリングしたデータを正規化して使用します。

n_samples = int(1e6)
sigma = 0.01

dist0 = torch.distributions.MultivariateNormal(torch.tensor([-2, -2], dtype=torch.float).to(device), sigma*torch.eye(2, dtype=torch.float).to(device))
samples0 = dist0.sample(torch.Size([n_samples//2]))
    
dist1 = torch.distributions.MultivariateNormal(torch.tensor([2, 2], dtype=torch.float).to(device), sigma*torch.eye(2, dtype=torch.float).to(device))
samples1 = dist1.sample(torch.Size([n_samples//2]))
samples = torch.vstack((samples0, samples1))

mean = torch.mean(samples, dim=0)
std = torch.std(samples, dim=0)

normalized_samples = (samples - mean[None, :])/std[None, :]

使用する\(\boldsymbol{x}\)を2Dのヒストグラムで可視化すると以下のようになります。

このデータを再現できるようにスコア関数を学習します。学習コードは以下の通りです。

import torch

batch_size = 512
n_steps = 100000

dataset = torch.utils.data.TensorDataset((normalized_samples))
dataloader = torch.utils.data.DataLoader(dataset, batch_size=512, shuffle=True, num_workers=0)
dataloader_iter = iter(dataloader)

score_model = ScoreModel().to(device)

optimizer = torch.optim.Adam(score_model.parameters())
lr_scheduler = torch.optim.lr_scheduler.OneCycleLR(optimizer, max_lr=0.001, total_steps=n_steps)

for i in range(n_steps):
  try:
    x = next(dataloader_iter)[0]
  except StopIteration:
    dataloader_iter = iter(dataloader)
    x = next(dataloader_iter)[0]
  x = x.to(device)

  optimizer.zero_grad()
  loss = dsm_loss(score_model, x, sigmas)
  loss.backward()
  optimizer.step()
  lr_scheduler.step()
  if (i % 1000) == 0:
    print(f"{i} steps loss:{loss}")

学習が終わったら、以下のようにして学習したモデルを利用してサンプリングします。

samples_pred = sbm_sample(n_samples=100000, score_model=score_model, sigmas=sigmas)

サンプリングされたデータの2Dのヒストグラムは以下の通りです。

可視化結果をみると元の分布の平均の近くにデータが集中しているので、うまくいっていると考えられます。

ただ、やってみるとわかるのですがちゃんとした結果を得るために人手で決めないといけないハイパーパラメータの選択が難しい印象です。この結果もかなり試行錯誤してなんとかこの結果を作ることができたというイメージです。

終わりに

今回は「拡散モデル　データ生成技術の数理」の2.2のスコアベースモデルの説明の部分のコードを書いたのでそのまとめの記事になります。最初、MNISTのデータでやろうとして、MNISTのデータを学習できるコードを説明するのは結構大変、ということでシンプルな混合ガウス分布にしました。ただ、それでも結構な分量になった印象です。ちなみに次のDDPMも紹介用のコードはできているので、近日中に記事を書いて公開しようと思います。

この記事が少しでもみなさんの理解の助けになれば幸いです。

参考文献

1. Song, Y., & Ermon, S. (2019). Generative Modeling by Estimating Gradients of the Data Distribution. ArXiv, abs/1907.05600.
2. Song, Y., & Ermon, S. (2020). Improved Techniques for Training Score-Based Generative Models. ArXiv, abs/2006.09011.
3. https://github.com/ermongroup/ncsnv2

The post [勉強ノート] 「拡散モデル　データ生成技術の数理」 2.2 スコアベースモデル first appeared on まったり勉強ノート.

今回はその第１弾として「1.5.1 ランジュバン・モンテカルロ法」の部分で紹介されているランジュバン・モンテカルロ法を実装をして動かすとどうなるのか？を試してみたので、そのまとめの記事になります。

今回示しているコードはここに上がっています。

https://github.com/shu65/diffusion-model-book/blob/main/diffusion_model_book_1_5_1_langevin_monte_carlo.ipynb

また、この本を買うか迷っている方はこちらの記事に私が読んだ感想を書いてますので参考にしてみてください。

ランジュバン・モンテカルロ法とは？

ランジュバン・モンテカルロ法は拡散モデルで登場するMCMC法の１つであり、 \(p(\boldsymbol{x})\)のスコアと呼ばれるものが計算できれば\(p(\boldsymbol{x})\)からサンプリングしたものに近い\(\boldsymbol{x}\)が得られる手法です。

ランジュバン・モンテカルロ法は拡散モデルを学習後、生成のフェーズにおいて使われるので拡散モデルを理解するうえでランジュバン・モンテカルロ法は重要な手法です。

このランジュバン・モンテカルロ法の疑似コードは以下の通りです。(「拡散モデル　データ生成技術の数理」Algorithm 1.1の引用)

1. \(\boldsymbol{x}_0\)を初期化(\(\boldsymbol{x}_0 \sim \mathcal{N}(0, \boldsymbol{I})\))
2. for \(k=1,…,K\) do
3. \(\quad \boldsymbol{u}_k\ \sim \mathcal{N}(0, \boldsymbol{I})\)
4. \(\quad \boldsymbol{x}_k\ := \boldsymbol{x}_{k-1} + \alpha \nabla_{\boldsymbol{x}}\log p(\boldsymbol{x}_{k-1}) + \sqrt{2 \alpha}\boldsymbol{u}_k\)
5. end for
6. return \(\boldsymbol{x}_K\)

ここで、\(\alpha\)はステップ幅、\(K\)はステップ回数、\(p(\boldsymbol{x})\)が\(\boldsymbol{x}\)の確率密度関数です。また、\(\nabla_{\boldsymbol{x}}\log p(\boldsymbol{x}) \)が\(\boldsymbol{x}\)に対するスコアです。

疑似コードを見ていただければわかる通り、最初、ランダムに\(\boldsymbol{x}_0\)を生成後、以下のランジュバン・モンテカルロ法の更新則を\(K\)回繰り返すことで\(p(\boldsymbol{x})\)からサンプリングしたようなデータを作ります。

$$ \begin{align*}
\boldsymbol{x}_k := \boldsymbol{x}_{k-1} + \alpha \nabla_\boldsymbol{x} \log p(\boldsymbol{x}_{k-1}) + \sqrt{2\alpha}\boldsymbol{u}_k \\
\end{align*} $$

Pythonでランジュバン・モンテカルロ法を実装する

ここからPythonでランジュバン・モンテカルロ法を実装して実際に動かした様子を見ていきます。

まずは実装に関してです。ここではスコア(\(\nabla_{\boldsymbol{x}}\log p(\boldsymbol{x}) \))を簡単に計算するために深層学習ライブラリのPyTorchを利用します。PyTorchには自動微分の機能がありこれにより\(\nabla\)が簡単に計算することができます。

先ほど示したアルゴリズムをそのままPyTorchで実装すると以下のようになります。

def langevin_monte_carlo_algorithm(log_p_dist, n_samples=100000, alpha=0.001, K=1000):
    x0 = torch.randn(n_samples, 2)
    x_k = x0
    for i in tqdm(range(K+1)):
        x_k.requires_grad_()
        log_p = log_p_dist(x_k)
        score = torch.autograd.grad(log_p.sum(), x_k)[0]
        with torch.no_grad():
          u_k = torch.randn(n_samples, 2)
          x_k = x_k + alpha * score + np.sqrt(2 * alpha) * u_k
    return x_k

ここでlog_p_distはサンプリングしたい分布の確率密度関数の出力を対数にしたものを返す関数です。後ほど具体例を示しますが、PyTorchで\(\nabla\)を計算したいのでlog_p_distは自動微分が計算できるPyTorchの関数だけで計算する関数である必要があります。

また、x_kは自動微分により\(\nabla\)を計算する必要があるのでx_k.requires_grad_()として、torch.autograd.grad()を呼んだときが自動微分の結果が入るようにしておきます。これに加えて実際に自動微分を計算する際にPyTorchの制約上、torch.autograd.grad()のoutputs（ここではlog_p.sum()のところ）がスカラである必要があります。このため、そのままlog_pをtorch.autograd.grad()に引数で渡すとエラーになってしまうため、log_p.sum()をでスカラにしてからtorch.autograd.grad()の引数にいれています。

Pythonで実装したランジュバン・モンテカルロ法

次に先ほど示した実装を動かしてみます。ここでサンプリングする分布の確率密度関数は以下のように平均が違うガウス分布二つの混合分布としました。

def log_p_dist(x):
  dist0 = torch.distributions.MultivariateNormal(torch.tensor([-2, -2]), torch.eye(2))
  p0 = torch.exp(dist0.log_prob(x))
  dist1 = torch.distributions.MultivariateNormal(torch.tensor([2, 2]), torch.eye(2))
  p1 = torch.exp(dist1.log_prob(x))
  return torch.log((p0 + p1)/2.0)

確率密度関数を可視化してみると以下のようになります。

次にランジュバン・モンテカルロ法を使ってサンプリングしてみます。実行方法としては以下の通りです。

samples_1em1 = langevin_monte_carlo_algorithm(log_p_dist, alpha=1e-1)

これでヒストグラムを描くと以下のようになります。

図を見ていただければわかる通り、確率密度関数に近いヒストグラムを得ることができていることがわかります。

ここまでできたので、折角なので２章あたりで指摘される問題点の１つである\(\alpha\)の値の設定によってうまくサンプリングできない例を確認してみます。

具体的には\(\alpha\)を以下のように極端に小さくして実行してみます。

samples_1em5 = langevin_monte_carlo_algorithm(log_p_dist, alpha=1e-5)

ヒストグラムを描くとこのようになります。

\(\alpha\)の値を変な値にすると入力した確率密度関数とは全く違う分布になってしまうことがわかりました。これで２章の説明の重要性を実感することができました。

終わり

今回はランジュバン・モンテカルロ法の解説とPythonの実装を示しました。今回の章は比較的簡単に実装できました。ただ、このあと紹介する予定の部分は説明を聞くと簡単そうですが、実装するのは大変な部分が多い印象です。ですが、頑張って紹介できればと思っています。

The post [勉強ノート] 「拡散モデル　データ生成技術の数理」 1.5.1 ランジュバン・モンテカルロ法 first appeared on まったり勉強ノート.

拡散モデル データ生成技術の数理

岡野原 大輔

Amazon

楽天市場

ポチップ

どんな本か？

Stable Diffusionをはじめとした画像生成技術が劇的に進化して実際に人が書いたような絵が生成できるようになり、画像生成を使ってみたという話をよく聞くようになりました。このStable Diffusionなどの画像生成技術を大きく進歩させた技術が今回の本で紹介されている「拡散モデル」です。

この本ではこの拡散モデルの背景にある数式をわかりやすく説明してくれています。また、拡散モデルの応用先として画像生成が有名ですが、他の生成モデルにはない拡散モデルならではの特徴とその特徴が生かされた他の生成モデル（例えば化合物とか）についても説明がされています。

どんな人におすすめか？

この本は以下のようなことに興味ある人が最初に読むと良い本かなと思います。

• Stable Diffusionはどういう数学的背景がどうなっているのか?知りたい
• 拡散モデルについて知らないけど、何に使えそうか？を知りたい
• VAEやGANなど他の生成モデルとの違いを知りたい

この本を読んでさらに詳しく知りたい人は本で引用されている拡散モデルの論文を読むという流れでいくとスムーズに論文が理解できるのではないか？と思います。

ただし、タイトルに「数理」と書かれている通り、数式が大量にでてきます。なので数式みたくないという人にはお勧めしません。

ただ、数式は図や丁寧な式変形のおかげで他の機械学習系の本も見習ってほしいレベルで分かりやすく説明がされてます。例えば、本の中で「暗黙的スコアマッチングの目的関数」がどういうものなのか式の項毎に図を使って説明されています。項ごとにここまで丁寧に式の説明をしてくれている本はなかなかないので個人的に感動しました。

このため、他の機械学習本の数式変形が難しすぎて躊躇している人も買って読んでみてほしいと思っています。

個人的によかった点

個人的によかった点としては以下の通りです。

• 拡散モデルの特徴とその特徴を活かした応用先としてどういうものがあるのか幅広く知れた
• 数式の説明で図を使ってわかりやすく書かれていた
• 随所に論文の引用がついていて、この部分詳しく知りたいというときどの論文を読めばよいかが分かりやすい

特に2つ目の式変形に関して、式変形を飛ばしすぎて意味わからないと思った本を最近たくさん読んだせいかすごく感動しました。

終わりに

今回は「拡散モデル データ生成技術の数理」について紹介しました。数理と書かれているのでちょっと身構えてましたが、思ったよりわかりやすくて素晴らしい本でした。

他にも何冊か本を読んだ＆読んでいるので、しばらくは本の紹介の記事を書いていく予定です。

The post [書評] 拡散モデル データ生成技術の数理 ー 目覚ましい画像生成の発展の裏側を知りたい人へ first appeared on まったり勉強ノート.