PaperlayerView exact revision on arXiv

Source-linked AI summary

SE(3)-Stochastic Flow Matching for Protein Backbone Generation

Avishek Joey Bose, Tara Akhound-Sadegh, Guillaume Huguet, Kilian Fatras, Jarrid Rector-Brooks, Cheng-Hao Liu, Andrei Cristian Nica, Maksym Korablyov, Michael Bronstein, Alexander Tong

arXiv:2310.02391v4cs.LGcs.AI

TL;DR

Computational protein design seeks novel structures with useful functions, but accurate backbone generation remains a central challenge. FoldFlow applies simulation-free flow matching on SE(3), with deterministic, optimal-transport, and stochastic variants, and generates protein backbones up to 300 amino acids that are designable, diverse, and novel.

  • Problem

    Computational design aims to create novel protein structures with desired functions, supporting applications including drug design.

  • Method

    FoldFlow is a family of simulation-free flow-matching generative models over SE(3), comprising deterministic, Riemannian-optimal-transport, and stochastic variants.

  • Results

    FoldFlow generates protein backbones up to 300 amino acids, outperforming FrameDiff-Improved across metrics and remaining competitive with RFDiffusion.

  • Takeaways & Limitations

    FoldFlow supports designable, diverse, and novel backbone generation and is more amenable to equilibrium conformation sampling.

  • Takeaways & Limitations

    Sampling the exact conditional stochastic bridge requires SDE simulation, which limits scalability.

Abstract

from arXiv · show

The computational design of novel protein structures has the potential to impact numerous scientific disciplines greatly. Toward this goal, we introduce FoldFlow, a series of novel generative models of increasing modeling power based on the flow-matching paradigm over $3\mathrm{D}$ rigid motions -- i.e. the group $\text{SE}(3)$ -- enabling accurate modeling of protein backbones. We first introduce FoldFlow-Base, a simulation-free approach to learning deterministic continuous-time dynamics and matching invariant target distributions on $\text{SE}(3)$. We next accelerate training by incorporating Riemannian optimal transport to create FoldFlow-OT, leading to the construction of both more simple and stable flows. Finally, we design FoldFlow-SFM, coupling both Riemannian OT and simulation-free training to learn stochastic continuous-time dynamics over $\text{SE}(3)$. Our family of FoldFlow, generative models offers several key advantages over previous approaches to the generative modeling of proteins: they are more stable and faster to train than diffusion-based approaches, and our models enjoy the ability to map any invariant source distribution to any invariant target distribution over $\text{SE}(3)$. Empirically, we validate FoldFlow, on protein backbone generation of up to $300$ amino acids leading to high-quality designable, diverse, and novel samples.

1 INTRODUCTION

FoldFlow addresses protein backbone generation as modeling invariant distributions over SE(3), using simulation-free flow matching to construct increasingly expressive continuous normalizing flows. Its models target backbones up to 300 residues and improve designability, novelty, or computational efficiency relative to cited diffusion baselines.

  • Problem: Protein backbones are represented as N residue-level rigid frames in SE(3), making protein design a sampling problem over the corresponding group product.Each residue frame encodes orientation-preserving rotations and translations applied to idealized backbone atoms.
  • Problem: Existing manifold generative models do not specifically exploit SE(3)'s geometric structure, while some alternatives require expensive SDE simulation or Riemannian-divergence objectives.These constraints motivate a purpose-built approach for protein backbone generation.
  • Approach: FoldFlow introduces SE(3)-tailored continuous normalizing flows trained with Conditional Flow Matching, which directly regresses time-dependent vector fields without simulation.The models learn SE(3)-invariant distributions for protein backbone generation and can start from an informative prior.
  • Approach: FoldFlow-Base uses a closed-form conditional vector field for SO(3), while FoldFlow-OT uses Riemannian optimal transport to construct shorter, simpler, and more stable flows.FoldFlow-SFM further learns stochastic dynamics over SE(3) through a simulation-free stochastic bridge.
  • Results: All FoldFlow models generate protein backbones of up to 300 residues; they outperform FrameDiff for in-silico designability, while FoldFlow-SFM is competitive with RFDiffusion for novelty using a fraction of its compute and data.FoldFlow-OT is the most designable model among the FoldFlow variants reported here.

2 BACKGROUND AND PRELIMINARIES

The paper formulates flow matching on Riemannian manifolds and parameterizes protein backbones as residue-level SE(3) frames. This geometric representation supports simulation-free learning of flows that transport one probability distribution to another.

  • Riemannian manifolds: A Riemannian manifold equips each tangent space with an inner product, enabling distances, volumes, angles, and length-minimizing geodesics.The metric supplies the geometric structure needed for manifold-valued probability modeling.
  • Lie groups: SE(3) is the matrix Lie group of orientation-preserving rigid motions, combining SO(3) rotations with R3 translations.Its Lie-group structure provides the state space used for protein residue frames.
  • Flow matching on Riemannian manifolds: A probability path ρ_t interpolates between endpoint distributions, while a time-dependent vector field generates the path through an ODE and pushes the source distribution toward the target.Flow matching learns this vector field by directly regressing a parametric field, yielding a simulation-free training objective.
  • Flow matching on Riemannian manifolds: During inference, the learned vector field propagates an ODE backward from samples of the target endpoint distribution along the learned probability path.Conditional and unconditional flow-matching objectives have the same gradients under the stated setup.
  • Protein backbone parametrization: Protein backbones are represented by N SE(3)-equivariant residue frames, each transforming idealized coordinates of four heavy atoms N, Cα, C, and O.An additional rotation angle constructs the backbone oxygen atom, producing full heavy-atom coordinates A ∈ R^(N×4×3).

3 FOLDFLOW FOR CONDITION FLOW MATCHING ON SE(3)

FoldFlow constructs equivariant flow-matching models on centered protein-backbone frames in SE(3)^N, progressing from baseline geodesic paths to Riemannian optimal-transport and stochastic bridges.

  • FOLDFLOW-BASE: Conditional Flow Matching learns time-dependent vector fields directly, avoiding simulation during training while transporting source and target distributions.The framework regresses vector fields associated with probability paths on the manifold.
  • SE(3) geometry: Protein backbones are modeled as centered products of residue-level rigid-motion frames in SE(3), with separate rotation and translation flows.The product-group geometry permits decomposing the flow across residues after removing global translations.
  • FOLDFLOW-OT: Riemannian optimal transport produces shorter, straighter conditional paths intended to reduce objective variance and stabilize training.FOLDFLOW-OT constructs its interpolation measure from an optimal transport plan on the product manifold.
  • FOLDFLOW-OT: FOLDFLOW-OT relies on Riemannian Monge or Kantorovich transport, with uniqueness guaranteed only under stated support and absolute-continuity assumptions.The proposition assumes compactly supported distributions and an absolutely continuous source distribution.
  • FOLDFLOW-SFM: FOLDFLOW-SFM extends the deterministic models with stochastic dynamics over SE(3), using guided diffusion bridges and a flow-matching objective.The bridge-based objective can match the unconditional flow objective in expectation when the coupling is valid.

4 MODELING PROTEIN BACKBONES USING FOLDFLOW

FOLDFLOW parameterizes velocity prediction through a predicted starting structure, enabling structure-specific modeling and auxiliary protein losses; generated examples are compared by RMSD across lengths.

  • Velocity parameterization: The velocity predictor consumes a protein state and time, then predicts the starting structure x̂0 used to parameterize the flow.The velocity is defined from the Riemannian gradient between the predicted start and current state.
  • Velocity parameterization: Predicting x̂0 permits specialized structure-prediction architectures and places auxiliary protein-specific losses directly on the predicted structure.These are presented as two principal benefits of the parameterization.
  • Training objective: The model applies auxiliary losses only for t < 0.25, scaled by λaux in the final loss.The supplied passage specifies the time restriction and scaling factor but not the complete displayed loss.
  • Qualitative evaluation: Figure 3 compares green FOLDFLOW-SFM structures with grey ProteinMPNN → ESMFold refolds for lengths 100, 150, 200, 250, and 300.Shown examples satisfy RMSD < 2 Å.

5 EXPERIMENTS

The experiments evaluate FOLDFLOW on synthetic SO(3) densities, protein backbone design, and equilibrium conformation generation. FOLDFLOW models accurately capture multimodal and protein conformational distributions while improving designability and efficiency against relevant baselines.

  • Synthetic densities: All proposed methods model every mode of the synthetic SO(3) ground-truth distribution, while FOLDFLOW-BASE exhibits mode shrinkage.FOLDFLOW-OT, FOLDFLOW-SFM, and simulated SDE achieve comparable performance, with the OT-based method performing best.
  • Protein backbone design: FOLDFLOW outperforms FrameDiff-Retrained on designability, diversity, and novelty when trained on 22,248 PDB proteins.The comparison includes pretrained FrameDiff variants, Genie, and RFDiffusion; RFDiffusion uses substantially more parameters, data, and compute.
  • Protein backbone design: Figure 4 evaluates novelty, designability, flow scaling, and feature ablations across models trained on PDB subsets.Novelty is assessed using TM-score thresholds, while designability uses scRMSD < 2.
  • Protein backbone design: Δ = 0.149 for FOLDFLOW-OT versus Δ = 0.357 for FrameDiff-Retrained on the reported designability comparison.FOLDFLOW-OT produces the highest fraction of designable proteins among non-RFDiffusion models, but has relatively low diversity and novelty.
  • Equilibrium conformation generation: FOLDFLOW can model both general BPTI conformations and highly flexible residues against unseen molecular-dynamics frames.The evaluation uses 200,000 training frames and tests against 20,000 unseen frames from a 1 ms BPTI trajectory.

6 RELATED WORK

Related work spans protein-design methods using sequence or structure information and generative models that incorporate symmetry constraints. These approaches motivate geometry-aware modeling for protein structures.

  • Protein design approaches: Protein design has developed through experimental libraries, machine-learning-assisted sequence methods, and structure-based biophysical approaches.The related work positions FOLDFLOW within a broader evolution from established protein-engineering techniques to machine learning.
  • Equivariant generative models: Equivariant generative modeling includes equivariant vector fields for continuous normalizing flows and finite flows using affine coupling transforms.The cited applications extend beyond protein design into theoretical physics and related generative-modeling settings.

7 CONCLUSION

FOLDFLOW is a family of simulation-free, SE(3)-invariant flow-matching models for protein backbone generation. Its variants use deterministic, optimal-transport, and stochastic dynamics, and are competitive with strong baselines while supporting equilibrium conformation sampling.

  • 7 CONCLUSION: FOLDFLOW-BASE learns deterministic dynamics over SE(3), FOLDFLOW-OT uses Riemannian optimal transport, and FOLDFLOW-SFM learns stochastic dynamics.The family is introduced under the simulation-free flow-matching framework for invariant generative modeling.
  • 7 CONCLUSION: Protein backbone generation is formulated over SE(3)^N, where each residue is represented by a rigid-body frame.The conclusion frames backbone design as sampling from a distribution over the product of residue-level rigid motions.
  • 7 CONCLUSION: FOLDFLOW models are competitive with RFDiffusion and significantly outperform the non-pretrained FrameDiff-Improved approach across reported metrics.The models are evaluated on PDB proteins containing up to 300 amino acids.
  • 7 CONCLUSION: FOLDFLOW is more amenable to equilibrium conformation sampling because it can be instantiated from an informative source distribution.The approach is motivated by Brownian bridges over SO(3) while avoiding simulation during training.
  • 7 CONCLUSION: The SE(3) metric decomposes into SO(3) and R3 components, allowing independent flows on rotations and translations.This construction supplies the geometric basis for applying flow matching to residue-level rigid transformations.

C.2 PROOF OF PROPOSITION 1

The appendix establishes optimal-transport conditions on SE(3)^N and compares simulation-free and simulated Brownian bridges on SO(3). The numerical comparison supports FOLDFLOW-SFM as a fast approximation to the guided bridge SDE.

  • C.2 PROOF OF PROPOSITION 1: The proof relies on SE(3)^N being connected, complete, smooth, and without boundary, with translations handled through a mean-subtracted product decomposition.The argument applies a standard optimal-transport theorem after verifying the required manifold and measure conditions.
  • Brownian bridge comparison: Brownian bridges on Lie groups are simulated with a guided diffusion SDE conditioned to reach a target at t = 1.For SO(3), the bi-invariant metric gives V0 = 0, and the paper uses a process-independent diffusion for its reversed formulation.
  • Brownian bridge comparison: The simulation-free FOLDFLOW-SFM SDE stays close to the simulated Brownian bridge across trajectories and diffusion coefficients.The comparison uses means and standard deviations over 1024 data points, with overlapping shaded standard-deviation regions.

D.3 PROOF OF PROPOSITION 2

The proposition establishes that conditional and unconditional FOLDFLOW-SFM training have identical parameter gradients, supporting stochastic flow learning on SE(3). The proof uses factorization into SO(3) and R3 and tangent-space vector-field construction.

  • The conditional and unconditional FOLDFLOW-SFM losses have equal gradients with respect to θ.
  • The proof reduces the SE(3) claim to SO(3) and R3 because both the measures and metric factorize across these components.
  • The SO(3) proof follows the established conditional-flow argument, while the R3 component follows a prior theorem.
  • The vector field is constructed intrinsically on the manifold so that it lies in the tangent space and is equivariant by construction.

F.2 ADDITIONAL RESULTS FOR SO(3) TOY EXPERIMENTS

The SO(3) toy experiments show that all three FOLDFLOW variants recover the modes of the ground-truth distribution, while FOLDFLOW-BASE exhibits slight mode shrinkage.

  • All FOLDFLOW-BASE, FOLDFLOW-OT, and FOLDFLOW-SFM models correctly learn the modes of the ground-truth distribution.
  • FOLDFLOW-BASE shows slight model shrinkage relative to the ground-truth distribution.

G.1 PROTEIN BACKBONE GENERATION EXPERIMENT ADDITIONAL RESULTS

Additional protein-generation results compare FOLDFLOW with prior models across designability, diversity, novelty, and computational cost. FOLDFLOW is especially competitive on longer backbones and trains faster than FrameDiff.

  • Protein generation metrics: For lengths 250 and 300, FOLDFLOW models are comparable in designability but more diverse and novel than other models, including RFDiffusion.
  • Protein generation metrics: For lengths ≤200, FOLDFLOW particularly closes the designability gap between non-pretrained models and RFDiffusion.
  • Protein generation metrics: FOLDFLOW exhibits a trade-off between designability and diversity or novelty across sequence lengths.
  • Training efficiency: FOLDFLOW is over 2x faster than FrameDiff per training step on a single GPU.
  • Training efficiency: The reported FOLDFLOW training run used four NVIDIA A100-80GB GPUs for roughly 2.5 days.

G.2 EQUILIBRIUM CONFORMATION GENERATION EXPERIMENT

The equilibrium-conformation experiment tests whether informed priors help FOLDFLOW model BPTI dynamics. The results support informed-prior benefits and show that FOLDFLOW improves over FrameDiff on this task.

  • Interpretation: FOLDFLOW’s ability to use an informative source distribution is relevant to equilibrium conformation generation and related structured-generation settings.
  • Experimental setup: BPTI is a 58-residue protein whose dynamics include backbone changes associated with water exchange and disulfide isomerization.
  • Experimental setup: The experiment evaluates 2-Wasserstein distances over all 58 residues and residue 56, plus KL divergence at residue 56.
  • Prior comparison: An informed prior improves FOLDFLOW performance both overall and on BPTI’s most flexible residue.
  • Prior comparison: Both informed- and random-prior FOLDFLOW improve over FrameDiff on the equilibrium-conformation task.
  • Interpretation: Flow matching can transport arbitrary source distributions to target distributions, unlike diffusion approaches that typically require Gaussian-like sources.
  • Model context: FOLDFLOW-OT and FOLDFLOW-SFM connect optimal transport with flow matching, while stochasticity in FOLDFLOW-SFM is associated with improved novelty in protein generation.
  • Model context: FrameDiff’s rotation-score computation contains a corrected bug in its improved version, and its data sampling scheme changes the protein distribution while increasing diversity.

I.5 DATA AND DATA SAMPLING

The evaluation dataset contains filtered protein monomers sampled uniformly over 30% similarity clusters, and the paper reports how this affects length distributions and metric-based assessment. Designability, diversity, and novelty are evaluated with structure-alignment measures and explicit dataset caveats.

  • Dataset construction: 22,248 monomers remain after filtering PDB structures by length, resolution, loop content, and other admissibility criteria.The retained proteins have lengths between 60 and 512 residues, resolution below 5Å, and no more than 50% loops.
  • Data sampling: Uniform sampling over clusters with 30% similarity supports diversity during training and changes the batch length distribution.The dataset, batch sampling distribution, and cluster-size distribution are shown in Figure 13.
  • Protein metrics: TM-score measures structural similarity after alignment, with values above 0.5 generally indicating similar folds and values below 0.2 corresponding to random comparisons.The metric ranges from 0 to 1, with 1 indicating perfectly aligned structures.
  • Protein metrics: RMSD measures paired-residue structural deviation, is length dependent, and is more sensitive to local errors than TM-score.The paper uses RMSD-based filtering alongside TM-score to assess designability.
  • Protein metrics: Novelty is assessed using minimum TM-score to training data and the fraction of proteins that are both designable and novel.The latter criterion uses scRMSD < 2 Å and average maximum TM-score < 0.5; larger training sets may overestimate novelty for some baselines.

J EXTENDED ABLATION OF EXPERIMENTS

The ablation study isolates stochasticity, optimal transport, auxiliary losses, and inference annealing across designability, diversity, and novelty. Each feature shows a distinct empirical benefit across most or all tested settings.

  • Stochasticity: Stochasticity improves robustness and novelty in 7/8 settings, measured by the fraction of novel and designable proteins.The stochastic FOLDFLOW-SFM model is also associated with improved novelty in protein generation.
  • Optimal transport: Optimal transport improves designability by reducing variance in the training objective in 7/8 settings.The paper contrasts non-crossing optimal-transport paths with potentially crossing FoldFlow-Base paths, which reduces objective variance.
  • Auxiliary losses: Auxiliary losses improve model designability in 7/8 settings.The model parameterization places protein-specific auxiliary losses directly on the predicted starting structure.
  • Inference annealing: Inference annealing improves every FOLDFLOW model across all evaluated metrics.The metrics are designability, diversity, and novelty.
Loading 2310.02391v4…