Stochastic localization

For the Gaussian channel the observation is a Markov diffusion whose drift is the denoiser. Two facts make the tilt usable. By Girsanov the likelihood of the whole path $Y_{[0,t]}$ given $x$ is $e^{\langle Y_t,x\rangle-t\|x\|^2/2}$, which sees the path only through its endpoint, so $Y_t$ is sufficient and $\mu_t$ is a function of $Y_t$ alone. And by the innovation theorem of filtering, $B_t=Y_t-\int_0^t m(Y_s;s)\,ds$ is a Brownian motion in the observation filtration, since its increments are the observation increments minus their conditional mean. Together they give $dY_t=m(Y_t;t)\,dt+dB_t$, which removes both obstacles to sampling this way: the process can be generated without first drawing $x$, and the infinite-dimensional $\mu_t$ never has to be represented, only its mean. Uniqueness needs only the finite second moment.

The reversed OU sampler is this SDE in another clock, and every reverse diffusion is a localization. For OU, $Z_s=e^{-s}x+\sqrt{1-e^{-2s}}\,G$, and Tweedie gives the reverse drift through $\mathbb E[x\mid Z_s]$. Take $t=1/(e^{2s}-1)$, the signal-to-noise ratio at OU time $s$, which runs from $\infty$ to 0 as $s$ runs from 0 to $\infty$. Then $e^{-2s}=t/(1+t)$ and $$\sqrt{t(1+t)}\,Z_{s(t)}=tx+\sqrt t\,G\ \overset{d}{=}\ Y_t,$$ so the reversed OU process, rescaled by $\sqrt{t(1+t)}$, is the localization SDE. The general statement needs no computation. Given any forward diffusion $(Z_s)_{s\in[0,S]}$ and a decreasing clock $s(t)$, let $Y_t=(Z_u:u\in[s(t),S])$ be the forward path from the current noise level to the end. Each earlier $Y$ is a later one with a piece deleted, so the reversed sequence is Markov, and $Y_T$ contains $Z_0=x$, so the process is complete; extending the observed segment by an infinitesimal piece is one step of the reverse SDE.

In general the posteriors form a martingale, and sampling means running the reversed chain of observations. $\mu_t(A)=\mathbb E[\mathbf 1_A(x)\mid\mathcal F^Y_t]$ is a Doob martingale, and completeness with Lévy's upward theorem gives $\mu_t(A)\to\mathbf 1_A(x)$; conversely, a measure-valued martingale converging to $\delta_x$ is $\mathbb P(x\in\cdot\mid\mathcal F_t)$ for its own filtration, which some observation process generates. Because $x,Y_{t_k},\dots,Y_{t_1}$ is Markov, so is $Y_{t_1},\dots,Y_{t_k},x$, with kernels $P_{t,t'}(y,A)=\mathbb P(Y_{t'}\in A\mid Y_t=y)$, and the algorithm is to discretize $I$, approximate each $P_{t_k,t_{k+1}}$, and sample forward. In the Gaussian case Itô's formula shows the martingale directly. With $D_Ym=\mathrm{Cov}(x\mid Y_t)=:Q$, $D_Y^2m=\mathbb E[(x-m_t)^{\otimes3}\mid Y_t]$ and $\partial_tm=-\tfrac12\mathbb E[(x-m_t)\|x\|^2\mid Y_t]$, every drift term cancels and $$dm_t=\mathrm{Cov}(x\mid Y_t)\,dB_t ,$$ so the posterior mean moves only along directions where the posterior is still uncertain. ($D_Ym=\mathrm{Cov}(x\mid Y_t)$ is the identity Kadkhodaie et al. use to read a denoiser's Jacobian.)

The known samplers are different observation channels, all trained by the same path-space KL. What changes from channel to channel is the conditional expectation the network must learn. Revealing coordinate $i$ at an independent uniform time (the erasure process) needs next-coordinate conditionals and is autoregressive sampling in a random order. Resampling coordinates at Poisson times (the symmetric processes) gives the discrete diffusions, whose flip rates are explicit in the posterior marginals $\mathbb E[x_i\mid Y_t]$. The linear channel $Y_t=tAx+B_t$ samples $Ax$ and decodes it, and Poisson counts, half-space indicators and combinations of channels (whose likelihoods multiply) fit the same frame. Training is the same for all: minimize the KL between the true and the model path laws. For the Gaussian channel Girsanov gives $D(P\|\hat P)=\tfrac12\int_0^T\mathbb E\|Q(t)^{1/2}(m-\hat m)(Y_t)\|^2dt$, and the Pythagorean identity $\mathbb E\|x-\hat m\|^2=\mathbb E\|x-m\|^2+\mathbb E\|m-\hat m\|^2$ turns this into the denoising regression $\tfrac12\int\mathbb E\|\hat m(Y_t,t)-x\|^2dt$ up to a constant, which is score matching. For erasure the KL is the next-coordinate cross-entropy of an autoregressive model.

The isotropic channel can get a mixture's weights wrong, because which cluster a trajectory joins is decided in a window of width $1/n$. Take $\mu=p\,N((1-p)a,I_n)+(1-p)\,N(-pa,I_n)$ with $\|a\|\asymp\sqrt n$. The denoiser is $m(y;t)=\frac{y}{1+t}+a\,\varphi\big(\langle a,y\rangle/\|a\|^2;t\big)$ for an explicit scalar $\varphi$, and for large $\|a\|$ it agrees with the denoiser of one of the two clusters up to $O(1/n)$, except in a window of width $\Theta(1/n)$ around $\langle a,y\rangle/\|a\|^2=(1-2p)t$, the midpoint between the scaled centers, where $\partial_s\varphi$ is of order $n$. A trajectory's cluster is decided by how it crosses that window, so the mixture weight rests on a steep feature of the denoiser in a region of low probability, where training data are scarce. With $n=128$, $p=0.7$ and two-layer ReLU denoisers, the histogram of $\langle X,a\rangle/\|a\|^2$ is visibly off target. The remedy reveals the cluster label first: split the samples by the sign of their top principal component, train a denoiser on each half, draw the sign with its empirical frequency, then run the isotropic diffusion with that half's denoiser. This combines the half-space and Gaussian channels, and with the same architecture the projection matches much better. The evidence is empirical, and the paper grants that a larger network might also have succeeded.

A local denoiser cannot create long-range correlations, and revealing one global coordinate fixes it. For $\Sigma=I+\alpha\mathbf 1\mathbf 1^\top$ the exact denoiser $(I+t\Sigma)^{-1}\Sigma\,y$ is linear; restrict it to symmetric circulant convolutions of window $2r+1\ll n$, fitted by least squares at each $t$. The generated law is then a centered circulant Gaussian whose spectrum at frequency zero is $\hat c^{\rm gen}(0)=1+(2r+1)\alpha$, since its correlations reach only to range of order $r$. So $\langle\mathbf 1,\Sigma^{\rm gen}\mathbf 1\rangle=(1+(2r+1)\alpha)n$ against $\langle\mathbf 1,\Sigma\mathbf 1\rangle=\alpha n^2+n$, and projecting any coupling onto $\mathbf 1/\sqrt n$ with the triangle inequality gives $$W_2(\mu_n,\mu^{\rm gen}_{n,r})\ \ge\ \sqrt{n\alpha+1}-\sqrt{(2r+1)\alpha+1}\ \ge\ \tfrac12\sqrt{n\alpha}.$$ TV tends to 1 because the coordinate mean $\langle x,\mathbf 1\rangle/n$ has variance $\alpha+1/n$ under $\mu_n$ but $(1+(2r+1)\alpha)/n\to0$ under the generated law. Proposition 7.2 adds the channel $Y_{0,t}=\frac{bt}{n}\langle x,\mathbf 1\rangle+B_{0,t}$, which reveals the mean $b$ times faster than the rest, and uses hand-built local estimators: $\hat m_0$, Bayes-optimal for the mean given $Y_{0,t}$ alone up to $O(1/n)$, and the known-mean denoiser with that estimate plugged in. Then $W_2$ and TV tend to 0. Images that are mostly red or mostly blue behave the same way.

Choosing the channel is a design problem. Localization reduces sampling to approximating posterior means or marginals given $Y_t$, and the paper gives rules for choosing what to observe. The channel should preserve the symmetries of $\mu$, so that $\mu_t$ stays invariant. It should avoid times where $m(\cdot,t)$ has an information-computation gap, since no polynomial-time denoiser then exists and sampling fails (shown in a simple case by Montanari and Vu), and more generally sharp drops in the MMSE curve, which co-occur with hardness. And it should reveal latent variables early when $\mu$ is multimodal, as in Step 5. These are heuristics, and the paper's own evidence for them is the two toy examples.