DreamFold: A World Model to efficiently generate protein folding pathways in the latent space.

Abstract

While there is an abundance of static data for the structure of biological macromolecules, the data regarding their folding mechanisms and dynamics is scarce, posing a challenge to the training of AI models. The World Model approach, which has been very successful in robotics, can come in handy in this regard. From limited data, it can build a latent, approximated representation of the spatio-temporal folding environment that can be used to train downstream AI models. Here, we developed DreamFold, a World Model-based generative framework to perform biomolecular simulations. DreamFold encodes protein structures and dihedral angle moves into latent vectors using variational autoencoders. Then, it uses another neural network to rapidly predict the next latent structure after a latent move, allowing the model to develop its own understanding of protein dynamics. Finally, in a "hallucinated" latent environment, an agent learns, through evolutionary algorithms, a policy to drive folding simulations toward a target. Latent configurations can then be decoded back into atomistic structures. DreamFold can compute protein folding pathways four orders of magnitude faster than molecular dynamics (MD), up to ~30,000×. Furthermore, for the monomeric proteins tested here (up to 449 aa), the computational cost scales linearly with the number of atoms (N) as O(N) (in comparison, MD scales as O(NlogN)). We have validated our results against established MD benchmarks and other available experimental data. Finally, we show how DreamFold allows us to identify folding intermediates with cryptic binding sites that are therapeutically valuable. Overall, DreamFold facilitates the study of protein dynamics by generative AI with applications in structure-based drug discovery.