Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

AlphaFold Meets Flow Matching for Generating Protein Ensembles

Bowen Jing, Bonnie Berger, Tommi Jaakkola

TL;DR

The paper tackles protein conformational heterogeneity by converting single-state predictors like AlphaFold and ESMFold into flow-matching, sequence-conditioned generators to sample structural ensembles. It introduces AlphaFlow and ESMFlow, which use a harmonic prior and an SE(3) quotient-space formulation with Fréchet-mean supervision to produce diverse yet accurate ensembles. Empirical results show superior precision-diversity tradeoffs on PDB ensembles compared with MSA subsampling, and faithful reproduction of MD-derived distributions and higher-order observables when trained on ATLAS MD ensembles, with faster convergence than explicit MD in some cases. The approach also yields efficient, distillable inference, offering a practical proxy for expensive physics-based simulations and broad applicability to structural biology tasks.

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.

AlphaFold Meets Flow Matching for Generating Protein Ensembles

TL;DR

The paper tackles protein conformational heterogeneity by converting single-state predictors like AlphaFold and ESMFold into flow-matching, sequence-conditioned generators to sample structural ensembles. It introduces AlphaFlow and ESMFlow, which use a harmonic prior and an SE(3) quotient-space formulation with Fréchet-mean supervision to produce diverse yet accurate ensembles. Empirical results show superior precision-diversity tradeoffs on PDB ensembles compared with MSA subsampling, and faithful reproduction of MD-derived distributions and higher-order observables when trained on ATLAS MD ensembles, with faster convergence than explicit MD in some cases. The approach also yields efficient, distillable inference, offering a practical proxy for expensive physics-based simulations and broad applicability to structural biology tasks.

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.
Paper Structure (36 sections, 1 theorem, 27 equations, 13 figures, 5 tables, 3 algorithms)

This paper contains 36 sections, 1 theorem, 27 equations, 13 figures, 5 tables, 3 algorithms.

Table of Contents

  1. Introduction
  2. Background
  3. Method
  4. AlphaFold as a Denoising Model
  5. Flow Matching for Protein Ensembles
  6. Comparison with Diffusion
  7. Experiments
  8. Training Regimen
  9. PDB Ensembles
  10. Molecular Dynamics Ensembles
  11. Conclusion
  12. Method Details
  13. Input Embedding Module
  14. Flow Matching on Protein Ensembles
  15. Unsuitability of MSE Loss
  16. ...and 21 more sections

Key Result

Proposition 1.1

Let $\mathbf{x}_1 \in \mathbb{R}^N$ and $\mathbf{x}_1 [i:i+M] \in \mathbb{R}^M$ be a crop of $\mathbf{x}_1$ of length $M \le N$ and define $p_t^{(M)}, p_t^{(N)}$ to be the conditional probability paths in dimensionalities $N, M$. Then for any $t, \tilde{\mathbf{x}} \in \mathbb{R}^M$, $p_t^{(N)}(\mat

Figures (13)

  • Figure 1: Conceptual overview of AlphaFlow / ESMFlow. (A) Samples are drawn from a harmonic (polymer-like) prior. (B) The sample is progressively refined or denoised under a flow field controlled by the structure prediction model (AlphaFold or ESMFold). (C) At each step, the denoised structure prediction parameterizes the direction of the flow and we interpolate the current sample towards it. (D) The final prediction is a sample from the learned distribution of structures.
  • Figure 2: AlphaFold as a denoising model. Just as (diffusion-based) text-to-image generative models are simply neural networks that denoise images (with text input), a modified AlphaFold that ingests noisy structures and predicts clean structures (with sequence input) immediately provides a sequence-to-structure generative model---when trained under an appropriate framework.
  • Figure 3: Evaluation on PDB ensembles---precision-diversity (left) and precision-recall (right) curves for all benchmarked methods (median taken over 100 test targets). The MSA subsampling curve is traced by reducing MSA depth (max 512, min 48) and joins to AlphaFold as they share the same weights (AlphaFold by default subsamples MSAs to a maximum depth of 1024 and thus has nonzero diveristy, unlike ESMFold). The AlphaFlow / ESMFlow curves are traced by truncating the initial steps of flow matching (described in Appendix \ref{['app:training_inference']}). Distilled models are marked by $\blacktriangle$. Tabular data is shown in Appendix \ref{['app:pdb_results']}, Table \ref{['tab:pdb_results']}
  • Figure 4: Efficiency of AlphaFlow vs replicate MD simulations. AlphaFlow+Templates with varying number of samples with distillation (green) and without distillation (orange); MD with varying trajectory lengths in blue. See Appendix \ref{['app:evaluation']} for further experimental details and Appendix \ref{['app:md_results']} for further results.
  • Figure 5: MD evaluations visualized. (A) Ensembles of PDB ID 6uof_A (transcriptional regulator from Streptococcus pneumoniae) from MD, AlphaFlow, and MSA subsampling (depth 48), with C$\alpha$ RMSF by residue index shown in insets. (B) $1-CDF$ of the distribution of (unsigned) cosine similarities between the top principal components of the predicted ensemble versus the MD ensemble. (C) Solvent exposure mutual information matrices computed from the ground truth MD ensemble and AlphaFlow ensemble for target PDB ID 7bwf_B (antitoxin from Staphylococcus aureus). (D, E, F) Deviations from the crystal structure in the MD simulation, corresponding to ensemble observables, which are correctly sampled by AlphaFlow. The probability of occurence in each ensembles is shown. (D) solvent exposure of a buried residue in PDB ID 6oz1_A (carboxylate reductase from M. chelonae). (E) association of a transient residue contact in PDB ID 6q9c_A (NADH-quinone oxidoreductase subunit E from Aquifex aeolicus). (F) dissociation of a weak residue contact in PDB ID 6d7y_B (immune protein from Enterobacter cloacae). Additional examples in Appendix \ref{['app:md_results']}$r$: Pearson correlation; $\rho$: Spearman correlation.
  • ...and 8 more figures

Theorems & Definitions (2)

  • Proposition 1.1
  • proof