Landing with the Score: Riemannian Optimization through Denoising

TL;DR

This work introduces a link function that connects the data distribution to the geometric operations needed for optimization, and shows that this function enables the recovery of essential manifold operations, such as retraction and Riemannian gradient computation.

Abstract

Under the data manifold hypothesis, high-dimensional data are concentrated near a low-dimensional manifold. We study the problem of Riemannian optimization over such manifolds when they are given only implicitly through the data distribution, and the standard manifold operations required by classical algorithms are unavailable. This formulation captures a broad class of data-driven design problems that are central to modern generative AI. Our key idea is to introduce a link function that connects the data distribution to the geometric operations needed for optimization. We show that this function enables the recovery of essential manifold operations, such as retraction and Riemannian gradient computation. Moreover, we establish a direct connection between our construction and the score function in diffusion models of the data distribution. This connection allows us to leverage well-studied parameterizations, efficient training procedures, and even pretrained score networks from the diffusion model literature to perform optimization. Building on this foundation, we propose two efficient inference-time algorithms -- Denoising Landing Flow (DLF) and Denoising Riemannian Gradient Descent (DRGD) -- and provide theoretical guarantees for both feasibility (approximate manifold adherence) and optimality (small Riemannian gradient norm). Finally, we demonstrate the effectiveness of our approach on finite-horizon reference tracking tasks in data-driven control, highlighting its potential for practical generative and design applications.

Figures (5)

• Figure 1: Optimized trajectory (orange) on the system trajectory manifold for the unicycle car model that is desired to track the reference trajectory (black, thin). The closest tracking trajectory in the set of available manifold samples is given in (green, dotted). See Section \ref{['sec:numericalExamples']} for more details.
• Figure 2: Objective value vs. flow time $t$ for the orthogonal manifold for $n=10$ (left, different $\sigma > 0$) and for $n=20$ (right, $\sigma = 0.05$). Here "exact landing" refers to equation \ref{['eq:projectedFlow']} with exact operations $\texttt{v} = \pi$ and $\texttt{v}' = \pi'$.
• Figure 3: Denoising Riemannian gradient descent: Angle of the first pendulum (left) and unicycle car position (right) with the optimized output trajectory $\boldsymbol{y}^*$ (blue, dashed), the true system trajectory $\boldsymbol{y}^{\operatorname{true}}$ (orange), the initial trajectory $\boldsymbol{y}_0$ (green, dotted) and the reference trajectory $\boldsymbol{r}$ (red)
• Figure 4: Denoising Riemannian gradient descent: Objective value $f$ vs the iteration count $j$ for double pendulum (left) and unicycle car model (right). Note the logarithmic scale on the y-axis. The current cost (blue, dashed) is the objective value $f(\boldsymbol{u}_j,\boldsymbol{y}_j)$ at the current (in general infeasible $(\boldsymbol{u}_j,\boldsymbol{y}_j) \notin \mathcal{M}_{\operatorname{IO}}$) iterate, while the true cost (orange) is the value $f(\boldsymbol{u}_j,\boldsymbol{y}_j^{\operatorname{true}})$, with $\boldsymbol{y}_j^{\operatorname{true}}$ obtained by simulating (\ref{['eq:systemDynamics']}) with input $\boldsymbol{u}_j$.
• Figure 5: Denoising Riemannian gradient descent: Unicycle car position (right) with the optimized output trajectory $\boldsymbol{y}^*$ (blue, dashed), the true system trajectory $\boldsymbol{y}^{\operatorname{true}}$ (orange), the initial trajectory $\boldsymbol{y}_0$ (green, dotted) and the reference trajectory $\boldsymbol{r}$ (red)

Theorems & Definitions (42)

• Theorem 1: Main
• Corollary 2
• Theorem 3
• Remark 4
• Theorem 5
• Remark 6
• Remark 7
• Theorem 8: Theorem 1 in alain2014regularized
• Lemma 9: Lemma A.1 in divol2022measure
• Lemma 10