机器学习粗粒化分子动力学中的Hessian匹配方法

# 面向机器学习粗粒化分子动力学的海森匹配  
来源：https://arxiv.org/html/2605.12823  

\\undefine@key newfloatplacement\\undefine@keynewfloatname\\undefine@keynewfloatfileext\\undefine@keynewfloatwithin  

Sanya Murdeshwar¹，Sanjit Shashi¹,²，Kevin Bachelor¹，William Noid³，Ashwin Lokapally²，Razvan Marinescu¹,²  
¹加州大学圣克鲁兹分校  
²GiwoTech Inc.  
³宾夕法尼亚州立大学  

\{smurdesh, sashashi, kwbachelor, ramarine\}@ucsc.edu  [email protected]  [email protected]  

###### 摘要  

粗粒化（CG）分子动力学能够模拟生物分子等原子体系，其时间尺度是全原子（AA）方法无法触及的。然而，现有通过力匹配训练的CG神经势只捕捉自由能面的梯度，其曲率未受约束。我们提出了一种框架，在力匹配的基础上增加随机海森-向量积（HVP）匹配，从而在不构建完整海森矩阵的情况下，将二阶曲率信息注入CG势中。我们推导了目标CG海森矩阵的分解：一部分是模型无关的投影AA海森矩阵，可在训练前一次性预计算；另一部分是模型依赖的协方差修正项，在前向传播过程中以可忽略的代价在线计算。我们利用随机探针向量构建了海森匹配目标的无偏随机估计器。我们在训练中未曾见过的九个快速折叠蛋白基准上，将我们的方法与力匹配进行了对比评估。在9个蛋白中，HVP匹配在慢模指标上优于纯力匹配，其中8个蛋白表现更佳，在最大蛋白最慢集体模式上，CG分布与参考分布之间的KL散度最多降低了85%。我们的结果表明，高阶物理监督是实现更准确、可迁移的CG生物分子势的实用途径。

## 1 引言

分子动力学（MD）模拟是理解生物分子动力学和功能特性的主要工具之一。然而，全原子（AA）MD可访问的时间尺度与研究者期望研究的生物过程在根本上存在差距。某些有意义的生物过程，如蛋白质折叠，可能在毫秒甚至秒级尺度上演化，远超专用AA MD硬件所能达到的微秒级动力学（Shaw2009; Shaw2021）。这一差距推动了粗粒化（CG）建模的发展（Noid2013; Kmiecik2016）。CG模型将一组原子表示为单个相互作用位点（珠子），从而减少自由度数量，并允许模拟更长的时间尺度。核心挑战在于学习一个CG势，使其自由能面能够再现底层AA系统的自由能面，包括亚稳态的位置、状态间势垒的高度，以及控制每个势阱内涨落的局部曲率。

传统上，CG建模采用力匹配方法（Ercolessi1994; Izvekov2005; Noid2008），通过拟合CG势来重现AA参考的平均力。最初的方法使用预定义的函数形式，但随后CG建模被推广到神经网络近似器，例如CGnet（Wang2019，一种前馈网络）和CGSchNet（Husic2020，一种图神经网络）。这些模型已成为学习生物分子CG势的最先进方法，并且可以训练出能够对快速折叠蛋白进行稳定、合理准确的模拟。

尽管取得了进展，力匹配仍存在一个根本局限性：它只能捕捉自由能面的梯度，而无法捕捉其曲率。两个势可以产生相同的力，但在这些力对扰动的响应上可能存在显著差异，而力匹配无法区分它们。在实践中，这导致亚稳态布居恢复不佳（Thaler\_Stupp\_Zavadlav\_2022; Köhler\_Chen\_Krämer\_Clementi\_Noé\_2023），并且力匹配验证损失无法反映全局自由能面质量（Thaler\_Stupp\_Zavadlav\_2022）。因此，仅凭力匹配训练的CG神经势会在长时间训练后，由于过度拟合梯度信号并失去能量景观的形状，使得慢构象模式的性能下降。此外，在一组蛋白上训练的模型对未见过的、分布以外的序列推广能力差，会产生训练期间未采样的构型中不切实际的低能量（Majewski2023），这直接源于局部曲率欠约束的自由能面。这些问题从根本上限制了CG神经势的数据效率、泛化能力和稳定性。

解决这一问题的自然方法是在匹配自由能面梯度（力）的同时，匹配其海森矩阵。海森矩阵编码了控制振动频率和势垒形状的局部曲率信息，提供了纯一阶力匹配所缺失的高阶物理信息。然而，在训练CG神经势时直接引入海森监督在扩展性方面带来了重大计算挑战。对于一个具有 \(d = 3N\) 个自由度的CG系统，完整海森矩阵有 \(d^2\) 个元素，并且需要 \(\mathcal{O}(d)\) 次力评估才能逐列构建，而梯度只有 \(d\) 个元素，仅需 \(\mathcal{O}(1)\) 次力评估。对于小型系统，这可行但速度较慢。对于较大系统，\(d\) 可上升到数千，显式构建在内存和计算上都变得难以处理。

### 1.1 相关工作

最近的一些工作将二阶曲率信息引入原子机器学习的原子间势（MLIP），并在外推、数据效率和过渡态优化方面显示出改进（Fang2024BeyondNumerical; Yuan2024AnalyticalHessian; Rodriguez2025HessianData）。与我们的方案最密切相关的是投影海森学习（PHL）（Rodriguez2026PHL），它用随机海森-向量积探针替代显式海森构建。PHL产生一个无偏的迹基损失，以不到8%的成本获得了完整海森监督的大部分收益。海森原子间势（HIP）（HIP2025）采取了不同的路线，通过等变读出直接预测海森矩阵，而不是通过自动微分。然而，所有这些方法都适用于在非常小分子的精确量子力学海森矩阵上训练的原子MLIP。CG系统可以扩展到数千自由度，这使得随机HVP探针变得必不可少而非可选。此外，CG海森分解为投影AA海森矩阵加上由积分掉的原子自由度产生的协方差修正项，后者在原子层面没有类似物。我们还注意到，CG建模的核心目标是对未见过的蛋白质构象和序列进行泛化，而现有任何海森知情方法尚未为此进行评估。

另一条并行研究方向致力于通过多样化的多蛋白训练实现可迁移的CG蛋白势。Majewski等人（Majewski2023）在十二个蛋白上联合训练了CG势，Charron等人（Charron2025Navigating）开发了一种基于CGSchNet的可迁移力场，在新序列上展示了外推MD。两者都仅依赖力匹配，通过数据规模而非更丰富的物理监督实现可迁移性。Ciccotti等人（Ciccotti2005）通过矢量的蓝月系综发展了具有一般（线性和非线性）CG映射的力匹配理论基础，并由Kalligiannaki等人（Kalligiannaki2015）重新表述为几何投影问题，他们还证明了对于一般CG映射，力匹配与相对熵最小化渐近等价。我们的工作将这个框架从自由能的一阶导数扩展到了二阶导数。

我们的工作也不同于优化器侧的曲率方法，如L-BFGS、K-FAC和海森自由优化（Martens2010; MartensGrosse2015; Pearlmutter1994），这些方法是在参数空间中近似训练目标的海森矩阵。相反，我们在构型空间中匹配CG势的海森矩阵，它编码了能量景观的物理曲率，而非优化动力学。

#### 我们的贡献。

我们提出了一种框架，通过海森-向量积（HVP）实现随机海森匹配，用于训练CG神经势。我们证明了CG海森恒等式可以清晰地分解为两项，它们具有根本不同的计算特性：一项是投影AA海森矩阵项，是模型无关的，可以在训练前一次性预计算；另一项是协方差修正项，是模型依赖的，但由前向传播过程中已经可用的量构建。我们将这个 \(d \times d\) 矩阵的作用与 \(K\) 个随机探针向量相匹配，避免了显式构建整个海森矩阵，同时保留了对完整海森匹配目标的无偏估计。由此产生的损失每个帧增加 \(\mathcal{O}(Kd)\) 的额外工作（对于 \(K\) 个探针），与系统大小成线性关系，并且不需要对底层模型进行架构更改。

我们在快速折叠基准蛋白上的实验结果表明，用HVP匹配补充力匹配提供了以下优势：
- (i) 在慢模准确性（由捕捉折叠/去折叠转变的时间滞后独立分量表示）上最多改进85%；
- (ii) 当在单独的单链数据集上训练时，对未见过的蛋白具有更好的泛化能力。

最终，这些发现表明，将高阶物理信息注入CG神经势是实现更准确、更可迁移生物分子模拟的实用且可扩展的途径。

## 2 推导粗粒化海森恒等式

我们推导自由能面的CG海森矩阵的符号表达式，该表达式作为我们模型的训练目标。对于蛋白质，标准的粗粒化机制将每个氨基酸残基的Cα碳原子作为一个珠子。注意，由于每个珠子是从AA原子中选择的，映射是线性的，这大大简化了我们的推导。一般非线性表达式见附录A（https://arxiv.org/html/2605.12823#A1）。我们使用蓝月系综（Carter1989; Kidder2021）来计算与约束分子系统相关的热力学量。

假设AA系统包含 \(n\) 个原子，CG系统由 \(N\) 1AAJ 1ACF 1AMM 1AYD 1BK2 1BZ4 1CM2 1DVN 1EQ6 1EYH 1FNA 1GYV 1I2T 1IO2 1J7X 1JMW 1JW4 1K40 1LPJ 1MJC 1N81 1NG6 1OGW 1P1L 1QTP 1QZM 1UKF 1W8V 1WP5 1WWI 1WY3 1X3O 1XGD 1XT0 1YP5 1YU5 1YW5 1ZEQ 1ZLM 2BFH 2EVB 2H0M 2ICT 2NSC 2OJ4 2RH3 2V14 2VQ4 2W9Q 3HVM 3ID4 3KJE 3LFO 3MX7 3MZZ 3P7K 3PG4 3R6D 3RJP 4BPF 4J5Q 4N6T 4QAJ 4QBE 4RZ9 4U3H 4Y2K 5BTH 5GY3 5IHW 5J2V 5UVR 5XEF 5Y8E 5YM7 6APK 6DR3 6GPM 6KND 6L7Q 6LG3 6MRR 6Q3V 6SYG 6TYY 6WEY 6WIN 7DMF 7DMS 7LIQ 7TGP 7XCD 8A8S 8ERE 8G22 8GQQ 8HQW 8I2D 8JED