ProtScape: A molecular structure and energy-aware representation for protein conformation generation

Understanding protein dynamics on microsecond to millisecond scales remains challenging, necessitating improved methodologies for analyzing and representing protein trajectories. ProtSCAPE introduces a novel deep learning architecture that leverages geometric scattering transforms alongside transformer-based attention mechanisms to capture protein dynamics from molecular dynamics (MD) simulations.

The gist: ProtSCAPE is a deep neural network that learns latent representations of MD trajectories through a novel transformer-based encoder-decoder architecture, which combines multiscale representation of the geometric scattering transform with dual attention structures to generate latent representations jointly informed by protein structure and its evolution.

ProtSCAPE Architecture and Components

ProtSCAPE consists of three primary parts: a learnable geometric scattering module, two transformer models (one for residues and one for amino acids), and a regularized autoencoder. The geometric scattering module uses a cascade of wavelet filters to identify meaningful multi-scale structural information of the input protein graphs. This is implemented using generalized diffusion wavelets, where the scales are chosen via a differentiable selection matrix.

Dual Attention Mechanisms

The architecture integrates dual attention structures to focus on different aspects of the protein. First, one uses attention across nodes (residues), and the other uses attention across signals (amino acids). The paper notes that The attention scores correlate with residue flexibility, allowing for the identification of the most flexible regions in the MD trajectory, which enables suggestions for flexible residues as candidates for generating protein mutants.

Latent Representation and Regression Heads

The embedding from the transformer modules is fed into an encoder network to compute a latent representation, denoted as z = E(p). This latent representation is controlled by two regression networks (N and M) which are used to predict specific properties:

1. Network N predicts the time point t corresponding to each z, utilizing a mean squared error loss.

2. Network M aims to predict the pairwise distances and the dihedral angle between residues, employing a mean squared error loss.

Training and Loss Formulation

The total loss function is defined as: L = αLt + βLs + (1 − α + β)Lc. This formulation ensures that the learned latent representations are temporally coherent, structure-aware, and capable of reconstructing the underlying protein structure. The model also incorporates a learnable node embedding strategy to manage memory requirements for large proteins.

Generalizability and Performance

ProtSCAPE demonstrates strong generalization capabilities:

1. It generalizes effectively from short to long trajectories.

2. It generalizes well to mutant trajectories when trained on wildtype trajectories, allowing the model to understand how mutants sample conformational space.

Experimental results show that ProtSCAPE excels in predicting pairwise distances and dihedral angle differences compared to other graph-based methods (GNNs), achieving the best, or second best, performance across all five proteins of interest. Furthermore, the latent representations exhibit low Dirichlet energy, indicating they are smooth and temporally organized, which is crucial for capturing stochastic switching between meta-stable conformations. The model was shown to uncover stochastic switching between two meta-stable conformations in the GB3 protein trajectory. Additionally, ProtSCAPE successfully reconstructs intermediate structures transitioning between open and closed states in the MurD protein, consistent with experimental findings.

Stability and Robustness

Theoretical proofs establish that the generalized geometric scattering transform is stable to small perturbations of the graph structure. Theorem B.1 proves that the generalized geometric scattering U transform is stable to small deformations of the graph structure, showing that it captures intrinsic structure rather than superficial artifacts like vertex ordering, even when considering minor perturbations in adjacency matrices. This stability is quantified through bounds involving the diffusion distance and spectral properties of the graphs. The architecture's robustness is further confirmed by its ability to maintain organized latent spaces when trained on short trajectories and evaluated on long ones, and its ability to align latent representations for wild-type and mutant proteins.

Interpolation Capabilities

ProtSCAPE can be used for structural interpolation. By constructing a linear interpolant between the cluster centroids of different conformations in the latent space, the model can decode their intermediate structures transitioning from open to close states, revealing hinge-like mechanisms consistent with experimental data. The attention scores learned by the transformer further provide interpretability, indicating that residues with high attention scores correspond to more flexible regions of the protein that play a pivotal role in conformational changes. This capability allows for visualizing the transition pathways between distinct protein conformations.

Reconstruction Metrics

In an additional variant, a model trained to directly reconstruct residue center of mass coordinates from the latent representation was also evaluated. ProtSCAPE achieved overall the top performing model with best, or second best, performance across all five proteins of interest when compared against MDTraj-derived features and Ramachandran plots using metrics like Spearman Correlation Coefficient (SCC) and Pearson Correlation Coefficient (PCC). The reconstruction metrics show that ProtSCAPE's performance is superior to baselines in capturing the spatial and angular relationships in protein conformations.

As a fastidious researcher, I have analyzed the provided paper on ProtSCAPE for mapping protein conformations in molecular dynamics. The core innovation lies in combining geometric scattering transforms with dual attention mechanisms within a transformer-based architecture to learn latent representations of MD trajectories.

Here are the specific improvements and capabilities this system enables for AI systems:

)Improved AI Systems & Specific Capabilities

The ProtSCAPE architecture, as described, can be integrated into downstream AI systems for several high-impact applications:

1. Dynamics Prediction and Phase Transition Mapping (Stochastic Switching):

A system trained on ProtSCAPE latent representations can identify and predict phase transitions between distinct protein states (e.g., open vs. closed conformations) by analyzing the clustering of latent embeddings in the PHATE visualization space (as seen in Figure 2 and 5).

-Capability: Predict the transition probability or time required for a protein to switch between metastable conformational states during a simulation, providing insight into functional mechanisms like allosteric regulation.

2. De Novo Protein Design via Targeted Mutagenesis (WL2M Generalization):

The model's ability to generalize from wild-type (WT) training data to mutant trajectories is a critical feature. By leveraging the attention mechanism that highlights flexible residues, the system can suggest precise locations for functional modifications.

-Capability: Automatically generate and predict the structural outcome of missense mutations (WL2M generalization). Specifically, it can rank potential mutation sites based on predicted impact on conformational space sampling (e.g., identifying residues whose substitution leads to a significant deviation in the latent embedding).

3. Structure Reconstruction from Dynamics (Latent Interpolation):

The architecture is trained to reconstruct pairwise residue distances and dihedral angles from a latent representation, allowing for the interpolation of protein conformations between two learned states.

-Capability: Generate novel intermediate protein structures that smoothly transition between two known functional states (e.g., an open and closed state) by decoding the latent space, which is crucial for understanding hinge mechanisms or conformational pathways.

4. High-Fidelity Conformational Search (Distance/Angle Prediction):

The dual regression heads predict pairwise distances and dihedral angles with high precision (Table 2 shows ProtSCAPE outperforming GNNs in these metrics).

-Capability: Perform rapid, high-accuracy distance and angle predictions on protein structures derived from MD simulations or experimental data, which can be used as a fast filter for identifying functionally relevant structural changes.

5. Feature Extraction for Dynamics (Attention Scoring):

The dual attention mechanism explicitly correlates attention scores with residue flexibility, allowing the system to pinpoint dynamic hotspots.

-Capability: Automatically identify the hotspots of conformational change within a protein structure, guiding researchers toward regions that are most critical for function or binding events during dynamic processes.

6. Robustness to Data Complexity (S2L Generalization):

The model generalizes from short MD trajectories (e.g., 1000 frames) to long trajectories (e.g., 10,000 frames).

-Capability: Extract meaningful conformational information and latent structure representations even when the input MD simulation data is truncated or incomplete, making it robust for analyzing long-term biological processes where full simulations are computationally prohibitive.

Molecular dynamics (MD) simulations are a principled but computationally expensive approach for studying protein conformational variability, making it challenging to generate large ensembles of structures or characterize transitions between metastable conformations. AI methods for upsampling MD simulations have been developed recently but struggle due to difficulties of sampling complex, high-dimensional molecular distributions. One avenue these methods overlook is to learn a latent space where such sampling becomes easier. To address this, we introduce ProtScape, a generative geometric deep learning framework that learns structure and energy-aware representations of protein conformational landscapes from MD simulations. ProtScape represents protein conformations using an equivariant graph neural network and a multiscale deep wavelet transform that captures local geometric interactions and nonlocal collective motions. This latent representation is dual-organized: 1) by structure captured by wavelet transform layers, via reconstruction error; 2) by energy using a Laplacian energy smoothness penalty. This structured latent space enables meaningful generation and exploration of protein conformations. ProtScape supports multiple generative modes: 1) ensemble generation, which upsamples ensembles given limited MD trajectories flow matching from noise to the organized manifold of conformations, 2) minimum-energy path generation between two high-energy conformations, guided by energy using a nudged elastic band, and 3) energy-descent, which generates trajectories toward lower energies from a high-energy conformation using gradient descent, thereby hypothesizing folding or other stabilizing trajectories. Organizing a latent space by structure and energy lets one representation support ensemble sampling, minimum-energy path finding and energy-guided descent on the same conformational landscape.