そもそもなぜ自作キーボードを基板から作ったか？

Youtubeでたまに自作キーボードの紹介を見て、自分も欲しくなったため、今年になっていくつか自作キーボードのキットを買って自作キーボードを作るということをしています。

ただ、自作キーボードのキットは個人が作っているものが多く、人気なキットは入手難易度が高い印象です。

例えば、roBaやmoNa2 などの無線＋分割+トラックボールのような人気のキットは、開発者の方が結構定期的に在庫補充してくれている印象ですが、在庫が入ったら数分で売り切れてまったく買える気配がありません。

また、この手の小型の分割キーボードは「自分ならここにキーがほしい！」という要求が使っているといろいろでてきます。

そんなわけで5月か6月ごろから「もうこれは自分でつくるか」ということで基板から設計をすることにしました。

2025年8月現在、基板から自作キーボードを作りたくなったらどうすればいいか？

自作キーボードに何を求めるか？で大分話が変わりますが、自分は以下のような要素がほしくて自作キーボードを作りました。

1. 分割キーボード
2. 完全無線（分割キーボード間も無線だし、PCとキーボードの間も無線）
3. 最低片側6カラム、合計12カラム
4. トラックボールあり（できればスクロールとカーソル操作を別々のボールで操作したいので二つ）

自分と同じような要素を盛り込もうと思った場合、実は2025年8月現在はnoteで最高の資料が公開されています。

それが先ほど紹介したmoNa2を開発している白湯さんが書いたnoteの記事になります。

https://note.com/pooh_polo/m/me4a885480f75

このnoteのマガジンの記事では無線分割キーボードであるTweakbitというキーボードを作る流れなどが紹介されています。一部有料記事ではあるりますが、すごく安い＋キーボードを基板から自作した場合にかかる費用からすると誤差みたいな金額です。

このため、キーボードを作りたいと思った場合、現在なら有料記事も含めて全部1読することがまず1歩かと思います。

それ以外の要素のキーボードを作りたいときの参考になりそうな資料は？

要素ごとに参考になりそうな資料が微妙に違うのでカテゴリーわけして紹介していきます。一部重複する資料があります。

Tweakbitからレイアウトを変えたい

写真を見ればわかりますが、Tweakbitは格子状にスイッチを配置するオーソリニア配列です。これは基板から作る場合、作成の難易度が下がるので初心者向けのキーボードとしては作りやすくてよいのですが、ここを変えたいという要求がまずでてくると思います。

この場合、スイッチのレイアウトを事前に考えてからキーボードの基板を作るという流れになると思います。この辺の流れは以下の本が一番詳しく載っていました

• 自作キーボード設計入門(電子版) (https://booth.pm/ja/items/1044084)

ケースにこだわりたい

同じ基板でもケースが違うと印象がガラリと変わります。このため、ケースにこだわりたいという要求は出てくると思います。その場合はこちらの本がお勧めです。

• 自作キーボード設計ガイド Vol2 ケース設計編 (https://salicylic-acid3.booth.pm/items/4982088)

Lipoバッテリーが怖いので有線にしたい

無線のキーボード向けのマイコンでよく使われるのはSeeed XIAO BLE nRF52840なのですが、多くの場合Lipoバッテリーが使われます。これは万が一を考えると火事の原因になりえるので、怖いということを考えます。この場合は以下の資料がお勧めです。

• 自作キーボード設計入門(電子版) (https://booth.pm/ja/items/1044084)
• 自作キーボード設計入門2(電子版) (https://booth.pm/ja/items/1315935)

注意点としては2025/8現在、公開されているサンプルが少し違うなどの問題があるので、適宜現状に合わせて読み替える必要があります。

初心者が基板から自作キーボードを作る際の注意点

注意点としては以下の二つです

1. 2025年8月から大分あとの場合、いろいろ情報が古くなっている可能性がある
2. 自分が盛り込みたい要素がTweakbitと違う部分がある場合、この記事の情報だけでは実現できない場合がある

まず1についてです。これはもうどうしようもないと思いますが、情報が古くなった場合、例えばツールの使い方が変わっていて、初めて触る人は苦労するみたいなことが発生する可能性があります。

例えば、基板の設計時に使われるKiCadや3D CADのFusionはちょっと前の資料と今とで使い方が変わっている部分がありました。このような変化はある程度知識がある人にとっては特に問題ないことが多いと思いますが、初心者は苦労するポイントかと思います。

二つ目として白湯さんが題材にしているTweakbitと違う要素を盛り込みたいという場合です。例えばキー数は同じでキー配置を変えたいなどであればそれほど問題ないのですが、キー数を増やしたいは実は難易度がすごく上がります。

この話自体はnoteの記事に書かれていますが、Seeed XIAO BLE nRF52840というマイコンを使って、かつトラックボールありだと、普通にやったらTweakbitと同じキー数が限界になります。いろいろ裏技はあるので、増やせないことはないのですが、難易度は高くなると思います。

このようにキーボードに盛り込みたい要素に応じて難易度が変わるので、この点は注意してどのようなキーボードを作るか考える必要があります。

基板から自作キーボードを作る際の苦労ポイント

どの辺が大変なのかわかると参考になると思ったので、基板から作ることが初めての自分がどこに苦労したのか紹介しておこうと思います。

1. キースイッチとトラックボールのレイアウトの試行錯誤に時間がかかった

多分一番時間がかかったのがこの部分です。実は当初は全く知識がなかった基板の設計に一番時間がかかるのでは？と思ってたのですが、基板の設計は実は数時間できました。一方、予想よりもはるかに時間がかかったのがこのキースイッチとトラックボールのレイアウトです。

予想ではこの部分は自分が今使っている「KeyBall」シリーズを参考にして変えるつもりがほとんどなかったので、すぐ終わると思っていました。

ただ蓋をあけると、KeyBallとキー数が違うので、その部分の調整をしていたら、いろいろ調整したくなり、3D プリンターで試作をしていたら、だんだんトラックボールの位置を変えたくなり、調整してたら思った以上に時間がかかりました。

おそらく自分はキースイッチのレイアウトにある程度要望があったため、この部分に時間がかかったと思われるので、盛り込みたい要素の部分に時間がかかるのを覚悟するのが良いかともいます。

2. ファームウェアの調整に時間かかった

これは基板から作らない自作キーボードでもある程度時間がかかるところですが、できたキーボードのキーマップやトラックボールの設定に結構時間がかかりました。

普通の自作キーボードの場合、開発者の方が練りこんだキーマップをデフォルトにしてくれているおかげで、ユーザーは結構キーマップ作製の時間短縮ができる印象ですが、基板から作るとなると、参考にするキーマップがないので、いろいろ考えて結局時間がかかるということがありました。

トラックボールに関しても同様で、感度をいろいろ調整するのに時間がかかった印象があります。

終わりに

今回は基板から自作キーボードを作ったので、自分と同じように基板から作ってみたいと思っている人向けに記事を書いてみました。

この記事を見て「自分もできそう！」と思ってもらえれば幸いです。

The post 基板から自作キーボードを作ったので自分も基板から作りたいと思った人へ first appeared on まったり勉強ノート.

エヌビディア 半導体の覇者が作り出す2040年の世界

著:津田 建二

Amazon

楽天市場

ポチップ

どんな内容の本か？

みなさん、NVIDIAと聞いてピンとくるでしょうか？
PCに詳しい人であれば結構昔からGPUでお世話になっていたと思いますが、一般の人にとってはちょっと前は「謎の半導体企業」なんていわれるくらい馴染みのない会社だったのではないかと思います。ただ、最近は普通の経済ニュースでもよくNVIDIAの名前を見る気がするくらい一般の人でも知られる企業になったのではないかと思っています。
この本ではそんなNVIDIAについて、創業から現在に至るまでどのように成長し、そして今後についてどのように考えられていそうかが紹介されています。
個人的にNIVIDAは2000年後半から研究でGPUを使ってたこともあり、それ以降はそこそこ知っているつもりですが、創業当時の話はあまり知りませんでした。この本ではそのような時期についてもどのような苦労があったのか、そもそもなぜNVIDIAがGPUを作り始めたのかなどが簡単にですが紹介されています。
また、半導体事業がどのように変化していったのかも書かれているので、半導体事業についての歴史が知りたい人も読むと参考になる部分はあるかと思います。

どんな人にお勧めか？

今回の本は以下のような人にお勧めな本です。

1. NVIDIAという会社について知りたいと思った方
2. 半導体事業の歴史について知りたい方

1に関しては特に創業期の話なんかは、本でしか見ることがないと思うので、その部分が知りたい方はお勧めだと思います。一方、NVIDIAが考える未来の部分はNVIDIAのイベントのGTCなどで語られている部分とかぶる部分が多く、GTCのkeynoteなどを見ている人は新しい部分はない気がしました。
2に関しては、NVIDIAの成長と半導体事業の変化がかなり密接に関係していることもあり、この本でもそこそこページを割かれて説明されています。この前紹介したTSMCの本と合わせて読むと理解がより深まると思います。

個人的に良かった点

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

1. NVIDIAの創業期の話が知れた
2. 日本の半導体事業の衰退要因として考えられるものが何なのか知れた

1に関しては内容の紹介でも書いた通り、創業期でどのような考えがあり、どのような失敗があったのか、実はあまり知りませんでした。
この本ではそのようなエピソードも語られていて、非常に参考になりました。
また、この本はちょこちょこ日本の半導体事業の批判が書かれているのですが、その中でなぜ日本の半導体事業が衰退していったのか、アメリカと何が違ったのかが書かれています。このあたり、ちゃんと調べようと思ったことがなかったので、このような意見もあるのかと勉強になりました。

終わりに

今回は「NVIDIA 半導体の覇者が作り出す2040年の世界」という本を読んだのその紹介になります。NVIDIAは知っているようで知らない部分も多くあり、興味深い本でした。
定期的にこのように本の紹介記事を書いてます。もし興味があれば他の本の記事も見ていただければ幸いです。

The post [書評] NVIDIA 半導体の覇者が作り出す2040年の世界 ー NVIDIAという会社がどういう会社なのか知りたくなったら読べき一冊 first appeared on まったり勉強ノート.

どんな内容の本か？

子育てをしていると「あれもやったほうがいい、これもやったほうがいい」という話にあふれていて、結局何をやったらいいかわからない、ということはないでしょうか？私自身、子供が生まれてから子育て本を何冊か読んでいますが、結局どうすればいいのかわからなくなる時があります。
今回の「引き算の子育て」では「親が先回りしてやることなんてなにもない」ということを二人の有名な教育者がインタビュー形式で紹介してくれている本になります。
子育てはあれこれやってあげたくなることが多いですが、この本ではむしろ何かしてあげることにより、子供が家畜化してよくなく、むしろ余計なことはしないほうがいいということが書かれています。
また、このような考え方に至った経緯なども書かれていて、特に経緯がなかなかぶっ飛んでいるので、面白かったです。

どんな人にお勧めか？

今回は以下のような人におすすめな本になるかと思います。

1. 子育てについてあれこれ調べて悩んでいる方
2. 子供に何かしてあげたいと思っている方

特に1のようなあれこれ調べて結局どうしよう？ってなっている方が特にお勧めな気がします。
また、これから何かしてあげたいと思い始めている人も考え方の指針の一つとなると思うので読んでもよいかもしれないと思っています。
ただ、この本はインタビュー形式をとっていることもあり、よくあるハウツー的な教育本とは全然雰囲気が違うので、ハウツー本みたいな本を想像していると「思っているのと違った」となりそうなので、注意してください。

個人的に良かった点

個人的にこの本を読んで良かったこととしては以下の通りです。

1. あれこれ考える前に子供をよく見るという当たり前のことをちゃんとやろうという考えに至った
2. あれこれ悩んでも子供が自分ごとに考えなくて良くないという考え方を知った

1は当たり前といえば当たり前ですが、この本を読んで改めて思ったことです。結局子供はいろいろ違うので、その子はどうなのかを見てちゃんと考えることが重要なんだということを改めて感じました。
また、2は子育てじゃなくて部下の教育という文脈にも近い話ですが、結局あれこれ考える必要があるのは本人で、本人が考えることができる環境を用意することが重要なんだということを再確認しました。

終わりに

今回は同僚から紹介された「引き算の子育て」を読んだので、その記事になります。久しぶりにこの手のインタビュー形式の本を読んだおともあり、少し新鮮で面白かったです。
定期的にこのように本の紹介記事を書いてます。もし興味があれば他の本の記事も見ていただければ幸いです。

The post  [書評] 引き算の子育て ー 子育てで思い悩んでたらお勧めな一冊 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 まったり勉強ノート.

この「事前に育休に関していろいろアドバイスをもらっておく」というのは思った以上に役立ったので、また育休を取るかもしれない自分への備忘録兼これから育休を取る方に向けて参考になれば思い、私の育休の過ごし方などのまとめを書きました。いつもよりもかなり分量が多い気がしますが、ご了承ください。

これから育休取るという人の参考になれば幸いです。

そもそもなぜ育休を取るのか？

今務めている会社の男性社員は結構子供が生まれると育休を取るという人が多い印象です。上司も育休取ったほうがいいということで、自分も育休はとるものと思っていました。ただ、「漠然と育休を過ごしているだけだともったいないなぁ」ということを思ってました。このため、同僚の先輩パパさんにいろいろ育休中にやっておいたほうがいいことを聞いて回り、それを受けて個人的には以下のことを頑張るつもりで育休を過ごすことにしました。

1. 妻と協力して家事、育児を行う
2. 1日ワンオペで問題ないように育児スキルを磨く
3. しっかりと勉強時間を取る

1は育休取ったらやる理由かと思いますが、生まれたばかりの赤ちゃんがいると時間が無限に吸われていきます。結果、一人だと赤ちゃんのお世話だけで1日終わってしまって家事が何もできない、ということが発生してしまいます。このため、赤ちゃんのお世話と家事を分担しておいて最低限度の生活を維持できるようにすることがまず育休の目標の一つ目になると思っています。

次に2つ目の「1日ワンオペで問題ないように育児スキルを磨く」です。これは先輩パパさんの話や本の内容、また実際に過ごしてみて、実は育休の最大理由はこれでは？と思ったものです。先ほど育児で無限に時間が吸われるとう話をしましたが、当然これは一人分の時間で、夫婦二人が同時に育休でいる場合、一人分の時間は空きます。そして家事も最近の家電を駆使すればかなり楽ができます。このため、夫婦で同レベルまで育児スキルがある状態なら、一人が子供を見て、もうひとりは家事がひと段落した昼前からずっと勉強したり、外に遊びに行ってリフレッシュするということが可能になります。このような自由な時間がある程度確保できるとメンタル的にも育児が楽になります。

このため、1日まるまる育児を任せられるくらい夫婦二人がそれぞれ育児スキルをあげておくというのはかなり重要なポイントかなと思っています。

最後に「しっかりと勉強時間を取る」です。私は仕事がら最新技術についてちゃんと知らないと仕事で困ることが多々あります。また、まとまった時間があるなら今まで勉強してこなかったことも勉強したいなぁということで、育休中はしっかり勉強時間を取って勉強するという目標を個人的に掲げて頑張っていました。

育休をどれくらいの期間とるか？

次に育休をどれくらいの期間取るかです。この育休の期間に関して誰もいい感じの答えを教えてくれなかったので私は適当に決めてしまったのですが、今思うとちゃんと考えたほうがいいポイントでした。

会社によっては育休は1か月しかとれないとかあるかもしれませんが、うちの会社は経営陣（かつ今の自分の上司）が育児の大変さをわかってくれる人なので、たぶん、言えば結構長い期間育休取れた気がします。ただ、私は何も考えずに他の先輩パパさんが2か月くらいとっていたので、「自分も2か月とります」みたいに言ってしまいました。

これに関して後から本などを読んで勉強して失敗したと思ったので、もし仮に育休を長くとれる場合、今の私なら育休期間を検討する材料として何を考慮するか？を以下にまとめます。

1. 保育園はいつ入れるか？
2. 子供の夜のお世話の負担がどこまで減ったタイミングで育休を終わりにするか？

まず、保育園に関してです。当然ですが、育休で休んだままだと保育園に入れることができないケースが存在します。このため、保育園にいつ入れたいか？が決まっている人（たとえば枠が多い4月に入れたいとか）であれば、育休は4月付近までになります。ここは地域にルールがあるかもしれないので自分の地域はどうかを予め確認し、いつ保育園にいれるのかを夫婦で相談したほうが良いかと思います。

次に夜のお世話に関してです。赤ちゃんは生まれたばかりは3時間おきにミルクをあげる必要があります。このため、生まれたばかりは夜に一人体制だと何回も夜に起きる必要があってかなりしんどいです。1度やればわかると思いますが、つら過ぎて二度とやるか！って気分になります。成長すると赤ちゃんが連続して寝る時間は増えていくので、ある程度成長すると夜の負担も減っていきます。個人的にはこの夜の負担がある程度減ったあと以降に育休が終わると良いと思います。

月齢がどれくらいだと夜に何回起きるのかは後ほど紹介する「ベネッセ・ムック 最新！初めての育児新百科」や「赤ちゃんにもママにも優しい安眠ガイド」に書いてあるので育休をどれくらい取るかを考えるうえで参考にするとよいと思います。

基本的な１日のスケジュール

次に実際の1日の大まかなスケジュールに関してです。ただ、この部分はどういう家庭なのかで話が大分変るのでまずは基本情報について説明していきます。

基本的な１日のスケジュールを考えるうえで重要そうな基本情報

先ほど述べたように我が家の家族構成やどういう家庭の方針、どういう性格の子供かで話が全然違うと思われる方がいると思ったので、ここで基本情報をまとめて説明します。

1. 子供は一人目
2. 妻は育休で1年お休み（つまり私が育休中は二人体制）
3. 生まれてから１か月は里帰り、２か月目から3か月を少し過ぎたところまで育休
4. 引き続き妻が育休をとるため保育園はもう少し先
5. 母乳主体でミルクの混合で育てる (主に育児分担のため)
6. 子供はお腹が空いたときと眠い時以外はあまり泣かないので手がかからないほうらしい

まずは今回生まれた子供が一人目です。また、妻のほうは私が取る1か月前から育休を取っていて、妻は育休を1年取る予定なので、私が育休中のタイミングは二人体制になっていました。

そして、私の育休を取り始めたタイミングですが、妻が里帰り出産ということもあり、最初の1か月は妻は実家にいて、その後月齢で1か月になったあたり（年末年始の休みの最初）で帰ってきました。そして、このタイミングで育休をスタートさせて、2月末まで取ってました。このため、生まれたばかりの方は少し変えたほうがいい部分があるかもしれないので注意してください。

また、できるだけ長く子供と一緒にいる時間を作りたいという妻の意見で、保育園は少し先にすることにしました。このため、保活に必要な時間はスケジュールに含まれていません。

そして、ミルクに関してですが、これも妻の意見でなるべく母乳主体で育てつつ、私もごはんがあげられるようにということで母乳とミルクの混合で育てることにしました。このあと話す育児の分担では混合、もしくはミルクだけじゃないと参考にならない部分があると思いますので注意してください。

また、どういう性格の子供か？ですが、これに関して私が他の赤ちゃんを知らないので、妻のお姉さんや会社の同僚に話を聞いた感想を基に判断すると、どうやら手がかからないタイプの子供らしいです。育休のころは大体決まったタイミングで泣くので過去の状況から簡単に泣いている理由が判断できて、精神的に楽でした。

夜の二交代制による楽な1日の過ごし方

ここから育休中に重要だった1日の過ごし方について詳しく説明します。基本的には以下のような過ごし方をしていました。

子供が寝ている夜の時間（19時から次の日の6時）に対処するために我が家では、19時から3時と3時から11時の二つに分けて、一人が就寝、一人が家事と育児もしくは暇なら自由時間という形にしていました。夜を完全に分担することで夜は必ず連続で8時間寝れるという状況にしていました。ちなみにこの夜の二交代制は先輩パパさんたちが口をそろえてよかったというので我が家も導入しました。

この夜の二交代制の懸念事項は夜は一人体制になるので二人とも十分な育児スキルがないと任せられない、寝る時間と起きる時間が変な時間になるので抵抗があるということが考えられます。私の場合は育児スキルのほうは最初は妻に見てもらいながら練習をすることでなんとかしました。また、寝る時間と起きる時間が変な時間になるという部分は海外に行ったらもっとひどい状態になるので、それよりまし！と思ってぐっとこらえてました。

この夜の二交代制のメリットは睡眠時間がちゃんと確保できる以外にもあります。例えば、夜の育児に関しては最初のうちは3時間おきにミルクをあげる以外はかなり自由にできます。私は二人目の夜の担当だったので3時から11時が担当でした。このため、赤ちゃんが起きなければ3時から6時までの間は集中して勉強する時間に充てることができました。

また睡眠時間がちゃんと確保できているので、昼間も勉強する気力が十分あります。このため、昼間に関しては赤ちゃんを抱っこしながらでも「聞くこと」と「しゃべること」はできるので、これを機に今までやってなかったスピーキングの練習（ELSA Speakを利用）をやっていました。それ以外に家事を終えて暇な時間は論文や教科書を読むなどの読書によって勉強してました。

このため、1日で集中して勉強できる時間は2時間から3時間、スピーキングは1時間、読書は1,2時間くらいで合計で4~6時間という感じでした。結果として人生で久々にたくさん勉強することができたという感じがしています。

工夫した点や意識した点

ここではいくつか育休中の過ごし方で工夫したり意識したりした点について順番に紹介します。

育児の記録アプリを使う

まず最初の一つ目は育児アプリの使用です。我が家では妻の提案で育児記録アプリの「ぴよログ」を使っていました。

このおかげで夜に交代した際、「次に何時にミルクをあげるべきか？」などの引継ぎ作業がいらなくなり、すぐさま交代することができます。

また、育児の最中一番つらかったのが「なんで赤ちゃんが泣いているのかわからなくて対処できない時」でした。これに関しては育児記録があれば、泣いている時間いつもは寝てるのかやミルクを飲むのかなどの傾向を調べることができ、対処がしやすかったです。

ちなみに、ぴよログはカスタムのイベントをいくつか登録できるのですが、我が家では「ぐずった」というイベントを作って、泣いたタイミングとその後何をしたかが分かるようにしておいて、泣いたタイミングで何をしたほうがよさそうか？をわかるようにしてあります。

生活リズムを赤ちゃんを中心に考え直して毎日そろえる

二つ目は生活リズムに関してです。我が家では生活リズムをそろえることを意識して生活してました。スケジュールを毎日そろえる利点としては家事はもちろん、育児に関しても赤ちゃんの泣くタイミングが同じになり、仕事がルーティーン化でき、例外的な予定（友達と遊びに行くなど）が立てやすくなります。

どういうスケジュールがいいかに関しては後ほど紹介する「赤ちゃんにもママにも優しい安眠ガイド」という本を参考にしつつ、育児アプリの記録を見て夫婦で相談し、赤ちゃんにとっても夫婦にとっても無理のないスケジュールを立てるようにしていました（大まかな1日のスケジュールは「夜の二交代制による楽な1日の過ごし方」参照）。また、赤ちゃんの成長とともに赤ちゃんの生活リズムが変化してないかや、1度立てたスケジュールに無理がないかを定期的に確認するようにもしていました。

部屋のレイアウトを赤ちゃん中心に変える

3つ目は次に部屋のレイアウトです。我が家は1LDKなのですが、どこで赤ちゃんを生活させるか？は結構もめました。結果としては「夜に寝ている人がいる部屋に入らなくて済むように赤ちゃんのために使用する部屋を1つにまとめる」という考えのもとLDK部分に赤ちゃんがずっといるようにし、一部屋は寝室にしていました。

これにより、寝ている人が物音や部屋の明かりで起きるようなことがないようにし、快適に過ごすことができました。

できる限り日々のやることを細かく書きだして、妻と自分、どちらがやっても同じようにできるようにする

4つ目は日々のタスクに関してです。我が家では日々の家事と育児を細かく書きだして妻と自分、どちらがやっても同じようにできるようにしました。例えば、最初は「赤ちゃんのお風呂の準備」を夕飯前の準備前あたりでやるという風になっていました。ただ、お風呂の準備と聞いて私は何を用意するべきなのか？が最初イメージできませんでした。これでは「赤ちゃんのお風呂の準備」というタスクを取り組むのが妻が担当するか、私が担当するかで準備の内容が変わってしまいます。このため、単に「赤ちゃんのお風呂の準備」ではなく、「着替え、おむつ、タオルの準備」のように細かく書き下して、妻と私で日々のやることリストを共有していました。

この結果、どちらがやっても同じようになるため、家事と育児の分担がかなり柔軟に行うことができるようになります。また、1日ワンオペの日も何をやればいいかは書きだしてあるので、妻に確認しなくても普段通り過ごすことができます。このあたりのやることを書きだすというのは仕事のときと同じかなと思っています。

育休中できるだけすぐに1日ワンオペできるようにする

5つ目はワンオペに関してです。最初に述べたように「1日ワンオペ」は育休中の目標にしていました。このため、最初の1週間は妻に教えてもらいながら育児（特にお風呂の入れ方）を練習していました。このおかげで妻も安心して昼間出かけることができるようになり、ストレス発散ができるようになったということを言っています。

また、妻が昼間でかけたら自分も出かけたいということが言いやすくなるため、自分が快適に過ごすためにも重要だったと思っています。このおかげで気兼ねなく妻は月に何度か友達と外でランチや飲み会にいけたり、私も趣味のカフェ巡りやライブにいくことができました。

赤ちゃんがいる状態で何ができるか？を思いつくものは片っ端から確認する

6つ目は育休後を意識して赤ちゃんがいる状態で何ができるのかの確認に関してです。具体的には電車やバスに乗るとどうなるのかなどの確認です。私が育休中であれば二人で対応ができるため問題があったとしても対処することが可能なことが多い印象です。一方、私の育休後は妻が一人でやらなければならない可能性がでてきます。そういうときを意識して、これは一人でできそうか？などを確認していました。

この結果、例えばベビーカーでバスに一人で乗るのは大変そうだとかがわかりました。このようなことが事前にわかればベビーカーを使ってバスは乗らないようにするなど予定を立てる段階で役に立つと思っています。

育休ラスト約2週間前でに夜の二交代制をやめて問題ないか確認する

7つ目は育休後を意識した生活に関してです。育休明けが近づいてきた段階で夜の二交代制はやめて夜は二人とも同時に寝るように変更して問題がなさそうかを確認しました。夜の二交代制は8時間睡眠時間が必ず確保できる一方、仕事をするとなると実行するのは時間的に難しいです。このため、私の育休が明けるタイミングでは夜の二交代制をやめる必要があります。ただ、育休が明けたタイミングで突然やめると何か問題があったときに対処が難しいということで2週間前から夜の二交代制をやめて二人とも21時から次の日の5時か6時に起きるという生活を試し始めました。

ちなみにわが子はたまたまですが、すでに連続で5時間から長いときは8時間寝るようになっていたため、二交代制をやめても起きる回数は多くて1回くらいという状態になっていました。

ただ、いざ二交代制をやめても、当初は妻は赤ちゃんが布団を蹴り飛ばして寒そうにしていないか心配で何度も起きてしまうという話をしていた。このため「スリーパー」を購入したりして、心配事を潰していってできるだけ安心して眠れるように改善を繰り返すようにしていました。結局、ちょうど2週間くらいは睡眠不足でつらいという日々が続いたので、前もって夜の二交代制をやめた生活を試し始めてよかった印象です。

買ってよかったもの

ここでは育休中に買っておいて本当によかったというものをまとめて紹介しておきます。

ルンバ j7

Amazon

楽天市場

ポチップ

良かったポイント：細かいところまでは掃除してくれないが部屋が7，8割くらいきれいになるので、平日の掃除にはちょうどいい

パナソニック　ななめドラム洗濯乾燥機 NA-VG760L

Amazon

楽天市場

ポチップ

良かったポイント：乾燥機のおかげで干す必要がないのですごい洗濯が楽

コンビ 電子レンジ除菌&保管ケース 除菌じょ~ずα バニラ

Amazon

楽天市場

ポチップ

良かったポイント：哺乳瓶を3本くらいまとめて電子レンジで消毒できるので消毒が楽

ディズニー ベビートイ くまのプーさん えらべる回転 6WAYジムにへんしんメリー

Amazon

楽天市場

ポチップ

良かったポイント：ごはんを食べたいときなど30分程度の時間を子供から手を放したいときすごく重宝する

パナソニック KX-HBC200-W ベビーカメラ ホワイト KXHBC200

Amazon

楽天市場

ポチップ

良かったポイント：別の部屋や視覚的に見えない位置に赤ちゃんがいるときでもスマホで様子が確認できるので安心

丹平製薬 ママ鼻水トッテ

Amazon

楽天市場

ポチップ

良かったポイント：電動のものと違ってうるさくないので、夜に鼻を詰まらせて泣いているとき、さっと取り出せて使えるので便利

アメジスト ママとベビーの水だけぬれコットン

Amazon

楽天市場

ポチップ

良かったポイント：目ヤニがよくついているのでさっと取りたいときに便利

Medela (メデラ)　ピュアレーン 7g

Amazon

楽天市場

ポチップ

良かったポイント：妻曰く、母乳で育てるなら絶対必要らしい

他にも便利だったものを以下の表にまとめておきます。

といった感じでした。特に「ディズニー ベビートイ くまのプーさん えらべる回転 6WAYジムにへんしんメリー」は個人的には半信半疑で買いましたが、30分程度とはいえ赤ちゃんが泣かない時間を親が任意に作れるのですごい重宝しました。

読んで良かった本

最後に育休を取るにあたり、何冊か本を読んで読みました。ここではその中でも個人的によかったと思う本をよかった順にまとめて紹介します。

最新！初めての育児新百科 新生児期から3才までこれ1冊でOK！

ひよこクラブ編集部

Amazon

楽天市場

ポチップ

良かったポイント：赤ちゃんの成長による変化やお世話の仕方、病気っぽい症状のときどうすればいいかなど幅広いことが書いてあるので困ったときに手元にあると安心

赤ちゃんにもママにも優しい安眠ガイド 0歳からのネンネトレーニング [ 清水悦子 ]

清水 悦子, 神山 潤

Amazon

楽天市場

ポチップ

良かったポイント：1日のスケジュールをどうするか考えるのに非常に参考になった

新しいパパの教科書

NPO法人ファザーリング・ジャパン

Amazon

楽天市場

ポチップ

良かったポイント：父親になるってどういうことか、何を気を付けないといけないかなどの心構えを勉強するのによかった

出産・育児ママのトリセツ 「子どもができて妻が別人になりました」というあなた

山本ユキコ

Amazon

楽天市場

ポチップ

良かったポイント：子供だけでなく妻に対して気を付けないといけないことが何かを勉強するのによかった

特に最初の「ベネッセ・ムック 最新！初めての育児新百科」は育児に関して知っておいたほうがいいことを網羅的に写真つきで分かりやすく説明してくれている本で、育児初心者の自分にとても勉強になることが多かった印象です。

終わりに

今回は育休中の自分の過ごし方をまとめてみました。振り返ってみるといろいろあったのですが、ちゃんと先輩パパさんにいろいろ聞いたり本で勉強しておいてよかったということが多かった印象です。

この記事がこれから育休を取る方の参考になれば幸いです。

The post 育休前の自分に伝えたい育休の過ごし方 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 まったり勉強ノート.

この部分も私はぱっと読んで意味がわからなくて苦労した点なので、自分であとで見返してもわかるように丁寧な説明を心がけて書いていこうと思います。

想定する因果グラフ

この記事で想定している因果グラフは以下の通りです。

本では媒介公式の式導出の部分（自然な効果を識別するための条件の節）で\(W\)として抽象的な変数を使ってまとめています。この記事でも基本同様に\(W\)でまとめていますが、この因果グラフでは\(W_1, W_2, W_3\)と具体的にして分けています。

直接効果と間接効果の定義

まずはこれらの定義がどういうものかを見ていきます。本に書かれている定義は以下の通りです。

総合効果TE (total effect)

$$ \begin{align*}
TE &= E[Y_1 – Y_0] \\
&=E[Y|do(T=1) – E[Y|do(T=0)]]\tag{4.44}
\end{align*} $$

ここで\(Y_x\)は本ででてくる表記で\(X=x\)の場合の\(Y\)を\(Y_x\)ということを表しています。具体的には本の4.1の「反事実」の例で下道を使った場合とフリーウェイを使った場合の運転時間の例で説明します。このとき、フリーフェイを使うケースを\(T=1\)とするとこの時の運転時間\(Y\)を\(Y_{T=1}\)、\(T\)が自明な時は\(Y_1\)ということを表しています。一方、下道を使った場合\(T=0\)の運転時間は\(Y_{T=0}\)または\(Y_0\)です。

このTEは処置\(T=0\)から\(T=1\)へ変化したときの\(Y\)の増加量です。重要な点として媒介変数\(M\)は\(T\)が変化するにつれて変化します。

制御された直接効果CDE (controlled direct effect)

$$ \begin{align*}
CDE(m) &= E[Y_{1,m} – Y_{0,m}] \\
&=E[Y|do(T=1, M=m)] – E[Y|do(T=0, M=m)]\tag{4.45}
\end{align*} $$

\(Y_{1,m}\)の部分は省略されていますが、\(Y_{T=1, M=m}\)の意味です。基本的には統合効果のときに説明したときとほぼ同じで、ここでは\(T\)だけでなく\(M\)についても固定した状態の\(Y\)の値を示しています。

このCDEは処置\(T=0\)から\(T=1\)への変化したときのYの増加量の期待値ですが、重要な点として媒介変数\(M\)は\(M=m\)で一律に設定されています。

自然な直接効果NDE (natural direct effect)

$$ \begin{align*}
NDE &= E[Y_{1,M_0} – Y_{0,M_0}] \tag{4.46}
\end{align*} $$

このNDEは処置が\(T=0\)から\(T=1\)へ変化したときの\(Y\)の増加量の期待値ですが、重要な点として媒介変数は変化前の状態、つまりT=0においてそれぞれの個体がとったであろう値(\(M_0\))に固定されています。この\(M_0\)は\(M_{T=0}\)を表していて、\(M\)が\(T=0\)のときの値を示しています。この媒介変数の値が一律で固定されているケースがCDEで、変化前の状態で個別にとったであろう値に固定されているのがNDEになります。

自然な間接効果NIE (natural indirect effect)

$$ \begin{align*}
NIE &= E[Y_{0,M_1} – Y_{0,M_0}] \tag{4.47}
\end{align*} $$

このNIEでは処置が\(T=0\)で固定されているときのYの増加量の期待値です。このとき媒介変数\(M\)は\(T=1\)においてそれぞれの個体がとったであろう値に変化させます。よって、\(Y\)が\(T\)で反応しないようにしたうえで媒介変数によって変化する量を表しています。

NDEとNIEを識別する十分条件 (自然な効果を識別するための条件)

先ほどのNDEとNIEまで計算するための十分条件として本には以下のように書かれています。

1. \(W\) の要素(最初の因果グラフでは\(W_1, W_2, W_3\))には\(T\)の子孫が一つもない
2. \(W\) は(\(T\rightarrow M\)と\(T\rightarrow Y\)を除いた)\(M\)から\(Y\)へのバックドアパスをすべてブロックする
3. \(T\)から\(M\)への特定\(W\)同時効果は(おそらく実験か調整により)識別可能である
4. \(\{T, M\}\)から\(Y\)への特定\(W\)同時効果は(おそらく実験か調整により)識別可能である

ここで\(W\)は因果グラフにおける\(\{W_1, W_2, W_3\}\)をまとめた変数になります。

これらの条件は言葉で言わんとしていることはわかるきがしますが、これを使って私はNDEとNIEの媒介公式を導出できるきがしませんでした。このため、別の表現がないかな？ということで探すと、別の論文[1]に別の表現の条件が書かれていたので今回はこの条件を使って証明していきます。

まず、一致性 (consistency assumption) と構成性 (composition assumption)という以下の仮定が成り立つとします。

• 一致性: \(T=t\)と\(M=m\)が条件付けされたとき\(Y=Y_{t, m}\)が成り立つ。同様に\(T=t\)が条件付けされたとき\(M=M_t\)が成り立つ
• 構成性: \(T=t\)が条件付けされたとき\(Y(t) = Y_{t, M(t)}\)が成り立つ

そしてNDEとNIEが識別可能とする十分条件として以下のものがあるとしています。

$$ \begin{align*}
Y_{t, m} & \perp\kern-5pt\perp T \mid W & \forall t, m \tag{A-1} \\
Y_{t, m} & \perp\kern-5pt\perp M \mid T, W & \forall t, m \tag{A-2} \\
M_t & \perp\kern-5pt\perp T \mid W & \forall t \tag{A-3} \\
Y_{t, m} & \perp\kern-5pt\perp M_{t^{\star}} \mid W & \forall t, t^{\star}, m \tag{A-4} \\
\end{align*} $$

変数に関しては今回の因果グラフに合わせて元論文から変更しています。

各効果の定義から式変形していく

ここから本では省略されている式変形をしていって、観測データから計算できる形に式を持っていきたいと思います。定義では期待値の形で書かれていますが、確率がわかれば十分だと思うので、確率の形で式変形していきます。TEに関しては３章で示されているdo演算子を使った場合の式変形を行えば目的の式まで変形できますが、他の効果の定義に合わせて先ほどの仮定と条件を使って式変形していきます。

ちなみに３章で示されているdo演算子を使った場合の式変形はこちらで詳しく書いていますので気になる方は見てください。

TE

TEの定義は以下の通りです。

$$ \begin{align*}
TE &= E[Y_1 – Y_0] \\
&=E[Y|do(T=1)] – E[Y|do(T=0)]]\tag{4.44}
\end{align*} $$

ここで\(P(Y|do(X))\)からdo演算子のない形にするというのでもいいのですが、先ほどの仮定と条件を使うために\(Y_1, Y_0\)から式変形していきます。ここで\(Y_1, Y_0\)をまとめて\(Y_t\)として計算していきます。

最終的には以下のような式になることを示します。

$$ \begin{align*}
P(Y_t) &= \sum_w P(Y|T=t, W=w)P(W=w) \tag{4.44.1}
\end{align*} $$

まず、構成性から以下の式に変形できます。

$$ \begin{align*}
P(Y_t) &= P(Y_{t, M_t}) \tag{4.44.2}
\end{align*} $$

ここで、確率の周辺化で以下の式が成り立つことを利用します。

$$ \begin{align*}
P(Y) =& \sum_x P(Y, X=x) \tag{4.44.3} \\
\end{align*} $$

この周辺化の式を使って\(W\)と\(M_t\)の同時確率の式に変形します。そして、その後、\(W\)と\(M_t\)の条件つき確率の形にもっていきます。

$$ \begin{align*}
P(Y_t) =& P(Y_{t, M_t}) \\
=& \sum_w P(Y_{t, M_t}, W=w) \\
=& \sum_w \sum_m P(Y_{t, M_t}, M_t=m, W=w) \\
=& \sum_w \sum_m P(Y_{t, M_t}, M_t=m|W=w)P(W=w) \\
=& \sum_w \sum_m P(Y_{t, M_t}| M_t=m, W=w)\\
& \qquad P(M_t=m|W=w)P(W=w) \tag{4.44.4}
\end{align*} $$

ここで\(M_t=m\)を利用すると以下のように変形できます。

$$ \begin{align*}
P(Y_{t, M_t}| M_t=m, W=w) &= P(Y_{t, m}| M_t=m, W=w) \tag{4.44.5}
\end{align*} $$

次に十分条件の中のA-1を利用すると式(4.44.5)は\(T=t\)を使った条件付き確率に変形できます。

$$ \begin{align*}
P(Y_{t, m}| M_t=m, W=w) &= P(Y_{t, m}| M_t=m, T=t, W=w) \tag{4.44.6}
\end{align*} $$

そして一致性を使うと\(P(Y_{t, m}| M_t=m, W=w)\)は以下のようになります。

$$ \begin{align*}
P(Y_{t, m}| M_t=m, T=t, W=w) &= P(Y| M_t=m, T=t, W=w) \tag{4.44.7}
\end{align*} $$

よって式(4.44.5)から(4.44.7)を利用すると式(4.44.4)は以下のように変形できます。

$$ \begin{align*}
P(Y_t) =& \sum_w \sum_m P(Y_{t, M_t}| M_t=m, W=w) \\
& \qquad P(M_t=m|W=w)P(W=w) \\
=& \sum_w \sum_m P(Y_{t, m}| M_t=m, W=w) \\
& \qquad P(M_t=m|W=w)P(W=w) \\
=& \sum_w \sum_m P(Y_{t, m}| M_t=m, T=t, W=w)\\
& \qquad P(M_t=m|W=w)P(W=w) \\
=& \sum_w \sum_m P(Y| M_t=m, T=t, W=w)\\
& \qquad P(M_t=m|W=w)P(W=w) \tag{4.44.8}
\end{align*} $$

最後に\(M_t=m\)の条件付き確率を同時確率の形に戻して、周辺化を適用して\(M_t=m\)の部分を取り除くと以下のようになります。

$$ \begin{align*}
P(Y_t) =& \sum_w \sum_m P(Y| M_t=m, T=t, W=w)\\
& \qquad P(M_t=m|W=w)P(W=w) \\
=& \sum_w \sum_m P(Y, M_t=m| T=t, W=w)P(W=w) \\
=& \sum_w P(Y| T=t, W=w)P(W=w) \tag{4.44.9}
\end{align*} $$

これで求めたい式(4.44.1)が導出できました。

CDE

次にCDEについても式変形していきます。CDEの定義は以下の通りです。

$$ \begin{align*}
CDE(m) &= E[Y_{1,m} – Y_{0,m}] \\
&=E[Y|do(T=1, M=m)] – E[Y|do(T=0, M=m)]\tag{4.45}
\end{align*} $$

これをTEのときと同じようにdo演算子がない状態に\(Y_{1,m}\)と\(Y_{0,m}\)を式変形していきたいと思います。\(Y_{1,m}\)と\(Y_{0,m}\)をまとめて\(Y_{t,m}\)という形で式変形していきます。

CDEの場合は最終的に以下の式の導出を目指します。

$$ \begin{align*}
P(Y_{t,m}) &= \sum_w P(Y|M=m, T=t, W=w)P(W=w) \tag{4.45.1}
\end{align*} $$

それでは\(P(Y_{t,m})\)を周辺化の式を使って\(W\)の同時確率の形にして、その後\(W\)の条件付き確率の形にします。

$$ \begin{align*}
P(Y_{t,m}) &= \sum_w P(Y_{t,m}, W=w) \\
&= \sum_w P(Y_{t,m}|W=w)P(W=w)\tag{4.45.2}
\end{align*} $$

次に\(P(Y_{t,m}|W=w)\)に注目するとA-1の条件から\(W\)が与えられると\(Y_{t,m}\)と\(T\)は独立なため、\(P(Y_{t,m}|W=w)\)を\(T=t\)の条件付き確率にできます。

$$ \begin{align*}
P(Y_{t,m}|W=w) &= P(Y_{t,m}|T=t, W=w) \tag{4.45.3}
\end{align*} $$

さらにA-2の条件から\(W\)と\(T\)が与えられると\(Y_{t,m}\)と\(M\)は独立なため、\(P(Y_{t,m}|T=t, W=w)\)を\(M=m\)の条件付き確率にできます。

$$ \begin{align*}
P(Y_{t,m}|T=t, W=w) &= P(Y_{t,m}|M=m, T=t, W=w) \tag{4.45.4}
\end{align*} $$

ここで一致性から以下のように変形できます。

$$ \begin{align*}
P(Y_{t,m}|M=m, T=t, W=w) &= P(Y|M=m, T=t, W=w)\tag{4.45.5}
\end{align*} $$

最後に式(4.45.2)を式(4.45.3)から式(4.45.5)を利用して変形していきます。

$$ \begin{align*}
P(Y_{t,m}) &= \sum_w P(Y_{t,m}, W=w) \\
&= \sum_w P(Y_{t,m}|W=w)P(W=w) \\
&= \sum_w P(Y_{t,m}|T=t, W=w)P(W=w) \\
&= \sum_w P(Y_{t,m}|M=m, T=t, W=w)P(W=w) \\
&= \sum_w P(Y|M=m, T=t, W=w)P(W=w) \tag{4.45.6}
\end{align*} $$

以上でCDEのほうも導出することができました。

NDEとNIE

最後にNDEとNIEです。それぞれ以下のような定義になっています。

$$ \begin{align*}
NDE &= E[Y_{1,M_0} – Y_{0,M_0}] \tag{4.46} \\
NIE &= E[Y_{0,M_1} – Y_{0,M_0}] \tag{4.47} \\
\end{align*} $$

これは\(Y_{1,M_0}, Y_{0,M_0}, Y_{0,M_1}, Y_{0,M_0}\)の４種類の確率が分かればよいことがわかります。これらの変数をまとめて\(Y_{t_y,M_{t_m}}\)としてこの後まとめて式変形していきます。最終的には以下の式に変形することを目指します。

$$ \begin{align*}
P(Y_{t_y,M_{t_m}}) =& \sum_w \sum_m P(Y|M=m, T=t_y, W=w)\\
& \qquad P(M|T=t_m, W=w)P(W=w) \tag{4.46.1}
\end{align*} $$

まずは\(W\)の同時確率の形にして、その後\(W\)の条件付き確率の形に変形します。

$$ \begin{align*}
P(Y_{t_y,M_{t_m}}) &= \sum_w P(Y_{t_y,M_{t_m}}, W=w) \\
&= \sum_w P(Y_{t_y,M_{t_m}}|W=w)P(W=w) \tag{4.46.2}
\end{align*} $$

次に\(P(Y_{t_y,M_{t_m}}|W=w)\)の部分に注目して\(M_{t_m}\)に関しても同様に周辺確率の形にしてその後\(M_{t_m}\)の条件付き確率の形に変形します。

$$ \begin{align*}
P(Y_{t_y,M_{t_m}}|W=w) &= \sum_m P(Y_{t_y,M_{t_m}}, M_{t_m}=m|W) \\
&= \sum_m P(Y_{t_y,M_{t_m}}|M_{t_m}=m, W=w)P(M_{t_m}=m|W) \tag{4.46.3}
\end{align*} $$

ここで\(M_{t_m}=m\)を利用すると以下のように変形できます。

$$ \begin{align*}
P(Y_{t_y,M_{t_m}}|W=w) &= \sum_m P(Y_{t_y,M_{t_m}}|M_{t_m}=m, W=w)P(M_{t_m}=m|W) \\
&= \sum_m P(Y_{t_y,m}|M_{t_m}=m, W=w)P(M_{t_m}=m|W) \tag{4.46.4}
\end{align*} $$

ここから\(P(M_{t_m}|W)\)の部分の式変形をしていきます。十分条件のA-3から\(W\)が条件付けされたとき、\(M_t\)と\(T\)は独立なので以下が成り立ちます。

$$ \begin{align*}
P(M_{t_m}|W) &= P(M_{t_m}|T=t_m, W) \tag{4.46.5}
\end{align*} $$

そして一致性から以下のようになります。

$$ \begin{align*}
P(M_{t_m}|W) &= P(M_{t_m}|T=t_m, W) \\
&= P(M|T=t_m, W) \tag{4.46.6}
\end{align*} $$

次に式(4.46.4)の\(P(Y_{t_y,m}|M_{t_m}=m, W)\)も式変形していきます。十分条件のA-4を利用すると\(W\)が条件付けされたとき、\(Y_{t_y,m}\)と\(M_{t_m}\)は独立なので以下のようになります。

$$ \begin{align*}
P(Y_{t_y,m}|M_{t_m}=m, W) &= P(Y_{t_y,m}|W) \tag{4.46.7}
\end{align*} $$

そして十分条件のA-1を使うと\(W\)が条件付けされたとき、\(Y_{t_y,m}\)と\(T\)は独立なので以下のようになります。

$$ \begin{align*}
P(Y_{t_y,m}|W=w) &= P(Y_{t_y,m}|T=t_y, W) \tag{4.46.8}
\end{align*} $$

同様に十分条件のA-2を使うと\(T\)と\(W\)が条件付けされたとき、\(Y_{t_y,m}\)と\(M\)は独立なので以下のようになります。

$$ \begin{align*}
P(Y_{t_y,m}|T=t_y, W) &= P(Y_{t_y,m}|M=m, T=t_y, W)\tag{4.46.9}
\end{align*} $$

ここで一致性から以下のようになります。

$$ \begin{align*}
P(Y_{t_y,m}|M=m, T=t_y, W) &= P(Y|M=m, T=t_y, W) \tag{4.46.10}
\end{align*} $$

式(4.46.7)から(4.46.10)までを利用すると最終的に以下のような関係が成り立ちます。

$$ \begin{align*}
P(Y_{t_y,m}|M_{t_m}=m, W) &= P(Y|M=m, T=t_y, W) \tag{4.46.11}
\end{align*} $$

式(4.46.6)と式(4.46.11)を使うと式(4.46.4)は以下のようになります。

$$ \begin{align*}
P(Y_{t_y,M_{t_m}}|W=w) =& \sum_m P(Y_{t_y,M_{t_m}}|M_{t_m}=m, W=w)\\
& \qquad P(M_{t_m}=m|W) \\
=& \sum_m P(Y_{t_y,m}|M_{t_m}=m, W=w)\\
& \qquad P(M|T=t_m, W) \\
=& \sum_m P(Y|M=m, T=t_y, W)\\
& \qquad P(M|T=t_m, W) \tag{4.46.12}
\end{align*} $$

最終的に式(4.46.12)を使うと式(4.46.2)は以下のようになります。

$$ \begin{align*}
P(Y_{t_y,M_{t_m}}) =& \sum_w P(Y_{t_y,M_{t_m}}, W=w) \\
=& \sum_w P(Y_{t_y,M_{t_m}}|W=w)P(W=w) \\
=& \sum_w \sum_m P(Y|M=m, T=t_y, W=w)\\
& \qquad P(M|T=t_m, W=w)P(W=w) \tag{4.46.13}
\end{align*} $$

これで求めたい式(4.46.1)の導出ができました。

終わりに

今回の記事は入門 統計的因果推論 の4.5.2に書かれているTE, CDE, NDE, NIEについて式変形をしていきました。一度理解してblogにかけそうと思ってからいざ書き始めたら、blogの記事を書く中でところどころ自分の理解が甘い部分が発見できました。このためblogの記事を書くことで自分の理解の再確認ができて非常に有意義でした。

他にもわからない部分がいくつかまだあるので、機会があれば勉強のために記事にしていきたいと思います。

参考文献

1. VanderWeele, Tyler J. and Stijn Vansteelandt. “Conceptual issues concerning mediation, interventions and composition.” Statistics and Its Interface 2 (2009): 457-468.

The post [勉強ノート] 入門 統計的因果推論 反事実における直接効果と間接効果、媒介公式の導出 (4.5.2 媒介についてのツール) first appeared on まったり勉強ノート.

前提として「入門 統計的因果推論」で事前に説明された用語に関して、わかりやすくて説明が不要な部分は説明せずに使っているので、この記事だけではおそらく内容がわからない部分もあるかもしれませんがご了承ください。

調整化公式の導出 (3.2 調整)

まず、新薬の効果を表す問題を考えます。考える因果グラフは以下の通り。

この問題において以下の平均因果効果（ACE: average causal effect）が知りたいことです。

$$ \begin{align*}
ACE=P(Y = 1 | do(X = 1)) – P(Y = 1 | do(X = 0))
\end{align*} $$

この式から\( P(Y| do(X)) \) の確率を観測データから計算できればACEの計算ができることがわかります。このため、これ以降は\( P(Y| do(X)) \)を観測データから計算する方法について説明していきます。具体的には以下の形に変形することを目指します。

$$ \begin{align*}
P(Y | do(X)) &= \sum_z P(Y|X, Z=z)P(Z=z) \\
\tag{3.1}
\end{align*} $$

この式が調整化公式と呼ばれるものです。

ここで、\( do()\) はdo演算子といいます。これは変数\(X\) が自然に値\(x\)をとる場合と \( X=x \) に固定する場合を区別して表記したいという理由から、値を固定する後者の場合を\( do(X=x) \) と表記するために用いています。例示した図3.1の因果グラフにおいては、\(Z\)に依存した何等かの理由により薬を使った人のうち回復する確率を \( P(Y = 1 | X = 1) \) 、\(Z\)によらず、全員に対して薬を使うように固定した場合の回復する確率を \( P(Y = 1 | do(X = 1)) \)と表しています。固定して薬を使うようなケースは介入試験で全員に薬を使った場合などを指します。

言葉だとわかりづらいと思うので因果グラフでも示します。\(do(X=x)\) というように\(X \) を固定した場合の修正後の因果グラフは以下のようになります。

この因果グラフにおける確率を本にならって\(P_m\) と表記します。この修正後の因果グラフの確率\(P_m\)を利用すると\( P(Y | do(X)) \) は以下のように表すことができます。

$$ \begin{align*}
P(Y | do(X)) = P_m(Y|X) \tag{3.2} \\
\end{align*} $$

この後はこの\(P(Y | do(X)) \)が式変形により\(P_m\)に依存しない形に式変形できることを示します。これができると介入試験のデータを使わずに観察データから\(P(Y | do(X)) \)が求められることとになります。

\(P(Y | do(X)) \)の式変形の前に後ほど式変形で重要な関係を以下に列挙しておきます。

$$ \begin{align*}
P_m(Z) =& P(Z) \tag{3.3} \\
P_m(Y|X, Z) =& P(Y|X, Z) \tag{3.4} \\
\end{align*} $$

式(3.3)は\(Z\) が\(X\)の介入による影響がないため、式(3.4)は\(X\)と\(Z\)が条件付けされた状況下において、\(Y\)の確率は修正前後で変化しないため成り立ちます。

それでは\(P(Y | do(X)) \)の式変形をしていきます。まず、先ほど示した通り、以下のように\(P_m\)を利用すると以下のように表すことができます。

$$ \begin{align*}
P(Y | do(X)) = P_m(Y|X) \tag{3.2} \\
\end{align*} $$

ここで、確率の周辺化が以下のような式で行えることを利用します。

$$ \begin{align*}
P(Y) =& \sum_x P(Y, X=x) \tag{3.5} \\
\end{align*} $$

この周辺化の式を利用して式(3.2)を\(P_m(Y, Z|X) \)の同時確率の式で書きます。

$$ \begin{align*}
P(Y | do(X)) &= P_m(Y|X) \\
&= \sum_z P_m(Y, Z=z | X) \tag{3.6}
\end{align*} $$

次に\(P_m(Y, Z | X)\)の部分を条件付き確率の定義を利用します。条件付き確率の定義より\( P_m(Y|X, Z) \)の部分は以下のようになります。

$$ \begin{align*}
P_m(Y|X, Z) &= \frac{P_m(Y, Z | X)}{P_m(Z|X)} \tag{3.7}
\end{align*} $$

後ほどの説明で分かりやすいように式(3.7)の左右を入れ替えて、両辺を\(P_m(Z|X)\)でかけると以下の式になります。

$$ \begin{align*}
\frac{P_m(Y, Z | X)}{P_m(Z|X)} &= P_m(Y|X, Z) \\
\frac{P_m(Y, Z | X)P_m(Z|X)}{P_m(Z|X)} &= P_m(Y|X, Z)P_m(Z|X) \\
P_m(Y, Z | X) &= P_m(Y|X, Z)P_m(Z|X) \tag{3.8}
\end{align*} $$

ここで修正後の因果グラフにおいて\(Z\)は\(X\)と独立なことを利用すると以下のように変形できます。

$$ \begin{align*}
P_m(Z|X) = P_m(Z) \tag{3.9}
\end{align*} $$

式(3.3)、 (3.4)、(3.9) を利用すると式(3.8)は以下のように変形できます。

$$ \begin{align*}
P_m(Y, Z | X) &= P_m(Y|X, Z)P_m(Z|X) \\
&= P_m(Y|X, Z)P_m(Z) \\
&= P_m(Y|X, Z)P(Z) \\
&= P(Y|X, Z)P(Z) \\
\tag{3.10}
\end{align*} $$

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

$$ \begin{align*}
P(Y | do(X)) &= \sum_z P_m(Y, Z=z | X) \\
&= \sum_z P(Y|X, Z=z)P(Z=z) \\
\tag{3.11}
\end{align*} $$

以上の結果から式(3.1)を導出できました。

調整化公式を傾向スコアを使った形に変形する (3.2.1 調整すべきか否か)

先ほど示した調整化公式を因果推論でよく使われる傾向スコアを使った以下の形の式に変形します。

$$ \begin{align*}
P(Y | do(X)) &= \sum_z \frac{P(X, Y, Z=z)}{P(X|Z=z)} \\
\tag{3.12}
\end{align*} $$

ここで\(P(Z|X)\)が傾向スコアです。

では、調整化公式の式(3.1)を式(3.12)の形を目指して変形していきます。まずは右辺の\(P(Y|X, Z=z)P(Z=z)\)に注目します。ここで、\(P(X|Z)/P(X|Z)\)をかけます。

$$ \begin{align*}
P(Y|X, Z)P(Z) &= \frac{P(Y|X, Z)P(X|Z)P(Z)}{P(X|Z)} \\
\tag{3.13}
\end{align*} $$

ここで分子の部分について条件付き確率の定義から\(X\)と\(Z\)を条件付けせずに同時確率の形で書くことができます。具体的には以下のように変形できます。

$$ \begin{align*}
P(Y|X, Z)P(X|Z)P(Z) &= P(X, Y|Z)P(Z) \\
&= P(X, Y, Z) \\
\tag{3.14}
\end{align*} $$

この式(3.13)、(3.14)を利用すると式(3.11)は以下のように変形できます。

$$ \begin{align*}
P(Y | do(X)) &= \sum_z P(Y|X, Z=z)P(Z=z) \\
&= \sum_z \frac{P(Y|X, Z=z)P(X|Z=z)P(Z=z)}{P(X|Z=z)} \\
&= \sum_z \frac{P(X, Y, Z=z)}{P(X|Z=z)} \\
\tag{3.15}
\end{align*} $$

これで求めたかった式(3.12)を導出することができました。

フロントドア公式の導出 (3.4 フロントドア基準)

次にフロントドア公式の導出について説明していきます。ここでは以下の因果グラフを考えます。

\(U\)は観測されていない交絡因子なので、この\(U\)を使わない形で\(X\)の効果、ここでは喫煙の肺がんへの影響があるかを考えます。これを考えるうえで、変数\(Z\)、ここではタール蓄積が観測できているとします。そして、ゴールとして因果効果を測るうえで、\(P(Y|do(X))\) という確率が求められれば良いと考えます。

本のほうの説明ではバックドア基準を２度にわたって適用すればよいと書かれています。しかし、この意味がなんとなくでしか私は理解できなかったので、私なりに\(P(Y|do(X))\)を式変形していってフロントドア公式を導出していきます。具体的には先ほどの因果グラフにおいて以下の式の導出を行います。

$$ \begin{align*}
P(Y | do(X=x)) &= \sum_z P(Y|do(Z=z))P(Z=z | do(X=x)) \\
&= \sum_z \sum_{x^{\prime}} P(Y|X=x^{\prime}, Z=z)P(Z=z | X=x) \tag{3.16}
\end{align*} $$

ではフロンドドア公式の導出を始めます。まずは式(3.16)の１行目のdo演算子を使った式を目指して\(P(Y|do(X))\)を変形していきます。まず、\(P(Y|do(X))\)を周辺化の式(3.6)を利用して\(P(Y, Z|do(X))\)という同時確率を使った式に変形し、その後、\(Z\)で条件付けした確率に変形させます。

$$ \begin{align*}
P(Y | do(X)) &= \sum_z P(Y, Z=z|do(X)) \\
&= \sum_z P(Y|do(X), Z=z)P(Z=z | do(X)) \tag{3.17}
\end{align*} $$

ここで、因果グラフより\(X\)から\(Y\)につながるパスは\(X \rightarrow Z \rightarrow Y \)しかありません。このため、\(Z\)で条件付けされていると変数\(Y\)は\(X\)とは独立となっています。このため、\(P(Y|do(X), Z=z)\)の部分は以下のようになります。

$$ \begin{align*}
P(Y|do(X), Z) &= P(Y|Z) \tag{3.18}
\end{align*} $$

次に、周辺化の式を利用して\(P(Y|Z)\)を\(X\)を加えた同時確率で表現したあと、\(X\)の条件付き確率の形を使って表します。

$$ \begin{align*}
P(Y|Z) &= \sum_{x^{\prime}} P(X=x^{\prime}, Y|Z) \\
&= \sum_{x^{\prime}} P(Y|X=x^{\prime}, Z)P(X=x^{\prime}) \tag{3.18}
\end{align*} $$

次に\(P(Y|X, Z)\)を変形していきます。ただ、ここからだと分かりにくいので\(P(Y|do(Z))\)から式変形していって、\(P(Y|X, Z)\)になることを示します。基本の流れは調整化公式の導出と同じ流れです。

まず、わかりやすいように\(do(Z)\)で修正された因果グラフを示します。

この因果グラフにおける確率を\(P_{z}\) とすると\(P(Y|do(Z))\)を\(P_{z}\)で表したあと、確率の周辺化と条件付きの定義を利用して変形していくと以下の通りです。

$$ \begin{align*}
P(Y|do(Z)) &= P_z(Y|Z) \\
&= \sum_{x^{\prime}} P_z(Y, X=x^{\prime}|Z) \\
&= \sum_{x^{\prime}} P_z(Y|X=x^{\prime}, Z)P_z(X=x^{\prime}|Z) \tag{3.19}
\end{align*} $$

また、修正前後の因果グラフより以下が成り立ちます。

$$ \begin{align*}
P_z(X|Z) &= P_z(X) \tag{3.20} \\
P_z(X) &= P(X) \tag{3.21} \\
P_z(Y|X, Z) &= P(Y|X, Z) \tag{3.22} \\
\end{align*} $$

この式(3.20)から(3.22)を利用して式(3.19)の\(P_z(Y|X, Z)P_z(X|Z) \)の部分を以下のように変形します。

$$ \begin{align*}
P_z(Y|X, Z)P_z(X|Z) &= P_z(Y|X, Z)P_z(X) \\
&= P_z(Y|X, Z)P(X) \\
&= P(Y|X, Z)P(X) \tag{3.23}
\end{align*} $$

この式(3.23)を利用すると式(3.19)は以下のようになります。

$$ \begin{align*}
P(Y|do(Z)) &= \sum_{x^{\prime}} P_z(Y|X=x^{\prime}, Z)P_z(X=x^{\prime}|Z) \\
&= \sum_{x^{\prime}} P(Y|X=x^{\prime}, Z)P(X=x^{\prime})
\tag{3.24}
\end{align*} $$

この式(3.24)と式(3.19)を利用すると式(3.18)は以下のようになります。

$$ \begin{align*}
P(Y|do(X), Z) &= P(Y|Z) \\
&= \sum_{x^{\prime}} P(Y|X=x^{\prime}, Z)P(X=x^{\prime}) \\
&= P(Y|do(Z)) \tag{3.25}
\end{align*} $$

よって式(3.25)を利用して式(3.17)を変形するとまず求めたかった式(3.16)のdo演算子を使った形になります。

$$ \begin{align*}
P(Y | do(X)) &= \sum_z P(Y|do(X), Z=z)P(Z=z | do(X)) \\
&= \sum_z P(Y|do(Z=z))P(Z=z | do(X)) \tag{3.26}
\end{align*} $$

ここまでで、式(3.16)の1行目のdo演算子を使った式ができました。ここから、do演算子のない形に変形していきます。

まず、式(3.24)の式を使い\(do(Z)\)を使わない形に式(3.26)を変形します。ここから\(X\)の値が2種類でてくるので区別できるように\(x\)と\(x^{\prime}\)の二つに分けます。

$$ \begin{align*}
P(Y | do(X=x)) &= \sum_z P(Y|do(Z=z))P(Z=z | do(X=x)) \\
&= \sum_{x^{\prime}} P(Y|X=x^{\prime}, Z)P(X=x^{\prime})P(Z=z | do(X=x)) \\
\tag{3.27}
\end{align*} $$

次に\(P(Z=z | do(X=x))\)の部分で\(do(X)\)がない形に変形することを目指します。ここで元の因果グラフを見ると\(X\) から\(Z\)につながるパスは\(X \rightarrow Z \)しかないため、因果グラフから以下の等式が成り立ちます。

$$ \begin{align*}
P(Z| do(X)) &= P(Z|X)
\tag{3.28}
\end{align*} $$

（\(do(X)\)によって修正された因果グラフを考えて\(P(Z=z | do(X=x))\)を修正後の因果グラフの確率から観測データから計算できる確率にするにはどうすればよいのか考えてもいいのですが、自明なので省略します。）

この式(3.28)を利用すると式(3.27)は以下のように変形できます。

$$ \begin{align*}
P(Y | do(X=x)) &= \sum_{x^{\prime}} P(Y|X=x^{\prime}, Z)P(X=x^{\prime})P(Z=z | do(X=x)) \\
&= \sum_{x^{\prime}} P(Y|X=x^{\prime}, Z)P(X=x^{\prime})P(Z=z | X=x) \tag{3.29}
\end{align*} $$

この式は求めたかった式（3.16）の式になっています。これでフロントドア公式を導出することができました。

終わりに

今回の記事では入門 統計的因果推論 の３章の部分で私がぱっとわからなかった部分の式変形を中心に説明する記事を書きました。本当はCDE (Controlled direct effect) (本で3.7 媒介で説明されていた部分)もぱっとわからなかったのですが、ここに関しては４章の内容とまとめて説明したほうが分かりやすいので、４章の記事のときに説明しようと思います。

この記事がこの本を読む他の人にも役に立てば幸いです。

The post [勉強ノート] 入門 統計的因果推論 3章 介入効果　 first appeared on まったり勉強ノート.