复旦大学向红军教授团队JMI最新研究论文 | 使用Kolmogorov-Arnold网络高效预测势能面和材料物理性质

题目：Efficient prediction of potential energy surface and physical properties with Kolmogorov-Arnold Networks

作者：Rui Wang, Hongyu Yu, Yang Zhong, Hongjun Xiang*

DOI: 10.20517/jmi.2024.46

Citation: Wang, R.; Yu, H.; Zhong, Y.; Xiang, H. Efficient prediction of potential energy surface and physical properties with Kolmogorov-Arnold Networks. J. Mater. Inf. 2024, 4, 32. http://dx.doi.org/10.20517/jmi.2024.46

图片摘要

首先，我们使用机器学习势能模型 Allegro评估了在机器学习势能模型的不同部分中，用采用不同基函数的 KAN 网络替换MLP所带来的影响。

我们尝试用不同基函数的 KAN 替换 Allegro 模型中的所有 MLP，如图1所示。我们测试了基于B 样条基函数、高斯函数和傅里叶函数的 KAN。我们使用 Ag 数据集评估了各种模型的准确性和效率，数据集包含 1,000 个不同的结构，其中 950 个用于训练，50 个用于验证。

图1. 用采用不同基函数的 KAN 替换机器学习势能模型 Allegro不同部分中的 MLP。

预测势能的平均绝对误差（MAE）和训练时间如表 1 和图 2所示。值得注意的是，所有采用 KAN 的 Allegro 模型训练得到的力和能量的 MAE均低于原始基于 MLP 的 Allegro 模型。具体而言，采用高斯基函数的 KAN 模型的力 MAE 为 0.014 eV/Å，比基于 MLP 的 Allegro 模型低 12.5%。采用 B 样条基函数的 KAN 模型在验证数据集上实现了最低的能量 MAE，为 0.029 eV/atom，比基于 MLP 的模型低 17.1%，但该模型的训练时间几乎是原模型的五倍。相比之下，采用高斯基函数的 KAN 模型的验证能量 MAE 也低于基于 MLP 的模型，为 0.032 eV/atom，并且训练时间与基于 MLP 的模型相当。采用傅里叶基函数的模型的能量 MAE 上与基于 MLP 的 Allegro 模型相近，但训练时间更长。采用 KAN 的三个 Allegro 模型在预测准确性方面均优于采用 MLP 的 Allegro 模型。这种性能改进可能归因于KAN在拟合能力上的优势，它们在解决复杂问题（如预测材料的势能面和物理性质）时具有显著的优势。

使用 B 样条基函数的 Allegro 模型表现出最高的预测精度，这可能归因于 B 样条作为分段多项式函数的灵活性，非常适合逼近复杂函数。使用高斯基函数的模型也表现出相当的精度。相比之下，傅里叶基函数虽然在捕捉数据中的周期性或振荡模式方面非常出色，但在预测势能面时可能不如前两种基函数实用。然而，使用 B 样条基函数的 Allegro 模型相比于其他基函数的模型需要显著更长的训练时间。这主要是由于计算 B 样条基函数和重新缩放网格所需的计算时间较多。采用更高效的基函数（如高斯和傅里叶函数）可以在保证相当精度的情况下显著加速模型计算。其中，基于高斯的 KAN 在精度和训练效率之间实现了最佳平衡，以显著更短的训练时间达到了与基于 B 样条的 KAN 相近的预测性能。在训练其他机器学习方法时，基函数的选择应依据应用的具体需求，如优先考虑精度还是计算效率。

表1.  采用 MLP 和使用 B 样条、高斯及傅里叶基函数的 KAN 的 Allegro 模型结果对比。最佳结果以粗体表示。

MAE: 平均绝对误差; F: 力; E: 能量。

图2.  将 Allegro 模型中的 MLP 替换为使用不同基函数的 KAN 的平均绝对误差（MAE）。采用 KAN 的三个 Allegro 模型在力和能量的 MAE 上均低于原始基于 MLP 的 Allegro 模型。

随后，我们用 KAN 替换 Allegro 模型不同部分中的 MLP，以确定最佳配置。我们选择了具有高斯基函数的 KAN，该配置在保证较高预测精度的同时保持了相对较短的训练时间。具体而言，我们评估了两种配置：在两体隐藏嵌入层和隐藏层 MLP 部分使用KAN，以及仅在输出块中使用 KAN，如图1所示。我们首先使用 Ag 数据集评估了各种模型的性能。我们还通过使用LAMMPS进行分子动力学模拟，评估了各种模型的推理速度和 GPU 内存使用情况。为了评估数据集复杂性对 KAN 和 MLP 相对性能的影响，我们进一步在更复杂的 HfO₂ 结构上评估了这些模型。HfO₂结构展现了复杂的原子间相互作用，包括由于 p-d 杂化引起的混合离子-共价特性，这使得开发准确的机器学习势模型变得具有挑战性。HfO₂ 数据集包括10,000 个结构，其中 9,000 个用于训练，1,000 个用于验证。在Ag数据集上的结果如表 2 和图 3所示。值得注意的是，在输出块中使用KAN的 Allegro 模型实现了最高的预测精度，其在验证数据集上的能量 MAE 为 0.022 eV/atom，比使用 MLP 的 Allegro 模型低 37.1%。而在两体隐藏嵌入层和隐藏层 MLP 部分使用KAN的 Allegro 模型在预测精度上也略有提升，并且相比使用 MLP 的模型训练时间有所缩短。

表 2.  在 Ag 数据集上，用高斯基函数的 KAN 替换 Allegro 不同部分的 MLP 的结果。最佳结果以粗体表示。

MAE: 平均绝对误差; F: 力; E: 能量。

图3.  在 Ag 数据集上，用 KAN 替换 Allegro 模型各部分中 MLP 的平均绝对误差（MAE）。采用 KAN 的三个 Allegro 模型在力和能量的 MAE 上均低于原始基于 MLP 的 Allegro 模型。值得注意的是，在输出块中使用KAN的 Allegro 模型实现了最高的预测精度。

我们还通过分子动力学模拟评估了各种模型的推理速度和 GPU 内存使用情况。各种模型的推理速度和 GPU 内存使用情况如表 3 所示。总体而言，采用不同基函数的 KAN 的 Allegro 模型相比使用 MLP 的模型，GPU 内存使用量略高。这表明，采用 MLP 的 Allegro 模型在模型设计和数据处理方面更高效，从而在计算资源效率上具有优势。而在 Allegro 模型中仅替换部分 MLP 为 KAN 可降低 GPU 内存使用量。具体而言，在输出块中使用KAN的 Allegro 模型需要 1945 MB 的 GPU 内存，仅比使用 MLP 的 Allegro 模型的 1941 MB 多 4 MB。采用 KAN 的 Allegro 模型的推理速度仅略慢于使用 MLP 的模型。在输出块中使用KAN的 Allegro 模型的推理速度为8.92 ms/time step，仅比使用 MLP 的 Allegro 模型慢 0.70 ms/ time step。仅在输出块中使用 KAN 相较于使用 MLP 提高了预测精度，同时相比在整个 Allegro 模型中使用 KAN，提高了推理速度和计算资源效率。

表3. 用高斯基函数的 KAN 替换 Allegro 不同部分的 MLP 的推理速度和 GPU 内存使用情况。

在 Ag 数据集上的结果显示，使用KAN的精度提升较为有限，这可能是因为该数据集过于简单，不足以显著发挥 KAN 的优势。因此，我们进一步在更复杂的 HfO₂ 结构上评估了这些模型。表 4 和图4展示的结果表明，将 Allegro 模型输出块中的 MLP 替换为 KAN，可显著提高能量和力的预测精度。验证集的力 MAE 降至 0.054 eV/Å，比使用 MLP 的 Allegro 模型降低了 27.0%；验证集的能量 MAE 降至 0.104 eV/atom，降低了 36.6%。此外，训练时间也显著缩短。对于完全使用 KAN 的 Allegro 模型，训练集的力 MAE 为 0.058 eV/Å，而能量 MAE 为 1.444 eV/atom。这种差异是因为模型收敛速度较慢，导致训练过程结束时尚未完全收敛。相比之下，仅在输出块中使用 KAN 的 Allegro 模型有效结合了 KAN 和 MLP 的优势。此混合配置利用了 KAN 的拟合精度，同时保留了 MLP 在模型其他部分的效率。因此，在输出块中使用 KAN 可加速训练收敛，并显著提高力和能量的预测精度。此外，在训练过程中，仅在输出块中使用 KAN 的 Allegro 模型的 GPU 内存分配为 45.63%，仅比使用 MLP 高 0.03%。而在 Allegro 模型的其他部分替换 MLP 对预测精度或训练时间的影响很小。

表4.  在HfO₂数据集上，用高斯基函数的 KAN 替换 Allegro 不同部分的 MLP 的结果。最佳结果以粗体表示。

MAE: 平均绝对误差; F: 力; E: 能量。

图4.  在 HfO₂ 数据集上，用 KAN 替换 Allegro 模型各部分中 MLP 的平均绝对误差（MAE）。将 Allegro 模型输出块中的 MLP 替换为 KAN，可显著提高能量和力的预测精度。

这些发现与 Ag 数据集上的结果总体一致。在 HfO₂ 数据集上，预测精度和训练时间的改进比在 Ag 数据集上更为显著。这种差异源于数据集复杂性对 KAN 和 MLP 相对性能的影响。在像 Ag 数据集这样较为简单的数据集中，KAN 和 MLP 的性能差异较小。然而，随着数据集复杂性的增加，KAN 相较于 MLP有更大的优势，因为它能够更好地表示数据中的复杂关系和依赖性。因此，在处理高复杂性的数据集时，在机器学习模型中引入 KAN 可能具有特别将 Allegro 模型输出块中的 MLP 替换为 KAN 显著提高了预测精度。在某些情况下，这种替换还缩短了训练时间。这种改进是因为 KAN 在函数拟合方面更有效。然而，像样条这样的基函数在利用组合结构方面能力较弱，因此在特征学习中不如 MLP。因此，用于从最终层输出中预测能量的输出块非常适合采用 KAN，以提高预测精度。而在 Allegro 模型的其他部分（如嵌入层）使用 KAN，相较于输出块，对预测精度的提升较小。

我们还研究了在 NequIP （Neural Equivariant Interatomic Potentials）模型中将 MLP 替换为 KAN 的效果。我们仅在输出块中将 MLP 替换为 KAN，如图5所示。我们测试了采用高斯和 B 样条基函数的 KAN。我们使用Ag 数据集对 NequIP 模型（包括使用 MLP 和 KAN 的模型）进行了测试。结果如表 5 所示。所有三个模型的精度相近，这可能是由于 Ag 数据集的简单性所致。此外，用高斯基的 KAN 替换 MLP 并未缩短训练时间。然而，用 B 样条基的 KAN 替换 MLP 则显著缩短了训练时间。

图5.  在 NequIP模型的输出块中，用采用 B 样条和高斯基函数的 KAN 替换 MLP。用 B 样条基函数的 KAN 替换 MLP 可提高预测精度，并显著缩短训练时间。

表 5.  使用高斯基函数和B样条基函数替换NequIP输出模块中的MLP为KANs在Ag数据集上的结果。最佳结果以粗体显示。

MAE: 平均绝对误差; F: 力; E: 能量。

我们还研究了在 MACE 模型中用 KAN 替换 MLP 的效果。与 Allegro 模型中的输出块类似，我们用 KAN 替换了最后一层读出部分中的 MLP。此外，我们在其他层的线性组合后引入了 KAN 层，如图 6所示。为了评估 KAN 模型在不同输入参数下的鲁棒性，我们在三个不同的随机种子（1111、2222 和 3333）下测试了每个模型。我们在碳数据集[53, 54]上测试了使用 MLP 和 KAN 的 MACE 模型。该数据集包括 4,080 个用于训练的结构和 450 个用于测试的结构。数据集涵盖了各种碳结构，包括非晶表面、体相晶体、液态和非晶碳。由于非晶材料缺乏规则的重复模式，在精确建模原子相互作用时具有挑战性，该数据集较为复杂，被用来评估基于 KAN 的模型在处理不同材料类型时的鲁棒性。

图6.  将 MACE 模型输出块中的 MLP 替换为采用 B-spline 基函数的 KAN。

测试集上的力、能量和应力的均方根误差（RMSE）汇总在表6中。在输出块中使用KAN和MLP的MACE模型在精度上表现相当。值得注意的是，使用KAN的MACE模型与使用MLP的模型相比，训练时间显著缩短。对于所有三个在输出块中使用KAN并且具有不同随机种子的MACE模型，结果始终与使用MLP的结果相当。特别地，这些基于KAN的模型在所有场景中相比于基于MLP的模型也展示了更短的训练时间。这个结果表明，KAN能够高效地学习和泛化，尽管初始化参数有所变化，突出了其鲁棒性、稳定性和适应性。

表6. 使用B样条基函数将MACE模型中不同部分的MLP替换为KAN的结果，数据集为碳数据集。最佳结果以粗体显示。

在本研究中，我们采用了基于边的张量预测图神经网络（ETGNN），用于预测晶体的张量性质。如图7所示，与我们对机器学习势模型的修改一致，我们仅替换了边更新模块和节点输出模块中的MLP投影部分。该部分与机器学习势模型中的输出块类似，将输出特征转换为标量。我们使用基于高斯和B样条基函数的KAN替换了这些MLP。我们比较了在SiO₂数据集上使用MLP和KAN（采用不同基函数）的ETGNN模型的精度。该数据集包含3,992个随机扰动SiO₂结构。数据集按6:2:2的比例划分为训练集、验证集和测试集。我们利用使用MLP和采用高斯及B样条基函数的KAN的ETGNN模型计算了Born有效电荷。结果显示在表7中。将输出块中的MLP替换为基于高斯基函数的KAN显著提高了预测精度，同时减少了训练时间。该结果与我们在Allegro模型上的发现一致。相比于MLP，使用B样条基函数的KAN缩短了训练时间，但仍比使用高斯基函数的KAN稍长。这些结果与Ag数据集、HfO₂数据集以及碳数据集上的结果一致，表明KAN在提高精度和计算效率等方面的优势在不同材料系统中具有一致性。

图7. 在ETGNN模型的输出块中使用采用B样条和高斯基函数的KAN替换MLP。用基于高斯基函数的KAN替换输出块中的MLP显著提高了预测精度，同时减少了训练时间。

表 7.  将ETGNN输出块中的MLP替换为使用高斯基函数和B样条基函数的KAN在SiO₂数据集上的结果。最佳结果以粗体表示。

本文研究了在机器学习势能和性质预测模型中用Kolmogorov-Arnold 网络（KANs）替代多层感知器（MLPs）的效果。结果表明，仅在输出层使用KANs即可显著提升预测精度，同时优化了计算效率和资源利用率，为机器学习势能和材料性质预测计算提供了新方法。未来研究可将KANs应用于高通量材料筛选、目标材料的逆向设计、领域适应及迁移学习等方向，并集成至通用模型中。KANs的强大潜力为加速新材料发现和机器学习模型开发提供了新机遇。

国家重点研发计划（No. 2022YFA1402901）、国家自然科学基金（No. 11991061, 12188101）、上海科技计划（No. 23JC1400900）、广东省基础与应用基础研究（极端条件下未来应用材料2021B0301030005）。

向红军，复旦大学物理学系谢希德特聘教授。2001年，2006年分别获得中国科学技术大学学士和博士学位，2006年—2007年美国北卡州立大学博士后，2007年—2009年美国国家再生能源实验室博士后。2018年获“国家杰出青年基金”，2018年获国际理论物理中心“ICTP”奖，2021年入选美国物理学会会士，2023年获黄昆物理奖。主要研究领域为计算凝聚态物理，包括磁性、铁电性、多铁性、纳米材料的理论计算研究、计算方法发展及程序开发、机器学习方法在计算凝聚态物理中的应用等。主要研究成果包括建立了自旋序诱导多铁性的普适模型，发现了多种多铁性新机制；建立了计算磁相互作用及磁电耦合强度的四态法，已被国内外同行广泛采用；自主开发了材料性质分析和模拟软件包（PASP），已被国内外多个研究组采用。迄今发表SCI论文200多篇，其中Phys. Rev. Lett. 30多篇。