Diffeomorphic Optimization
Ludwig Winkler∗ Microsoft Research AI4Science
Andrew Leaver Fay Prescient Design, Genentech Roche
arXiv:2607.00947v1 [cs.LG] 1 Jul 2026
Joseph Kleinhenz Prescient Design, Genentech Roche
Pan Kessel Prescient Design, Genentech Roche
Abstract Generative models learn data distributions that reside on a low-dimensional manifold within a higher-dimensional ambient space. Optimizing differentiable objectives on this manifold is challenging: the ambient loss landscape is highdimensional, rugged, and non-convex. Direct gradient descent, blind to the manifold’s geometry, quickly drifts off it. Diffeomorphic optimization starts from the observation that diffusion and flow models provide a map from the data manifold to a much simpler base space in which we perform gradient descent. Using differential geometry, we show this is equivalent to Riemannian gradient descent on the data manifold up to O(λ2 ) corrections, keeping trajectories on-manifold by construction and yielding a smoother optimization surface. For protein design, we extend diffeomorphic optimization to the matrix Lie groups SO(3) and SE(3), deriving an autograd-compatible SO(3) gradient and a generalized adjoint-state method for backpropagation through Lie-group ODE solvers. Diffeomorphic optimization improves over tuned guidance on secondary-structure targeting with FrameFlow (91.3% vs. 63.3% of residues in the Ramachandran target), outperforms OC-Flow on peptide binding affinity at 2× the speed, and reduces Rosetta energies by thousands of units across the PDB test set for structures with hundreds of residues.
1
Introduction
Machine learning data is typically concentrated on a low-dimensional manifold (Fefferman et al., 2016; Brown et al., 2022; Kiani et al., 2024). In practice, these manifolds are not known explicitly. Many objectives we care about, such as the stability of a protein or the fidelity of a generated image, are in fact defined on low-dimensional data manifolds embedded in high-dimensional spaces. Direct optimization in this space is difficult: the objective landscape is highly non-convex, the data manifold is implicit, and unconstrained updates, with e.g. direct gradient descent in the data space, tend to produce out-of-distribution solutions. Diffusion and flow-based generative models provide a means to address this difficulty. These models learn a diffeomorphic map from a tractable base distribution to the data distribution. This transformation reparameterizes the data manifold: complex, multimodal regions in data space correspond by design to smoother and more regular regions in the base space of the prior distribution. The geometry of the manifolds required for optimization are already implicitly learned by these models. While generative models have been highly successful in sampling from such manifolds, their use as a foundation for subsequent optimization is less developed. In this work, we propose diffeomorphic optimization - a novel method to optimize arbitrary differentiable cost functions on the data manifold. Our method harnesses the recent advances in flow matching and diffusion models, is applicable to a wide range of tasks, and scales to high dimensional ∗
Work done during internship at Genentech
Figure 1: Optimization of protein with undesired properties (red) to a protein of desired properties (green). Right: Direct optimization in target space gets stuck in local minima. Left: Optimization in base space leads to smoother objective that allows for mode-switching. data, such as proteins with several hundreds of residues. The proposed framework yields a inference time compute method for frozen, pretrained diffusion and flow models. For a visual intuition of this idea, we refer to Figure 1. Existing methods that try to optimize objectives with pretrained models come with caveats: guidance approximates an intractable conditional density Dhariwal & Nichol (2021); reward fine tuning requires retraining Clark et al. (2023), and optimal control methods inject control into the flow and can pull trajectories off manifold Wang et al. (2024). None of these apply to diffeomorphic optimization. The approach diffeomorphic optimization takes is the following. Diffusion and flow models learn a diffeomorphic (smoothRand invertible) map g : Z → X by integrating the ordinary differential equa1 tion g(z) ≡ x1 = z + 0 dτ vθ (xτ , τ ) that maps a base space Z, equipped with a simple probability density qZ , to the target space X , for example, the space of atom positions. Consider a cost function L : X → R, e.g. the energy or thermostability of a protein, that we want to optimize on the data manifold D. Since the map g is diffeomorphic, we change coordinates to the base space variables, L ◦ g : Z → R, z 7→ L(g(z)) .
(1)
We can then perform gradient descent with learning rate λ ∈ R in the base space variables z (i+1) = z (i) − λ∇z L(g(z (i) )) ,
for i ∈ {1, . . . , n} ,
(2)
and map the final base point z (n) to its corresponding x(n) = g(z (n) ). This parameterization has several advantages: the map g is bijective and thus no information is lost. The data distribution is often highly multimodal. However, the density of the base space qZ is chosen to be a simple unimodal distribution, such as a normal distribution N (0, 1). Thus, the base space reparameterization makes it easier to switch modes leading to a smoother loss landscape. The optimization task is complicated by the fact that generative models for proteins often do not sample the atom positions directly but rather the SE(3) backbone frames and SO(2) sidechain dihedral angles Yim et al. (2023b,a); Jing et al. (2024); Jumper et al. (2021). We will therefore derive two efficient methods to facilitate backpropagation through ODE solvers on matrix Lie groups in Section 4. The first relies on repurposing existing autograd engines for calculating the Riemannian gradient. For the second, we derive a suitable generalization of the adjoint-state method for matrix Lie groups. We showcase diffeomorphic optimization in several numerical experiments. Specifically, we demonstrate that we can start from a given protein and then controllably change (parts of) its secondary structures using the flow matching model FrameFlow Yim et al. (2023b) as the diffeomorphic map. Using the DiffDock diffusion model Corso et al. (2023), we show that we can optimize the Vina docking score Trott & Olson (2010) for a protein-ligand pocket and peptide, and improve the stability and affinity of peptides Wang et al. (2024). Furthermore, we demonstrate that we can minimize the Rosetta energy function of a given protein in the base space of the AlphaFlow model Jing et al. (2024). This relaxation protocol leads to significantly lower energies than the state-of-the-art Rosetta Relax. 2
Our work coincides with a recent rise of inference computation for protein design. Practitioners often sample thousands of designs, rank, and submit only a handful for further wetlab experiments Bennett et al. (2024); Abramson et al. (2024); Watson et al. (2023); Frey et al. (2025). The reason for this is that experimental wet lab capacity (and not sampling costs) is the main bottleneck. Furthermore, extensive sampling has been shown to lead to pronounced performance improvements, for example, for antibody-antigen coupling Abramson et al. (2024). Diffeomorphic optimization provides a more targeted approach to obtain high quality samples compared to brute-force sampling followed by ranking and selection. It does not require finetuning (unlike RL or reward finetuning), auxiliary networks (as in guidance) or rejection sampling (best-of-N). It can be applied to any frozen, pretrained, flow or diffusion model with a differentiable reward. We therefore believe that it will become an important part of the inference compute toolbox for proteins. Briefly summarized, our key contributions are the following. We propose diffeomorphic optimization that is applicable to any differentiable optimization objective. As we show theoretically, it automatically enforces the manifold constraint and leads to a smoother optimization landscape. To enable this, we propose two effective methods to facilitate backpropagation through matrix Lie group ODE solvers and will provide efficient PyTorch implementations for it. Experimentally, we apply diffeomorphic optimization to tasks of high practical relevance, i.e. protein ligand docking, Rosetta energy relaxation, and secondary structure modification, by combining it with state-of-the-art generative models in the field of protein generation, specifically FrameFlow, DiffDock, and AlphaFlow.
2
Related Works
Protein hallucination is a version of computational protein design which uses backpropagation through a folding model to its input sequence Anishchenko et al. (2021); Kosugi & Ohue (2022); Goverde et al. (2023); Pacesa et al. (2024); Cho et al. (2025). Gradient descent and normalizing flows has been explored in the explainability literature to generate counterfactual explanations Joshi et al. (2019); Dombrowski et al. (2021, 2023); Dhurandhar et al. (2018) although not for flow matching and diffusion models. Ben-Hamu et al. (2024) proposes to differentiate through flows for controlled generation. Our work builds on this reference by generalizing it to matrix Lie groups, which is of high relevance for proteins. We also provide a detailed theoretical analysis of the method and derive an efficient adjoint state method. Wang et al. (2024) explores related ideas for matrix groups in the framework of optimal control. Specifically, the authors propose a nice framework which adds an additive control to the vector field of the flow. Our approach does not use a control but rather optimizes the initial condition following Ben-Hamu et al. (2024). We discuss the relationship to this reference in more detail in Appendix B and compare in detailed numerical experiments to their approach. Liu et al. (2023) similarly relies on control variables but directly applies gradient descent to them. This reference does however not consider matrix groups. The adjoint state method on manifolds has been discussed in other works for charts Lou et al. (2020); Mathieu & Nickel (2020) or particular manifolds Rezende & Mohamed (2020); Winkler et al. (2024); Bacchio et al. (2023); Albergo & Vanden-Eijnden (2022). Guidance is a widely used method to bias diffusion and flow-matching models towards certain desiderata is guidance. There exist various flavors of it, such as classifier-based guidance Dhariwal & Nichol (2021), classifier-free guidance Ho & Salimans (2022), and universal guidance Bansal et al. (2023). Further information can be found in the appendix B.
3
Diffeomorphic Optimization stays on Manifold
In this section, we analyze the diffeomorphic optimization procedure theoretically using differential geometry. In particular, we will demonstrate its relation with gradient descent on the manifold. 3.1
Basics Concepts of Differential Geometry
Manifolds and coordinates: a manifold M is a space that locally takes the form of RD , similar to the earth which can locally approximated by flat three-dimensional space. More formally, for each point p ∈ M, there exists a chart φ : U → RD where U is an open subset of M containing the point p. In this sense, the manifold is locally described by a Euclidian D space. The pair (U, φ) is called coordinate chart and the component functions xi of φ(p) = (x1 (p), · · · , xD (p)) are called coordinates. 3
d Tangent space: the tangent space Tp M contains the velocity vectors dt γ(t)|t=0 of curves γ : R → M with γ(0) = p. It can be shown that the tangent space Tp M is a D-dimensional vector space. Let (U, φ) be a coordinate chart on M with coordinates x. We can then define φ ◦ λk (t) = (x1 (p), . . . , xk (p)+t, . . . , xD (p)) with k ∈ {1, . . . , D}. This implicitly defines curves λk : R → M d through p. We denote the corresponding tangent vectors as ∂x∂ k = dt λk (t)|t=0 and it can be shown that they form a basis of the tangent space Tp M. The differential dgp : Tp M → Tf (p) N for a map g : M → V between two manifolds M and V is a linear map between the respective tangent spaces. The curve γv : R → M, corresponding to the vector v ∈ Tp M, is mapped by the differential dgp to d the curve g ◦ γv : R → N . Thus, the differential acts as dgp [v] = dt g ◦ γv (t)|t=0 .
Riemannian metric: A Riemannian manifold is endowed with an inner product ⟨·, ·⟩Gp : Tp M × Tp M → R for each tangent space, which allows us to define a notion of the length of tangent vectors. G is also known as the Riemannian metric. We refer to the excellent textbook Lee (2018). Riemannian gradient: The Riemannian gradient gradG p f ∈ Tp M at point p ∈ M of a function f : M → R is the vector uniquely defined by the relation ⟨v, gradG p f ⟩Gp = dfp [v]
∀v ∈ Tp M :
(3)
and crucially depends on Riemannian metric G of the manifold. In coordinates, the Riemannian P ij \ G gradient is grad f= G ∂ f where Gij is the inverse of the metric tensor G = G( ∂ , ∂ ). p
j
j
ij
∂xi
∂xj
Exponential map: On a Riemannian manifold, we can further define geodesics, i.e. curves γ : I ⊂ R → M that have vanishing acceleration with respect to the covariant derivative induced by the Riemannian connection. For each p ∈ M and v ∈ Tp M, there exists a unique geodesic γv that d satisfies γv (0) = p and dt γv (t)|t=0 = v. The exponential map is then defined by expp : Tp M → M,
v 7→ γv (1) .
(4)
Intuitively, the exponential map formalizes the notion of taking a step in the direction of the tangent vector v and then finding the closest point on the manifold. Gradient descent on manifolds: adding and scaling points on the manifold will not necessarily lead to points on the manifold. In contrast, such operations acting on tangent vectors do lead to other tangent vectors, since a tangent space is a vector space. This suggests a generalization of gradient descent on manifolds: we first scale the gradient of a loss L : M → R, which in particular is a tangent vector, by the learning rate λ ∈ R. To obtain the updated point pi+1 ∈ M, we apply the exponential map at the previous point pi ∈ M to the rescaled gradient: pi+1 ← exppi (−λ gradG pi L) .
(5)
In Rn , the exponential map is expp (v) = p + v leading to standard gradient descent. 3.2
Diffeomorphic optimization is equivalent to Gradient Descent on the Data Manifold
Let g : Z → D be a generative model that maps its latent space Z to the data manifold D. Given such a model, we can then perform gradient descent in its latent space and map the resulting point on the data manifold. We then show: Theorem 1. Let g : Z → D be a diffeomorphic generative model and L : D → R is the loss. Then, up to quadratic corrections in the learning rate λ, performing gradient descent in the latent space Z and then mapping it go the data space D with g is equivalent to gradient descent on the data manifold, i.e., G 2 g(expz (−λ gradG̃ z L ◦ g)) = expg(z) (−λ gradg(z) L + O(λ )) ,
(6)
where G(u, v) = G̃(dg −1 u, dg −1 v) denotes the pushforward of the Riemannian metric G̃ on Z. Proof. See Appendix C.1. In coordinates, the gradient in data space is given by d d grad g(ẑ) L = Jg gradz L ,
with Jg =
4
∂ĝ . ∂ ẑ
(7)
We can perform a SVD decomposition of the Jacobian Jg . In the basis of its left singular values, the learning rate will be scaled by the corresponding singular vectors. The data manifold D is heavily concentrated and thus we expect low values for the singular vectors along these contracted directions. This gives us a mechanism to extract the implicitly learned tangent space from the generative diffusion or flow-matching model, which is (approximately) spanned by the remaining left-singular vectors.
Backpropagation through SE(3) ODE solvers
4
The frame representation of proteins considers an idealized backbone geometry of its heavy atoms [N, Cα , C, O] ∈ R3,4 determined by fixed bond lengths and angles Jumper et al. (2021); Yim et al. (2023b). A protein consists of multiple residues and for each residue i, the main backbone atoms can be described by a simple rotation Ri ∈ SO(3) and translation ti ∈ R3 of these idealized backbone coordinates, as seen in Figure 2, [Ni , Ci , , (Cα )i ] = Ti [N, C, Cα ]
(8)
where Ti = (Ri , ti ) is an element of the three-dimensional special Euclidean group SE(3) = SO(3) ⋉ R3 . The positions of the remaining heavy atoms (backbone oxygen and sidechain carbons) can be fixed by dihedral angles. Therefore, the vector field vθ in the ODE takes value in the tangent space of the corresponding Lie groups SE(3) and SO(2). Integration on SO(2) can simply be performed by standard numerical integration in R3 and then wrapping the result on the circle S 1 . The case of SE(3) is, however, highly non-trival. For example, the group SE(3) has the product rule T1 T2 = (R1 , t1 )(R2 , t2 ) = (R1 R2 , R1 t2 + t1 ) . 3
(9)
3
The translation group R ≃ {(I, t)|t ∈ R }, where I is the unit elFigure 2: SE(3) rotoement of SO(3), is an abelian normal subgroup of SE(3) since translation of the idealized ′ −1 ′ (R, t)(I, t )(R, t) = (I, Rt ). Therefore SE(3) is not semibackbone positions. simple and there is no canonical left- and right-invariant Riemannian metric induced by the Killing form. In the generative model literature, one typically chooses the Riemannian metric to be the sum of the metrics of SO(3) and R3 Yim et al. (2023b). This choice is right-invariant but not left-invariant2 . For this choice, the exponential maps and Riemannian gradients for the translations and rotations decouple expT ((ω, v) = (expR (ω), expt (v)) ,
gradT f = (gradR f, gradt f ) ,
(10)
for ω ∈ so(3), v ∈ r3 ≃ R3 , and f : SE(3) → R. The explicit form of these exponential maps and gradients for both translations and rotations are: Translations: since R3 is a vector space, it holds that gradt f = ∇t f (t, R)
expt (v) = t + v , 3
(11)
3
where we use the canonical isomorphism r ≃ R and ∇t denotes the standard gradient operator ∇ = (∂1 , ∂2 , ∂3 ). As a result, we recover standard gradient descent for the translational part ti+1 = t(i) − λ∇ti L(ti , Ri ) .
(12)
Rotations: as we discuss in Appendix D, the right multiplication Rg : G → G, h 7→ hg induces a isomorphism between the tangent spaces Tg G and the Lie algebra g ≃ Te G for any Lie group G. It is often easier to work with the exponential map exp : g → G in terms of Lie algebra elements. For matrix Lie groups, such as SO(3), this exponential map is simply given in terms of the matrix P∞ vn exponential exp(v) = i=0 n! where v is the matrix Lie algebra element. Similarly, we can uniquely associate the Riemannian gradient gradR f of a function f : SO(3) → R with the Lie algebra element X d ∇f (R) = T a ∇a f ∈ so(3) , ∇a f (R) = f (exp(tT a )R) , (13) dt t=0 a 2
The question of whether the metric is left- or right-invariant is a matter of convention. We choose the Riemannian metric to be induced by push-forward of the Killing form with respect to the right multiplication.
5
where T a denote the antisymmetric generators of the Lie algebra so(3). Therefore, gradient descent on SO(3) amounts to Ri+1 = exp(−λ∇Ri L(Ri , ti ))Ri . Summary: Gradient descent on SE(3) can be performed by i+1 exp −λ∇Ri L(Ri , ti ) Ri R i+1 i T = i+1 = expT i (−λ∇T i L(T )) = t t(i) − λ∇ti L(ti , Ri )
(14)
(15)
In this sense, the gradient descent of the rotational and translational part decouples. 4.1
Gradients of SE(3) Solvers
Let g : SE(3)n → SE(3)n be a diffeomorphism. We restrict to n = 1 for notational simplicity but the results immediately generalize to n > 1. We use the notation (R, t) = g(Z, z) with Z ∈ SO(3) and z ∈ R3 denoting the base space variables. We can thus reparameterize the loss function L : SE(3) → R by L(g(Z, z))
(16)
for which we want to perform gradient descent with respect to the base variables (Z, z) ∈ SE(3). In particular, we can consider flow matching and diffusion models for which the diffeomorphism g is defined by g(Z, z) ≡ (Z1 , z1 ) where the right hand side is the solution of the following initial value problem on SE(3): dTτ = (dZτ , dzτ ) = ( Vθ (Zτ , zτ ), vθ (Zτ , zτ ) ) dτ ,
with
T0 = (Z0 , z0 ) = (Z, z) ,
(17)
where τ ∈ [0, 1] and Vθ ∈ so(3), vθ ∈ R3 are parameterized by neural networks. As a result, we need to backpropagate through the numerical solver of the ODE. Specifically, we need to calculate the Riemannian gradient ∇Z L(g(Z, z)) with respect to rotation Z ∈ SO(3). In the following, we will propose two methods to do so. Repurposing Autograd for SO(3): the Riemannian gradient can be expressed in terms of a simple matrix derivative: Theorem 2. The Riemannian gradient equation 13 on SO(3) of a loss function L : SO(3) → R can be written as df ⊤ ∇L(R) = 2 , (18) R dR A df where dR denotes the standard matrix-calculus gradient with respect to the matrix R ∈ SO(3). Furthermore, we denote the antisymmetric component of a matrix M by [M ]A = 12 (M − M ⊤ ).
Proof. See Appendix C.2. This suggests a straightforward way to calculate Riemannian gradients with modern autograd frameworks. Specifically, one wraps any SO(3)-valued leaf variable in an autograd method acting as the df identity in the forward pass but performing the multiplication equation 18 in its backward where dR is the output gradient. The gradient of this variable will correspond to the antisymmetric Riemannian gradient matrix equation 13, see Appendix A for a very concise pytorch implementation of this. We emphasize that this approach is completely general and ensures that the Riemannian gradient is seamlessly integrated in existing autograd functionality. In particular, we can facilitate backpropagation through ODE solvers on SO(3) by combining this repurposing trick with standard activation checkpointing at intermediate points of the integration trajectory. This allows us to limit the memory footprint of the computational graph to the desired degree. Adjoint State Method on SE(3): the repurposing method has the advantage that it is rather general. However, for flow matching and diffusion models, the gradients can be written in terms of a particular adjoint state ODE. For the translational part, this is well known Chen et al. (2018) and given by dL(g(Z, z)) = a0 dz 6
(19)
Figure 3: Diffeomorphic Optimization on SO(3): left hand side visualizes the gradient descent trajectory in the base space Z. Right hand side visualizes the same trajectory when mapped to the target space X . Diffeomorphic optimization clearly stays on manifold. The green plane is spanned by the left-singular vectors of the Jacobian ∇Z g(Z) scaled by the corresponding singular vectors. This shows that the diffeomorphic map indeed captures the tangent space of the data manifold to good approximation. where the adjoint state a0 ∈ R3 is obtained by solving the following terminal value problem 3 X daτ dv i (Zτ , zτ ) aiτ θ =− , dτ dzτ i=1
a1 =
dL(R, t) , dt
(20)
with τ ∈ [0, 1]. In the appendix, we derive a generalization of the adjoint state method for SO(3): Theorem 3. The Riemannian gradient of the reparameterized loss function ∇Z L(g(Z, z)) = A0
(21)
can be obtained by the following terminal value problem for the Lie-algebra-valued adjoint state P3 Aτ = i=1 Aiτ T i ∈ so(3): 3 X dAτ Aiτ ∇Zτ Vθi (Zτ , tτ ) , = [Vθ (Zτ , tτ ), Aτ ] − dτ i=1
A1 = ∇R L(R, t) ,
(22)
where τ ∈ [0, 1] and [A, B] = AB − BA is the commutator of matrices A and B. Proof. See Appendix C.3. Wang et al. (2024) used the adjoint method on SO(3) to calculate gradients with respect to deterministic optimal control which they used as guidance. As shown in the appendix, this result also holds for any matrix Lie group. A notable advantage of the adjoint state method is that it turns the gradient calculation into a terminal value problem and thus can benefit from modern ODE solvers. However, the adjoint state method has the downside that it assumes that the forward pass has no discretization errors. Depending on the stiffness of the ODE, one may therefore accumulate a systematic error by the discretization effects of both the forward ODE equation 17 and the backward adjoint state ODE equation 22. In practice, we recommend relying on numerical experiments to benchmark the relative performance of the adjoint state method or the repurposed autograd method for the problem at hand. We will provide a modular integration library that implements both methods.
5
Experiments
We demonstrate the generality of diffeomorphic optimization in several numerical experiments. We refer to Appendix H for more details. SO(3) Manifold: In Figure 3, we visualize diffeomorphic optimization for a toy dataset on SO(3). Specifically, we create a dataset of S-shape in the angle-axis vectors of the rotations. We then train a flow matching model to sample SO(3) elements on this manifold. The diffeomorphic optimization 7
Figure 4: We show various snapshots of the trajectory of diffeomorphic maximization of the distance between the termini of the proteins (marked in red). The optimization traverses a number of plausible conformations.
aims to push the sample as close as possible to the point marked by a star on the S-shaped data manifold. As the figure demonstrates, diffeomorphic optimization leads to a gradient descent trajectory on the data manifold. This illustrates that diffeomorphic optimization can harness the differential geometric notions learned by the flow model.
Figure 5: Diffeomorphic optimization is used to change the secondary structure of the protein. For this, we use the ABEGO classification shown on the left. We start from the green protein shown on the top right whose residues are mostly clustered in the A region of the Ramachandran plot which corresponds to α-helices. We then optimize the structure such that these residues are pushed to the B region corresponding to β-sheets. The resulting protein structure is shown in orange in the top right. The final and initial Ramachandran plot are shown on the left. The lower right shows comparison to guidance for 50 samples obtained from the model. Diffeomorphic optimization significantly outperforms guidance despite careful tuning of the baseline hyperparameters described in the appendix.
Table 1: Evaluation of OC-Flow peptide design. Ground-truth PepFlow OC-Flow(trans) OC-Flow(rot) OC-Flow(trans+rot) DiffeoOpt
MadraX ↓ -0.588 -0.195 -0.229 -0.221 -0.263 -0.309
RMSD ↓ – 1.645 1.774 1.643 2.127 1.605
SSR % ↑ – 0.794 0.797 0.794 0.797 0.796
BSR % ↑ – 0.874 0.876 0.872 0.869 0.881
Stability ↓ -84.893 -45.660 -48.380 -48.636 -48.853 -49.417
Affinity ↓ -36.063 -26.538 -27.328 -27.211 -27.468 -28.409
Diversity ↑ – 0.310 0.323 0.310 0.338 0.340
Peptide Design: we compare diffeomorphic optimization to the OC-Flow framework Wang et al. (2024) R 1 in which a control parameter θt is added to the flow in the form of g(z) = z + 0 vθ (xτ ) + θτ dτ , see Appendix B for more details. We kept the experimental setup identical to Wang et al. (2024) and compare to our proposed diffeomorphic optimization of both translation and rotations of the peptide’s backbone. We use the repurposing approach and summarize the results in Table 1. We stress that this improvement can be obtained while being 2× faster in terms of runtime. Secondary Structure Modification with FrameFlow: we use FrameFlow Yim et al. (2023a), a widely used backbone generation flow matching model, to optimize the backbone structure with respect to a given cost function. Figure 4 illustrates an optimization trajectory for which the distance between the protein termini is maximized. This demonstrates that our method can optimize a given 8
protein structure while staying on the data manifold. Similarly, Figure 5 shows that we can also optimize the secondary structure of a protein as specified by the ABEGO classification scheme Wintjens et al. (1996); Kim et al. (2009). For this, we minimize an energy function that is based on the Ramachandran-PAAPP-combined term in the Rosetta energy Alford et al. (2017); Leaver-Fay et al. (2025) which interpolates energies derived from statistics of backbone dihedral preferences in the PDB on a toroidal grid. The grid values are modified such that the undesired regions are disfavored. We compare to guidance with carefully tuned hyperparameters. Diffeomorphic optimization significantly outperforms the while maintaining designability with results in Appendix H.2. Protein-Ligand Docking with DiffDock: DiffDock is a diffusion model that samples both the center of mass translation and rotation (i.e. an SE(3) element) and the torsion angles of the ligand for a given protein Corso et al. (2023). The model is trained on the PDBBind dataset Wang et al. (2005), which contains experimentally determined protein-ligand docking structures. We then apply diffeomorphic optimization to maximize the popular VinaSF score of the protein-ligand complex, utilizing its OpenDock implementation Hu et al. (2024).
Figure 6: Diffeomorphic optimization applied to AlphaFlow exhibits strong improvement as measured in the Rosetta energy score.
To enable backpropagation through the solver, we modify the DiffDock generation process to use probability ODE sampling through our custom integration library. We then optimize the score function on pdb test set samples. We compare diffeomorphic optimization to iid sampling followed by selecting the sample with the best docking score using the same computational budget. For this, each gradient descent step is counted as three sampling trajectories, which is quite generous for the baseline. Figure 7 shows the improvement Sdiffeo − Siid obtained by diffeomorphic optimization which significantly outperforms this baseline. Our method can thus efficiently combine physics-based dockingr scores with trained generative models. Further analysis is provided in Appendix H.3. Minimization of Rosetta Energy with AlphaFlow: we consider the Rosetta energy function Alford et al. (2017) which is widely used in the protein community using the tmol pytorch implementation Leaver-Fay et al. (2025) of beta_nov2016_cart. We select the same pdb test set as in the AlphaFlow publication Jing et al. (2024). This reference proposed several flows that are obtained from protein structure prediction models. We select the ESMFold model for convenience, as it does not require MSA processing. We compare diffeomorphic optimization to the state-of-theart Rosetta Relax protocol. From our theoretical analysis, we expect that diffeomorphic optimization will lead to better mode mixing. Once an energetically favorable mode is reached, it is cheaper to use standard relaxation to find its minimum. We thus Figure 7: Diffeomorphic optimizafirst minimize the Rosetta energy function in the base space of tion leads to better/lower vina score the flow which is then followed by standard Rosetta Relax. than iid sampling with the same budget. Figure 6 shows that diffeomorphic optimization leads to substantially lower energy values throughout the pdb test set. Furthermore, the energy values of the baseline Rosetta Relax do not improve significantly by running a longer optimization or using more seeds. Thus, the relative advantage is not a function of the higher cost of diffeomorphic optimization. Further analysis is provided in Appendix H.4.
9
6
Conclusion and Limitations
We have proposed diffeomorphic optimization which allows us to minimize arbitrary differentiable cost functions on the data manifold. This is achieved with off-the-shelf autograd engines with simple gradient wrappers making it a readily available tool for practitioners. A downside of this approach are the considerable numerical costs due to the backpropagation through the generation. However, sampling costs are of lower concern in protein design where the main bottleneck is experimental wet-lab verification of the designs. It is completely standard in this setting to sample thousands of designs, rank them, and submit only a handful for wetlab experiments. Diffeomorphic optimization provides a more targeted approach to obtain high quality samples.
10
References Josh Abramson, Jonas Adler, Jack Dunger, Richard Evans, Tim Green, Alexander Pritzel, Olaf Ronneberger, Lindsay Willmore, Andrew J Ballard, Joshua Bambrick, et al. Accurate structure prediction of biomolecular interactions with alphafold 3. Nature, 630(8016):493–500, 2024. Michael S. Albergo and Eric Vanden-Eijnden. Building normalizing flows with stochastic interpolants. arXiv preprint arXiv:2209.15571, 2022. URL https://arxiv.org/abs/2209.15571. Rebecca F Alford, Andrew Leaver-Fay, Jeliazko R Jeliazkov, Matthew J O’Meara, Frank P DiMaio, Hahnbeom Park, Maxim V Shapovalov, P Douglas Renfrew, Vikram K Mulligan, Kalli Kappel, et al. The rosetta all-atom energy function for macromolecular modeling and design. Journal of chemical theory and computation, 13(6):3031–3048, 2017. Ivan Anishchenko, Samuel J Pellock, Tamuka M Chidyausiku, Theresa A Ramelot, Sergey Ovchinnikov, Jingzhou Hao, Khushboo Bafna, Christoffer Norn, Alex Kang, Asim K Bera, et al. De novo protein design by deep network hallucination. Nature, 600(7889):547–552, 2021. Simone Bacchio, Pan Kessel, Stefan Schaefer, and Lorenz Vaitl. Learning trivializing gradient flows for lattice gauge theories. Physical Review D, 107(5):L051504, 2023. Arpit Bansal, Hong-Min Chu, Avi Schwarzschild, Soumyadip Sengupta, Micah Goldblum, Jonas Geiping, and Tom Goldstein. Universal guidance for diffusion models. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 843–852, 2023. Heli Ben-Hamu, Omri Puny, Itai Gat, Brian Karrer, Uriel Singer, and Yaron Lipman. D-flow: Differentiating through flows for controlled generation. arXiv preprint arXiv:2402.14017, 2024. Nathaniel R Bennett, Joseph L Watson, Robert J Ragotte, Andrew J Borst, Déjenaé L See, Connor Weidle, Riti Biswas, Ellen L Shrock, Philip JY Leung, Buwei Huang, et al. Atomically accurate de novo design of single-domain antibodies. biorxiv, 2024. Bradley CA Brown, Anthony L Caterini, Brendan Leigh Ross, Jesse C Cresswell, and Gabriel Loaiza-Ganem. Verifying the union of manifolds hypothesis for image data. arXiv preprint arXiv:2207.02862, 2022. Gabriel Cardoso, Yazid Janati El Idrissi, Sylvain Le Corff, and Eric Moulines. Monte carlo guided diffusion for bayesian linear inverse problems. arXiv preprint arXiv:2308.07983, 2023. Ricky TQ Chen, Yulia Rubanova, Jesse Bettencourt, and David K Duvenaud. Neural ordinary differential equations. Advances in neural information processing systems, 31, 2018. Yehlin Cho, Martin Pacesa, Zhidian Zhang, Bruno Correia, and Sergey Ovchinnikov. Boltzdesign1: Inverting all-atom structure prediction model for generalized biomolecular binder design. bioRxiv, pp. 2025–04, 2025. Kevin Clark, Paul Vicol, Kevin Swersky, and David J Fleet. Directly fine-tuning diffusion models on differentiable rewards. arXiv preprint arXiv:2309.17400, 2023. Gabriele Corso, Hannes Stärk, Bowen Jing, Regina Barzilay, and Tommi Jaakkola. Diffdock: Diffusion steps, twists, and turns for molecular docking. In International Conference on Learning Representations (ICLR), 2023. Prafulla Dhariwal and Alexander Nichol. Diffusion models beat gans on image synthesis. Advances in neural information processing systems, 34:8780–8794, 2021. Amit Dhurandhar, Pin-Yu Chen, Ronny Luss, Chun-Chen Tu, Paishun Ting, Karthikeyan Shanmugam, and Payel Das. Explanations based on the missing: Towards contrastive explanations with pertinent negatives. Advances in neural information processing systems, 31, 2018. Ann-Kathrin Dombrowski, Jan E Gerken, and Pan Kessel. Diffeomorphic explanations with normalizing flows. In ICML Workshop on Invertible Neural Networks, Normalizing Flows, and Explicit Likelihood Models, 2021. 11
Ann-Kathrin Dombrowski, Jan E Gerken, Klaus-Robert Müller, and Pan Kessel. Diffeomorphic counterfactuals with generative models. IEEE Transactions on Pattern Analysis and Machine Intelligence, 46(5):3257–3274, 2023. Zehao Dou and Yang Song. Diffusion posterior sampling for linear inverse problem solving: A filtering perspective. In The Twelfth International Conference on Learning Representations, 2024. Charles Fefferman, Sanjoy Mitter, and Hariharan Narayanan. Testing the manifold hypothesis. Journal of the American Mathematical Society, 29(4):983–1049, 2016. Nathan C Frey, Isidro Hötzel, Samuel D Stanton, Ryan Kelly, Robert G Alberstein, Emily Makowski, Karolis Martinkus, Daniel Berenberg, Jack Bevers III, Tyler Bryson, et al. Lab-in-the-loop therapeutic antibody design with deep learning. bioRxiv, pp. 2025–02, 2025. Casper A Goverde, Benedict Wolf, Hamed Khakzad, Stéphane Rosset, and Bruno E Correia. De novo protein design by inversion of the alphafold structure prediction network. Protein Science, 32 (6):e4653, 2023. Jonathan Ho and Tim Salimans. Classifier-free diffusion guidance. arXiv preprint arXiv:2207.12598, 2022. Qiuyue Hu, Zechen Wang, Jintao Meng, Weifeng Li, Jingjing Guo, Yuguang Mu, Sheng Wang, Zheng Liangzhen, and Yanjie Wei. OpenDock: A pytorch-based open-source framework for protein-ligand docking and modelling. Bioinformatics, pp. btae628, 10 2024. ISSN 1367-4811. doi: 10.1093/bioinformatics/btae628. URL https://doi.org/10.1093/bioinformatics/ btae628. Bowen Jing, Bonnie Berger, and Tommi Jaakkola. Alphafold meets flow matching for generating protein ensembles. In Forty-first International Conference on Machine Learning, 2024. Shalmali Joshi, Oluwasanmi Koyejo, Warut Vijitbenjaronk, Been Kim, and Joydeep Ghosh. Towards realistic individual recourse and actionable explanations in black-box decision making systems. arXiv preprint arXiv:1907.09615, 2019. John Jumper, Richard Evans, Alexander Pritzel, Tim Green, Michael Figurnov, Olaf Ronneberger, Kathryn Tunyasuvunakool, Russ Bates, Augustin Žídek, Anna Potapenko, et al. Highly accurate protein structure prediction with alphafold. nature, 596(7873):583–589, 2021. Bobak Kiani, Jason Wang, and Melanie Weber. Hardness of learning neural networks under the manifold hypothesis. Advances in Neural Information Processing Systems, 37:5661–5696, 2024. Patrick Kidger, James Foster, Xuechen Chen Li, and Terry Lyons. Efficient and accurate gradients for neural sdes. Advances in Neural Information Processing Systems, 34:18747–18761, 2021. David E Kim, Ben Blum, Philip Bradley, and David Baker. Sampling bottlenecks in de novo protein structure prediction. Journal of molecular biology, 393(1):249–260, 2009. Takatsugu Kosugi and Masahito Ohue. Solubility-aware protein binding peptide design using alphafold. Biomedicines, 10(7):1626, 2022. Andrew Leaver-Fay, Jeff Flatten, Alex Ford, Joseph Kleinhenz, David Solberg, Henry amd Baker, Andrew M Watkins, Brian Kuhlman, and Frank DiMaio. tmol: a gpu-accelarated, pytorch implementation of rosetta’s relax protocol (manuscript in preparation), 2025. URL https:// github.com/uw-ipd/tmol. John M Lee. Introduction to Riemannian manifolds, volume 2. Springer, 2018. Xuechen Li, Ting-Kam Leonard Wong, Ricky TQ Chen, and David K Duvenaud. Scalable gradients and variational inference for stochastic differential equations. In Symposium on Advances in Approximate Bayesian Inference, pp. 1–28. PMLR, 2020. Yaron Lipman, Marton Havasi, Peter Holderrieth, Neta Shaul, Matt Le, Brian Karrer, Ricky TQ Chen, David Lopez-Paz, Heli Ben-Hamu, and Itai Gat. Flow matching guide and code. arXiv preprint arXiv:2412.06264, 2024. 12
Xingchao Liu, Lemeng Wu, Shujian Zhang, Chengyue Gong, Wei Ping, and Qiang Liu. Flowgrad: Controlling the output of generative odes with gradients. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 24335–24344, 2023. Aaron Lou, Derek Lim, Isay Katsman, Leo Huang, Qingxuan Jiang, Ser Nam Lim, and Christopher M De Sa. Neural manifold ordinary differential equations. Advances in Neural Information Processing Systems, 33:17548–17558, 2020. Emile Mathieu and Maximilian Nickel. Riemannian continuous normalizing flows. Advances in Neural Information Processing Systems, 33:2503–2515, 2020. Martin Pacesa, Lennart Nickel, Christian Schellhaas, Joseph Schmidt, Ekaterina Pyatova, Lucas Kissling, Patrick Barendse, Jagrity Choudhury, Srajan Kapoor, Ana Alcaraz-Serna, et al. Bindcraft: one-shot design of functional protein binders. bioRxiv, pp. 2024–09, 2024. Danilo Jimenez Rezende and Shakir Mohamed. Normalizing flows on tori and spheres. In International Conference on Machine Learning, pp. 8083–8092, 2020. Raghav Singhal, Zachary Horvitz, Ryan Teehan, Mengye Ren, Zhou Yu, Kathleen McKeown, and Rajesh Ranganath. A general framework for inference-time scaling and steering of diffusion models. arXiv preprint arXiv:2501.06848, 2025. Brian L Trippe, Jason Yim, Doug Tischer, David Baker, Tamara Broderick, Regina Barzilay, and Tommi Jaakkola. Diffusion probabilistic modeling of protein backbones in 3d for the motifscaffolding problem. arXiv preprint arXiv:2206.04119, 2022. Oleg Trott and Arthur J Olson. Autodock vina: Improving the speed and accuracy of docking with a new scoring function, efficient optimization, and multithreading. Journal of Computational Chemistry, 31(2):455–461, 2010. doi: 10.1002/jcc.21334. Luran Wang, Chaoran Cheng, Yizhen Liao, Yanru Qu, and Ge Liu. Training free guided flow matching with optimal control. arXiv preprint arXiv:2410.18070, 2024. Renxiao Wang, Xueliang Fang, Yipin Lu, Chao-Yie Yang, and Shaomeng Wang. The pdbbind database: methodologies and updates. Journal of medicinal chemistry, 48(12):4111–4119, 2005. Joseph L Watson, David Juergens, Nathaniel R Bennett, Brian L Trippe, Jason Yim, Helen E Eisenach, Woody Ahern, Andrew J Borst, Robert J Ragotte, Lukas F Milles, et al. De novo design of protein structure and function with rfdiffusion. Nature, 620(7976):1089–1100, 2023. Ludwig Winkler, César Ojeda, and Manfred Opper. Stochastic control for bayesian neural network training. Entropy, 24(8):1097, 2022. Ludwig Winkler, Lorenz Richter, and Manfred Opper. Bridging discrete and continuous state spaces: Exploring the ehrenfest process in time-continuous diffusion models. In International Conference on Machine Learning, pp. 53017–53038. PMLR, 2024. René T Wintjens, Marianne J Rooman, and Shoshana J Wodak. Automatic classification and analysis of αα-turn motifs in proteins. Journal of molecular biology, 255(1):235–253, 1996. Luhuan Wu, Brian Trippe, Christian Naesseth, David Blei, and John P Cunningham. Practical and asymptotically exact conditional sampling in diffusion models. Advances in Neural Information Processing Systems, 36:31372–31403, 2023. Yu Xie, Ludwig Winkler, Lixin Sun, Sarah Lewis, Adam Foster, Jose Jimenez-Luna, Tim Hempel, Michael Gastegger, Yaoyi Chen, Iryna Zaporozhets, et al. Enhanced diffusion sampling: Efficient rare event sampling and free energy calculation with diffusion models. In ICML 2026 Workshop on Structured Probabilistic Inference {\&} Generative Modeling, 2026. Jason Yim, Andrew Campbell, Andrew YK Foong, Michael Gastegger, José Jiménez-Luna, Sarah Lewis, Victor Garcia Satorras, Bastiaan S Veeling, Regina Barzilay, Tommi Jaakkola, et al. Fast protein backbone generation with se (3) flow matching. arXiv preprint arXiv:2310.05297, 2023a. Jason Yim, Brian L Trippe, Valentin De Bortoli, Emile Mathieu, Arnaud Doucet, Regina Barzilay, and Tommi Jaakkola. Se (3) diffusion model with application to protein backbone generation. arXiv preprint arXiv:2302.02277, 2023b. 13
A
Riemannian Gradient Implementation for SO(3)
As derived in Theorem 2, we can efficiently modify existing autograd engines to calculate the SO(3) Riemannian gradient. Specifically, we can wrap the tensor by an operation that acts like the identity in the forward pass and adjusts the gradient to coincide with the Riemannian gradient in the backward pass using Theorem 2. In PyTorch, this can most conveniently be implemented with a backward hook as shown below. We emphasize that this implementation is completely general and ensures that the Riemannian gradient is seamlessly integrated in existing autograd functionality. 1 2
import torch from scipy . spatial . transform import Rotation
3 4 5
def so3tensor (* args , ** kwargs ) : x = torch . tensor (* args , ** kwargs )
6
def so3_backward_hook ( grad ) : so3_grad = grad @ x . transpose ( -1 , -2) return so3_grad - so3_grad . transpose ( -1 , -2)
7 8 9 10
if x . requires_grad : x . register_hook ( so3_backward_hook ) return x
11 12 13 14 15 16
x = so3tensor ( Rotation . random (1) . as_matrix () , requires_grad = True ) y = so3tensor ( Rotation . random (1) . as_matrix () )
17
loss = ( x - y ) . pow (2) . sum () so3_grad = torch . autograd . grad ( loss , x ) [0] 20 print ( so3_grad ) 18 19 21
>> tensor ([[[ 0.0000 , [ -0.8490 , 24 [ -0.1187 , 22 23
B
0.8490 , 0.1187] , 0.0000 , -2.2494] , 2.2494 , 0.0000]]] , dtype = torch . float64 )
Expanded Related Work
Backpropagation through folding models: protein hallucination is a version of computational protein design which uses backpropagation through a folding model to its input sequence Anishchenko et al. (2021); Kosugi & Ohue (2022); Goverde et al. (2023); Pacesa et al. (2024). Interestingly, even the most recent work Cho et al. (2025) based on AlphaFold3 avoids explicit backpropagation through its diffusion structure module and instead relies on the pair representation - a limitation that could potentially be overcome by our proposed methods. Gradient descent and normalizing flows has been explored in the explainability literature to generate counterfactual explanations Joshi et al. (2019); Dombrowski et al. (2021, 2023); Dhurandhar et al. (2018) although not for flow matching and diffusion models. Ben-Hamu et al. (2024) proposes to differentiate through flows for controlled generation. Our work builds on this reference by generalizing it to matrix Lie groups, which is of high relevance for proteins. We also provide a detailed theoretical analysis of the method and derive an efficient adjoint state method. Wang et al. (2024) explores related ideas for matrix groups in the framework of optimal control. Specifically, the authors propose to add an additive control to the vector field of the flow. Our approach does not use a control but rather optimizes the initial condition following Ben-Hamu et al. (2024). We discuss the relationship to this reference in more detail in the appendix and compare in detailed numerical experiments to their approach. Liu et al. (2023) similarly relies on control variables but directly applies gradient descent to them. This reference does however not consider matrix groups. Adjoint state method on manifolds: the adjoint state method on manifolds has been discussed in other works - however only in terms of charts Lou et al. (2020); Mathieu & Nickel (2020) or particular manifolds Rezende & Mohamed (2020); Bacchio et al. (2023); Albergo & Vanden-Eijnden (2022). The adjoint state method derived in our work is applicable to any matrix Lie group and is 14
particularly efficient for the group SE(3) which plays a crucial role the frame-based protein backbone representation. Guidance: a widely used method to bias diffusion and flow-matching models towards certain desiderata is guidance. There exist various flavors of it, such as classifier-based guidance Dhariwal & Nichol (2021), classifier-free guidance Ho & Salimans (2022), and universal guidance Bansal et al. (2023). There has also been substantial recent interest in biasing the generation with additional reward functions Wu et al. (2023); Trippe et al. (2022); Cardoso et al. (2023); Dou & Song (2024); Xie et al. (2026); Singhal et al. (2025) often relying on sequential Monte Carlo. However, these methods struggle for guidance potentials that are sensitive to fine-grained details of the final sample, such as force fields or certain biochemical properties, as they tend to rely on few-shot denoising and require carefully hyperparameter finetuing of their starting point and time-dependent weighting factor. We discuss the limitations of guidance in more detail in Appendix F. Relation to Wang et al. (2024): a closely related reference is Wang et al. (2024) which outlines a very nice approach to maximize a terminal reward using deterministic optimal control Winkler et al. (2022). Specifically, the authors propose to add an additive control θτ to the vector field of the flow dxτ dxτ = vθ (xτ , τ ) → = vθ (xτ , τ ) + θτ dτ dτ The authors then propose to optimize the control with respect to a (regularized) loss Z 1 1 θ θ 0 L(x1 ) + d(x1 , x1 ) + dτ ||θτ ||2 2 0
(23)
(24)
where xθ1 corresponds to the terminal value under the control θτ and x01 is the terminal value for no control. Furthermore, d(·, ·) denotes a distance. The optimal control approach defines control terms θτ at intermediate time steps τ during the generative process. By construction, the control thus acts upon the generation not only at the initial condition τ = 0 but also for τ > 0. A change in initial condition x0 → x0 + δx0 can then only be interpreted as a control if the latter is not a function ∀τ but rather concentrates all modelling capacity at only τ = 0. In the continuous time setting this would imply a control proportional to a Dirac impulse δ(τ ) in time, seen as follows Z 1 Z 1 x1 = x0 + δx0 + v(xτ )dτ = x0 + (vθ (xτ ) + δ(τ )δx0 ) dτ (25) 0
0
and thus implying a distributional control θτ = δ(τ )δx0 . Note that this distributional control has highly undesirable properties from a numerics standpoint, i.e., it is necessarily divergent for τ = 0 and vanishing for all other values of flow time τ . Intuitively, it would require the control to voluntarily constrain all its modelling capacity to the time of the initial condition τ = 0, something that is hard to achieve in practice with unconstrained optimization. Diffusion and flow models are trained to map a simple base space variable to a target distribution. Changing the base space variable ensures that the samples remain on the learned manifold in the target space. This can not be guaranteed if the trajectory is changed at intermediate points. The control parameters can exploit the loss function L(xθ1 ) while moving the sample off manifold. To counteract such movements, a distance loss term d(xθ1 , x01 ) in the target space is applied but this but simultaneously constraints how much the original loss function can be optimized. Diffeomorphic optimization solves both problems conveniently at the same time by allowing free, unconstrained movement over the manifold while also remaining on the learned manifold at all times.
C
Proofs
C.1
Proof of Theorem 1
Lemma 1. For g : M → N , it holds that g(expp (λv)) = expg(p) (λ dgp (v) + O(λ2 )) for any p ∈ M , v ∈ Tp M , and small λ ∈ R. 15
(26)
Proof. For proving this statement, it is useful to consider normal coordinates. For p ∈ M and a neighborhood U of 0 ∈ Tp M , these coordinates are given by the chart ψ : U ⊂ Tp M → M with ψ(v) = expp (v). The coordinates y α of a point q ∈ M in this chart are obtained by q = expp (y α eα )
(27)
where eα is an orthonormal basis of Tp M and we have used the Einstein summation convention. This implies that y α (expq (v)) = v α
(28)
Using normal coordinates for both M and N , we can expand the function g using Taylor’s theorem α
∂g β 2 α β α β α \ d v + O(λ2 ) = exp \ g(exp g(p) (λ dgp (v) + O(λ )) p (λv)) = g (y ) = g (λv ) = g (0) + λ ∂y β (29) where we have used that the connection Γ vanishes in normal coordinates and the hat symbol denotes the coordindate representation. Since the left and the right hand side are written in terms of covariant objects, the result stated in the theorem follows. Lemma 2. Let g : Z → X be a diffeomorphism with Z being a Riemannian manifold with metric G. Then, it holds that G̃ dgz (gradG z L ◦ g) = gradg(z) L
(30)
where G̃ = g∗ G is the pushforward metric of G.
Proof. This statement is easily shown in coordinates. Let z µ denote coordinates on Z and xα = g α (z). Then, the coordinate representation of the pushforward metric G̃αβ follows by Gµν dz µ dz ν = Gµν
∂z µ ∂z ν α β dx dx = G̃αβ dxα dxβ ∂xα ∂xβ
(31)
In these coordinates, dgz (gradG z L ◦ g) is given by ∂xα µν ∂(L ◦ g) ∂xα µν ∂xσ ∂L ∂L G = G = G̃ασ σ µ ν ∂z ∂z ∂z µ ∂z ν ∂xσ ∂x
(32)
which corresponds to the coordinate representation of the right hand side.
Theorem 1 follows almost immediately using the the two lemmas:
Proof. By Lemma 1, it follows that G 2 g(expz (−λ gradG z L ◦ g)) = expg(z) (−λ dgz (gradz L) + O(λ )) .
(33)
The right hand side can then be rewritten by Lemma 2 as G̃ 2 2 expg(z) (−λ dgz (gradG z L) + O(λ )) = expg(z) (−λ gradg(z) L + O(λ ))
which shows the stated result. 16
(34)
C.2
Proof of Theorem 2
We consider the definition of the Riemannian gradient on the Lie algebra d f (exp(τ T a )R)|τ =0 (35) dτ df d(exp(τ T a )R)αβ = |τ =0 (36) dRαβ dτ df (T a R)αβ (37) = dRαβ df = T a Rσβ (38) dRαβ ασ df T a (RT )βσ | T a = −(T a )T (39) = dRαβ ασ df =− T a (RT )βσ (40) dRαβ σα df T R (41) = −tr T a dR df T 1 R (42) = − tr T a 2 2 dR df T = ⟨T a , 2 R ⟩ (43) dR where we have used the Einstein summation convention and the definition of inner product ⟨A, B⟩ = 1 T 2 trA B for Lie algebra elements A, B ∈ so(3). Recall that the Riemannian gradient is given by X T a ∂a f (R) (44) ∂f ≡ ∂ a f (R) =
a
Due to the orthonormality of the generators, ⟨T a , T b ⟩ = δ ab and denoting the antisymmetric part [M ]A = 12 (M − M T ) of a matrix M , it thus holds that df T R , (45) ∂f = 2 dR A The antisymmetric component arises in this step because for an arbitrary matrix M , we have ⟨T a , M ⟩ = ⟨T a , [M ]A ⟩. C.3
Proof of Theorem 3
The adjoint state is given by Aτ = ∇Zτ L
(46)
for which we want to derive a differential equation for its time evolution. For this, we define the time evolution operator Tϵ : SO(3) → SO(3) ,
Zτ 7→ Zτ +ϵ = exp(ϵVθ (Zτ )) Zτ .
(47)
We use the chain rule equation 90 to derive that Aaτ = ∇aZτ L =
X
∇bZτ +ϵ L Dba Tϵ
(48)
Abτ +ϵ Dba Tϵ .
(49)
b
=
X b
The differential equation 89 of the time evolution operator is given by: Dba Tϵ = ⟨T b , ∇aZτ Tϵ Tϵ⊤ ⟩ = ⟨T b , ∇aZτ (eϵVθ (Zτ ) Zτ ) Zτ⊤ e−ϵVθ (Zτ ) ⟩ , 17
(50) (51)
where we have used that Vθ ∈ so(3) is antisymmetric. We now expand this expression up to quadratic order in the step size ϵ to obtain Dba Tϵ = ⟨T b , ∇aZτ ( (I + ϵVθ (Zτ )) Zτ ) Zτ⊤ (I − ϵVθ (Zτ )) ⟩ + O(ϵ2 ) ba
b
a
b
(52)
, ∇aZτ Vθ (Zτ )⟩ + O(ϵ2 )
= δ + ϵ⟨T , [Vθ , T ]⟩ + ϵ⟨T (53) We now use the fact that 1 1 ⟨T b , [Vθ , T a ]⟩ = − tr T b (Vθ T a − T a Vθ ) = − tr T a T b Vθ − Vθ T b = ⟨T a , [T b , Vθ ]⟩ . (54) 2 2 P Expanding Vθ = c Vθc T c and similarly for [T b , Vθ ] ∈ so(3), we then obtain that the differential is given by Dba Tϵ = δ ba − ϵ[Vθ , T b ]a + ϵ∇aZτ Vθ (Zτ )b + O(ϵ2 ) Plugging this result into the chain rule for the adjoint state equation 49, we obtain X Aaτ = Aaτ+ϵ − ϵ[Vθ , Aτ +ϵ ]a + ϵ Abτ +ϵ ∇aZτ Vθ (Zτ )b + O(ϵ2 ) .
(55) (56)
b
We now rearrange this to isolate the finite time difference on the right-hand-side X a a a 1 Abτ +ϵ ∇aZτ Vθ (Zτ )b + O(ϵ) . ϵ Aτ +ϵ − Aτ = [Vθ , Aτ +ϵ ] −
(57)
b
Taking the limit ϵ → 0 gives X d a Aτ = [Vθ , Aτ ]a − Abτ ∇aZτ Vθ (Zτ )b . dτ
(58)
b
Multiplying by the generator T a and summing over the index a, we then obtain: X d Aτ = [Vθ , Aτ ] − Abτ ∇Zτ Vθ (Zτ )b , dτ
(59)
b
which is the claimed time-evolution equation of the theorem.
D
Lightning Review of Differential Geometry and Lie Groups
Right multiplication: Lie groups are smooth manifolds endowed by an additional group multiplication. There is therefore a natural diffeomorphism given by right multiplication Rg : G → G , Rg : h 7→ hg , (60) for h, g ∈ G. Let us denote the unit element of the group by e. Since right multiplication Rg is a diffeomorphism, its differential (dRg )e : Te G → Tg G is a linear isomorphism (linear bijective map). We can therefore uniquely identify tangent vectors v ∈ Tg G and Lie algebra elements ṽ ∈ g ≃ Te G. Matrix groups and right multiplication: For matrix groups, the differential (dRg )e : Te G → Tg G takes a particularly simple form (dRg )e v = vg (61) where we take the standard matrix product of the Lie algebra element v and the group element g on the right hand side. This statement can be easily checked by noticing that the differential acts on any function f : G → R by d [(dRg )e v]f = f (γv (t)g ) , (62) dt | {z } t=0 ≡γ̃(t)
d where γv denotes the curve associated with v ∈ Te G, i.e. γv (0) = e and dt γv (t)|t=0 = v. As a result, the curve γ̃ obeys d d γ̃(0) = g , γ̃(t)|t=0 = γ(t)|t=0 g = vg , (63) dt dt and is therefore a curve associated with vg ∈ Tg G as claimed.
18
Killing form:
we first define the adjoint map adu (v) = [u, v] (64) where u, v ∈ g. The adjoint map is linear and a Lie algebra homomorphism ad[u,v] = [adu , adv ] . (65) Then the Killing form is given by B(u, v) = Tr(adu ◦ adv ) (66) for u, v ∈ g is a bilinear symmetric form that is non-degenerate if the group is semi-simple and compact. We can use the Killing form to define a inner product for Te G by ⟨u, v⟩e = −B(u, v) . (67) The negative sign ensures positive definiteness. Riemannian metric: Often it is more convenient to work with the Lie algebra instead of the various tangent spaces. For example, an inner product ⟨·, ·⟩e on the Lie algebra induces a Riemannian inner product on the group G by ⟨v, w⟩g ≡ ⟨(dRg−1 )g v, (dRg−1 )g w⟩e = ⟨vg −1 , wg −1 ⟩e , for v, w ∈ Tg G , (68) where we have used that dRg−1 = (dRg )−1 to pull tangent vectors back to the Lie algebra. Using (dRg )−1 = dRg−1 , it immediately follows that the inner product is right-invariant, i.e., ⟨(dRg )e ξ, (dRg )e η⟩g = ⟨ξ, η⟩e (69) a for Lie algebra elements ξ, η ∈ Te G. For SO(3), we choose antisymmetric generators T that are orthonormal with respect to the inner product induced by the Killing form ⟨T a , T b ⟩ = 12 tr(T a )⊤ T b = δ ab . Riemannian gradient: Recall that the Riemannian gradient gradg f ∈ Tg G is the the unique element of the tangent space which obeys ⟨gradg f, v⟩g = dfg v (70) for all v ∈ Tg G. Due to the isomorphism between the tangent space Tg G and the Lie algebra g ≃ Te G, the Riemannian gradient gradg f ∈ Tg G uniquely corresponds to the Lie algebra element ∇f ≡ (dRg−1 )e gradg f ∈ g . (71) It is often more convenient to work with this Lie algebra representative of the Riemannian gradient. We will now derive a simple expression for this Lie algebra gradient. By definition of the Riemannian gradient, it holds that ⟨gradg f, (dRg )e v⟩g = dfg ((dRg )e v) , (72) for any v ∈ g. We can invert equation 71 to obtain gradg f = (dRg )e ∇f (73) Using this, we can rewrite equation 72 as follows ⟨(dRg )e ∇f, (dRg )e v⟩g = ⟨∇f, v⟩e , (74) where we have used the right-invariance of the metric equation 69. We thus conclude that ⟨∇f, v⟩e = dfg ((dRg )e v) (75) By definition of the right hand side, it therefore holds for a matrix group that d ⟨∇f, v⟩e = dfg (vg) = f (exp(tv)g) . (76) dt t=0 P a a We can expand the Lie algebra representative in terms of generators ∇f = a T ∇ f whose components are given by d ∇a f = ⟨∇f, T a ⟩e = f (exp(tT a )g) , (77) dt t=0 where we have used equation 76 and assumed that the generators are chosen to be orthonormal, i.e., ⟨T a , T b ⟩e = δ ab . For brevity, we typically refer to the Lie algebra representative of the Riemannian gradient ∇f as simply the Riemannian gradient in the main part of the paper. 19
Exponential map:
For a Lie group, the exponential map takes the form expg (v) = exp((dRg−1 )g v) g
(78)
for g ∈ G and v ∈ Tg G where exp : g → G is the matrix exponential for matrix Lie groups. We note that, by the Lie algebra expression of the Riemannian gradient equation 73, this implies that gradient descent on a Lie group with learning rate λ ∈ R and loss L : G → R is given by g i+1 = exp(−λ∇L)g i .
(79)
This result is used extensively in the main part of the paper. Differential: Right multiplication Rg : G → G induces a isomorphism between the tangent spaces Tg G and the Lie algebra g. Consider the differential dfg : Tg G → Tf (g) G of the map f : G → G. Using the right multiplicative isomorphism equation 61, we can write any tangent space element in terms of a Lie algebra element Tg G ∋ ΩG ↔ ω ∈ g
(80)
Specifically, we can define a Lie algebra representative of the differential Dfg ≡ d(Rf (g)−1 ◦ fg ◦ Rg )e : Te G ≃ g → Te G ≃ g
(81)
Using the fact that the right isomorphism amounts to right multiplication equation 61, we can derive a explicit expression for this representative Dfg (ω) = dfg (ωg)f (g)⊤ d = f (exp(tω)g) t=0 f (g)⊤ dt df = (ωg)f (g)⊤ dg X df = ωa (T a g)f (g)⊤ dg a X d = ωa f (exp(tT a g) t=0 f (g)⊤ dt a X = ωa ∇a f (g) f (g)⊤ .
(82) (83) (84) (85) (86) (87)
a
P Since Dfg (ω) ∈ g, we can expand it as Dfg (ω) = b T b Db fg (ω). Hence, it holds that X X Db fg (ω) = ⟨T b , ∇a f (g) f (g)⊤ ⟩ ωa ≡ Dfgba ωa , a
(88)
a
with Dfgba ≡ ⟨T b , ∇a f (g) f (g)⊤ ⟩ .
(89)
We can think of Dfgba as being the matrix representation of the differential in the generator basis T a . Chain rule: Let’s also consider a function F : G → R and its composition F ◦ f : G → R. Then it holds that X ∇a (F ◦ f )g = ∇b Ff (g) Dfgba . (90) b
This follows directly from F ◦ f ◦ Rg = F ◦ Rg ◦ Rg−1 ◦ f ◦ Rg .
(91)
Taking the differential of this and using the standard chain rule of the differential then gives d(F ◦ Rf (g) ◦ Rf (g)−1 ◦ f ◦ Rg )e = d(F ◦ Rf (g) )e d(Rf (g)−1 ◦ f ◦ Rg )e = ∇Ff (g) Dfg . (92) 20
E
SO(3) Conventions
Any group element g ∈ SO(3) can be written as g = exp(A)
(93)
with A ∈ so(3) taking value in the Lie algebra so(3) = {A⊺ = −A | A ∈ R3,3 } .
(94) a
A Lie algebra is, in particular, a vector space. We will choose the basis T with a ∈ {1, 2, 3} with # " # " # " 0 −1 0 0 0 1 0 0 0 3 2 1 (95) T = 1 0 0 . T = 0 0 0 , T = 0 0 −1 , 0 0 0 −1 0 0 0 1 0 In the context of Lie theory, the basis vectors T a are also referred to as the generators of the Lie algebra. In particular, any antisymmetric three-by-three matrix can be written as a linear combination of the generators. Furthermore, the generators obey the following commutation relations [T a , T b ] =
3 X
ϵabc T c ,
(96)
c=1
where ϵabc denotes Levi-Civita symbol. It can be checked that Tr(T a ) = 0 , a
b
(97)
Tr(T T ) = −2δ
ab
,
(98)
where δ ab is the Kronecker symbol. Using equation 98, we can equip the Lie algebra with an inner product ⟨A, B⟩ =
3 X 1 Aa B a , Tr(A⊺ B) = 2 a=1
(99)
where we have used that we can express an arbitrary element A of the Lie algebra so(3) as A = P3 a a a a=1 A T . Note that our basis T is orthonormal with respect to this inner product.
F
Limitations of Guidance
Guidance is a technique to bias the generation process, i.e., instead of sampling from the model distribution p, one aims to sample from ptilted (x) = p(x)
1 Zguide
e−βEguide (x)
(100)
where Eguide denotes the guidance energy of interest. For score based models, this amounts to a modification of the score ∇x log ptilted (x) = ∇x log p(x) − β ∇x Eguide (x) .
(101)
A challenging aspect of guidance is diffusion involves a one-parameter family of density which is typically parameterized in terms of diffusion time t or noise level σ(t). For many guidance energies, such as the Rosetta energy function, it is highly non-trivial to calculate the score of the guidance density at time t, Z 1 pguide (xt ) = p(x0 |xt )e−βEguide (x0 ) dx0 , (102) Zguide as it involves an untractable marginalization over the unnoisy datasample x0 . Here, p(x0 |xt ) denotes the backward kernel, i.e. the density of obtaining the unnoisy data sample x0 when integrating the reverse SDE starting from xt at diffusion time t. The fundamental challenge of all guidance schemes is to approximate the above marginalization in a suitable manner for the application at hand. This approximation becomes particularly challenging in the limit of β → ∞ which corresponds to 21
sampling minimizers of the guidance energy Eguide . In this limit, the integrand is heavily dominated by the minimizers (and its neighborhoods for very large but finite β as used in numerical experiments). As such, one has to estimate the conditional density p(x∗ |xt ) for the minimizers x∗ with high precision. For this reason, guidance cannot be expected to work well for minimizing generic guidance energies. Let us consider the simpler case in which we do not want to minimize the guidance energy, i.e., we do not consider the challenging β → ∞ limit. In this case, there are various strategies to approximate the marginalization. None of them are perfect, and each comes with a certain set of tradeoffs. We will only focus on methods which work for a pretrained unconditional model as this is the relevant setting for diffeomorphic optimization. In principle, one could use Monte-Carlo to estimate the guided score N
∇xt log pguide (xt ) ≈ ∇xt log
1 X −βEguide (x(i) 0 (xt )) e N i=1
(103)
(i)
with x0 ∼ p(x0 |xt ). For this, we have to repeatedly integrate the reverse SDE starting from xt and backpropagate through it. There are methods for refined backpropagation through SDEs, such as suitable generalizations of the adjoint state methodKidger et al. (2021); Li et al. (2020) but we are not aware of any work in the protein space applying these techniques. This is possibly due to the considerable technical difficulty, such as the appearance of Stratonovich stochastic integrals, as well as scaling concerns to larger problem sizes. Furthermore, one would need to derive generalizations of these methods for matrix Lie groups. Universal guidance Bansal et al. (2023) is a technique inspired by classifier guidance Dhariwal & Nichol (2021). In this setting the marginalization is approximated by a zeroth-order saddle point approximation around the expectation value x̂0 (xt ) = E[x0 |xt ] of the kernel p(x0 |xt ), i.e., Z 1 1 −βEguide (x̂0 (xt )) pguide (xt ) ≈ δ(x̂0 (xt ) − x0 ) e−βEguide (x0 ) dx0 = e , (104) Zguide Zguide which can be a good approximation at low noise levels (small diffusion time) but tends to be very poor at high noise. To alleviate this problem, guidance is often combined with a heuristic schedule which switches the guidance score on when the diffusion time is lower than some threshold (often in an adiabatic manner). However, this results in a fundamental tension: we want the guidance signal to kick in as early as possible in the generation process to meaningfully change the selected mode of the sample. On the other hand, we do not want to start the guidance too early as the zeroth-order saddle point approximation would be very poor. Whether such a delicate balance can be struck often heavily depends on the problem setting and requires careful (and costly) hyperparameter tuning. As a general rule of thumb, energy functions that depend only on the coarse grain details of the sample x0 seem to be easier to deal with. Examples include the classification of image types or the sentiment of a sentence. However, energy functions in physics (which are highly relevant in the protein space) are typically not of this type. For example, the Rosetta energy is extremely sensitive to the repulsive component of the atomic interactions resulting in numerically destabilizing energy values even if the global fold and secondary structure of the predicted unnoisy protein is roughly correct. Recently, there has been considerable interest in combining diffusion models with Sequential Monte Carlo (SMC) to facilitate sampling from the tilted distribution ptilted . We base our discussion on the very recent Singhal et al. (2025) which introduces Feynman-Kac steering unifying a number of previous approaches in a common theoretical framework. The basic idea is to modify the transition kernel p(xt−1 |xt ) of the reverse diffusion p(xt |xt+1 ) → p(xt |xt+1 )G(x1 , . . . , xt )
(105)
by a guidance potential G which fullfills the consistency condition 0 Y t=1
G(x1 , . . . , xt ) =
1 −βEguide (x0 ) e . Zguide
(106)
This ensures that the joint distribution of the sampling trajectory (x1 , . . . , xt , . . . x0 ) is given by pguide (x1 , . . . , x0 ) = p(x1 , . . . , x0 )
22
1 Zguide
e−βEguide (x0 )
(107)
and thus has the desired marginal pguide (x0 ). In principle, any guidance potential satisfying the consistency condition will lead to asymptotic gurantees but in practice a careful choice of potential is vital to avoid prohibitive variance. This is because, at each step of the reverse diffusion one uses importance sampling to reweight the sample xt according to the tilted transition kernel equation 105. In practice, one often has to choose guidance potentials G that involve challenging marginalization. For example, a widely used guidance potential is the expected guidance energy Eguide . Thus this approach faces analogous challenges to universal guidance, particularly in the β → ∞ limit. Another important point of differentiation between guidance and diffeomorphic optimization is that the latter is particularly suited in applications in which we want to improve a particular sample x0 with respect to some cost function. For guidance, we would need to incorporate a constraint in the guidance potential that ensures similarity to x0 . For diffeomorphic optimization, we perform gradient descent updates that iteratively refine the sample.
G
Complexity and Efficiency of Gradient Estimations through Flows
A forward pass evaluates the network layer by layer, applying linear transformations and nonlinearities to propagate an input through the model. Its cost is determined by the size and number of layers, with matrix multiplications dominating the total FLOPs. A general backward pass is more expensive because it must compute gradients with respect to all parameters. This involves propagating gradients backward through each layer and forming gradient tensors for every weight matrix in each layer and the layers input for further backward propagation. These extra operations roughly double the compute relative to the forward pass, and they also require storing all intermediate activations, which increases memory consumption and can become a limiting factor on modern hardware. A backward pass that computes gradients only with respect to the input is cheaper. It is comparable to a single backward pass as the parameter gradients do not have to be computed. The computational effort becomes close to that of the forward pass, with only a small overhead for the backward sweep. Gradient checkpointing reduces memory consumption during backpropagation by storing only a subset of intermediate activations from the forward pass and recomputing the missing ones when needed. In standard backpropagation, all activations must be kept in memory because each layer’s gradient depends on its forward output; this creates a memory footprint that grows linearly with depth and often becomes the bottleneck in training large models. Gradient checkpointing breaks the computation graph into segments, saves only the boundary activations, and discards the rest. During the backward pass, the discarded activations are recomputed on the fly by running partial forward passes within each segment. This trades additional compute for a substantial reduction in memory usage. The memory cost can drop from linear to roughly the square root of the number of layers in optimal schemes, enabling deeper networks or larger batch sizes on the same hardware. The extra compute cost is typically a factor of 1.5–2×, but the trade-off is advantageous when memory is the limiting resource. Modern implementations generalize this idea through customizable checkpoint policies, selective recomputation, and integration with automatic differentiation frameworks, making gradient checkpointing a standard technique in large-scale training, especially for transformer architectures and diffusion models where activation tensors dominate memory load. ODE solvers compute trajectories by applying numerical integration steps, with cost driven by repeated evaluations of the vector field. Higher-order or adaptive methods improve accuracy but increase per-step compute, so total complexity depends on both the solver order and the number of required steps. In neural ODEs, most of the expense comes from evaluating the neural network that defines the dynamics. Using the autograd engine to compute gradients through a ODE solution keeps the entire computational graph in memory, resulting in O(N L) scaling with a second pass over the graph in reverse direction required for the backward pass with a total cost of O(N L). The adjoint method recomputes the state trajectoy and backpropagates the gradients by integrating a separate reverse-time ODE, avoiding storage of intermediate states and keeping memory use low. Compared to gradient checkpointing it does not require storing any values besides the vectorfield as it recomputes both the trajectory and the gradients during the backward pass. This leads to a memory complexity of O(L) as only the network parameters and their activations need to kept in memory at any time. This comes at the price of additional solver calls and potential numerical instability or even divergence from the trajetory of the forward pass, since the backward integration can amplify errors and often requires tighter tolerances. After O(N L) steps of compute in the forward pass, the 23
adjoint method recomputes the forward pass to reconstruct the trajectory as well as one backward pass to compute the gradient, resulting in an additional O(2N L) compute complexity during the backward pass. As a result, adjoint-based gradients are memory-efficient O(L) but computationally more expensive with O(3N L). Solving a reverse time ODE comes with potential numerical issues as the recomputed trajectories diverge and subsequently lead to diverging gradients. Checkpointing provides a middle ground by storing selected solver states and recomputing segments during backpropagation. This reduces memory usage to S = N/K at any given time where S is the segment length. Storing K checkpoints in the form of a state of the ODE is negligible as it amounts to K samples. For K checkpoints, we only have to solve S steps at any one time, while the recomputation between the K checkpoints amounts to going forward and backward on the segments compute complexity of O(2KSL) = O(2N L). It balances compute and memory more predictably, making it a practical choice for many neural ODE applications. To quantitatively measure the performance of both the gradient checkpointing and adjoint method, we evaluated each gradient against the ground truth gradient obtained with numerical differentiation. Figure 10 shows the gradient error over increasing modelling dimensions as a funtion of the step size and the checkpointing interval. For the Euclidean part, we integrated the function dxt = cos(t/2π)·xt P3 and for the SO(3) part, we integrated dRt = cos(t/2π) a va T a θ(v) with v = [1.0, −0.2, 0.1]T and θ(v) = ∥v∥. The initial condition was sampled from a Gaussian distribution in Euclidean space and the respective axis-angle space for the SO(3) elements. The dimensionality was determined by the number of SE(3) ODE’s integrated in parallel for t ∈ [0, 1]. We used this toy ODE’s to establish the numerical behavior of the autograd checkpointing and the adjoint methods for backpropagating gradients through ODE solvers. The analysis suggests that the adjoint method with a checkpointing at every step yields gradients very close to its autograd gradient. Decreasing the checkpointing frequency requires recomputing longer parts of the trajectory which accumulates errors. Overall a trend emerges that decreasing the step size is beneficial for both adjoint and autograd gradients, and particularly for very small stepsizes, the difference reduces significantely.
Figure 8: Gradient Error as a function of dimensionality and the number of steps We also extended this analysis to the FrameFlow model Yim et al. (2023a) model but due to the computational cost of numerical differentiation for the ground truth gradient refrained from evaluating multiple checkpointing frequencies. For this reason we kept the checkpointing frequency fixed at every five steps. One can observe that the adjoint gradients backpropagated through the SE(3) ODE generally more imprecise than their autograd counterparts. For commonly used stepsizes of 1/25 to 1/100, the total deviation can be quite significant, but correcting for the per dimension gradient deviation, the gradient error is small. 24
Figure 9: Gradient Error as a function of dimensionality and the number of steps Finally, we compare the gradient error of the AlphaFlow architecture Jing et al. (2024). Again, adjoint gradient estimators perform worse than their autograd counterparts. It should be noted that AlphaFlow is a flow defined on the Cβ carbon atoms and does not model a rotation group, compared to the previous two models. We want to note that the evaluation of these gradients were on the order of hours, particularly for the highest step sizes.
Figure 10: Gradient Error as a function of dimensionality and the number of steps
H
Experiments
H.1
SO(3) Manifolds
We generate a manifold in the space of SO(3) matrices. To plot the rotation matrices easily, we parameterize the SO(3) elements in their axis-angle vectors. The manifold in the axis-angle representation was generated with the S-curve function of sklearn, sklearn.datasets.make_s_curve(n_samples=250000, noise=0.1). 25
The flow matching model was parameterized with a four layer MLP with SELU activation functions and 64 hidden units in each of the hidden layers. For training, the model was optimized with Adam using a learning rate of 1e−3 over 20 epochs on the training data with conditional flow matching to enforce that the predicted flow on the manifold of SO(3) matrices closely matches the target logarithmic map on the manifold. During sampling, the vector field was integrated with the Crouch-Groussmann integrator with a step size of dt = 1/200 from 0 to 1. For the sample optimization, a target sample z ∗ was experimentally determined and a second random sample was drawn as the initial value x(0) = g(z (0) ) of the optimization problem. The mean squared error (MSE) between the target sample x∗ and the sample x(0) being optimized served as the loss function L. The sample z (i) was then optimized by calculating the gradient dL(g(z (i)) )/dz (i) and performing 20 steps of gradient descent in the SO(3) base space with a base space learning rate of 0.1. The tangent space was calculated by evaluating the Jacobian matrix Jg with torch.nn.functional.jacobian at point z and its corresponding target space point x = g(z). Subsequently, the two largest eigenvalues with the corresponding left columns were computed with a SVD and plotted with a scaling of 1.5 the point x. H.2
Secondary Structure Optimization
The secondary structure experiments were conducted on top of the FrameFlow codebase Yim et al. (2023a). For training and architectural details, we refer to the manuscript. We used the model parameters inn the provided model checkpoints by the original authors. The samples x with the time index t ∈ [0, 1] consist of backbone elements of each residue centered on the Ca carbon atom. Each of the L backbone elements is parameterized by its translation and rotation and are elements of the SE(3) group. The flow generates an element (t, R) of the SE(3) group for each residue. Upon completion of the sampling process, these transformations are applied to the idealized backbone coordinates. For the experiment described in the main part of the paper, the loss function L was calculated as a function of the distance between the atom positions in the first and the last residue. The distance loss was implemented by taking the ReLU function over the difference between the desirable target 2 (i) (i) distance and the actual distance, i.e. L = max(0, d∗ − d(x0 , xL )) where the target distance (i)
(i)
d∗ between the atom positions in the two termini x0 and xL in the backbone structure of length L ∗ at optimization step (i) of the structure qPwas chosen as d = L/2. The distance metric was chosen as 2 the mean squared error d(x, y) = j (xj − yj ) over the atoms j in two residues x and y. The order of subtraction max(0, d∗ − d) ensures a loss that becomes zero if d ≥ d∗ and that the two termini are repelled from each other. We also experimented with an attraction loss in the sense of max(0, d − d∗ ) which worked equaly well. In fact, we were able to dynamically switch losses during optimization and the termini were pulled apart and together smoothly as determined by the loss with all intermediate optimization states x(i) staying on the data manifold. When switching between attracting and repelling loss function it was beneficial to reset the momenta of the adaptive optimizers like Adam to accelerate the optimization. For sampling we integrated with the Euler integration method with a step size of dt = 1/25. The learning rate for the optimization was chosen at a constant 0.5 for 300 optimization steps. For the optimization of the dihedral angles according to the ABEGO scheme, we utilized a differentiable implementation of Rosetta energy function in PyTorch, i.e. tmol Leaver-Fay et al. (2025). We use a specific component of the full energy function that takes the dihedral angles of neighboring backbone elements as an input. We modified this component such that it increases the contribution for dihedral angles corresponding to β-sheets. In more detail, the combined Ramachandran/p_aa_pp energy term in tmol interpolates the energies (the negative-log probabilities) stored in a toroidal-grid lookup-table over phi and psi dihedrals. To push the generated samples away from the alpha-helical region of conformation space, we created our own custom lookup table where we dramatically increased the energies for an elliptical region of the Ramachandran map that covers the alpha-helical region. The energies at the edge of this ellipse are not uniform, so our first crude assignment of energies inside this ellipse that sloped away from 26
a central peak to a particular value at its edge frustrated the minimizer, trapping conformations in the weird minima at the edge of the ellipse. In our second attempt, we instead labeled each grid cell inside the ellipse by the closest cell outside of the ellipse and then interpolated the energies along the distance from the center of the ellipse (to which we assigned an energy of +20) to the ellipse’s edge. This led to the desired behaviour and we therefore chose this approach. The resulting grid is shown in Figure 11.
Figure 11: Grid used for Ramachandran guidance.
Figure 12: Self-consistent RMSD and the optimized ABEGO percentage over the course of diffeomorphic optimization. 27
Three important metrics in machine learning aided protein design are designability, diversity and novelty. We consider designability the most pertinent to our task as diversity and novelty are of lesser interest when trying to converge on a single optimal structure. The optimization towards a single structure makes diversity collapse as expected. Similarly, since our aim is not to generate diverse samples but rather optimize towards a singular structure, novelty is a misplaced metric in this case. Empirically, the designability of the optimization trajectories agree well with the designability of the initial values drawn from the generative model. Only for larger learning rates of 1.0 does the designability experience a significant drop while at the same time only marginally improving the percentage of dihedral angles classified as B in the ABEGO scheme. Since the FrameFlow model does not predict side chains, we interpreted all residues as alanine. We applied the sampled translations and rotations to the idealized backbone atom14 positions to obtain the sampled backbone geometry. The score function was minimized with gradient descent in the base space with a learning rate of 0.1 for 500 steps. We implemented several versions of guidance to obtain maximally competitive baseline results. Specifically, we implemented both a single step and and a multi step denoiser combined with the loss function described above. The guidance was started at t = 0.1. We selected the starting point of the guidance signal by hyperparameter line search over t = 0.8, 0.1, 0.25, 0.5, 0.75. We also found that it was benefical to stop sampling early at t = 0.98 to prevent distortive numerical pathologies that arose due to the divisor of the optimal transport transport path close to t = 1. The guidance vector field was obtained by evaluating the loss on the atom14 positions of each backbone residue and computing the gradient with respect to the corresponding element in the base space. For weighting the guidance term, we tested the theoretically motivated weighting function bt from Lipman et al. (2024) but obtained better results with a constant, time-independent weight of w = 25. This value was selected using a line-search hyperparameter sweep for w = 0.1, 1, 10, 25, 50, 100, 1000. For evaluating the guidance term on a sample xt , we tested both taking the direct gradient of the loss function, ∂L(x)|x=g(xt ) /∂x, and backpropagating back through the denoiser ∂L(g(xt ))/∂xt . We chose the latter as it yielded better results. Experimentally, we also found that only using the translation term (and not the rotational part) as guidance improved the performance. This may be due to the fact that on SO(3), there is no equivalence between score and denoising as in Euclidean space. We furthermore carefuly ablated the number of denoiser steps used in multistep denoising. This is because fair comparison to diffeomorphic optimization should allow for a generous budget of denoising steps as the method is numerically costly. As shown in Table 2, even significant increase in the denoising steps does not lead to a pronounced gain in performance. To demonstrate that the numerical cost of diffeomorphic gradient calculations can be reduced by coarse-grained backpropagation through the solver, we performed the following ablation study: we measured the performance of diffeomorphic optimization as a function of the number of steps used for sampling and calculating gradients during optimization. To this end, we ran diffeomorphic optimization with 10, 25, 50, 75 and 100 integration steps during optimization, and evaluated the optimized sample after optimization with 1000 integration steps. The results are shown in Figure 13. Compared to guidance, even a very small number of steps used for sampling during optimization yields good performance. For this experiment, we measured the performance over 10 seeds as compared to 50 seeds in the main part. H.3
Docking Optimization
For docking optimization we used the DiffDock code base. This diffusion model samples from the manifold of permissible rotational degrees of freedom of a small molecule on SO(2) and its global translation and rotation on SE(3). In order to make the sampling process differentiable, we implemented the probability flow formulation of the diffusion models which provided a deterministic map between base space and target space. For a differentiable loss, we considered two different metrics: the confidence head of the diffusion model and the OpenDock scoring function Hu et al. (2024). Opendock is a physics-inspired scoring function and implements the Vina Score in a differentiable manner in PyTorch. One of the technical challenges of incorporating OpenDock as a suitable loss function were posed by its reliance on PDBQT file formats. In addition to the positional information stored in the PDB file format, the PDBQT file format additionally stores the partial charges, torsional bonds, and autodock atom types. To work around this constraint, we first saved the 28
Figure 13: The performance of diffeomorphic optimization (higher is better) as a function of the number of ODE integration steps taken for each evaluation during optimization. After optimization, we evaluated all optimized samples with a 1000 steps. Diffeomorphic optimization is robust even when using just 10 integration steps during optimization. Secondary Structure A B Unguided Guidance
Diffeomorphic
53.5 ± 23.3
32.8 ± 20.1
ODE Steps
Denoiser Steps
200
1 50 100 200
63.3 ± 25.1 63.5 ± 24.8 63.8 ± 24.5 64.1 ± 24.1
27.7 ± 23.3 27.4 ± 23.1 27.1 ± 22.9 26.8 ± 22.6
500
1 50 100 200
64.2 ± 24.4 64.4 ± 24.0 64.7 ± 23.7 65.0 ± 23.4
26.6 ± 22.4 26.3 ± 22.1 26.0 ± 21.9 25.8 ± 21.6
1000
1 50 100 200
65.0 ± 23.6 65.1 ± 23.4 65.2 ± 23.2 65.3 ± 23.1
25.5 ± 21.5 25.3 ± 21.3 25.1 ± 21.0 24.9 ± 20.8
ODE Steps
Optim Steps
25
500
91.3 ± 8.4
7.4 ± 7.3
Table 2: The addition of more compute for guidance leads to diminishing returns. Increasing the the number of integration steps both in the flow sampling as well as in the denoiser does not automatically lead to a higher target metric as measured with the percentage of A classified dihedral angles. All guidance terms were weighted with a constant w = 25. Diffeomorphic optimization samples with 25 integration steps and applies 500 gradient descent steps in the base space. Diffeomorphic optimization yields higher A classification scores with a substantially lower standard deviation in performance.
ligand conformations generated by DiffDock to disk in the SDF file format, converted it to the PDBQT and reloaded it with the additionally generated partial charges created during the conversion. At this 29
point the computational graph connecting the base space variables to inputs of the differentiable loss scoring function was suspended due to the reading and writing to disk. To reconnect the computational graph with the input of the Opendock scoring function, we substituted the atom positions in the loaded PDBQT data structure with the atom positions originating from the DiffDock model. The confidence head of the diffusion model is a neural network and was trained to predict whether the generated conformation is within 2 Angstroms of the ground truth. We additionally include on the scoring of OpenDock into our experimental setup which is a differentiable implementation of Vina. Following the original authors recommendation on the confidence prediction, we screened for the entries in the publicly available PDBBind test data set that were assigned a confidence logit of larger than 0 by the DiffDock confidence head. This resulted on a subset of roughly 78 PDBBind ligand-protein pairs on which the results were reported. We perform 10 optimization steps with an Adam-style optimizer with a learning rate of 0.1. Experimentally, we observe that the gradients of the rotation and torsional degrees of freedom play a dominant role in the norm of the total gradient. To counteract this, we introduced separate scalings for the learning rates of the rotation and torsions and multiply the learning rates for these two sets of degrees of freedom by an additional 0.1. To ensure fair comparison, we draw more samples for the sampling baseline in order to match the computational budget used by diffeomorphic optimization. Since both checkpointing and the adjoint method recompute the activations during the backward pass through the ODE solver, we used a budget corrector of 3× a single sampling step. This budget corrector number could be reduced more with coarser checkpointing and fewer steps, but we observed favorable experimental results even for this setting and thus did not explore this direction further. H.4
Tertiary Structure Optimization
We built on the alphaflow codebaseJing et al. (2024) and chose the ESMFold-based architecture as it does not require costly MSA. We ensured that diffeomorphic optimization works both for distilled and non-distilled case but large-scale testset-wide experiments were conducted with the latter as it leads to significant speed-up. The model ingests the pairwise distances between random atom positions to predict the backbone and the side chain atom positions conditioned on the sequence information of the protein. As a differentiable loss function, we used the tmol implementation of of beta_nov2016_cart Rosetta score function which support pytorch autograd to compute gradients. As the architecture contains non-differentiable distograms, we implemented a differentiable straight-through version thereof harnessing the sigmoid function σ, dsof t = σ(β · (d − lower)) · (1 − σ(β · (d − upper)) dgram = detach_grad(d) + (dsof t − detach_grad(dsof t )) where lower and upper denote the lower and upper bound of the discrete bin in which the distance d should have been placed. While the forward pass is unaffected by the soft distogram implementation, the gradients are affected by the smooth approxiation. Higher values of the β temperature result in a steeper sigmoid function at the cost of a stiffer gradient surface. As a baseline we used the state-of-the-art Rosetta Relax protocol Alford et al. (2017) without any diffeomorphic optimization. For random sidechain repacking, Rosetta Relax uses n seeds which are commonly referred to as structures in the Rosetta community. This is followed by gradient descent in which the repulsive terms in the energy function are adiabatically increased with a linear schedule. Packing and gradient-based minimization is then repeated k times. In Figure 15 we evaluated the Rosetta Relax baseline with an increasing amount of compute measured as the product of the number of seeds 1, 3, 5 with the number relaxation cycles 3, 5. We were not able to discern an improvement compared to diffeomorphic optimization in conjunction with Rosetta Relax. For diffeomorphic optimization, we used Adam with a learning rate of 0.1 and otherwise standard pytorch hyperparameters. This was then followed by standard Rosetta relax to reduce numerical cost. The performance of this protocol as a function of diffeomorphic optimization steps is shown in Figure 16. In addition to the experiments reported in the main part of the paper, we also compared the performance of our diffeomorphic Rosetta Relax to a sampling-based baseline. For this baseline, we simply sample from AlphaFlow with the equivalent computational budget used by diffeomorphic 30
optimization. We select the sample with the lowest energy which is then minimized by the same standard Rosetta relax as used in the last step of the diffeomorphic protocol. As shown in Figure ??, diffeomorphic optimization again outperforms this baseline.
Figure 14: Diffeomorphic optimization of the Rosetta energy function: each point denotes a prediction of ESMFlow for an element of the pdb test set of Alphaflow for which we measure the improvement Ediffeo − Ebaseline of diffeomorphic Rosetta Relax over standard Rosetta relax. Left: Diffeomorphic Rosetta Relax outperforms Rosetta Relax and improves the Rosetta score with the default Rosetta Relax values of three structures and three repeated relaxations. Right: This improvement holds across an increasing computational budget as measured in terms of the total number of relaxations.
Figure 15: Diffeomorphic optimization in combination with Rosetta Relax yields improved (lower) Rosetta energies when directly compared to Rosetta Relax. This statement holds even with increasing computational budget for the baseline. In this experiment, we varied the number of Rosetta Relax random seeds (referred to as structures in the Rosetta community) and the number of relaxations for each random seed. The product of both Rosetta Relax random seeds and relaxations is plotted on the horizontal axis. We can observe no substantial benefits of providing Rosetta Relax with additional compute budget compared to diffeomorphic optimization which consistently improves the final Rosetta energy.
31
Figure 16: Diffeomorphic Optimization in combination with Rosetta Relax can achieve substantial energy reductions in a small number of steps. This is exemplified by the high energies at the beginning of the optimization curves which corresponds to running Rosetta Relax without a diffeomorphic optimization step which are subsequently is minimized by several orders of magnitudes while remaining stable.
32