Diffeomorphic Optimization

Generative models learn data distributions that reside on a low-dimensional manifold within a higher-dimensional ambient space. Optimizing differentiable objectives on such a 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. Our approach 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 instead. Using concepts from differential geometry, we show this is equivalent to Riemannian gradient descent on the data manifold, keeping trajectories on-manifold by construction and yielding a smoother optimization surface. We extend diffeomorphic optimization to matrix groups, such as SO(3), SE(3) and SU(N) with autograd-compatible Lie-group ODE solvers, which allows us to empirically demonstrate the effectiveness of our approach in the highly relevant task of protein design and lattice gauge theory.