Landing with the Score: Riemannian Optimization through Denoising

This paper proposes a novel framework for Riemannian optimization on implicitly defined data manifolds by leveraging the score function from diffusion models to recover essential geometric operations, resulting in two efficient algorithms (DLF and DRGD) with theoretical guarantees for feasibility and optimality in data-driven design and control applications.

The Gist

Imagine you are trying to find the best possible route for a delivery truck. You have a map, but it's not a flat piece of paper; it's a bumpy, winding mountain trail.

In the world of traditional math and AI, finding the best route usually requires you to know the exact shape of the mountain beforehand. You need a blueprint of every curve and slope to calculate the fastest path. This is called Riemannian Optimization.

But what if you don't have a blueprint? What if the "mountain" is just a collection of millions of photos of delivery trucks that successfully made it up the hill? You don't know the shape of the trail, but you know where the "good" spots are because you've seen them before.

This is the problem the paper "Landing with the Score" solves.

Here is the breakdown of their solution using simple analogies:

1. The Problem: The Invisible Mountain

In modern AI (like generating images or controlling robots), data often lives on a hidden "shape" or manifold.

• The Analogy: Imagine a giant, invisible sheet of fabric floating in a 3D room. All the "real" data (like realistic faces or valid robot movements) lies on this fabric. Everything else in the room is just "noise" or nonsense.
• The Challenge: You want to find the best point on this fabric (e.g., the most aerodynamic car shape), but you don't know where the fabric is. You only have a bag of marbles (data points) that happened to land on it. Traditional math tools can't work here because they need a blueprint of the fabric, which you don't have.

2. The Secret Weapon: The "Score" Function

The authors use a trick from Diffusion Models (the technology behind AI image generators like DALL-E or Midjourney).

• The Analogy: Imagine the fabric is covered in a thick, invisible fog. If you stand on the fabric, the fog is thin. If you step off, the fog gets thicker.
• The "Score": In AI, there is a tool called a Score Function. Think of it as a magnetic compass or a smell detector.
  • If you are in the fog (off the fabric), the compass points strongly toward the fabric.
  • If you are on the fabric, the compass stops pointing because you're already there.
  • Crucially, this compass is trained on your bag of marbles (the data). It learns the shape of the fabric just by smelling the fog.

3. The Magic Link: From "Smell" to "Map"

The paper's biggest breakthrough is proving that this "compass" (the score) isn't just pointing to the fabric; it actually reveals the geometry of the fabric.

• The Gradient (Direction): The direction the compass points tells you exactly how to project yourself back onto the fabric. It's like a "Landing Gear" that automatically lowers you onto the nearest point of the trail.
• The Hessian (Curvature): The way the compass wiggles tells you the slope of the trail. It tells you which way is "uphill" and which is "downhill" along the fabric itself.

In short: They proved that you can turn the "smell" of the data into a perfect map of the terrain, even without ever seeing the terrain directly.

4. The Two Algorithms: The Hikers

Once they have this "compass" and "map" derived from the data, they built two ways to find the best spot:

• DLF (Denoising Landing Flow): Imagine a hiker who is allowed to wander a little bit off the trail to look around, but has a strong magnetic pull (the "Landing" force) that constantly drags them back to the fabric. They slide down the hill, and the magnetic pull ensures they never fall off the edge.
• DRGD (Denoising Riemannian Gradient Descent): This is a hiker who takes small, careful steps. At every step, they use the compass to check: "Am I still on the fabric? If I'm slightly off, I'll correct my step immediately." This mimics the classic, perfect math method but uses the AI compass to do the heavy lifting.

5. Why This Matters (The Real World)

Why should you care?

• Design: Imagine designing a new airplane wing. Instead of testing thousands of random shapes, the AI knows the "shape of reality" (what wings actually work). It can slide along that invisible shape to find a wing that is better than anything in the training data.
• Robotics: A robot learning to walk doesn't need to be told the laws of physics. It just needs to know the "shape" of valid movements (from data) and can optimize its walk to be faster or more energy-efficient.
• Efficiency: You don't need to retrain the AI for every new problem. If you already have a pre-trained AI (like a general image generator), you can use its "compass" to solve specific optimization problems instantly.

The Bottom Line

The authors took a complex math problem (optimizing on a shape you can't see) and solved it by borrowing a tool from generative AI (the score function). They showed that AI doesn't just generate data; it understands the geometry of reality.

They turned a "black box" AI into a precise navigation tool, allowing us to find the absolute best solutions in complex, real-world problems without needing a manual or a blueprint.

Technical Summary

1. Problem Statement

The paper addresses Riemannian optimization over manifolds that are implicitly defined by data rather than explicitly known via mathematical formulas.

• Context: In many modern applications (e.g., generative AI, data-driven control, aerodynamic design), high-dimensional data concentrates near a low-dimensional manifold $\mathcal{M}$. The goal is to minimize an objective function $f(x)$ subject to $x \in \mathcal{M}$.
• The Challenge: Classical Riemannian optimization requires explicit geometric operations: projecting onto the manifold ($\pi(x)$), projecting onto the tangent space ($P_{T_x\mathcal{M}}$), retractions, and exponential maps. When $\mathcal{M}$ is only available as a set of samples from a distribution $\mu_{data}$, these operations are unavailable.
• The Gap: Existing methods either require explicit manifold definitions or rely on parametric latent space learning (e.g., VAEs, GANs), which introduces challenges like unknown manifold dimension, difficult atlas parameterization, and the need for invertible maps.

2. Core Methodology

The authors propose a framework that bridges diffusion models and Riemannian optimization by establishing a theoretical link between the score function of a smoothed data distribution and the geometric operators of the manifold.

A. Theoretical Foundation: The Link Function

The authors define a smoothed data distribution $p_\sigma$ by convolving the data distribution $\mu_{data}$ with a Gaussian kernel of variance $\sigma^2$:
$$p_\sigma = \mathcal{N}(0, \sigma^2 I) * \mu_{data}$$
They introduce a link function $\ell_\sigma(x)$:
$$\ell_\sigma(x) = \frac{1}{2}\|x\|^2 + \sigma^2 \log p_\sigma(x)$$

Key Theoretical Insight (Theorem 1):
As the noise scale $\sigma \to 0$, the derivatives of $\ell_\sigma$ recover the fundamental geometric operations of the manifold:

1. Gradient as Projection: The gradient $\nabla \ell_\sigma(x) \approx x + \sigma^2 \nabla \log p_\sigma(x)$ converges to the closest-point projection $\pi(x)$ onto the manifold $\mathcal{M}$.
2. Hessian as Tangent Projection: The Hessian $\nabla^2 \ell_\sigma(x) \approx I + \sigma^2 \nabla^2 \log p_\sigma(x)$ converges to the projection onto the tangent space $P_{T_x\mathcal{M}}$.

Crucially, $\nabla \log p_\sigma$ is the score function, which is the central quantity learned by diffusion models.

B. Practical Implementation: Score Networks

Instead of deriving geometric formulas, the authors leverage pretrained score networks $s(x, \sigma) \approx \nabla \log p_\sigma(x)$ trained via Denoising Score Matching (DSM).

• Approximate Projection: $v(x) = x + \sigma^2 s(x, \sigma)$ approximates $\pi(x)$.
• Approximate Tangent Projection: The Jacobian $v'(x)$ approximates $P_{T_x\mathcal{M}}$.
• Advantage: This allows the use of powerful, pretrained generative models without additional training for the optimization task.

C. Proposed Algorithms

Based on these approximations, the authors propose two inference-time algorithms:

1. Denoising Landing Flow (DLF):
   A continuous-time flow defined by:
   $$\dot{x} = -v'(x)\nabla f(v(x)) + \eta(v(x) - x)$$

   • The first term mimics Riemannian gradient descent using the approximate tangent projection.
   • The second term is a "landing" penalty that pulls the iterate back toward the manifold (approximated by $v(x)$), ensuring feasibility.
   • Theoretical guarantee: Converges to approximate stationary points with bounded distance to the manifold.

2. Denoising Riemannian Gradient Descent (DRGD):
   A discrete-time algorithm:
   $$x_{k+1} = v(x_k - \gamma_k v'(x_k)\nabla f(x_k))$$

   • Here, $v$ acts as an approximate retraction map.
   • The update step uses the approximate tangent projection $v'$ to compute the gradient step, followed by a projection back to the manifold via $v$.
   • Theoretical guarantee: Non-asymptotic convergence to approximate stationary points with bounds on the Riemannian gradient norm.

3. Key Contributions

1. Theoretical Link: Proved that the score function and its Jacobian of a smoothed data distribution asymptotically recover the closest-point projection and tangent space projection, respectively, under the manifold hypothesis.
2. First Score-Based Optimization Framework: Introduced the first algorithms (DLF and DRGD) that perform Riemannian optimization using only a pretrained score network, bypassing the need for explicit manifold definitions or latent space parameterizations.
3. Non-Asymptotic Guarantees: Established convergence rates and feasibility bounds for both algorithms as $\sigma \to 0$ and the score approximation error $\epsilon \to 0$.
4. Inference-Time Efficiency: The methods require only forward and backward passes through the neural network (no parameter updates), making them computationally efficient and applicable to any task where a score network is already available.

4. Experimental Results

The authors validated their approach on two types of problems:

• Synthetic Optimization (Orthogonal Group $O(n)$):

  • Task: Minimize Brockett's cost function on the orthogonal group.
  • Result: The method successfully found solutions with lower objective values than the best points in the training data, demonstrating the ability to generalize and optimize beyond the sample set. Accuracy improved as $\sigma \to 0$.

• Data-Driven Control (Reference Tracking):

  • Task: Control a discrete-time dynamical system (Double Pendulum and Unicycle Car) to track a reference trajectory, where the system dynamics are unknown and only input-output data is available.
  • Setup: The "manifold" is the set of feasible input-output trajectories (System Behavior Manifold).
  • Result: The DRGD algorithm generated control inputs that tracked the reference significantly better than the best trajectory in the training dataset. The resulting trajectories remained close to the true system behavior manifold, validating the feasibility of the approach in a real-world control setting.

5. Significance and Impact

• Paradigm Shift: Moves Riemannian optimization from a domain requiring explicit geometric knowledge to a data-driven regime where geometry is inferred from samples.
• Generative AI Integration: Unifies the fields of generative modeling (diffusion models) and constrained optimization. It allows practitioners to leverage the strong inductive biases of pretrained diffusion models for constrained design and control problems.
• Practical Applicability: Since the algorithms operate at inference time, they can be immediately applied to any domain where a diffusion model has been trained (e.g., image generation, molecular design, robotics) to solve constrained optimization problems without retraining the generative model.
• Robustness: The "Landing" mechanism ensures that even if the approximation is imperfect, the iterates remain close to the feasible set, providing a robust alternative to classical constrained optimization methods that often struggle with complex, implicit constraints.