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

1 INTRODUCTION

• Protein backbones are modeled as products of SE(3) residue frames, making protein design a sampling problem over this group.The paper motivates computational design as a route to proteins with desired structures and functions, with examples of designed binders for influenza, COVID-19, and cancer targets.Source: §1 · pp. 1–2
• FOLDFLOW-BASE extends Riemannian flow matching with a closed-form SO(3) conditional vector field, increasing training speed and stability.The expression is used in loss computation.Source: §1 · p. 2
• FOLDFLOW-OT uses Riemannian optimal transport to construct shorter, simpler flows and proves the existence of a Monge map on SE(3)^N.Source: §1 · p. 2
• FOLDFLOW-SFM learns a stochastic bridge on SE(3)^N within the paper’s simulation-free conditional flow matching framework.The three models learn SE(3)^N-invariant distributions for protein backbone generation, and FOLDFLOW can start from an informative prior, unlike the SE(3) diffusion approach of Yim et al. (2023b).Source: §1 · p. 2
• All FOLDFLOW models outperform FrameDiff (Yim et al., 2023b) in in-silico designability, with FOLDFLOW-OT best; FOLDFLOW-SFM is competitive with RFDiffusion for novelty.The models generate backbones of up to 300 residues; the paper also reports equilibrium conformation generation by simulating molecular dynamics trajectories from a reference empirical distribution, compared with an uninformed prior.Source: §1 · p. 2

2.1 RIEMANNIAN MANIFOLDS AND LIE GROUPS

• A Riemannian metric equips each tangent space with an inner product, enabling distances, volumes, angles, geodesics, and gradients on a manifold.Source: §2.1 · p. 2
• SE(3) is the matrix Lie group of orientation-preserving rigid motions, represented as the semidirect product of rotations and translations.Lie groups combine smooth group operations with differentiable-manifold structure.Source: §2.1 · pp. 2–3

2.2 FLOW MATCHING ON RIEMANNIAN MANIFOLDS

• Riemannian flow matching learns a time-dependent vector field to connect endpoint distributions, with conditional fields enabling simulation-free training when the path has the required boundary distributions.The conditional path and fields are marginalized to recover the target unconditional path and its vector field.Source: §2.2 · p. 3
• At inference, sampling from the terminal distribution and integrating the learned vector field’s ODE backward generates samples from the data distribution.The FM and CFM objectives have the same gradients.Source: §2.2 · p. 3

2.3 PROTEIN BACKBONE PARAMETRIZATION

• The backbone representation assigns each residue an SE(3) frame that transforms an idealized four-heavy-atom frame, with an additional rotation angle used to construct the oxygen atom.For a protein of length N, the representation has N frames; the full heavy-atom coordinates are denoted A ∈ RN×4×3.Source: §2.3 · p. 3, Figure 2 · p. 3

3 FOLDFLOW FOR CONDITION FLOW MATCHING ON SE(3)

• Centering inputs at their center of mass yields a translation-invariant subspace, while the product metric on SE(3) lets FoldFlow build independent flows on SO(3) and R3.The section focuses on SO(3), since conditional flow matching on R3 is already well studied.Source: §3 · pp. 3–4

3.1 FOLDFLOW-BASE

• FOLDFLOW-BASE transports rotations along geodesics between target and Haar-uniform source samples, using axis-angle and parallel transport to avoid costly matrix logarithms and autograd.Its conditional vector field follows the geodesic’s constant-velocity flow on SO(3).Source: §3.1 · p. 4

3.2 FOLDFLOW-OT

• FOLDFLOW-OT accelerates FOLDFLOW-BASE by constructing conditional paths with Riemannian optimal transport, seeking shorter, straighter, and more stable paths.The method samples endpoint pairs from an optimal transport plan and uses factorized rotational and translational losses across backbone frames.Source: §3.2 · p. 5

3.3 FOLDFLOW-SFM

• The SO(3) stochastic bridge lacks a closed-form conditional path and requires reverse-time SDE simulation, so FOLDFLOW-SFM uses a simulation-free approximation based on IGSO(3).For translations, the reverse-time Brownian bridge is simulation-free and maps arbitrary samplable marginals in expectation.Source: §3.3 · pp. 5–6
• FOLDFLOW-SFM learns stochastic flows between arbitrary supported source and target distributions using samples, with conditional and unconditional losses having equal gradients.Stochastic dynamics are motivated by greater robustness to high-dimensional noise and the ability to sample beyond the training distribution’s support.Source: §3.3 · pp. 5–6

Modeling protein backbones

• The velocity model predicts a starting structure x̂_0 = wθ(t, x_t), then defines velocity as ∇gd(x̂_0, x_t)^2/t.This parameterization supports specialized structure-prediction architectures and auxiliary protein-specific losses on the predicted structure.Source: §4 · p. 6
• FoldFlow uses an AF2 structure module with time-dependent node and edge embeddings, invariant point attention, and an MLP head to predict oxygen torsion angle φ.The architecture follows Anand & Achim (2022) and Yim et al. (2023b) (FrameDiff).Source: §4 · p. 6
• Training combines SO(3) and R3 flow-matching losses with backbone-position and local pairwise atomic-distance losses, applying auxiliary losses only for t < 0.25.The pairwise-distance tensor covers the four heavy atoms per residue.Source: §4 · p. 6

5.1 SO(3) SYNTHETIC DATA

• On synthetic SO(3) densities, all FoldFlow methods recover every mode; FoldFlow-Base shows mode shrinkage, while FoldFlow-OT, FoldFlow-SFM, and simulated SDE perform comparably, with OT best.The comparable performance indicates that simulation-free SDE approximation does not hinder modeling, supporting its use alongside its training speedup.Source: §5.1 · p. 7

5.2 PROTEIN BACKBONE DESIGN

• FoldFlow outperforms FrameDiff-Retrained on all three metrics; its designability gap to RFDiffusion is ∆=0.149 versus ∆=0.357 for FrameDiff-Retrained.The protein-design models were trained on 22,248 PDB proteins; RFDiffusion uses pretraining and substantially more parameters, training data, and compute, while Genie uses a larger dataset.Source: §5.2 · pp. 7–8, Table 2 · p. 7
• FoldFlow-SFM beats FrameDiff-Improved on every metric and produces the most novel structures; FoldFlow-OT has the highest designability among FoldFlow models but relatively low diversity and novelty.FoldFlow-OT and FoldFlow-SFM are the most diverse FoldFlow models; the novelty gain of SFM comes with worse designability.Source: §5.2 · p. 8
• Inference annealing scales rotational velocity by i(t)=ct; larger c increases designability, and c ≈10 yields structures that are both designable and diverse.Annealing removes an unwanted increase in flow norm near the end of inference and is critical to higher designability in the ablation study.Source: §5.2 · p. 8, Figure 4 · p. 8

5.3 EQUILIBRIUM CONFORMATION GENERATION

• For BPTI equilibrium conformations, FoldFlow captures the distribution’s modes and flexible residues; with an informed prior, it outperforms a random prior and FrameDiff on W2 for all angles and residue 56.Training uses 200,000 frames sampled every 5 ns from a 1 ms molecular-dynamics trajectory; evaluation uses 20,000 unseen frames. AlphaFold2 does not model the flexibility well.Source: §5.3 · pp. 8–9, Table 3 · p. 8, Figure 5 · p. 9

Protein design approaches

• Protein design spans libraries, sequence-based ML, structure-based biophysics, and increasingly prominent diffusion approaches.Design libraries developed over decades, alongside ML assistance. Sequence ML yielded successful designs, while biophysical methods produced drug candidates; diffusion approaches improved biological experiment success over previous SOTA.Source: §6 · p. 9
• SE(3) diffusion has been applied to protein-ligand binding, docking, and robotics, while FrameFlow concurrently investigates a model similar to FoldFlow-Base.Applications include protein-ligand binding (Jin et al., 2023), docking, and robotics (Brehmer et al., 2023); FrameFlow is described as similar to FoldFlow-Base.Source: §6 · p. 9

Equivariant generative models

• Equivariant generative work includes vector-field CNFs, affine-coupling finite flows, physics applications, and newer equivariant score-matching and diffusion models.Examples include equivariant CNFs (Köhler et al., 2020; Katsman et al., 2021; Garcia Satorras et al., 2021; Klein et al., 2023), finite flows (Dinh et al., 2017; Bose & Kobyzev, 2021; Midgley et al., 2023), and physics applications (Boyda et al., 2020; Kanwar et al., 2020; Abbott et al., 2023). Score matching and diffusion work includes De Bortoli et al. (2022), Brehmer et al. (2023), Hoogeboom et al. (2022), Xu et al. (2022), and Igashov et al. (2022).Source: §6 · p. 9

7 CONCLUSION

• FOLDFLOW-BASE learns deterministic dynamics over SE(3), FOLDFLOW-OT uses Riemannian optimal transport for more stable flows, and FOLDFLOW-SFM learns stochastic dynamics as an SDE.The stochastic model is motivated by learning Brownian bridges over SO(3) in a simulation-free manner.Source: §7 · p. 9
• On PDBs with up to 300 amino acids, FOLDFLOW is competitive with RFDiffusion and significantly outperforms FrameDiff-Improved on all metrics.The paper also identifies equilibrium conformation sampling as a protein-design subtask for which FOLDFLOW is more amenable.Source: §7 · p. 9
• FOLDFLOW is more amenable to equilibrium conformation sampling, an important subtask in protein design.Source: §7 · p. 9
• A natural future direction is extending FOLDFLOW to conditional generation using target sequence and structure during training.Source: §7 · p. 9

A.1 SO(3) LIE GROUP

• On SO(3), rotations can be represented by axis-angle vectors, and a canonical choice of metric is bi-invariant.The rotation-vector magnitude gives the angle, its direction gives the axis, and the canonical metric uses Q = 1/2I.Source: §A.1 · p. 17

A.2 SE(3) LIE GROUP

• With the stated metric, SE(3) decomposes into SO(3) and R3, so its geodesics follow those of the product manifold.The Lie algebra is isomorphic to R6, and the metric is not bi-invariant.Source: §A.2 · p. 18

A.3 THE ISOTROPIC GAUSSIAN DISTRIBUTION ON SO(3)

• IGSO(3) sampling draws a uniform rotation axis and uses inverse-transform sampling for the angle, with a closed-form approximation making sampling efficient.The angle’s cumulative density function is approximated using the density described for IGSO(3).Source: §A.3 · p. 19

B FLOW MATCHING IN Rd

• FoldFlow’s Euclidean flow-matching setup uses Gaussian conditional paths, while OT-CFM uses transport-plan pairs to obtain straighter flows for faster inference and more stable training.The standard path runs from a standard normal to a Gaussian centered at a data point; the OT-CFM path sets its Gaussian mean by interpolating source and target points.Source: B FLOW MATCHING IN Rd · p. 19

C RIEMANNIAN OPTIMAL TRANSPORT

• The section situates optimal transport in generative modeling and introduces its formulation on Riemannian manifolds.It also states that it will discuss empirical computation and prove Proposition 1.Source: C RIEMANNIAN OPTIMAL TRANSPORT · p. 20

C.1 BACKGROUND ON RIEMANNIAN OPTIMAL TRANSPORT

• Riemannian Kantorovich transport minimizes displacement cost, and under stated manifold and density conditions its optimal map exists uniquely and induces McCann interpolation.The condition is a smooth compact manifold without boundary and a source distribution with a density; the optimal plan is supported on the graph of the Monge map.Source: §C.1 · p. 20
• Minibatch OT reduces the cubic-time, quadratic-memory burden of exact transport, and small batches can still perform well empirically.For each source and target minibatch, training samples pairs from the minibatch transport plan; this approximation introduces error relative to exact OT.Source: §C.1 · p. 20

C.2 PROOF OF PROPOSITION 1

• With compactly supported distributions and an absolutely continuous source, the Kantorovich and Monge problems have a unique solution represented by an almost-everywhere-unique transport map.The map has the form Ψ(x) = exp_x(∇ϕ(x)) for a d^2-concave function ϕ.Source: §C.2 · p. 21
• The proof checks that SE(3)_0 is connected, complete, smooth, and without boundary by expressing it through mean-subtracted SE(3) components and a mean component.The argument uses disintegration of measures to define a measure proportional to R3 for the final component.Source: §C.2 · p. 21

D.1 BROWNIAN BRIDGE ON SO(3)

• A Brownian bridge on a metric Lie group is simulated with a guided diffusion conditioned to reach v at t = 1.For a bi-invariant metric, as on SO(3), the Riemannian logarithm coincides with the Lie group logarithm.Source: §D.1 · p. 21
• On SO(3), the bi-invariant metric gives V0 = 0; a process-independent guided diffusion has identical Stratonovich and Itô formulations and yields the reversed process in Eq. 7.Source: §D.1 · p. 21

D.2 SIMULATION-FREE APPROXIMATION OF BROWNIAN BRIDGES ON SO(3)

• FoldFlow-SFM closely tracks simulated Brownian bridges across full SO(3) trajectories, supporting its use as a fast, drop-in approximation.The comparison covers three diffusion coefficients; the mean and standard deviation of the SO(3)-norm are plotted over 1024 data points, with overlapping standard-deviation regions.Source: §D.2 · pp. 21–22, Figure 6 · p. 22
• The SO(3)-norm comparison reports trajectory means and standard deviations over 1024 data points for three diffusion coefficients, with overlapping shaded standard-deviation regions.Source: §D.2 · p. 21, Figure 6 · p. 22

D.3 PROOF OF PROPOSITION 2

• The conditional and unconditional FoldFlow-SFM losses have equal gradients with respect to θ.The passage states this equality for N 0.Source: §D.3 · p. 22
• The proof reduces the claim to SO(3) and R3 by factorizing the probability distributions and the SE(3) metric.The passage says the factorized probabilities follow from disintegration of measures and cites Pollard (2002) and Proposition 3.5 from Yim et al. (2023b).Source: §D.3 · p. 22
• The SO(3) proof follows a structure similar to Chen & Lipman (2023), while the R3 proof directly follows Theorem 3.2 from Tong et al. (2023a).The cited proof approaches are described separately for the SO(3) and R3 components.Source: §D.3 · p. 22

E EXTENDED FIGURE INFORMATION

• FoldFlow-Base paths may cross, whereas conditional FoldFlow-OT paths do not cross, reducing objective variance and stabilizing training (Pooladian et al. (2023b); Tong et al. (2023b)).The comparison is shown for probability paths projected onto S2.Source: E EXTENDED FIGURE INFORMATION · p. 23
• FoldFlow-SFM adds stochasticity that improves novelty in protein generation, and achieves optimal transport only as noise goes to zero, recovering FoldFlow-OT.The passage notes that the associated bias may still help reduce objective-function variance even without achieving optimal transport.Source: E EXTENDED FIGURE INFORMATION · p. 23

F.1 TOY MODEL PARAMETERIZATION

• The vector field is parameterized by applying Π_r to the MLP output, projecting its 3 × 3 matrix onto T_rSO(3), which computes its skew-symmetric component M − Mᵀ.Source: §F.1 · p. 23
• The projection extracts the matrix's skew-symmetric component and transports it to the tangent space of r by left matrix multiplication.The passage identifies left matrix multiplication as the group operation on SO(3).Source: §F.1 · p. 23

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

• All three FoldFlow models learn the modes of the SO(3) toy distribution, with slight model shrinkage in FoldFlow-Base.Figure 7 visualizes the data and model distributions using Euler angles.Source: §F.2 · p. 23, Figure 7 · p. 23

G.1 PROTEIN BACKBONE GENERATION EXPERIMENT ADDITIONAL RESULTS

• Supplementary visuals show rotation-flow norms with and without t scaling and five SFM-generated proteins per length, paired with best-of-eight ProteinMPNN sequences refolded by ESM that retain diverse folds.The refolded structures vary in secondary structure and overall 3D conformation.Source: §G.1 · p. 23, Figure 8 · p. 24
• For 250- and 300-residue backbones, FoldFlow models are comparable in designability but more diverse and novel than other models, while RFDiffusion remains more designable.The comparison includes RFDiffusion; Figure 9 plots designability, diversity, and novelty across backbone lengths.Source: Figure 9 · p. 24, §G.1 · p. 24
• At shorter lengths (≤200), FoldFlow closes the designability gap with models trained without pretraining, but designability trades off against diversity and novelty.The reported trade-off occurs at short lengths and as sequence length increases.Source: §G.1 · pp. 23–24, Figure 9 · p. 24

G.1.1 TIMING COMPARISON IN TABLE 2 AND TABLE 4

• FoldFlow is over 2x faster per step than FrameDiff, largely because it avoids the costly IGSO(3) score computation required by FrameDiff.A step is a forward and backward pass on the effective batch size on one GPU.Source: §G.1.1 · p. 24
• FoldFlow was trained for roughly 2.5 days on four NVIDIA A100-80GB GPUs, compared with RFDiffusion’s pretrained weights from a four-week run on 64 V100 GPUs.The paper describes FoldFlow’s training time as substantially less than that of comparable models.Source: §G.1.1 · p. 24

G.2 EQUILIBRIUM CONFORMATION GENERATION EXPERIMENT

• On BPTI, informed-prior FoldFlow captures both conformational modes and improves performance over FoldFlow-Rand, while both variants outperform FrameDiff, which focuses on one mode.The experiment uses a noisy mixture of folded conformations as its informed prior; the authors report low, uniform KL divergence for FoldFlow’s generated angle distribution and describe training for one day on four A100 GPUs.Source: §G.2 · pp. 26–27, Figure 12 · p. 26

H FURTHER DISCUSSION OF FOLDFLOW AND RELATED MODELS

• Intrinsic tangent-space parameterization makes the SE(3) vector field equivariant; flow matching transports arbitrary source-to-target distributions, supports optimal transport, and generally enables more efficient ODE inference than diffusion SDE simulation.The discussion also notes that diffusion typically requires a Gaussian-like source, while entropic-regularized OT for diffusion yields stochastic paths rather than optimal-transport paths.Source: H FURTHER DISCUSSION OF FOLDFLOW AND… · pp. 27–28
• FoldFlow additions improve all metrics over FrameDiff-Improved; FrameDiff-Improved fixes a rotation-score bug, changes cluster sampling and separates axis-angle loss, while FrameDiff retains costly IGSO(3) scoring and imperfect training-time translation invariance.The authors say cluster-based sampling increases diversity and the separate axis-and-angle rotation loss reduces variance; they note FrameDiff centers structures at inference but not during training.Source: H FURTHER DISCUSSION OF FOLDFLOW AND… · p. 28

I.1 TRAINING AND INFERENCE

• FoldFlow-SFM training uses uniform rotations and centered Gaussian translations, OT-resampled endpoints, geodesic rotation and linear translation interpolants, then simulation-free noise; zero noise recovers OT, and removing resampling recovers Base.The rotational perturbation uses IGSO(3), while the translational perturbation is Gaussian.Source: §I.1 · pp. 28–29

I.2 SDE TRAINING AND INFERENCE

• For SO(3), FoldFlow-SFM trains on OT-resampled endpoints with geodesic interpolation and simulation-free IGSO(3) perturbation, then infers by integrating the vector field with rotational Brownian noise.Inference also accepts an annealing schedule, noise scale, and integration step size.Source: §I.2 · p. 29

I.3 VECTOR FIELD PARAMETRIZATION

• For protein modeling, the vector-field output lies in the tangent space, providing the stated geometric constraint for the architecture.Source: §I.3 · p. 29

I.3.1 PROTEIN MODEL PARAMETERIZATION

• Protein parameterization uses the FrameDiff architecture (Yim et al., 2023b), based on AlphaFold2’s structure module, and transforms predicted x_0 into a tangent-space vector.The rotation component uses the log map, and the translation component is centered; the resulting product-space vector lies in the tangent space of SE(3)^N.Source: §I.3.1 · p. 30

I.4 PROTEIN TASK HYPERPARAMETERS

• Protein models use Adam at learning rate 10^-4, length-dependent effective batch sizes, auxiliary-loss weight 0.25, and rotation-loss weight 0.5 versus translation weight 1.0.Separating the rotation-vector loss into axis and angle components is reported to reduce variance and numerical instability.Source: §I.4 · p. 30

I.5 DATA AND DATA SAMPLING

• The dataset contains 22,248 PDB monomers of length 60–512 and resolution <5Å, after filtering out proteins with >50% loops; batches contain one length and sample uniformly across 4,268 clusters.The structures were downloaded July 20, 2023, using FrameDiff criteria; cluster similarity is 30%.Source: §I.5 · p. 30
• Figure 13 shows length and cluster-size distributions; uniform cluster sampling increases batch diversity but slightly shifts length frequencies, and two cluster-size bins contain about 14% of proteins.The cluster-size plot uses logarithmic axes.Source: §I.5 · p. 30, Figure 13 · p. 31

I.6 PROTEIN METRICS

• TM-score ranges from 0 to 1, with 1 indicating perfect alignment; scores >0.5 generally indicate similar folds, while scores <0.2 indicate unrelated structures.Its length-normalizing factor is d0(N) = 1.24(N −15)^1/3 −1.8.Source: §I.6 · p. 31
• RMSD uses distances between paired heavy atoms [N, Cα, C, O], depends on protein length, and is more sensitive to local errors and less to global misalignments than TM-score.RMSD is described as a more stringent designability filter than TM-score >0.5.Source: §I.6 · p. 31
• Designability requires at least one of eight ProteinMPNN sequences to refold with ESMFold to scRMSD <2.0 against the generated structure.The protocol generates 50 structures at each length in {100, 150, 200, 250, 300}; the threshold may be overly strict for longer backbones, and self-consistency can yield false negatives because ProteinMPNN and ESMFold are imperfect.Source: §I.6 · p. 31
• Diversity is the mean pairwise TM-score among designable generated structures at each length, with lower scores indicating greater diversity.Restricting the calculation to designable proteins avoids inflating diversity with dissimilar structures that may not be refoldable.Source: §I.6 · pp. 31–32
• Novelty uses minimum TM-score to training data and the fraction of proteins both designable and novel, defined by avg. max TM-score <0.5.Training datasets differ across models; larger datasets for Genie and RFdiffusion may inflate novelty relative to the paper’s test set.Source: §I.6 · p. 32
• Reported standard errors average per-length estimates for novelty and designability, but diversity standard errors are invalid because pairwise-distance estimates are correlated.The paper notes that Genie and RFdiffusion use substantially larger training data than FoldFlow and the FrameDiff comparators.Source: §I.6 · p. 32

• Inference annealing improves the performance of all FoldFlow models in all metrics.Source: Table 6 · p. 32, J EXTENDED ABLATION OF EXPERIMENTS · p. 32
• Stochasticity improves robustness and novel-protein generation in 7/8 settings, as measured by the fraction of novel and designable proteins.Source: Table 6 · p. 32, J EXTENDED ABLATION OF EXPERIMENTS · p. 32
• Optimal transport improves designability in 7/8 settings by reducing variance in the training objective.Source: Table 6 · p. 32, J EXTENDED ABLATION OF EXPERIMENTS · p. 32
• Auxiliary training losses improve model designability in 7/8 settings.Source: Table 6 · p. 32, J EXTENDED ABLATION OF EXPERIMENTS · p. 32