ChatGPTが初めてリリースされたのは2022年の11月、つまりほんの少し前のことなのですが、もう「ChatGPT以前・以後」と呼び分けられそうなほど強いインパクトを多くの人に与えています。そうした「AI」と呼ばれるシステムが日々われわれの日常生活に入り込み、そのたびに誰もが「また1つ、新しいことができるようになった」と感じています。

本記事で紹介するのもこうしたAI、その実体であるディープラーニング技術の「新しい」展開についてです。取り上げる事例はスーパーコンピュータ「富岳」による数値計算によって行ってきたシミュレーションの一部を、ニューラルネットの推論処理で加速するという、東京大学の平島敬也さん（博士課程2年）と藤井通子准教授による研究[1]です。つまり数値シミュレーションをニューラルネットによって代替するのです。

今回、藤井氏とディープラーニング処理の実装を担当したMorpho社のスタッフにお話を聞く機会を得ました。日々新しい領域に切り込んでいくAI技術の展開と、その応用を支えるソフトウェア・エンジニアの仕事の両面をお伝えできればと思います。

研究の内容

藤井氏の研究は「銀河形成シミュレーション」と呼ばれるものです。つまり空間に粒子をいくらか密度に揺らぎを与えて配置し、自己重力によって変化していくさまをシミュレーションします。すると星間ガスから星が生まれ、また銀河（星が多く集まった天体）などができていく様子が見られます。

藤井: これ（写真1の左側）がたとえば我々のシミュレーションが出した画像ですが、真ん中にポツポツとある白いのが星、モヤモヤとしているのが星間ガスと呼ばれる物質です。それ以外にもダークマターと呼ばれる重力はあるが正体のわからない物質があり、だいたいこの3つの成分をシミュレーションで表現して宇宙の進化を計算していきます。

これで「それなりに現実に近いもの」ができあがることがまず驚きですが、ただ細かいところまで全部重力の計算だけで行っているわけではない、とのこと。

藤井: たとえば星ができることに関しても、実際に星ができるところまで星間ガスが集まり、温度が上がってきて、といったところまで全部は計算しません。密度が高いところに星ができそうな条件というのがもうわかっていますから、（⁠全部計算せずに）もう星ができたと仮定して星を置くような処理をします。星が死んだときに超新星爆発で周りのガスを吹き飛ばすところもそうです。

こうしたモデルを入れて処理しないと、今のコンピュータでは（処理量的に）計算できないところがたくさんあるのです。

ただそうするとモデルに依存する部分が出てくるので、今だといろんなモデルを試して、これがより現実に近いとか、そういう議論（研究）をしています。現実に近い結果が出せればそこで仮定したモデル、この物理現象を考えたメカニズムの正しさ（妥当性）が証明できる、ということです。

できるだけ高い分解能で、その数値誤差をなるべく排除して、かつモデル依存性を下げる、つまり「入れた物理がそのままちゃんとシミュレーションに反映される」ようなものを目指しています。

実装上の限界

安田: 藤井さんはこのソフトを書いておられるのですね。（⁠はい）何人くらいで書いているのですか？

藤井: 2、3人です。コアで書いているのは1人です。おおまかに何種類かのコードがあって、私は粒子を扱うN体計算をやっています。その中でも星がたくさん集まってくると頻繁に重力評価してやらないといけないので、その並列計算の効率の良いアルゴリズムを考えて書く、といったことを主にやっています。

シミュレーションソフト全体では重力だけを扱うN体計算以外に、星が生成される時のガスの性質を扱うための流体計算をSPH法[2]で処理しているとのこと。取材時は「富岳」のような多数ノードによる大規模並列計算で処理する際の工夫などについても丁寧に説明していただいたのですが、本稿ではディープラーニング処理との関係に集中するために割愛します。

藤井: 私たちシミュレーション側（の人間）としては、できるだけ小さいスケールまでちゃんと計算したい、という希望があるのですが、今は（ノードがたくさんあっても）計算パワーを使い尽くせないような段階に来ています。

安田: 何が理由ですか？

藤井: 時間刻みの幅ですね。より小さいスケールまで分解していくと、より短い時間刻みの現象が出てきます。数百年で変わるような現象を扱うためには、その時間刻みで周りのノードたちと通信をします。するともう通信の割合が計算に対して高くなってしまい、それ以上ノード数を増やしてもスケールしないところまで結構来ています。

安田: なるほど。

藤井: そこで今挑戦しているのは、（⁠写真2の右側図）一部を切り出してそこだけ少ない並列度で細かな時間刻みで繰り返して計算し、その一方で全体はより粗い時間刻みで1ステップ分やり、後で切り出した部分の結果を戻す、と言う方法です。

つまりこの切り出した小部分の細かな時間刻みのシミュレーションを、数値計算でなくディープラーニングの推論で行うことで加速する、という手法が今回のテーマです。

SoftNeuroによる実装

これ以降は、切り出した部分の処理の実装についてMorphoのスタッフ（代表取締役社長 平賀督基氏とCTO室のシニアリサーチャー松尾恒氏、同リサーチャー三宅博史氏他）にお話を伺いました。

まず藤井氏ら自身で富士通から出ていた「富岳」向けTensorFlowを用いてディープラーニングによる推論処理を試したところ、速度が出ませんでした。このとき平島さんがこの論文[3]に出会い、Morpho社と連絡をとったとのこと。

arXiv経由というのがなんとも今どきです。論文のタイトルにもなっているSoftNeuroはディープラーニングの推論処理に特化した、CPU, GPU, DSP などさまざまなデバイスで高速に動作するエンジンです。今回の処理はSoftNeuroを用いて実装されています。

平賀：もともと平島さんはTensorFlowで推論処理を書かれていました。「⁠富岳」で標準的に動く推論環境もTensorFlowなので、それを使って「富岳」上で動くようにしていました。これを我々がSoftNeuro向けにCで書き直した、という形です。

平島さんの推論モデルでは3Dコンボリューションを多用しており、計算の90％か、それ以上を占めていました。ただ3Dコンボリューションはディープラーニング分野ではメジャーなオペレーションではなく、「⁠富岳」のTensoFlow エンジンの中であまり最適化されておらず、ナイーブな実装しかないような状態でした。

松尾: またあまり効率的にSVE[4]が使われていませんでした。

安田: SVEはArmがHPC用に作ったのではなかったですか？

松尾: そうです。

平賀：もともとSoftNeuroは組み込み用なのでプラットフォームにはArmが多く、neonなどのSIMD命令を使って作り込んだものを我々は持っていたのですが、SVE向けに作り込んだものはなかったのです。「⁠富岳」のアーキテクチャ向けにはSVEをフル活用する必要があったので、SoftNeuroのそこの部分はもうスクラッチから書き直しました。

安田: 元々あった平島さんのモデルを残してソフト周りは書き替えた、と。そこでSVEを用いたのであれば作業自体はArm汎用の機能で行ったわけですから、とくに「富岳」用に何かをしたと言うことは無さそうですね。

松尾: はい。そこまで「富岳」に特化したことはしていないので、SVE があれば何にでも使える状態です。ただパフォーマンスチューニングするために、intrinsic 命令を手作業で入れて、キャッシュのプリフェッチなどを調整しました。

安田: 細かくキャッシュプリフェッチすると言っても、普通にやると余分なデータがキャッシュラインに道連れに上がってきちゃいますよね。つまり元々のデータが並んでいないと意味がない。そのあたりの調整もやっているのですか？

平賀：そこも取り組んでいるのがSoftNeuroの高速な理由の1つです。あらかじめメモリ上の位置が計算の順になるようウェイトのパラメタを並び替えるようにしています。

三宅: （キャッシュプリフェッチを含めて）プロファイラを使ったチューニングももちろんやったのですが、結局やってみないとわからない部分もあります。

SoftNeuroだとそうした試行錯誤をそれほど気にしなくても、とりあえず何種類かのアルゴリズムで実装を書いておけば、実際に「富岳」の上で走らせることができます。それで推論ネットワーク全体の中にある、この3Dコンボリューションはウェイトのサイズや入力のサイズなどの条件からこのアルゴリズムが最適だ、また別のネットワークの中の別のコンボリューションでは、別のアルゴリズム・別のパラメタが最適だ、と言ったことを自動的に選択してくます。

だから開発者としては結構スケーラブルというか、とりあえず書いておけばあとは選択してくれるというところがあります（それがうまく働きました⁠）⁠。

平賀：もともとSoftNeuroのアーキテクチャには、実機で実際にいろんなパターンのアルゴリズムを動作させて、一番速く動作したものを採用する、といった仕組みがあるのです。

安田: 選択するためには何かモニタリングの情報がいりますよね。つまりSoftNeuro自体がモニタリング情報を内部的にたくさん出力するようになっているのでしょうか。

平賀：そのとおりですね。そのためのフレームワークのアーキテクチャを arXivに論文としてあげておいたら平島さんが見つけてくれた、と。

安田: ああなるほど。SoftNeuroは元々組み込み用なのでプラットフォームによる違いが大きく、チューニングについてはやってみないとわからないことが多いから、そういった状況に対応するための仕組みが元々あったのですね。

平賀：それがたまたま「富岳」にも活きた、ということかと。

安田: 組み込み向けでもこういったプロファイラを見ながらのチューニングをされるのですか？

平賀：そういったケースもありますが、組み込みだと最近は AI DSP といったAIアクセラレータのようなものが増えてまして、、

安田: ああ、Movidiusでしたか（はい⁠）⁠、そうかアクセラレータはそういう情報が見えないんですね。

平賀：結構ブラックボックスですね。

松尾: 各レイヤーでどのくらい時間がかかったか分かるくらいですね。

安田: すると「試してみないと分からない」と言うことになりそうですね。（⁠笑い）

松尾: 計算機のアーキテクチャのトレンドとしては、そうした AI、推論のアクセラレータを共存・乗せていってほしいのですがね。GPU や DSP は AI フレームワークに特化したものが出ているし、最近 CPUの中でもコンボリューションを高速に計算するためのアセンブラ命令みたいなものが追加されてきています。CPUの汎用性にプラスアルファで何か特化する、たとえば行列演算がちょっと速いと言ったものが追加されるような構造ですね。

最後は自社製品のことからも離れて今どきのアーキテクチャの方向性の話になりましたが、現場で最先端の技術を扱うエンジニアの視点を少し共有させてもらえた感じがして、ちょっと幸せになれました。

ディープラーニングの推論で解く未来

この、「⁠数値シミュレーションの一部をディープラーニングで置き換える」というアプローチは、今回の銀河形成シミュレーションだけに限らない、一般性の高そうな話題です。取材の最後はそのあたりに議論が及びました。

安田: Morpho では過去にも数値計算でやっていたものを推論で当てにいく方法でやってきたのですか？

松尾: 今回初めてですね。トレンドとしてはあるので注目はしていたのですが。

安田: そういうものが増えているんですね。それでも学習はアプリケーションごとに個別にやるしかないように思えます。たとえば数値風洞というアプリケーションのためには、何かの数値風洞の結果をたくさん学習させることになる。

松尾: ただこれまでの手法に比べると、ディープラーニングの部分が共通しているので、その枠の中で最適化されたものができると、ディープラーニングの方が（従来的な数値計算による手法に比べて）汎用性がある（新しい応用に適用しやすい）ように思います。

つまり従来は方程式ごとに違うシミュレーションコードを書かなければなりませんでした。しかし周辺がディープラーニングで統一されると、行列積のように効率よく計算できる形をしているので、そこを共通的な基盤として最適化しておけば、いろんな応用に対しても効率的な計算ができる枠組みになっている、と言った状況です。

また、微分方程式などをディープラーニングで解こう、といったものも出始めています。単純にシミュレーションした結果を学習するのではなく、方程式ベースでディープラーニングに学習させる、というようなものです（詳細についてはコラム1を参照ください⁠）⁠。

いろいろ出始めていて、そうするとSoftNeuroも乗せやすくなるのでありがたいことかなと。

安田: いろんなところで応用が進みそうですね。

藤井: ただ、今はまだシミュレーション屋としてはその精度に確証がない状況のように思えます。

安田: 研究では精度が非常に重要というのはわかりますが、社内利用でそれほど問題にならないところもありそうに思えます。

藤井: そうした数値シミュレーションを研究でなく産業側でやっている知人がいるのですが、そこでは数値計算だと細かなパラメータ調整については経験値（知）のようなものがあるそうなのです。

ディープラーニングを使って、こういう経験値のようなものがなくてもそれなりの結果が出せるようになるのは、ちょっと良いかなと思っています。

まだまだ話したいことはあったのですが、取材としてはこのあたりでおしまいとしました。

数値シミュレーションをディープラーニングで置き換える、という手法には多様な応用の可能性があると思えます。藤井氏に伺うと、このアプローチを考えたのは3年ほど前で、その頃はまだ天文業界にAIが普及していなかったとか。使っている人はいたが、今のように天文のいろんな研究で使われるようになったのはこの2年ほどだそうです。

筆者はHPC分野にもディープラーニング分野にも詳しくないのですが、それでも従来的な（いわゆる）認識・識別系や生成系AI以外の「ディープラーニングがもたらす、もう1つの新しい応用」の入り口を見た、と感じる取材となりました。