Neural Network Distillation for Protein Folding Dynamics: Reaction Coordinate Extraction from LSTM to Transformer

This article is an undergraduate thesis project at the University of Leeds. Its core research direction is to explore whether LSTM and Transformer neural networks can extract physically meaningful reaction coordinates from protein folding dynamics, and evaluate their effectiveness using committor theory and Zq validation methods. Reaction coordinates are key low-dimensional descriptions for understanding protein folding mechanisms. This study combines machine learning and biophysical theory to provide a new perspective for solving the protein folding problem.

Protein folding is a core problem in molecular biology. Linear polypeptide chains need to fold into three-dimensional structures within specific time scales, and their mechanisms are crucial for drug design, disease treatment, etc. Traditional molecular dynamics simulations generate high-dimensional data that is difficult to analyze intuitively, so reaction coordinates (RC) have emerged as low-dimensional description tools to capture key features of the folding process. Traditional RCs rely on expert experience for design and may miss important information, while machine learning methods can automatically learn hidden patterns from raw trajectories.

This study compares two sequence modeling architectures:

1. LSTM: Processes long-range dependencies in conformation sequences through gating units (input gate, forget gate, output gate) to capture key transition points in the folding process. The input is conformation features (e.g., atomic coordinates, dihedral angles), and the output is the learned RC values.
2. Transformer: Uses self-attention mechanisms to process sequences in parallel, capturing dependencies at any position. Multi-head attention can characterize the folding process from multiple perspectives (e.g., secondary structure formation, hydrophobic core collapse), overcoming the gradient vanishing problem of LSTM.

To verify the physical meaning and predictive ability of RCs, two methods are used:

• Committor theory: A statistical mechanics tool that defines the probability of a conformation folding into the folded state first. A high-quality RC should be highly correlated with the committor, avoiding expensive transition path sampling and learning directly from equilibrium trajectories.
• Zq validation: Quantifies the predictive ability of RCs, measuring the accuracy of current RC values in predicting the distribution of the system's future q-step conformations. The closer Zq is to 1, the stronger the predictive ability, providing an objective comparison standard.

The term 'distillation' in the project has two meanings:

1. Information distillation: Compresses low-dimensional RCs from high-dimensional molecular dynamics trajectories, retaining key folding information and discarding irrelevant details.
2. Model distillation: Transfers knowledge from complex models (e.g., deep Transformers) to simple models (e.g., shallow LSTMs), reducing computational costs while maintaining RC quality.

This study reflects the cutting-edge trend in computational biology: the combination of deep learning and physical theory. Its applications include:

• Drug design: Identifying drug targets and intervention strategies;
• Disease research: Revealing the molecular mechanisms of diseases related to protein misfolding (e.g., Alzheimer's disease);
• Synthetic biology: Guiding the design of new proteins.

Challenges faced by the research and future directions:

• Data requirements: Training requires a large amount of long-time simulation data, which can be alleviated through transfer learning/pretraining;
• Interpretability: Need to use explainable AI techniques (e.g., attention visualization) to understand the conformation features focused on by the network;
• Generalization ability: Explore the transferability of models between different proteins, and use local structural commonalities (e.g., α-helices) to improve generalization.

This project integrates machine learning and molecular biophysics, extracting reaction coordinates through LSTM and Transformer and verifying their effectiveness, providing a new tool for understanding protein folding mechanisms. With the advancement of computing power and algorithms, data-driven RC learning methods will play a more important role in protein science. The ultimate goal is to achieve 'explainable AI' and reveal the physical principles behind folding.