[Paper Review] Fast protein backbone generation with SE(3) flow matching

Jason Yim, Andrew M. Campbell+9 more|arXiv (Cornell University)|Oct 8, 2023

Protein Structure and DynamicsBiochemistry, Genetics and Molecular Biology21 citations

TL;DR

FrameFlow adapts SE(3) flow matching to protein backbone generation, achieving comparable designability with FrameDiff at far fewer sampling timesteps and outperforming GENIE in designability. It significantly speeds up sampling while maintaining diversity and novelty characteristics.

ABSTRACT

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.

Motivation & Objective

• Motivate faster, more efficient generative modeling for de novo protein design using flow matching on SE(3).
• Adapt SE(3) flow matching to backbone generation and train an equivariant vector field to regressor clean frames.
• Demonstrate that FrameFlow achieves comparable or superior designability with fewer sampling steps versus diffusion-based baselines.
• Explore architectural and training tweaks (equivariance, priors, pre-alignment) to maximize performance and stability.

Proposed method

• Apply SE(3) flow matching to protein backbones using a product space SE(3)^N with additive metrics.
• Use conditional vector-field learning to match the conditional flow, enabling a CNF via ODE integration.
• Parameterize translations and rotations with appropriate geodesics (linear interp. for translations; Rodrigues-based exp/log for rotations).
• Train with SE(3)-equivariant FramePred architecture and IPA updates to ensure SE(3) equivariance.
• Employ pre-alignment with Kabsch to improve training efficiency and an IGSO3 prior to improve rotation sampling.
• Use an SO(3) inference scheduler with exponential ramp during sampling and tune loss weights to stabilize training.
• Evaluate on SCOPe data, compare designability, diversity, and novelty against GENIE and FrameDiff, using 10 samples per length and ProteinMPNN/ESMFold designability checks.

Experimental results

Research questions

• RQ1 Can SE(3) flow matching produce high-quality protein backbones with significantly fewer sampling steps than diffusion models?
• RQ2 Does FrameFlow offer improved designability without sacrificing diversity or novelty compared to FrameDiff and GENIE?
• RQ3 What architectural and training adjustments most effectively stabilize SE(3) flow matching for protein backbones?
• RQ4 How do different priors and sampling schedulers affect performance and efficiency in SE(3) flow matching for proteins?

Key findings

• FrameFlow achieves designability of 0.81 at 500 timesteps and 0.77 at 100 timesteps, outperforming FrameDiff and GENIE at corresponding settings.
• FrameFlow attains diversity scores around 0.22–0.28 (clusters) across timesteps, showing more designable samples and more modes than FrameDiff in several settings.
• FrameFlow reaches 0.69–0.67 novelty (lower is better), indicating competitive similarity to known structures while maintaining high designability.
• Compared to GENIE, FrameFlow with 100 timesteps yields higher designability (0.77 vs 0.22 at 1000 timesteps for GENIE SDE) and similar or better designability at fewer steps.
• Sampling time is dramatically reduced: GENIE at 1000 timesteps takes 128 seconds for a length-100 protein, while FrameDiff/FrameFlow at 100 timesteps take 5.7 seconds on similar hardware.
• FrameFlow and FrameDiff use the same backbone architecture; FrameFlow’s ODE-based sampling shows straighter trajectories and faster convergence than FrameDiff’s SDE-based sampling.