Gradient Descent on Point Clouds and Applications in Learned Operator Correction

Technical Summary: Gradient Descent on Point Clouds and Applications in Learned Operator Correction

Problem Statement
The paper addresses the challenge of minimizing an energy functional $E(x)$ over an unknown manifold $\mathcal{M} \subset \mathbb{R}^D$. While data in many applications (particularly imaging) satisfies the manifold assumption (low intrinsic dimension despite high extrinsic dimension), the manifold is typically unknown and only accessible via a finite set of samples (a point cloud) $\{x_i\}_{i=1}^n \subset \mathcal{M}$.

Standard gradient descent in the ambient Euclidean space fails in this context for two primary reasons:

1. Geometric Constraint: Iterates may drift away from the manifold, entering regions where the energy landscape or gradient information is unreliable or undefined.
2. Learned Operator Correction: In inverse problems, learned corrections (e.g., neural networks correcting a cheap forward model) are trained on specific data manifolds. These corrections often fail or produce erroneous gradients when iterates leave the vicinity of the training data. Existing methods that attempt to correct this by adding points along the optimization trajectory are computationally prohibitive.

Methodology
The authors propose a two-step gradient descent scheme that operates directly on the point cloud without requiring an explicit parameterization of the manifold. The method assumes the manifold is a compact, $d$-dimensional, twice continuously differentiable Riemannian manifold with positive reach.

1. Continuum Formulation (Known Manifold):
   The authors first define a gradient flow for a known manifold $\mathcal{M}$:
   $$\dot{\xi}(t) = -\Pi_{T_{\xi(t)}\mathcal{M}} \nabla E(\xi(t)) + \lambda n(\xi(t))$$
   where $\Pi_{T_{\xi(t)}\mathcal{M}}$ is the projection onto the tangent space, $n(\xi(t)) = P_\mathcal{M}(\xi(t)) - \xi(t)$ is the normal vector pointing to the manifold, and $\lambda > 0$ is a penalty parameter. This flow ensures the sequence converges to the manifold while minimizing energy.

2. Discrete Scheme (Unknown Manifold):
   Since the manifold is unknown, the tangent spaces and normal vectors are approximated using local Principal Component Analysis (PCA) on the point cloud samples.

   • Step 1 (Tangential Descent): The gradient is projected onto an estimated tangent space. This is achieved by averaging projections onto the tangent spaces of the $K$-nearest neighbors of the current iterate $\xi_k$. The tangent space at a sample $x_i$ is estimated via the eigenvectors of the local empirical covariance matrix.
   • Step 2 (Normal Correction): The iterate is projected back toward the manifold. The normal vector is estimated by averaging the difference between the current point and the local average of the manifold samples, projected onto the estimated normal space (orthogonal complement of the tangent space).

   The algorithm iterates between these steps, ensuring the trajectory remains within a tubular neighborhood of the manifold.

Theoretical Contributions
The paper provides rigorous convergence analysis under suitable stability and sampling assumptions:

• Convergence to Manifold: The authors prove that the discrete scheme converges to the solution of the continuum gradient flow PDE as the time step $\tau \to 0$.
• Error Bounds: They establish that the distance of the iterates to the manifold vanishes as the time step, sampling error, and tangent plane estimation error vanish.
• Local Minimizer Convergence: Under assumptions of geodesic convexity and local exponential stability, the discrete scheme converges to a local minimizer. The error bound depends on the consistency of the learned tangent and normal approximations ($\beta$ and $\gamma$).
• PCA Consistency: The paper includes a proof (Theorem 2.6) showing that the PCA-based tangent plane approximation converges to the true tangent plane with a rate dependent on the sampling density and the smoothness of the manifold density.

Applications and Results
The methodology is applied to Learned Operator Correction in inverse problems, specifically Photoacoustic Tomography (PAT).

• Context: In PAT, the exact forward model $A$ is computationally expensive, while a cheap approximation $\tilde{A}$ introduces artifacts. A neural network is trained to correct the normal operator $\tilde{A}^*\tilde{A}$ to approximate $A^*A$.
• Problem: Standard optimization using the learned correction diverges or produces artifacts when iterates leave the training data manifold because the correction is unreliable outside this region.
• Experiments:
  • Synthetic 2D Manifold in 3D: The authors demonstrate that point-cloud gradient descent is robust to noisy gradients (where noise scales with distance to the manifold), whereas standard Euclidean gradient descent fails significantly under the same noise conditions.
  • PAT Reconstruction:
    • Without manifold constraints, the learned correction fails to converge to a good solution, producing artifacts worse than the uncorrected approximation.
    • With the proposed point-cloud gradient descent, the optimization is constrained to the data manifold. The results show that the corrected model converges to a solution with accuracy comparable to the expensive exact model, significantly outperforming the uncorrected approximation.
    • The method is shown to be robust to the number of samples, though performance degrades slightly with fewer samples due to increased noise in the gradient estimation.

Significance and Claims
The authors claim that their primary novelty lies in developing a provably convergent gradient descent method that operates directly on unknown manifolds represented solely by point clouds.

• Unlike classical approaches requiring explicit geometric knowledge, this method combines local tangent estimation and projection to keep iterates close to the data manifold.
• The framework is particularly effective for learned operator correction, where it stabilizes optimization by preventing excursions into regions where learned models are inaccurate.
• The paper positions this as a bridge between manifold-based optimization and data-driven modeling, enabling gradient-based methods in settings where the underlying geometry is not explicitly known but data samples are available.

The authors remain modest, noting that while the method stabilizes optimization, the corrected gradients are not perfect even on the manifold, necessitating strategies like early stopping. They also acknowledge that fine sampling of the manifold is necessary for high accuracy, suggesting that adaptive sampling near the minimizer may be a future requirement for broader application.