数値シミュレーションデータの低次元潜在空間における時間発展ダイナミクスの学習

Keisuke Fukuda

Engineer

本記事は、2022年度PFN夏季インターンシップで勤務された助田一晟さんによる寄稿です。

こんにちは、東京大学大学院修士1年の助田です。大学では数理最適化と理論統計学の研究に取り組んでいます。

PFNインターンシップでは、「気象学の物理知識をもとにした気象予測モデルの機械学習」に取り組みました。その中で、「数値シミュレーションデータの低次元潜在空間における時間発展ダイナミクスの学習」に挑戦しましたので、こちらを簡単にご紹介します

背景

現在の数値モデルによる気象予測は主にスーパーコンピュータを用いたシミュレーションを基に行われています。数値モデルは、流体力学、熱力学、化学反応などの物理過程を数値的に計算していますが、小さいメッシュ（高い解像度）でシミュレーションしようとすればするほど非常に大きな数のボクセルを扱う事となり計算コストが高くなってしまうゆえ、スーパーコンピュータが不可欠となっています。そこで、この物理過程の数値計算を省計算・省メモリで模倣できるニューラルネットワークを学習できないかと考えました。省メモリで高速な計算によるシミュレーション結果の再現が可能となれば、実行回数や解像度の制限が緩和され、気象予測手法の選択肢が大きく広がります。

気象シミュレーションでは、計算の対象となる地域を、例えば東西方向・南北方向・高さ方向の３次元に格子（メッシュ）状に分割し、それぞれの格子点に物理量（気温、気圧、湿度など）を割り当てます。仮に、東西方向に128分割、南北方向に128分割、高さ方向に16分割すると、128 * 128 * 16 = 262,144 個の格子が必要です。実用的なシミュレーションではより大規模な格子数が必要となります。例えば、気象庁の台風予報等[3]で用いられる格子間隔20kmの全球モデルであれば、数千万個以上の格子が要求されます。

ところで、大気を構成する物理量は、ランダムに分布しているわけではなく、ある程度の規則性と滑らかさをもっているはずです。格子に分布している物理量のデータの情報量は、メモリ上の表現より大幅に小さいことが予想されます[5]。このようなデータに対し、機械学習を用いて低次元の潜在表現への変換が可能であれば、情報の圧縮が達成できます。

さらに、シミュレーションの計算結果（格子上の物理量から格子上の物理量への演算結果）を学習データとして、同等の潜在表現から潜在表現への時間発展計算の学習を行えば、大量のメモリを消費することなく、省メモリで高速な計算によるシミュレーション結果の再現が可能となります。流体力学などのシミュレーションを機械学習モデルを用いて代替することをSurrogate Modelと呼びます。今回の取り組みでは、このSurrogate Modelを潜在空間（Latent）上で実現しようとしており、我々はこれを “Latent-based surrogate modeling” と称しています。

このような、現実を潜在空間上で表現し、シミュレーションを始めとして様々な操作や演算をする手法はState Space Modelingと呼ばれ、最近の例としては DayDreamer[4] などがあります。DayDramerは、ロボットの強化学習において、従来は時間がかかる処理であった物理シミュレーションを潜在空間上で行うことで、大幅な学習の高速化を達成しています。

なお、シミュレーション×深層学習というテーマは、PFNでもいくつかの先行事例[1][2]があります。

方針

今回のプロジェクトでは、3次元空間の物理量を低次元の潜在空間の変数にEncodeし、また潜在変数を元の3次元空間の物理量にDecodeするモデルAutoEncoderと、観測された物理量を変換した潜在変数を時間発展させるDynamicsのモデルを学習します。

学習データは、理化学研究所計算科学研究センター様に提供いただいた、スーパーコンピューター富岳を用いた気象シミュレーションデータを用いました。そのうち、2021年7月29日東京上空の17時時点における30秒おき30分先までの気象シミュレーションデータ（３次元空間、11物理変数）を、それぞれ16枚ずつに分割した、合計1952枚のデータを用いて実験を行いました。学習データのある時刻について、AutoEncoderの学習を行い、そのEncoderによって得られた潜在表現について、時刻 t から t+1 への変換を学習することによって潜在空間上のDynamicsの演算を獲得します。

実装

モデル構成

上図はモデルの全体像です。大きく２つの部分で構成されています。

1. AutoEncoder部
３次元ボクセルデータを低次元に圧縮するためのEncoderと、元のデータを再構成するためのDecoderです。最も性能が良かった3D UNetベースのネットワークを採用しました。ただし圧縮という目的を達成する観点から、オリジナルのUNetに特徴的な「特徴マップのconcatenate」は最深層部を除いて削除しました。

今回は、元のボクセル空間全体をEncode/Decodeするのではなく、空間を部分空間に切り分け、それら1つ1つに対してEncode/Decode、さらにDynamicsを適用することにしました。切り分けた部分空間の次元は (11, 23, 64, 64) です。各点の物理量が11個、高さ方向に23次元、東西南北の水平方向に64×64のサイズに切り分けたということになります。一方、それを変換した後の潜在空間（UNetの特徴量マップ）の次元は (512, 2, 8, 8) としました。
今回の実装では、 (11, 23, 64, 64) 次元であった元のVoxel空間の部分区間を、Encoderを用いて潜在空間へと変換すると、(512, 2, 8, 8) 次元になります。これにより、約15.8倍の圧縮が実現できます。

2. Dynamics部
圧縮された潜在ボクセルを、一期先の潜在ボクセルに変換するネットワークです。今回は試験的に１層のUNetを用いました。一期先（３０秒先）では物理量の分布が大きく変化しているとは考えにくく、恒等写像に近い変換となっているためUNetが有効と予想できます。また、数値モデルは現在の状態のみから次の時刻の状態を計算するため、LSTMのような履歴を利用した予測はとらず、現在の状態のみから一期先の予測を行いました。

【学習方法】

End to Endで学習する方法と、AutoEncoder部とDynamics部をそれぞれ分けて学習する方法が考えられます。今回は後者を採用しました。そのためタスクは

AutoEncoder部：入力画像を精度良く再構成する。
Dynamics部：一期先の潜在ボクセルを精度良く予測する。

となりました。Lossは両者ともMSE(Mean Squared Error)を用いました。なお、前者のタスクの学習ではvalidationを行なっていません。

実験結果

11 channelを有する３次元ボクセルデータの再構成を行いましたが、可視化の都合上適当に1つのchannelおよび1つのある高さのデータを２次元画像として表示した図を参考例として掲載します。

AutoEncoder部評価（0 timestep）

まず時間発展を考慮しない、単純なAutoEncoderの性能を評価します。

上が元の観測の画像、下にUNetで低次元に落とした潜在変数から再構成した画像になります。MSElossは0.006程度まで低下しており、観測画像の次元を(11, 23, 64, 64) 次元から(512, 2, 8, 8)次元に落としたUNetでも十分な表現力を有していることが示唆されます。

Dynamics部

次に時間発展を1 step適用した、一期先予測の性能を評価します。

trainデータに対する例 / testデータに対する例

trainデータに対してはMSElossが0.02程度と比較的良い再現が達成されましたが、testデータに対しては概観は一致しているもののMSElossと0.1程度と誤差が広がってしまっており、汎化性能が十分でないという結果が観察されました。

今後の課題

Dynamicsを表現したモデルは一期先予測のみならず、複数回適用することで任意の時刻のシミュレーション結果を予測することが理屈の上では可能です。しかし、ステップ数を重ねるほど誤差が広がってしまうので、十分な実用性を得るためには汎化性能の獲得が必要であり、今後の課題となっています。原因としてはUNetを用いた圧縮のプロセスで過学習していると考えられます。

総括

今回のインターンシップでは大変貴重なデータを用いて0からのテーマで研究を行うことができ、非常に良い経験になりました。短いインターンシップ期間中に挑戦できたことは限られてしまいましたが、研究としての今後の発展が大いに期待できるトピックであると強く感じることができました。

謝辞

今回必要となったデータを快く提供いただいたのみならず、実験計画の相談にも乗っていただきました理化学研究所計算科学研究センターの雨宮様、大塚様、国立研究開発法人情報通信研究機構の佐藤様、メンターはじめPFN社員の皆様、およびその他ご協力いただいた皆様にこの場を借りて御礼申し上げます。貴重な機会をいただきまして、誠に有難うございました。