もうスパコンは要らない!?AIが物理学を学習し物質の動きをシミュレーション、DeepMindの最新研究成果から

先々週、AI・機械学習の学会International Conference on Machine Learning(ICML)が開催され最新の研究成果が発表された。今年はコロナ感染拡大のためデジタル学会となり、欧米及びアジア諸国の研究者がオンラインで参加し、Zoomで講演する形式となった。

出典: Alvaro Sanchez-Gonzalez et al.

DeepMindの研究概要

この中でDeepMindはAIをシミュレータとして使う技法を発表した。シミュレータは物理現象をグラフィカルに表示する機能を持ち、水槽に水を注ぐと、AIがその動きを予測し、水の動きをビデオで表現する(上の写真右側)。実際の水の動き(上の写真左側)と比べると、複雑な動きをAIが正しく予測していることが分かる。

ニューラルネットワークでシミュレーション

これは「Graph Network-based Simulators」と呼ばれ、ニューラルネットワークでシミュレータを構築する。上の事例は、ニューラルネットワークが水槽に注がれた水の動きを予測したもので、初期条件を入力すると、ニューラルネットワークがその後の動きを計算する。つまり、ニューラルネットワークで水を表現し、それを動かすと、その後の挙動を推測する。

汎用のシミュレータ

Graph Network-based Simulatorsは、水のような液体だけでなく、砂やゼリーなど物理特性の異なる物質の動きを予測できる。水槽に水の塊を落とすと、その後の水の動きを予測する(下の写真上段)。同じニューラルネットワークが、ゼリーの塊を重ねると、それが崩れる動きを計算する(下の写真中段)。また、砂の塊を落とすと、それがタンク内に広がる動きを予測する(下の写真下段)。

出典: Alvaro Sanchez-Gonzalez et al.  

ニューラルネットワークの教育

ニューラルネットワークは実際の物質の動きを見て物理法則を学習する。教育の過程で、物質の動きを1ステップだけ教えると、ニューラルネットワークは数千ステップ先まで予測する。つまり、AIは物理法則を習得し、水槽に水の塊を落とすと、水が波打ちそれが鎮まるまで、遠い先の動きまで予測する。

シミュレーションの規模

更に、少量の物質(例えば水の分子2000個)を使ってニューラルネットワークを構成すると、ネットワークは大量の物質(水の分子85,000個)の動きを予測する。このため、少量の水で流れ方を教えると(下の写真、右上の箱)、ニューラルネットワークは大量の水の流れ方を学習する(下の写真、全体部分)。

出典: Alvaro Sanchez-Gonzalez et al.  

スパコンによるシミュレーション

物理現象のシミュレーションにはスパコンが使われる。スパコンは物質の動きをシミュレーションするために開発されたといっても過言ではない。事実、米国国立研究所Oak Ridge National LabはIBMのスパコン「Summit」を使って様々なシミュレーションを実行している。原子炉内部をスパコンでシミュレーションし、原子炉の耐用期間を延長する研究を展開している。

AIがスパコンを置き換える

スパコンによるシミュレーションで社会は多大な恩恵を受けているが、そのための対価が大きいのも事実である。IBM Summitのコストは2憶ドルといわれ、また、シミュレーショアプリの開発では数多くの研究者が必要となる。これに対し、Graph Network-based Simulatorsは汎用シミュレータで安価なAIプロセッサ(Google Cloud TPU)で動き、幅広い分野に適用できる。今すぐにSummitを置き換えることはできないが、AI開発が進むことでスパコンの一部をニューラルネットワークで代行できると期待されている。

【技術情報:Graph Network-based Simulators】

ニューラルネットワークの構成

Graph Network-based Simulatorsはニューラルネットワークで構成され、ネットワークのニューロンに物質の最小単位(例えば水の分子)を割り当てる。更に、ニューロン間の物理状態(分子の位置や速度、物質の特性、重力など)を指定する。これを実際の物理現象で教育すると、ニューラルネットワークは物質の動きを理解する。完成したニューラルネットワークに初期条件(水槽に水を灌ぐなど)を入力すると、その後の動きを予測する。

ニューラルネットワークの機能

Graph Network-based Simulatorsは物質の分子をEncodeし、これをProcessorで実行し、その結果をDecodeする(下のグラフィックス)。Encodeとは物質の状態(位置や速度や特性など)を凝縮しベクトルで表示する処理を指す。Processorは入力された分子の状態を元に、次の動きを予測する。Processorはこのプロセスを繰り返し、将来の動き(Mステップ先)まで予測する。DecodeとはProcessorの予測結果(ベクトル)を物質の状態に戻す処理をする。

出典: Alvaro Sanchez-Gonzalez et al.  

Message Passingという手法

Graph Network-based Simulatorsは物質の分子をネットワークのニューロンに割り当てるが、これら分子間の相互作用をメッセージ交換(Message Passing)として表現する(下のグラフィックス、中央)。メッセージを交換することで、分子は次の状態に移る。このプロセスを繰り返し分子の動きをMステップ先まで予測する。メッセージは分子の特性(物質の特性や重力など)と隣の分子との相互関係(距離や速度など)で構成される。

出典: Alvaro Sanchez-Gonzalez et al.