Diffusion Models for Inverse Problems on Riemannian Manifolds

Ivan Lee, Xingyu Xu, Yuejie Chi

Keywords: Diffusion models, Riemannian manifolds, Inverse problems, Posterior sampling

Abstract:

Score-based diffusion models have become powerful learned priors for solving
Bayesian inverse problems in Euclidean spaces, but existing diffusion-based
posterior samplers do not extend to data that is intrinsically supported on
a Riemannian manifold. We introduce Manifold DPnP, a generalization
of the diffusion plug-and-play (DPnP) sampler to arbitrary complete
Riemannian manifolds, by reformulating each sampling iteration as a
manifold-valued SDE driven by Brownian heat flow and using the
Bismut--Elworthy--Li formula to estimate the score of the likelihood under
this flow. We further derive Manifold DPS, an extension of the DPS
sampler obtained from a Riemannian Tweedie's formula based on Varadhan's
short-time heat-kernel asymptotics. Experiments on the flat torus, on real
earthquake locations on the sphere, and on a random-walk trajectory
reconstruction problem show that both samplers substantially outperform a
prior-free MALA baseline. Our results close a gap between Euclidean
diffusion-based inverse problem solvers and intrinsic manifold-valued
posterior sampling.

Submission Number: 60