[论文解读] 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

一句话总结

该论文提出了一种替代的黑箱矩阵-矩阵乘法（AltBBMM）算法，以加速基于分子轨道的机器学习（MOB-ML）中的精确高斯过程（GP）训练，使超过6,500个分子（超过一百万个成对能量）的可扩展训练成为可能，相较于原始BBMM算法实现四倍速度提升，同时保持了最先进水平的准确性和可迁移性。该方法在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推理和高预测准确度的前提下，将MOB-ML扩展至大规模数据集——特别是超过一百万个成对能量。
• 开发一种替代的BBMM实现（AltBBMM），以提升训练效率，同时不牺牲模型的准确度或可迁移性。
• 证明AltBBMM可在包括更大分子在内的多种化学体系中实现分子能量预测的最先进性能。

提出的方法

• 论文采用黑箱矩阵-矩阵乘法（BBMM）算法，将GP训练复杂度从O(N³)降低至O(N²)，实现多GPU加速和精确推理。
• 提出一种替代的BBMM实现（AltBBMM），通过重构矩阵运算以改善内存访问模式并减少计算开销，实现更快的收敛速度。
• 使用共轭梯度（CG）和块共轭梯度（BCG）方法求解GP推理中的线性系统，并通过预处理加速收敛。
• 在MOB-ML中将对角和非对角成对能量的训练分离，为每类使用独立模型，并通过学习到的超参数进行基于核的GP回归，最大化对数边际似然。
• 通过添加白噪声（σ²_add）进行正则化以减少矩阵病态性，选择σ²_add = 10⁻⁵以在准确度和训练速度之间取得平衡。
• 在QM7b-T和GDB-13-T数据集上评估该框架，模型性能通过平均绝对误差（MAE）和每7个重原子的MAE（MAE/7HA）进行衡量。

实验结果

研究问题

• RQ1 BBMM算法能否有效应用于将MOB-ML中的精确GP训练扩展至此前约220个分子的限制之外？
• RQ2 所提出的AltBBMM实现是否在保持相同预测准确度和可迁移性水平的前提下，相较于原始BBMM实现实现了显著的速度提升？
• RQ3 在QM7b-T和GDB-13-T等基准数据集上，MOB-ML结合AltBBMM是否能实现最先进水平的准确度，尤其是在大规模数据场景下？
• RQ4 在准确度和跨分子尺寸的泛化能力方面，AltBBMM与现有最先进方法（如MOB-ML/Full GP和DeePHF/扩展神经网络）相比表现如何？

主要发现

• 在6,500个QM7b-T分子上训练时，AltBBMM将训练时间缩短为6.24小时（相比原始BBMM的26.52小时），实现四倍速度提升，且在GDB-13-T上的测试MAE仅增加1%。
• 使用AltBBMM训练的模型在QM7b-T上的测试MAE为0.193 kcal/mol，在GDB-13-T上的MAE/7HA为0.493 kcal/mol，与或优于以往最先进方法的性能。
• BBMM和AltBBMM将MOB-ML的训练规模扩展至超过一百万个成对能量（1,152,157个非对角元素和124,973个对角元素），是此前约220个分子限制的30倍。
• 使用σ²_add = 10⁻⁵可确保最优准确度，因为更高的值（如10⁻¹）会使预测MAE翻倍，尽管训练时间减半。
• AltBBMM保持了高可迁移性，在GDB-13-T上实现了迄今为止报告的最佳MAE/7HA（0.493 kcal/mol），优于DeePHF/扩展神经网络方法。
• 图1中的学习曲线表明，无论在低数据量还是大规模数据场景下，AltBBMM均使MOB-ML更快达到化学精度（1 kcal/mol），且泛化能力更优。