Source-linked AI summary
Fast protein backbone generation with SE(3) flow matching
Jason Yim, Andrew Campbell, Andrew Y. K. Foong, Michael Gastegger, José Jiménez-Luna, Sarah Lewis, Victor Garcia Satorras, Bastiaan S. Veeling, Regina Barzilay, Tommi Jaakkola, Frank Noé
TL;DR
Protein backbone diffusion models can be expensive because high-quality sampling often requires many forward passes. The paper adapts FrameDiff into FrameFlow, an SE(3) flow-matching model, and reports substantially fewer sampling steps with better designability and lower sampling cost.
Problem
Diffusion-based protein backbone generation can require ∼1000 model forward passes, making large-scale inference prohibitively expensive.
Method
FrameFlow adapts the FrameDiff diffusion model to flow matching on the SE(3) manifold for protein backbone generation.
Results
Using five times fewer sampling timesteps than FrameDiff, FrameFlow achieves two-fold better designability; compared with GENIE, it achieves a 23-times sampling speedup with higher designability.
Takeaways & Limitations
FrameFlow demonstrates the potential of flow matching to reduce inference costs while maintaining or improving protein backbone generation quality.
Abstract
from arXiv · showhide
We present FrameFlow, a method for fast protein backbone generation using SE(3) flow matching. Specifically, we adapt FrameDiff, a state-of-the-art diffusion model, to the flow-matching generative modeling paradigm. We show how flow matching can be applied on SE(3) and propose modifications during training to effectively learn the vector field. Compared to FrameDiff, FrameFlow requires five times fewer sampling timesteps while achieving two fold better designability. The ability to generate high quality protein samples at a fraction of the cost of previous methods paves the way towards more efficient generative models in de novo protein design.
1 Introduction
FrameFlow applies flow matching to protein backbone generation to address the high inference cost of diffusion models. It adapts FrameDiff and reports improved designability and sampling efficiency relative to prior approaches.
- Diffusion models can require ∼1000 forward passes to generate high-quality protein samples, making large-scale inference prohibitively expensive.
- Flow matching offers straighter sampling trajectories that can be simulated with fewer integration steps than diffusion-based approaches.
- FrameFlow adapts the FrameDiff diffusion model to flow matching for protein backbone generation on SE(3).
- Using five times fewer sampling timesteps than FrameDiff, FrameFlow achieves two-fold better designability, comparable diversity, and equal novelty scores.
- Compared with GENIE, FrameFlow achieves a 23-times sampling speedup while maintaining a significantly higher designability score.
2 Method
FrameFlow applies flow matching to protein backbone generation on the SE(3) manifold, learning an equivariant vector field that transforms simple priors into backbone structures. The method adapts FrameDiff and adds training and inference modifications for effective, efficient generation.
- Flow matching on Riemannian manifolds: Flow matching learns a time-dependent vector field on a Riemannian manifold and generates samples by integrating an ordinary differential equation.The learned field is trained through conditional vector fields and then used to simulate the flow from a prior toward the data distribution.
- Flow matching on SE(3): Protein backbones are represented as residue-wise rigid transformations in SE(3), combining rotations in SO(3) with translations in R3.The construction extends from one frame to all residues using the product space SE(3)^N and an additive metric.
- Flow matching on SE(3): Conditional paths connect prior and data frames along geodesics, with linear translation interpolation and geodesic rotation interpolation on the selected manifold.The priors are uniform on SO(3) and Gaussian on R3, while exponential and logarithmic maps define the rotational geodesics.
- FrameFlow architecture: FrameFlow uses an SE(3)-equivariant neural network and FramePred components to predict clean residue frames while preserving the required symmetry.The architecture incorporates invariant point attention and transformer layers, while the model operates in the zero center-of-mass subspace.
- FrameFlow modifications: Training and inference are modified with rotation-loss weighting, clipped late-time scaling, pre-alignment, an exponential SO(3) inference scheduler, and an alternative IGSO3 training prior.Pre-alignment removes global rotations associated with increased ODE kinetic energy; the exponential scheduler is used during inference while linear scheduling is retained during training.
3 Experiments
FrameFlow is evaluated for protein backbone generation using designability, diversity, novelty, and sampling efficiency, with comparisons against GENIE and FrameDiff. Its straighter ODE trajectories support strong performance with fewer timesteps, while GENIE degrades sharply when sampling is shortened.
- Evaluation: FrameFlow is evaluated on SCOPe proteins below length 128 using designability, diversity, and novelty metrics.The evaluation samples backbones across lengths 60–128, designs sequences with ProteinMPNN, and assesses structures with ESMFold.
- Sampling setup: FrameFlow uses Euler ODE sampling, while GENIE and FrameDiff use Euler-Maruyama SDE sampling.The experiments vary integration timestep counts across methods.
- Results: GENIE’s performance deteriorates rapidly below 1000 timesteps and is unusable at 500 timesteps because it produces no designable samples.Its diversity and novelty can also be skewed because those metrics are conditioned on designable samples.
- Results: FrameFlow with 100 timesteps outperforms FrameDiff on designability, despite using the same architecture.FrameFlow’s performance deteriorates rapidly at 10 timesteps, suggesting the tested timestep reduction has a practical limit.
- Results: FrameFlow discovers more structural modes than FrameDiff despite apparently lower diversity under the reported metric.The paper attributes the apparent diversity difference to diversity being inversely proportional to the number of designable samples.
- Efficiency: Sampling a length-100 protein takes 5.7 seconds for FrameDiff/FrameFlow at 100 timesteps versus 128 seconds for GENIE at 1000 timesteps on an NVIDIA V100.GENIE has fewer parameters but uses computationally expensive triangle updates.
- Trajectory analysis: FrameFlow and FrameDiff exhibit different sampling trajectories, with FrameFlow motivated by straighter and faster paths.The trajectory visualization compares FrameFlow’s ODE sampling with FrameDiff’s SDE sampling for a length-100 protein.
4 Discussion
The paper adapts FrameDiff into FrameFlow, an SE(3) flow-matching model, and reports superior performance over FrameDiff and GENIE. The authors present these experiments as preliminary evidence that flow matching may help scale protein generative models, while identifying further exploration as necessary.
- Discussion: FrameFlow adapts the SE(3) diffusion model FrameDiff into a flow-matching model and is reported to outperform FrameDiff and GENIE.The discussion characterizes the experiments as preliminary demonstrations rather than a completed account of flow matching’s potential.
- Implications: The authors identify scaling generative models for protein design as a potential application of flow matching.They state that additional flow-matching techniques should be explored for real-world protein-design applications.