Protein Language Model Embeddings Improve Generalization of Implicit Transfer Operators

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.

Understanding how proteins change structure over time is essential for understanding biological processes and developing new drugs. However, simulating protein dynamics on a computer is prohibitively slow, especially for larger systems, where observing a single protein folding event that happens in microseconds can require weeks of large-scale simulations. PLaTITO is a deep learning model trained to take much longer steps in time than traditional simulations, dramatically accelerating the overall process. To help the model generalize to unseen proteins, we incorporate knowledge from available AI systems that are trained on billions of protein sequences rich in evolutionary information. This allows PLaTITO to achieve state-of-the-art performance using roughly ten times less training data and compute than previous approaches. By making protein dynamics simulations far cheaper, this work could accelerate drug discovery and make protein simulation accessible to research groups with limited resources.