[论文解读] Two for One: Diffusion Models and Force Fields for Coarse-Grained Molecular Dynamics

Marloes Arts, Victor García Satorras+7 more|arXiv (Cornell University)|Feb 1, 2023

Protein Structure and Dynamics被引用 10

一句话总结

本论文在粗粒度蛋白质结构上学习去噪扩散模型，以同时获得独立同分布的CG样本和用于分子动力学的保守CG力场，从而实现准确的平衡分布和保持的动力学。

ABSTRACT

Coarse-grained (CG) molecular dynamics enables the study of biological processes at temporal and spatial scales that would be intractable at an atomistic resolution. However, accurately learning a CG force field remains a challenge. In this work, we leverage connections between score-based generative models, force fields and molecular dynamics to learn a CG force field without requiring any force inputs during training. Specifically, we train a diffusion generative model on protein structures from molecular dynamics simulations, and we show that its score function approximates a force field that can directly be used to simulate CG molecular dynamics. While having a vastly simplified training setup compared to previous work, we demonstrate that our approach leads to improved performance across several small- to medium-sized protein simulations, reproducing the CG equilibrium distribution, and preserving dynamics of all-atom simulations such as protein folding events.

研究动机与目标

• 推动在训练中不使用原子力的情况下学习热力学一致性的CG力场。
• 证明基于分数的扩散模型可以逼近CG势能表并可直接用于CG MD。
• 证明方法可扩展至更大蛋白质并保持全原子级别的动力学特性在CG模拟中的体现。
• 提供一个简单的一阶段训练方法，实现i.i.d. CG样本生成与确定性CG力场两者。

提出的方法

• 在从CG Boltzmann分布抽取的CG样本上训练去噪扩散概率模型（DDPM）.
• 将DDPM损失与去噪得分匹配相关联，以提取近似CG力的得分函数。
• 将扩散模型的噪声预测器参数化为能量函数的梯度，以确保保守、平移不变、并且在SO(3)等变性下具备旋转等变性（通过数据增强实现）。
• 在CG朗之万动力学中使用提取的去噪力场来模拟CG分子动力学。
• 允许选择扩散步骤i，以在力场的精度和MD期间的数值稳定性之间取得平衡。

实验结果

研究问题

• RQ1 一个在CG平衡样本上训练的扩散模型是否能够学习一个适用于MD模拟的保守CG力场？
• RQ2 该方法是否再现CG平衡分布并保持与全原子模拟相当的动力学？
• RQ3 该方法是否可扩展到比以往基于流的方法更大的CG蛋白质？
• RQ4 学习得到的力场是否在不需要明确能量先验或师生蒸馏的情况下提供稳定的CG MD？
• RQ5 该方法在i.i.d.样本生成与基于MD的模拟两种模式下的表现如何？

主要发现

Method	Chignolin TIC JS	Chignolin PWD JS	Trp-cage TIC JS	Trp-cage PWD JS	Bba TIC JS	Bba PWD JS	Villin TIC JS	Villin PWD JS	Protein G TIC JS	Protein G PWD JS
Reference	.0057	.0002	.0026	.0002	.0040	.0002	.0032	.0004	.0014	.0002
Flow i.i.d.	.0106	.0022	.0078	.0057	.0229	.0073	.0109	.0142	n/a	n/a
DFF sim.	.0096	.0005	.0052	.0007	.0111	.0017	.0073	.0009	.0131	.0009
Flow-CGNet sim.	.1875	.1271	.1009	.0474	.1469	.0594	.2153	.0535	n/a	n/a
DFF sim.	.0335	.0067	.0518	.0403	.1289	.0408	.0564	.0244	.2260	.0691

• DFF方法在再现CG平衡分布方面优于基于流的方法，量化为关节扭转角与成对距离的JS散度。
• DFF仿真在通过基于TIC的亚稳态的转移概率等动力学行为上比Flow-CGNet仿真更准确，DFF的i.i.d.在采样质量上接近Flow i.i.d.。
• 扩散模型比基于流的方法更好地捕捉全局结构与远程接触，即使对于如Protein G（56个珠子）这类更大蛋白质也如此。
• 所学习的DFF是一个保守、平移不变、并且在SO(3)等变性下具备等变性的力场，适用于稳定的CG MD，而无需显式能量先验。
• 该框架可扩展至快速折叠蛋白，并保持关键的动力学特征，包括折叠/未折叠事件和亚稳态种群，且具备有利的基于TIC的指标。