Protein Language Model Embeddings Improve Generalization of Implicit Transfer Operators

TL;DR

This paper introduces PLaTITO, a novel method that leverages protein language model embeddings to enhance the generalization and data efficiency of implicit transfer operators in molecular dynamics simulations, especially for out-of-distribution protein systems.

Contribution

The work demonstrates that incorporating protein language model embeddings into transfer operators significantly improves their transferability and out-of-distribution performance in molecular dynamics.

Findings

Coarse-grained TITO models are more data-efficient than Boltzmann Emulators.

Protein language model embeddings enhance out-of-distribution generalization.

PLaTITO achieves state-of-the-art results on protein equilibrium sampling benchmarks.

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.