Learning the Kohn-Sham map with neural operators for quasi-linear scaling density functional theory

60年以上にわたり、科学者たちは原子がどのように結合して周囲の物質を形成するかを理解するために、強力な一連の規則に頼ってきました。密度汎関数理論として知られるこの枠組みは、原子を取り囲む電子雲を、個々の粒子の集まりとしてではなく、滑らかで連続的な密度として扱います。このアプローチにより、研究者は新しい医薬品から高度な電池に至るまで、あらゆるものの特性を予測することができます。しかし、そこには落とし穴があります。正確な結果を得るためには、すべての電子の挙動を解くという、複雑で反復的な計算が必要になります。システムが大きくなるにつれて、これらの計算に必要な時間は爆発的に増加し、その速度は非常に速いため、最も強力なスーパーコンピュータを用いたとしても、中規模の金属片や大きな生物学的分子をシミュレートすることは不可能になります。このボトルネックが多くの重要な科学的問いを手の届かないものにしており、研究者は精度とスケールの間で選択を迫られてきました。

カリフォルニア工科大学の研究チームは、精度を損なうことなく、計算の中で最もコストのかかる部分をコンピュータにスキップさせる方法を見つけ出し、この障壁を打破しました。最終的な答えを直接予測しようとしたり、失敗しやすい難しい数学的なショートカットを学習させたりする代わりに、彼らは人工知能に、毎回の計算サイクルで発生する特定の反復ステップを実行するように教えました。この中間ステップに焦点を当てることで、彼らは、以前は数千のプロセッサを必要としていたタスクを、単一のグラフィックスカードで数万個の原子を持つシステムとしてシミュレートできる手法を作り上げました。その結果、最も単純な近似法に近い速さを持ちながら、最も厳密な手法と同等の高い精度を維持する、物質の新しいモデリング手法が誕生しました。

問題の核心は、この理論が電子をどのように扱うかにあります。材料の安定した状態を見つけるために、コンピュータは電子密度のマップが変化しなくなるまで、繰り返し調整を行う必要があります。従来の方法では、この調整には巨大な固有値問題を解く必要があり、これは本質的に、システム内のすべての原子に対する電子の動きの特定のパターンを見つけ出す作業です。この操作は非常に計算負荷が高く、そのコストは電子の数に対して立方体に比例して増加します。つまり、システムのサイズを2倍にすると、計算の難易度は8倍になります。数十年にわたり、科学者たちは個々の電子のパターンを追跡する必要性を排除することで、この回避を試みてきました。これは軌道フリー密度汎関数理論として知られる戦略です。しかし、機械学習を用いてこのショートカットを学習させようとするこれまでの試みは、苦戦してきました。電子密度とエネルギーの関係を学習させようとしたものもありますが、この関係は数学的に不安定でエラーが発生しやすいものでした。また、原子の配置から最終的な電子密度を直接予測しようとしたモデルもありましたが、それらは、学習時よりも大きなシステムや化学的に異なるシステムを予測しようとすると失敗しました。

研究者たちは、最終的な目的地を予測しようとするのではなく、その「過程（ジャーニー）」を学習することこそが鍵であると気づきました。従来の計算の各ステップにおいて、コンピュータは電子に作用する力のマップを取り込み、電子がどこに見出される可能性が高いかを示す新しいマップを生成します。この前進ステップは、数学的に安定しており、明確に定義されています。チームは、空間的なパターンを理解するように設計された特殊なニューラルネットワーク（一種の人工知能）を用いて、まさにこの変換を学習させました。彼らは、有機化合物、絶縁体、金属を含む、8,504種類の分子と固体材料から抽出された、これら力のマップと、それらが生成する対応する電子密度マップの数千の例をネットワークに投入しました。ネットワークは、力から新しい電子密度を直接予測することを学習し、これにより、低速で重い計算を、高速で学習された予測へと効果的に置き換えたのです。

このアプローチをユニークなものにしているのは、異なるシステムの複雑さをどのように扱うかという点です。研究者たちは、学習したルールを、再学習させることなく、小さな分子にも巨大な結晶にも適用できるように、ネットワークをドメイン不変（領域不変）となるよう設計しました。また、ネットワークが空間の基本的な対称性を尊重するように設計されており、分子を回転させたり移動させたりしても、その物理的特性が変わらないことを理解しています。未知のシステムに対してテストを行った際、このモデルは驚くほど堅牢であることを証明しました。学習データに含まれていない元素やサイズを持つ薬物様分子のセットにおいて、新しい手法は高い精度を維持しましたが、従来の直接予測モデルは、分子が大きくなるにつれてエラーが急増しました。この自己整合的な性質により、ナビゲーターがルート全体を推測するのではなく、頻繁に現在地を確認するように、各ステップで自らの間違いを修正することが可能になります。

この手法の威力は、結晶のような周期系や、金属における拡張欠陥に適用した際に明らかになりました。結晶の標準的な計算では、コンピュータは結晶の繰り返しのパターンを記述する数学的空間内の多くの異なる点に対して、重い計算を繰り返さなければなりません。新しい手法は、これらの点を個別にサンプリングすることなく、単位胞上の電子密度を進化させることで、この反復を完全にスキップします。これにより、研究者は単一のモデルで金属系や半導体の計算を収束させることができ、従来の計算手法と同じ精度で電子スペクトルと構造特性を再現することができました。さらに、このモデルは、異なる理論的近似間の転移についても驚くべき能力を示しました。あるタイプの計算で学習されたモデルは、入力される力を入れ替えるだけで、再学習なしに、より正確な近似を用いた計算に使用することができました。

究極のテストは、研究者が金属結晶における欠陥の一種であり、材料の変形や破壊を理解する上で極めて重要な、マグネシウムの転位にこの手法を適用した際に訪れました。これらの欠陥は長距離の弾性場を生み出すため、正確にモデル化するには数千の原子のシミュレーションが必要です。このようなシステムを高精度でシミュレートしようとする従来の試みでは、数千個のグラフィックスプロセッシングユニットを備えたスーパーコンピュータが必要であり、実行には数日を要しました。しかし、最新のグラフィックスカード1枚で動作するこの新手法は、8,250個のマグネシウム原子と82,500個の価電子を含むシステムの計算を正常に収束させました。計算にかかる時間はシステムのサイズに対してほぼ線形に増加し、計算のボトルネックが取り除かれたことが確認されました。得られた電子密度は高精度なリファレンスと区別がつかず、金属の強度を決定する異なる種類の欠陥間の微妙なエネルギー差を解明することを可能にしました。

この研究は、機械学習が物理理論全体を置き換えるのではなく、その中で最も高価で繰り返される操作を加速させることによって、電子構造計算の範囲を広げられることを示しています。力から密度へのフォワードマップを学習することで、研究者たちは、安定しており、異なる種類の物質間で転移可能で、かつ以前はアクセス不可能であったスケールを扱うことができるツールを作り上げました。この手法は、計算が収束しない場合には、そのシステムがモデルの信頼範囲外であることを示す直接的な診断機能を提供しており、これは直接予測モデルには欠けている特徴です。現在のモデルは電子密度を予測するものですが、全エネルギーやその他の特性は、依然として単一の最終ステップを経て得ることができ、完全な理論の精度を保持しています。このアプローチは、かつては理論的に不可能であった詳細さとスケールで、複雑な材料、生物学的システム、および化学反応をシミュレートするための扉を開いています。