[论文解读] A Note on Learning Rare Events in Molecular Dynamics using LSTM and Transformer

Wenqi Zeng, Siqin Cao+2 more|arXiv (Cornell University)|Jul 14, 2021

Machine Learning in Materials Science参考文献 20被引用 15

一句话总结

本文研究了LSTM和Transformer模型在从高维分子动力学（MD）轨迹中学习罕见事件时的局限性。结果表明，时间分辨率和状态划分至关重要：将构象空间粗粒化为亚稳态并去除来回穿越事件可显著提升对慢动态过程的预测性能；而若无此类预处理，两种模型均无法有效分离快动态与慢动态。

ABSTRACT

Recurrent neural networks for language models like long short-term memory (LSTM) have been utilized as a tool for modeling and predicting long term dynamics of complex stochastic molecular systems. Recently successful examples on learning slow dynamics by LSTM are given with simulation data of low dimensional reaction coordinate. However, in this report we show that the following three key factors significantly affect the performance of language model learning, namely dimensionality of reaction coordinates, temporal resolution and state partition. When applying recurrent neural networks to molecular dynamics simulation trajectories of high dimensionality, we find that rare events corresponding to the slow dynamics might be obscured by other faster dynamics of the system, and cannot be efficiently learned. Under such conditions, we find that coarse graining the conformational space into metastable states and removing recrossing events when estimating transition probabilities between states could greatly help improve the accuracy of slow dynamics learning in molecular dynamics. Moreover, we also explore other models like Transformer, which do not show superior performance than LSTM in overcoming these issues. Therefore, to learn rare events of slow molecular dynamics by LSTM and Transformer, it is critical to choose proper temporal resolution (i.e., saving intervals of MD simulation trajectories) and state partition in high resolution data, since deep neural network models might not automatically disentangle slow dynamics from fast dynamics when both are present in data influencing each other.

研究动机与目标

• 研究LSTM和Transformer模型在从分子动力学（MD）轨迹中学习罕见事件时的性能表现。
• 识别影响深度学习模型在复杂随机系统中捕捉慢动态过程准确性的关键因素。
• 评估反应坐标维度、时间分辨率（保存间隔）以及状态划分对模型性能的影响。
• 比较LSTM与Transformer在处理具有快慢混合动态的高维MD数据时的有效性。
• 提出预处理策略（如动力学归并和来回穿越事件去除），以提升模型在MD模拟中学习罕见事件的能力。

提出的方法

• 将丙氨酸二肽的MD轨迹投影到二维反应坐标（φ, ψ）上，并采用基于RMSD的聚类方法实现高维状态划分。
• 使用交叉熵损失函数，基于序列状态转移训练LSTM和Transformer模型，以学习路径熵和动力学特性。
• 通过改变保存间隔（0.1 ps, 1 ps, 2 ps, 10 ps）评估时间分辨率对模型性能的影响。
• 应用动力学归并与来回穿越事件去除，以减少快速动态带来的噪声，并增强状态空间表示中的慢动态信号。
• 在Transformer中采用双向与单向编码，并使用三角形掩码模拟单向行为。
• 通过与参考MD数据对比，利用平均首次通过时间（MFPT）、首次通过时间分布及隐含时间尺度（ITS）评估模型性能。

实验结果

研究问题

• RQ1 MD轨迹的时间分辨率（保存间隔）如何影响LSTM和Transformer在学习慢动态过程中的表现？
• RQ2 反应坐标的维度在深度学习模型捕捉罕见事件的能力中起到何种作用？
• RQ3 状态划分与来回穿越事件去除是否能提升LSTM和Transformer对慢动态过程预测的准确性？
• RQ4 Transformer中的注意力机制是否在从高维MD数据中学习长期动态方面相比LSTM具有性能优势？
• RQ5 LSTM和Transformer模型在原始MD轨迹中在多大程度上能自动分离慢动态与快动态？

主要发现

• 在0.1 ps保存间隔下，LSTM性能显著下降，尽管其自由能景观准确，但快慢动态混杂导致性能受损。
• 在10 ps保存间隔下，数据量不足限制了模型学习，导致慢动态预测性能差。
• LSTM在1 ps和2 ps间隔下能成功预测自由能景观与动力学，但在0.1 ps（噪声干扰）和10 ps（数据稀缺）间隔下均失败。
• 尽管采用双向注意力机制，Transformer模型在慢动态预测方面并未优于LSTM，原因在于序列长度有限（100个token）且难以有效处理长程依赖。
• 去除来回穿越事件并对接收状态分区实施动力学归并，可显著提升LSTM与Transformer在高维RMSD基数据上的性能。
• 正确的时间分辨率与状态划分组合至关重要——深度学习模型无法在无此类预处理的情况下自动分离慢动态与快动态。