Allegro-Legato：基于锐度感知最小化的可扩展、快速且鲁棒的神经网络量子分子动力学

基于机器学习的神经网络量子分子动力学（NNQMD）模拟正以量子力学精度但快数个量级的优势革新材料的原子尺度模拟，这体现于 ACM Gordon Bell 奖（2020）及入围（2021）。建立在群论之上、具有旋转等变性与局域描述子的当前最优（SOTA）NNQMD 模型比那些模型提供了高得多的精度与速度，故命名为 Allegro（意为快速）。然而，在大规模并行超级计算机上，它遭受保真度随规模退化的问题：越来越多的原子间力的非物理预测禁止了涉及更多原子更长时间的模拟。在此，我们通过将 Allegro 模型与锐度感知最小化（SAM）结合来解决该问题，借助改善损失景观的光滑性来增强模型的鲁棒性。所得的 Allegro-Legato（意为快速且“光滑”）模型被证明延长了失效时间 $t_\textrm{failure}$，且不牺牲计算速度或精度。具体而言，Allegro-Legato 的失效时间对问题规模的依赖弱得多，$t_{\textrm{failure}} \propto N^{-0.14}$（$N$ 为原子数），而 SOTA Allegro 模型为 $\left(t_{\textrm{failure}} \propto N^{-0.29}\right)$，即系统性地延迟了失效时间，从而允许更大规模更长时间的 NNQMD 模拟而不失效。该模型在 Argonne Leadership Computing Facility 的 Polaris 超级计算机上还展现出优异的计算可扩展性与 GPU 加速。此类可扩展、精确、快速且鲁棒的 NNQMD 模型很可能在新一代 exaflop/s 计算机上的 NNQMD 模拟中找到广泛应用，其中一个具体例子是在氨的动力学中计入核量子效应。