[论文解读] Simulate Time-integrated Coarse-grained Molecular Dynamics with Multi-Scale Graph Networks

Xiang Fu, Tian Xie+3 more|arXiv (Cornell University)|Apr 21, 2022

Machine Learning in Materials Science被引用 14

一句话总结

本文提出了一种可微分的多尺度图神经网络（GNN），可直接以皮秒级时间步长模拟粗粒化分子动力学，避免了力的积分计算。通过结合用于粗粒化粒子表示的分层GNN与基于扩散的精炼模块，该模型在聚合物和电池电解质体系中实现了相较于经典力场10^3–10^4倍的速度提升，即使在未见过的组成和超过训练数据长度的轨迹上，也能准确预测长期的结构与动力学性质。

ABSTRACT

Molecular dynamics (MD) simulation is essential for various scientific domains but computationally expensive. Learning-based force fields have made significant progress in accelerating ab-initio MD simulation but are not fast enough for many real-world applications due to slow inference for large systems and small time steps (femtosecond-level). We aim to address these challenges by learning a multi-scale graph neural network that directly simulates coarse-grained MD with a very large time step (nanosecond-level) and a novel refinement module based on diffusion models to mitigate simulation instability. The effectiveness of our method is demonstrated in two complex systems: single-chain coarse-grained polymers and multi-component Li-ion polymer electrolytes. For evaluation, we simulate trajectories much longer than the training trajectories for systems with different chemical compositions that the model is not trained on. Structural and dynamical properties can be accurately recovered at several orders of magnitude higher speed than classical force fields by getting out of the femtosecond regime.

研究动机与目标

• 为克服经典力场与机器学习力场在模拟大规模、长时间分子动力学时面临的计算瓶颈。
• 实现在皮秒级时间步长下的粗粒化分子动力学直接模拟，避免力计算与飞秒级积分约束。
• 通过一种新颖的基于扩散的精炼模块，解决粗粒化模型在长时间模拟中的不稳定性问题。
• 实现对复杂体系（如单链聚合物和多组分锂离子聚合物电解质）的高速、可泛化模拟，其性能超越训练数据范围。

提出的方法

• 设计了一种多尺度图神经网络架构，包含两个组件：用于在细粒度层级学习粗粒化粒子表示的嵌入GNN，以及用于在粗粒度层级模拟时间积分的粗粒化动力学GNN。
• 模型直接预测大时间步长（最高达纳秒级）内的状态转移，无需力积分，从而实现高速模拟。
• 引入一种受扩散模型启发的基于分数的精炼模块，用于校正不稳定的动力学行为并提升长时间模拟的准确性。
• 模型在短MD轨迹（例如5 ns）上进行训练，并可泛化至模拟比训练数据长数个数量级的轨迹（例如50 ns），无需重新训练。
• 该框架完全可微分，支持粗粒化动力学的端到端训练与优化。
• 该方法在两个真实体系上进行了评估：隐式溶剂中的单链聚合物和多组分锂离子聚合物电解质。

实验结果

研究问题

• RQ1 学习得到的、可微分的模拟器是否能够绕过力计算，实现在皮秒级时间步长下的粗粒化分子动力学直接模拟？
• RQ2 当模拟轨迹远长于训练数据时，机器学习模型如何保持稳定性和准确性？
• RQ3 模型是否能在未见化学组成下实现泛化，同时保持关键的结构与动力学性质？
• RQ4 与标准GNN基线相比，基于扩散的精炼模块在多大程度上提升了模拟的稳定性？
• RQ5 与经典力场和机器学习力场相比，模拟长时间轨迹在实际运行时间上可实现多大的加速？

主要发现

• 该模型在模拟长时间轨迹时，相较于最快的经典力场实现10^3至10^4倍的实际运行时间加速。
• 仅使用5 ns的训练轨迹，该模型即可成功预测聚合物电解质中锂离子和TFSI离子的扩散系数，优于监督型GNN基线模型，后者在缺乏长轨迹时无法收敛。
• 若不使用扩散精炼模块，模型表现出动力学不稳定性，且无法恢复长期扩散系数，凸显了精炼步骤的关键作用。
• 该模型可泛化至未见的化学组成，并能模拟远超训练数据长度的轨迹，准确恢复平衡态与动力学性质。
• 该方法使大规模体系的高通量筛选成为可能，而这些体系因时间和系统规模限制，对现有机器学习力场而言计算成本过高。
• 该模型在两个复杂体系上得到验证：单链聚合物和多组分锂离子聚合物电解质，展现出强健性与可迁移性。