[论文解读] Transfer learning for atomistic simulations using GNNs and kernel mean embeddings

Isak Falk, Leo H. Bonati+3 more|arXiv (Cornell University)|Jun 2, 2023

Machine Learning in Materials Science被引用 5

一句话总结

该论文提出均值嵌入核岭回归（MEKRR），一种迁移学习框架，结合在OC20数据集上预训练的图神经网络（GNN）表征与核均值嵌入，用于建模原子体系中的势能面。通过将化学种类信息融入新型核函数，MEKRR在催化体系中实现了最先进性能，显著优于独立的GNN或岭回归方法，在分布外数据上的测试误差最高降低了70%。

ABSTRACT

Interatomic potentials learned using machine learning methods have been successfully applied to atomistic simulations. However, accurate models require large training datasets, while generating reference calculations is computationally demanding. To bypass this difficulty, we propose a transfer learning algorithm that leverages the ability of graph neural networks (GNNs) to represent chemical environments together with kernel mean embeddings. We extract a feature map from GNNs pre-trained on the OC20 dataset and use it to learn the potential energy surface from system-specific datasets of catalytic processes. Our method is further enhanced by incorporating into the kernel the chemical species information, resulting in improved performance and interpretability. We test our approach on a series of realistic datasets of increasing complexity, showing excellent generalization and transferability performance, and improving on methods that rely on GNNs or ridge regression alone, as well as similar fine-tuning approaches.

研究动机与目标

• 为解决在原子模拟中训练精确的原子间势能函数所面临的高数据与计算成本挑战。
• 利用预训练的GNN表征，实现复杂催化体系下的数据高效且快速的迁移学习。
• 通过结合GNN特征与蕴含物理信息的核函数（编码化学种类），提升泛化能力与迁移能力。
• 开发一种在保持物理对称性（旋转平移对称性与排列对称性）的同时，实现可解释性与性能提升的方法。
• 在预训练分布之外、日益复杂的分布外数据集上，验证该方法的有效性。

提出的方法

• 在包含约2.6亿个催化体系DFT弛豫构型的大规模OC20数据集上预训练GNN（如SchNet或SCN）。
• 提取固定、预训练的GNN特征作为原子环境的表征，保持旋转平移不变性。
• 基于核均值嵌入定义核函数，强制实现排列对称性，并通过可学习权重引入化学种类信息。
• 使用带有所提核函数的核岭回归（KRR）从少量、体系特异的数据中回归势能面。
• 通过交叉验证优化核权重，以突出显示在表面催化中代表性不足的物种，特别是吸附原子。
• 利用核矩阵进行谱聚类，检测如键断裂等反应事件，展示其超越回归任务的可解释性。

实验结果

研究问题

• RQ1 预训练的GNN表征与核均值嵌入相结合，能否在原子体系的势能面建模中实现有效的迁移学习？
• RQ2 将化学种类信息融入核函数，对模型性能与可解释性有何影响？
• RQ3 所提方法在预训练分布之外的数据集上，其泛化能力达到何种程度？
• RQ4 基于核的相似性度量能否在无显式标注的情况下检测如键断裂等物理事件？
• RQ5 在数据效率与精度方面，MEKRR相较于微调的GNN或独立的核方法表现如何？

主要发现

• 使用SchNet特征的MEKRR在D1→D2迁移任务中，相比监督微调的SchNet，测试误差降低了68%，达到8.0 meV/atom。
• 在最复杂的数据集（D4）上，使用SCN特征的MEKRR实现55 meV/atom的测试误差，优于微调的SchNet-FT（116 meV/atom）。
• 该方法在分布外数据上表现出强大的泛化能力，包括预训练集中未出现的N-N键断裂构型。
• 多权重核（α=1）与N-N键长等物理可观测量的相关性更高，表明对关键反应事件的敏感性增强。
• 利用核矩阵进行谱聚类成功识别出键断裂轨迹中的反应物与产物状态，验证了核函数的物理可解释性。
• 引入化学种类感知的核权重显著提升了催化体系的性能，尤其在吸附原子等代表性不足物种上。