GomalizingFlow.jl 入門(Part3)

本日は

計算物理 春の学校 ２０２３ に参加された皆さんお疲れ様でした. ここでは GomalizingFlow.jl
のパッケージをベースに 4 次元の理論を動かすことを考えます. そこで得たアイデアは affine coupling layer を構成する畳み込みネットワークの構造を変更すること. 学習時の受容率の向上スピードを上げることができました. これは従来よりも短い時間で学習を効率できることを意味しています.

Part 1 からの復習

Part 1 の復習をします. https://zenn.dev/terasakisatoshi/articles/introduction-to-gomalizingflow-part1

$L$ を各軸の格子点の数とします. $d$ 格子上の場の理論では時空を離散化し有限サイズにカットオフする立場をとります. 周期境界条件を取っているとすれば格子 $\mathcal{L}$ は $(\mathbb{Z}/L\mathbb{Z})^d$ と思うことができます.

格子 $\mathcal{L} = (\mathbb{Z}/L\mathbb{Z})^d$ 上の場 $\phi$ は写像 $\phi : \mathcal{L} \to \mathbb{R}$ ですが, $\mathcal{L}$ は有限集合なので $\phi$ を多次元配列とみなすことができます. 例えば $d=2$ ならば行列です. 以下 $\phi$ は多次元配列と思ってください. また $\phi$ を配位と呼ぶことにします.

いま配位 $\phi$ が 1 次元のベクトルとして並べた時に $D(= L ^ d)$ (便宜上偶数)次元になってるとします. $\phi$ を二つに分割します. 添え字が even なのか odd なのかで $\phi_e$, $\phi_o$ のように分けましょう. 各々 $D/2$ 次元のベクトルと思えます.

Normalizing Flow は複雑な分布のデータからのサンプリング問題を微分可能で可逆な写像 (flow) を用いて簡単な分布(例えばガウス分布)
からのサンプリングに帰着させます.

$\varphi = T^{o}(\phi)$ を定義します:

$$\begin{aligned} \varphi_e &= \exp(s(\phi_o)) \odot \phi_e + t(\phi_o), \\ \varphi_o &= \phi_o. \end{aligned}$$

ここで $\odot$ は要素ごと (element-wise) の積を表すとします. $\exp(\phi_o)$ も $\phi_o$ の要素ごとに $\exp$ を適用していると思ってください. $s$, $t$ はパラメータ $\theta = (\theta^{(s)}, \theta^{(t)})$ を持つ機械学習モデルです.
数学的には $D/2$ 次元から $D/2$ 次元への(一般に非線形な)写像です. 奇数番目の値を変えて偶数番目の値は変えない変換 $T^e$ も同様に定義できます. この flow を affine coupling layer と呼びます.

Part 1 の議論で上記の変換は微分が可能で可逆な変換でした. また, 変換のヤコビアン(ヤコビ行列の行列式)も効率よく計算することができます.

GomalizingFlow.jl では上記のような変換を複数用意($T_i^o$, $T_i^e$, $i = 1,2,\dots, M$) します. $T_1^o$, $T_^e$ を交互に適用させた合成写像によって表現が豊かな flow を構築することを狙います.

ネットワークの構造

$T_i^o$ の定義にでてきた $s=s_i$, $t=t_i$ は $\theta_{i} = (\theta_{i}^{(s)}, \theta_{i}^{(t)})$ をパラメータとして持つ機械学習モデル, もっと詳しく言えばニューラルネットワーク, です. 具体的な実装に興味がある読者を対象にモデルの構造を述べていきます.

簡単のために二次元の理論, すなわち $d = 2$, の場合で考えます. このとき配位 $\phi$ というのは $L \times L$ の形状を持つ行列です. 行列という言葉を用い”ない”とすれば, 形状がタプル $(L, L)$ で輝度値が実数 $\mathbb{R}$ をとるチャンネル数が 1 の画像データとみなすことができます. 便宜上空間方向のサイズが $(L, L)$ でチャンネル数が 1 のデータを $(L, L, 1)$ という形状を持つ 3 次元配列のデータとみなします.

これの変換を考えようとすると畳み込みネットワークが思い浮かびます. $\mathrm{Conv_{1 \to 2}}$ をチャンネル数 1 から チャンネル数 2 の特徴量への変換だとします.

$$\mathrm{Conv_{1 \to 2}}: \phi \mapsto C(\phi)$$

$D$ 次元上の写像 $m_o=m_o(x)$ として次のような写像を考えます:

$$\mathbb{R}^D \ni x=(x_1, \dots, x_D) \mapsto (1, 0, 1, 0, \dots) \in \mathbb{R}^D$$

これはベクトルの添字の奇数番目の情報を残し, 偶数番目の項を無視するものとします. この $m_o$ を使うと $\phi_o \in \mathbb{R}^{D/2}$ は $m_o(\phi) \in \mathbb{R}^D$ と同一視することができます. $\mathrm{Conv_{1 \to 2}}(m_o(\phi))$ を計算すると $(L, L, 2)$ を形状とする配列ができます. 1 チャンネル目のデータ $(L, L, 1) \sim (L, L)$ を $s$ と置きましょう. これに $m_o$ でマスクをかけて $m_o(s) \in \mathbb{R}^{D}$ を $\mathbb{R}^{D/2}$ 次元のベクトルと同一視します. これで得られたものを $s(\phi_o)$ と書くことにします. 同様に $t$ についても $t(\phi_o)$ を定義できます. これによって affine coupling layer を計算することができます.

特徴量のデータレイアウトについて

深層学習フレームワークに慣れている読者は畳み込み計算においてデータレイアウトとして一番最後の次元がチャンネルを表すようにしたんだなと思っていただければOKです. これは $(H, W, C)$ のデータレイアウトを採用する TensorFlow と見た目似ていますが, どちらかというと $(C, H, W)$ のデータレイアウトを採用する PyTorch(バージョン 1) を逆順, すなわち右から左に並べた, $(W, H, C)$ だと思ってください. GomalizingFlow.jl が採用している機械学習フレームワーク Flux.jl が採用するデータレイアウトに合わせています. ミニバッチサイズをあらわす $B$ もこめるとデータレイアウトは $(W, H, C, B)$ というスタイルを採用しています. PyTorch では $(B, C, H, W)$ でした. 逆順になるのは Flux.jl が実装が PyTorch の実装にインパイアされている(Issue などを読むと PyTorch と比較している議論がある)のと Julia の配列のメモリレイアウトがカラムメジャーになってることに起因します. メモリレイアウトの関係でループの順がひっくり返るからです.
メモリレイアウトの話は 永井さんの１週間で学ぶ Julia の本にも簡単にですが触れられていますので必要に応じて参照してください.

畳み込みのパディングについて

畳み込みを計算する際に用いるパディングは場が周期境界条件を採用していることに対応して circular padding を採用しています. つまり左端の領域を畳み込み計算をする場合はぐるっと回って右端の情報を利用します. 上端は下端の情報を採用しています. これで画像としての空間方向の次元のサイズを保った変換を得ることができます. PyTorch だと padding_mode のオプションを指定すれば簡単に利用できます. ところが Flux.jl では未実装だったのと, Zygote.jl で微分可能な関数を定義するために特長量の右側の領域を切り取って左側にくっつけると言った配列操作を行なって実装しています. もっと良い方法があると思いますが，できたら教えてくれると嬉しいです. 私だけじゃなく Flux.jl の開発者が喜ぶと思います.

2 次元で説明した例の一般化をしたい

上記の例は $d=2$ の場合で説明しましたが, $d=3$ でも 3 次元畳み込みを使えば実装ができます. $d=2$ の場合に比べると必要とする GPU のメモリが増え, 学習時間は長くなりますが, ESS, acceptancce_rate を向上させるように学習が進むことを確認しています. となると $d=4$ の場合でも計算したくなるわけです. ただ, 4 次元畳み込みが Flux.jl で用意されていたなったのと, 必ずしもその次元のカーネル/フィルタを使った畳み込みである必要はない(原理的には, プログラムの上で自動微分可能のフレームワークにて計算できるに耐えうる写像であれば良い)というのもあって, 既存の部品
で作る方法を考えました. そこで得た知見をまとめたのが

https://twitter.com/TomiyaAkio/status/1583304573208645633

です.

Combinational-convolution for flow-based sampling algorithm

Combinational-convolution について

与えられた特長量空間において $d$ 本の軸から $k < d$ 本の軸を選びます. 選ばれた $k$ 本の軸に関して畳み込みを行います. $k$ 本の軸に関して畳み込みは次の例を挙げた方が良いのでそれで説明します.

ミニバッチサイズ が B, チャンネル数が inC, そして特長量の空間のサイズが (L1, L2, L3, L4) だとします. 各チャンネルが 4 次元のデータを持っているという設定です. タプル (L1, L2, L3, L4, inC, B) という形状を持つデータに対して次の手続きを経て outC をチャンネル数にもつ (L1, L2, L3, L4, outC, B) を形状とするデータを作ります. データの形状の移り変わりに注目してステップバイステップで説明すると下記のようになります:

(L1, L2, L3, L4, inC, B) # L1, L2 の軸に関して畳み込みをしたい
->
(L1, L2, inC, L3, L4, B) # データの軸を入れ替える
->
(L1, L2, inC, (L3, L4, B)) # L3, L4, B を一つの軸に注目する
->
(L1, L2, inC, (L3 * L4 * B)) # L3, L4, B を一つの軸にまとめる. これは reshape をすることに相当
->
(L1, L2, outC, (L3 * L4 * B)) # 2 次元の畳み込みを計算する. これで `outC` をチャンネルとする特長量が得られる.
->
(L1, L2, outC, L3, L4, B) # 最後の軸を (L3 * L4 * B) から (L3, L4, B) に分解する. reshape の操作をする
->
(L1, L2, L3, L4, outC, B) # L3, L4 の軸を移動させる.

これは選ばれなかった軸のフィルタサイズが 1 の 4 次元の畳み込みを実行していることになります. 上記の例では (f1, f2, 1, 1) をフィルタとする 4 次元の畳み込みを実行しており, この操作を 2 次元畳み込みで実現している様子になります.

軸の選び方

元々は 4 次元の計算ができたら嬉しいなというモチベーションで行っていましたが, デバッグ目的のために $d = 2, 3$ に向けての実装を行なった結果, 学習の効率化にも貢献していることがわかりました.

GomalizingFlow.jl では Normalizing Flow によって非自明化写像を計算しそれを提案確率とするメトロポリス・ヘイスティングアルゴリズムでサンプルを生成します. その際に受容確率が高くなるように学習するのが望ましいです. 実験における経験則から $L$ の軸, 次元 $d$ が増えていくと学習がしにくくなり, 特定の受容確率を実現するために必要な計算時間が増えていく傾向があります. そのため効率的に学習する方法も望まれていました.

Combinational-convolution のアーキテクチャを使用すると $d=3$ の場合, 標準的な畳み込みを使わずに低い次元の畳み込みを使う方が 50% の受容率を得るために必要な epoch 数を削減することができています. (横軸の表示がバグっていますが, 500 まで epoch を回したものだと読み替えてください.)

興味深いことに $d = 3, 4$ の場合, 軸を 1 本だけ選ぶ ($k=1$ に相当) ようにした場合, それ以外の選び方に比べ効率よく学習できている様子がわかります(下図参照).

大雑把に言えば $k=1$ の場合は十字架, テトラポットのような形をしているフィルターを用いて畳み込みを実施していることとに相当しています. なぜこれが良いかというのは十分な議論ができてませんが, 配位の分布を決める作用 $S$ に出現する離散ラプラシアン

$$\partial^2\phi(n) = \sum_{\mu=1}^d(\phi(n + \hat{\mu}) + \phi(n - \hat{\mu}) - 2\phi(n)).$$

が関係してるのかもしれません. つまり $n$ が固定されているとき $n \pm \hat{\mu}$ は上下左右に伸びている座標でその形状が $k=1$ でのフィルターのそれと類似するからです. 通常の畳み込みだと斜め方向の位置にあるピクセルの値も計算するのでそれが悪さをしてたのかなど色々想像はできますが自分が保有している計算資源の関係の制約で十分検証はできていません（そういう意味でオープンプロブレムです）.

動かし方

メモリが 16 GB 以上ある GPU の環境が望ましいです.

$ git clone https://github.com/AtelierArith/GomalizingFlow.jl
$ julia --project=@. -e 'using Pkg; Pkg.instantiate()'
$ cp playground/notebook/julia/4d_flow_4C2.md .
# jupytext を使って md 形式から jl 形式へ変換.
$ jupytext --to jl 4d_flow_4C2.md
$ julia -e 'include("4d_flow_4C2.jl")'

最近ではデバイスの進歩に伴ってメモリが潤沢なハードウェアが出てきています. 生成モデルや大規模言語モデルの実行のためにGPUを買う人が増えていますが, 残念ながら私は資金的な意味でその波に乗れずにいます.

まとめ

• GomalizingFlow.jl における affine coupling layer の詳細に触れました.
• 4 次元の理論での flow-based sampling アルゴリズムのアイデアを紹介しました.
• Combinational-convolution の解説をしました.
• 上記アーキテクチャによる知見をまとめました.
• 動かし方書きました.