SE(3) STOCHASTIC FLOW MATCHING FOR PROTEIN BACKBONE GENERATION

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 3D rigid motions-i.e. the group 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 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 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 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.