Manifold Aware Denoising Score Matching (MAD)

Alona Levy-Jurgenson; Alvaro Prat; James Cuin; Yee Whye Teh

cs.LGMar 2, 2026

Manifold Aware Denoising Score Matching (MAD)

Authors: Alona Levy-Jurgenson, Alvaro Prat, James Cuin, Yee Whye Teh

Organizations: Department of Statistics University of Oxford · Department of Mathematics, Imperial College London, London, United Kingdom

Abstract

A major focus in designing methods for learning distributions defined on manifolds is to alleviate the need to implicitly learn the manifold so that learning can concentrate on the data distribution within the manifold. However, accomplishing this often leads to compute-intensive solutions. In this work, we propose a simple modification to denoising score-matching in the ambient space to implicitly account for the manifold, thereby reducing the burden of learning the manifold while maintaining computational efficiency. Specifically, we propose a simple decomposition of the score function into a known component $s^{base}$ and a remainder component $s-s^{base}$ (the learning target), with the former implicitly including information on where the data manifold resides. We derive known components $s^{base}$ in analytical form for several important cases, including distributions over rotation matrices and discrete distributions, and use them to demonstrate the utility of this approach in those cases.

Figures & tables

Figure 1: Introduction to MAD through a toy discrete distribution over

\mathcal{M}=\{(-1,0),(1,0)\}

(red) with

p=\{0.1,0.9\}

, respectively, and

p^{base}=\{0.5,0.5\}

. Vector fields are shown for the actual score

s

, the base score

s^{base}

and the difference

s-s^{base}

for

\sigma_{t}=0.8

. The magnitude of

s

(DSM’s learning target) and

s-s^{base}

(MAD’s learning target) as a function of

\sigma_{t}

is shown for

x=(1,0)

Method	Volcano	Earthquake	Flood	Fire
RSGM-divfree	0.1272 ± 0.0289	0.0602 ± 0.0101	0.0579 ± 0.0094	0.0779 ± 0.0172
RSGM-ambient	0.1114 ± 0.0090	0.0803 ± 0.0067	0.0751 ± 0.0087	0.0830 ± 0.0152
DSM	0.1082 ± 0.0217	0.0488 ± 0.0096	0.0503 ± 0.0092	0.0456 ± 0.0165
MAD (ours)	0.1030 ± 0.0185	0.0526 ± 0.0108	0.0500 ± 0.0100	0.0452 ± 0.0131
Data size	827	6120	4875	12809

Table 1: Earth MMD results (mean ± std) computed using the minimum between 1000 samples and the size of the available test set.

Method	$K=16$	$K=32$	$K=64$	Sample time (s/1k)
FFF	0.0230 ± 0.0023	0.0242 ± 0.0016	0.0204 ± 0.0018	0.0040 ± 0.0003
RSGM	0.0155 ± 0.0025	0.0171 ± 0.0024	0.0190 ± 0.0043	0.4628 ± 0.0190
DSM	0.0222 ± 0.0021	0.0209 ± 0.0046	0.0178 ± 0.0031	0.2216 ± 0.0350
MAD (ours)	0.0229 ± 0.0026	0.0220 ± 0.0052	0.0180 ± 0.0029	0.2305 ± 0.0360

Table 2:

\mathrm{SO}(3)

MMD (mean ± std) computed between generated sets and test set using 5,000 samples, and sampling time for 1K samples.

Panel of Figure 2 — Figure 2: Earth data (left panel) vs. generated samples from different methods.The colour map, defined by the polar axis, is left for visualisation purposes.

Panel of Figure 3 — Figure 3: $\mathrm{SO}(3)$ results for $K=16$ in the form of a Mollweide projection for different methods and data. All figures are from seed $0$ . Points are coloured according to angle, not component.

Figure 4: MMD vs cumulative training-step time. Shaded areas are ± std across

5

runs. MMD is computed against the training set for

K=64

using 1,000 samples for efficiency (the table uses 5,000). Note that for better clarity of comparison between all methods, the time shown is up to 80, but FFF ends after 400.

Figure 5: 2D image snapshots of a randomly rotated cone, cylinder, cube, icosahedron (a-d) and Mollewide projections of the corresponding pose distributions in $\mathrm{SO}(3)$ . Black circles in (a-b) denote samples from MAD and are excluded from (c-d) for visual clarity.

Figure 6: Top Row: Manifold Drift (

|1-\|x\|_{2}|

). MAD (blue) maintains consistently lower drift, indicating predictions stay closer to valid rotations. DSM (orange) exhibits higher drift, indicative of ’ghost rotation’ averaging. Bottom Row: Sample Spread (lower is better).

Method	Cone	Cyl.	Tet.	Cube	Ico.	Avg.
Deng et al. (2020)	10.1	15.2	16.7	40.7	29.5	22.4
IPDF ( Murphy et al., 2022 )	1.4	1.4	4.6	4.0	8.4	4.0
DNF ( Liu et al., 2023 )	0.5	0.5	0.6	0.6	1.1	0.7
DSM + Quotient	0.7	0.9	1.2	0.7	0.9	0.9
MAD + Quotient (ours)	0.9	0.6	0.7	0.6	0.6	0.7

Table 3: Spread estimation on SymSol I . Values are in degrees (

\downarrow

). Baseline results are taken from Liu et al. (2023) .

Panel of Figure 7 — Figure 7: Generating discrete distributions – uniform (top) and skewed (bottom). Data (left) and generated samples from each method (middle, right). Samples are not projected onto the support.

Appendix figures & tables3 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 8: Demonstration of the base score for $S^{2}$ – samples generated using only the base score (the score function for a uniform distribution on $\mathcal{S}^{2}$ ) with $\sigma_{\min}=1e^{-4}$ . Samples are not projected onto the manifold. The colour is by $z$ -axis values and is for ease of visualisation only.

Figure 9: Training loss for MAD is consistently lower throughout training compared to that of DSM. Losses are for $\mathrm{SO}(3)$ with $K=64$ . Shaded area represents ±std across $5$ random seeds.

Category	Value
Optimization
Optimizer	AdamW (default betas)
Batch Size	64
Training Steps	128,000
Peak Learning Rate	$4\times 10^{-4}$
Min Learning Rate	$1\times 10^{-4}$

Appendix

Table 4: Hyperparameter configuration for SymSol I experiments.

Score-Based Generative Modeling

Score-based generative modeling learns data-distribution gradients to generate samples through reverse diffusion, flows, and related stochastic dynamics.

See all 110 papers

Rao-Blackwellized Score Matching on Manifolds
May 25, 2026Divit RawalScore Matching Score-Based Generative Modeling
Department of Statistics, University of California, Berkeley, Berkeley CA, U.S.A..
Asymptotic Learning Curves for Diffusion Models with Random Features Score and Manifold Data
Mar 24, 2026Anand Jerry George, Nicolas MacrisIntrinsic Dimensionality Random Feature Methods
École Polytechnique Fédérale de Lausanne (EPFL), Lab for Statistical Mechanics of Inference in Large Systems (SMILS), CH-1015 Lausanne, Switzerland
Riemannian Denoising Diffusion Probabilistic Models
May 7, 2025Zichen Liu, Wei Zhang, Christof Schütte +1Generative Modeling Score-Based Generative Modeling
Center for Data Science, Peking University, Beijing 100871, P.R. China · Zuse Institute Berlin, Takustrasse 7, Berlin 14195, Germany · Institute of Mathematics, Freie Universität Berlin and Zuse Institute Berlin, Takustrasse 7, Berlin 14195, Germany +1

May 25, 2026J/K move · Enter open · S save

Rao-Blackwellized Score Matching on Manifolds

Divit Rawal

Department of Statistics, University of California, Berkeley, Berkeley CA, U.S.A..

Score Matching Score-Based Generative Modeling

Open paper AskPDF arXiv

We study denoising score matching (DSM) when the latent distribution is supported on a smooth embedded manifold

M \subset \mathbb{R}^D

. Under ambient Gaussian corruption, the tangent denoising target contains a singular normal-fiber noise channel whose variance diverges as

d/σ^2

σ\to 0^+

. We show that conditioning on the nearest-point projection

π(X)

canonically removes this singularity: the resulting conditional expectation is the unique

L^2

-optimal Rao-Blackwellized predictor of the tangent DSM target among all estimators depending only on the projected observation

π(X)

. We then compute the small-noise expansion of this canonical target and show that it equals the intrinsic Riemannian score up to an explicit order-

σ^2

correction that decomposes into an intrinsic Tweedie term and an extrinsic curvature term involving the Weingarten and Ricci operators. In the flat case, the construction reduces exactly to ordinary lower-dimensional Gaussian DSM, while on

S^d

the extrinsic correction simplifies to the scalar factor

(1-d/2)\nabla_M \log q

; this extrinsic

σ^2

correction cancels identically on

S^2

, though the intrinsic Tweedie term remains.

Abstract

Score-Based Generative Modeling

Explore similar work

Rao-Blackwellized Score Matching on Manifolds

Asymptotic Learning Curves for Diffusion Models with Random Features Score and Manifold Data

Riemannian Denoising Diffusion Probabilistic Models