Diffeomorphic Optimization
2026-07-01 • Machine Learning
Machine Learning
AI summaryⓘ
The authors explain that generative models work with data that lives on a complicated, curved surface inside a higher-dimensional space, which makes optimization hard because simple methods can easily go off this surface. They introduce a method called diffeomorphic optimization that uses a special mapping to a simpler space, allowing gradient descent that respects the curved shape of the data. For protein design, they adapt this method to work with rotations and movements in 3D space and develop tools to efficiently compute gradients. Their approach improves protein structure targeting, speeds up peptide binding predictions, and lowers energy scores for large protein structures.
Generative modelsLow-dimensional manifoldGradient descentDiffeomorphic optimizationRiemannian gradient descentMatrix Lie groupsSO(3)SE(3)Protein designAdjoint-state method
Authors
Ludwig Winkler, Andrew Leaver-Fay, Joseph Kleinhenz, Pan Kessel
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.