[論文レビュー] 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

ひとこと要約

本稿は、高次元の分子動力学（MD）軌道から希少事象を学習する際のLSTMおよびTransformerモデルの限界を調査する。時間分解能と状態分割が重要であることが示された。コンformational空間を遷移状態に粗視化し、再通過イベントを除去することで、遅いダイナミクスの予測が著しく向上するが、両モデルはこのような前処理なしでは高速と遅いダイナミクスを分離できない。

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の有効性を比較すること。
• キネティックラumpingと再通過除去を含む前処理戦略を提案し、MDシミュレーションにおける希少事象の学習を向上させること。

提案手法

• アラニンジペプチドのMD軌道を2次元反応座標（φ, ψ）に投影し、RMSDに基づくクラスタリングを用いて高次元の状態分割を実施した。
• 交差エントロピー損失を用いて、パスエントロピーとキネティクスを学習するため、LSTMおよびTransformerモデルを順序付き状態遷移で学習させた。
• 時間分解能の影響を評価するため、保存間隔（0.1 ps, 1 ps, 2 ps, 10 ps）を変化させた。
• 高速ダイナミクスに起因するノイズを低減し、状態空間表現における遅いダイナミクスの信号を強化するために、キネティックラumpingと再通過イベントの除去を適用した。
• Transformerでは双方向および単方向符号化を用い、三角マスクを用いて単方向動作を模擬した。
• 平均第一到達時間（MFPT）、到達時間分布、および実効スケール（ITS）を、基準MDデータと比較してモデルを評価した。

実験結果

リサーチクエスチョン

• 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モデルはLSTMを上回る性能を示さない。これは、シーケンス長が限界（100トークン）であり、長距離依存性を効果的に扱えないためである。
• 再通過イベントの除去と状態パーティションへのキネティックラumpingの適用により、高次元RMSDベースのデータにおいてLSTMおよびTransformerの両方の性能が著しく向上する。
• 適切な時間分解能と状態分割の組み合わせが不可欠である。深層学習モデルは、このような前処理なしでは、高速と遅いダイナミクスを自動で分離できない。