[1906.01563] Hamiltonian Neural Networks [paper-reading]
Opened this issue · 12 comments
論文リンク
https://arxiv.org/abs/1906.01563
公開日（yyyy/mm/dd）
2019/06/04
概要
ハミルトン形式の議論を参考にして、保存量であるべきエネルギーが保存されるように教師なしで学習するモデルを考案（モデルが保存量とするのは正確にはエネルギーではない）。
位相空間での $ q, p $ を入力としてハミルトニアンの値を出力し、それを微分で戻していって $ \dot{q}, \dot{p} $ を得て、ハミルトン方程式を満足するように loss function を設定する、というのがモデルの作り方となっている。
いくつかの一体問題や二体問題を解くことで、baseline となる通常の NN ($ q,p $ を入力として $ \dot{q}, \dot{p} $ を出力とする) では保存量がうまく保存されない問題を正しく解くことができることを示した。
著者実装のソースコード: https://github.com/greydanus/hamiltonian-nn
物理におけるシミュレーションを NN でやるという話を聞いたりして、それらがどんなもんなのかというのはちょっと調べてみたいと思っていた。
そもそも保存則とかの知識をうまく取り入れる必要があるんじゃないのかなぁと思っていたが、そこで発見したのがこの論文で、いかにもエネルギーが保存できそうな感じだったので読んでみることにした。
chaotic three-body problems を従来手法の 1 億分の 1 の計算時間で解いたという 論文 が話題になってた。ただこの論文は眺めたらちょっと表層的な感じがしたので、もう少しちゃんと理解できそうな論文ということで今回の論文を選んだ。
論文のモチベーションは保存則をうまく取り入れたネットワークを構築することだが、そのアプローチは以下とのこと。
Instead of crafting the Hamiltonian by hand, we propose parameterizing it with a
neural network and then learning it directly from data.
これの意味するところは初見ではちょっとよく分からなかった。
物理系を定めればハミルトニアンを計算することができるので、手で計算できるというのはそういうことを言っている。やり方としては初期状態のエネルギーを計算して、それが一定であるという hard-constraint を設定して解いていくということもできるだろう。
ただし NN で hard-constraint を入れるというのは筋は良くない（と思う）。この論文では、ハミルトニアン（全エネルギー）を予測対象とする NN を構築して、それを微分して一般化座標と**量の微分を求めてエネルギー一定値を満足する方向を定める loss function を設定する、という手法を採用している。
学習データを作る際には予測対象のエネルギーを計算して与えているわけなので、そういう意味では hand-crafted だとは思うのだが、とにかくこの論文での戦略はそうなっている。
本題に入る前にハミルトン力学の復習。特に知らないことはないので理由とかはすっ飛ばして簡単に概略だけ書いておく。
相空間を考えて一般化座標と**量を $ q = (q1, q2, ..., qN), p = (p1, p2, ..., pN) $ と書いて、これらを系の「座標」と呼ぶ。
$ S(q, p) $ を座標の時間微分とすれば、系の時間発展 t0 -> t1 は以下のように書ける。
ここで、ハミルトニアン $ H(q, p) $ を以下を満たすように定義すれば、これが系の全エネルギーになっている。これらの方程式を正準方程式（ハミルトン方程式）と呼ぶ。
これで時間発展がわかり、先ほどの $ S(q, p) $ は $ S(q, p) = (∂H/dp, -∂H/dq) $ とすればよい。この微分を symplectic gradient と呼び、これは（正準方程式を考えれば明らかに $ dH/dt = 0 $ なので）系のエネルギーを保ったまま**する方向を与えるものになっている。
ハミルトン力学はラグランジュ力学と同様にニュートン力学の再定式化したものであるので、ごちゃごちゃ言わずニュートン力学でやっとけばいいのでは？と思うかもしれないが、一回微分の方程式に落ちてるので数値計算上有利だったり、上で見たようにエネルギー保存がわかりやすい形になっていたりするので利点がある。
モデルの構築について既に軽く触れているが、ハミルトニアンを出力として、それを微分して $ \dot{q}, \dot{p} $ を得て正準方程式を満たすように loss を設定する、という戦略になっている。
概念図は以下の図の下段にある Hamiltonian NN。
loss function は以下のように設定する。
以上。
ということで物理をある程度分かっていれば理解するのはそんなに難しくない（本当にうまくいくかは置いておいて）。
このモデリングの利点として以下の三点を挙げている。
- loss function によって全エネルギーの保存が担保されるようになっている
- (q_{t}, p_{t}) <-> (q_{t+1}, p_{t+1}) というように map できる
- 全エネルギーを操作して、系にエネルギーをもっと与えたら（もしくは取り除いたら）どうなるかというようなこともシミュレートできる
ここまで読んできて論文では「全エネルギー」とは書いてなくて「全エネルギー的なもの」と書いてあって違和感がある。というのも軽くコードを眺めた段階では total energy とかを求めていたので、それがモデルの出力となるような学習もしているものだと思っていたからだ。
結論から言うとこれは間違いだった。Hamiltonian NN ではデータの total energy の情報は使っていない。
あくまで、
- 相空間での入力を受けて何かしらのスカラー値を forward で出す
- そのスカラー値を q, p で微分したものを (∂H/dp, -∂H/dq) として扱う（ここでの H は系のハミルトニアンであることは要請されておらず、この段階では上のステップで計算した単なるスカラーの微分）
- データから q, p の時間微分を計算してそれを教師データとして、モデルが計算した (∂H/dp, -∂H/dq) と合わせて L2 loss を計算して学習
というステップになっている。
つまり、ここでの $ H $ は物理で扱うところのハミルトニアンと完全に一致する保証はない。実際の物理的な**を教師データにしているのでかなりそれに近くなるはずだが、完全には一致しなくてもよくなる。ここはのちの結果とかを見てもちょっと疑問も残る部分だが、まあ原理的にはそうだと納得して進む。
細かいことを言っておくと、スカラー値を算出すると言っていたところはもう少し入り組んだ構造になっていて、モデルの出力は保存力場と湧き出しなし場で分けられて、それぞれの微分を計算して足し合わせている。ただし実験としてはポテンシャルなしのものを考えているのがほとんどかな。
まずは簡単な一体問題として以下の問題を解く。
- 調和振動子
ハミルトニアンは以下。実験では $ k = m = 1 $ で初期値のエネルギーは [0.2 ,1] からサンプルしている。30 個の観測点を有する軌跡を 25 個作り、それらには $ σ^2 = 0.1 $ のノイズを加えている。
- 摩擦なし振り子
ハミルトニアンは以下。実験では $ m = l = 1 $ で $ g= 3 $ にしていて初期値のエネルギーは [1.3, 2.3] からサンプルしている。あとは調和振動子の場合と同じ。
- 摩擦あり振り子
上の二つは理想的な系なので、教師データは初期条件を与えて微分方程式を解く（scipy の scipy.integrate.solve_ivp を使用）ことで学習データを作成している。摩擦あり振り子の場合は この論文 から持って来てるとのこと。この論文自体は観測データから保存則を求めるというものらしいが、データが公開されている。
比較手法として、直接 \hat{q}, \hat{p} を出力する NN を使っている。提案手法も比較手法も 3-layers で 200 hidden units MLP のモデルを使っている。
Ground Truth は上で述べたように真面目に**方程式を解いたものになっている。
結果は以下。
それぞれテストデータの結果で、一番左が相空間のプロット、二番目が Ground Truth との MSE loss、三番目がモデルの forward の出力値、四番目はハミルトニアンの値になっている。
まず、baseline NN はうまく物理的振る舞いを捉えきれておらず、保存則を満たしていないことが分かる。
また、Hamiltonian NN はハミルトニアンと同じ（しかし少し異なる）保存量を出力できていることが分かる。これがなぜハミルトニアンと違うのかというと、論文では基準系が異なればポテンシャルエネルギーは定数倍分だけ異なるし、そしてそれは物理的には完全に acceptable という説明をしている。
ここは先にも述べたようにちょっと納得いってない部分だが、なぜデータでは周期的な違いが現れているかがこの基準系の説明とうまくマッチしない。baseline も周期的になっているので何かしら意味のある解釈はできると思うのだが。調和振動子の周期に合わせてうまく学習できてない部分があるのかな。基準系の話との結び付きで言えば、調和振動子の周期のそれぞれで本来のハミルトニアンの学習がうまくできていない要素を、適宜基準系を変えることでハミルトニアンとは別の保存量という形で学習してくれているということかなぁ。
最後の摩擦ありの場合は、Hamiltonian NN は保存量を保っているという結果になってしまう。これは loss の構成が保存量を持つように方向づけているのが問題であって、これは Hamiltonian NN の弱点であるので散逸系を考える場合は別のアイデアで対策をする必要がある。
次はもう少し複雑なケースとして、重力で引き合う二粒子が回転している連星系のような状況を考える。ハミルトニアンは以下。
重心系を考えることにして $ p_{CM} = 0 $ とする。
$ m_1 = m_2 = g = 1 $ として、$ μ $ は換算質量なので 1/2 となる。
$ x, y $ のそれぞれにおいて二次元平面で $ q, p $ を考えるので自由度は 8 となり、難しさ的には前の例よりも高くなる（保存則などを拘束条件として実際の物理系では自由度はこれより小さいが、いま Hamiltonian NN は物理のことを知らずにただ loss function を minimize するだけなので自由度は 8 考えてよい）。
データの作成は以下のようにしている。
- 重心系の**量は 0
- 半径 |q_2 - q_1| = r は [0.5, 1.5] からサンプリング
- 数値計算の安定性のため、初期値の速度は完璧に円軌道を描くものとして与えて、そこに σ^2 = 0.5 のガウシアンを加える
- 4 次のルンゲクッタで 50 個の観測点を持つ 200 個の軌道のデータを作成し、8:2 = train:test とする
結果は以下。
baseline は全然ダメですぐ collapse してしまう動きを見せる。
Hamiltonian NN は二週目に入るところで軌道が膨らむので、baseline よりは遥かに良いが完璧にはシミュレートできていない。しかし total energy は保存する振る舞いを見せていて保存則はちゃんと守ろうとしていることが分かる。また、total energy は調和振動子の場合と違ってちゃんと物理的なエネルギーになっているっぽい。
これは結構うまくいってる場合で、seed 値によっては total energy がちゃんと保存してくれないケースも出てくるが、baseline と比べると常に良い結果を返してくれることを Appendix で示している。
二体問題も完璧ではないがかなりうまく学習できるということが分かった。
次に、より challenging な問題として振り子の画像データから直接学習をするというタスクを実施する。
データは OpenAI Gym’s Pendulum-v0 environment から持ってきていて、元の 400 × 400 × 3 を 28 × 28 × 1 にして、次の frame のデータとセットにして 28 × 28 × 2 をバッチ毎に扱う。2 frame を合わせて扱っているのは速度を扱うためである。
AutoEncoder を使ってモデル自体は複雑ではないが、これまでとはだいぶ趣が異なることをするので注意が必要である。
latent は 2 次元で $ z_q, z_p $ に対応するものと考え、画像から一般化座標と**量の情報を抽出するものとする。そこから decode して画像を再構成するところは普通の AutoEncoder になっている。
latent の情報を時間微分すると $ dq/dt = ∂H/dp $ などが得られるので、それを使えば time step を一つ進めることができるので、2 frame で扱うときに 1 つめの time frame で時間発展を一つ進めた結果と 2 つめの time frame の結果を使って HNN loss を計算するということになっている。
この辺は頭での理解とコードがちょっと結びついてないところがある。あとでもうちょっとちゃんとコードを追って理解しよう。
loss の作り方は工夫が必要で、以下の三つの項から成る。
- 上で説明した HNN loss
- 普通の AutoEncoder loss (画像ピクセル値の L_2 loss)
- latent space で以下の loss を入れる。これは z_q, z_p が一般化座標と**量の性質を満たすように要請している。これを入れても AutoEncoder loss に悪さはしなかったとのこと。
実験結果として、再構成した画像と振り子の角度を示したものが以下。
今回は latent space が $ q, p $ の働きをするので、時間発展を解いているのはここで、再構成画像は学習した decoder で求めていることになる。baseline では振り子がすぐ止まってしまっているが、提案手法ではちゃんと揺れ続けていることが見て取れる。
ここはコードも見直してみてほぼ理解したが一点理解できてない点がある。
まず、latent space の出力を一般化座標と**量として扱っているというのはその通りである。
画像を入力にして latent を出力するのは AutoEncoder の encoder の仕事だが、それを入力にしてハミルトニアンを出力するのはこれまでの実験で使ってきたものと同じ HNN である。これに z を入力して得られるハミルトニアンを微分すればこれまで同様の HNN loss が計算できる。
問題は $ L_{CC} $ の loss だ。これは $ z_q^t - z_q^{t+1} = \dot{q} = ∂H/∂p $ である。ただこれを $ z_p^t $ と同じだと考えるならば、$ m = l = 1 $ の摩擦なし振り子の H の式を陽に使わないといけないと思うんだよな。
そう考えればまあ矛盾はないのだが、論文のモチベーションとずれる気もするし、ここだけはちょっとどういう意味なのか理解できていない。
これまでの実験の loss と energy の （Ground Truth との差）結果をまとめると以下の表となる。
興味深いのは、loss の意味では差がないが、energy で比べると提案手法が圧倒的に良いということが分かる。これはそういう風にモデルを作ったのでそれはそうなのだが、物理シミュレーションを NN で解くには、やはり人間の持つ知識をうまくモデルに取り入れてやることが重要だという意味で興味深い。
最後に、物理的でない動きも誘発できるということで、symplectic gradient ではなく、$ R_H = (∂H/dq, ∂H/dp) $ で発展をさせてみる。これはエネルギー保存則を破ってエネルギーを注入したり取り除いたりすることになる。
以下は pixel 振り子の latent space の実験で、途中の紫が $ R_H $ の方向で、これによってエネルギーが注入されてより振り幅の大きな振り子へと変換されたことが見て取れる。
あとは $ (q_t, p_t) \rightarrow (q_t+1, p_t+1) $ が invertible でこれはリウビルの定理が成り立つからだとか述べているが、特に何か意味のある主張をしているわけではない。
計算の途中結果を保持しなくても逆解きできるのでメモリ効率がよくなる (Neural ODEとか) という話だが、この論文では HNN に関して何か実験をしていたりするわけではない。
ということで一通り読んでみた。
ハミルトン力学を参考にエネルギー保存則を満たすようなモデルを考案して、実験的にも物理系を良くシミュレートできたという話だった。
個人的にはやっぱり物理との絡みだと興味深いところだし、最近はそういう話も増えてきているので、今後もちょこちょこウォッチしていこう。