Matlantis CSPにおけるOptunaを使った結晶構造探索

はじめに

結晶構造探索は、材料開発や物性予測の入口にあたる重要なステップです。一方で、その探索空間は広く、評価には計算コストがかかり、元素系ごとにどのような組成のどのような構造が実際に安定なのかが大きく変わります。結果として、探索の成否や効率が“経験と試行錯誤”に強く依存しがちで、同じ問題に取り組む人ほど「もっと賢く回せないか」と感じる場面が多く生じます。

我々は先日、材料探索のためのクラウド型汎用原子レベルシミュレーションプラットフォームMatlantis︎™上に結晶構造探索サービスMatlantis CSP (MTCSP)をリリースしました。 MTCSPでは、結晶構造探索を少ない試行で効率的にうまく進めるためのアプローチとして、PFNが開発を主導するブラックボックス最適化フレームワーク Optuna™が利用されています。本記事では、 MTCSPにOptunaがどのように取り入れられているかを紹介します。

ブラックボックス最適化フレームワークOptunaとは

Optunaは、「評価はできるが、どう最適化すればよいかが分かりにくい」タイプの問題に対して、少ない試行回数で良い解を見つけるためのブラックボックス最適化フレームワークです。典型例は機械学習のハイパーパラメータ探索ですが、本質的には「入力となるパラメータを与えると、出力となるスコアが返ってくる」仕組みさえあれば、対象はハイパーパラメータ探索に限りません。

結晶構造探索の文脈では、与えられた元素系や組成比の範囲等の制約に対してエネルギー的に安定な構造を探索することになりますが、入力は結晶構造で、出力は構造の形成エネルギーです。従来は時間のかかるDFT計算を用いていましたが、Matlantis上では汎用ニューラルネットワークポテンシャルPFPを用いて高速に形成エネルギーを評価できます。

結晶構造探索サービスMTCSPとは

Matlantis CSP(MTCSP)は、クラウド型汎用原子レベルシミュレーションプラットフォームMatlantis上で提供される結晶構造探索サービスです。ユーザーが与えた組成や条件に対して、候補となる結晶構造(原子配置)を生成し、エネルギーなどの指標で評価しながら、より安定(あるいは有望)な構造を探索していきます。

結晶構造探索が難しいのは、探索空間が広いことに加えて、構造の評価に一定の時間がかかり、しかも「どの探索の進め方が効くか」が元素系や条件で大きく変わるところにあります。つまり、探索は単に候補を増やせば良いわけではなく、限られた計算予算の中で、

どう候補を作るか
どう構造を評価するか
どう探索ループを構成するか

を工夫する必要があります。MTCSPは、重い探索を現実的に回すために、探索の実行や結果の分析をサービスとして扱えるようにしたもので、ユーザーは目的に応じた条件指定と各々の応用に応じた分析に集中できます。

MTCSPにおけるOptunaの利用

MTCSPでは結晶構造探索を効率的に実施し、結果を簡単に分析する機能を多数備えていますが、Optunaが深く関わっているのは探索の部分です。特に構造生成アルゴリズム(どう評価すべき結晶構造の候補を作るか)と探索ループの構成に関連が深いので、それらについて説明します。構造のエネルギーを評価する部分はMatlantis上で利用可能なPFPを用いることで高速に実現できるため、ここではスコープ外とします。MTCSPの概念図を図1に示します。MTCSPは構造生成・探索を行うフェーズと構造緩和・エネルギー評価を行うフェーズを交互に処理し、最終的に相図評価を行って新規結晶構造を出力してくれるサービスです。この中でOptunaが深く関わっているのは「構造生成・探索」になります。

図1: MTCSPの概念図

MTCSPにおけるOptunaを用いた探索アルゴリズムは、単にOptunaで実装されている探索アルゴリズムを用いているのではなく、結晶構造探索向けに専用にデザインして実装されたものです。詳細は公開しているテクニカルペーパーをご覧ください。結晶構造探索においては歴史的に確率的なサンプリングアルゴリズムやそれに基づく遺伝的アルゴリズムが用いられてきたという背景があり、我々もOptunaに実装されているNSGA-IIという遺伝的アルゴリズムをベースとした手法をデザインし実装しています。特に、Optunaは大規模な非同期並列最適化による高速な探索をサポートしており、その恩恵を受けるためには非同期並列にアルゴリズムが動作する必要があるため、MTCSPにおいてもそのように実装されています。

MTCSPにおけるOptunaを用いた探索ループの構成は、単にOptunaを用いるのではなく必要に応じて独自に機能を追加しつつOptunaの様々な機能を結びつけて実現しています。まず以下に、実現している探索ループの初期化フェーズのシーケンス図を示します。探索はExperimentと呼ばれるクラスによって制御されます。これはOptunaのStudyと呼ばれる探索を制御するクラスのラッパーとなっており、MTCSPの利用者に向けて様々な便利な機能を提供します。例えば指定された元素系に対応する単原子結晶を追加するadd_pure_atomsやMTCSP専用の遺伝的アルゴリズムの初期集団を追加するcreate_initial_populationなどです。それらは内部でOptunaのAPIを呼びながら構造の生成・永続化を行います。構造を永続化する際は、Optunaのストレージに直接保存するのではなく、別途外部に構造を保存するためのStructures Storeを構築して用いています。Structure Storeはファイルベースのストレージです。Optunaのストレージは典型的にはMySQLなどのRDBであり、そこに比較的データ量の大きい結晶構造を文字列化して直接保存することは効率的ではないため結晶構造向けにはこちらを利用しています。初期化フェーズは単一プロセス・単一スレッドによって行われます。

図2: MTCSPにおける初期化フェーズのシーケンス図

次に、以下に探索ループ本体のシーケンス図を示します。探索ループはExperimentのsearchメソッドによって起動されます。このループ全体がOptunaのStudyクラスによって制御されており、内部では上述のMTCSP専用のアルゴリズムやアルゴリズムが生成した構造を評価するRelaxer、枝刈りを行うRejecterなどと連携しながら非同期並列に効率的に結晶構造の候補を生成、評価、不要なものは枝刈りしていきます。繰り返しになりますが、この探索ループは非同期並列に実行することができ、単一ノード上のスレッド・プロセス並列から、大規模な計算インフラストラクチャ上のマルチノード並列実行までを同一のフレームワークでサポートしています。

図3: MTCSPにおける探索ループのシーケンス図

このようにMTCSPはOptunaと密接に関わっており切っても切り離せない関係にあります。このような構成にしている理由はいくつかあり、

PFNの保有する大規模な計算クラスタとOptunaの相性が良いこと。Optunaの非同期処理のサポートによって数十万構造の評価を容易に行うことができます。
いくつか設定を書き換えるだけで計算クラスタや別のプラットフォーム(すなわちMatlantis)の両方で動かすことができること
OSSで持続的に開発されているOptunaの発展の恩恵をMTCSP側でも受けられること。実際データベースの高速化等の恩恵を受けてMTCSPのワークロードが継続的に効率化されています。

などがあります。

おわりに

本記事では、Matlantis CSP (MTCSP) において結晶構造探索を「少ない試行で賢く回す」ために、ブラックボックス最適化フレームワーク Optuna をどのように取り入れているかを紹介しました。結晶構造探索は探索空間の広さと評価コスト、そして元素系ごとに最適な探索戦略が変わるという難しさを抱えていますが、MTCSPでは結晶構造探索向けに専用設計したアルゴリズムと、Optunaが得意とする非同期並列の探索ループ制御を組み合わせることで、現実的な計算予算の中でも効率よく探索を進められるようにしています。

我々は今後もOptunaとMTCSP両方の開発を継続し、両者のシナジーを十分に高めることでサービスレベルの向上に努めていきます。材料探索の試行錯誤を少しでも前に進める道具として、MTCSPが皆さまの研究開発の現場で役に立てば幸いです。

投稿 Matlantis CSPにおけるOptunaを使った結晶構造探索は Preferred Networks Tech Blog に最初に表示されました。