Google Colabの無料枠で推論を回す方法は前回記事にしましたので、そもそもPLaMo 2 1Bって何と思った方や推論を回してみたいという方はそちらをご覧ください。

また、今回説明に使うコードはこちらに置いてありますので、適宜参照してください。

https://github.com/shu65/plamo-2-1b-sft-example

Google Colabにおける一連の実行に関してはJupyter Notebookにまとめてありますので、細かい実行方法がわからないという方はこちらをご覧ください

https://github.com/shu65/plamo-2-1b-sft-example/blob/main/run_sft_google_colab.ipynb

Supervised Fine-Tuning(SFT)とは？

SFTを知らない方に簡単に説明すると、SFTは指示と想定されている回答のペアを用意し、LLMに対して学習を行い、指示に従いやすいモデルを作る方法になります。

特にPLaMo 2 1Bのような事前学習モデルでは、特に指示に従うように学習されていないケースもあり、そのまま利用した際、余計なことをだらだらと出力し続けたり、頓珍漢な回答が返ってきたりという問題が発生することがあります。

このため指示に適切にこたえてもらうための技術がいろいろあるのですが、そのうちの一つにSFTというものがあります。

Google ColabでPLaMo 2 1BをSFTする

それでは本題のGoogle ColabでPLaMo 2 1BをSFTする方法について説明します。今回はGPUメモリの関係上、おそらく無料で使えるT4だと無改造では実行できない気がするのでL4を使った説明をします。

L4 GPUの利用

まず、Google ColabでL4が使えるように、課金が必要になります。

課金についてはこちらをご覧ください。

https://colab.research.google.com/signup?hl=ja

今回のコードを動かすだけであれば「Pay As You Go」で100 コンピューティング ユニットを購入すれば十分です。この記事を執筆時点では1200円に満たない程度で購入できます。

課金が済んだら、メニューバーから「ランタイム」→「ランタイムのタイプを変更」をクリックします。すると無料枠では選択できないL4 GPUが選択できるようになっていると思うので、L4 GPUを選択します。

これでGPUを使う準備ができました。

実行環境準備

L4を利用するようにしたら、実行するコードのダウンロードやPythonパッケージのインストールを行います。

まずGithubよりコードをcloneしてきます

!git clone https://github.com/shu65/plamo-2-1b-sft-example.git

次に、PyTorchのバージョンを現在の最新版よりも前の以下のものに変更します。

!pip install torch==2.4.1 torchvision==0.19.1 torchaudio==2.4.1 --index-url https://download.pytorch.org/whl/cu124

この後は以下のようにPyTorch以外のPLaMo 2 1Bの実行に必要なパッケージやSFTに必要なパッケージなどをインストールします。

!pip install -r plamo-2-1b-sft-example/requirements.txt

ここまで実行すると2025/02/12現在以下のようなバージョンがインストールされました。

causal-conv1d                      1.5.0.post8
fastrlock                          0.8.3
mamba-ssm                          2.2.4
numba                              0.61.0
numba-cuda                         0.0.17.1
sentence-transformers              3.4.1
torch                              2.4.1+cu124
torchaudio                         2.4.1+cu124
torchsummary                       1.5.1
torchvision                        0.19.1+cu124
transformers                       4.48.2
trl                                0.14.0

これであとはSFTのコードを実行すれば、SFTをすることができます。このSFTの中身に関しては次で紹介していきます。

PLaMo 2 1BをSFTする

SFTをする部分はsft.py　というスクリプトにまとめてあります。このスクリプトの重要な部分について簡単にですが説明していきます。

まず、今回はすぐに実行が終わるように少量の質問と回答のペアのデータを用います。

今回は日本語の指示学習でよく使われるkunishou/databricks-dolly-15k-jaというデータセットのうち、input がなくinstruction とoutput のペアになっているデータのみを取り出しその一部だけを利用します。一つ例を見せると以下のようなデータを利用します。

{
  "output": "イコクエイラクブカ",
  "input": "",
  "index": "1",
  "category": "classification",
  "instruction": "魚の種類はどっち？イコクエイラクブカとロープ"
}

一部だけ取り出すコードは以下の通りです。

    dataset = datasets.load_dataset("kunishou/databricks-dolly-15k-ja")
    train_dataset = dataset["train"].filter(lambda data: data["input"] == "")

次にSFTConfig というSFTの実行の設定のクラスのインスタンスを用意します。具体的には以下の通りです。

    sft_args = SFTConfig(
        output_dir="./outputs",
        evaluation_strategy="no",
        per_device_train_batch_size=1,
        gradient_accumulation_steps=4,
        learning_rate=5e-5,
        num_train_epochs=0.1,
        lr_scheduler_type="cosine",
        warmup_ratio=0.3,
        logging_steps=10,
        save_strategy="epoch",
        report_to="tensorboard",
        bf16=True,
        max_seq_length=1024,
        gradient_checkpointing=True,
        deepspeed='./deepspeed_config.json',
    )

重要なこととして、今回はGPUのメモリが少ないため、DeepSpeedのStage 3という学習時に一部のデータをCPU側に置いておくモードを利用します。

これによりGPUメモリが少ない環境でもSFTを回すことができます。

DeepSpeed周りの設定はdeepspeed_config.json に書いてありますので気になる方はご覧ください。

また、今回は学習データの10%だけを利用するようにしています。これはこの学習を早く終わらせるためであり、本来はもっと回す必要があると考えられますので、本気でSFTをする場合は注意してください。

次にデータをどのようなフォーマットでLLMに入力するかを指定するformatting_func という関数を用意します。今回は以下のようにしました。

INSTRUCTION_TEMPLATE = string.Template(
    """### Question:
${input}

### Answer:
${response}<|plamo:eos|>
"""
)


def formatting_func(examples):
    output_texts = []
    for i in range(len(examples['instruction'])):
        text = INSTRUCTION_TEMPLATE.substitute(input=examples['instruction'][i], response=examples['output'][i])
        output_texts.append(text)
    return output_texts

INSTRUCTION_TEMPLATE が今回のフォーマットで、### Question:\n の後に指示、### Answer:\n のあとに回答が続き、最後にend of sequenceである<|plamo:eos|> が来るようになっています。

また、学習時には回答部分だけを学習してほしいので、どこからが回答かがわかるように‎DataCollatorForCompletionOnlyLM のインスタンスも用意します。これは以下の通りです。

    data_collator = DataCollatorForCompletionOnlyLM(
        response_template=tokenizer.encode(" Answer:\n", add_special_tokens=False),
        tokenizer=tokenizer
    )

response_template のところで回答前の部分がどのようなtoken idになるかを指定する部分があるので、上記のように指定します。前後の文字の影響で指定したtoken idが出現しないケースがあるので、その時はいろいろresponse_template に指定する文字列を調整してみてください。

最後にSFTを実行するためのクラスの‎SFTTrainer を以下のように用意します。

    trainer = SFTTrainer(
        model=model,
        args=sft_args,
        data_collator=data_collator,
        train_dataset=train_dataset,
        tokenizer=tokenizer,
        formatting_func=formatting_func,
    )

そして、以下のように実行し、結果を保存します。

    trainer.train()
    trainer.save_model()

これで学習が終わるとSFTConfig のoutput_dir で指定した./outputs に結果が出力されます。試しに私がGoogle Colabで実行した際は13分程度で学習が終わりました。コンピューティングユニットとしてはパッケージなどのインストールも含めて4だけ消費しました。

SFTされたモデルで推論してみる

最後にSFTされたモデルで推論するというのを行います。

これはPLaMo 2 1Bのexampleとほぼ同じでpromptだけ少し変えたものを例として用います。コードとしては以下の通りです。

from transformers import AutoTokenizer, AutoModelForCausalLM
import torch


model_name = "./plamo-2-1b-sft-example/outputs"

tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(model_name, trust_remote_code=True)


# プロンプトの準備
prompt = "### Question:\n埼玉の県庁所在地は何市？\n\n### Answer:\n"

# 推論の実行
inputs = tokenizer(prompt, return_tensors="pt")
generated_tokens = model.generate(
    **inputs,
    max_new_tokens=64,
    pad_token_id=tokenizer.pad_token_id,
)[0]
generated_text = tokenizer.decode(generated_tokens)
print(generated_text)

出力結果は以下のようになります。

<|plamo:bos|>### Question:
埼玉の県庁所在地は何市？

### Answer:
埼玉県の県庁所在地はさいたま市です。<|plamo:eos|>

ちゃんと学習で指定されたように### Answer:\n の後に質問に対する回答をし、その後<|plamo:eos|> を出力するということができています。

ちなみにSFTしていないモデルではどうなるかというと、以下のように余計なことを出力するうえ、出力が止まらないという状態になっています。

<|plamo:bos|>### Question:
埼玉の県庁所在地は何市？

### Answer:
さいたま市

### 解説
「県庁所在地」とは、都道府県庁が置かれている都市のことです。
「さいたま市」は埼玉県の県庁所在地です。

### 関連記事
### 取り急ぎお知らせ
「埼玉の県庁所在地は何市？」の解説は以上です。
「埼玉の県庁所在地は何市？」の解説は以上です。

このため、SFTでうまくフォーマットに従うよう学習できたと考えられます。

終わりに

今回はPLaMo 2 1Bを使ってSFTをする例を示しました。今回示したように簡単なSFTなら十分Google Colabで実行することができます。みなさんもぜひいろいろ試していただければと思います。

The post 小型LLM PLaMo 2 1BをGoogle ColabでSFTしてみる first appeared on まったり勉強ノート.

どんな内容の本か？

この本を一言でまとめると「Transformersを使った推論、学習など幅広くまとめた本」という感じかと思います。
「Transformer”s”」って何?という方向けに説明すると、ChatGPTなどで使われているTrasnfomerというモデルを扱いやすくしたPythonライブラリです。おそらく、この記事を執筆している現在、自然言語系のタスク向けにTransfomerのモデルを使って学習したり、推論したりしようと思ったら多分使うことになるライブラリかと思います。

このTransformersについて開発しているHugging Faceの人たちが自ら解説した本がこの本になります。扱っているテーマは幅広く、Transformerの仕組みから、Transformersを使ったテキスト分類などいくつかの応用タスクを実際に実行する方法、Transformersの高速化、学習などが書かれています。Transformersについて知りたいと思ったら、このを本をまず読んでみると全体を俯瞰できてよいかと思います。

どんな人にお勧めか？

この本は以下のような人に向いている本かなと思っています。

1. Transformerの自然言語応用について幅広く勉強したい人
2. Transformersを使ったコードについていろいろ知りたい人

特にTransformerの自然言語応用について知りたい方はちょうどよい本かと思います。一方、Transfomerの言語以外の応用、例えば画像なんかについては簡単な紹介はありますが、詳しくは書かれていません。このため、自然言語以外について知りたい人には向かない本だと思います。

個人的に良かった点

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

1. Transformersを使ったpretrainingについてちゃんと書いてある
2. 備考的なことについてもいろいろ言及があり、しかも参考文献がしっかりついているので、詳しく知りたい場合は論文にあたりやすい

Transformers + 自然言語については最近話題なこともあり、何冊か本が出ています。私自身、数冊読んだのですが、どれも応用よりなことが多く、pretrainingなどまで書いてない、もしくは書いてあったとしてもちょっとしかないみたいな本が多い印象です。この点、この本はpretrainingのやり方までちゃんと具体例を示しながら説明してあって良かったです。

また、単純にTransformersの使い方の説明にとどまらず、例えばデータセットの課題やTokenizerごとの違いについても簡単な言及がちゃんと書かれています。また、これらにちゃんとどの論文に書かれているのか示されているので、より詳しく知りたい場合は論文を読んで勉強するということができるようになっています。

終わりに

今回はTransfomrersについて書かれた「機械学習エンジニアのためのTransformers ―最先端の自然言語処理ライブラリによるモデル開発」について紹介する記事を書きました。

今後もこのように読んだ本の紹介を毎週月曜日に投稿しようと思いますので、興味がある方は見に来てみてください。

この記事が皆様の役に立てば幸いです。

The post [書評] 機械学習エンジニアのためのTransformers ー 自然言語のTransformerについてより知りたい人向けな一冊 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 まったり勉強ノート.

時間計測の部分で測定に使ったコードはここにあげてあります。

https://github.com/shu65/pytorch_2_compile_example/blob/main/torch_2_0_compile.ipynb

torch.compileとは？

torch.compileはPyTorch 2.0の新機能で、PyTorchの複数の機能を組み合わせて使い関数や深層学習のモデルを実行時に最適化して、その後の呼び出して高速に実行できるようにする機能です。

torch.compileの中身の詳しい説明はここにかかれています。

https://pytorch.org/get-started/pytorch-2.0/#technology-overview

簡単に説明するとtorch.compileの中身としては以下の３つで構成されています。

1. Graph acquisition: 計算グラフの構築
2. Graph lowering: PyTorchのオペレーションをバックエンドのデバイス（CPUやGPU）に特化した細かい命令に分解
3. Graph compilation: バックエンドのデバイス特化の命令を呼び出し

これらのステップを経ることで、より効率よく計算リソースを使えるようにし、高速化を実現しています。

また、この機能のすばらしいところは使い方も非常に簡単であるというものがあります。以下にデコレータで使う方法とtorch.compileの関数を呼び出して使う方法を示します。

デコレータで使うやり方

まずデコレータで使う方法です。これは以下のようになります (このチュートリアルの例：https://pytorch.org/tutorials/intermediate/torch_compile_tutorial.html#basic-usage)

@torch.compile
def opt_foo2(x, y):
    a = torch.sin(x)
    b = torch.cos(x)
    return a + b
opt_foo2(torch.randn(10, 10), torch.randn(10, 10))

torch.jit.scriptを使ったことがある方は、それと同じ感覚で使えるというと使い方がイメージしやすいかもしれません。

torch.compileの関数を呼び出して使うやり方

torch.compileの関数を呼び出してコンパイルする場合は以下のようにやります。(このチュートリアルの例：https://pytorch.org/tutorials/intermediate/torch_compile_tutorial.html#basic-usage)

class MyModule(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.lin = torch.nn.Linear(100, 10)

    def forward(self, x):
        return torch.nn.functional.relu(self.lin(x))

mod = MyModule()
opt_mod = torch.compile(mod)
opt_mod(torch.randn(10, 100))

こちらもtorch.jit.scriptのときと同じような使い方だと思います。

torch.compileによるパフォーマンスの評価

次にtorch.compileを実際に使ってみたときの計算時間を計測したので、その紹介です。今回は以下の二つのGPUで測定しました。

1. T4
2. V100

T4はTuringなので公式のドキュメントでtorch.compileのサポートが書かれてないものになっています。ただ、やってみたら少し早くなったので、測定結果を載せています。GitHubにあげたコードはT4で測定したほうです。

また、CUDAのバージョンはどちらのケースも12.0利用し、測定に使ったモデルはチュートリアルにあったtorchvisionのResNet18を使用しました。

また、torch.compileにはモードが以下の３つあります。

1. デフォルト
2. reduce-overhead
3. max-autotune

これらと何もしてない場合も含めて合計４つパターンの測定をしています。

具体的な測定方法が分かりやすいようにコードの一部を紹介します（torch.compleのデフォルトの場合）。

import time 

import torch
import torchvision.models as models
import torch._dynamo

batch_size = 64
n_warmup_iters = 10
n_iters = 500

x = torch.randn(batch_size, 3, 224, 224).cuda()

def get_mode():
    return models.resnet18()

torch._dynamo.reset()

model = get_mode().cuda()
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)

# compile
compiled_model = torch.compile(model)
for i in range(n_warmup_iters):
    optimizer.zero_grad()
    torch.cuda.synchronize()
    start = time.time()
    out = compiled_model(x)
    torch.cuda.synchronize()
    forward_elapsed_time = time.time() - start
    torch.cuda.synchronize()
    start = time.time()
    out.sum().backward()
    backward_elapsed_time = time.time() - start
    print(f"with compile {i} iter forward: {forward_elapsed_time/1000:.3e} msec., backward: {backward_elapsed_time/1000:.3e} msec.")
    optimizer.step()

print("-"*10)

torch.cuda.synchronize()
start = time.time()
for i in range(n_iters):
    optimizer.zero_grad()
    out = compiled_model(x)
    out.sum().backward()
    optimizer.step()
torch.cuda.synchronize()
elapsed_time = time.time() - start

print(f"with compile total:{elapsed_time:.3e} sec. {batch_size*n_iters/elapsed_time:.3e} imgs/sec.")

最初に、モデルの入力とモデルを作ったあと、コンパイルする場合はtorch.compile(model)でコンパイルします。このときコンパイルのモードを変える場合は引数のmodeにモードの名前を渡します。

その後、最初の数回はforward、backwardの呼び出し時にコンパイルなどのオーバーヘッドが入って遅いので、あらかじめ何度か呼びます。そして最後に実際に時間を計測します。今回は10回あらかじめforwardとbackwardを呼んでおいて、その後500回イテレーションを回したときの時間を測定しています。バッチサイズに関しては変化させると高速化率が変化することはわかっていますが、今回固定で64で実行しています。

T4, V100ともに同様の方法でtorch.compileのありなし等を測定しています。

では、時間計測の結果です。500回イテレーションを回したときの実際の計算時間を順番に示していきます。まずはT4の場合です。

T4はtorch.compileのサポートが書かれてない世代のGPUなので、効果が全くでないのかと思ったのですが、そんなことはなかったです。ただ、10％は満たない高速化にとどまっているという印象です。ちなみにT4を使ったケースではtorch.compileのmodeをmax-autotuneに変えると以下のようにサポートされてないGPUであると警告がでてきます。

[2023-03-17 18:31:06,314] torch._inductor.utils: [WARNING] not enough cuda cores to use max_autotune mode

次にV100のResNet18の結果です。

V100のほうはtorch.compileのサポートされていると書かれているGPUです。実際、V100はtorch.compileのmodeをmax-autotuneに変えると確かにより速くなり、高速化率も最大値は10%台に入っています。

現状のtorch.compileの注意点

最後にtorch.compileの注意したほうがよさそうな点を書いておきます。

まず、公式で書かれいたものの紹介です。基本的な注意点はこのドキュメントに書いてあります。

https://pytorch.org/get-started/pytorch-2.0/#pytorch-2x-faster-more-pythonic-and-as-dynamic-as-ever

重要なものとして、現在提供されているtorch.compileの機能を最大限活かせるのはCPU、NVIDIAのVoltaとAmpere世代のGPUのみになっています。他のGPUでは使おうとすると警告が出てきます。ただ、私が試した範囲では警告がでるだけで現状では使えないわけではなさそうです。

また、私が使ったときに感じた注意点としては

1. おそらくforwardとbackwardで別々にコンパイルが走るので、forward、backwardの両方とも最初は遅い
2. 実行が遅いのは最初の１回目だけでなく、最初の数回の呼び出しが遅いケースがある
3. Google ColabなどでCellの実行を一度止めて再度実行しようとするとエラーがでて、ランタイムの再起動をしないと復帰できないケースがある

1と２は時間計測をしようとしたときにはまったポイントです。まず、１に関してです。torch.compileの直後の呼び出しはコンパイルが走るので、遅いというのはドキュメントにも書かれています。ただ、forwadだけがおそいのかな？と思ってました。ただ、torch.compileの説明をちゃんと読めば想像できると思いますが、backwardも最初の実行のときは遅いです。なので、時間を計測するときは、forwardとbackwardの両方が遅いことを考慮して測定する必要があります。

次に２です。これに関しては私が見逃してなければドキュメントに明示的に説明が書いてあるわけではないのですが、チュートリアルの時間計測の結果や実際に測定してみるとどうやら遅いのは最初の１回目の呼び出しだけではなく、そのあと数回遅いケースが存在しているようです。このため、計算時間の測定の際、最初に数回呼び出してから測定しないとtorch.compileを使ったときよりも遅いみたいな誤った結果になるので注意してください。

最後に３です。これは何度かはまったのですが、どこかにキャッシュか何か残っているのか変なところで止めるとコード的には問題ないはずなのに、エラーがでるようになるときがあります。調べても解決方法が分からなかったので、エラーがでるようになったらランタイムごと再起動するということを何度かやりました。Google Colabでやるときは注意してください。

終わりに

今回はtorch.compileについて使ってみたのでまとめを書きました。去年発表があったときから楽しみにしていましたが、期待通りのものとなっていました。なにより使い方が非常に簡単なことには驚きました。

今回はT4とV100の測定結果でしたが、A100だとどうなるのかも今度測定しようかなと思っています。

この記事がみなさんのお役に立てば幸いです。

The post PyTorch 2.0の新機能「torch.compile」使ってみた 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 まったり勉強ノート.

今回の記事では私が読んだ本の軽い解説とどういう人向けかのまとめです。今後同じように統計的因果推論関係の勉強をしたいという人の参考になれば幸いです。

因果推論の科学 「なぜ?」の問いにどう答えるか

因果推論の科学 「なぜ？」の問いにどう答えるか

ジューディア・パール , ダナ・マッケンジー, 夏目 大

Amazon

楽天市場

ポチップ

まずは私が統計的因果推論おもしろい！と感じさせてくれた「因果推論の科学 「なぜ?」の問いにどう答えるか」です。本の内容としては、因果に関する研究が如何に難しいかやどれほど役立つのかはもちろん、因果に関する研究の歴史にも言及し、他の分野、とりわけ統計学とどういう関わりがあってどのように発展してきたか？について書かれています。個人的にはPearsonやFisherなど手法名で名前をしっている人たちがどういうことをしていた人なのかも少ししれて面白かったです。

こちらの本は一般人向けに因果推論がどういうものかを書いた本になっています。書いた人が研究者ということもあってか、一部数式も出てきますが、基本、誰でもわかるような表現で書かれていて、「因果推論ってなに？」ということを背景的なことを重点的に知りたいという方はとっつきやすい１冊ではないかと思います。一方、数式がでてきますが、詳しい説明はそれほどないので、式変形に関してや、統計的因果推論の実際の応用例をたくさん知りたいという方には向かない本です。また、統計的因果推論についてすでにご存じの方向けの説明になってしまいますが、この本は構造的因果モデルとdo演算子を利用した因果推論についての本で、潜在的結果変数の枠組みを使うRubin流の因果推論の話はほとんどでてこないので、後者のほうを知りたいという人にも向かないので注意してください。

入門 統計的因果推論

入門 統計的因果推論

Judea Pearl, Madelyn Glymour, Nicholas P. Jewell, 落海 浩

Amazon

楽天市場

ポチップ

次は「入門 統計的因果推論」です。こちらは「因果推論の科学 「なぜ?」の問いにどう答えるか」の方が著者の一人に入っている統計的因果推論に関する教科書です。こちらは「因果推論の科学 「なぜ?」の問いにどう答えるか」とは違い、一般人向けではなく、技術者、研究者向けの本になっています。内容としては統計的因果推論の話をするのに必要な確率と統計、回帰、グラフィカルモデルの基本的な話から始まり、統計的因果推論で重要な介入効果や反実仮想とその応用について解説しています。「因果推論の科学 「なぜ?」の問いにどう答えるか」でも登場した因果グラフの具体例はもちろん、他にもいろいろ出てくるので、因果グラフのイメージはつきやすい印象です。一方、数式の式変形に関する記述が少なくて、この式どうしてでてきたの？という疑問が結構読んでてありました。また、因果グラフの構築の仕方に関する具体的な説明があまりなく、具体的な数値まで使った例というのはすくないので、実務に向けて勉強したい方はこの本だけだと、自分の持っているデータにどう適用すればいいのか？というのがわからないと思っています。

こちらの本、「因果推論の科学 「なぜ?」の問いにどう答えるか」の内容をちゃんと勉強したいという向け方はおすすめの本です。一方、構造的因果モデルとdo演算子を利用する法の因果推論の紹介が主で潜在的結果変数の枠組みを使うRubin流の因果推論の話を期待していると思っているのと違った、ということが起きると思っていますので注意してください。また、先ほども述べたように、この本だけでは統計的因果推論を使いこなすのは難しいと思っていて、実務で統計的因果推論を使いたいという方は後ほど紹介する具体的なコード付きの本を参考を読むとよいかと思っています。

統計的因果探索 (機械学習プロフェッショナルシリーズ)

統計的因果探索 （機械学習プロフェッショナルシリーズ）

清水 昌平

Amazon

楽天市場

ポチップ

次は「統計的因果探索」です。前の２冊は統計的因果”推論”に関する本でしたがこちらは統計的因果”探索”に関する本です。前の本では構造的因果グラフは与えられているものと仮定して説明されている部分が多く、どうやって作るか？についての言及がほとんどありません。この本は構造的因果グラフをどうやって作るか？ということに関して、LiNGAMという手法の説明をした本になります。

内容としては簡単な統計的因果推論の話はあるものの、基本は統計的因果探索がメインです。また、統計的因果探索に関してもLiNGAMに関する記述が多く、それ以外に関しては簡単な解説が少しある程度です。ただ、LiNGAMに関しては説明が幅広く、未観測共通原因がある場合のLiNGAMや、LiNGAMで仮定を緩めることに関する記述もあります。なので、LiNGAMを深堀したい人にはちょうどよい１冊かと思います。

一方、他の機械学習プロフェッショナルシリーズの本の例に漏れることなく、コードを使った説明がなかったり、細かい数式変形までは書いてないので、詳細に理解するのは結構大変な印象です。

こちらの本は他の本でLiNGAMに触れて面白そうだからもっと勉強したい！という人向けかなと思っています。

つくりながら学ぶ! Pythonによる因果分析 ~因果推論・因果探索の実践入門 (Compass Data Science) 

つくりながら学ぶ! Pythonによる因果分析 因果推論・因果探索の実践入門

小川雄太郎

Amazon

楽天市場

ポチップ

次は「つくりながら学ぶ! Pythonによる因果分析 ~因果推論・因果探索の実践入門」です。

これまで紹介した本の内容をPythonで実際に実装するか？を勉強したい人には現状ベストな本です。内容としては構造的因果グラフとdo演算子を使った統計的因果推論の基本的な説明に加えて、LiNGAMとベイジアンネットワーク、最新の深層学習を使った因果探索の話に関する説明＋Pythonを使ったコードになっています。特にPythonコードはJupyter Notebookで書かれていて、Google Colabを使えばすぐに実行できるように工夫されています。このため、これまで紹介した本のどれよりも具体的でかつ分かりやすい印象です。コードも一通り読みつつ実行もしてみましたが、きれいなコードで分かりやすかった印象です。

このため、今すぐ統計的因果推論や統計的因果探索を実務で使いたい人に対して一番最初に読むといい本はどれですか？と聞かれたらこれが良いと答えると思います。

統計的因果推論の理論と実装 (Wonderful R)

統計的因果推論の理論と実装 潜在的結果変数と欠測データ （Wonderful R　5）

高橋 将宜

Amazon

楽天市場

ポチップ

最後に「統計的因果推論の理論と実装 (Wonderful R)」です。

今までの本と同じく「統計的因果推論」と書かれていますが、こっちは潜在的結果変数の枠組みを使うRubin流の因果推論の本です。このため、因果グラフは一部出てきますが、do演算子についてはでてきません。ただ、かなり幅広い内容の本となっているので、Rubin流の因果推論について詳しく知りたいという人です。説明にはRによる具体的なコードが付いていて、言葉だけではわかりづらいところもコードで補えるような形になっています。

PythonはわからないけどRを使える人や、Rubin流の因果推論についても詳しく知りたいという人にはお勧めな本かと思います。

終わりに

年末年始のお休みで読んだ本のうち、統計的因果推論、統計的因果探索に関する本をまとめて紹介しました。最初に紹介した「因果推論の科学 「なぜ?」の問いにどう答えるか」が思いのほか面白くてその勢いのまま読んだという感じがしてますが、因果推論はいつかまとめてちゃんと勉強したいと思っていたのでちょうどよかったです。

今後、今日紹介した一部の本に関しては本に書かれていた説明だけだと意味がわからないところがあったので、そういうところに関してはまたメモとして記事にしておこうと思います。

The post 統計的因果推論、統計的因果探索の勉強で読んだ本まとめ(2023/1版) first appeared on まったり勉強ノート.

ちなみに、自分の手法に関してはこちらをご覧ください。

https://www.kaggle.com/competitions/open-problems-multimodal/discussion/366961

Open Problems – Multimodal Single-Cell Integration とは？

コンペのサイトのoverviewに書かれていることをざっくり日本語に訳すと「骨髄幹細胞が血液細胞になるにつれて単一の細胞のDNAとRNA、タンパク質がどのように変化するかを予測するコンペ」ということになります。もう少し専門家の人がわかりやすいように説明すると以下の二つのsingle-cellのデータを使ったモデルを構築するコンペです。

1. Chromatin accessibility(peak countをTF-IDFで変換したデータ)からRNAの発現量 (library-sizeでノーマライズされたcountデータ)の予測
2. RNAの発現量 (library-sizeでノーマライズされたcountデータ)からsurface protein levels (dsb でノーマライズされたデータ)の予測

そもそもなぜ参加したのか？

参加してがんばろうと思ったきっかけをせっかくなので書いておくと

1. Bioチームの同僚がこういうコンペがあるよとslackで教えてくれた (参加したきっかけ)
2. 育休に入る前にやっておくべきことは？ということを会社の先輩パパさんに相談したら、「solo gold medalは取っておけ」と言われた (がんばったきっかけ)

ということがあります。特に2つ目は「確かに！」と思いました。なので、今回が初ソロ参加にして、solo gold medalを取る最後のチャンスということで頑張りました。アドバイスをくれた先輩パパさんには感謝しかありません。

やってよかったこと

さて、ここからやってよかったことについて忘れないように書いておきます。

まずはとにかく簡単な方法でいいのでsubmitする

社内のkaggle強い人に前に言われた気がするので、まず意識したことがこれです。やってみて思ったのですが以下のような効果があることを実感しました。

1. submitして順位やスコアがでるようになるとモチベーションが上がる。
2. 何かベースラインがあると手法開発がしやすい

特に1の効果がすごかった気がします。ちなみに私の場合、KaggleのCodeで公開されてたシンプルな手法をそのままコピペしてsubmitして一番最初のスコアを出しました。シンプルな手法だったため、最初は200位にも入れなかったと記憶していますが、それでもモチベーションはそれまでと比べてすごく上がりました。

毎日決めた数の改良を試す

これはコンペに限らず重要なことだと思いますが、とにかくコンスタントに改良を続けていき、最低限local CV scoreを出すということを意識してました。

今回のコンペでは3，4個の改良を毎日試すことを目標にしてやってました。私の場合、gitの1 branchが1改良になっていて、9/1からがんばり始めて11/15までに最終的には351個branchができていました。なので、通算4から5個くらいの改良を毎日試していたことになります。

このとき改良としてうまくいきそうなものはもちろんですが、思いついたタイミングではあまり筋が良くなさそうだけど他にやることがないというときは、筋が良くないアイディアもダメ元で試すようにしてました。

結果として、やってよかったと感じた理由としては、仮説を立てて試すとうまくいかなかったときの問題への理解度がすごい上がるため、とにかくいろいろ試すことで、そこからいいアイディアをひらめくということが多くあったからです。

今回のコンペで何度か大きくスコアを上げたタイミングがありましたが、総じて何かの失敗から気が付いたアイディアをもとにしたことが多かった印象です。ちなみにどれくらい失敗し続けたのかの体感ですが、1週間スコアが上がらないということがよくあったので、うまくいった改良というのはこれだけ試しても数えるくらいしかなかった印象です。

ensembleはすぐに試さない

これも社内のkaggle強い人に前に言われたことなのですが、ensembleをすぐに試さないというのはやってよかったなと思います。
今回はensembleを試さなくてよかったなと感じた理由としては、ensembleはやれば簡単にスコアが上がるのでいいように感じますが、ensembleは始めるとすごいいろいろ試せることがある一方、おそらくそこまで高いスコア上昇をしていなかったような印象があります。

このためensembleを頑張ることに時間使うよりも1つのモデルのスコアが上がるように頑張るという作戦でいたのですが、結果としてそれが良かった印象です。

最後まで諦めない

最後はこれです。特に今回はsolo gold medalが欲しかったので、一人で出ていたのですが、途中全然順位が上がらず何度ももうやめようかなぁと思うタイミングがありました。結果としてそのときダメもとで試した改良やそこから思いついたアイディアがうまくいってスコアを伸ばすということが何度もあったので諦めなくてよかったです。

課題

次にコンペの参加中、もしくは終わったあと振り返って感じた課題的な部分も列挙しておきます。

モチベーションの維持が大変（特に一人のとき）

「最後まで諦めない」のところに似たようなことを書きましたが、とにかくモチベーション維持が大変でした。今まではチームででていたのと、すごい応援してくれる上司がいてくれたりとこの部分はそこまで問題にならなかった印象でした。

今後一人ででるならこの部分はまだまだどうにかしないと最後まで戦うのは難しい気がしています。

public scoreがどうしても気になってしまう

今回のコンペはデータセットの説明を読んだ段階でpublic scoreとprivate scoreのギャップが激しそうだなぁということを思ってたので、それほどpublic scoreを気にしないほうがいいかも？ということを最初思ってました。ただ、それでも最後はpublic scoreを気にしてモデルの改良をしてしまっていました。

コンペのサイトのデータセットの説明のところに詳しくかかれていますが、今回public scoreは4人中一人の最終日一つ前までのデータを、private scoreは4人全員の最後の日のデータを予測するというものでした。この説明を見ると予測する対象が全然違うことがわかります。実際、public scoreの上位陣が軒並みprivate scoreでは順位を落としていました。なので、後から振り返るとやっぱり最初に思った通りpublic scoreをそれほど気にせずモデルの改良をしていて正解だったと思います。ただ、そうはいっても、public scoreの順位が気になってしまって、結局最後はpublic scoreを気にして最終submittionを決めていました。ただ、ふたを開けたらlocal CVがベストなものがprivate scoreも一番よかったので、public scoreを気にしすぎたなぁと思っていました。

ただ、この部分はふたを開けてみないとわからないところなので、難しいポイントな気がしています。

論文を読んで最新研究をコンペの期間中に試すことが心理的に難しい

期間の前半ならまだましですが、どうしても後半になればなるほど、心理的に追い込まれていきます。なので、試すアイディアがないと思いつつも、自分の全然読んだことない論文の手法を試すということが難しかったです。この結果、特に後半はアイディアが前に試したもののちょっと変更したものばかりになり、結果としてそれほど精度が向上しないということが起きました。ただ、この部分一人でやっていると意識しても難しい部分な気がしています。

このため、日ごろから論文を読んでいろいろ手法を勉強し、できれば実装を動かして感触を確かめておくことが重要かなと思いました。実際、今回のコンペではReactomeのpathwayデータを使ったのですが、これは論文の追試のついでにいろいろ試したことがあったからこそできたことだったと思っています。

ちなみにその時の記事はこちらにあります。

ensembleの準備をせずにいると最後の土壇場で困る

ensembleをすぐに試さないことを意識していた結果、実は今回期間最後のほうでensembleをしようとしたときにいろいろ準備できてないことに気が付いてすごい慌てました。具体的には以下の二つが締め切り1週間前の段階でできていませんでした。

1. 予測結果をどのように集計して最終的なファイルを作るか？
2. ensembleした結果をどのように評価するか？

結局、1はギリギリできたのですが、2の評価方法はとても間に合いそうになかったのでpublic scoreを見てうまくいっているかいっていないか？を判断していました。ただ、これは評価としては微妙なので、次からはそんなに頑張らないにしても中盤くらいにはensembleの準備はしておこうと思います。

最後に

折角なので個人的な振り返りをblog記事にしました。仕事ではない状態でのkaggle参加は初めてだったのですが、思ったよりも疲れたのと、無事子供がコンペ締め切りの前日に生まれて、これから子育てもあるのででしばらくはコンペにでない気がします。ただ、次出たときに今回のコンペで何を思ったのか忘れないようにしておければと思っていたので、記事にできてよかったです。

これが他の人の参考になれば幸いです。

The post Kaggleの「Open Problems – Multimodal Single-Cell Integration」の振り返り first appeared on まったり勉強ノート.

最近読んだ論文にVariational Graph Auto-Encoders (VGAE) を使ったモデルがあったので、自分でもやってみようと思い、作ってみました。本日はそのまとめになります。

本日紹介する使うコードは以下のものです。

https://github.com/shu65/pytorch_geometric_examples/blob/main/PyTorch_Geometric_Variational_Graph_AutoEncoder.ipynb

また、このコード自体、以下のPyTorch Geometricのexampleのコードとほぼ同じです。

https://github.com/pyg-team/pytorch_geometric/blob/ee509ad65aefa679047356bb00bc498f35ce7e20/examples/autoencoder.py

このblog記事ではVGAEで必要な機能がPyTorch Geometricでどう実装されているのかわからなかった部分がいくつかあるのでその部分を解説していく記事になります。

PyTorch Geometricとは

PyTorch GeometricはPyTorchを使って構築されたGraph Neural Network向けのライブラリになります。

GitHubのURLは以下の通りです。

https://github.com/pyg-team/pytorch_geometric

最新のPyTorchやCUDAにもちゃんと対応しており、Graph Neural Networkで必要な基本的な機能はそろっている印象です。

Variational Graph Auto-Encoders (VGAE)とは

VGAEはVariational Auto-Encoder (VAE) というモデルをGraphデータ向けに拡張したモデルです。VAEの説明を始めるとそれだけですごく長くなりますので、今回はVGAEを実装するうえで必要なところだけ紹介します。

VAEは以下のようにEncoderとDecoderという二つのモデルを組み合わせたモデルになります。

このうち、EncoderとDecoderは以下のようなモデルになります。

• Encoder: 入力Xを受け取って潜在変数Zの分布のパラメータを出力する
• Decoder: 潜在変数Zを受け取って入力Xを再構成する

VAEで重要なのがEncoderの部分と潜在変数Zのサンプリングの部分です。この潜在変数Zの分布が標準正規分布という仮定のもと学習させながら、Encoderで潜在変数Zの分布のパラメータを出力し、その分布のパラメータを使って潜在変数ZをサンプリングしてDecoderに渡すということを行います。

このVAEをGraph データに拡張するためにVGAEはEncoderとDecoderを以下のようなモデルにしています。

• Encoder: ノードの特徴ベクトルXと隣接行列Aを入力として受け取り、潜在変数Zの分布のパラメータを出力する
• Decoder: 潜在変数Zを受け取り隣接行列Aを再構築する

図にすると以下のようなイメージです。

VGAEとVAEとの違いはEncoderでグラフの情報であるノード情報と隣接行列を受け取れるようにしたことと、Decoderが出力するものが隣接行列になることです。

VGAEをPyTorch Geometricを使って実装する

VGAEの概略を説明したので次は実際に実装を紹介していきます。まずはEncoderであるVariationalGCNEncoderから見ていきます。EncoderではPyTorch Geometricに実装されている GCNConv を使って実装します。

from torch_geometric.nn import GCNConv

class VariationalGCNEncoder(torch.nn.Module):
    def __init__(self, in_channels, out_channels):
        super().__init__()
        self.conv1 = GCNConv(in_channels, 2 * out_channels)
        self.conv_mu = GCNConv(2 * out_channels, out_channels)
        self.conv_logstd = GCNConv(2 * out_channels, out_channels)

    def forward(self, x, edge_index):
        x = self.conv1(x, edge_index).relu()
        return self.conv_mu(x, edge_index), self.conv_logstd(x, edge_index)

GCNConv はノードのインプットのチャンネル数、アウトプットのチャネル数を引数にとってインスタンスを作ります。そしてforwardではノードのtensor x と隣接行列のかわりにどのノード同士がつながっているか？を示すedge_indexを渡します。GCNConv の中身についてはドキュメントに詳しく書かれているのでそちらをご覧ください。

https://pytorch-geometric.readthedocs.io/en/latest/modules/nn.html#torch_geometric.nn.conv.GCNConv

このEncoderではVGAEの概要でも説明した通り、潜在変数の分布のパラメータを返します。ここではガウス分布の平均を表すmuと標準偏差にlogを適用したlogstdを返しています。

モデルの実装としてはあとはPyTorch Geometricで実装されているVGAEというクラスに渡せば終わりになります。

from torch_geometric.nn import VGAE

model = VGAE(VariationalGCNEncoder(in_channels, out_channels))

ただ、これだとさすがに初見だと何が何だかわからなかったので、少し説明します。

まず、Decoderについてです。DecoderはVGAE のデフォルトではInnerProductDecoderというものが使われます。これはVGAEの元論文でも使われていたDecoderの実装で、エッジの両端のノードに対応する潜在変数の各要素の積を取って総和を取り、sigmoidを適用して0-1の値にして出力します。出力値が0-1の値になっているのでDecoderの出力値は計算に使った二つのノードの間にエッジがある確率とみることができます。

詳しくは以下のドキュメントをご覧ください。

https://pytorch-geometric.readthedocs.io/en/latest/modules/nn.html#torch_geometric.nn.models.InnerProductDecoder

また、ロス関数についてですが、VGAE の中にVGAEで必要な以下の二つが実装されています。

• recon_loss: 潜在変数zとノード同士のつながりを示すpos_edge_indexを入力にとり、Decoderを利用して各エッジのある確率を計算、その確率に対してbinary cross entropyを計算してlossとして返す関数
  https://pytorch-geometric.readthedocs.io/en/latest/_modules/torch_geometric/nn/models/autoencoder.html#GAE.recon_loss
• kl_loss: Encoderの出力したmuとlogstdを使って標準正規分布とのKLダイバージェンスを計算しlossとして返す関数
  https://pytorch-geometric.readthedocs.io/en/latest/modules/nn.html#torch_geometric.nn.models.VGAE.kl_loss

これを以下のように学習ループで利用して学習をおこないます。

for epoch in range(0, 400):
    model.train()
    optimizer.zero_grad()
    z = model.encode(train_data.x, train_data.edge_index)
    recon_loss = model.recon_loss(z, train_data.pos_edge_label_index)
    kl_loss = (1 / train_data.num_nodes) * model.kl_loss()
    loss = recon_loss + kl_loss
    loss.backward()
    optimizer.step()

最後に上のコードではノード間にエッジがあるところの情報はtrain_data.pos_edge_label_indexで渡しているのですが、ノード間にエッジがないという情報はどこで渡しているか？ということについて説明します。

コードを読むと実はrecon_lossの中で自動的にエッジがないという情報を生成してそれを込みでロスが計算されています。具体的には以下の部分です。

https://github.com/pyg-team/pytorch_geometric/blob/d2b2e662488eae07d153de6d4b8c56c24bf413d9/torch_geometric/nn/models/autoencoder.py#L101

ここで引数でneg_edge_indexがNoneのときは自動でエッジが存在しないノードのペアをサンプリングするという処理になっています。

以下です。その他の部分で気になるところがある場合は全体のコードを以下のところに置いてありますのでご覧ください。

https://github.com/shu65/pytorch_geometric_examples/blob/main/PyTorch_Geometric_Variational_Graph_AutoEncoder.ipynb

終わりに

今回はPyTorch Geometricの練習として、VGAEを実装してみたのでまとめの記事を書きました。PyTorch Geometricを今回初めて使ったのですが、Graph Neural Networkに必要な基本的な機能はそろっていそうなので、今後もGraph Neural Networkを使う機会があれば使ってみようと思います。

The post PyTorch Geometricを使ってVariational Graph Auto-Encodersを作って学習してみる first appeared on まったり勉強ノート.

最近ベイズ推論の勉強をしていて機械学習スタートアップシリーズの「ベイズ推論による機械学習入門」を読んでいます。今回はこの本の5.7 のニューラルネットワークの章で紹介されていたモデルをPyTorchで実装したので、実装と苦労した点を紹介していきます。

参考にしたJuliaのサンプルコードはこちらです。

https://github.com/sammy-suyama/BayesBook/blob/master/src/demo_BayesNeuralNet.jl

今回はこのサンプルコードをもとにして、PyTorchで実装したものを作り、以下に公開しました。

https://github.com/shu65/blog-bayes-book/blob/main/%E3%83%99%E3%82%A4%E3%82%BA%E6%8E%A8%E5%AE%9A%E3%81%AB%E3%82%88%E3%82%8B%E6%A9%9F%E6%A2%B0%E5%AD%A6%E7%BF%92%E5%85%A5%E9%96%80%E3%80%805_7_%E3%83%8B%E3%83%A5%E3%83%BC%E3%83%A9%E3%83%AB%E3%83%8D%E3%83%83%E3%83%88%E3%83%AF%E3%83%BC%E3%82%AF.ipynb

今回はこのPyTorch実装についての紹介になります。パラメータは少しいじっていますが大体Juliaのサンプルと合わせています。

PyTorchで実装することで、より複雑なモデルを作ることが簡単になるのと、単純作業だけどやるのは大変な偏微分の計算のところをPyTorchの自動微分に任せることができるため、より本質的なことろが理解しやすくなると思ったためです。

それではまず本に書かれているモデルについて簡単に説明して、その後実際のPyTorch実装の説明という流れで説明していきたいと思います。

本で紹介されているモデルと変分推論

モデルの定義

まずは本で紹介されているモデルについて説明していきます。訓練データの入力値と出力値の集合をそれぞれ \( \boldsymbol{X} \)、 \( \boldsymbol{Y} \) と置きます。この集合の要素数を \( N \) とし \( n \in N \) 番目のデータの入力値を \( \boldsymbol{x}_n \in \mathbb{R}^M \) 、出力値を \( \boldsymbol{y}_n \in \mathbb{R}^D \) とするとき、ガウス分布によってモデル化すると以下のようになります。

$$ \begin{align*}
p(\boldsymbol{Y}|\boldsymbol{X}, \boldsymbol{W})
=& \prod_{n=1}^N p(\boldsymbol{y}_n|\boldsymbol{x}_n, \boldsymbol{W}) \\
=& \prod_{n=1}^N \mathcal{N}(\boldsymbol{y}_n|f(\boldsymbol{W}, \boldsymbol{x}_n), \lambda_y^{-1}\boldsymbol{I}_D) \tag{5.257}
\end{align*} $$

ここで、\( \boldsymbol{W} \) は モデルのパラメータの集合、\( \lambda_y^{-1} \) 固定の精度パラメータです。また、非線形関数\(f\)は次のように定義することにします。

$$ \begin{align*}
f(\boldsymbol{W}, \boldsymbol{x}_n) = {\boldsymbol{W}^{(2)}}^\mathrm{T} \text{Tanh}({\boldsymbol{W}^{(1)}}^\mathrm{T}\boldsymbol{x}_n) \tag{5.258}
\end{align*} $$

ここでモデルパラメータ \( \boldsymbol{W} \) の要素を\( \boldsymbol{W}^{(1)} \) と \( \boldsymbol{W}^{(2)} \)として、\( \boldsymbol{W}^{(1)} \in \mathbb{R}^{M \times K}\) 及び\( \boldsymbol{W}^{(2)} \in \mathbb{R}^{K \times D}\) という行列としています。このモデルパラメータの各要素は次のようなシンプルなガウス事前分布を仮定することにします。

$$ \begin{align*}
p (w_{m,k}^{(1)}) = \mathcal{N} (w_{m,k}^{(1)}|0, \lambda_w^{-1}) \tag{5.259} \\
p (w_{m,k}^{(2)}) = \mathcal{N} (w_{m,k}^{(2)}|0, \lambda_w^{-1}) \tag{5.259}
\end{align*} $$

また \( \text{Tanh}(\cdot) \) は以下のように定義されます。

$$
\text{Tanh}(a) = \frac{\text{exp}(a) – \text{exp}(-a)}{\text{exp}(a) + \text{exp}(-a)} \tag{5.261}
$$

以上が本に書かれたモデルの定義になります。

変分推論

ここではニューラルネットワークモデルのパラメータ \( \boldsymbol{W} = \{ \boldsymbol{W}^{(1)}, \boldsymbol{W}^{(2)} \} \) の事後分布を推論する問題を考えます。こちらは本で一つ前に紹介されていたロジスティック回帰のときとほぼ同じように行っていきます。

今回は細かい式の説明は省いていますが、ニューラルネットワークの勾配の計算は5.6ロジスティック回帰と似たような式変形になります。5.6ロジスティック回帰の式の導出はこちらの記事に詳しく書いてあるので参考にしてみてください。

事後分布を推定するために、対角ガウス分布を使った以下のような近似事後分布を利用します。

$$ \begin{align*}
q(\boldsymbol{W}^{(1)}; \boldsymbol{\eta}^{(1)}) =& \prod_{m=1}^M \prod_{k=1}^K \mathcal{N}(w_{m,k}|\mu_{m,k}^{(1)},{\sigma_{m,k}^{(1)}}^2) \tag{5.262.1} \\
q(\boldsymbol{W}^{(2)}; \boldsymbol{\eta}^{(2)}) =& \prod_{m=1}^M \prod_{k=1}^K \mathcal{N}(w_{m,k}|\mu_{m,k}^{(2)},{\sigma_{m,k}^{(2)}}^2) \tag{5.262.2} \\
\end{align*} $$

ここで \( \boldsymbol{\eta} = \{ \boldsymbol{\eta}^{(1)}, \boldsymbol{\eta}^{(2)} \} \) が変分パラメータの集合となります。今回は以下のような近似事後分布と真の事後分布のKLダイバージェンスを最小化するような変分パラメータ \( \boldsymbol{\eta} \)を見つけることを目指します。

$$ \begin{align*}
&\text{KL} [q(\boldsymbol{W};\boldsymbol{\eta})||p(\boldsymbol{W}|\boldsymbol{Y},\boldsymbol{X})] \\
& \ = \langle \ln q(\boldsymbol{W};\boldsymbol{\eta}) \rangle _{q(\boldsymbol{W};\boldsymbol{\eta})}
– \langle \ln p(\boldsymbol{W}) \rangle _{q(\boldsymbol{W};\boldsymbol{\eta})} \\
& \qquad – \sum_{n=1}^N \langle \ln p(\boldsymbol{y}_n | \boldsymbol{x}_n, \boldsymbol{W}) \rangle _{q(\boldsymbol{W};\boldsymbol{\eta})}
+ \text{const} \tag{5.236}
\end{align*} $$

この式 (5.236)の最小化するにあたり、以下のような再パラメータ化トリックを利用して、 \( \boldsymbol{W} \) の各要素 \( w \) (インデックスは省略してます。) を以下のように置きます。

$$ \begin{align*}
\tilde{w} = \mu + \sigma \tilde{\epsilon} \tag{5.237} \\
\text{ただし} \tilde{\epsilon} \sim \mathcal{N} (\epsilon|0,1) \tag{5.238}
\end{align*} $$

これを利用すると以下のような \( g(\boldsymbol{\tilde{W}}, \boldsymbol{\eta}) \) を最小化することになります。

$$ \begin{align*}
& \text{KL} [q(\boldsymbol{W};\boldsymbol{\eta})||p(\boldsymbol{W}|\boldsymbol{Y},\boldsymbol{X})] \\
& \ \approx \ln q(\boldsymbol{\tilde{W}};\boldsymbol{\eta})
– \ln p(\boldsymbol{\tilde{W}}) \\
& \qquad – \sum_{n=1}^N \ln p(\boldsymbol{y}_n | \boldsymbol{x}_n, \boldsymbol{\tilde{W}}) + \text{const} \\
& \ = g(\boldsymbol{\tilde{W}}, \boldsymbol{\eta}) \tag{5.239}
\end{align*} $$

ただし、本ではすべてのデータで尤度の勾配を計算する方法ではなく、確率的勾配降下法（stochastic gradient descent, SGD）にも触れられているので、この記事ではSGDを使って最適化します。ただ、後ほどまた説明しますが、本で書かれているデータを1つ1つ逐次的に与えて勾配を計算するオンラインのSGDではうまく収束してくれなかったので、ミニバッチを用いるSGDを使います。これは基本的に本で書かれているように式(5.239) の事前分布と近似事後分布の項の影響をデータ数に応じて抑えます。今回はミニバッチを利用するので、ミニバッチ内の訓練データを\( \{\boldsymbol{X}_B, \boldsymbol{Y}_B\}\)とし、 \( b \in B \) 番目のデータの入力値を \( \boldsymbol{x}_b\) 、出力値を \( \boldsymbol{y}_b \)として式 (5.239) を変形した以下の式の勾配を利用します。

$$ \begin{align*}
& \text{KL} [q(\boldsymbol{W};\boldsymbol{\eta})||p(\boldsymbol{W}|\boldsymbol{Y}_B,\boldsymbol{X}_B)] \\
& \ \approx \frac{B}{N} \lbrace \ln q(\boldsymbol{\tilde{W}};\boldsymbol{\eta})
– \ln p(\boldsymbol{\tilde{W}}) \rbrace \\
& \qquad – \sum_{b=1}^B \ln p(\boldsymbol{y}_b | \boldsymbol{x}_b, \boldsymbol{\tilde{W}}) + \text{const} \\
\tag{5.269.1}
\end{align*} $$

本ではこれ以外に勾配を計算するための式変形が細かく書いてありますが、今回はPyTorchの自動微分の機能を使うため、式 (5.269.1) の値を計算し、この値をlossとしてbackward() を実行するため、説明はここまでにします。

実装について

今回はPyTorchの自動微分の機能を使って変分パラメータの最適化に必要な勾配を計算します。このため、処理の基本的な流れとしては式 (5.269.1) の値を計算し、この値をlossとしてbackward() を実行するというのを指定した回数繰り返して最適化します。

モデル部分の実装

式 (5.258) をPyTorchで実装します。具体的なものは最初に示したJupyter Notebookの BayeNNModel クラスの実装をご覧ください。ここでは重要な部分だけ示します。まず、式 (5.258)を forward() に実装します。

    def forward(self, X):
      W1 = self.sample_W(self.mu1, self.rho1)
      h1 = torch.nn.functional.linear(X, W1, bias=None)
      h2 = torch.tanh(h1)
      W2 = self.sample_W(self.mu2, self.rho2)
      Y_est = torch.nn.functional.linear(h2, W2, bias=None)
      return Y_est, W1, W2

今回は \( \boldsymbol{\tilde{W}} \) はサンプリングしてくる必要があるので、sample_W() という関数でサンプリングしてそれをtorch.nn.functional.linear()に入れるということをしています。

\( \boldsymbol{\tilde{W}} \) はサンプリングしてくる部分は以下のようにします。

    def sample_W(self, mu, rho):
      eps = torch.randn(mu.shape)
      sigma = self.rho2sigma(rho) + self.approximate_posterior_sigma_eps
      W = mu + sigma * eps
      return W

基本的には式(5.237) の実装になります。ただ、後ほど説明する \( \ln q(\boldsymbol{\tilde{W}};\boldsymbol{\eta}) \) の計算のところで、ガウス分布に入れる \( \sigma \) が0になってしまいエラーになるケースが発生してしまうため、0にならないようにするための補正 (self.approximate_posterior_sigma_eps) を加算しています。

変分推論部分の実装

勾配を計算するための式(5.269.1) の各項を計算します。これらの項はすべてガウス分布になっているのでPyTorchの torch.distributions.normal.Normal() を使えば簡単に実装できます。

\( \ln q(\boldsymbol{\tilde{W}};\boldsymbol{\eta}) \) と\( \ln p(\boldsymbol{\tilde{W}}) \) 、\( \sum_{b=1}^B \ln p(\boldsymbol{y}_b | \boldsymbol{x}_b, \boldsymbol{\tilde{W}}) \)のそれぞれの項は以下の関数で計算するようにしています。

    def _compute_approximate_posterior_log_prob_core(self, W, mu, rho):
      sigma = self.rho2sigma(rho) + self.approximate_posterior_sigma_eps
      W_dist = torch.distributions.normal.Normal(mu, sigma)
      log_prob_W = W_dist.log_prob(W)
      sum_log_prob_W = torch.sum(log_prob_W)
      return sum_log_prob_W

    def _compute_prior_log_prob_core(self, W, sigma_w):
      W_dist = torch.distributions.normal.Normal(0, sigma_w)
      log_prob_W = W_dist.log_prob(W)
      sum_log_prob_W = torch.sum(log_prob_W)
      return sum_log_prob_W

    def compute_approximate_posterior_log_prob(self, W1, W2):
      log_prob_W1 = self._compute_approximate_posterior_log_prob_core(W1, self.mu1, self.rho1)
      log_prob_W2 = self._compute_approximate_posterior_log_prob_core(W2, self.mu2, self.rho2)
      return log_prob_W1 + log_prob_W2

    def compute_prior_log_prob(self, W1, W2):
      log_prob_W1 = self._compute_prior_log_prob_core(W1, self.sigma_w)
      log_prob_W2 = self._compute_prior_log_prob_core(W2, self.sigma_w)
      return log_prob_W1 + log_prob_W2

    def compute_log_prob_p(self, Y, Y_est):
      Y_dist = torch.distributions.normal.Normal(Y_est, self.sigma_y)
      log_p = Y_dist.log_prob(Y)
      return torch.sum(log_p)

これらを使ってKLダイバージェンスの勾配に関係する部分だけ計算して、backward()とOptimizerのstep()を以下のように呼びます。

def vi_step(X, Y, model, optimizer, N, max_grad_norm=1e2):
  model.zero_grad()

  batch_size = X.shape[0]
  Y_est, W1, W2 = model(X)
  prior_log_prob_W = model.compute_prior_log_prob(W1, W2)
  posterior_log_prob_W = model.compute_approximate_posterior_log_prob(W1, W2)
  log_prob_p = model.compute_log_prob_p(Y, Y_est)
  kl_divergence = batch_size/N * (posterior_log_prob_W - prior_log_prob_W) - log_prob_p

  kl_divergence.backward()
  torch.nn.utils.clip_grad_norm_(model.parameters(), max_grad_norm)
  optimizer.step()
  return kl_divergence.item()

これらを使ってJuliaのサンプルと同様の訓練データで事後分布の推定を行って本の図5.18と似たような図を以下のように作ってみました。

図を見ると十分な精度がでていると思っています。

実装で苦労した点について

いざ実装してみると以下の点で工夫が必要だったので簡単に紹介します。

最初の数イテレーションの勾配が大きすぎてモデルのパラメータがnanになる

最初の数イテレーションはどうしても勾配が大きくなりがちです。今回は特にバッチサイズが大きいとモデルのパラメータが途中でnanになってしまうという問題が発生しました。この手の問題は深層学習ではよくあるためいくつか対処する手段はありますが、今回はシンプルなgradient clippingを利用しています。具体的には vi_step() でtorch.nn.utils.clip_grad_norm_()を呼んでいる部分がそれにあたります。

近似事後分布のσが0になる

バッチサイズや他のパラメータによっても発生したりしなかったりしますが、時々\( \ln q(\boldsymbol{\tilde{W}};\boldsymbol{\eta}) \)を計算する部分で \( \sigma \) が計算誤差で0に丸められてしまうとうケースが発生しました。このため、以下のようにして小さい値を加算して0になるのを防ぐ必要がありました。

      sigma = self.rho2sigma(rho) + self.approximate_posterior_sigma_eps
      W_dist = torch.distributions.normal.Normal(mu, sigma)

終わりに

「ベイズ推論による機械学習入門」の5.7 ニューラルネットワークのモデルをPyTorchで実装したので、実装についてと苦労した点についての記事を書きました。

実装する前はすぐできるだろうと思っていましたが、苦労した点に書いたような問題が出てきて思ったより時間がかかった印象です。ただ、実際に実装してみてベイズ推論の理解が深まったのでやってよかったです。実はこの本の他の実装もいつくかしてあるので機会があればまたblogの記事にしようと思います。

The post [勉強ノート]「ベイズ推論による機械学習入門」5.7 ニューラルネットワークで紹介されたものをPyTorchで実装してみた first appeared on まったり勉強ノート.