[論文レビュー] Molecular Energy Learning Using Alternative Blackbox Matrix-Matrix Multiplication Algorithm for Exact Gaussian Process

Jiace Sun, Lixue Cheng+1 more|arXiv (Cornell University)|Sep 20, 2021

Machine Learning in Materials Science参考文献 25被引用数 5

ひとこと要約

本稿は、分子軌道に基づく機械学習（MOB-ML）における正確なガウス過程（GP）学習を高速化するための代替ブラックボックス行列積算（AltBBMM）アルゴリズムを導入する。これにより、6,500分子（100万以上のペアエネルギーを含む）のスケーラブルな学習が可能となり、元のBBMMと比較して4倍の高速化が達成され、最先端の精度と一般化性能を維持する。本手法はQM7b-TおよびGDB-13-Tデータセットにおいて、報告された最高性能を達成し、分子エネルギー予測分野における新たなベンチマークを設定する。

ABSTRACT

We present an application of the blackbox matrix-matrix multiplication (BBMM) algorithm to scale up the Gaussian Process (GP) training of molecular energies in the molecular-orbital based machine learning (MOB-ML) framework. An alternative implementation of BBMM (AltBBMM) is also proposed to train more efficiently (over four-fold speedup) with the same accuracy and transferability as the original BBMM implementation. The training of MOB-ML was limited to 220 molecules, and BBMM and AltBBMM scale the training of MOB-ML up by over 30 times to 6500 molecules (more than a million pair energies). The accuracy and transferability of both algorithms are examined on the benchmark datasets of organic molecules with 7 and 13 heavy atoms. These lower-scaling implementations of the GP preserve the state-of-the-art learning efficiency in the low-data regime while extending it to the large-data regime with better accuracy than other available machine learning works on molecular energies.

研究の動機と目的

• MOB-MLにおける正確なガウス過程回帰のO(N³)計算ボトルネックを克服すること。これは、従来220分子程度までの学習に制限されていた。
• 正確なGP推論と高い予測精度を維持したまま、大規模データセット（特に100万以上のペアエネルギー）へのスケーリングを実現すること。
• モデルの精度や一般化性能を損なわず、学習効率を向上させる代替BBMM実装（AltBBMM）を開発すること。
• AltBBMMが、より大きな分子を含む多様な化学系において、分子エネルギー予測で最先端の性能を達成できることを示すこと。

提案手法

• 本稿では、ブラックボックス行列積算（BBMM）アルゴリズムを用いて、GP学習の複雑度をO(N³)からO(N²)に低減し、マルチGPU加速と正確な推論を可能にする。
• 代替BBMM実装（AltBBMM）を提案し、行列演算の再構造によりメモリアクセスパターンを最適化し、計算オーバーヘッドを低減することで、より速い収束を実現する。
• GP推論における線形方程式系の解法に共役勾配法（CG）およびブロックCG（BCG）を用い、収束を加速するための前処理を施す。
• MOB-MLにおける対角および非対角ペアエネルギーの学習を分離し、それぞれ別々のモデルを用いる。また、学習されたハイパーパrameterを用いたカーネルベースのGP回帰を実施する。
• 行列の特異性を軽減するため、追加の白色ノイズ（σ²_add）による正則化を適用し、精度と学習速度のバランスを最適化するためにσ²_add = 10⁻⁵を選択する。
• フレームワークはQM7b-TおよびGDB-13-Tデータセットで評価され、モデル性能は平均絶対誤差（MAE）および1分子あたり7個の重原子あたりのMAE（MAE/7HA）で測定される。

実験結果

リサーチクエスチョン

• RQ1 BBMMアルゴリズムは、従来の約220分子の制限を超えて、MOB-MLにおける正確なGP学習をスケーリングするために効果的に適用可能か？
• RQ2 提案されたAltBBMM実装は、元のBBMMと同等の予測精度および一般化性能を維持しつつ、顕著な高速化を達成できるか？
• RQ3 AltBBMMを用いたMOB-MLは、特に大規模データ環境下において、QM7b-TやGDB-13-Tといったベンチマークデータセットで最先端の精度を達成できるか？
• RQ4 分子サイズの広がりにわたる一般化性能や精度の観点から、AltBBMMはMOB-ML/Full GP やDeePHF/拡張ニューラルネットワークといった既存の最先端手法と比較して優れているか？

主な発見

• 6,500個のQM7b-T分子で学習する場合、AltBBMMは学習時間を6.24時間（元のBBMMの26.52時間）に4倍高速化し、GDB-13-TでのテストMAEはわずか1%増加にとどまる。
• AltBBMMで学習されたモデルは、QM7b-Tで0.193 kcal/molのテストMAE、GDB-13-Tで0.493 kcal/mol MAE/7HAを達成し、先行する最先端手法と同等またはそれを上回る性能を示す。
• BBMMおよびAltBBMMにより、MOB-MLの学習は115万2,157個の非対角ペアエネルギーと12万4,973個の対角ペアエネルギー（合計127万7,130ペア）にまでスケーリングされ、従来の約220分子の制限から30倍の増加を達成する。
• σ²_add = 10⁻⁵を用いることで最適な精度が確保され、より高い値（例：10⁻¹）を用いると予測MAEが2倍に増加する一方で、学習時間が半減する。
• AltBBMMは高い一般化性能を維持し、GDB-13-Tで報告された最高のMAE/7HA（0.493 kcal/mol）を達成し、DeePHF/拡張ニューラルネットワーク手法を上回る。
• 図1の学習曲線から、AltBBMMは低データおよび大規模データの両環境下で、化学的精度（1 kcal/mol）に到達するまでの速度と一般化性能が、従来手法を上回ることを示している。