PaperlayerView exact revision on arXiv

Source-linked AI summary

AlphaFold Meets Flow Matching for Generating Protein Ensembles

Bowen Jing, Bonnie Berger, Tommi Jaakkola

arXiv:2402.04845v2q-bio.BMcs.LG

TL;DR

Protein functions depend on dynamic structural ensembles, whereas accurate predictors such as AlphaFold primarily model single states. This work fine-tunes AlphaFold and ESMFold with custom flow matching to generate sequence-conditioned ensembles, achieving improved precision-diversity trade-offs and accurate modeling of MD-derived flexibility and distributions.

  • Problem

    Accurate single-state predictors do not account for the conformational heterogeneity underlying many protein functions.

  • Method

    The authors repurpose AlphaFold and ESMFold as sequence-conditioned generative models using a custom flow-matching framework and train them on protein structure ensembles.

  • Results

    AlphaFLOW and ESMFLOW surpass MSA-subsampling baselines in PDB precision-diversity trade-offs and accurately reproduce flexibility, positional distributions, and higher-order MD ensemble observables.

  • Takeaways & Limitations

    The models provide a potential proxy for expensive molecular-dynamics simulations, retaining ensemble quality with fewer samples and faster convergence for certain equilibrium properties.

  • Takeaways & Limitations

    Prior direct ensemble models had not shown convincing validation or comparisons with MSA subsampling on PDB test sets.

Abstract

from arXiv · show

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.

1. Introduction

Proteins rely on dynamic structural ensembles, but accurate single-state predictors and MSA-based diversification have important limitations. AlphaFLOW and ESMFLOW repurpose AlphaFold and ESMFold with custom flow matching to generate sequence-conditioned protein ensembles.

  • Motivation: Protein functions often depend on distinct structural states, collective motions, and disordered fluctuations.These ensemble properties are relevant to transporters, channels, enzymes, and molecular interactions.
  • Motivation: AlphaFold excels at single-state modeling but fails to account for conformational heterogeneity.
  • Existing limitations: MSA subsampling diversifies AlphaFold predictions but cannot generalize to protein-language-model predictors such as ESMFold.
  • Approach: AlphaFLOW and ESMFLOW combine AlphaFold or ESMFold with flow matching to model and sample protein conformational landscapes.The approach applies to PLM-based predictors and can be trained or fine-tuned on arbitrary ensembles.
  • Results: On PDB and molecular-dynamics ensemble evaluations, the models improve the combination of precision, diversity, flexibility modeling, positional distributions, and higher-order observables relative to MSA baselines.

2. Background

Protein structure prediction has progressed from single-structure models and MSA-based diversification toward direct generative modeling of ensembles. Flow matching provides a framework for learning a vector field that transports a prior distribution to protein structures, while prior ensemble models lacked convincing validation against MSA subsampling.

  • Protein structure prediction: AlphaFold predicts all-atom coordinates from sequence, MSA, and optionally a template, while ESMFold replaces MSA information with protein-language-model embeddings.
  • Existing ensemble methods: MSA subsampling became a standard way to obtain multiple AlphaFold predictions representing conformational heterogeneity.
  • Generative ensemble models: Direct ensemble generators such as EigenFold, Distributional Graphormer, and SENS use diffusion-based or local generative approaches.
  • Open gap: Prior ensemble models had not shown convincing validation or comparisons with MSA subsampling on PDB test sets.
  • Flow matching: Flow matching defines a conditional probability path from a shared prior distribution to an approximate data point and learns its marginal vector field with a neural network.
  • Flow matching: At convergence, the learned vector field forms a neural ODE that evolves the prior distribution into the data distribution.

3. Method

The method repurposes AlphaFold and ESMFold as denoising networks within a protein-specific flow-matching framework. It interpolates polymer-like prior structures toward data structures while handling SE(3) invariance, missing residues, and all-atom supervision.

  • Objective: The objective is to model the distribution of 3D protein structures conditioned on an amino-acid sequence.
  • AlphaFold as a denoiser: AlphaFold can be repurposed nearly out of the box as a generative model by embedding it in an iterative denoising framework.
  • Flow matching: Flow matching linearly interpolates between a harmonic-prior sample and a data structure along a conditional probability path.
  • AlphaFold as a denoiser: The denoising predictor receives a noisy structure and time embedding, then predicts a clean structure conditioned on sequence and, for AlphaFold, MSA information.
  • Protein-specific framework: The harmonic prior keeps intermediate inputs polymer-like and physically plausible.
  • Protein-specific framework: The framework uses SE(3)-quotient structures, RMSD-aligned interpolation, and squared FAPE while learning the distribution technically over β-carbon coordinates.
  • Comparison with diffusion: Flow matching imposes the prior at t = 0 for every dimensionality and handles missing residues by omitting them from interpolation.

4. Experiments

AlphaFLOW and ESMFLOW are fine-tuned on PDB structures and then ATLAS molecular-dynamics ensembles to generate diverse protein conformations. Across PDB and MD evaluations, they improve the precision–diversity trade-off and reproduce flexibility, positional distributions, and ensemble observables, while AlphaFLOW can be more computationally efficient than short replicate MD simulations.

  • Training regimen: AlphaFLOW and ESMFLOW fine-tune AlphaFold and ESMFold with flow matching on PDB structures, then continue training on ATLAS all-atom MD ensembles.The PDB stage uses 1.28M and 720k examples, while ATLAS provides 1265/39/82 train/validation/test ensembles.
  • PDB ensembles: AlphaFLOW traces a substantially superior precision-diversity Pareto frontier than MSA subsampling on 100 heterogeneous PDB test proteins.Both methods increase diversity relative to baseline AlphaFold, but AlphaFLOW remains clustered around ground-truth conformations while reaching comparable or greater diversity.
  • PDB ensembles: ESMFLOW injects substantial diversity relative to deterministic ESMFold and increases recall at little to no cost in precision.ESMFold-family methods retain lower precision than AlphaFold-family methods overall.
  • MD ensembles: AlphaFLOW ensembles match aggregate MD diversity and achieve the strongest Pearson correlation with ground-truth flexibility across 82 unseen ATLAS test proteins.For within-target RMSF, AlphaFLOW reaches a median Pearson correlation of 0.85, whereas MSA subsampling does not meaningfully exceed baseline AlphaFold.
  • MD ensembles: AlphaFLOW is more accurate than every MSA-subsampling level for positional distributions and substantially outperforms MSA subsampling in modeling atomic variance.It slightly improves on AlphaFold for mean atomic positions and nearly doubles the success rate for obtaining > 0.5 cosine similarity between principal components.
  • Efficiency: AlphaFLOW retains much of its ensemble quality with up to a 10x reduction in samples, while short MD trajectories require substantially longer to converge to comparable quality.Distilled AlphaFLOW supplies 10x as many samples for the same runtime, despite converging to a lower performance level.

5. Conclusion

AlphaFLOW and ESMFLOW combine AlphaFold and ESMFold with flow matching to sample protein ensembles, achieving precise and diverse PDB structures and reproducing distributions and properties of MD ensembles.

  • AlphaFLOW and ESMFLOW combine AlphaFold and ESMFold with flow matching to sample protein ensembles.
  • The method uses training-time modeling of structural diversity rather than inference-time input modifications alone.
  • Experiments demonstrate precise and diverse PDB structure prediction and replication of MD ensemble distributions and properties.
  • The authors anticipate applications in structural biology, including settings involving cryo-EM structural heterogeneity.

A. Method Details

The method adapts AlphaFold and ESMFold to flow matching over protein structures while addressing SE(3) invariance, all-atom supervision, and manifold-based inference.

  • Architecture: AlphaFLOW and ESMFLOW attach an input embedding module to AlphaFold and ESMFold, producing pair embeddings from coordinates and time.The module outputs an input pair embedding z ∈ R^N×N×64 and adds time information through Gaussian Fourier embeddings.
  • SE(3) invariance: Standard MSE is unsuitable because an SE(3)-invariant network cannot predict randomly rotated targets from invariant inputs.The paper notes that this mismatch can cause degenerate prediction of zero coordinates.
  • SE(3) invariance: Flow matching therefore operates on the quotient space R^3N/SE(3), where structures related by rigid transformations are treated as identical.This produces a non-Euclidean Riemannian manifold with tangent-space vector fields.
  • Flow construction: Protein structures are interpolated after RMSD alignment, and inference advances by linear interpolation toward the denoised prediction.The inference step moves a fraction dt/(1 − t) toward the predicted structure.
  • Training targets: The model is supervised on all-atom coordinates and residue frames using a squared FAPE-based target while learning the distribution over β-carbon coordinates.At the final inference step, the predicted all-atom structure is returned rather than only its β-carbon component.

A.3. Comparison with Harmonic Diffusion

Compared with harmonic diffusion, the flow-matching framework offers fixed-time and crop-aligned noising properties, handles missing residues more flexibly, but loses crop invariance after RMSD alignment.

  • Harmonic diffusion: Harmonic diffusion selects a finite maximum time by requiring the conditional path to approach its stationary noisy prior within an acceptable KL-divergence.
  • Convergence: For larger proteins, harmonic diffusion convergence becomes difficult because the smallest drift-matrix eigenvalue decreases with protein size.Training on 256-residue crops may not expose intermediate states where longer proteins are only partially denoised.
  • Convergence: Flow matching imposes the noisy prior at t = 0 for all protein lengths and crops, avoiding the intermediate-state mismatch described for harmonic diffusion.
  • Crop invariance: Before RMSD alignment is introduced, crop invariance makes a crop’s marginal distribution equal to independently noising the intact crop.The argument follows because both the prior and linear interpolation preserve the corresponding crop distribution.
  • Crop invariance: The final flow-matching formulation loses crop invariance because of its RMSD alignment step, although the authors posit that initial distributional alignment may aid generalization to unseen large proteins.
  • Missing residues: Unlike harmonic diffusion’s globally coupled Gaussian perturbation, flow matching can treat missing coordinates without requiring the same matrix-exponential imputation procedure.

B.1. Training and Inference

The models are trained by fine-tuning AlphaFold and ESMFold with flow matching, then sampled through a multi-step inference process whose diversity and precision can be adjusted. Self-conditioning and distillation provide additional precision and speed options.

  • Training: Training uses an updated PDB snapshot clustered at 40% sequence identity with MMSeqs2, while the total training cost is reported separately in Table 2.All training runs use 8 NVIDIA A100 GPUs and 2 Intel Xeon Gold 6258R processors.
  • Inference: Default inference uses 10 evenly spaced steps from t = 0 to t = 1, beginning with an unnoised prediction.Merging the first K > 1 steps reduces sampled-distribution variance and increases precision by stepping directly to an intermediate time.
  • Inference: The original and modified interpolation processes differ because the modified process moves toward the initial point estimate rather than sampling the full data distribution.The original uses x1 sampled from pdata(x1), whereas the modified process interpolates toward Epdata[x1].
  • Inference: Self-conditioning reuses an initial forward pass during later predictions, while distillation learns the full inference map so one forward pass can produce a denoised output.Self-conditioning is used to increase AlphaFLOW precision, and distillation exploits the deterministic inference process after the initial noisy sample.

B.2. Datasets

The evaluation uses two complementary datasets: a PDB test set of recent structurally heterogeneous proteins and the ATLAS collection of all-atom molecular-dynamics ensembles.

  • PDB test set: The PDB test set groups chains by shared UniProt segments and clusters the resulting segments at a Jaccard similarity threshold of 0.75.SIFTS residue-level mappings associate PDB chains with UniProt references before complete-linkage clustering defines distinct proteins.
  • MD dataset: ATLAS contains 1390 non-membrane proteins representing eligible ECOD structural classes, each with three 100 ns all-atom explicit-solvent simulations.Each replicate contains 10k frames, providing multiple trajectories for training, validation, and evaluation.
  • Dataset characterization: Figure 6 compares sequence-length distributions for the PDB test set and the ATLAS test set.The two histograms provide a direct view of dataset size ranges rather than structural or ensemble measurements.

B.3. Evaluation Procedures

Evaluation compares generated ensembles with reference structures or MD trajectories using alignment-aware structural, distributional, and dynamical measures, including controlled replicate-MD and normal-mode baselines.

  • Structural similarity: Symmetrized lDDT-Cα aligns sequences and scores whether selected Cα pairwise distances agree within four distance tolerances.Pairs are selected when within 15 Å in either structure, and agreement is assessed within 0.5 Å, 1 Å, 2 Å, and 4 Å.
  • MD ensemble metrics: MD evaluations align generated and reference ensembles to the static structure that initialized the simulation before analyzing Euclidean atomic coordinates.Most procedures use 1000 randomly subsampled MD frames to reduce analysis runtime.
  • MD ensemble metrics: RMWD decomposes into translation and variance contributions, which are reported separately alongside aggregate distributional comparisons.The squared distance separates displacement of ensemble means from differences in ensemble variance.
  • MD ensemble metrics: Joint W2 distance is computed after PCA projection because thermal fluctuations make the full-dimensional metric unsuitable without very large sample counts.The evaluation also accounts for deviations along degrees of freedom outside a PCA basis learned only from the MD reference.
  • Baselines: Replicate-MD comparisons use one ATLAS trajectory as the comparison run and pool the other two as the reference ensemble.This avoids using the same replicate both as the evaluated trajectory and as part of the pooled reference.
  • Baselines: Normal-mode baselines draw 250 samples from Gaussian and anisotropic network models constructed from PDB Cα coordinates.The models use adjusted Γ = 0.15 and default 10 Å and 15 Å cutoffs.

C. Additional Results

Additional results examine PDB ensemble quality, MD-ensemble fidelity, contact and solvent-exposure behavior, and runtime convergence relative to replicate MD.

  • PDB ensembles: Table 3 reports median precision, recall, diversity, and per-sample runtime over 100 PDB test targets.Figures 7 and 8 visualize PCA relationships and conformational changes sampled by AlphaFLOW.
  • PDB ensembles: MSA subsampling produces diversity similar to AlphaFLOW but drifts away from the true PDB structures in PCA space.Figure 7 contrasts the generated ensembles with the true structures marked by triangles.
  • Ablations and baselines: Table 5 shows that two-step AlphaFLOW+MD training matters, while normal mode analysis often fails to outperform AlphaFLOW+MD and underperforms the template-conditioned method.The table notes that some normal-mode RMSF and RMWD results are Cα-only rather than all-atom.
  • Runtime convergence: Figure 13 compares AlphaFLOW variants with replicate MD across runtime, using GPU-hours averaged over targets and reference-MD dashed lines for Pairwise RMSD and RMSF.The comparison includes template-conditioned sampling, single-forward-pass distillation, and MD trajectories of varying lengths.
Loading 2402.04845v2…