AdaptNTK：用于神经网络势的自适应不确定性量化与主动学习

# AdaptNTK：神经网络势的自适应不确定性量化与主动学习  
来源：https://arxiv.org/html/2609.00488  
舒文岳  
所属单位：美国纽约州伊萨卡市康奈尔大学化学与生物分子工程系  
通讯作者：[shuwen\.yue@cornell\.edu](mailto:[email protected])  

###### 摘要  
机器学习原子间势桥接了量子化学精度与经典计算速度之间的鸿沟，使得具有第一性原理精度的分子动力学模拟成为可能。其可靠性通常通过主动学习来提升，即通过识别不确定的、分布外的构型来迭代扩展训练集。现有的不确定性量化方法通常在计算成本与可靠性之间存在权衡，并且在构建采集批量时通常无法考虑冗余性。本文引入了 AdaptNTK，一个单模型框架，通过经验神经切线核（NTK）特征空间中的正则化马氏距离来度量不确定性。由于 NTK 特征在采集过程中保持固定，不确定性仅取决于已采集的构型，而不依赖其参考标签。这使得不确定性可以在每次选择后无需重新训练即可递归更新，从而减少采集批量内的冗余。在留出的 rMD17 数据上，AdaptNTK 实现了与力误差最高的平均相关性（斯皮尔曼系数 0.68，皮尔逊系数 0.71），并在误差保留率上匹配三成员集成模型。在主动学习实验中，AdaptNTK 在 rMD17 和 Transition-1X 上实现了最低的力误差，特别是在 Transition-1X 的过渡态构型上表现突出。AdaptNTK 每个 Transition-1X 循环比集成模型快 2.6 倍，为数据高效的主动学习提供了高效的单模型不确定性估计和顺序更新能力。  

## 1 引言  
分子动力学（MD）模拟广泛应用于研究原子分辨率下的化学过程。然而，其精度和适用范围传统上受限于计算成本与底层势能面（PES）质量之间的权衡。电子结构方法能准确描述键断裂、电荷重新分布和小能量差，但计算成本依然高昂；而经验力场可访问更大的长度和时间尺度，但通常无法泛化到新的化学环境。机器学习原子间势（MLIPs）通过从参考计算中学习势函数并近似第一性原理的能量和力，解决了这一局限性，其计算成本可随系统规模线性扩展（Behler and Parrinello, 2007；Bartók et al., 2010；Unke et al., 2021；Batzner et al., 2022；Batatia et al., 2022；Musaelian et al., 2023）。尽管高效，MLIPs 强依赖于其训练数据。它们通常在训练集所代表的构型空间区域内表现良好，但当轨迹遇到这些区域之外的构型时，其预测性能可能下降。这种外推误差可能导致不准确的力，进而产生错误的轨迹。因此，可靠的不确定性估计对于识别模型预测可能不可信的构型非常重要。高斯过程（GP）势通过其后验方差自然地提供不确定性估计，该方差依赖于训练构型但不依赖其参考标签，从而允许在训练集变化时解析评估和更新不确定性。另一方面，神经网络势更具表达力，但通常缺乏类似的闭式不确定性估计。其不确定性通常通过集成（Seung et al., 1992；Smith et al., 2018）或近似贝叶斯方法（Gal and Ghahramani, 2016；Maddox et al., 2019；Amini et al., 2020；Soleimany et al., 2021）来近似，这些方法可能计算成本高昂或在训练分布外不可靠。因此，为单个训练好的 MLIP 开发可靠且高效的不确定性估计仍是一个重要挑战。  

主动学习是不确定性量化的天然下游应用。可以将具有高不确定性的构型选择用于参考标记并添加到训练集中（Podryabinkin and Shapeev, 2017；Smith et al., 2018；Vandermause et al., 2020；Zhang et al., 2020）。然而，当在批量中选择多个构型时，仅在批量开始时计算的不确定性估计无法考虑已选择的构型，可能导致冗余采集。GP 后验方差天然支持在每次选择后更新不确定性，而无需对应的参考标签（Vandermause et al., 2020）。为神经网络势实现类似能力更具挑战性，因为传统的不确定性估计通常需要额外的模型拟合或重新训练来纳入新选择的构型。与神经切线核（NTK）相关的参数梯度特征（Jacot et al., 2018）为在训练好的神经网络势中恢复这种类 GP 能力提供了一个原理性框架。本文介绍了 AdaptNTK，一个用于神经网络势不确定性量化和主动学习的框架。AdaptNTK 使用每个构型预测能量对所有模型参数的梯度来表示该构型，并在此切线特征空间中分配从训练数据出发的正则化距离。由于不确定性取决于选择的构型而非其参考标签，因此可以在采集批量组装过程中，每次选择后无需重新训练网络即可解析更新。我们的贡献包括：  

- • 我们将正则化经验 NTK 不确定性与特征空间几何联系起来，评估其排序、保留和校准性能，并证明了一个由岭有效维数控制的高斯基保障。  
- • 我们使用无标签的秩一更新，在每次选择后调整剩余的不确定性分数，减少批量内的冗余，而无需额外的模型重新训练。  
- • 我们在 rMD17（Christensen and von Lilienfeld, 2020）和 Transition-1X（Schreiner et al., 2022）上，针对集成和单模型基线，对 AdaptNTK 的力误差排序、误差保留、不确定性校准和主动学习性能进行了基准测试。  

总体而言，AdaptNTK 在评估的不确定性量化方法中实现了与力误差最高的平均相关性，并在 rMD17 和 Transition-1X 的主动学习实验中实现了最低的力误差。其归一化风险-覆盖性能与三成员集成模型相当，而在报告的 Transition-1X 时间实验中，其实际耗时成本不到集成模型的一半。这些结果共同表明，草绘经验 NTK 不确定性为神经网络势的数据高效主动学习提供了一个实用的单模型基础。  

## 2 相关工作  
已经开发了几种不确定性量化方法来估计深度神经网络中的预测不确定性（Seung et al., 1992；Gal and Ghahramani, 2016；Lakshminarayanan et al., 2017；Maddox et al., 2019；Amini et al., 2020；Soleimany et al., 2021）。这些方法也已应用于原子机器学习，特别是通过集成和委员会查询（QBC）方法（Smith et al., 2018；Podryabinkin and Shapeev, 2017；Zhang et al., 2020；Kahle and Zipoli, 2022）。最近的工作开发了基于 Dropout 和学习特征分布（Wen and Tadmor, 2020；Zhu et al., 2023；Bigi et al., 2024；Chong et al., 2025）、证据回归（Xu et al., 2026）、共享主干集成（Kellner and Ceriotti, 2024；Beck et al., 2025）、变分推断（Coscia et al., 2026；Mamun et al., 2026）以及事后不确定性估计与校准（Vita et al., 2025；Perez et al., 2025；Ho et al., 2026）的单模型或低成本替代方案。然而，这些方法通常在计算可扩展性与不确定性估计的可靠性之间存在权衡（Tan et al., 2023；Bilbrey et al., 2025）。最近，Wilson et al. (2025) 表明，基于经验神经切线核（NTK）的不确定性量化可以在一般回归设置中以远低于深度集成的成本匹配或超越深度集成的性能。他们的方法是事后且基于采样的，通过线性化网络的梯度下降采样构建集成，而我们的不确定性估计则直接从经验 NTK 特征空间中的闭式正则化二次型获得。批量主动学习方法也在机器学习应用中得到了广泛研究，包括图像分类和表格回归。最后一层和梯度特征表示通常与批量采集的多样性标准配对（Holzmüller et al., 2023；Ash et al., 2020；Sener and Savarese, 2018），相关方法也已应用于原子间势（Zaverkin et al., 2022）。Varga-Umbrich et al. (2026a) 和 Varga-Umbrich et al. (2026b) 的并发工作也使用神经切线核特征和具有多样性意识的选择来进行 MLIPs 的主动学习。这些工作证明了基于特征的表示和批量内多样化对于选择信息性训练数据的效用。然而，它们的主要关注点是固定预算采集，其中特征分数用于对候选构型进行排序，而不是为每个构型提供一个可校准的定量不确定性。我们的关注点是互补的：我们使用经验 NTK 特征构建一个点态不确定性度量，该度量可以在构型被采集时进行校准和自适应更新。这种区别对于分子模拟非常重要，因为在轨迹演化过程中，不确定性需要在单个构型上进行评估。这种估计可以识别偏离训练分布的情况，为数据采集提供停止标准，并在长时间运行的模拟中指示模型预测可能不可靠的时机。  

## 3 背景  
### 3.1 高斯过程  
高斯过程（GP）定义了函数 f:Rd→R 上的先验，其均值函数为 m(·)，协方差核为 k(·,·)（Rasmussen and Williams, 2006）。给定一个带有观测噪声 σ² 的数据集 D=(X,y)，测试输入 x⋆ 处的后验预测分布是高斯分布，其方差为 Σ(x⋆)=k(x⋆,x⋆)−k_{x⋆,X}(K_{X,X}+σ²I)⁻¹k_{X,x⋆}。在本文中，我们重点关注预测方差，它作为不确定性度量并构成了我们采集策略的基础。  
### 3.2 神经切线核  
考虑一个参数化为 θ∈RP 的神经网络 f(·,θ):Rd→R。其在参数 θ 处的神经切线特征映射为 φθ(x)=∇θf(x,θ)∈RP，它诱导出经验神经切线核（NTK）Kθ(x,x′)=φθ(x)⊤φθ(x′)。在无限宽度极限下，NTK 收敛为一个确定性核，在标准假设下，神经网络训练可以用相应的核动力学来描述（Jacot et al., 2018；Lee et al., 2019）。本文中，我们使用的是一个有限、训练好的网络的经验 NTK。将该核视为高斯过程协方差，可以得到一个预测不确定性，该不确定性可以直接用切线特征矩阵 Φ∈Rn×P 来表示。这种构造对应于学习参数周围的局部线性化，且不假设模型在无限宽度区间下运行。  

## 4 方法  
### 4.1 动机与理论分析  
我们首先通过将经验神经切线核（NTK）诱导的高斯过程（GP）方差与 NTK 特征空间中的距离联系起来，为预测不确定性发展一种几何解释。令 Φ∈Rn×P 表示训练数据集的经验 NTK 特征矩阵，其中每行对应于一个训练构型的切线特征表示。将经验 NTK 视为协方差核，其