Diffeomorphic Optimization

Ludwig Winkler, Pan Kessel, Andrew Leaver-Fay, Joseph Kleinhenz

Keywords: Diffusion Models, Geometric Machine Learning, Differentiable Geometry

TL;DR: Optimization of samples with respect to a target space loss function while staying on the learned manifold.

Abstract: Optimization is a challenging task due to the rugged nature of the optimization landscape and the concentration of data on a low-dimensional manifold. Our approach starts from the observation that flow and diffusion models map the data manifold to a smooth and simple base space. We thus propose to reparameterize the optimization problem in terms of these simple base-space variables. Using concepts from differential geometry, we demonstrate that this reparameterization naturally constrains optimization to the data manifold and results in a smoother optimization surface. We extend diffeomorphic optimization to matrix groups, such as $SO(3)$ and $SE(3)$, which allows us to empirically demonstrate the effectiveness of our approach in the highly relevant task of protein design.

Primary Area: generative models

Submission Number: 2754