本日は
計算物理 春の学校 ２０２３ に参加された皆さんお疲れ様でした. ここでは GomalizingFlow.jl に対するチュートリアルを書きます.
https://github.com/AtelierArith/GomalizingFlow.jl
GomalizingFlow.jl について
GomalizingFlow.jl は格子上の場の理論, 特にスカラー場 \phi^4 理論, における flow-based サンプリングを用いた配位生成アルゴリズムを提供するパッケージです．
https://twitter.com/MLPhysJP/status/1560741857222467584?s=20
これはプログラミング言語 Julia で書かれています. 実装のベースになったのは下記のツイートで紹介されている論文です:
https://twitter.com/MLPhysJP/status/1554395703429906433?s=20
上記の論文のチュートリアル実装が PyTorch を使って公開されています:
この実装を参考に GomalizingFlow.jl は書かれています. PyTorch 実装を単に移植しただけではなく, 物理系のパラメータを指定する設定ファイルを与えることでそれに対応した物理系をシミュレーションするように整えました. そのため, 2 次元の理論だけではなく 3 次元の理論も試すことができるようになりました. Flow-based サンプリングはニューラルネットワークを用いています. このネットワークをカスタマイズすることもでき, 4 次元の理論におけるサンプリングアルゴリズムを計算することも可能となっています.
https://twitter.com/AkioTomiya/status/1599079877168496640?s=20
格子上の場の理論
場の理論，(または場の量子論) は素粒子物理学の世界を記述する際に用いられる理論です. 格子上の場の理論は, 連続的な時空間を離散化し時空間を有限個の格子点として定式化し場はその格子から定まるものとして捉えます. この記述の立場に立つと我々はその格子を用いて定義されたもとで数値計算や理論を展開することができます.
つまり下記で述べるような数学的な定式化が可能になり, コードとして落とし込めることができます. かなり天下りですが, 記号を導入しましょう.
いま d を時空の次元とします. L を正の整数とします. 各方向の(= 軸の)サイズが L である格子点の集合を次で定めます:
\mathcal{L} \coloneqq \{n = (n_1, \dots, n_d) \mid n_\mu \in \{1,2,\dots, L\}, \mu \in \{1,\dots, d\} \} \subset \mathbb{R}^d .
この \mathcal{L} が d 次元の離散化された時空に対応します. ここではスカラー場 \phi の場合のみを考えます. スカラー場 \phi は形式的には n \in \mathcal{L} 上の関数 \phi=\phi(n)ですが, 計算機の扱いに馴染みのある方は実際にコードを読むときは格子 \mathcal{L} の位置 n に実数値 \phi(n) が配置されているオブジェクトとみなした方がわかりやすいかもしれません. プログラムでスカラー場を記述すると d=2 の場合行列, d=3, 4 では多次元配列として解釈できます. このオブジェクトを配位(configuration) と呼び, これも単に \phi と書くことにします. スカラー場 \phi^4 理論によれば配位 \phi は次の分布に従って与えられます:
\phi \sim P[\phi] = \exp(-S[\phi])/Z.
ここで S はユークリッド作用と呼ばれるものです. ここでは配位 \phi を与えたら実数値を与えるものだと思いましょう. 後で具体的な S は後で与えます. Z は分配関数とよばれる量であり次の式で定義されます:
Z = \int \mathcal{D}\phi \exp(-S[\phi]).
ここで \mathcal{D}\phi は
\mathcal{D}\phi = \prod_{n\in\mathcal{L}} d\phi(n)
のことを表します. 例えば 3 変数の関数 f=f(x_1, x_2, x_3) を変数 x_1, x_2, x_3 に関して積分するときに
\int dx_1dx_2dx_3 f(x_1, x_2, x_3)
のような記号を使うと思います. その時の dx_1dx_2dx_3 を \prod_{n} dx_n のように書いていると思ってください. d\phi(n) は n \in \mathcal{L} で添字づけられた変数であると理解しましょう.
放置していたユークリッド作用 S の説明に戻ります. 配位 \phi に対してスカラー場 \phi^4 理論に対応する作用 S[\phi] を次で定義します:
S[\phi]= - \sum_{n\in\mathcal{L}}\phi(n)\partial^2\phi(n) + \sum_{n\in \mathcal{L}} V[\phi](n)
ここで, V[\phi] はポテンシャル項と呼ばれ n\in\mathcal{L} に対して
V[\phi](n) = m^2\phi(n) + \lambda \phi^4(n)
で与えられるものだとします. m^2 や \lambda は定数です. 文献によっては 1/2 や 1/4! をつけることが多いと思いますがここではつけていません.
\partial^2\phi は離散化されたラプラシアンです. 具体的には下記のようになります:
\partial^2\phi(n) = \sum_{\mu=1}^d(\phi(n + \hat{\mu}) + \phi(n - \hat{\mu}) - 2\phi(n)).
さらに n + \hat{\mu} の記号の説明も必要です. n は格子 \mathcal{L} の点であったことに注意しましょう. n+\hat{\mu} は d 次元空間の \mu 軸の方向に対し n を出発点とし一歩進んだ点を表します. n - \hat{\mu} は n を出発点とし \mu 軸のマイナス方向に一歩進んだ点だとします. なお, 格子 \mathcal{L} には境界条件は周期境界条件を採用することにします.
これで格子上のスカラー \phi^4 理論を導入することができました. 定数 d, L, m^2, \lambda を与えればそこから機械的に上記の定めた定義式を求め，配位の分布の定義を数学的に(形式的にかつ厳密に)与えられることがわかるでしょう.
時空間を格子にせず連続したまま取り扱うと Z の定義式は無限次元の多重積分になります. つまり \mathcal{D}\phi は
\mathcal{D}\phi = \prod_{x \in \mathbb{R}^d} \phi(x)
という計算を行う必要が出てきます. このような問題点を回避するために格子の場の理論が有効です. また格子上で考えることで上記で述べたように理論を形式的に定義することができるのも魅力的です. この時点で場の理論そもそもわからずチンプンカンプンになってるかもしれませんが， とりあえず問題設定を形式的に理解することができれば OK です.
もし読者が Ising 模型を触ったことがあれば, S の部分がハミルトニアン, \phi が格子上のスピンの値を格納する配列と思っておけば Ising 模型のアナロジーとして捉えれば雰囲気は掴みやすいと思います.
格子上の場の理論における計算
ところで, 場の理論は量子力学と特殊相対性理論を土台にした理論です. 特に量子論の要請から物理量 O は確率的に振る舞います.
\langle O \rangle = \int \mathcal{D}\phi P[\phi]O[\phi]
によって期待値を得ます. どうやらこの計算ができれば人類は幸せになれるようです.
P[\phi] を計算するには原理的に Z と S が計算できれば良いです. S に関しては
S[\phi]= - \sum_{n\in\mathcal{L}}\phi(n)\partial^2\phi(n) + \sum_{n\in \mathcal{L}} V[\phi](n)
で与えられるものでした. 配位 \phi は格子上の点 n\in \mathcal{L} に対して実数値が乗っかっているオブジェクトとおもえばOKです. あとは右辺の式一つ一つを思い出せば機械的に計算できるようになってます. ここでの機械的というのは適切な定式化のもとでプログラムに落とし込めるという意味です. ひとまず理論物理がチンプンカンプンでも色々割り切って積分を数値計算を行えば人類に貢献できそうという一つの拠り所ができたのではないでしょうか.
高次元の積分の壁
上記の方針は間違ってないですが, Z を直接求めることが難しいという壁があります. 例えば d=4, L=10 の場合, 格子点は L^d = 10000 点になります. 我々が考えている積分の定義上, その数を変数とした多次元の多重積分を実行していることになります.（とても膨大ですが非加算無限の濃度の変数に関する多重積分を考えるよりはマシですね）.
この規模の高次元積分を台形公式などを使って真面目にすると現実時間では終わらないことが知られています(富谷さんの格子QCDの講義ノートを参照).
さてさて困りました. この手の計算では乱数を用いたモンテカルロ法をよく使います.
\langle O \rangle = \int \mathcal{D}\phi P[\phi]O[\phi]
を計算するために分布 P に従う配位をサンプルします. 例えば \phi_1, \phi_2, \dots, \phi_N のように得られたとしましょう.
\langle O \rangle の計算を下記によって実現させるのです:
\frac{1}{N}\sum_{i=1}^N O[\phi_i]
この方法のメリットは配位をいっぱい生成する手段が確立できれば積分対象の次元の高さに依らずサンプル \phi_i から定まる物理量 O[\phi_i] の平均を計算することに帰着できることです. この平均値のの期待値は \langle O \rangle と一致し 誤差は 1/\sqrt{N} のオーダで減少することがわかります.
注意したいのは上記の誤差評価は \phi_1, \dots, \phi_N が独立にサンプルされた場合であることです. 後述するマルコフ連鎖モンテカルロ法 によるサンプリングはその構成法からサンプルに相関が出てしまます. 特定の誤差に収めるために必要なサンプル数は理論的にもとまるものよりも多くの量を要請することになります.
マルコフ連鎖モンテカルロ法による配位の生成
格子上の場の理論のテキストを読むと配位の生成のやり方としてマルコフ連鎖モンテカルロ法 (MCMC) の一つである ハミルトニアンモンテカルロ法(HMC) がよく使われます. 例えば LatticeQCD.jl の README に書いてあるように LatticeQCD.jl は配位生成アリゴリズムとして HMC を採用した機能を提供しています. 自己学習モンテカルロ法 (SLMC) を使った手法も提供しているようです. SLMC については LatticeQCD.jl の開発者の一人である永井さんの解説 自己学習モンテカルロ法：機械学習を用いたマルコフ連鎖モンテカルロ法の加速 がとてもわかりやすいです. Flow-based サンプリングアルゴリズムも配位生成をする際に SLMC のアイデアを用いています.
簡単にですが, マルコフ連鎖モンテカルロ法を思い出します. これは点列
x_0, x_1, x_2, \dots, x_k, x_{k+1}, \dots
を逐次構成し所望の分布 P=P(x) に収束させるアルゴリズムでした. 点と言いましたが, 実際にはスカラー場やスピン配列などの多次元データです. x_{k+1} は x_{k} にのみによって決まります(マルコフ連鎖). ある x から x^\prime へ遷移する確率を T(x \rightarrow x^\prime) のようにして書くとしましょう. 遷移確率と呼ばれる T は所望の分布 P との詳細つりあい条件
P(x) T(x \rightarrow x^\prime) = P(x^\prime) T(x^\prime \rightarrow x)
を満たすように設計することで関数 O=O(x) の期待値を
\langle O \rangle = \int dx\ O(x) P(x) = \lim_{K\to\infty}\frac{1}{K}\sum_{k=1}^K O(x_k)
によって得ることができます. メトロポリス・ヘイスティング法は遷移確率 T を提案確率 f(x \to x^\prime) と 採択確率(または受理確率) A( x \to x^\prime) に分解します:
T(x\to x^\prime) = f(x \to x^\prime)A( x \to x^\prime)
A( x \to x^\prime) は詳細つりあい条件を満たすように
A( x \to x^\prime) = \min\left(1, \frac{P(x^\prime) f(x^\prime \to x)}{P(x)f(x \to x^\prime)}\right)
のようにして与えられます. プログラムを書いているとアルゴリズムの k ステップ目において x_k があるとした時に f 君が x_{k+1} を「作ってええか？」と A さんに尋ねる姿を思い浮かべることになるでしょう. A さんは \min の中にある式を計算の計算次第で「ええよ！ f 君がくれた x_{k+1} を採択するよ！」, 「やだ! x_{k+1} は x_k にするもん！」という採択または棄却するやりとりが聞こえるようになります. ゼロからできるMCMC に触れられているようにメトロポリス法, ギブスサンプリング, HMC などは今述べたメトロポリス・ヘイスティングの特別な場合と理解することもできます.
自己相関の壁
さて，採択する確率が低い(棄却される頻度が多い)と
x_0, x_1, x_2, \dots, x_k, x_{k+1}, \dots x_N
という点列の間に相関が出てきます. これは次の意味で人間にとって都合が悪くなります.（極端な例として）例えば 1,1,4,4,2,2,3,3 として得られたデータの平均を求めたいとします. 偶数版目のデータは 一つ手前のデータと相関があります. これは実質 1,4,2,3 という独立に得た 4 点のデータから求めるのと同じ程度の精度しか得られません. 得られる乱数の質が悪いと所望の誤差内にて計算するのに必要なサンプル数を増やす必要があります.
誤差はルートNのオーダーで制御できます. これは精度を１桁改善するには必要なサンプル数の桁を2つ増やす必要を意味します:
julia> 1/sqrt(100)
0.1
julia> 1/sqrt(10000)
0.01
julia> 1/sqrt(1000000)
0.001
したがって相関が少ないサンプル方法が望まれています. ゼロからできるMCMC でも例を交えて説明しています. Ising 模型で相転移点に近いパラメータで計算した際に自己相関が長くなるケースを示しています. これを臨界減速とびます.
場の理論でも 富谷さんの格子QCDの講義ノート, スライドの94p 付近を参照 にて述べられているように相転移点が連続極限に対応しています. (格子上の場の理論(青木) の繰り込み群と連続極限にその議論が書いてあります). 計算機上だと格子サイズ L を大きくする極限を考えていることに対応します.
経路積分の正当化のために時空を離散化したわけなので一旦格子で考えたら格子間の長さを狭める連続極限を考えるのは自然な要請です. ところが HMC などのマルコフ連鎖モンテカルロ法の効率と離散化誤差にはトレードオフが知られています. Flow-based な方法ではこの課題を解決するための手法の一つです.
まだ試していませんが, U(1) のゲージ理論だと Equivariant flow-based sampling for lattice gauge theory の FIG. 1 のようにサンプリングが改善しているようです.
Flow-based の文脈では「これでええか？」と尋ねる提案確率 f(x\to x^\prime) として目標とする分布 P を近似する \tilde{P}=\tilde{P}(x^\prime) を設計するアプローチを取ります. マルコフ連鎖は直前のサンプルから一定の手続きで新しいサンプルを作成しますが, \tilde{P} はその直前のサンプルにもよらずサンプルするという気持ちが込められてます. さらに A さんが「ええよ！」と採択する確率が 1 に近づくようにすれば自己相関の問題を改善することが期待できます. もし人間が適切な \tilde{P} を設計できれば採択確率 A の式は
A(x \to x^\prime) = \min\left(1, \frac{P(x^\prime)\tilde{P}(x)}{P(x)\tilde{P}(x^\prime)} \right)
となります. \min の 2 番目の引数が なるべく 1 になるようになれば良いわけです. 例えば P = \tilde{P} の場合 A は常に 1 の値をとります. ところで勘の良い読者は A の右辺で P(x) そして P(x^\prime) が出てきますが, 分配関数, すなわち, P(x) = \exp(-S(x))/Z の右辺に出てくる Z の計算が難しいのにどうやって計算すれば良いのかと疑問に思うかもしれません. これは良い疑問です. P(x), P(x^\prime) を各々計算するのではなく比 P(x)/P(x^\prime) を計算しています. 比の計算は Z を直接計算することを回避できます. 実際次のようになるわけです:
P(x)/P(x^\prime) = (\exp(-S(x))/Z)/(\exp(-S(x^\prime))/Z) = \exp(-S(x))/\exp(-S(x^\prime))
GomalizingFlow.jl ではスカラー場のための \tilde{P} の設計として Normalizing Flow による手法を実装しています. ここまでくると場の理論の話をしていたのですがいつの間にか機械学習の文脈に移ってきました. 実際, 後述するように \tilde{P} がニューラルネットワークを用いた機械学習モデルとなっており, 学習によりスカラー場の配位生成を実現しています.
ボックス・ミュラー変換から学ぶ Flow-based models の気持ち
Flow-based の flow の気持ちに寄り添うためにボックス・ミュラー変換を思い出します. 下記のような (x, y) \mapsto (z_1, z_2) を考えます:
\begin{aligned}
z_1 &= \sqrt{-2\log(x)} \cos (2\pi y), && \\
z_2 &= \sqrt{-2\log(x)} \sin (2\pi y).
\end{aligned}
ここで x, y は各々開区間 (0, 1) を動くものとします.
r = \sqrt{-2\log(x)}, \theta = 2\pi y とみなせば (z_1, z_2) は (r, \theta) を媒介変数によって曲座標表示がされていると見做せます. 後々のことを見据えてボックス・ミュラー変換を次の 2 つのステップからなる変換だとしましょう:
(x, y) \mapsto (r, \theta) \coloneqq (\sqrt{-2\log(x)}, 2\pi y) \mapsto (z_1, z_2) \coloneqq (r\cos\theta, r\sin\theta).
今 x と y が独立に区間 (0, 1) 上の一様分布から得られたとき z \coloneqq (z_1, z_2) はどのように振る舞うかを調べましょう.
確率変数 X と対応する確率密度関数が p_X = p_X(x) があるとします. Y = f(X) というように変換されたものは X がランダムに振る舞うから Y もランダムに振る舞うので Y も確率変数になります. この確率変数 Y の確率密度関数 p_Y = p_Y(y) はディラックのデルタ関数 \delta を用いて次のように与えることができます:
p_Y(y) = \int \delta(y - f(x)) p_X(x)\, dx\ .
上記の式は次のように解釈できます. 変換 f の値域に属する y が与えられたとき y = f(x) を満たす (f の定義に属する) x をかき集めてその p_X(x) を足す(積分) とみなせます.
特に f が可逆(逆写像 x = f^{-1}(y) を定義できる)であり(ヤコビ行列の行列式 \partial f^{-1}/\partial y が定義できるという意味で)微分可能だとします. 多重積分の置換積分の公式を思い出せば x = f^{-1}(\tilde{y}) という変換で形式的に
\begin{aligned}
p_Y(y) &= \int \delta(y - f(x)) p_X(x)\, dx \\
&= \int \delta(y - \tilde{y}) p_X(f^{-1}(\tilde{y})) |\partial f^{-1}(\tilde{y})/\partial \tilde{y}|\, d\tilde{y} \\
&= p_X(f^{-1}(y)) |\partial f^{-1}(y)/\partial y|
\end{aligned}
と変形できます. |\bullet| は絶対値を表します. この密度関数の関係式は X やその確率密度関数 p_X が人間にとって(解析的に記述できるという意味で)容易だけれど Y の(分布の形状が複雑で)取り扱いに難しい時に役に立ちます.
Normalizing Flow はこのような可逆な変換 f をいくつか繰り返すことで単純な形状の密度関数から出発し密度関数を得る(分布を得る)方法です. このときの f を flow と呼びます. f には学習パラメータが付いており学習によって適切な flow を作ることを目指します.
ここまで読んだ読者は想像が容易だと思いますが, GomalizingFlow.jl のパッケージ名の由来になってます. 読者は Normalizing Flow の方を覚えてください. Goma といっても通じません! このページの読者は全体を通じて十分なユーモアセンスを持っていると仮定します.
以下では変数を Y = f(X) のように変換させるという操作を反対方向から捉えることにします. すなわち, f に注目する代わりに g \coloneqq f^{-1} に注目し Y を X = g(Y) という変換で簡単な X に帰着させるという立場をとることにします. GomalizingFlow.jl に対応するプレプリント では g を自明化写像 (trivializing map) と呼んでいます.
このようにすると変数変換 f から定まる密度関数の変換式は p_Y(y) = P_X(g(y)) |\partial g(y)/\partial y| と書き換えられます. 逆元の記号 {\bullet}^{-1} を省くことで記述が楽になります.
もう一度ボックス・ミュラー変換の話に戻って z = (z_1, z_2) の分布を調べることにしましょう.
(x, y) \overset{f_1}{\longmapsto} (r, \theta) \overset{f_2}{\longmapsto} (z_1, z_2)
と捉える代わりに自明化写像を主軸に考えます.
(z_1, z_2) \overset{g_2}{\longmapsto} (r, \theta) \overset{g_1}{\longmapsto} (x, y)
この時 (z_1, z_2) \coloneqq f_2(r, \theta) \coloneqq (r\cos\theta, r\sin\theta) であり (r, \theta) \coloneqq g_2(z_1, z_2) \coloneqq f_2^{-1}(z_1, z_2) を与えることができることに注意します. (z_1, z_2) の密度関数 p_Z=p_Z(z_1, z_2) は (r, \theta) に関する密度関数 p_\Phi = p_\Phi(r, \theta) を用いて
\begin{aligned}
p_Z(z_1, z_2)
&= p_{\Phi}(g_2(z_1, z_2)) |\partial g_2 /\partial (z_1, z_2)| \\
&= p_{\Phi}(g_2(z_1, z_2)) |\partial f_2 /\partial (r, \theta)| ^ {-1} \\
&= \frac{1}{r}p_{\Phi}(g_2(z_1, z_2))
\end{aligned}
と与えることができます. ヤコビ行列の行列式は曲座標表示の計算でお馴染みのように
\begin{aligned}
\partial f_2 /\partial (r, \theta)
= \frac{\partial(z_1, z_2)}{\partial (r, \theta)}
= \det \begin{bmatrix}
\frac{\partial z_1}{\partial r} & \frac{\partial z_1}{\partial \theta} \\ \\
\frac{\partial z_2}{\partial r} & \frac{\partial z_2}{\partial \theta}
\end{bmatrix}
= \det \begin{bmatrix}
\cos\theta & \sin\theta \\
-r\sin\theta & r\cos\theta
\end{bmatrix}
= r
\end{aligned}
で与えられます. 続いて (r, \theta) と (x, y) の関係を議論します. (r, \theta) \coloneqq f_1(x, y) \coloneqq(\sqrt{-2\log x}, 2\pi y) のように定めていました. この逆写像である
自明化写像 g_1 は
(x, y) \coloneqq g_1(r, \theta) \coloneqq (\exp(-r^2/2), y/2\pi)
で与えることができます. これによって p_\Phi は (x, y) における密度関数 p_{(X, Y)} を用いて
\begin{aligned}
p_{\Phi}(r, \theta)
&= p_{(X, Y)}(g_1(r, \theta)) |\partial g_1/\partial(r, \theta)| \\
&= p_{(X, Y)}(\exp(-r^2/2), y/2\pi) \left |\det \begin{bmatrix} -r \exp(-r^2/2) & 0 \\ 0 & \frac{1}{2\pi} \end{bmatrix}\right| \\
&= p_X(\exp(-r^2/2)) p_Y(y/2\pi) \frac{1}{2\pi} r \exp(-r^2/2) \\
&= \frac{1}{2\pi} r \exp(-r^2/2)
\end{aligned}
のように変形できます. 区間 (0, 1) 上の一様分布 U(0, 1) の密度関数の定義と \exp(-r^2/2), y/2\pi が (0, 1) の範囲に収まることから p_X, p_Y の値は常に 1 として得らます. 結果としてヤコビ行列の行列式の項が分布を特徴付けるようになります.
まとめると次のように変形できます:
\begin{aligned}
p_Z(z_1, z_2)
&= \frac{1}{r} p_\Phi(r, \theta) \\
&= \frac{1}{r}\frac{1}{2\pi} r \exp(-r^2/2) \\
&= \frac{1}{2\pi} \exp(-r^2/2) \\
&= \frac{1}{2\pi} \exp(-(z_1^2 + z_2^2)/2).
\end{aligned}
これは二変数の標準正規分布の密度関数そのものです. このようにして一様分布から得たデータを使って正規分布に従うデータを生成することができます. これが有名なボックス・ミュラー法の動作原理になります. ボックス・ミュラー変換は 2 つの flow からなる Normalizing Flow とみなせるでしょう.
ヤコビ行列の行列式の計算の壁を解決
ところで, Normalizing Flow は単純な確率分布(例えば多次元の正規分布)からスタートしいくつかの flow を経て複雑な確率分布を得る手法でした. 格子上の場の理論の文脈であれば作用 S から定まる P(x) = \exp(-S(x))/Z またはそれを近似する \tilde{P}(x) を構成するために使います. 実際, 下記で述べるように, 格子サイズ分の次元の正規分布を flow によって \tilde{P} を作成することをしています.
ボックス・ミュラー変換を Normalizing Flow という立場からもう一度眺めてみましょう. この変換の良い理由として自明化写像の微分(ヤコビ行列)が簡単に計算できることがあります. つまり，密度関数の計算式
p_Y(y) = p_X(g(y)) |\partial g(y)/\partial y|
における |\partial g(y)/\partial y| が大学１年生が習う微積分程度の知識と手計算で厳密な値を求められることです.
素直に Normalizing Flow による計算を実装しようとすると「適切な変換の実装」,「変換の微分」と得られた行列の「行列式の計算」を強いられることになります. 特に 2, 3 番目は計算コストが高く数値計算の観点から技術的な問題がありました.
Flow-based generative models for Markov chain Monte Carlo in lattice field theory では flow の作成に Density estimation using Real NVP で導入されている affine coupling layer を採用し上記の問題点を解決しています.
affine coupling layer の説明をしましょう. いま配位 \phi が 1 次元のベクトルとして並べた時に D (便宜上偶数)次元になってるとします. \phi を二つに分割します. 添え字が even なのか odd なのかで \phi_e, \phi_o のように分けましょう. 次のようにして i 番目の自明化写像 \varphi = T_i^{o}(\phi) を定義します:
\begin{aligned}
\varphi_e &= \exp(s_i(\phi_o)) \odot \phi_e + t_i(\phi_o), \\
\varphi_o &= \phi_o.
\end{aligned}
ここで \odot は要素ごと(element-wise)の積を表すとします. \exp(\phi_o) も \phi_o の要素ごとに \exp を適用していると思ってください. s_i, t_i は学習パラメータ \theta_i を持つ機械学習モデルです. もう少し詳しくいうと画像でも使う畳み込みネットワークをいくつか並べた非線形変換になります.
このようにして得られる変換のヤコビ行列の行列式を計算すると行列の行と列の適当な置換によって対角成分に 1 またはベクトル \exp(s_i(\phi_o)) の成分が並んだ三角行列を構成することができます. D=2, 4 の場合で実際に手を動かしてみるとよいでしょう. \det の計算は列，行の置換で符号を除いて同じになることと, 三角行列の行列式は対角成分の積を取れば良い性質を使うとヤコビ行列の行列式の絶対値は \prod_j \exp(s_i(\phi_o))_j と等しくなります. ベクトル \exp(s_i(\phi_o)) の成分の積を計算するだけにフォーカスすれば良いわけです. ちなみにヤコビ行列の行列式の「絶対値」を求めれば良いので行と列の置換で生じる±の符号の心配は気にしなくて大丈夫です.
また \varphi = T_i^o(\phi) は構成法から \phi = (T_i^o)^{-1}(\varphi) を構成することができます.
\begin{aligned}
\phi_e &= \exp(-s_i(\varphi_o)) \odot \varphi_e - t_i(\varphi_o), \\
\phi_o &= \varphi_o.
\end{aligned}
\exp(a) の逆数が \exp(-a) となることをうまく利用していますね.
T_i^o およびその逆変換は配位の偶数番目値を更新し奇数番目の値は変えず変換されます. 奇数番目の値を変えて偶数番目の値は変えない変換 T_i^e も同様に定義できます.
GomalizingFlow.jl では T_{\bullet}^e, T_{\bullet}^o を交互に適用した自明化写像 \mathcal{F}を構成しその逆写像 \mathcal{F}^{-1} によって単純な分布から物理的に意味のある分布 P を近似する \tilde{P} を構成しています. \mathcal{F}^{-1} を非自明化写像(un-trivializing map)と呼ぶことにします.
なんだか難しそうなことを言っていて「よくわかんない」と感じるかもしれません. その場合はひとまず「affine coupling layer というレイヤーをいくつか積み上げたディープニューラルネットワークを定義したんだな・・・」と思っていただければOKです. \mathcal{F}^{-1} をプログラムで書くことになります.
学習法
非自明化写像つまり \tilde{P} には学習パラメータが入っています. そのパラメータを勾配法で学習し受容確率
A(x \to x^\prime) = \min\left(1, \frac{P(x^\prime)\tilde{P}(x)}{P(x)\tilde{P}(x^\prime)}\right)
がなるべく 1 に近くなることを目指します. 元々は自己相関長を改善するものですがそのためには受容確率を増やすことが必要条件となります.
損失関数は shifted KullbackLeibler (KL) divergence を採用します:
\begin{aligned}
D_{KL}(\tilde{P}||P) - \log Z
&=
\int \mathcal{D}\phi \tilde{P} (\log \tilde{P}(\phi) - \log P(\phi) - \log Z) \\
&=
\int \mathcal{D}\phi \tilde{P}(\log \tilde{P}(\phi) + S(\phi))
\end{aligned}
shifted のココロは -\log Z を付与することで -\log P を計算する時に出てくる分配関数 \log Z を打ち消すというものです.
上記の関数は \int が入ってて難しそうな雰囲気がありますが, 次の近似計算で実現します:
\frac{1}{M}\sum_{i=1}^M \left(\log \tilde{P}(\phi^{(i)}) + S(\phi^{(i)}) \right)
M はミニバッチサイズだと思ってください. \tilde{P}(\phi^{(i)}) は次のようにして計算できます:
\tilde{P}(\phi^{(i)}) = r(\mathcal{F}(\phi^{(i)})) \left|\frac{\partial \mathcal{F}}{\partial \phi}\right|_{\phi=\phi^{(i)}}
r は格子点数と同じ次元の標準正規分布の密度関数です. z^{(i)} = \mathcal{F}(\phi^{(i)}) となる z^{(i)} は
標準正規分布を生成するプログラムがあれば作れます. \phi^{(i)} は非自明化写像 \mathcal{F}^{-1}(z^{(i)}) を計算すればOKです. これはプログラムでかけます. ヤコビ行列の行列式の絶対値の部分は affine coupling layer の議論でベクトルの要素を掛け合わせるプログラムを書けばOKです. つまり，上記の式は実装可能な量でありコードとして落とし込むことができます. もし読者がニューラルネットワークを用いたコードの実装に自信があれば, 物理に詳しくなくても物理学の世界に貢献できるようになります.
これまでの流れをまとめる
詳細は置いておいてプログラムを組むと次のような流れになります:
- 格子上の場を定義します. 配位は多次元配列(時空 d 次元配列)と思えます.
- 作用は配位を入力とする関数です.
- 非自明化写像を定義します. affine coupling layer によって作られるニューラルネットワークを実装します.
- 多次元配列の要素次元数の正規分布から出発し非自明化写像を計算しその点 x における \tilde{P}(x) を計算します.
- shifted KullbackLeibler (KL) divergence を計算, 勾配法で非自明化写像のパラメータを学習.
- 学習後: A(x \to x^\prime) = \min\left(1, \frac{P(x^\prime)\tilde{P}(x)}{P(x)\tilde{P}(x^\prime)}\right) を受容確率とするメトロポリスヘイスティング法によって P(x) = \exp(-S(x))/Z に従うサンプルを作成する.
まとめ
GomalizingFlow.jl が理論上何をしているのかの解説を書きました.
Part2 はこちら
https://zenn.dev/terasakisatoshi/articles/introduction-to-gomalizingflow-part2
Discussion
続きとしてどうぞ:
Idea ジャンルもどうぞ