[논문 리뷰] A Note on Learning Rare Events in Molecular Dynamics using LSTM and Transformer

Wenqi Zeng, Siqin Cao외 2명|arXiv (Cornell University)|2021. 07. 14.

Machine Learning in Materials Science참고 문헌 20인용 수 15

한 줄 요약

이 논문은 고차원 분자 동역학(MD) 궤적에서 희귀 사건을 학습하는 데 있어 LSTM 및 Transformer 모델의 한계를 조사한다. 시간 해상도와 상태 분할이 핵심임을 보여주며, 혼합된 빠른 동역학과 느린 동역학을 분리하기 위해 구성공간을 준안정 상태로 군집화하고 재교차 사건을 제거함으로써 느린 동역학 예측 성능을 크게 향상시킬 수 있음을 밝힌다. 반면, 이러한 전처리 없이 둘 다 빠른 동역학과 느린 동역학을 분리하지 못한다.

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.

연구 동기 및 목표

• 분자 동역학(MD) 궤적에서 희귀 사건을 학습하는 데 있어 LSTM 및 Transformer 모델의 성능을 조사하기 위해.
• 복잡한 확률적 시스템에서 느린 동역학을 포착하는 데 있어 딥러닝 모델의 정확도에 영향을 미치는 핵심 요인을 규명하기 위해.
• 반응좌표의 차원 수, 시간 해상도(저장 간격), 상태 분할이 모델 성능에 미치는 영향을 평가하기 위해.
• 고차원 MD 데이터에서 빠르고 느린 동역학을 동시에 다루는 데 있어 LSTM과 Transformer의 효과성을 비교하기 위해.
• 모델 학습 성능을 향상시키기 위한 전처리 전략(예: 운동학적 뭉치화 및 재교차 제거)을 제안하기 위해.

제안 방법

• 알라닌 다펩타이드의 MD 궤적을 2차원 반응좌표(φ, ψ)에 투영하고, 고차원 상태 분할을 위해 RMSD 기반 클러스터링을 수행하였다.
• 경로 엔트로피와 동역학을 학습하기 위해 교차 엔트로피 손실를 사용하여 순차적 상태 전이에 대해 LSTM 및 Transformer 모델을 훈련시켰다.
• 시간 해상도 영향을 평가하기 위해 저장 간격(0.1 ps, 1 ps, 2 ps, 10 ps)을 다양하게 설정하였다.
• 빠른 동역학에서 유래하는 노이즈를 줄이고 상태공간 표현에서 느린 동역학 신호를 강화하기 위해 운동학적 뭉치화 및 재교차 제거를 적용하였다.
• Transformer에서는 이방향 및 단방향 인코딩을 사용하였으며, 삼각형 마스킹을 통해 단방향 동작을 시뮬레이션하였다.
• 기준 MD 데이터와의 비교를 위해 평균 제1통과 시간(MFPT), 통과 시간 분포, 그리고 함의된 시간스케일(ITS)을 사용하여 모델 성능을 평가하였다.

실험 결과

연구 질문

• RQ1 MD 궤적의 시간 해상도(저장 간격)가 LSTM 및 Transformer의 느린 동역학 학습 성능에 어떤 영향을 미치는가?
• RQ2 반응좌표의 차원 수가 딥러닝 모델이 희귀 사건을 포착하는 데 있어 어떤 역할을 하는가?
• RQ3 상태 분할 및 재교차 제거가 LSTM 및 Transformer의 느린 동역학 예측 정확도 향상에 기여하는가?
• RQ4 Transformer의 어텐션 메커니즘이 고차원 MD 데이터에서 장기적 동역학을 학습하는 데 있어 LSTM보다 성능 우위를 보이는가?
• RQ5 LSTM 및 Transformer 모델이 원시 MD 궤적에서 빠른 동역학과 느린 동역학을 자동으로 분리할 수 있는 정도는 어느 정도인가?

주요 결과

• 0.1 ps 저장 간격에서 LSTM 성능이 심각하게 저하되며, 이는 자유 에너지 표면가 정확하게 추정되더라도 빠른 동역학과 느린 동역학의 혼합으로 인한 것이다.
• 10 ps 저장 간격에서는 데이터 수량 부족으로 인해 모델 학습에 악영향을 미쳐 느린 동역학 예측 성능이 열악하다.
• 1 ps 및 2 ps 간격에서는 LSTM이 자유 에너지 표면과 동역학을 성공적으로 예측하지만, 0.1 ps에서는 노이즈로 인해, 10 ps에서는 데이터 부족으로 인해 성능이 저하된다.
• Transformer 모델은 양방향 어텐션을 사용하더라도, 시퀀스 길이 제한(100 토큰)과 장거리 의존성 처리 능력 부족으로 인해 LSTM보다 느린 동역학 예측에서 성능이 뛰어나지 않는다.
• 재교차 사건 제거 및 상태 분할에 대한 운동학적 뭉치화 적용은 고차원 RMSD 기반 데이터에서 LSTM 및 Transformer 양쪽의 성능을 크게 향상시킨다.
• 적절한 시간 해상도와 상태 분할의 조합이 필수적이다 — 딥러닝 모델은 이러한 전처리 없이 빠른 동역학과 느린 동역학을 자동으로 분리할 수 없다.