Enabling Denoising Score Matching Type Training for Manifold Diffusion Model via Momentum and Splitting

Denoising score matching (DSM) is the cornerstone of scalable diffusion model training. However, DSM becomes inapplicable on curved manifolds: the curved geometry makes the forward diffusion process nonlinear and the transition kernel is no longer tractable. Existing approaches therefore revert to implicit score matching (ISM) with expensive Brownian simulation, approximate the intractable heat kernel, or bypass the transition kernel altogether via flow matching with the logarithm map, each incurring significant computational cost or approximation error on general manifolds. We propose splitting diffusion, a new framework that restores DSM-like training on general Riemannian manifolds by replacing the intractable Brownian forward process with a splitting scheme that alternates two steps: (1) an Ornstein-Uhlenbeck velocity update in the tangent space, whose transition kernel has a closed-form solution; (2) a deterministic geodesic transport step. This splitting yields a fully tractable transition kernel locally, enabling DSM-style training while requiring only the exponential map as a geometric oracle. The effectiveness of splitting diffusion is demonstrated on complicated manifolds and high-dimensional data.