High Accuracy Protein Structural Ensemble Prediction with Generative Flow Models
Abstract
Deep learning has enabled accurate prediction of protein structures, yet under native conditions proteins exist as thermodynamic ensembles of conformations that are critical to understanding their function. Deep generative models have recently emerged for sampling directly from this distribution, offering a cheaper alternative to molecular dynamics simulations. However, current state-of-the-art approaches largely rely on fine-tuning pre-trained folding models, are restricted to backbone-only structures, and require generating many samples to estimate any ensemble-level observable. We instead propose a flow-based generative model that captures the distribution of all-atom residue displacements from a reference structure, avoiding the need to re-learn the fold from scratch. Building on this, we introduce a method to predict ensemble-level observables directly from the model's shared latent space, without sampling any conformations, and an amortized two-stage architecture that substantially accelerates inference. Our method achieves state-of-the-art performance in generating accurate all-atom protein conformational ensembles across structurally diverse proteins.