Protein Language Model Embeddings Improve Generalization of Implicit Transfer Operators

Technical Summary: Protein Language Model Embeddings Improve Generalization of Implicit Transfer Operators

Problem Statement
Molecular dynamics (MD) simulations are essential for predicting biophysical observables by sampling from high-dimensional molecular distributions, such as Boltzmann equilibrium distributions and transition densities. However, conventional MD is computationally prohibitive because it requires femtosecond-scale time steps to simulate processes that occur on microsecond-to-second timescales. While Generative Molecular Dynamics (GenMD) methods, including Boltzmann Generators/Emulators and Implicit Transfer Operators (ITO), offer efficient sampling surrogates, they typically suffer from limited data efficiency and poor generalization to unseen molecular systems. Existing approaches often require extensive collections of long MD trajectories, restricting their applicability to out-of-distribution (OOD) protein systems.

Methodology
The authors introduce PLaTITO (Protein-Language-aware Transferable Implicit Transfer Operator), a coarse-grained model designed to approximate long-timescale transition densities $p(x_{t+\Delta t} | x_t, S, T, \Delta t)$ for protein molecular dynamics. The methodology builds upon the Implicit Transfer Operator framework using Flow Matching.

1. Coarse-Grained Representation: The model operates on C$\alpha$ backbone coordinates ($x_t \in \mathbb{R}^{3L}$), ignoring side chains to enable scaling to larger protein systems.
2. Flow Matching for ITO: The model learns a velocity field to transform Gaussian noise into future conformations $x_{t+\Delta t}$ given the current state $x_t$. It minimizes a conditional flow-matching loss using linear interpolants between noise and the target conformation.
3. Architecture: A two-stage non-equivariant Transformer architecture is employed:
   • Conditioning Network ($f_c$): Computes per-residue representations conditioned on the current backbone coordinates, amino-acid sequence ($S$), temperature ($T$), and time step ($\Delta t$).
   • Velocity Network ($f_v$): Predicts the velocity field at flow-matching time $s$ using the latent flow variable and the conditioning representations.
4. Auxiliary Conditioning Signals: To improve data efficiency and generalization, the authors incorporate three types of auxiliary information:
   • Sequence Embeddings: Pre-trained embeddings from Protein Language Models (pLMs), specifically ESM Cambrian, are integrated into the conditioning network.
   • Structure Embeddings: Pre-trained structure-aware embeddings from the Proteina model are used to provide geometric context.
   • LLM-Derived Annotations: Metadata regarding protein stability and simulation suitability is extracted using Large Language Models (DeepSeek Reasoner) to flag potentially unphysical simulations lacking native context.

Key Contributions

1. Development of Coarse-Grained TITO: The authors develop transferable implicit transfer operators that generalize to OOD protein systems without system-specific fine-tuning, trained on diverse off-equilibrium MD trajectories across multiple temperatures.
2. Integration of Auxiliary Information: The paper demonstrates that incorporating pLM embeddings (PLaTITO) significantly boosts data efficiency and performance compared to base TITO models and Boltzmann Emulators.
3. Non-Arrhenius Kinetics: The authors show that the temperature-dependent kinetics learned by PLaTITO exhibit non-Arrhenius behavior, consistent with the complex, rugged free energy landscapes of protein folding.

Results
The models were trained on the mdCATH dataset, comprising ~4,482 protein domains with approximately 56 ms of aggregate MD simulation time. Performance was evaluated on equilibrium sampling benchmarks for fast-folding proteins and cryptic binding pocket formation.

• Equilibrium Sampling: PLaTITO-19M (a 19M-parameter variant using ESM Cambrian 6B embeddings) achieved state-of-the-art performance on equilibrium sampling metrics (MAE, RMSE, Coverage) for OOD fast-folding proteins. It outperformed BioEmu, a recent Boltzmann emulator, across all metrics despite using substantially less training data (56 ms vs. larger datasets for BioEmu) and lower computational cost (~1,100 GPU hours vs. ~9,216 GPU hours for BioEmu).
• Data Efficiency: Coarse-grained TITO models were found to be substantially more data-efficient than Boltzmann Emulators. Incorporating pLM embeddings reduced MAE and RMSE while increasing coverage compared to the base TITO model.
• Kinetics and Dynamics: PLaTITO-19M successfully reproduced folding-unfolding transition kinetics. The model predicted non-Arrhenius temperature-dependent folding and unfolding rates, aligning with physical expectations for rugged energy landscapes. However, the model systematically underestimated timescales (e.g., mean first passage times) compared to MD references, consistent with variational principles for imperfect approximations.
• Cryptic Binding Pockets: In evaluating the formation of cryptic binding pockets, PLaTITO-19M improved coverage over baselines, successfully recovering apo states in difficult instances where BioEmu failed. However, it did not form clearly metastable basins in all cases, suggesting limitations in energetic separation.
• Ablation Studies:
  • Structure embeddings provided modest improvements.
  • LLM-derived annotations did not improve performance, suggesting the prompting strategy or metadata was insufficient.
  • Scaling model capacity (up to 19M parameters) and using larger pLM embeddings yielded complementary benefits, with performance plateauing or degrading slightly at 100M parameters without additional data.

Significance
The paper claims that PLaTITO demonstrates that leveraging the auto-correlation structure of MD data via Implicit Transfer Operators, combined with auxiliary conditioning from pre-trained models, can lower the data and compute requirements for learning MD surrogates. This approach allows for the training of data-efficient models that generalize to unseen systems, offering a computationally cheaper alternative to Boltzmann Emulators for tasks such as conformational sampling and folding kinetics analysis. The authors note that while the model captures qualitative kinetics and equilibrium distributions effectively, further scaling of data and model capacity, as well as extension to all-atom representations, are required to fully capture meta-stable phases and accurate timescales.