Manifold Aware Denoising Score Matching (MAD)

TL;DR

This paper introduces a simple modification to denoising score matching that implicitly accounts for data manifolds, reducing computational complexity while maintaining effectiveness, especially for distributions over rotations and discrete data.

Contribution

It proposes a decomposition of the score function into known and learned components, enabling manifold-aware density modeling without explicit manifold learning.

Findings

Derived analytical forms of known score components for rotations and discrete distributions.

Demonstrated the approach's utility on specific manifold-structured data.

Reduced computational burden compared to explicit manifold learning methods.

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.