SE(3)-Stochastic Flow Matching for Protein Backbone Generation
Authors: Joey Bose
Tara Akhound-Sadegh
Guillaume Huguet
Kilian FATRAS
Jarrid Rector-Brooks
Cheng-Hao Liu
Andrei Cristian Nica
Maksym Korablyov
Michael M. Bronstein
Alexander Tong
Venue: ICLR 2024 spotlight (2024)
Keywords: Proteins; Equivariance; Riemannian; Flow Matching; Generative models
Primary area: generative models
TL;DR: A Flow Matching for the de novo design of protein backbones

Abstract: 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 $3\mathrm{D}$ rigid motions---i.e. the group $\mathrm{SE(3)}$---enabling accurate modeling of protein backbones. We first introduce $\text{FoldFlow-Base}$, a simulation-free approach to learning deterministic continuous-time dynamics and matching invariant target distributions on $\mathrm{SE(3)}$. We next accelerate training by incorporating Riemannian optimal transport to create $\text{FoldFlow-OT}$, leading to the construction of both more simple and stable flows. Finally, we design \foldflowsfm, coupling both Riemannian OT and simulation-free training to learn stochastic continuous-time dynamics over $\mathrm{SE(3)}$. Our family of $\text{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 $\mathrm{SE(3)}$. Empirically, we validate $\text{FoldFlow}$, on protein backbone generation of up to $300$ amino acids leading to high-quality designable, diverse, and novel samples.

Published as a conference paper at ICLR 2024
SE(3) STOCHASTIC FLOW MATCHING FOR PROTEIN BACKBONE GENERATION
Avishek (Joey) B 2.3: , Tara Akhound-Sadeghh.220 Jarrid Rector-Brooks 4.2,3, Cheng-Hao Liu Andrei Cristian Nica Maksym Korablyov", Michael Bronstein Alexander To.g4.2.3
1.233 Guillaume Huguet 4.2.3, Kilian Fatras 1.2,3
McGill University, "Mila, 3 'Dreamfold, *Université de Montréal, University of Oxford
ABSTRACT
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 flo-mmatching paradigm over 3D rigid motions- i.e. the group SE(3)-enabling 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 FOI.FFL.W.. to the construction of both more simple and stable flows. Finally, we design FOI.DFL.OW FFMMS. coupling both Riemannian OT and simulation free training to learn stochastic continuous- time dynamics over SE(3). Our family of FOLDFLOWW 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 up to 300 amino acids leading to high-quality designable, diverse, and novel samples.
leading
INTRODUCTION
Proteins are one of the basic building blocks of life. Their complex geometric structure enables specific inter-molecular interactions that allow for crucial functions within organisms, such as acting as catalysts in chemical reactions, transporters for molecules, and providing immune responses. Normally, such functions arise as result of evolution. With the emergence of computational techniques, it has become possible to rationally design novel proteins with desired structures that program their functions. Such methods are now seen as the future of drug design and can lead to solutions to long-standing global health challenges. Some recent examples include rationally designed protein binders for receptors related to influenza (Strauch al., 2017), COVID-19 (Cao et al., 2020a; Gainza et al., 2023), and cancer (Silva et al., 2019).
FoldFlow - Base
FoldFlow OT
FoldFlow SFM
Method
Any
QT Stochastic Score free
RFDiffusion
x
x
FrameDiff
x
X
ROLDFLOW-BASHH
FOLDFLOW-Or
FoLDFLOW-SFM
Figure 1: Left: Conditional prohability paths learned by WaI.FF.o--B... BASE (left), FOL.DFLOW-OT (mid), and FOLDFLOW-S FFM (right). We visualize the rotation trajectory of a single residue by the action of SO(3) on its homogenous space 8² Right: Table with the properties of each model: whether they can map from general source distribution. perform optimal transport, are stochastic, or require the score the IGso(3) density.
In protein engineering, the term de novo design refers a setting when new protein is designed to satisfy pre-specified structural and functional properties (Huang et al., 2016). Chemically, a protein is
"Equal Contribution. Corresponding authors: (joey. bose, t. tara, ex @drremfolldd

---

Published as a conference paper at ICLR 2024
sequence of amino acids (residues) linked into chain that folds into complex 3D structure under the influence of electrostatic forces. The protein backbone can be seen as N rigid bodies (correspond- ing to N residues) that contain four heavy atoms N C. -C O. Mathematically, each residue can be associated with a frame that adheres to the symmetries of orientation irrenttaioo-presevinne rigid transformations (3D rotations and translations), forming the special Euclidean group SE(3) (Jumper et al., 2021); the entire protein backbone is described by the group product SE(3)* x The problem of protein design can be formulated as sampling from the distribution over this group, which a perspective used in our paper. Recently, generative models have been generalized to Riemannian manifolds (Mathieu & Nickel, 2020; De Bortoli et al., 2022). However, they are not purpose-built to exploit the rich geometric structure of SE(3) Furthermore, several approaches require numerically expensive steps like simulating a Stochastic Differential Equation (SDE) during training or using the Riemannian divergence in the objective (Huang et al., 2022; Leach et al., 2022; Ben-Hamu et al., 2022).
Our approach. We introduce FOLDFLOW, a family of continuous normalizing flows (CNFs) tailored for distributions on SE(3)* (fig. 1). We use the framework of Conditional Flow Matching (CFM), a simulation- approach to learning CNFs by directly regressing time-dependent vector fields that generate probability paths (Lipman et al., 2022; Tong et al., 2023b). In particular, we introduce three new CFM-based models that learn S3) N -invariant distributions to generate protein backbones. In contrast to the previous SE(3) diffusion approach of Yim et al. (2023b), our FOLDFLOW is able to start from an informative prior. This enables new applications of generative models for protein design such as equilibrium conformation generation (Zheng al., 2023).
Main contributions. Our first model FOLDFLOW-BASE extends the Riemannian flow matching approach (Chen & Lipman, 2023) by introducing a closed-form expression of the ground truth condi- tional vector field for (3) needed in the loss computation theputtttioothuus greatly increasing speed and stabil- ity of training. Next, in FOLDFLOW-OT we accelerate the training of our base model by constructing shorter and simpler flows using Riemannian Optimal Transport (OT) by proving the existence of Monge map on SE(3)". ". Finally, we present our most complex simulation-frce model, FOLDFLOW SFM, which learns a stochastic bridge on SF(33) " Empirically, we validate our proposed models by learning to generate protein backbones of up to 300 residues. We observe that all FOLDFL.OW models outperform the current SOTA non-pretrained diffusion model in FrameDiff (Yim et al., 2023b) for ir-silico designability with FOLDFLOW-OL FOLDFLOW-STT SFM is competitive with the current gold-standard RFDiffusion (Watson et al., 2023) with a fraction of the compute and data resources. We highlight the importance of novel and designable proteins as a key goal in Al Al-powered drug discovery where useful drug candidates are necessarily be- yond the available training set (Schneider, 2018 8; Schneider et al., 2020; Marchand et al., 2022). Finally, we show the utility of FOL.FL.OW on equilibrium conformation generation by learning to simulate molecular dynamics trajectories starting from reference empirical distribution in comparison to an uninformed prior. Our code can be found at https: / /gi thub. , wwm/ cww/DreamFoll d/Fol lo.
being the most designable. Moreover, for novelty
BACKGROUND AND PRELIMINARIES
2.1 RIEMANNIAN MANIFOLDS AND LIE GROUPS
Riemannian manifolds. Informally, an n-dimensional manifold M is topological space locally equivalent (homeomorphic) to R". This implies that one has the notion of 'neighbourhood" but not of 'distance' ''agle' on M. The manifold is said to be smooth if it additionally has a C differential structure. At every point x E M, one can attach a langent space Ta- The disjoint union of tangent spaces forms the tangent bundle. A Riemannian manifold' (M,g) g) is additionally equipped with an inner product (Riemannian metric) 9s 'TMM x TzM + R on the tangent space T&M at each M. The Riemannian metric g allows to define key geometric quantities on .M such distances, volumes, angles, and length minimizing curves (geodesics). We consider functions defined on M and the tangent bundle, referred to as scalar- and vector fields, respectively. The Riemannian gradient is an operator V, C (M) + ((M) between the respective functional spaces. Given a smooth scalar field f E C (M), its gradient V,, E x(M) is the local direction of its steepest change.
Lie groups. A Lie group is group that is also differentiable manifold, in which the group operations o: : x G - G of multiplication and inversion are smooth maps. It has a left action Ln : + G defined by x + hoz that is topological isomorphism and whom the derivative is also an isomorphism between the tangent spaces on G. Since a group has an identity element, its tangent space is of special interest and is known as the Lie algebra G. The Lie algebra is a vector space with an associated bilinear
We tacitly assume M to he orientable, connected, and complete and admit volume form denoted dr.
2

---

Published as a conference paper at ICLR 2024
operation called the Lie bracket that is anticommutativve and satisfies the Jacobi identity. Lie algebras elements can be mapped to the Lie group via the exponential map exp: : 1 - G which has an inverse called the logarithmic map log G - 6. For matrix Lie groups where the group action is the matrix multiplication, the exp and log maps correspond to the matrix exponential and matrix logarithm. The orientation prreetrvinn rigid motions form the matrix Lie group (3) _ SO(3) x (R², +),, semidirect product of rotations and translations (see $A.1 $A.2, and Hall (2013) for details).
2.2 FLOW MATCHINN ON RIEMANNIAN MANIFOLDS
Analogous to Euclidean spaces, probability densities can be defined on Riemannian manifolds as continuous non-negative functions MM- - R that integrate to JM p(z)ddx =
Probability paths on Riemannian manifolds. Let P(M) be the space of probability distributions on M. A probability path Pt : 10 - P(M) is an interpolation in probability space between two distributions Po: P1 (MM indexed by continuous parameter f. A one- one-parameter diffeomorphism M is known as a flow on M and is defined as the solution of the following ordinary
4/2 :M -
differential equation (ODE): z the (xx) = uy (v)()))), with initial conditions (1) = 2, for u [0, x M + .M a time-dependent smooth vector field. We say the flow is generates pt if it pushes forward Po to P1 by following the iime-dependent vector field u, Mi.iee. P: [Wa] # (Po). As the is diffeomorphism, Pt verifies the famous continuity equation and the density can be calculated using the instantaneous change of variables formula for Riemannian manifolds (Mathieu & Nickel, 2020).
Riemannian flow matching. Given a probability path Pt that connects Po to Pl. and its associated flow 4/2. we can learn CNF by directly regressing the vector field u, with parametric one vy E (.M.. This technique is termed flow matching (Lipman et al., 2022, FM) and leads to a simulation-free training objective as long as Pt satisfies the boundary conditions Po Paata and P1 Pprior- Unfortunately, the vanilla flow matching objective is intractable as we generally do not have access to the closed-form of ue that generates pt. Instead, we can opt regress vo against conditional vector field uz (zx])), generating a conditional probability path pa (xxx)22) and use it to recover the target unconditional path: pe (2+x) = JM Pe (=))=) The vector field uy can also
be recovered by marginalizing of conditional vector fields: us (xx) = SM (x(x)))) (())=))(((s)))))) pe(aals)((s) dz. The Riemannian CFM objective (Chen & Lipman, 2023) is then,
pe(2)
(I) As FM and CFM objectives have the same gradients (Tong et al., 2023b), at inference, we can generate by sampling from ()i and using Uy to propagate the ODE backward in time.
Cyetm (0) = (x. (l xx) - (x,z,<²))3,
t ~~ 10)
2.3 PROTEIN BACKBONE PARAMETRIZATION
Our protein backbone parameterization follows the seminal work of AlphaFold2 (Jumper et al., 2021, AF2) in that we associate frame with each residue in the amino acid sequence. For a protein of length N this results in N frames that are Each frame maps a rigid transformation starting from idealized coordinates of four heavy atoms N ", O* E IR3, with C; = 0, being Figure 2: Protein backbone parametrization. centered at the origin, and is measurement of experimental bond angles and lengths (Engh & Huber, 2012). Thus, residue E [N] is represented as an action of x = =rr,, s²) E SE(() applied to the idealized frame [N, C., C, Oj' x' o [N", C., C O"]. To construct the backbone oxygen atom O, we rotate about the axis given by the bond between C, and C using an additional rotation angle s. Finally, we denote the full 3D coordinates of all heavy atoms as A E N4x3, An illustration for this backbone parametrization is provided fig. with rotations being parametrized as r v x w.
V
W
O'
FOLDFL.WW FOR CONDITION FLOW MATCHING ON SF(3)
We seek to learn an SE(3)* invariant density Pt by training a flow using the objective in eq. (1). To do so we can pushforward an SE(3)* (FF(3))-invaa))) -invariant source distribution P1 to the empirical distribution of proteins Po using an equivariant flow. One way to guarantee the existence of a canslation-invartee measure is construct subspace that is invariant to global translations. This can be achieved by simply subtracting the center of mass of all inputs to the flow (Rudolph et al., 2021: Yim et al., 2023b). Formally, this leads to invariant measure on SE(3), which is subgroup of SE(3)". We then note that (E)))) is product group and thus the Riemannian metric extends in natural way the product
3

---

Published as a conference paper at ICLR 2024
space: the exponential and logarithmic maps decompose across each manifold, and the geodesic distance in ( is simply the sum of each individual distance in the product. As such, a flow on SE(3), can be built from separate flows for each residue in the backbone, on SE(3), after centering.
Decomposing SE(3) into SO(3) and IR.. As Lie groups are manifolds, they can be equipped with a metric obtain a Riemannian structure. In the case of SE)(3 = SO(3) x (IR3, +) there are multiple possible choices, but a natural one is (F, F))()()))) = (t, (t,r')so(3) + (s, (s) s²) (see $A.2). Moreover, the disintegration of measures implies that every -in3aiinnrriant measure can be broken down to SO(3)- measure and measure proportional the Lebesgue measure on R.3 (Pollard, 2002). Thus, we may simply build independent flows on SO(3) and TR.. In this section, we focus on designing FOLDFLOW models on SO(3), as CFMs on IR are well-studied in Albergo & Vanden- Eijnden (2023); Lipman et al. (2022); Tong et al. (2023b) (we provide complete description in $B). We use the notation Pt. q, and I for densities whose support is determined by its context.
3.1 GOLDFLOW-BASE
To construct a flow on SO(3) that connects the target distribution Py to a source distribution Pr. must first choose parametrization of the group elements. The most familiar and natural parametrization is by orthogonal matrices with unit determinant (see $A.1 for a discussion on other parametrizations e.g. rotation- vector). The Lie algebra so(3) contains skew-symmetric matrices v that are tangent vectors at the identity of SO(3). The last important component that we require is choice of Riemannian metric for (3)).. A canonical bi-invariant metric for SO(3) can be derived from the Killing form (see $A.1), and is given by: (t, (t,)3))(() = tr(rr"T)
(3(3) conditional vector fields and flows. We seek to construct conditional vector field (relz)) lying on the tangent space (O)). that transports ru ~ Po to T1 ~ Pr. Following, Tong et al. (2023b) we set the conditioner to (ru, Next, we construct a flow th that connects to ri- We follow the most natural strategy which is to build the flow using the geodesic between ny and 71- For general M, the geodesic interpolant between two points, indexed by t. has the following form:
Tt = exp,, (t logr, (r1)).
(2)
For rotation matrices, eq. (2) involves computing the exp and log maps which are both infinite matrix power series. Unfortunately, controlling the approximation error of log,, map is computationally expensive as the de facto inverse scaling method for computing matrix logarithms requires estimating and calculating fractional matrix powers (AI-Mohy & Higham, 2012). Instead, we use a numerical trick by converting T1 to its axis-angle representation which gives a vector representation of t E (3) and, by definition, lives at the tangent space at the identity and is equivalent to log, (ri). Next, we can parallel transport i to the tangent space of ro since Lie algebras of all tangent spaces are isomorphic and SO(3) carries a free action which gives us the desired end result logr, (ri).
Given Fr. we can build constant velocity vector fields by directly leveraging the ODE associated with the conditional flow: & (r) = rt (Chen & Lipman, 2023). As result, computing u, (ri(z) boils down to computing the point r; along the ODE and taking its time derivative. In practice, taking the time-derivative to compute i4 = f amounts using autograd to compute the gradient during for- ward pass. We can overcome this unnecessary overhead without relying on automatic differentiation but instead by exploiting the geometry of the problem. Specifically, corresponding to the relative rotation between to and r., given by r. ro. We divide by / to get vector which is an element of (3) and corresponds to the skew-symmetric matrix representation of the velocity vector pointing towards the target r'. Finally, we parallel-transpee the tangent space (3) using left matrix multiplication by rt- These operations can be concisely written log,, (ro). where we our numerical trick to calculate the matrix logarithm. The closed form expression of the loss to train the 50(3) component of FOLDFLOW-B.. BASE is thus
calculate the s0( element
the velocity vector to
(FOLDPR.OwBBsa-S0003) (0) (1.-
[1, r() log,, (ro)/1|[socs)
(3)
In eq. (3) the conditioning distribution g(z) = <(ro, r1) is the independent coupling g(ro, r1) = Po/1 where P1 = (((((3)) and is left-invariant w.r.t. to the Haar measure. Also, note that the vector field in eq. (3) is the tangent space Uy E T,, SO(3) and the norm is induced by the metric on SO(3).
3.2 FOLDFLOW-OT
The conditional vector field (r(r1[2) generates the conditional probability path pr(rx (2) which deterministically evolves po to p1- However, there is no reason to believe the conditional probability
4

---

Published as a conference paper at ICLR 2024
path is optimal in the sense that it is teiett-minimizing curve, under appropriate metric, in the space of distributions P(SO(3)). We seek to rectify this by constructing conditional probability paths that are not only shorter and straighter, but also more stable from an optimization perspective. This is motivated by previous research ( (Tong et al., 2023b; Pooladian et al., 2023a) which has shown optimal transport to lead to faster training with lower variance training objective in Euclidean spaces.
To this end, we propose FOLDFLOW-
a model that accelerates FOLDFLOW- BASE by constructing conditional probability paths using Riemannian optimal transport. The interpolation measure Pt con- nects Pa - Pi and is built from Riemannian OT which solves the Monge optimal transport problem:
OT(po: pi) =
inf
SE(M) 2 (x,, V (xx))² ddo(2:).
Here is the geodesic cost induced by the metric (ef. eq. (25) in $A.2) and v pushforward map: Po + Pl. A related problem, called the OT-Kantorovich formulation, relaxes the Monge problem by looking for joint probability distribution ı minimizing the displacement cost of transporting po to P1 (see $C.1). The uniqueness of the Monge map over SLi(3,, N is guaranteed under some assumptions on the measures Pu. P1 as stated in the following proposition and proven in SC.2.
Proposition 1. Let us consider SF(3), with the product distance deEes) and two compactly supported probability distributions po: Pt E P(SE(3),) In addition, suppose that pa is abso- lutely continuous with respect to Riemannian volume form (i.e., Po < dx). Then for the distance c= 1 SE(3) the Kantorovich and Monge problems admit unique solution that is connected as follows π = (id x V # Po. where Vi is almost uniquely determined everywhere po- Furthermore, we have that (x) = exp_ exp=(Vx(x)) (Vo(x)) for some de SF(3)) (SR -concave function é.
Following this proposition, we define the McCann interpolants as p+(x)) (exp,(- (-(Vo(ar)) # Po While it is possible to approximate the Monge map and McCann interpolants using e-concave functions, it imposes practical limitations on the architecture of the flow (Cohen et al., 2021). Instead, use the correspondence between the Monge and Kantorovich problems and rely on the optimal transport plan т. Formally, we draw two samples from (() = q(Io, x1) == T( 21) and we compute for given frame (c(rt (ro r1) = 5(expr, (t logro (ri))), where ô is a Dirac. Since our choice of metric for SE( factorizes into metrics on SO(3) and 3R, we can use independent losses on rotations and translations simiila to FOLDFLOW-BAL long as each geometric quantity in ı is coupled properly. Defining ( TI) as the projection of T (Zo, 2 on SO(3), we present the SO(3) loss for single frame in FOLDFLOW-OF as
repeat this over N frames, as
CFOLDHLOW-OT-SOrs) (0) E4.. (1 [[,((trr)))
logr, (ro)/ (5)
3.3 FOLDFLOW-SFM
We finally present FOLDFLOW-SFN. which builds on the foundations of both FOLDFLOW-BASE and FOLDFLOW-OT. Departing from the deterministic dynamics of the previous models, we aim to build stochastic flow over SE(3), by replacing these deterministic bridges with guided stochastic bridges. Previous research has shown that, compared to ODEs, SDEs have the important benefit of being more robust noise in high dimensions (Tong et al., 2023a: Shi al., 2023; Liu al., 2023a). For proteins in SE(3), this means the generative process has greater empirical calibre to sample outside the support of the training distribution-crucii for generating designable proteins that are also novel. In Euclidean space, we can build translation invariant flow on (R))) by using a (reverse time) Brownian bridge as the conditional flow between the points,
dS, S. 80 dt + (t)dW,, S =1=
(6)
This flow, also known as Doob's 's h-transform (Doob, 1984). is easy sample from in a simulation-free manner and correctly maps between arbitrary samplable marginals in expectation. Specifically, we can build simulation-free bridge by sampling from the conditional probability (r( (z0, xi) A (1sg lot + (1 l) 1²1(1 - 1)) (Shi et al., 2023; Albergo et al., 2023). See $B for more details.
Brownian Bridge on SO(3). On SO(3), we aim to model the dynamics between rotations matrices by a guided diffusion bridge (Jensen et al., 2022; Liu et al., 2022), leading to the following dynamics,
logR, To dt + (t). R1 = I1,
dR. =
(7)
t
5

---

Published as a conference paper at ICLR 2024
for B, the Brownian motion on (3) (see $D.1 for further technical details on this SDE). Despite the close resemblance of this SDE to the translation one in eq. (6), the corresponding Brownian bridge-to the best of our knowledge- does not have closed-form expression for pu(rs) (ro T1). Thus, to sample from Pe (Tr 71) correctly, we start at T1 and simulate the bridge backward in time using algorithm 3 in $1.2. Given this form of the conditional bridge, matching loss to optimize a flow between source and target distributions on SO(3) as follows:
make use of
SSFM-S0(3) (0) = (1t.
noo( Fx) logr, (ro) (003)
Here Fy is a sample from the bridge between ro and r1. When π T( is a valid coupling between po and Pi-and thus F(ro,7)) on SO(3) this objective is equivalent in expectation to matching directly the (computationally intractable) marginal loss LUSFM Ex- ((..) [|va(L, ft) =u(l,Fr)llson (r))lssa))))) The correctness of this approach is established in the next proposition and proved in $D.3.
Proposition 2. Given pc(a)) > O, Vx E SE(3) the conditional and unconditional FOLDFLOW- SFM losses have equal gradients w.r.t. 0: VoLusEM
V voLsFM(0)
This result allows us to learn a stochastic flow from any source any target distribution supported on SE(3), only requiring samples from both distributions. However, it does require simulation of an SDE sample from the conditional probability (r(relro, r1), limiting scalability.
An Efficient Simulation-free Approximation. Unfortunately, sampling from the correct conditional bridge requires simulation and is thus computationally expensive for training. In practice, we use a simulation-free approximation that closely matches the true conditional probability path on SO(3), (e (Fo 71). Specifically, we approximate Pt with the simulation-frr free alternative,
(F r1) IGso(a) Fr; expro (t logr, (()))(1(()((11--))))))))))))))))
where IBso(3) denotes the isotropic Gaussian distribution SO(3). This distribution can be seen as an analog of the Gaussian distribution in IR.. is the heat kernel on SO(3) (Nikolayev & Savyolov, 1900) and it can be seen as the limit of small i.i.d. rotations in 3D (Qiu, 2013). Additionally, has some of the desirable properties of the normal distribution, such as being closed under convolution (Nikolayev & Savyolov, 1900). We provide both the training and sampling algorithms for FOLDFLOW-SFI in SI, details on IGso(3) in $A.3, and results on the approximation error in $D.2.
4 MODELING PROTEIN BACKBONES USING FOLDFLOW
To model protein backbones using FOLDFLOW models we parameterize the velocity prediction (t, xx) as a function that consumes protein o on the conditional path at time and predicts the starting point åo. Specifically,, the predicted velocity is (t xx) VV )2/², with = (t, 2t). This choice of parameterization has two principal benefits. (1) It allows the usage of specialized architectures specifically designed for structure prediction, and (2) it allows for auxiliary protein- -spccffc losses to be placed directly on the èo to improve performance.
Architecture. Following by Anand & Achim (2022); Yim et al. (2023b) (FrameDiff) we use the structure module of AF2 model Wg. This begins with time-dependent node and edge embeddings (o, 24)) and 2t, followed by layers of invariant point attention. We use a small MLP head on top of the node embeddings to predict the torsion angle of the oxygen (p و MLP 2:)).
Full Loss. Our FOLDFILOW models are trained to optimize a flow-matching loss on SO(3) and R3 for each residue i E [N] in the backbone. These are denoted LFOLDFLOW and LFOLDFLOWEES (see $B for the complete expression in IR²) respectively. In addition to the flow-matching losses, we also include auxiliary losses from Yim et al. (2023b) which enforce good predictions at the atomic level in the R atomic representation A. These include a direct regression on the backbone (bb) positions Lss and a loss on the pairwise atomic distance in a local neighbourhood LAD.
|1(D 6A) D)1l²
Laux Eg (Lob + LaD), Lab
1 4N
Σ
1140 doll",
LaD
EIDeGA N
Here Q(t, To. 21, ..) == u(0 1) & ( 21) 8 (TT x1) is the factorized joint distribution, 1 is the indicator function, A is in Angstroms (Å), D is an N x N x x 4 tensor containing the pairwise distances between the four heavy atoms, i.e. Dijah Ai A and Đ is defined similarly from Å. We only apply auxiliary losses for t < 0.25, with scaling Aaux for a final loss of:

---

Published as a conference paper at ICLR 2024
Figure 3: FOLDFLOW- -FLLO generated structures in green compared to ProteinMPNN - ESMFold refolded
structures in grey. Samples with RMSD 2Å for lengths 100, 150, 200, 250, 300 from left right.
FFOLDFLOW (0) CFOLDILOW-SOO LFOLDHLOW R + 1(t 0.25) Aaux
(10)
EXPERIMENTS
5.1 SO(3) SYNTHETIC DATA
We evaluate all our FOLDFLOW models on synthetic multimodal densities on (3(3) as done by Brofos et al. (2021) (see SF.1 for details). We report the Wasserstein distance between generated and ground truth samples in table 1 and visualize the generated samples in fig. 7 in $F.2. We find that our proposed methods correctly model all the modes of the ground truth distribution. However, FOLDFLOW- BASE exhibits mode shrinkage relation to the ground truth. FOLDFL.OW- -O.TOOTTT and the simulated SDE results in comparable performance, with the OT-based method being the best. Importantly, this shows that our simulation-free approximation of the SDE does not hinder model performance, and combined with its significant speedup justifies use in protein experiments.
Table 1: Mean and std of the 1- and Wasserstein distances, computed against 5000 points the test set, over seeds.
(4 o)
Wi (x10 ?)
Wy (x10
FOLDFLOW-BAL FOLDFLOW-OT
0.27
1.96 + 0.27
1.25 + 0.12
FOLDFLOW-SFI
4.92
0.49
Simulated SDE
5.13 1. 1.36
1.33 1. 0.44
5.2 PROTEIN BACKBONE DESIGN
Table 2: Comparison of Designability (fraction of proteins with scRMSD 2.0 and mean scRMSD), Diversity (avg. pairwise TMscore), Novelty (max. TM-score to PDB and fraction of proteins with averaged max. TMscore 0.5 and scRMSD < 2.0Å). Designability and Novelty metrics include standard errors." *RFiiffusion and Genie have larger training sets likely overestimate novelty with respect our dataset.
Designability
Novelty
Diversity (1)
iters / (1)
Fraction (*)
scRMSD (L)
Fraction (?)
avg. max TM (L)
RFDiffusion
0.969 + 0.023
0,650
0.136
*0.708
0.060
*0.449
0.012
0.256
Genie
0,581 1: 0.064
2,968 1 0.344
*0.556
0.093
*0.434 1 0.016
0.228
FrameDiff ICMI.
0.402 + 0.062
3,885
+ 0.415
0.176
0.124
0.542 + 0.046
0.237
FrumeDiff-Improved
0.555 + 0.071
2.929 + 0.354
0.296
0.112
0.457 0.026
0.278
FrameDift- Retrained
0.612 + 0.060
2,990 + 0.307
0,108 + 0.083
0.684 + 0.032
0.403
1.278
BOLDFLOW-BAST
0.657 + 0.042
3,000 = 0.271
0.432 0.074
0.452 = 0.024
0.264
2.674
FOLDFLOW-OT
0.820 1: 0.037
1,806
1 0.249
0.484 0.068
0.460 + 0.020
0.247
2,673
FOLDFLOW-SFM
0.716 + 0.040
2.296 + 0.391
0.544 0.061
0.411 = 0.023
0.248
2.647
We evaluate FOLDFLOW models in generating valid, diverse, and novel backbones by training subset of the Protein Data Bank (PDB) with 22.248 proteins. We compare FOLDFLOW to pretrained versions FrameDiff (Yim et al., 2023b) (FrammDiffIIMML) the improved version on the authors' GitHub (FrameDDifIImprooeel Improved), Genie (Lin & AlQuraishi, 2023), and RFDiffusion, which is the gold standard (Watson et al., 2023). We also retrain FrameDiff (FrameDiff-Retrained) on dataset, which contains 100% more admissible structures, while inheriting the majority of the hyperparameters of FOLDFL.OW. We provide a detailed deseription of all the metrics $1.6. Figures 3 and 10 visualize generated samples and ESM-refolded structures.
We report our findings in table and observe that FOL.D outperforms SifmeeDiff-Reariinnng on all three metrics. We identify FrameDiff as the most comparable baseline as it the current SOTA model that does not utilize pre-training while using comparable resources. In contrast to FOLDFLOW we highlight that RFDiffusion uses a pre-trained backbone and a significantly larger model (60m vs. 17m parameters), training set, and compute resources (1800 vs. 10 GPU days). We also note that Genie is trained on a larger dataset (195k vs. 22k), which hinders rigorous comparisons with FOLDFLOW. Next, we analyze the performance of FOLDFLOW on each metric in detail.
Designability. We measure designability using the self-consistency metric with ProteinMPNN (Dau- paras al., 2022) and ESMFold (Lin al., 2022), counting the fraction of proteins that refold (C.- RMSD (scRMSD) 2.0À and mean scRMSD) over 50 proteins at lengths (100, 150, 200, 250, 250.
7

---

Published as a conference paper at ICLR 2024
200
175
1.6
Designahility (?)
Stochas, OT
Aux. Loss
lnf. nealing
1.5
120
0.716
/
0.820
x
0.657
x
x
0.648
x
0.228
x
x
X
X
0-2
Nax to Thirrst
(a)
6. 0.8 10
Sequercc Length
(b)
(c)
Figure 4: (a) Count of novel and designable (scRMSD < 2) proteins out of 250 at various novelty thresholds
with TM-score to the training set < for models trained on subsets of PDB. (b) scRMSD of designed proteins vs. ESMFold under flow scaling (c) Table showing an ablation study of FOLDFLOW features against designability.
In table 2, we find that all FOLDFLOW models achieve significantly higher Frac. designability score than all Frame models, and appreciably close the gap to RFDiffusion, c.g. A = 0.149 vs. A 0.357 for FOLDFLOW-OT and FrameDiff-Retrained respectively. When retrained on our dataset with 10% more samples, we find that Frame is more designable, but is still below all FOLDFLOW models. We also note that while FOLDFLOW- (excluding RFDiffusion), it has relatively low diversity and novelty. We find that adding stochasticity with FOLDFLOW-SFT results in a model that beats PrameDiff-Improvee dramatically improve novelty at the cost of worse designability (table 2). In fig. 9 in $G.1, we plot designability versus sequence length and observe the largest gains on sequence lengths < 300.
creates the highest fraction of designable proteins
on every metric and can
Diversity. We use the average pairwise TM-score of the designable generated samples averaged across lengths as our diversity metric (lower is better). We find an inverse correlation between performance on designability and diversity metrics for FrameDiff models, which interestingly does not hold for FOLDFLOW models. We note that FOLDFLOW models have comparable diversity to the baselines with FOLDFL.OW-OT OT and FOLDFI.OW- SFM being the most diverse.
Novelty. Designing novel but realistic protein structures compared to the training data is also an important goal. Unlike conventional generative modeling problems, e.g. images, the novelty of proteins is particularly important since the entire premise of ML-driven drug discovery requires developing original drugs that may be vastly different than current human knowledge (training data) but also synthesizable (designable) (Marchand et al., 2022; Schneider al., 2020; Schneider, 2018). We measure novelty using two metries: I.) the fraction of designable proteins with TM-score < 0.5 as used in Lin & AlQuraishi (2023) (higher is better) and 2.) the average maximum TM-score of designable generated proteins to the training data (lower is better). In fig. 4a count the number of designable proteins
function of the Max TM-score the training set. We see that FOLDFLOW- SFM designs the most novel structures against all methods including RFDiffusion and Genie. This substantiates the hypothesis that the stochasticity of the learned SDE is crucial to the improved robustness in high dimensions of FOLDFLOW SFM versus FOLDFLOW-BAAS -BSSE, and FOLDFLOW-OT which allows it to sample designable proteins far outside of the support of the training set.
Inference Annealing. We now describe numerical trick during inference that greatly improves the designability of FOLDFLOW which we term inference annealing. Instead of following the theoretical ODE SDE for generating rotations, use multiplicative scaling of the velocity- dR, = i(t)vo (t, R.) ) dt + (t)dB, for some positive function it)). In practice, we use i(t) = ct for some constant c. This annealing removes unwanted increase in the flow norm during the end of inference (fig. 8). We observe that larger drastically increases the designability of FOLDFLOWW (fig. 4b). In practice, we values of 10, which leads to designable yet diverse structures.
Ablation study of FOLDFLLOW. Next, we ablate various additions to the FOLDFLOW-BB BASE model in fig. 4c and report the full extended ablation in table 6. We find FOLDFLOW-OT creates the most designable model, but adding stochasticity helps increase novelty and diversity. We also find that infer- ence annealing is critical to the performance of the model in terms of achieving higher designability.
5.3 EQUILIBRIUM CONFORMATION GENERATION
Table 3: Wi in angle space be- tween generated and test samples.
Modeling various protein conformations is crucial in determining bi- ological behaviours such as mechanisms of actions or binding affinity to other proteins. Unlike diffusion models, FOLDFLOW can easily be instantiated from any sampleable source distribution. To test this,
Wo
We856
FOLDFLOW
4,379
0.406
Rodd
4,446
0.557
FrameDiff
4.844
0.800
model the equilibrium distribution of protein given initial predicted structures from pre-trained
8

---

Published as a conference paper at ICLR 2024
folding models including OmegaFold (Wu et al., 2022b) and ESMFold (Lin et al., 2022). The training target distribution consists of 200,000 frames at ōns intervals a Imss molecular dynamics trajectory of the BPTI protein (Shaw et al., 2010); the inference of FOLDFLOW is tested against 20,000 unseen frames in that trajectory. FOLDFLOW can successfully model both the general set of conformations, as indicated by ICA of the dihedral angles in fig. 5b, well the highly flexible residues, in the 2D Ramachandran plot in fig. 5a. We note that our approach can capture all the modes of distribution in contrast AlphaFold2, which does not model the flexibility well (fig. 5c). In table 3, we observe under the 2- 2-Wasserstein (W2) metric for all angles and only residue 56, FOLDFLOW with an informed prior outperforms random prior, and FrameDiff which can only use uninformed prior which further highlights a key advantage of FOLDFLOW over diffusion beadee approaches.
20
TUP
W
1.5
3/6
930
10
326
-0.5
316
-1.0
308
16
-25 -10 -05
15
28
300
(a)
(c)
Figure 5: (a) Ramachandran plot of & and V/ of the most flexible residue (56) in BPTI (b) ICA of all dihedral angles of BPTI (c) 1000 BPTI conformations sampled by FOLDFLOW with C. alignment highlighted in yellow and AlphaFold2 samples in green. FOLDFLOW reproduces test MD frames while AlphaFold2 samples do not.
6 RELATED WORK
Protein design approaches. The field of protein design has evolved over the course of several decades with many useful libraries (McCafferty et al., 1990; Winter et al., 1994; Romero & Arnold, 2009; Wang al., 2021) with subfields being impacted by assitance (Yang et al., 2019). Sequence- based machine learning approaches resulted in multiple successful protein design cases (Madani et al., 2023; Verkuil et al., 2022; Alamdari et al., 2023; Hie et al., 2022). Structure- approaches resulted in several drug candidates (Röthlisberger et al., 2008; Fleishman et al., 2011; 1: Cao et alaa,, 2020b). Moreover, diffusion-based approaches have risen in prominence (Wu et al., 2022a; Yim et al., 2023b), including significantly improved biological experiment success compared to previous SOTA (Watson et al., 2023). SE(3) diffusion has also seen applications et al., 2023), docking (Somnath et al., 2023), as well as robotics (Brehmer et al., 2023). Lastly, FrameFloww (Yim et al., 2023a) concurrently investigates a model similar to FOLDFL.OW- BASA.
biophysics
protein-ligant binding (Jin
Equivariant generative models. There have been several efforts to incorporate symmetry constraints in generative models. These include building equivariant vector fields for CNFs (Köhler et al., 2020; Katsman et al., 2021: Garcia Satorras et al., 2021: Klein et al., 2023) and finite flows using the affine coupling transform (Dinh et al., 2017: Bose & Kobyzev, 2021: Midgley et al., 2023). Applications in theoretical physies have also been impacted by equivariant flows (Boyda et al.. 2020; Kanwar et al., 2020; Abbott et al., 2023), Lastly, beyond flows a new genre of models coming into prominence based on the idea of equivariant score matching (De Bortoli et al., 2022: Brehmer et al., 2023) and diffusion models (Hoogeboom et al., 2022; Xu et al., 2022; Igashov et al., 2022).
7 CONCLUSION
In this paper, we tackle the problem of protein backbone generation using SE(3)* N -invariant generative models. In pursuit of this objective, we introduce FOLDFLOW, family of simulation ffree generative models under the flow matching framework. Within this model class, we introduce WOLDFLOW-BASS which learns deterministic dynamics over SE(3), FOLDFLOW- OT which learns more stable flows using Riemannian OT. To learn stochastic dynamics, we propose FOLDFLOW-SFFN which learns an SDE over SE(3)* and is motivated by learning Brownian bridges over SO(3), but in a simulation-free manner. We investigate the empirical caliber of FOLDFL.OW models on PDBs that contain up to 300 amino acids and find that our proposed models are competitive with RFDiffusion while significantly outperforming the current non-pretrained SOTA approach, FrameDiff- Improee. on all metrics. Finally, FOLDFLOW is more amenable for equilibrium conformation sampling which is an important subtask in protein design. Beyond generating 3D structures, a natural direction for future work is to extend FOLDFLOW to conditional generation by using target sequence and structure during training.

---

Published as a conference paper at ICLR 2024
ACKNOVLEDGEMENTS
AJB was supported by the Ivado Phd fellowship. KF supported by NSERC Discovery grant (RGPIN- 2019-06512), CIFAR AI Chairs, and a Samsung grant. CHL is supported by the Vanier scholarship The authors would like to thank Clément Bonet and Gauthier Gidel for fruitful discussions on Riemannian geometry and optimal transport. We also thank Riashat Islam for providing feedback on early versions of this work. We thank Alexander Stein and the entire DreamFold team for providing a vibrant workspace that enabled this research. The authors would like to thank the restaurant Le Don Donburi for keeping our stomachs full and sustaining this research. Finally, the authors would like to acknowledge Anyscale and Google GCP for providing computational resources for the protein experiments.
REFERENCES
R Abbott, MS Albergo, A Botev, D Boyda, K Cranmer, DC Hackett, G Kanwar, AGDG Matthews, Racaniere, A Razavi, et al. Normalizing flows for lattice gauge theory in arbitrary space-time dimension (2023). arXiv preprini arXiv:2305.02402, 2023. (Cited on page 9)
Gustaf Ahdritz, Nazim Bouatta, Sachin Kadyan, Qinghui Xia, William Gerecke, Timothy J O' O'Donnell, Daniel Berenberg, Ian Fisk, Niccolò Zanichelli, Bo Zhang, Arkadiusz Nowaczynski, Bei Wang, Marta M Stepniewskaa-ziibbnnska Shang Zhang. Adegoke Ojewole, Murat Efe Guney, Stella Biderman, Andrew M Watkins, Stephen Ra, Pablo Ribalta Lorenzo, Lucas Nivon, Brian Weitzner, Yih-En Andrew Ban, Peter K Sorger, Emad Mostaque, Zhao Zhang, Richard Bonneau, and Mohammed AlQuraishi. OpenFold: Retraining AlphaFold2 yields new insights into its learning mechanisms and capacity for generalization. bioRxiv, 2022. doi: 10.1 11.11//222..1.0.5511111
(Cited on page 30)
Awad H Al-Mohy and Nicholas J Higham. Improved inverse scaling and squaring algorithms for the matrix logarithm. SIAM Journal on Scientific Computing. 34(4):C153-CI6) C169, 2012. (Cited on page 4)
Sarah Alamdari, Nitya Thakkar, Rianne van den Berg, Alex Xijie Lu, Nicolo Fusi, Ava Pardis Amini, and Kevin K Yang. Protein generation with evolutionary diffusion: sequence is all you need. biokxiv, pp. 2023-09 2023. (Cited page 9)
Michael Albergo and Eric Vanden- Rijnden. Building normalizing flows with stochastic interpolants. International Conference Learning Representations (ICIR), 2023. (Cited page 4)
Michael S. Albergo, Nicholas M. Boffi, and Eric Vanden-Bijnden. Stochastic interpolants: A unifying framework for flows and diffusions. arXiv preprint 2303.08797, 2023. (Cited on page 5)
Namrata Anand and Tudor Achim. Protein structure and sequence generation with equivariant denoising diffusion probabilistic models, 2022. (Cited on pages 6 and 30)
Martin Arjovsky, Soumith Chintala, and Léon Bottou. Wasserstein gan, 2017. (Cited on page 20)
Heli Ben-Hamu, Samuel Cohen, Joey Bose, Brandon Amos, Aditya Grover, Maximilian Nickel, Ricky TQ Chen, and Yaron Lipman. Matching normalizing flows and probability paths on manifolds. arXiv preprint arXiv:2207.04711, 2022. (Cited page 2)
Helen M. Berman, John Westbrook, Zukang Feng, Gary Gilliland, N. Bhat, Helge Weissig. Ilya N. Shindyalov, and Philip E. Bourne. The protein data bank. Nucleic Acids Research, 28(1):235-242, January 2000. ISSN 1305-1048, (Cited on page 30)
Avishek Jocy Bose and Ivan Kobyzev. Equivariant finite normalizing flows. arXiv preprint arXiv:2/ 1 10.08649, 2021. (Cited pages 9 and 27)
Denis Boyda, Gurtej Kanwar, Sébastien Racanière, Danilo Jimenez Rezende, Michael Albergo, Kyle Cranmer, Daniel C Hackett, and Phiala E. Shanahan. Sampling using su(n) gauge equivariant flows. arXiv preprint arXiv:2008. 05456, 2020. (Cited on page 9)
Johann Brehmer, Joey Bose, Pim De Haan, and Taco Cohen. Edgi: Equivariant diffusion for planning with embodied agents. Neural Information Processing Systems (2023), 2023, (Cited page 9)
10

---

Published as a conference paper at ICLR 2024
Yann Brenier. Polar factorization and monotone rearrangement of vector-valued functions. Commu- nications on Pure and Applied Mathematics, 44:375 417, 1991. (Cited on page 20)
James A. Brofos, Marcus A. Brubaker, and Roy R. Lederman. Manifold density estimation via generalized dequantization. 2021. (Cited on page 7)
Longxing Cao, Inna Goreshnik, Brian Coventry, James Brett Case, Lauren Miller, Lisa Kozodoy, Rita E. Chen, Lauren Carter, Alexandra C. Walls. Young-Jun Park, Eva-Maria Strauch, Lance Stewart, Michael S. Diamond, David Veesler, and David Baker. De novo design of picomolar sars-cov-2 miniprotein inhibitors. Science, 37066515):426-4 2020a. (Cited on page I)
Longxing Cao, Inna Goreshnik, Brian Coventry, James Brett Case, Lauren Miller, Lisa Kozodoy, Rita Chen, Lauren Carter, Alexandra C Walls, Young-Jun Park, et al. De novo design of picomolar sars-cov-2 miniprotein inhibitors. Science, (7)(6515):426-44) 431, 2020b. (Cited on page 9)
Ricky TQ Chen and Yaron Lipman. Riemannian flow matching general geometries. arXiv preprint arXiv:2302.03666 2023. (Cited on pages 2, 3, 4, 22, and 23)
Samuel Cohen, Brandon Amos, and Yaron Lipman. Riemannian Convex Potential Maps. In Marina Meila and Tong Zhang (eds.), Proceedings of the 38th International Conference Machine Learning, volume 139 of Proceedings of Machine Learning Research, pp. 5028-2038. PMLR, 18 24 Jul 2021. (Cited page 5)
Bharath Bhushan Damodaran, Benjamin Kellenberger, Remi Flamary, Devis Tuia, and Nicolas Courty. Dee:JDOT: Deep joint distribution optimal transport for unsupervised domain adaptation. European Conference on Computer Vision (ECCV), 2018. (Cited on page 20)
Dauparas, I. Anishchenko, N. Bennett, H. Bai, R. J. Ragotte, L. F. Milles, B. I. M. Wicky, A. Courbet, R. J. de Haas, N. Bethel, P. J. Y. Leung, T. E. Huddy, S. Pellock, D. Tischer, F. Chan, B. Koepnick, H. Nguyen, A. Kang, B. Sankaran, A. K. Bera, N. King, and D. Baker. Robust deep learning-based protein sequence design using proteinmpnn. Science, 378(6615))49 2022. (Cited on page
Valentin De Bortoli, James Thornton, Jeremy Heng, and Arnaud Doucet. Diffusion schrl"odinger bridge with applications to score-based generative modeling. NeurlPS, 2021. (Cited on page 20)
Valentin De Bortoli, Emile Mathieu, Michael Hutchinson, James Thornton, Yee Whye Teh, and Arnaud Doucet. Riemannian score-based generative modelling. Advances in Neural Information Processing Systems, 35:2406-2422, 2022. (Cited pages and 9)
Laurent Dinh, Jascha Sohl-Dickstein, and Samy Bengio. Density estimation using real nvp. In The 5th International Conference on Learning Representations (ICLR), Vancouver, 2017. (Cited on page 9)
Joseph Doob. Classical Potential Theory and Its Probabilistic Counterpart, volume 549. Springer, 1984. (Cited on page 5)
RA Engh and R Huber. Structure quality and target parameters. 2012. (Cited on page 3)
Kilian Fatras, Younes Zine, Rémi Flamary, Remi Gribonval, and Nicolas Courty. Learning with minibatch wasserstein: asymptotic and gradient properties. In Silvia Chiappa and Roberto Calandra (eds.), Proceedings of the Twenty Third International Conference on Artificial Intelligence and Statisties, volume 108 of Proceedings of Machine Learning Research, pp. 2131-2144 Online, 26-28 Aug 2020. PMLR. (Cited on page 20)
Kilian Fatras, Thibault Sejourne, Rémi Flamary, and Nicolas Courty. Unbalanced minibatch optimal transport; applications to domain adaptation. 139:3186-3197, 18-24 Jul 2021a. (Cited on page 20)
Kilian Fatras, Younes Zine, Szymon Majewski, Rémi Flamary, Rémi Gribonval, and Nicolas Courty. Minibatch optimal transport distances; analysis and applications, 2021b. (Cited page 20)
Chris Finlay, Jöm--enrik Jacobsen, Levon Nurbekyan, and Adam M. Oberman. How train your neural ode: The world of jacobian and kinetic regularization. International Conference on Machine Learning (ICMI), 2020. (Cited page 20)
11

---

Published as a conference paper at ICLR 2024
Rémi Flamary, Nicolas Courty, Alexandre Gramfort, Mokhtar Z Alaya, Aurélie Boisbunon, Stanislas Chambon, Laetitia Chapel, Adrien Corenflos, Kilian Fatras, Nemo Fournier, et al. Pot: Python optimal transport. The Journal of Machine Learning Research, 22(1):3571-3578 2021. (Cited on page 20)
Sarel J Fleishman, Timothy A Whitehead, Damian C Ekiert, Cyrille Dreyfus, Jacob E Corn, Eva- Maria Strauch, lan A Wilson, and David Baker. Computational design of proteins targeting the conserved stem region of influenza hemagglutinin. Science, 33: 16-811-21, 2011. (Cited on page 9)
Pablo Gainza, Sarah Wehrle, Alexandra Van Hall- Balll-BBaavais, Anthony Marchand, Andreas Scheck, Zander Harteveld, Stephen Buckley, Dongchun Ni, Shuguang Tan, Freyr Sverrisson, et al. De novo design of protein interactions with learned surface fingerprints. Nature, pp. 1-9, 2023. (Cited on page !)
Victor Garcia Satorras, Emiel Hoogeboom, Fabian Fuchs, Ingmar Posner, and Max Welling. E (n) equivariant normalizing flows. Advances in Neural Information Processing Systems, 34:4181-4192. 2021. (Cited page 9)
Aude Genevay, Gabriel Peyre, and Marco Cuturi. Learning generative models with Sinkhorn divergences. Artificial Intelligence and Statistics (AISTATS), 2018. (Cited page 20)
Brian Hall. Lie groups, Lie algebras, and representations. Springer, 2013. (Cited on page 3)
Brian Hie, Salvatore Candido, Zeming Lin, Ori Kabeli, Roshan Rao, Nikita Smetanin, Tom Sercu, and Alexander Rives. high-level programming language for generative protein design. bioRxiv, pp. 2022-12, 2022. (Cited on page 9)
Emiel Hoogehoom, Victor Garcia Satorras, Clément Vignac, and Max Welling. Equivariant diffusion for molecule generation in 3d. In International Conference on Machine Learning, pp. 8867-8887. PMLR, 2022. (Cited on page 9)
Chin- Wei Huang, Milad Aghajohari, Joey Bose, Prakash Panangaden, and Aaron C Courville. Riemannian diffusion models. Advances Neural Information Processing Systems, 35:2750-2761, 2022. (Cited page 2)
Po-Ssu Huang, Scott E Boyken, and David Baker. The coming of age of de novo protein design. Nature, 537(7620):320-33 327, 2016. (Cited on page 1)
Ilia Igashov, Hannes Stärk, Clément Vignac, Victor Garcia Satorras, Pascal Frossard, Max Welling, Michael Bronstein, and Bruno Correia. Equivariant 3d-conditional diffusion models for molecular linker design. arXiv preprint arXiv:2210.05274, 2022. (Cited on page 9)
Mathias Højgaard Jensen, Lennard Hilgendorf, Sarang Joshi, and Stefan Sommer. Bridge simulation and metric estimation on lie groups and homogeneous spaces. 2022. (Cited on pages and 21)
Wengong Jin, Siranush Sarkizova, Xun Chen, Nir Hacohen, and Caroline Uhler. Unsupervised protein-ligand binding energy prediction via neural euler's rotation equation. Neural Information Processing Systems (NeurIPS), 2023. (Cited on page 9)
John Jumper, Richard Evans, Alexander Pritzel, Tim Green, Michael Figurnov, Olaf Ronneberger, Kathryn Tunyasuvunakool, Russ Bates, Augustin Židek, Anna Potapenko, al. Highly accurate protein structure prediction with alphafold. Nature, 596(7873):583-589, 2021. (Cited on pages 2, 3, 30, and 32)
Gurtej Kanwar, Michael S Albergo, Denis Boyda, Kyle Cranmer, Daniel C Hackett, Sébastien Racaniere, Danilo Jimenez Rezende, and Phiala E Shanahan. Equivariant flow-based sampling for lattice gauge theory. Physical Review Letters, (25(12):12160) 2020. (Cited on page 9)
Isay Katsman, Aaron Lou, Derek Lim, Qingxuan Jiang, Ser-Nam Lim, and Christopher De Sa. Equivariant manifold flows. In ICML Workshop on Invertible Neural Networks, Normalizing Flows, and Explicit Likelihood Models, 2021. (Cited page 9)
12

---

Published as a conference paper at ICLR 2024
Leon Klein, Andreas Krämer, and Frank Noé. Equivariant flow matching. arXiv preprint arXiv:2306. 15030, 2023. (Cited on page 9)
Jonas Köhler, Leon Klein, and Frank Noé. Equivariant flows: exact likelihood generative learning for symmetric densities. In International conference machine learning, pp. 5361-5370. PMLR. 2020. (Cited on pages and 27)
Adam Leach, Sebastian M Schmon, Matteo Degiacomi, and Chris Willcocks. Denoising diffusion probabilistic models so (3) for rotational alignment. In ICLR 2022 Workshop on Geometrical and Topological Representation Learning, 2022. (Cited on pages 2 19)
Ziyao Li, Xuyang Liu, Weijie Chen, Fan Shen, Hangrui Bi, Guolin Ke, and Linfeng Zhang. Uni-fold: An open-source platform for developing protein folding models beyond alphafold. bioRxiv, 2022.
(Cited on page 26)
Yeqing Lin and Mohammed AlQuraishi. Generating novel, designable, and diverse protein structures by equivariantly diffusing oriented residue clouds, 2023. (Cited on pages 7, 8, and 32)
Zeming Lin, Halil Akin, Roshan Rao, Brian Hie, Zhongkai Zhu, Wenting Lu, Nikita Smetanin, Allan dos Santos Costa, Maryam Pazel-Zarandi Tom Sercu, Sal Candido, et al. Language models of protein sequences at the scale of evolution enable accurate structure prediction. bioRxiv, 2022. (Cited on pages and 9)
Yaron Lipman, Ricky T. Q. Chen, Heli Ben-Hamu, Maximilian Nickel, and Matt Le. Flow matching for generative modeling, October 2022. (Cited on pages 2, 3, 4, 19, 20, and 28)
Guan-Horng Liu, Arash Vahdat, De-An Huang, Evangelos A. Theodorou, Weili Nie, and Anima Anandkumar. 1esb: Image-to-image schrödinger bridge. International Conference Machine Learning (ICML), 2023a. (Cited on pages 5 and 20)
Xingchao Liu, Lemeng Wu, Mao Ye, et al. Learning diffusion bridges on constrained domains. In The Eleventh International Conference on Learning Representations, 2022. (Cited on page
Xingchao Liu, Chengyue Gong, and Qiang Liu. Flow straight and fast: Learning generate and transfer data with rectified flow. International Conference on Learning Representations (ICLR), 2023b. (Cited on pages 20 and 27)
Ali Madani, Ben Krause, Eric R Greene, Subu Subramanian, Benjamin P Mohr, James M Holton, Jose Luis Olmos Jr, Caiming Xiong, Zachary Sun, Richard Socher, et al. Large language models generate functional protein sequences across diverse families. Nature Biotechnology, pp. 1-8, 2023. (Cited on page 9)
Anthony Marchand, Alexandra K Van Hall-Beauvais, and Bruno E Correia. Computational design of novel protein interactions-an overview methodological approaches and applications. Current Opinion in Structural Biology, 74:102370, 2022. (Cited on pages 2 and 8)
Emile Mathieu and Maximilian Nickel. Riemannian continuous normalizing flows. Advances in Neural Information Processing Systems, 33:2503-2515, 2020. (Cited on pages and 3)
S. Matthies, J. Muller, and G. W. Vinel. On the normal distribution in the orientation space. Textures and Microstructures, 10:595067, Jan 1900. ISSN 1687-5397. (Cited on page 18)
John McCafferty, Andrew D Griffiths, Greg Winter, and David Chiswell. Phage antibodies: filamentous phage displaying antibody variable domains. nature, 348(6301):552-55,, 1990. (Cited on page 9)
Robert J. McCann. Polar factorization of maps on riemannian manifolds. Geometric & Functional Analysis GAFA, 11:589-608, 2001. (Cited on page 20)
Laurence Midgley, Vincent Stimper, Javier Antorán, Emile Mathieu, Bernhard Schölkopf, and José Miguel Hernández-Lobato. Se (3) equivariant augmented coupling flows. arXiv preprint a30iv:2308, 10364, 2023. (Cited on page 9)
Dmitry I. Nikolayev and Tatjana Savyolov. Normal distribution on the rotation group so(3). Textures and Microstructures, 29:173236, Jan 1900. ISSN 1687-5397. (Cited page 6)
13

---

Published as a conference paper at ICLR 2024
F. C. Park and R. W. Brockett. Kinematic dexterity of robotic mechanisms. The International Journal of Robotics Research, 13(1):1-15, 1994. (Cited on page 18)
Erik Persson and Bertil Halle. Nanosecond to microsecond protein dynamics probed by magnetic relaxation dispersion of buried water molecules. Journal of the American Chemical Society, 130 (5):1774-1787, 2008. doi: 10. 10110//j0777887. (Cited on page 26)
Gabriel Peyré and Marco Cuturi. Computational Optimal Transport. TtVV:8803.00560 2019, (Cited on page 20)
David Pollard. A user's guide to measure theoretic probability. Number 8, Cambridge University Press, 2002. (Cited on pages 4 and 22)
Aram-Alese Alexandre Pooladian, Heli Ben-Hamu, Carles Domingo-Enrich. Brandon Amos, Yaron Lipman, and Ricky Q. Chen. Multisample flow matching: Straightening flows with minibatch couplings, 2023a. (Cited on page 5)
wram-lexandree Pooladian, Heli Ben-Hamu, Carles Domingo-Enrich, Brandon Amos, Yaron Lipman, and Ricky T.Q. Chen. Multisample flow matching: Straightening flows with minibatch couplings. International Conference Learning Representations (ICLR), 2023b. (Cited pages 19, 20, and 23)
Yu Qiu. Isotropic Distributions for 3-dimension Rotations and One-sample Bayes Inference. Phl) thesis, Iowa State University, 2013. (Cited page 6)
Philip A Romero and Frances H Arnold. Exploring protein fitness landscapes by directed evolution. Nature reviews Molecular cell biology, 10( 10(12)866-876, 2009. (Cited on page 9)
Daniela Röthlisberger. Olga Khersonsky, Andrew M Wollacott, Lin Jiang, Jason DeChancie, Jamie Betker, Jasmine L Gallaher, Eric A Althoff, Alexandre Zanghellini, Orly Dym, et al. Kemp climination catalysts by computational enzyme design. Nature, 453(7192) 2008. (Cited
on page 9)
Marco Rudolph, Bastian Wandt, and Bodo Rosenhahn. Same same but differnet: Semi-supervised defect detection with normalizing flows. In Proceedings of the IEEE/CVE winter conference on applications of computer vision, pp. 1907-1916, 2021. (Cited on page 3)
Tim Salimans, Han Zhang, Alec Radford, and Dimitris Metaxas. Improving gans using optimal transport, 2018. (Cited page 20)
Gisbert Schneider. Automating drug discovery. Nature reviews drug discovery, 17(2):97-113,, 2018. (Cited on pages 2 and 8)
Petra Schneider, W Patrick Walters, Alleyn T Plowright, Norman Sieroka, Jennifer Listgarten, Robert Goodnow Jr, Jasmin Fisher, Johanna M Jansen, José S Duca, Thomas S Rush, al. Rethinking drug design in the artificial intelligence era. Nature Reviews Drug Discovery. 19(5): 353-364, 2020. (Cited on pages and 8)
David E. Shaw, Paul Maragakis, Kresten Lindorff-Larsen Stefano Piana, Ron O. Dror, Michael P. Eastwood, Joseph A. Bank, John M. Jumper, John K. Salmon, Yibing Shan, and Willy Wriggers. Atomic-level characterization of the structural dynamics of proteins. Science, 33066000):341-346 2010. (Cited pages and 26)
Yuyang Shi, Valentin De Bortoli, Andrew Campbell, and Arnaud Doucet. Diffusion Schrödinger bridge matching. arXiv preprint 1303. 2023. (Cited pages 5 and 20)
Dadiel-Adriano Silva, Shawn Yu, Umut Y Ulge, Jamie Spangler, Kevin M Jude, Carlos Labão- Almeida, Lestat R Ali, Alfredo Quijano-Ruu Rubio, Mikel Ruterbusch, Isabel Leung, et al. De novo design of potent and selective mimics of il-2 and il-15, Nature, 565(7738): 191, 2019. (Cited on page 1)
Vignesh Ram Somnath, Matteo Pariset, Ya-Ping Hsieh, Maria Rodriguez Martinez, Andreas Krause, and Charlotte Bunne. Aligned diffusion Schrödinger bridges. arXiv preprint 2302. / 1419, 2023. (Cited on pages 9 and 27)
14

---

Published as a conference paper at ICLR 2024
Eva-Maria Strauch, Steffen M Bernard, David La, Alan J Bohn, Peter S Lee, Caitlin E Anderson, Travis Nieusma, Carly A Holstein, Natalie K Garcia, Kathryn A Hooper, et al. Computational design of trimeric inftuenza-nnutraaliingg proteins targeting the hemagglutinin receptor binding site. Nature biotechnology, 35(7):667-67,, 2017. (Cited on page 1)
Hannes Stärk, Bowen Jing, Regina Barzilay, and Tommi Jaakkola. Harmonic self-conditioned flow matching for multi-ligand docking and binding site design, 2023. (Cited on page 27)
Alexander Tong, Jessie Huang, Guy Wolf, David van Dijk, and Smita Krishnaswamy. TrajectoryNet: dynamic optimal transport network for modeling cellular dynamics. International Conference on Machine Learning (ICML), 2020. (Cited on page 20)
Alexander Tong, Nikolay Malkin, Kilian Fatras, Lazar Atanackovic, Yanlei Zhang, Guillaume Huguet, Guy Wolf, and Yoshua Bengio. Simulation-free schrödinger bridges via score and flow matching, 2023a. (Cited on pages 19, 20, and 22)
Alexander Tong, Nikolay Malkin, Guillaume Huguet, Yanlei Zhang, Jarrid Rector- -roook,, Kilian Fatras, Guy Wolf, and Yoshua Bengio. Improving and generalizing flow-based generative models with minibatch optimal transport. arXiv preprint 2302.00482, 2023b. (Cited on pages 2, 3, 4, 5, 19, 20, 23, 27, and 28)
Mihaly Varadi, Stephen Anyango, Mandar Deshpande, Sreenath Nair, Cindy Natassia, Galabina Yordanova, David Yuan, Oana Stroe, Gemma Wood, Agata Laydon, Augustin Zidek, Tim Green, Kathryn Tunnyasuvunakool, Stig Petersen, John Jumper, Ellen Clancy, Richard Green, Ankur Vora, Mira Lutfi, Michael Figurnov, Andrew Cowie, Nicole Hobbs, Pushmeet Kohli, Gerard Kleywegt, Ewan Birney, Demis Hassabis, and Sameer Velankar. AlphaFold Protein Structure Database: massively expanding the structural coverage protein-sequence space with high- accuracy models. Nucleic Acids Research, 500D1):D439-D444, 11 2021. ISSN 0305-1048. doi: 10. n0.1093/nar/gkabl001 URL htttps://ddee org/ 1093/ nar http://0/0/////aaaaaaaaoo 061. (Cited on page 32)
Robert Verkuil, Ori Kabeli, Yilun Du, Basile IM Wicky, Lukas Milles, Justas Dauparas, David Baker, Sergey Ovchinnikov, Tom Sercu, and Alexander Rives. Language models generalize beyond natural proteins. bioRxiv, pp. 2022-12, 2022. (Cited on page 9)
Villani. Topics in Optimal Transportation. American Mathematical Society, 2003. (Cited on pages 20 and 21)
C. Villani.
Optimal Transport: Old and New. Grundlehren der mathematischen Wissenschaften.
Springer Berlin Heidelberg, 2008. ISBN 97835407 10509. (Cited on page 20)
Yajie Wang, Pu Xue, Mingfeng Cao, Tianhao Yu, Stephan Lane, and Huimin Zhao Directed evolution: methodologies and applications. Chemical reviews, 121 12444, 2021. (Cited on page 9)
Joseph L. Watson, David Juergens, Nathaniel R. Bennett, Brian L. Trippe, Jason Yim, Helen E. Eisenach, Woody Ahern, Andrew Borst, Robert J. Ragotte, Lukas Milles, Basile I. M. Wicky, Nikita Hanikel, Samuel J. Pellock, Alexis Courbet, William Sheffler, Jue Wang, Preetham Venkatesh, Isaac Sappington, Susana Vázquez Torres, Anna Lauko, Valentin De Bortoli, Emile Mathieu, Sergey Ovchinnikov, Regina Barzilay, Tommi S. Jaakkola, Frank DiMaio, Minkyung Baek, and David Baker. De novo design of protein structure and function with rfdiffusion. Nature, (2077976):1089-1111 2023. ISSN 1476-4687. (Cited pages 2, 7, 24, and 31)
Greg Winter, Andrew Griffiths, Robert Hawkins, and Hennie R Hoogenboom. Making antibodies by phage display technology. Annual review of immunology, 12(1):433-455 1994. (Cited on page 9)
Kevin Eric Wu, Kevin K Yang. Rianne van den Berg, James Zou, Alex Xijie Lu, and Ava P Amini Protein structure generation via folding diffusion. 2022a. (Cited on page 9)
Ruidong Wu, Fan Ding, Rui Wang, Rui Shen, Xiwen Zhang, Shitong Luo, Chenpeng Su, Zuofan Wu, Qi Xie, Bonnie Berger, Jianzhu Ma, and Jian Peng. High-resolution de novo structure prediction from primary sequence. bioRxiv, 2022b.
(Cited page
15

---

Published as a conference paper at ICLR 2024
Minkai Xu, Lantao Yu, Yang Song, Chence Shi, Stefano Ermon, and Jian Tang. Geodiff: A geometric diffusion model for molecular conformation generation. arXiv preprint arXiv:2203.02923, 2022. (Cited on page
Kevin K Yang, Zachary Wu, and Frances H. Arnold. Machine-learning -uidedd directed evolution for protein engineering. Nature methods, 16(8):687-694, 2019. (Cited on page 9)
Jason Yim, Andrew Campbell, Andrew Y. K. Foong, Michael Gastegger, José Jiménez-Luna, Sarah Lewis, Victor Gareia Satorras, Bastiaan Veeling, Regina Barzilay, Tommi Jaakkola, and Frank Noé. Fast protein backbone generation with se(3) flow matching, 2023a. (Cited on page 9)
Jason Yim, Brian L Trippe, Valentin De Bortoli, Emile Mathicu, Arnaud Doucet, Regina Barzilay, and Tommi Jaakkola. Se (3) diffusion model with application to protein backbone generation. arXiv preprint 2322022302.02277, 2023b. (Cited on pages 2, 3, 6, 7, 9, 21, 22, and 30)
Shuxin Zheng, Jiyan He, Chang Liu, Yu Shi, Ziheng Lu, Weitao Feng, Fusong Ju, Jiaxi Wang, Jianwei Zhu, Yaosen Min, et al. Towards predicting equilibrium distributions for molecular systems with deep learning. arXiv preprint arXiv:23000 2.4...5455.0555 2023. (Cited on page 2)
16

---

Published as a conference paper at ICLR 2024
A THEORETICAL PRELIMINARIES
A.1 SO(3) LIE GROUP
The Special Orthogonal group in 3 dimensions, SO(3) consists of the 3D rotation matrices:
SO(3) {r R : rr de = 1
(11)
is a matrix Lie group with the lie algebra given by:
s0( = (1 E Rax3 :T = -r]
(12)
Parametrizations of (3) The skemesymmetric matrices E SO(3) can be uniquely identified with vector R such that Vv E IR,, rv x v, where indicates the cross product. This vector is known as the rotation vector. The magnitude this vector, w = ||w|| is the angle of rotation its direction, ew [wi
is the axis of rotation.
Mapping the R vector to the saemmsymmtttic matrix is known as the hat operation, (-)
Another parametrization of SO(3) with Euler angles, described using three angles (ø. 0, 1²) A common convention is to use the z-convention, where the rotation is given by: a rotation about the z-axis by Φ a second rotation about the former r-axis by 0 and a last one about the former 2-axis by v/.
Metric on (3)... First, we recall that metric is a bilinear function (:.-) :
R" x R + that is both symmetric and positive definite. Additionally, we recall that a quadratie form manifold M a bilinear map T.. x T. - R that is smooth and symmetric. A positive-define quadratic form is, therefore, a metric. Let us consider the following symmetric positive definite quadratic form defined:
A
B
)
Q
B
C
(13)
canonical choice for the metric of (O(3) is obtained by taking Q 1/21, resulting in bi-invariant metric on (3))). Therefore, the metric is given by:
1 (t1, 12)50(3) tr(t] Qt2) ur(t - 12)
(14)
Note that the inner product on Lie groups consumes elements of the Lie algebra and, because the left action is transitive, this inner product is well-defined for all tangent spaces of the group elements.
The distance induced by this metric is given by:
(r1, r2) = || log (r] r2)
(15)
for r1, r2 E SO(3) and where the Frobenius matrix norm is used.
The exponential and logarithmic maps on SO(3). Generally speaking, the exponential and logarithmic maps of a Lie group G relate the elements in the group to the lie algebra, G. In the case of matrix Lie groups, these coincide with the matrix exponential:
g exp(g)
(16)
n!
and the matrix logarithm:
N
1) (g 1)-
g = log(g)
(17)
For (3), since the elements of the lie algebra are skew-symmetric matrices, eq. (16) for the matrix exponential can be simplified significantly obtain closed-form, known as Rodrigues formula. Given w a rotation vector, and e E 50(3), the corresponding element of the lie group, r E SO(3) is given by:
exp ف = cos(w)) + sin(w)ew ew + (1 cos(ww
(18)
where w and ew are the angle and axis of rotation for w.
Similarly, the matrix logarithm be expressed using the rotation angle:
17

---

Published as a conference paper at ICLR 2024
ifr I, 2 (r
log(r)
(19)
r I.
A.2 SE(3) LIE GROUP
The special Euclidean group, SE(3) is used to represent rigid body transformations in dimensions:
SE(3) {
5O(3), E (TR
(20)
Represented by this 4 x matrix and with the group operation defined by matrix multiplication, this group can be seen subgroup of the general linear group GI, IR). The lie algebra of the group se(3) is given by:
sc(3) -1=- I : ( R3
(21)
Note that the tangent space of R$ is isomorphic the space itself so we can simply use the notation instead of This lie algebra is isomorphic to R using the map: + (w, s), where we have identified the symmmettttiic matrix r s0(3) with its axis-angle representation.
E R As the group of
translations, (R3 +) is a normal subgroup of SE(3), the group can be understood as a semi-direct product: SE(3) SO(3) x (R².+)
Metric SE(3). Although there many possible choices for metrics SE(3), none of them are bi-invariant. Instead, one can choose to build a left-invariant or right-invariant metric. A simple choice for the quadratic form Q from eq. (13) is setting the matrices A C Iz and B 0 (Park & Brockett, 1994), which gives:
13 (
(22)
Ta
Using this metric we can define an inner product on SE(3) as F2)s03))) = tr(p] Qr2), where is the trace operation. Writing out the inner product explicitly for F1 E se(3) get,
tr(x, T QIz) tr
[
r $2
(23)
$
$ 82
Thus, we have tr(rī Q₁2) metric on SO(3) and R3:
ri 12 + s2. Therefore, the metric on SE(3) decomposes into the
(T1 (1.F2)SE(S) = (t1, 12)so(3) + ($1, $2)22
(24)
This means that we can obtain the geodesics on SE(3) from the geodesics on the product manifold SO(3) x 3
(SE(3)(21, / dsoc3) (21.12) dga ($1, (s1)82)²
(25)
where 2 (11, x2 = (r2,82) E {EE(3), dso(3) is defined in eq. (15) and d3a is the usual Euclidean distance.
A.3 THE ISOTROPIC GAUSSIAN DISTRIBUTION ON SO(3)
IGSo(3) density. The isotropic Gaussian distribution on SO(3) is parametrized by a mean, r E SO(3) and concentration parameter, IR. It can be parametrized in axis-angle, where the axis of rotation is sampled uniformly and the angle of rotation w has probability density function (pdf) given by:
f(we, e) Σ (21 + 1)e 1(+1)) sin ((1 + 1/2)wz)
(26)
sin (wz/2)
1=0
Although this expression contains an infinite sum, Matthies et al. (1900) has shown that for S 1, can be approximated by closed-form equation:
f(wz,e) c) VTT -3/2 -2 (w - ((w 2 mw/e + (w + 2 2Tr me--////))
(27)
2sin (
18

---

Published as a conference paper at ICLR 2024
Sampling from I9s0(3)- Sampling from IGso;3) is done following Leach al. (2022). The angle of rotation is obtained by inverse transform sampling, where the cumulative density function is approximated using the pdf above, scaled by uniform density SO(3) with density f(w) " the axis is sampled uniformly from s? We note that the closed-fore form approximation of eq. (27) makes the computation the edf, and hence the sampling process very efficient.
1-c0sw
B FLOW MATCHING IN Rd
To perform FOLDFLOW on SE(3), we consider two different flows. One on SO(3) that we described in the main paper and another one on R that we describe depending on the consider FOLDFLOW method.
Riemannian Flow Matching is generalization of Flow Matching on Riemannian manifold. Therefore, the setting as well as the main ideas are similar and are straightforward to adapt to the Euclidean case. This means that the objective is also to regress conditional vector field built from conditional probability paths. In this section, we described the conditional probability paths and conditional vector fields that were used respectively by Lipman al. (2022) and Tong et al. (2023b).
The main difference is that the conditional probability path is now Gaussian conditioned on latent variable (==) with variance os (s(x) = (sslz, The conditional vector field has closed
form derived from the following Theorem:
Theorem 1 (Theorem 3 of Lipman al. (2022)). The unique vector field whose integration map satisfies (e(s) = He + 028 has the form
of (s me) of
ue(s)
p/b
(28)
Ot
We now describe the Flow Matching method Lipman et (2022) and OT-CFM Tong et al. (2023b;a); Pooladian et al. (2023b).
Flow Matching. In the context of data living in the Euclidean space Rd. Identifying the condition = with single datapoint 81, and choosing smoothing constant 0, sets
((() = A(s ts (to - -t -t+11 1)
(29)
$1 1- o)
-1-
(1(5]2)
(30)
which is probability path from the standard normal distribution (po(x)2) W(x;; 0, 1))
Gaussian distribution centered at 21 with standard deviation a (pr (x|22 = (x; o²)). one sets q(=) ((x1) be the uniform distribution over the training dataset, the objective introduced by Lipman et al. (2022) is equivalent to the CFM objective (1) for this conditional probability path.
OT-Conditional Flow Matching (Tong et al., 2023b), As explained the main paper, the probability path used in FM is not the optimal transport probability paths between the distributions Po and p1. Therefore, we want to get straighter flows for faster inference and more stable training. To achieve that, we leverage the optimal transport theory and want the probability path to be the Euclidean McCann interpolants defined as P: = (<n) + (1 ()sg- However, as the map V intractable in practice, we rely on the Brenier theorem which makes a connection between the map V and the optimal transport plan æ. Therefore we set the mean of Gaussian conditional probability path as #4 ts1 + (1 t)so and the latent distribution g(so. $1) = T(S. S1).
pe(s)so, 81) [Nss (sl (1 --)))s,0²²)
o²),
(31)
pr(s) / N(s ts + (1-t)ss  ( s1 s1)dspds,,
(32)
(s|=)) = 81 so-
(33)
In the case of the Euclidean space, the FM loss is equal LFOLDFLOW (1v0(L, s) (=]=]))]
This can be simplified to down to LROLDFLOw-3s [[vo 8) - (8) - 8o) This method is the main inspiration to develop FOI.DFL.O-OO.
19

---

Published as a conference paper at ICLR 2024
C RIEMANNIAN OPTIMAL TRANSPORT
Optimal transport in generative models. OT has been used in generative models for several approaches. For GANs, it was used as a loss function (Genevay et al., 2018; Fatras et al., 202 lb; Salimans al.. 2018; Arjovsky al., 2017). More recently, it was used to speeding up training and inference for continuous normalizing flows (Finlay et al., 2020: Liu et al., 2023b; Lipman et al., 2022; Tong et al., 2020; 2023b;a), Schrödinger bridge models (Shi al.. 2023; Liu et al., 2023a; De Bortoli et al., 2021). In this section, we recall its basic definition over a Riemannian manifold. Then we discuss its empirical computation and we finish this section by proving Proposition 1.
C.I BACKGROUND ON RIEMANNIAN OPTIMAL TRANSPORT
Optimal transport Riemannian manifold was first studied in the seminal work of McCann (2001) and we refer to Villani (2003; 2008) for a review of all results. Recently, optimal transport has also drawn attention from the machine learning community, and we now give longer introduction this topic.
The (static) Kantorovich optimal transport problem seeks a mapping from one measure another that minimizes a displacement cost. Formally, we define the 2- W-ssssrsteiin distance between distributions Po and P1 on M with respect to the cost c(x, y) = 2 1 dx, y)² as:
/ c(r, v) dr (x, v):
W(o: pI)
inf TEll(po.Pr) M22
(34)
where [poo P1) denotes the set of all joint probability measures on M x .M whose marginals are Po and P1. To compute the optimal transport plan, we rely on the POT library Flamary al. (2021). This problem is a relaxation of the well-known Monge formulation described in the main paper and that we recall now for the sake of readability.
The Monge optimal transport problem is defined as
J. c(x, (x)) dpo(.))
OT(po.
inf V:V
(35)
When M is a smooth compact manifold with no boundary and po has density, (McCann, 2001, Proposition 9) shows that the map T' exists and is unique. This is an extension to Riemannian manifold of the well-known Brenier Theorem (Brenier, 1991). The optimal transport map v and the McCann interpolation have then the following form:
(x) expx
(x(x)) = exvx((x\x)V
(36)
where ø is a c-concave function. Furthermore, we have that the optimal transport plan is supported on the graph of the Monge map, i.e., (id, 1) # Po- Therefore, knowing the transport plan leads to the Monge map. The connection between the two formulations for SE(3), is explicitly stated in Proposition I which is proved in SC.2. However, we first discuss the computation of the OT plan x.
Minibatch OT approximation. For empirical distributions, the Kantorovich problem is linear program and can be efficiently solved with the simplex algorithm. We refer to (Peyré & Cuturi, 2019, Chapter 3) for a review on how to solve the Kantorovich problem. However, when we deal with large datasets, computing and storing the transport plan π for Optimal Transport (OT) can be challenging due to its cubic time and quadratic memory complexity with respect the number of samples. To address this, a minibatch OT approximation is often employed. While this approach introduces some error compared the exact OT solution (Fatras et al., 2020), it has been proven effective in various applications such as domain adaptation and generative modeling (Damodaran et al., 2018; Genevay et al., 2018). Specifically, during training, for each source and target minibatch, pairs of points are sampled from the optimal transport plan computed between the pair (z, y) ~ Tbatch- We empirically show that the batch size can be small compared to the full dataset size and still give a good performance, which aligns with prior studies (Fatras et al., 2021b;a). This strategy is also at the heart of the OT-CFM methods (Tong et al., 2023b;a; Pooladian al., 2023b).
20

---

Published as a conference paper at ICLR 2024
C.2 PROOF OF PROPOSITION 1
We recall the proposition statement here for convenience and then prove it below.
Proposition Let us consider SE(3), with the product distance deEes and two compactly supported probability distributions po, Pr E )(SE(3)))) In addition, suppose that po is abso- lutely continuous with respect to Riemannian volume form (i.e., Po & dr). Then for the distance c= 220 SE(3), the Kantorovich and Monge problems admit unique solution that is connected as follows π = (id x 4)# Po. where Vi is almost uniquely determined everywhere po- Furthermore, we have that V (x) = cxp2 (V (Vo(x)) for some de (SE(s)--Concav function ø.
Proof. The manifold SE(3) is a connected, complete, (C*)) smooth manifold without boundary. SE(3)* (equipped with the usual product distance) is a finite Cartesian product of connected, complete, smooth manifolds without boundary and therefore it is itself connected, complete, smooth manifold without boundary. We only need to check these assumptions are also satisfied by SE(3), We do so by noting that SE(3), can written as N copies of SE(3) where the RS component is mean subtracted-i.e. 5f = 1/N E s", and the final Nth element in the product is the mean, 1 I/N 1/NE Ei= 8j. Certainly, the first N components satisfy connectedness, and completeness, and are manifolds that are smooth without boundary. Furthermore, the disintegration of measures on SF(3), (Proposition 3.5 (Yim al., 2023b)) allows us to define measure já proportional to R&3 for the final component Nth component. Therefore, by our assumptions on the measures Po, Pl. we can apply the following Theorem from Villani (2003) to get the desired results.
Theorem 2 (Theorem 2.47, Villani (2003)). Let M be connected, complete and smooth (C3) Riemannian manifold without boundary, equipped with its standard volume measure. Let po. P be two compactly supported distributions and set the ground cost c(2,y) = _d((x,))) y)² with the geodesic distance on M. Further, assume that po is absolutely continuous with respect to the volume measure on M. Then the Kantorovich and Monge problems admit a unique solution that is connected as follows π (id x V Po. where V is almost uniquely determined everywhere po. Furthermore we have that V/ (xx exp (Vo(z)) for some da de-coccave function é.
D STOCHASTIC RIEMANNIAN FLOW MATCHING
D.1 BROWNIAN BRIDGE ON SO(3)
We follow the presentation in Jensen et al. (2022) to define the Brownian bridge a Lie group G endowed with a metrie. We note that log is the inverse of the Riemannian exponential map. However, if the metric is bi-invariant, which is the case for (3))), it coincides with the Lie group logarithm. We can simulate a bridge on G via the guided diffusion SDE (using o for the Stratonovich integral), for process conditioned to reach at /
logR, (v) 1-t di
dR. (R.).) dt Vi(R.) dB
Ro ro:
(37)
where Vx = de2)) with {uI..... Vd} } an orthonormal basis T. and where B, is a Brownian motion on G. On (3()). since the metric is bi-invariant, we have V, = 0. In this work, we model the guided bridge with a diffusion that does not depend on the process R,, In this case, the Stratonovich and Itô formulations the same, yielding the reversed process defined in Eq. 7.
D.2 SIMULATION-FREE A APPROXIMATION OF BROWNIAN BRIDGES ON SO(3)
We now numerically investigate the fidelity of our simulation-free SDE which is employed in FOLDFLOW-SFNE in relation to the guided drift SDE in eq. (8). In fig. 6 we plot the mean and the standard deviation (over 1024 data points) of the distribution of the SO(3)-norm along the trajectory against time, for three different values the diffusion coefficient, 7. We find the true simulated Brownian bridge (bold black line) is in close proximity the simulation-free FOLDFL.OW-FFN SDE (red dotted lines). We further note that this holds for the entire trajectory and leads to overlapping shaded regions that correspond to the standard deviation of the norm. This result adds empirical
21

---

Published as a conference paper at ICLR 2024
substantiation to using the FOLDFLOW-SPE as drop-in and fast approximation to the Brownian bridge SDE on (3(3).
Sin Aetal SXX
7.5
3.3
310
330
225
4.0
12
200
3:80
114
3.15
1.20
122
04
08
04
24
60
02
04
10
(a) 0.1
(b)
0.5
(c) 1.0
Figure 6: Numerical comparison between the simulation- of the SDE FOLD vs. simulated Brownian bridge on SO)3), for different values of the diffusion coefficient," 7.
D.3 PROOF OF PROPOSITION 2
Proposition Given pt(a)) > O, Va E SE(3)) the conditional and unconditional FOLDFLOW- SFM losses have equal gradients w.r.t. 0: VoLusEM(O) VoLSFM(O),
(e(wll)) mi(a) ue [e(cl2)] for E SE(3), We claim that:
Proof. Let uu = Ep(x)
Va Vobsapta))<mmittl) () u=(=l=)]
Volzmpi(o) [lloo(t.x) u
(38)
From disintegration of measures (Pollard (2002) and Proposition 3.5 form Yim et al. (2023b)), we know the probabilities (=()) & ((r="))...p))))) pr(r"")..-p(rr">) p.(:-
(arN)))Pp (s*) ..st))...((()) ). and similar for the conditional probability p (x]2). Given that by eq. (24), the metric on SF(3) also factorizes into metric on SO(3)
and 3 it suffices prove the claim for SO(3) and IR3 The claim can therefore be stated follows,
where we have written r² as r and s¹ as for conciseness.
Vo VoRenoteeeesserlle [Iloo(t,)) (r(r
[[loo(t,)) u4 (r))
VoRrwpe(r)
(39)
Vell (t, 5) (s|=,))))
Vellampe(s) (1]oa(t, [e
(40)
The proof of this claim follows similar structure to Chen & Lipman (2023). We proceed by proving eq. (40). Dropping the distributions for conciseness, we have:
V (l (r(22)))) -...)l)))--...lll (oan))-u((rll) (r)
(41)
Vo 2T (rott). (r( )ss(()) 2 (vo ur (r)))s0(3)
Now:
(.(vo(t,)), L. /sors, (v)(t,T), (((r))sosspp((())vvl)
L.' (oo(t.)) / pc(r)az) ddvol, pr(r)d,
[SO(3) pr(r)
(42)
1/'' [ |sora) htw.
Er.. (va(t, r)ux(r(z.))
The proof for 3 directly follows Theorem 3.2 from Tong et al. (2023a).
22

---

Published as a conference paper at ICLR 2024
E EXTENDED FIGURE INFORMATION
E.
FIGURE
depicts the probability paths of FOLDFLOW-BASE FOLDFLOW-OT, and FOLDFLOW- SFM projected onto 8². Where FOLDFLOW BASE paths may cross, FOLDFLOW-OT conditional paths do not cross reducing the variance in the objective stabilizing training as studied in Pooladian et al. (2023b); Tong et al. (2023b). FOLDFLOW-S. LFM-SFM adds in stochasticity which improves novelty in our protein generation task. Figure also contains a table summarizing the differences between methods. Showing whether or not they can map from general source distribution, can perform optimal transport under some conditions, are stochastic or deterministic, and do not require calculation the score. We note that there is for FOLDFLOW-SEP SFM performing OT, as it only achieves OT when noise goes zero and it recovers FOLDFLOW OT However, this bias may still be helpful in reducing the variance of the objective function even if OT is not achieved.
F SO(3) TOY EXPERIMENT
F.1 ToY MODEL PARAMETERIZATION
For the vector-field parametrization, the goal to create a function that by construction lies on tangent space of the manifold. For the toy experiments, this is done by using 3-layer MLP, and projecting the output of the network to the tangent space of the input. That is, similar to Chen & Lipman (2023), we have:
up (4( -1,,ML((,(())
(43)
where II, (M) projects 3 x 3 matrix onto (3). This operation essentially computes the skew- symmetric component of M, given by M-M and parallel transports it the tangent space of using left matrix multiplication which is the group operation SO(3).
F.2 ADDITIONAL RESULTS FOR SO(3) TOY EXPERIMENTS
In this section, present the qualitative results of our toy experiments. In fig. 7, the three models, FOLDFLOW-BASe BASE,,, FOLDFLOW-OT FOLDFLOW-ST FFM learn to correctly model the modes of the ground-truth distribution with slight model shrinkage in the FOLDFLOW-BAS
that all
(a)
(b)
(c)
(d)
Figure 7: (a) Data distribution (b) FOLDFLOW-BASE (c) FOLDFLOW-OT (d) FOLDFLOW-SFM.
The data visualized using the Euler-angle representation of the rotation matrices.
G ADDITIONAL RESULTS AND ANALYSIS FOR THE PROTEIN EXPERIMENTS
G.1 PROTEIN BACKBONE GENERATION EXPERIMENT ADDITIONAL RESULTS
Empirical investigation of rotation norms for Inference Annealing.
In fig. 10, we show five proteins generated by FOLDFLOW-SFN, from each backbone length. Here we show the generated structure in green and the best ESM-refolded structure out of eight sequences generated using ProteinMPNN. We can see that FOLDFLOW-SFI generates diverse folds that refold with diversity in secondary structure and overall 3D conformation.
In fig. 9 we compare the performance of models on designability, diversity, and novelty tasks for that FOLDFL.Ow closes the gap between models
different backbone lengths. In particular,
23

---

Published as a conference paper at ICLR 2024
103
Rotation Flow Norm
10a
10
109
10-1
200
300
500
Inference Step
Figure 8: Norm of the rotation flow with and without scaling.
13.5
1.5
6.0
330
250 300 180 153
2UE
250
3.0
189 155
200
250
300
(a)
(b)
(c)
Figure 9: (a) Designability quantified by scRMSD (lower is better). (b) Diversity quantified by average pairwise TMScore (lower better), and (c) Novelty of proteins as quantified by maximum TMScore PDB (lower better), designed across lengths for FOLDFLOW models and previous state the art models.
without pretraining (Genie, Frfme FOLDFLOW) and RFDiffusion terms of designability, particularly shorter sequences (5 200).
We note there is a trade-off between designability and diversity/novelty, both at the short sequence lengths and as sequence length increases. For longer sequences (250, 300), while FOLDFLOW models are comparable in terms of designability, they generate significantly more diverse and novel structures as compared all other models, even RFDiffusion (although RFDiffusion still generates significantly more designable proteins at these lengths.
G.1.1 TIMING COMPARISON IN TABLE 2 AND TABLE 4
In Table 2 we compare the number of steps per second for each model, where step corresponds to forward and backwards pass on effective batch size as defined Equation (46) on a single GPU. Here we find that FOLDFLOW is over 2x faster than FrameDiff per step. This drastic improvement is due to a number of optimizations, with the largest being that we can avoid the costly IGso(3) score computation which is necessary for their method.
We train our model in Pytorch using distributed data data-paralle parallel (DDP) across four NVIDIA A100-80GB GPUs for roughly 2.5 days. We note that this substantially less than comparable models (table 4). RFDiffusion requires the use of pre-trained weights from RosettaFold which trained for 4 weeks on 64 V100 GPUs (Watson al., 2023).
G.2 EQUILIBRIUM CONFORMATION GENERATION EXPERIMENT
As described in Section 5.3 proteins take many different physical conformations in the real world. These conformations dictate many important attributes of protein's 's behaviour, e.g., how one protein might bind to another. As a protein's 's conformations generally do not deviate greatly from a desirable approach would be to start from a noised version a known conformation of the protein generate another conformation. We hypothesize that the flows required to do this are easier to learn than starting from an uninformed source distribution. We find FOLDFLOW is an ideal candidate for this setting, and show its efficacy in fig. 5.
another,
24

---

Published as a conference paper at ICLR 2024
Figure 10: More FOLDFLOW-SFN. generated structures in green compared to ProteinMPNN -> ESMFold refolded structures in grey. 5 samples all with RMSD < 2Å for lengths 100, 150, 200, 250. 300 from left to right. FOLDFLOW-SFI generates designable diverse proteins.
Figure 11: KL divergence per residue of the 2D dihedral angle ($ and w) distributions between the samples from FoLO and test MD frames (blue) and AlphaFold 2 and the test MD frames (orange).
25

---

Published as a conference paper at ICLR 2024
Table 4: Training resources for protein generation models.
Model
Training time
Optimization Steps
#gpus
Distributed Training
RFDiffusion
28 3 days
64
8
RFDiffusion w/o pretraining
3 days
8
Genie (SwissProt)
8 days
~000
DP
FrameDiff-ICML
~7 days
~119m
DP
TrameDiif-Improvd
+7 days
+1.9m
DP
FrameDiff-Retiee Retrained
10 days
2.2m
2
DP
FOL.DFLOW (BASE, OT and SFM)
2.5 days
600k
4
DDP
12
(a) FOLDFLOW
(b) FOLDFLOW-ran
(c) FifmeDiff
Figure 12: Ramachandran plot of & and v for FOLDFLOW with informed prior (b) FOLDF1.OW with uniformed prior and (c) FrameDiff steps. With the informed prior, FOLDFLOW able better capture the mode on the left.
We chose bovine pancreatic trypsin inhibitor (BPTI) to study this experiment. The 58-residue protein is the first protein whose dynamics were studied experimentally by nuclear magnetic onance (NMR) and was the first protein that was simulated by molecular dynamics for 1ms. On timescales ranging from nanoscconds milliseconds, the dynamics of BPTI involve protein back- bone structural changes that, for example, accommodate water molecule exchange and disulfide isomerization (Persson & Halle, 2008). We used the 1-ms MD simulation trajectory of 300 K (Shaw et al., 2010) to reproduce and interpret the kinetics of folded BPTI. To construct our source distribution, Rose Fool,,, and Unifold (Li et al., 2022). Our source distribution was then added small amount of noise from the standard Gaussian and [9so(3)- We trained FOLDFLOW for one day A100 GPUs.
temperature
first generated four folded conformations from each of AlphaFold2, ESMFold,
Results in Section 5.3 show that FOLDFLOW generates conformations covering all modes of the true conformation distribution. Moreover, we sample different conformations of BPTI from AlphaFold2 and plot them on the Ramachandran and ICA plots, observing while FOLDFLOW can capture all modes of the distribution AlphaFold2 only captures one. Further fig. 11 shows that the KL divergence between the distribution has learned the distribution of the target. We believe this is an exciting direction meriting larger experiments on more proteins in the future.
angles generated by FOLDFLOW is low and uniform, conveying
Table 5: Quantitative performance on the equilibrium conformation generation on BPTI protein. Measures the 2-Wasserstein W3 angle space between generated and test samples over all residues (W2)), most flexible residue (W2 @W6)), and the Kullback-Leibe divergence also at the most ficxible residue (KL @ 56). We also measure the distance of distribution of AlphaFold2 structures (AlphaFold2) well as the distance samples from the trainset.
W2(1)
W2@56(1)
KL@56(L)
FOLDFLOW
4.379
0.406
0.441
FOLD
4.446
0.557
1.205
FrameDiff
4.844
0.800
3.051
RandomPrior
18.752
1.993
0.746
AlphaFold2
7.298
1.917
5.724
Trainset
4.140
0.198
0.487
26

---

Published as a conference paper at ICLR 2024
Comparison to Uninformed Prior and FrameDiff. Next we further describe the setup of the comparison experiment in table 3 and table 5 which compares FOLDFLOW with an informed prior, FOLDFLOW with uniform random prior (FOLDFLOW FLoW-Radd) and FrameDiff. We measure the 2- 2-Wasserstein (W2) distance between generated and test samples either on all 58 residues (denoted W2) or just on residue 56 (denoted W2 @56), which is the most flexible residue. The W2 is calculated using the (2)))* distance on & and V. We also compare the generated and test samples over the Ramachandran plot of residue 56 with a Kullback-Leibler (KL) divergence (depicted in fig, 12). Here we compute an empirical histogram from samples over a 100 by 100 grid in angle space. We then compute the KL divergence between the smoothed empirical distributions where the minimum of each bin is clipped at 10 10 for stability.
We train each model for 10k steps on 2 A100 GPUs using either a random prior (uniform over rotations and Gaussian translation) or a mixture of Gaussians / IG5o(3) distributions using centers defined by the four folded prior conformations and with standard deviation and IG5o(3) concentration 0.5. We do not use inference annealing for this experiment. We generate 250 samples from the model and test against 1000 samples from the test set for computational efficiency reasons.
To contextualize these results, we compare the performance of these models with various baselines such as 250 samples from a random prior (used as the prior in FOLDFL.WWRRnd and FrameDiff where each residue is sampled from (0, 10)), 160 conformations sampled from AlphaFold 2, and 250 samples from the training set (Trainset). Results are averaged over 10 seeds for the random prior and the train set. The trainset represents well-trained model the Wasserstein distance is not zero even for empirical distributions drawn from the same distribution. The RandomPrior and the AlphaFold2 represent the random and informed priors respectively. All models are significantly better than these two priors, and FOLDFLOW approaches the performance of samples from the training set.
Figure 12 depicts the Ramachandran plots for residue 56 with seatter plots for FOLDFLOW FOLD- FLOw-Rand and FrameDiff against kernel density estimate (KDE) of the test set. We see that FOLDFLOW with the informed prior is able to model both modes where MLLDFLOW-Rand and FrfmeDiff both focus on the mode in the bottom right, centred at Ø- 7/2.
We have two major findings from this experiment:
An informed prior helps improve performance both overall and on the most flexible residue as seen by comparing the performance of FOLDFLOW and FOLDFLOW- -and in table 3.
FOLDFLOW (with both informed and random priors) improve over FrameDiff on this task.
The equilibrium conformation generation task, studied here, is an example of a setting where an informed prior may be useful. Recent work has explored other applications of starting from an informed prior, such as protein docking (Somnath et al., 2023; Stärk et al., 2023), single-cell (Tong et al., 2023b) and image-to-image translation (Liu et al., 2023b).
H FURTHER DISCUSSION OF FOLDFLOW AND RELATED MODELS
Symmetries as an inductive bias for flow matching. Leveraging symmetries as an inductive bias in deep learning models (for example by data augmentation or design equivariant models) has been shown to improve data efficiency and lead to better generalization. In the context of flow matching for proteins, the goal is to learn the vector field generating the flow, which maps an invariant source to an invariant target distribution, guaranteeing the existence of an equivariant vector field (Köhler et al.. 2020; Bose & Kobyzev, 2021). Therefore, one way to exploit this symmetry would be to parameterize the vector field with an equivariant network, taking as input the 3D coordinates of the protein. Alternatiely, since protein backbones can be parametrized by elements of SE(3)", we can directly construct the vector field by taking an intrinsic perspective by using charts on the manifold and their coordinate system. In this case, as the vector field lies on the tangent space of SE(3) it is equivariant by construction.
Comparison between flow matching and diffusion approaches. While flow matching and diffusion models bear many similarities they also have key differences which we highlight in this appendix.
1. Flow matching based approaches enjoy the property of transporting any source distribution to any target distribution. This is in contrast to diffusion where one typically needs a Gaussian-like source distribution.
27

---

Published as a conference paper at ICLR 2024
2. Flow matching approaches are readily compatible with optimal transport due to the same property of being able to transport and source to any target. Optimal transport which itself has the advantage of providing faster training with a lower variance training objective and reducing the numerical error in inference due to straighter paths. Diffusion models by themselves not amenable to optimal transport but instead one can do entropic regularized OT. In Euclidean space, this corresponds to a Schrodinger bridge but this is not an optimal transport path is it stochastic.
3. In general, simulating ODE is much more efficient than simulating SDE during inference. Conditional fow-matching and OT-conditional flow matching (Lipman et al., 2022; Tong et al., 2023b) both learn ODEs as the Iearned flow corresponds to continuous normalizing flow. Diffusion models on the other hand are SDEs and while being more robust to noise in higher dimensions require more challenging inference.
Comparison to FrameDiff. While model uses similar setup to FrameDiff, we introduce a number of improvements that help to stabilize training and improve performance. Indeed our additions lead to improvements all metrics over the PrameDDif-Improoedd model released on GitHub which substantially improves on the designability over FrameDiff-ICM We first recap the improvements made in CoaneDiffImpprovee over FrameDiff-ICML detected in the code:
1. A bug in the score calculation for rotations means that there is a stop gradient in the rotation score calculation and FrameDiff-ICML is not trained to match the rotation score, which makes its performance quite impressive given this limitation. This bug is fixed in the CrumeDiff-Improeer model which uses different score calculation.
2. The dataloader was switched from sampling uniform over proteins in the dataset, to uniform over clusters, then uniform within clusters. As we explore in Appendix I.5, this changes the distribution of proteins but overall increases diversity as there are many similar proteins in a small number of clusters Figure 13c.
3. The rotation loss was changed to use separate axis and angle component reduce variance in the loss.
While these items improve the performance of FrameDifff especially in terms of designability, there are still few potential areas for improvement.
One we focus is the costly loss function of FrameDiff which relies calculating the pdf of [Gsoo13) for sampling and computing the score. In the setting used Frame-Diff, the infinite-sum formulation of the density from eq. (26) had to be used, leading an expensive score loss.
We also noticed that FrameDiff does not exactly follow theory in that the model is not exactly translation invariant: As mentioned in section 3 in order to obtain SE(3) invariant distributions, the center of mass has to be removed. However, in the FrameDiff code, this was only done at inference and not during training. It is unclear to what extent this impacts the performance, as the model remains translation invariant in expectation and during inference.
I IMPLEMENTATION DETAILS AND EXPERIMENTAL SETUP
I.1 TRAINING AND INFERENCE
To describe the precise algorithm for training FOLDFLOW models over distributions in SE(3)* N Our starting distribution in SE(3)" x is rı (goo3) iee. uniform over rotations and 81 ~ N 7). After centering (i.e. subtracting the mean) this distribution will be uniform over rotations and with translations distributed according to the centered normal (r1, E SE(3), with sf ~ (0, 1). In algorithm 1, also slightly abuse the notation and denote the output of the rotation part of Up as (t. re) and similarly the translation part of Ug as (t, sz). We do not include separate algorithms for DOL.FLLoW-BAE and FOLDFL.OW-OT
they are simple modifications to FOLDFLLOW-SFF If we set 7r (1) = 0 and 7s (1) = 0, then we recover the FOL.FLLOW-OY algorithm. If in addition we remove the resampling in lines 4 and then recover the FOL.DFLOO BASE algorithm.
28

---

Published as a conference paper at ICLR 2024
Algorithm FOL.FL.OW-SFM training on SE(3)*
1: Input: Source and target P1 (21), (o flow network Ug. and diffusion scalings 7,()), (t).
2: while Training do
3:
2,00 70, 1), Po:
4:
(ZT((z,, 21
D OT resampling step to obtain FOLDFLOW-OT
5:
(0), ~
6:
sf f- su N Σ, sn-
$1 N Σ st
mean subtract: (56, ru), (1) E SE(3),
7:
rt - expro (florr, (ri))
geodesic interpolant from eq. (2)
8:
5t - tsi +1 t)s
iinterpolant (Euclidean)
9:
Fr ~ (Gsors) (t) -1) 1))
simulation-free approximation from eq. (9)
10:
S, ~ Nss (t) -t)
FOLDFLOW
1o (t, logr, (ro)
+ oo kt)
81-28
(S(3)
12:
+ Up VoLFOLFFI.w)
13:
return Up
1.2
SDE TRAINING AND INFERENCE
In this section, we outline our training and inference algorithms for the SO(3) component of FOLDFLOW-SFIN The training algorithm is detailed in algorithm 2 while the inference algorithm is provided in algorithm
Algorithm 2 FOLDFLOW-SFII training on SO(3)
1: Input: Source and target Po, diffusion schedule 7.), flow network vo
2: while Training do
3: t, 20, Ø1 (0, 1), Po: Po.
P1
4:
- OT
5:
To: F1 ~ π
6:
Ti expro lo (ri))
geodesic interpolant from eq. (2)
7:
A IGso(3) (r 72 (t))(( - t)) (.(Felro,ri) log,, (ro)
simulation-free approximation from eq. (9)
8:
9:
CFOLDFLOW-SFM (t, (f r1)l3)))))
10:
+ Update(0, VoLFoLnFiow.SFM)
11:
return 00
Algorithm FoldFlow-S -FM Inference
1: Input: Source distribution P1, flow network Up. diffusion schedule 7-), inference annealing i(-),
noise scale, Ç, integration step size At.
2: Sample r1 ~ Pi
3: for in [0, 1/At) do
4:
t+ 1 sAt
5:
Sample z ~ (0,
6:
dB. - S7h vdl.
7:
dB, hat(dBi)
>> map rotation vector to so(3)
8:
ur r, vo(,ott
parallel-transport the vector field to so
9:
Titar t re (urid dBe))
10:
return
I.3 VECTOR FIELD PARAMETRIZATION
Similar to the toy experiment, for the protein modelling case, the architecture is constructed such that the output vector lies the tangent space.
29

---

Published as a conference paper at ICLR 2024
1.3.1 PROTEIN MODEL PARAMETERIZATION
For proteins, we use the FrameDiff architecture (Yim et al., 2023b) over SE(3),, which is based on the structure module of AlphaFold2 (Jumper et al., 2021) following the initial work diffusion models with AF2-like architectures (Anand & Achim, 2022). As described in the main text, this architecture outputs a predicted Cn, which we can then deterministically transform into a vector located at the tangent space of 24. This transformation can be split into 2NV components, the N R components, and the N, SO(3) components. For the SO(3) components we calculate
logr, Ñ 1
(o rt)
(44)
and for the R components we calculate after centering,
8t 50
St 80
vo(4, 8x)
(45)
t
where (fu. Sp) uu wy (t, x.). For R33 (g s;) clearly on the tangent space as R isomorphic to its own tangent space. This is because Euclidean space is flat space. For SO(3), (t rz) is also on the tangent space of r, due to the definition of the log map. Since all components of the product space SE(3)," are on the tangent space, co((, (1)) is on the tangent space of SE(3),,
1.4 PROTEIN TASK HYPERRARAMETERS
FOLDFLOW is implemented in Pytorch, and uses the invariant point attention (IPA) implementations from OpenFold (Ahdritz et al., 2022) the backbone. We use the Adam optimizer with constant learning rate 10-1, 31 = 0.9, 32 0.99. The batch size depends on the length of the protein to
maintain roughly constant memory usage. In practice, we set the effective batch size to eff bs
x 500, 000/N²),
(46)
for each step. We set Apuz = 0.25 and weight the rotation loss with coefficient 0.5 as compared to the translation loss which has weight 1.0.
We also used trick from FrameDiff- Frame-Dif-Improeedd to stabilize the rotation loss. Instead of the L loss on the rotation vector, separate the loss into two components: for the rotation vector. This seemed to reduce variance and numerical instability in the training.
the axis and the angle
1.5 DATA AND DATA SAMPLING
We use a subset of PDB filtered with the same criteria as FrameDiff, specifically, we filter for monomers of length between 60 and 512 (inclusive) with resolution < 5Å downloaded from PDB (Berman et al., 2000) on July 20, 2023. After filtering out any proteins with > 50% loops we are left with 22248 proteins. To support diversity, we sample uniformly over clusters with similarity of 30% suggested in Frame Diffeeiff-mpppp Improved model?. Our model functions most efficiently with batches of proteins of the same length, so we each batch contains proteins of a single length. There 4268 clusters in dataset.
To assess the effects of our sampling methods on protein diversity and length distribution, we present three plots. fig. 13a illustrates the variability and range of protein lengths in the dataset, giving an overview of available lengths for sampling. fig. 13b shows the batch fraction per length, highlighting alterations in sequence length distribution during training due uniform cluster sampling. Fig. 13c, which uses a log scale on both axes, unveils the variation in cluster sizes and the skewness in protein distribution across clusters. Uniform cluster sampling enhances batch diversity, aiding model generalization over various protein sequences and structures. However, as observed in fig, 13b, it slightly modifies the sequence length distribution during training. fig. 13c reveals an unevenness in protein distribution across clusters, with two bins containing approximately 14% of proteins.
txt
30

---

Published as a conference paper at ICLR 2024
1193
0483
Mirber Quees
32:0
10
Sequerce ineeth
Secuwnco Langth
hunber
(a)
(b)
(c)
Figure 13: (a) Distribution of protein lengths our dataset. (b) Distribution of protein lengths a batch (when sampling uniformly by cluster) (c) Distribution of number of proteins per cluster.
1.6 PROTEIN METRICS
The TM-score.
The template modeling score (TM-score) measures the similarity between two
protein structures. The TM score can be expressed for two protein backbones Io, Ti E SE((3) as
Nommmon
1
TM-score(Zo, .1) max
(47)
Ntarget
where Ntarget is the length of the target sequence, Neommon is the length of the common sequence after 3D structural alignment, d: is the distance (post alignment) of the jtth residues in zo, and 21, and (N) = 1.24(N - 15)1/3 - 1.8 is a scaling factor to normalize across protein lengths. The TM-score ranges between (0, 1] with TM-score of 1 indicating perfectly aligned structure. In general TM-score > 0.5 are considered roughly similar folds, with TM-score < 0.2 corresponding to randomly chosen unrelated proteins.
The RMSD metric. The poot-mean-squee deviation (RMSD) is simple metric over paired residues expressed as
L
d² L
RMSD(20. x2)
(48)
where d; again the distance between the 7th residues heavy atoms [N, C., CJ. The RMSD score is length dependent unlike the TM-score, but has been shown to be a more stringent filtering step then TM-score 0.5 for designability (Watson et al., 2023). In general, compared TM-score the RMSD metric is more sensitive local errors and less sensitive to global misalignments.
Designability. generated protein structure considered designable if there exists an amino acid sequence which refolds to that structure. We first generate 50 proteins at lengths ( (100, 150, 200, 250, 300 ). then apply ProteinMPNN with sampling temp = 0.1 8 times to generate sequences for every generated structure. Finally apply default ESMFold and aligned RMSD of the Ca backbone atoms to calculate alignment of each ESMFold-refolded structure with the generated structure. We determine protein designable if at least one of the 8 refolded structures has scRMSD score < 2.0. While threshold of < 2.0 for designability is standard, this threshold may be unreasonably strict for longer backbones. However, it is unclear how this threshold should decay with increasing sequence length.
Finally, we note the imperfection of the ceff-consistency designability metric: when ESMFold does not produce the same structure as FOLDFLOW it does not imply FOLDFLOW's structure is wrong, especially for longer sequences where protein folding models are known perform worse. Both ProteinMPPNN and ESMFold are imperfect, and the failure cases of these models has not been well characterized. While the false positive rate of this metric appears be low, the false negative of this metric has not been quantified.
Diversity. We calculate all pairwise TM-scores for all generated structures that achieve the des- ignability threshold of scRMSD < 2 for each length of protein. We then compute the mean over all of these pairwise TM-scores as our diversity metric. For this metric, a lower score is better. We choose compare diversity designable proteins as do not want the designability score to be
31

---

Published as a conference paper at ICLR 2024
inflated by models which produce poor, proteins that may be very dissimilar to the space of refoldable proteins at that length.
Novelty. We calculate novelty using two metrics. The first the minimum TM-score of designable generated proteins to the training data as described in Appendix I.5. The second metric is motivated by previous research (Lin & AlQuraishi, 2023) and is the fraction of proteins that are both designable (scRMSD 2 A) and novel (avg. max TM-score 0.5). We note that all models are not trained on the same dataset: FOLDFLOW and CraameDif-RRetraadd FrameDiif-ICML use very similar training datasets (only differing in about 10% of structures). However, Genie and RFdiffusion use substantially larger datasets. Genie is trained on the Swissprot high-
share their dataset andFOLDFLOW and
database (Jumper al., 2021; Varadi et al., 2021) and REdiffusion is least pretrained confidence AlphaFold2 structures. These larger training sets may cause novelty to be overestimated for these models as there are structures their training sets that are far from the training set we use to test novelty against.
Error bounds in Table 2. We also report the standard error of the novelty and designability metrics in table 2. This is calculated by taking the standard error for each metric per sequence length, and then taking the mean over sequence lengths. We note that as the diversity is calculated as the averaged pairwise distances of designable proteins, each estimate of the mean is correlated resulting in invalid estimate of the standard error.
EXTENDED ABLATION OF EXPERIMENTS
Table 6: Ablation study of FOLDFLOWw features (stochasticity, optimal transport, auxiliary losses and inference annealing) against designability, diversity and novelty metrics.
Designability(1)
Diversity(()
Novelty
Stochas.
OT
Aux. Loss
Inf. an- nealing
max(()
fractionn())
0.228
0.230
0.440
0.172
X
X
X
0.648
0.267
0.447
0.412
X
X
0.132
0.235
0.432
0.096
X
J
X
0.657
0.264
0.452
0.432
X
J
0.112
0.209
0.414
0.088
J
X
0.592
0.247
0.419
0.424
X
J
X
X
0.152
0.190
0.443
0.108
J
0.820
0.247
0.460
0.484
X
0.128
0.198
0.394
0.100
0.580
0.253
0.439
0,416
0.164
0.196
0.427
0.120
0.684
0.253
0,412
0.500
0.188
0.215
0.449
0.136
0.632
0.257
0.433
0.432
0.268
0.210
0.446
0.188
0.716
0.251
0.411
0.544
We perform a complete ablation experiment for the four features of the FOLDFLOW models: stochas- ticity. optimal transport, auxiliary losses in training and inference annealing. For each of the experiments, we evaluate the performance of the model on the designability, diversity and novelty metrics. These results be seen in table 6. Overall, we observe the following trends:
Stochasticity improves robustness and the ability of the model to generate novel proteins in 7/8 settings, as observed in the fraction of novel and designable proteins (Novelty-fran fraction....
Optimal transport improves the designability of the model by reducing the variance in the training objective in 7/8 settings.
The auxiliary losses improve the designability the models in 7/8 settings.
Inference annealing improves the performance of all FOLDFL.Ow models in all metrics.
32


Reviews:
[{"review_id": "LINjKZhUZ9", "reviewer": "Reviewer_4M3T", "summary": "The paper introduces a series of novel generative models called FoldFlow, which are designed to accurately model protein backbones. The models are based on the flow-matching paradigm over rigid motions, allowing for the construction of stable and fast training methods.\nThe first model introduced is FoldFlow-Base, a simulation-free approach that learns deterministic continuous-time dynamics and matches invariant target distributions. This model is then improved upon by incorporating Riemannian optimal transport, resulting in FoldFlow-OT, which creates more simple and stable flows.\nLastly, the authors design FoldFlow-SFM, which combines Riemannian optimal transport and simulation-free training to learn stochastic continuous-time dynamics.", "questions": "See above. If my concerns are addressed, I am willing to increase my score.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "The paper introduces a novel approach called FoldFlow for generative modeling of protein structures. This approach is based on the flow-matching paradigm over rigid motions. The authors claim that incorporation of Riemannian optimal transport and simulation-free training techniques improves the stability and efficiency of the models compared to diffusion-based approaches.\nThe authors present a series of models within the FoldFlow framework, each with increasing modeling power.", "weaknesses": "The work raised a couple of questions.\n\n1. Why should optimal transport be a good prior for protein backbone generation? Section 3.2 on FoldFlow-OT provides little explanation why this should overall improve the quality of the generated samples. Could you elaborate on that? Similarly, it would be great if you could provide an intuition for motivating the introduction of the stochastic version. The results (Table 2) also do not suggest that there is any benefit in introducing three variants.\n\n3. Studying Table 2, the authors confirm that FoldFlow significantly underperforms RFDiffusion. Also, as raised in Question 1, there is no pattern of improvement between each method of FoldFlow. My main concern regards the stability of the results? You seem to report a single value, is this a lucky single shot. I would suggest to align the analysis to previous literature, i.e., as conducted in FrameDiff, and report results on several runs and examples. It would also be important to also extend Figure 4b to include a comparison to RFDiffusion, FrameDiff, and Genie in terms of scRMSD. Similarly, why not repeat the analysis in Figure 5 with the considered baselines?\n\n3. You claim that FoldFlow models are designed to be fast, stable, and efficient training. Why are the runtimes of the baselines in Table 2 missing? Regarding Table 3, what are the time differences between -base, -OT, and -SFM? What does FoldFlow in Table 3 refer to? You claim in Appendix B that you \"want to get straighter flows for faster inference and more stable training\". Is this something you observe?\n\n4. Using flow matching methods on protein structures is not new. It would be important to cite previous work here and highlight the differences [1].\n\nI very much enjoyed reading the paper and the direction is exciting. The current experimental validation, however, raises several concerns that need to be addressed. This in particular concerns the absence of a pattern in performance of different FoldFlow variants and chosen reported metrics and comparisons.\n\n[1] Somnath, Vignesh Ram, et al. \"Aligned Diffusion Schrödinger Bridges.\" Conference on Uncertainty in Artificial Intelligence (UAI), 2023.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "gDUXwEL2jz", "reviewer": "Reviewer_PWT4", "summary": "Flow matching, and OT flow matching are applied to the problem of learning generative models of protein backbones. Three variants of the method are provided (1) train a CNF via flow matching, (2) train a CNF via conditional flow matching with minibatch optimal transport, and (3) they replace the deterministic flow model with a model that learns to map between the base and target sources via stochastic dynamics. Strong empirical results are shown when training on structures from the PDB:\n- The OT version shows significant improvement above the base model\n- They improve upon FrameDiff substantially. \n- Their model does not rely on a pretrained model, and is substantially cheaper to train than RosettaDiff\n\nAdditionally, the flow matching paradigm allows for an arbitrary base distribution to be used - and this is demonstrated in an experiment on conformer generation of the BPTI protein where noised samples from pre-trained models are used as the base distribution.", "questions": "- What is the advantage of the FOLDFLOW-SFM over the other models in theory? \n- I find the term “stochastic flow” a bit confusing as typically flow refers to deterministic dynamics - could the authors please clarify?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "- Well motivated: The presentation of the method is clear, the method (flow matching/OT flow matching) is well suited to the problem of protein generation, and is the first to explore this. The ability to use arbitrary base distributions is well suited to confomer generation where we can choose informed base distributions (e.g. using cheminformatics methods, or with existing models e.g. AlphaFold). \n- Strong experimental section: The results (improving upon the FrameDiff model) look really good. The models are trained for substantially less time than the RFDiffusion alternative, and do not rely on pre-training. Good baselines are presented, showing the benefit of the OT-flow over the flow trained with vanilla flow matching.", "weaknesses": "- The FOLDFLOW-SFM does not have much motivation provided (besides empirical results e.g.  noting that they have higher novelty). \n- The experiment on equilibrium confomer generation utilizing the informed base distribution does not compare against an uninformed base distribution. Thus, it is hard to judge exactly how helpful this is (although intuitively it seems like it would be very helpful!).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "wQjnqYQD6B", "reviewer": "Reviewer_zh9j", "summary": "The authors extend previous work on conditional flow matching on Riemannian manifolds and apply it to generate protein backbones. Their theoretical results lead to fast and stable training, and they empirically generate more diverse and designable structures than non-pretrained denoising diffusion probabilistic models of protein backbone. \n\nIn general, I think this is an interesting and strong paper. Most of my concerns (detailed below) are around specifics of the evaluation, and I believe that they are addressable", "questions": "- Are there applications for the arbitrary starting distribution other than trying to sample from the conformational ensemble via generation?\n- Is there a reason to prefer scRMSD instead of scTM, as was seen in some earlier work? \n- Is DP vs DDP something intrinsic to the models, or could, for example, FrameDiff be trained with DDP?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 4, "strengths": "- \\[originality\\] This is in the first cohort of work applying conditional flow matching to generate protein backbones. The paper makes a number of theoretical contributions, such as proving the existence of a Monge map on $SE(3)_0^N$ and a closed-form expression of the target conditional vector field for SO(3). \n- \\[significance\\]  Conditional flow matching is a family of generative models that relaxes many of the constraints of normalizing flows, and extending them to SE(3) enables many applications in biology, chemistry, and materials. The specific application studied here is of great importance and interest. \n- \\[quality\\] The theoretical results seem sound, although I did not carefully check the proofs. Empirically, FoldFlow does generate more designable, diverse, and novel structures than previous comparable methods. The comparisons between FoldFlow-Base, OT, and SFM are thorough and sound. \n- \\[clarity\\] The paper is generally well-written, the contributions are clear, and the empirical results are easy to follow. Despite an obvious need to compress a lot of material into the page limit, the paper is well-cited and well-situated in the current literature.", "weaknesses": "The main empirical claim in this paper is that using flow matching instead of diffusion results in better generated protein structure backbones. Therefore, the major weaknesses of the paper are in explaining why we see this result and in showing, with as many other parameters held constant as possible, that the flow matching objective is responsible for improvements in generation quality. In addition, the section on conformation generation lacks the proper ablations to support the claim that flow matching is helpful there, as well as a justification for why this is a useful task. \n\nI will now try to detail these weaknesses in clarity and quality and suggest ways to address them where applicable. Within each section, I will move from more important to less important. \n\n### Clarity\n\n- While the theoretical contribution lies in enabling stochastic flow matching for protein backbone generation, there is no theoretical or qualitative discussion of why flow matching should result in more efficient training and higher-quality generations. \n- In general, the authors have the unenviable task of explaining the background for both protein generation and Riemannian flow matching within the page limit. I think people from outside the fields will have a lot of trouble following the more technical parts of the exposition, but I'm not sure it's possible to do much better given the depth of material and the page limit. \n- Is generating from the equilibrium distribution over conformations after training on data from molecular dynamics useful? I can see how this would be useful if it generalized to new structures, but it seems that if you've already done the MD for that particular model, using another model to sample from the distribution is superfluous. \n- There is some information relegated to the supplement that should be in the main text. For example, that the authors use the same neural architecture as FrameDiff but a slightly different training set is very important for interpreting the empirical results. \n- The two paragraphs beginning \"Our approach\" and the final list of main contributions at the end of the introduction are fairly redundant, with several sentences of repeated content. Given how much theory, experiment, and background the authors need to cover, reducing redundancy here could help clarify other parts of the paper. \n- In section G.1.1, I think \"Here we find that FOLDFLOW is over 2x faster than FoldFlow-Improved per step\" should be \"Here we find that FOLDFLOW is over 2x faster than FrameDiff-Improved per step.\"\n\n\n### Quality\n\n- The claim that FoldFlow outperforms FrameDiff in particular and diffusion in general would be stronger if the training set and other train compute were matched between FrameDiff-improved and FoldFlow. As it is, they are very similar, but the slight discrepancy in train sets and hyperparameters such as batching weaken the comparison. \n- The section on equilibrium conformation generation doesn't make a strong case that FoldFlow is better than (1) diffusion from a simple prior trained on the same data or (2) flow matching from a simple prior trained on the same data. \n- In Table 2, the numbers for Diversity and novelty should include uncertainties. It would also be nice to have scRMSD in this table. \n- Ablations of FoldFlow-OT and -SFM without the Aux loss or inference annealing would strengthen the case that the auxiliary loss and inference annealing are universally helpful. \n- FoldFlow-Base and FoldFlow-SFM do poorly on designability for longer proteins. This may be because the deduplicated dataset does not have very many longer proteins.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "DdZb67zqC1", "reviewer": "Reviewer_Nm3A", "summary": "This paper presents FoldFlow models, extensions of Riemannian flow matching objective for distributions over $SE(3)^N$ with applications to protein backbone generation. Specifically, FoldFlow-Base leverages Riemannian FM defined on $SO(3)$. FoldFlow-OT further enhances the approach by Riemannian optimal transport. FoldFlow-SFM introduces stochasticity into the flow model. The models have been comprehensively benchmarked on protein backbone generation tasks and achieved strong performance compared with the diffusion-based baselines.", "questions": "1. Is it possible to analyze how the sampling steps affect the quality of the generated samples? Presenting such results would give the readers better insight how the generation process works for this flow-matching based model, especially compared with diffusions (e.g., FrameDiff).\n\n2. For the ablation study, it would be great to have the ablate the four parts (Stochas., OT, Aux loss, inf. annealing) on top of the full version, besides the incremental approach adopted in this paper. That is, it would be interesting to see how the performance would be if we remove  one of the four parts respectively from the full version with all four techniques.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "1. The method is very well motivated as an application of flow matching on data with the form of $SE(3)^N$, e.g., protein backbones.\n\n2. The presentation is very clear and the paper is well prepared.\n\n3. The experimental evaluations are interesting and the results are promising.", "weaknesses": "1. It would be better if some detailed results on the generation process are presented (see Q1).\n\n2. Ablation studies could be improved to provide a full picture on the proposed techniques/tricks (see Q2).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (spotlight)