Diffeomorphic Optimization

2026-07-01

AI Summary

This work addresses differentiable optimization on low-dimensional data manifolds embedded in high-dimensional spaces, where conventional gradient descent often deviates from the manifold and struggles with non-convex loss landscapes. The authors propose a novel approach that leverages diffusion and flow models to construct a diffeomorphic mapping, pulling the manifold back to a simple base space for optimization. Using tools from differential geometry, they prove this procedure is equivalent to Riemannian gradient descent, inherently preserving trajectories on the manifold. This is the first integration of diffeomorphic mappings with Riemannian optimization, extended to the Lie groups SO(3) and SE(3), yielding an automatic differentiation–compatible SO(3) gradient and a generalized adjoint-state backpropagation for Lie group ODE solvers. In protein design tasks, FrameFlow achieves a 91.3% secondary structure targeting success rate (versus 63.3% baseline), doubles the peptide binding affinity optimization speed compared to OC-Flow, and significantly reduces Rosetta energy by thousands of units.

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 high-dimensional, 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 $\mathcal{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 $\mathrm{SO}(3)$ and $\mathrm{SE}(3)$, deriving an autograd-compatible $\mathrm{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\times$ the speed, and reduces Rosetta energies by thousands of units across the PDB test set for structures with hundreds of residues.

Problem

• manifold optimization
• generative models
• gradient descent
• diffeomorphic mapping
• Lie groups

Innovation

• Diffeomorphic Optimization
• Riemannian Gradient Descent
• Lie Groups
• Protein Design
• Generative Models