[Paper Review] AlphaFold Meets Flow Matching for Generating Protein Ensembles

Bowen Jing, Bonnie Berger+1 more|arXiv (Cornell University)|Feb 7, 2024

Advanced Proteomics Techniques and Applications74 citations

TL;DR

The paper repurposes AlphaFold and ESMFold into flow-matching based generative models to sample realistic protein structural ensembles, trained on PDB and MD-derived data, outperforming MSA subsampling in precision-diversity and accurately capturing MD-like distributions and ensemble observables.

ABSTRACT

The biological functions of proteins often depend on dynamic structural ensembles. In this work, we develop a flow-based generative modeling approach for learning and sampling the conformational landscapes of proteins. We repurpose highly accurate single-state predictors such as AlphaFold and ESMFold and fine-tune them under a custom flow matching framework to obtain sequence-conditoned generative models of protein structure called AlphaFlow and ESMFlow. When trained and evaluated on the PDB, our method provides a superior combination of precision and diversity compared to AlphaFold with MSA subsampling. When further trained on ensembles from all-atom MD, our method accurately captures conformational flexibility, positional distributions, and higher-order ensemble observables for unseen proteins. Moreover, our method can diversify a static PDB structure with faster wall-clock convergence to certain equilibrium properties than replicate MD trajectories, demonstrating its potential as a proxy for expensive physics-based simulations. Code is available at https://github.com/bjing2016/alphaflow.

Motivation & Objective

• Motivate and enable accurate sampling of conformational ensembles for proteins, addressing conformational heterogeneity beyond single-state predictions.
• Leverage high-accuracy single-state predictors (AlphaFold, ESMFold) as denoising models within a flow-matching framework.
• Enable training on ensembles beyond PDB, such as MD simulations, to capture dynamic conformations.
• Provide a principled training-time approach to diversity and precision in protein structure generation.
• Demonstrate that generated ensembles can serve as efficient proxies for costly MD simulations.

Proposed method

• Repurpose AlphaFold and ESMFold as denoising models conditioned on sequence inputs within a flow-matching framework.
• Define a conditional probability path in R^3 coordinates via harmonic-prior-based interpolation and learn the conditional expectation of data points using a Fréchet mean with FAPE loss.
• Project structure space to SE(3)-quotient to handle symmetry and missing residues, using RMSD-aligned linear interpolation in the quotient space.
• Train with a modified loss (squared FAPE) to align with the Fréchet-mean target in the quotient space.
• Fine-tune on PDB ensembles to achieve Alpha Flow and ESM Flow, then further fine-tune on MD ensembles (ATLAS) to obtain Alpha Flow-MD and ESM Flow-MD.
• Optional distillation to accelerate sampling and enable templates to improve convergence.

Experimental results

Research questions

• RQ1 Can AlphaFold/ESMFold be repurposed as generative denoising models under flow matching to sample protein ensembles conditioned on sequence?
• RQ2 Does training on PDB ensembles generalize to MD ensemble data and capture conformational flexibility and higher-order ensemble observables?
• RQ3 How does flow matching compare with MSA subsampling and traditional MD in precision, diversity, and distributional accuracy for protein ensembles?
• RQ4 Can the approach serve as an efficient proxy for replicate MD simulations and diversify solved structures with fewer computational resources?

Key findings

• Alpha Flow and ESM Flow achieve a better precision-diversity Pareto frontier than MSA subsampling on PDB test proteins.
• Alpha Flow ensembles show strong correlation with MD-derived flexibility metrics and improve on mean/variance of atomic positions versus MSA baselines.
• Alpha Flow methods accurately reproduce complex ensemble observables such as intermittent/transient contacts and solvent exposure patterns.
• When trained on ATLAS MD ensembles, Alpha Flow variants better approximate MD distributions across RMSF, W2 distances, and MI matrices than MSA subsampling.
• Sampling with Alpha Flow from a static structure converges faster to equilibrium properties than replicating MD trajectories, demonstrating potential as a proxy for expensive simulations.