Protein Language Model Embeddings Improve Generalization of Implicit Transfer Operators

AI summary

Overview

Research area: Generative machine learning for molecular simulation, specifically learned surrogates of protein molecular dynamics (MD), combining flow matching, implicit transfer operators, and pretrained protein language model (pLM) representations.

Technical level: Advanced. The paper assumes familiarity with molecular dynamics, Boltzmann distributions, flow matching, and transfer operator learning.

Scope: The paper studies how auxiliary conditioning signals — especially pretrained protein language model embeddings — affect the data efficiency and out-of-distribution generalization of transferable implicit transfer operator (TITO) models trained on coarse-grained protein MD trajectories.

What This Paper Is About

Conventional molecular dynamics simulates proteins with femtosecond-scale integration steps, but biologically relevant processes such as folding take microseconds to seconds, making simulation prohibitively expensive. Generative molecular dynamics (GenMD) learns neural surrogates that sample molecular distributions directly, but existing methods need large collections of long MD trajectories and generalize poorly to unseen protein systems. This paper asks whether injecting pretrained protein language model embeddings (and other auxiliary signals) into transferable implicit transfer operators improves data efficiency and generalization to out-of-distribution proteins, without system-specific fine-tuning.

Key Contributions

1. Coarse-grained transferable implicit transfer operators (TITO). The authors develop TITO models for protein molecular dynamics that generalize to out-of-distribution protein systems without system-specific fine-tuning, trained from scratch on diverse off-equilibrium MD trajectories across multiple temperatures from the mdCATH dataset.

2. Conditioning on auxiliary information for data efficiency. They investigate the impact of conditioning on several auxiliary sources of information and introduce Protein-Language-aware TITO (PLaTITO), which achieves state-of-the-art performance on equilibrium sampling benchmarks while being trained with substantially less data and computational resources.

3. Temperature-dependent, non-Arrhenius kinetics. They show that the temperature-dependent folding and unfolding kinetics learned by PLaTITO are non-Arrhenius, consistent with complex rugged folding free energy landscapes.

4. Scaling and comparison to Boltzmann Emulators. They demonstrate that a scaled 19M-parameter variant, PLaTITO-19M, surpasses the BioEmu Boltzmann emulator across all reported equilibrium sampling metrics at a fraction of the training data and compute, and that TITO outperforms an architecturally identical Boltzmann Emulator (Emu) trained with the same data and compute.

Main Findings

• pLM embeddings improve out-of-distribution generalization. PLaTITO (conditioned on residue-level embeddings from the pretrained ESM Cambrian 300M pLM) improved over the base TITO model on all three reported equilibrium sampling metrics: MAE 0.949 ± 0.269 versus 1.068 ± 0.272, RMSE 1.228 ± 0.328 versus 1.382 ± 0.302, and Coverage 0.651 ± 0.151 versus 0.590 ± 0.111, at no additional test-time cost.

• Structure embeddings add a small, consistent gain. PLaTITO+Struct, which additionally conditions on embeddings from the pretrained Proteina 60M structure model, reached MAE 0.938 ± 0.321, RMSE 1.213 ± 0.348, and Coverage 0.655 ± 0.158, suggesting structure and language embeddings provide complementary information.

• LLM-derived annotations hurt performance. PLaTITO+Struct+LLM, which conditions on suitability and confidence annotations obtained by prompting the DeepSeek Reasoner LLM, degraded performance (MAE 1.066 ± 0.270, RMSE 1.346 ± 0.292, Coverage 0.570 ± 0.087). The authors suggest the prompting strategy or the available metadata were insufficient to provide useful conditioning signals.

• Scaling capacity and pLM size helps. PLaTITO-19M, a 19M-parameter variant using embeddings from the larger ESM Cambrian 6B pLM, reached MAE 0.824 ± 0.170, RMSE 1.099 ± 0.212, and Coverage 0.666 ± 0.136, improving further under the same training compute budget.

• State-of-the-art versus BioEmu at far lower cost. PLaTITO-19M outperformed BioEmu (MAE 1.110 ± 0.292, RMSE 1.389 ± 0.346, Coverage 0.594 ± 0.175) across all reported equilibrium sampling metrics, despite converging in approximately 1,100 GPU hours on a single NVIDIA A100 80GB GPU versus 9,216 GPU hours for BioEmu. BioEmu was additionally trained on 131k AFDB structures and 502k experimental ΔG measurements.

• TITO beats a matched Boltzmann Emulator. The Emu baseline, architecturally identical to TITO but trained to sample directly from p(x|S, T), scored MAE 1.305 ± 0.378, RMSE 1.639 ± 0.406, and Coverage 0.529 ± 0.112 — worse than TITO on all metrics under the same architecture, data, and compute budget. The authors interpret this as evidence that exploiting the auto-correlation structure of MD data lowers data and compute requirements.

• Qualitative reproduction of folding free energy landscapes. For a representative subset of three fast-folding proteins (Villin, WW domain, and A3D), PLaTITO-19M reproduced both folded and unfolded regions of the reference MD free-energy landscapes and consistently outperformed BioEmu. PLaTITO-19M also maintained the best performance against ConfDiff and Str2Str in Appendix D.5, and against various pLMs and model sizes in Appendices D.3 and D.4.

• Long-timescale dynamics are captured but timescales are underestimated. A 120 µs PLaTITO-19M trajectory for A3D showed repeated folding and unfolding. Estimated mean first passage times were faster than MD: folding 5.5 ± 0.9 µs versus 27 ± 8 µs, and unfolding 14.8 ± 1.2 µs versus 31 ± 9 µs. The authors attribute this direction of error to a variational principle showing imperfect approximations of MD systematically underestimate timescales.

• Non-Arrhenius temperature dependence. Across five temperatures from 320 K to 440 K, predicted folding and unfolding timescales deviated from simple Arrhenius behavior, consistent with prior experimental and computational studies of protein folding kinetics. Reference MD timescales at selected temperatures were consistently larger than PLaTITO-19M predictions.

• Partial success on cryptic binding pockets. Across four experimentally characterized cryptic pocket cases curated by Lewis et al. (2025), PLaTITO-19M sampled broad ensembles and, when initialized from the holo state, generated apo-like conformations. BioEmu failed to sample the apo state in three of four cases; PLaTITO-19M recovered the apo state in one of these difficult instances, failing in only two of the three baseline failure cases (systems Q29495 and Q58L87). The authors note PLaTITO-19M does not yet form the clearly metastable basins seen in BioEmu, with its limitation appearing to be energetic separation rather than mode coverage.

Methodology in Plain English

The authors reframe MD sampling as a generative modeling problem. They use only the Cα backbone coordinates of proteins (a coarse-grained representation, x_t ∈ ℝ^{3L} for L residues) and train a model to predict long-time transitions directly: given a conformation at time t, the amino-acid sequence S, the temperature T, and a time step Δt, the model generates a plausible conformation at time t + Δt.

Training uses conditional flow matching. A latent variable interpolates linearly between Gaussian noise and the future conformation, and a neural network learns to predict the velocity field that transports noise into the future state. Pairs of conformations separated by Δt ≫ τ (the femtosecond integration step) are sampled randomly from the same MD trajectory, so a single trajectory yields many training pairs at different time lags.

The architecture is a two-stage transformer design based on Proteina. One network produces a per-residue conditioning representation from the current coordinates, sequence, temperature, and time step. A second network predicts the velocity field from the noisy latent, the flow-matching time, and that conditioning representation. Both use conditioned and biased multi-head self-attention with pair representations as attention biases.

To test whether outside knowledge helps, the authors augment the conditioning with three extra signal types: (I) residue-level sequence embeddings from a pretrained pLM such as ESM Cambrian, which can be precomputed offline at no training overhead; (II) residue-level structure embeddings from a pretrained structure model such as Proteina, which must be computed online for each generated backbone; and (III) LLM-derived annotations about whether a protein domain is suitable for simulation in isolation, generated by prompting the DeepSeek Reasoner LLM with curated metadata.

Training data comes from mdCATH, restricted to domains of at most 200 residues (4,482 domains). To enforce a strict train-test split, any protein with at least 40% sequence similarity to a test protein over an alignment of at least 20 residues was removed, leaving 4,471 domains and approximately 56 ms of aggregate MD simulation time.

Evaluation uses the fast-folding proteins benchmark: 12 systems of 10 to 80 residues simulated with the CHARMM22* force field in explicit solvent, with configurations saved every 200 ps. The data is proprietary but available upon request for research purposes. Free-energy landscapes are compared after projecting onto the four slowest TICA components (10 ns lag time) computed from pairwise Cα–Cα distances, binning, and converting histograms to free energies via G_i = −k_B T ln p_i. Metrics are mean absolute error, root mean squared error, and Coverage (the fraction of reference MD bins that the model samples). Sampling used 1,000 trajectories initialized from an unfolded state with 1,000 iterative roll-outs at Δt = 1 ns (1 µs total), keeping the final 10 ns of each trajectory for 10,000 generated samples per system.

Why This Matters

Impact on research. The paper reframes the trade-off in generative molecular dynamics: rather than scaling data and compute the way Boltzmann Emulators do, it shows that exploiting the temporal auto-correlation structure of MD data and injecting pretrained biological representations can deliver better out-of-distribution generalization at roughly an order of magnitude lower cost. It also provides a rare test of whether learned surrogates recover physically meaningful kinetics — non-Arrhenius temperature dependence — rather than only matching equilibrium densities.

Real-world applications (as suggested by the paper's scope and benchmarks):

• Drug discovery targeting cryptic binding pockets, which the paper evaluates on four experimentally characterized systems.
• Protein engineering and design, where rapid estimates of conformational ensembles and stability are useful.
• Biophysical characterization of fast-folding proteins and folding-related disease mechanisms.
• Replacing or accelerating expensive MD workflows for computing equilibrium observables such as free-energy differences.

Industry relevance. Pharmaceutical and biotechnology companies invest heavily in MD for lead optimization and target assessment. A surrogate that generalizes to unseen proteins without system-specific fine-tuning, runs on commodity GPUs, and was trained on roughly 56 ms of aggregate MD time is directly relevant to lowering simulation costs. The explicit comparison against BioEmu — a model trained with far more data and 9,216 GPU hours — frames the result as a compute-and-data efficiency argument that matters for any organization paying for simulation clusters.

Future Directions

1. Scaling model capacity and data. The authors state that PLaTITO-19M still fails to capture certain metastable phases in the fast-folders and the apo structures of some cryptic pocket systems, and systematically underestimates timescales, and that scaling both model capacity and data is needed.

2. Moving to all-atom representations. The current coarse-grained Cα-only model ignores side-chain rearrangements, which the authors note may be critical for ligand specificity or allosteric regulation. Extending to all-atom representations is flagged as future work.

3. Extending to complexes and other biomolecules. The conclusion (truncated in the provided content) begins to mention including data with complexes, other proteins, and biomolecules, alongside the limitation that generalization across chemical space and thermodynamic condition will likely remain limited given finite training data.

4. Fixing the conditioning signals that did not help. The LLM-derived annotations degraded performance, leaving open whether better prompting strategies or richer metadata could make that conditioning channel useful. Similarly, the absence of formal guarantees — unbiased dynamics, detailed balance, semigroup self-consistency (Chapman-Kolmogorov), and stability — remains an open theoretical gap for this class of surrogate.

Target Audience

Machine learning researchers working on generative models for scientific data, computational chemists and biophysicists interested in learned MD surrogates, and practitioners in structure-based drug discovery evaluating whether transferable neural samplers can replace or supplement conventional MD. Readers need working knowledge of flow matching and Markov state or transfer operator concepts to follow the methods and metrics; the high-level findings and benchmark table are accessible to a broader computational biology audience.

Authors’ abstract

Molecular dynamics (MD) is a central computational tool in physics, chemistry, and biology, enabling quantitative prediction of experimental observables as expectations over high-dimensional molecular distributions such as Boltzmann distributions and transition densities. However, conventional MD is fundamentally limited by the high computational cost required to generate independent samples. Generative molecular dynamics (GenMD) has recently emerged as an alternative, learning surrogates of molecular distributions either from data or through interaction with energy models. While these methods enable efficient sampling, their transferability across molecular systems is often limited. In this work, we show that incorporating auxiliary sources of information can improve the data efficiency and generalization of transferable implicit transfer operators (TITO) for molecular dynamics. We find that coarse-grained TITO models are substantially more data-efficient than Boltzmann Emulators, and that incorporating protein language model (pLM) embeddings further improves out-of-distribution generalization. Our approach, PLaTITO, achieves state-of-the-art performance on equilibrium sampling benchmarks for out-of-distribution protein systems, including fast-folding proteins. We further study the impact of additional conditioning signals such as structural embeddings, temperature, and large-language-model-derived embeddings on model performance.