Rao-Blackwellized Score Matching on Manifolds

Setup#

Generally, the paradigm of generative modeling is as follows: given a dataset, can we find some way to generate new data that seems as if it was pulled from the original dataset? A common example of this is images; a generative modeling task here might be that one is given a set of 100 images of cats and after some "learning," is asked to generate a new image of a cat. More formally, what we usually want here is to approximate some kind of probability distribution given samples from that distribution, then generating new data is as simple as sampling from our distribution.

Score Matching#

At first glance, this is quite a difficult task. In particular, one of the main obstacles here is normalization. Recall that a probability distribution must always have total probability mass 1 over its entire domain (Kolmogorov's unit measure axiom). But, when we sample, we are only given relative weightings between regions of the domain (e.g. I might notice that a calico cat is roughly twice as likely as a black cat). This is easy to get around conceptually, we can just renormalize (in our cat example, if calico and black are the only allowed cat colors, it follows immediately that the probability of a given cat being calico is 2/3 and black 1/3), but this is often difficult to do in an efficient manner. This obviously presents a pretty big problem for our generative modeling task: if we can't normalize the distribution efficiently, we can't learn/sample from it efficiently!

So, to get around this, [Hyv05] proposed the following paradigm: instead of approximating the distribution, just approximate its gradient (or its "score"). So now, instead of us learning some probability law $\tilde{\mathcal{P}}$ so that $\tilde{\mathcal{P}} \approx \mathcal{P}$ (here $\mathcal{P}$ is the true probability distribution) for every $x$ in our domain, instead, we want $\nabla_x \log \tilde{\mathcal{P}} \approx \nabla_x \log \mathcal{P}$ for every $x$, where of course $\nabla_x$ is the gradient with respect to $x$. This method, called score matching, gets rid of our earlier problem: since normalization must sum over all $x$, it drops out in the derivative and isn't something we need to worry about.

Manifolds#

Okay so now we've got a handle on the score matching part of the title, now let us turn our attention to manifolds. Manifolds are a really tricky subject and the formal definition of them is quite nasty¹, but intuitively, they're just a type of space that looks flat if you zoom in enough. A really good example here is Earth; Earth is a sphere, but to us, it seems flat since we are local enough. So now why do we care about manifolds? Well, manifolds tend to be really useful for representing structure.² And generally, the data we care about has some kind of structure to it (e.g. in our cat example, a form of structure is that every cat has whiskers). This structure imposes constraints on our data, and oftentimes, we can represent this as our data living on some manifold (in fact, a pretty well known theory in machine learning is that the data we learn is drawn from a manifold, the aptly named manifold hypothesis). So in these structured settings, we know that our data is drawn from a manifold, and thus, our probability law is defined as 0 off manifold. A concrete example is your computer screen. Your computer screen has only 2 dimensions (height and width), but every pixel technically exists in 3 dimensions. Because the pixels all exist on your computer screen, in the local coordinate system of the computer screen, we need only 2 dimensions to describe its position (height and width), but in the coordinate system of the room, we need 3 (height, width, and length; here we have fixed some reference point). But now, if I ask you to draw the distribution of where the pixels are located in the room, it makes no sense to assign any probability to the chair; a pixel can't exist there.

A careful reader will have noticed an issue here: if our data lives on a manifold, and we take the gradient with respect to ambient coordinates, there is no density in the ambient space at all, so the score is undefined everywhere on the manifold: the density falls off a cliff in every direction pointing off of it, and there is nothing to differentiate. Indeed, returning to our computer example, the screen has no thickness: on the screen, there is a nonzero probability of finding an on pixel, but immediately in front of or behind the screen, the probability is zero. So how do we deal with that? The standard method is just by corrupting the data — we just add some noise to the data so it lies slightly off the manifold, and then we no longer have this issue (intuitively, you can think of this as almost smoothing a peak). The problem is, now we're not learning the true law, we're learning a corrupted surrogate!

Thus, the natural question is What exactly is this method learning?

Rao-Blackwell Theorem#

So far, we know that score matching is a way of approximating a distribution by matching its derivatives and that manifolds are some type of mathematical way of representing structured data. Now, let's take a brief detour into statistics to study the Rao-Blackwell Theorem.

Consider the following setup: we have a (possibly unfair) coin that lands heads with probability $p$. Given that we observe $n$ flips of the coin, we wish to estimate $p$. Any rule that takes in these $n$ observations and outputs a prediction $\hat{p}$ is an estimator. There's two main desirable properties that we want from an estimator:

• Unbiasedness, i.e. it should be correct on average;
• Low Variance, i.e. it shouldn't vary wildly from one dataset to the next.

Getting unbiasedness is quite easy, and there's a number of rules that allow us to get an unbiased estimator. The harder part here is finding an estimator with low variance.

To illustrate this point, let's consider a deliberately bad rule. Write $X_i = 1$ if flip $i$ came up heads and $X_i = 0$ otherwise, and consider $\hat{p}_{\text{bad}} \doteq X_1$: look at the first flip, report 1 if heads and 0 if tails, and throw the other $n-1$ flips away. This is an unbiased estimator since the first flip is heads with probability $p$. However, its variance is quite large: it only outputs 0 or 1 and its variance is $p(1-p)$ regardless of $n$.

On the other hand, the natural estimator here is to count the heads and report the fraction. Writing $T \doteq X_1 + \cdots + X_n$ for the number of heads, the sensible rule is $\hat{p} = T/n$, which is also unbiased but has variance $p(1-p)/n$. The main point of Rao-Blackwell is that the sensible rule can be constructed by "cleaning up" the bad rule.

A full record of flips, say $HTTHT$, contains two separable pieces of information: how many heads there were, and in what order they arrived. The count obviously says something about $p$. The order is totally irrelevant: if I tell you 2 of 5 flips were heads, all $\binom{5}{2} = 10$ arrangements are equally likely, regardless of what $p$ is. So the ordering is pure noise as far as $p$ is concerned. A summary like $T$ that contains all the relevant information so that the remaining information is randomness whose distribution doesn't depend on the quantity we are estimating is called a sufficient statistic.

Our bad estimator is contaminated by the irrelevant part: it asks which flip came first, so it inherits variance from the ordering. The fix is to average that dependence away. Freeze the count at $T = t$, look at every arrangement of flips consistent with it, and average $\hat{p}_{\text{bad}}$ over all of them. By symmetry, each of the $n$ positions is equally likely to hold one of the $t$ heads, so $$\mathbb{E}\bigl[X_1 \mid T = t\bigr] = \frac{t}{n}.$$ Averaging our bad rule over all the irrelevant information (the permutations) gives back our good estimator.

That's the Rao-Blackwell Theorem. Given an unbiased estimator $\delta(X)$ and a sufficient statistic $T(X)$, define $$\delta^{\mathrm{RB}}(X) \doteq \mathbb{E}\bigl[\delta(X) \mid T(X)\bigr],$$ the average of $\delta$ over all datasets sharing the same summary. Then $\delta^{\mathrm{RB}}$ is still unbiased, and $\operatorname{Var}(\delta^{\mathrm{RB}}) \leq \operatorname{Var}(\delta)$, with equality exactly when $\delta$ was already a function of $T$ alone.³

Intuitively, it just says to find the part of your data that is pure noise with respect to the thing you want, and integrate it out.

Rao-Blackwellization on Manifolds#

Equipped with an understanding of the main ideas, we can now start to answer the question we posed.

The first step is to identify our analog of a sufficient statistic here. Recalling that our probability law is defined only on the manifold, we really only care about points on the manifold. So, a natural sufficient statistic type thing to condition on here is the nearest point projection: the map $\pi$ that sends every noisy point back to whichever point on the manifold it's closest to.

In order to condition on anything, we must first define our estimator. In denoising score matching (score matching on our corrupted data), we take a clean sample $Z$ from our data distribution $q$ (which lives on the manifold $M$), corrupt it into $X \doteq Z + \sigma\xi$ with a bit of Gaussian noise $\xi \sim \mathcal{N}(0, I_D)$, and train a network to predict the direction that undoes the corruption, $(Z-X)/\sigma^2$. Observe that this direction splits into two pieces: the part pointing along the manifold, and the part pointing off of it. The off-manifold piece is, to leading order, $(\pi(X) - X)/\sigma^2$, which we can compute from $X$ without knowing $Z$. All of the signal is in the tangent piece: $$T_\sigma \doteq \frac{1}{\sigma^2} P_{T_{\pi(X)} M}(Z - X) \in T_{\pi(X)} M,$$ where $P_{T_{\pi(X)}M}$ keeps the component pointing along $M$ at $\pi(X)$.

This is very similar to our coin flipping example. Our observation $X$ carries two kinds of information: roughly where on $M$ we came from, and how far off $M$ the noise shifted us. The second kind is irrelevant for what we want, just like the ordering of the flips was pure noise with respect to $p$. So, we just average over it. Define the Rao-Blackwellized target $$r_\sigma(z) \doteq \mathbb{E}\bigl[T_\sigma \mid \pi(X) = z\bigr],$$ the average of $T_\sigma$ over every noisy point that projects back to the same $z$. Averaging over a fiber is our analog of averaging over the permutations.

In fact, this reduces the variance just like in the Rao-Blackwell Theorem. $\pi(X)$ isn't a sufficient statistic in the classical sense, but we don't need it to be: the law of total variance already gives $\operatorname{Var}(r_\sigma) \leq \operatorname{Var}(T_\sigma)$. In the below image, $r_\sigma$ is the Rao-Blackwellized target, while $T_\sigma$ is the non-RBized target.

The main benefit here is that we are able to estimate what we want with significantly less variance, and that becomes more important as the amount we corrupt our data by (represented by $\sigma$) gets smaller.

What Does DSM Learn?#

We now have a target that has low variance, but we still have not answered the question we posed earlier. Because we had to corrupt the data to make score matching well defined at all, we aren't learning the true law, we're learning some corrupted surrogate. So, what's the surrogate?

One can show (via a lot of tedious math), that expanding $r_\sigma(z)$ as a power series in the noise level $\sigma$ gives $$r_\sigma(z) = \nabla_M \log q(z) + \sigma^2\bigl[b_q(z) + g_M^{\text{ext}}(z) + g_M^{\text{inh}}(z)\bigr] + O(\sigma^4).$$

The first term is exactly what we want to see: $\nabla_M \log q$ is the intrinsic score: the gradient of the log density taken along the manifold, which is what we wanted to learn. So at leading order, the naive ambient approach of corrupting with ordinary Gaussian noise and then cleaning up with Rao-Blackwellization recovers the right thing. All good there.

Things get a bit messier with the $\sigma^2$ term; in some sense, the "price" of corrupting the data. It splits neatly into three pieces, but before we break them down, we need to introduce a useful idea from differential geometry.

Curvature. In differential geometry, there are two notions of curvature. Extrinsic curvature describes how $M$ bends inside the ambient space it sits in. Intrinsic curvature is what one could measure while living on $M$, with no knowledge of the ambient space at all. These are quite different. For example, take a cylinder: it obviously bends in $\mathbb{R}^3$, so it has extrinsic curvature, but you can unroll it flat onto a table without stretching or tearing anything, so an ant living on it could never distinguish it from a plane. The cylinder is extrinsically curved and intrinsically flat. Alternatively, a sphere is neither intrinsically nor extrinsically flat (which is why you can't flatten an orange peel without tearing it⁴). It's the exact reason that every flat world map ends up distorting something.

With that distinction in mind, let's characterize the three corrections.

1. Tweedie bias $b_q(z)$. This part actually has nothing to do with the embedding. Gaussian smoothing blurs the density a little bit, and that changes its score [Efr11], so this is just the flat correction, transplanted onto $M$ by using the gradient and Laplacian along the manifold. It is the price of smoothing that any method pays, intrinsic ones included.
2. Extrinsic curvature term $g_M^{\text{ext}}(z)$. This is a linear operator built from ways of measuring both types of curvature and applied to the score.⁵ This is the term that measures how the embedding of $M$ in the ambient space affects our process.
3. Inhomogeneity drift $g_M^{\text{inh}}(z)$. This term measures how much the curvature of $M$ changes as you move around it, so it vanishes when the manifold curves the same way everywhere.⁶ A canonical example is the sphere, since every point looks like every other.

Both geometric terms are built from the embedding alone: the curvature operator in $g_M^{\text{ext}}$ and the drift $g_M^{\text{inh}}$ are determined by how $M$ sits inside the ambient space, never by $q$. The data distribution enters only through the score that the operator is applied to, which is the thing we already trained a model to predict. So given a trained score model $\hat{s}(z) \approx \nabla_M \log q(z)$, we can simply subtract the bias off afterwards: apply the curvature operator to $\hat{s}(z)$, add the drift, scale by $\sigma^2$, subtract. No retraining, and no access to $q$ required.⁷

A Happy Coincidence on $S^2$#

Let's now restrict to spheres, a really good toy example. On $S^d$ (the sphere in $d+1$ dimensions), the curvature is the same in every direction at every point, and plugging the values in collapses the entire extrinsic operator down to a single scalar: $$g_{S^d}^{\text{ext}}(z) = \bigl(1 - \tfrac{d}{2}\bigr)\nabla_{S^d}\log q(z).$$ The inhomogeneity drift is already zero, since a sphere curves identically everywhere. So on a sphere the entire geometric correction is one scalar $\alpha_d \doteq 1 - d/2$ times the score:

• $d = 1$: $\alpha_1 = +1/2$;
• $d = 2$: $\alpha_2 = 0$;
• $d = 3$: $\alpha_3 = -1/2$;
• $d \geq 4$: you get the point.

A neat result here is that at $d=2$ the coefficient is exactly zero. On the round 2-sphere both geometric terms vanish, and ambient DSM recovers the intrinsic score up to only the flat Tweedie bias that every method pays anyway.

This is demonstrated in the top row of the below figure. The bottom row shows the sample-complexity gain from Rao-Blackwellization: the RB target gives lower score error at small and moderate training set sizes, with the gap closing as $N$ grows.

Ambient DSM is often observed to work about as well as intrinsic methods on data living on $S^2$ (e.g. earth-science event data like earthquake or wildfire locations), even though score matching on manifold-supported data is ill-posed. The $\alpha_2 = 0$ cancellation offers an explanation: on the round 2-sphere, the leading-order geometric bias vanishes, so ambient and intrinsic methods agree up to the Tweedie term.

P.S. Apparently some account also made one of those AI generated reels of voiceovers with Minecraft gameplay in the background, if that's more your style, you can check it out here :).