FoldFlow-SFM: Stochastic SE(3) Protein Design

• FoldFlow-SFM is a generative modeling framework that leverages stochastic flow matching on SE(3) to design diverse protein backbones.
• It integrates simulation-free training with Riemannian optimal transport to align drift fields with geodesic interpolations, reducing regression variance.
• Empirical evaluations reveal that FoldFlow-SFM achieves high novelty in protein design while maintaining competitive designability and faster convergence.

FoldFlow-SFM is a stochastic generative modeling framework for protein backbone generation that operates on the Lie group of 3D rigid motions, $\text{SE}(3)$, integrating simulation-free flow matching, Riemannian optimal transport (OT), and stochastic bridge dynamics. Designed to address challenges in generating accurate, designable, diverse, and novel protein structures, FoldFlow-SFM learns continuous-time stochastic differential equations (SDEs) on $\text{SE}(3)$, coupling rotational and translational dynamics. It enables stable and efficient mapping from arbitrary source distributions to invariant target distributions, improving on the limitations of previous diffusion-based approaches in both training speed and expressiveness (Bose et al., 2023).

1. Mathematical Formulation on $\text{SE}(3)$

FoldFlow-SFM operates in the space of 3D rigid motions, with each element $X = (R, s)$ where $R \in SO(3)$ and $s \in \mathbb{R}^3$. The model defines time-indexed SDEs that decompose into a rotational Brownian bridge on $SO(3)$ and a Euclidean bridge on $\mathbb{R}^3$. The forward-inference SDEs, used in generation and trained in reverse, are:

• Rotation ($SO(3)$):

$dR_t = v_\theta(t, R_t)dt + \gamma_r(t) dB_t$

• Translation ($\text{SE}(3)$0):

$\text{SE}(3)$1

where $\text{SE}(3)$2 are trainable drift fields, $\text{SE}(3)$3 and $\text{SE}(3)$4 are Brownian motions on $\text{SE}(3)$5 and $\text{SE}(3)$6 respectively, and $\text{SE}(3)$7, $\text{SE}(3)$8 are time-dependent noise scales. The guided bridge dynamics, conditioned on initial and terminal states, are expressed using the Riemannian log-map:

• Rotation Bridge:

$\text{SE}(3)$9

• Translation Bridge:

$\text{SE}(3)$0

These dynamics enable direct sampling and interpolation between boundary conditions, critical for protein backbone design.

2. Flow-Matching Objective and Loss

FoldFlow-SFM utilizes a flow-matching paradigm, regressing a neural drift field $\text{SE}(3)$1 against the conditional bridge drift $\text{SE}(3)$2 under a coupling plan $\text{SE}(3)$3 on $\text{SE}(3)$4. The stochastic flow-matching loss is:

$\text{SE}(3)$5

where $\text{SE}(3)$6 represents the Riemannian norm on $\text{SE}(3)$7. This objective ensures the learned flow aligns with the true dynamics of the stochastic bridges over time, facilitating the generative process for structured protein backbones.

The 2-Wasserstein (OT) cost between marginals $\text{SE}(3)$8 and $\text{SE}(3)$9 is defined as:

$X = (R, s)$0

where $X = (R, s)$1 denotes the geodesic distance on $X = (R, s)$2.

3. Incorporation of Riemannian Optimal Transport

FoldFlow-SFM incorporates optimal transport to efficiently couple source and target distributions. Geodesic distances on $X = (R, s)$3 with the standard left-invariant metric split into rotation and translation components:

$X = (R, s)$4

with

$X = (R, s)$5

Replacing independent couplings with the OT plan $X = (R, s)$6 in the flow-matching loss results in OT-regularized flow-matching, empirically reducing variance in the regression loss and accelerating convergence. Conditional paths between paired samples become geodesic ("straight") in probability space, yielding more stable learning trajectories.

4. Algorithmic Structure and Training

The FoldFlow-SFM training loop operates over protein backbones of $X = (R, s)$7 frames (elements of $X = (R, s)$8), utilizing both OT coupling and stochastic perturbations. The high-level training procedure is as follows:

1. Sample $X = (R, s)$9.
2. Draw $R \in SO(3)$0 via a mini-batch OT routine.
3. For each frame $R \in SO(3)$1, compute the geodesic interpolant $R \in SO(3)$2.
4. Sample bridge perturbations:
   • $R \in SO(3)$3
   • $R \in SO(3)$4
5. Compute drift target:
   • $R \in SO(3)$5
6. Compute flow loss: $R \in SO(3)$6
7. Add auxiliary loss: $R \in SO(3)$7 (structure + distance-map losses on predicted $R \in SO(3)$8)
8. Update parameters: $R \in SO(3)$9

Hyperparameters include a learning rate of $s \in \mathbb{R}^3$0 (Adam optimizer), noise scales such as $s \in \mathbb{R}^3$1 with $s \in \mathbb{R}^3$2, auxiliary weight $s \in \mathbb{R}^3$3 (for $s \in \mathbb{R}^3$4), and an effective batch size inversely scaling with $s \in \mathbb{R}^3$5 ($s \in \mathbb{R}^3$6). At generation time, inference annealing is used, scaling the drift by $s \in \mathbb{R}^3$7 with $s \in \mathbb{R}^3$8 (Bose et al., 2023).

5. Empirical Evaluation and Comparative Results

FoldFlow-SFM was evaluated on protein backbone generation tasks with sequence lengths $s \in \mathbb{R}^3$9, using metrics of designability (fraction of generated backbones refolded by ProteinMPNN→ESMFold with C$SO(3)$0-RMSD < 2Å), diversity (average pairwise TM-score among designable samples), and novelty (fraction and mean maximum TM-score to the training set).

Key results, reported as mean ± standard error over 5 seeds, are summarized below:

Model	Designable %	scRMSD (Å)	Novel frac	avg max TM	Diversity (TM)
FoldFlow-Base	65.7 ± 4.2	3.00 ± 0.27	43.2 ± 7.4	0.452 ± 0.024	0.264
FoldFlow-OT	82.0 ± 3.7	1.81 ± 0.25	48.4 ± 6.8	0.460 ± 0.020	0.247
FoldFlow-SFM	71.6 ± 4.0	2.30 ± 0.39	54.4 ± 6.1	0.411 ± 0.023	0.248
FrameDiff-Improved	55.5 ± 7.1	2.93 ± 0.35	29.6 ± 11.2	0.457 ± 0.026	0.278
RFdiffusion*	96.9 ± 2.3	0.65 ± 0.14	70.8 ± 6.0*	0.449 ± 0.012*	0.256

*RFdiffusion results correspond to a pretrained model on a larger dataset.

FoldFlow-OT achieves the highest non-pretrained designability. FoldFlow-SFM, with its stochastic bridge construction, results in the highest novelty, indicating improved off-support sampling due to the stochastic dynamics. Training is more than twice as fast per step compared to FrameDiff, as repeated score computations are unnecessary.

6. Advantages, Limitations, and Analysis

FoldFlow-SFM leverages OT couplings to generate "straight" McCann interpolants in probability space, empirically reducing regression variance and expediting convergence. Stochastic bridges confer robustness in high dimensions and foster sampling that extends beyond the strict support of deterministic ODE flows, reducing mode collapse and increasing sample novelty.

Unlike classical diffusion models, FoldFlow-SFM's flow-matching framework permits arbitrary source distributions, not restricted to Gaussian priors, supporting applications such as equilibrium sampling. The method is simulation-free, bypassing computational bottlenecks typical of score-matching and diffusion-based models.

A plausible implication is that stochasticity in the flow-matching setup allows FoldFlow-SFM to explore protein conformational space more broadly, achieving superior novelty at a modest cost to designability relative to its deterministic and deterministic-OT counterparts.

7. Broader Context and Related Research

FoldFlow-SFM is part of the broader FoldFlow family, which also includes FoldFlow-Base (deterministic, simulation-free) and FoldFlow-OT (deterministic with OT regularization). These models substantially improve over FrameDiff and offer competitive non-pretrained performance against RFdiffusion, a leading pretrained approach. The integration of SE(3)-stochastic flow matching with Riemannian optimal transport represents an overview of manifold generative modeling, optimal transport theory, and protein structure generation (Bose et al., 2023).

The formulation, scalability (up to 300 amino acids), and empirical efficiency position FoldFlow-SFM as a robust approach for high-dimensional generative problems involving rigid-body dynamics and symmetry constraints, with implications across computational protein design and related fields.