SALTED: a symmetry-adapted machine-learning program for predicting electron-densities in molecules and materials

技術要約：SALTED：電子密度予測のための対称適応型機械学習プログラム

問題提起
密度汎関数理論（DFT）は、精度と計算コストのバランスに優れた、材料特性をシミュレーションするための標準的な手法である。しかし、データ駆動型の材料科学がより大規模な系やより長い時間スケールへと向かうにつれ、DFTの控えめなコストでさえも、実行不可能なものとなりつつある。そこで、DFTにおける最も基本的な量である実空間の電子密度 $n(\mathbf{r})$ を、高い精度を維持しつつ、極めて低い計算コストで予測できる機械学習（ML）手法への切実なニーズが存在する。既存のMLアプローチの多くは、多様な化学環境に対する転移性に課題があったり、膨大な量の訓練データを必要としたりするが、高レベルの電子状態計算手法からこれだけのデータを生成することは困難である。

手法
SALTED（Symmetry-Adapted Learning of the Electron Density）は、特定の表現形式と学習アルゴリズムを通じてこれらの課題に対処する：

• 原子基底表現： 3Dグリッド上で $n(\mathbf{r})$ をサンプリングするグリッドベースのMLモデルとは異なり、SALTEDは線形な原子中心展開を採用している。電子密度は、各原子を中心とする径方向関数 $R^\lambda_n$ と球面調和関数 $Y^\lambda_\mu$ の基底を用いて展開される：
  $$n(\mathbf{r}) = \sum_{i n \lambda \mu} c_{i n \lambda \mu} \sum_{\mathbf{u}} R^\lambda_n(|\mathbf{r} - \mathbf{r}_i - \mathbf{u}|) Y^\lambda_\mu\left(\widehat{\mathbf{r} - \mathbf{r}_i - \mathbf{u}}\right)$$
  ここで、$i$ は原子指数を表し、$\mathbf{u}$ は周期系におけるセル並進ベクトルを表す。係数 $c_{i n \lambda \mu}$ がターゲット変数となる。この表現は、原子軌道ベースの電子状態計算コードと自然に互換性があり、分解のアイデンティティ（RI）法を介して参照係数を直接抽出することが可能である。

• 対称適応型ガウス過程回帰（GPR）： 学習アルゴリズムには、対称適応型の拡張GPRを利用している。等変ニューラルネットワークの代わりに、SALTEDは正確な $O(3)$ 対称性を満たすために球面テンソルカーネル（$K_\lambda$）を計算する。これらのカーネルは、featomic ライブラリによって生成された原子密度特徴量から導出される。

  • データの効率性： これらのカーネルによって定義される制約付き関数空間と、特徴量の物理的な解釈可能性により、過学習を防ぐことができる。これにより、GPRは小規模データ領域（例：$\sim 10^2-10^3$ 回の計算）において特に効果的であり、これは高レベルの電子状態データを用いて訓練を行う際に極めて重要である。
  • ベクトル場カーネル： ベクトル場カーネル関数を実装することで、SALTEDは印加された電場に対する電子密度の一次応答（$\partial n(\mathbf{r})/\partial \mathbf{E}$）を学習でき、3Dベクトル場の正しい変換特性を保証する。
  • 長距離相互作用： featomic の使用により、長距離等変（LODE）特徴量の導入が可能となり、これは金属フレームワークのような高度に分極しやすい系における非局所的な密度再分布を捉えるために不可欠である。

• ソフトウェア・アーキテクチャ： SALTEDは、データセット準備、モデル訓練、予測の3段階のワークフローで構成されたオープンソースのPythonパッケージである。本プログラムは、訓練データを統一するために主要な電子状態計算コード（CP2K、FHI-aims、PySCF）と連携する。ワークフローはハイパフォーマンス・コンピューティング（HPC）向けに設計されており、チェックポインティング、記述子縮退およびヘッセ行列構築のためのMPI並列化、およびパフォーマンス向上のためのNumbaコンパイル済みカーネルを備えている。

主な貢献

1. オープンソース・フレームワーク： SALTEDは、第一原理データから電子密度およびその線形応答のMLモデルを構築するための、実用的かつ拡張可能なフレームワークを提供する。
2. 相互運用性： 既存の電子状態計算パッケージ（CP2K、FHI-aims、PySCF）とシームレスに統合されており、研究者がコードを切り替えることなく、既存のワークフローや設定を利用することを可能にする。
3. 応答特性： 本ソフトウェアは電場に対する密度応答の予測を独自にサポートしており、分極率テンソルや分極ベクトルの計算を可能にする。
4. 小規模データ領域における効率性： 対称適応型カーネルを用いたGPRを活用することで、データ要求量の多いニューラルネットワーク構造と比較して、大幅に少ない訓練点数で高い精度を実現する。

結果と応用
本論文では、SALTEDが計算ワークフローに統合されたいくつかの成功事例を紹介している：

• QM/MMシミュレーション： 金属電極における電荷密度応答を予測することで、モデルイオンキャパシタのDFTレベルの静電力を計算する際の速度を約$10^3$倍向上させた。
• 大規模電子構造： 予測された密度を単一のコーン・シャモンのハミルトニアン対角化への入力として用いることで、極めて大規模なスーパーセルにおけるツイスト二層モアレ材料の電子特性の研究を可能にした。
• 物性計算： 本手法は、クーロン力、分極ベクトル、および分極率テンソルの正確な評価を容易にする。
• 結晶学： 予測されたRI密度係数は、X線回折や電子回折データからの結晶構造精緻化に直接適用可能である。

意義
SALTEDは、物理的に意味のある観測量へのアクセスを維持しながら、電子状態計算を加速させるための重要な一歩である。電子密度を直接予測することで、DFTのわずかなコストで多数の下流の材料特性を計算する経路を提供する。本ソフトウェアは再現可能な研究をサポートするように設計されており、機械学習とマルチスケールシミュレーションツールの架け橋として、効率的な予測を求めるユーザーと、手法の開発を進める開発者の両方に資するものである。分子系と凝縮相系の両方を扱う能力と、標準的な電子状態計算コードとの互換性を備えた本手法は、計算化学および材料科学コミュニティにとって多用途なツールである。