2. Scores and denoisers
Tweedie's formula and the learning objective Draft
The score and its posterior form
The transport velocity in Section 1.3 depends on a quantity we do not know: the gradient of the log density. This score points toward increasing density, but evaluating it seems to require the very distribution we are trying to learn. The way out is to express it as an average of something we do know—the noise we add during training. Least squares can learn that average from samples.
The score describes a local change in density. At the top of a smooth peak it is zero, even if that peak contains much of the probability. On either side it points uphill. Its magnitude tells us how quickly the log density changes with position. A sampler uses this local information repeatedly as it moves through the changing densities.
Definition 2.1.1 (score). The score of a density \(p\) on \(\mathbb R^d\), positive and differentiable, is the vector field \(\nabla \log p\). For the noised laws of a forward process the series writes \(s_t(x) = \nabla \log p_t(x)\).
Proposition 2.1.2 (three properties). (a) The score does not see the normalization: if \(p = \tilde p / Z\) then \(\nabla \log p = \nabla \log \tilde p\). (b) Under \(p\) the score has mean zero: \(\int p\, \nabla \log p\, dx = \int \nabla p\, dx = 0\) whenever \(p\) vanishes at infinity. (c) Its second moment is the Fisher information of Definition 1.4.1: \(\int p\, \|\nabla \log p\|^2 dx = I(p)\).
Property (a) makes the score useful when a density's normalizing constant is too expensive to compute: differentiating the log density removes that constant. Property (b), its zero mean, will be used in Theorem 2.3.4. Property (c) connects its size to the Fisher information introduced in Remark 1.4.6.
top log p_sigma of R2 at sigma = 0.2 (alpha = 1) over p_sigma as a wash, and the tangent at the probe
x = -1.00: log p_sigma = -1.36 and its slope, the score, is -3.90 (the slope of p_sigma itself is -1.00)
bottom the score s_sigma = (log p_sigma)': zero at the modes -1.400, 0.200 and 1.600 (heights 0.561, 0.559
and 0.221, basins of mass 0.449, 0.354 and 0.197) and at the valleys -0.515 and 0.872; steepest at
x = -0.68, where it is -6.50 and p_sigma is 0.047; dotted, the score +0.53 of the Gaussian with R2's
mean -0.24 and variance 1.40; under p_sigma the score averages 0.000 and its mean square is I = 9.93
(the live figure needs JavaScript)
Example 2.1.3 (the score of a Gaussian, and of the kernel). For \(p = \mathcal N(m, v I)\), \(\nabla \log p(x) = -(x - m)/v\): the score points to the mean, with a slope that is the precision. For the forward kernel \(p_t(x \mid x_0) = \mathcal N(x; \alpha_t x_0, \sigma_t^2 I)\),
\[\nabla_x \log p_t(x \mid x_0) \;=\; -\frac{x - \alpha_t x_0}{\sigma_t^2} \;=\; -\frac{z}{\sigma_t} \qquad \text{when } x = \alpha_t x_0 + \sigma_t z .\]For a known clean point, the kernel's score is easy to compute: take the negative of the added noise and divide by the noise level. When we observe only the noisy point, many clean points could have produced it. The next theorem says to average their kernel scores using the posterior probabilities.
There is a useful distinction here: the noise is independent of the clean point, but it is not independent of the noisy observation. Seeing a large positive observation makes some combinations of clean point and noise more plausible than others. The network can therefore predict something about the noise from the observation, even though it could predict nothing about a fresh noise draw in isolation.
Theorem 2.1.4 (the score is a posterior average). Let \(x_t = \alpha_t x_0 + \sigma_t z\) with \(x_0 \sim p_0\) and \(z \sim \mathcal N(0, I)\) independent, and \(\sigma_t > 0\). Then for every \(x\),
\[s_t(x) \;=\; \mathbb E\bigl[\nabla_x \log p_t(x \mid x_0) \,\big|\, x_t = x\bigr] \;=\; -\frac{1}{\sigma_t}\, \mathbb E\bigl[z \,\big|\, x_t = x\bigr]. \tag{2.1}\]Proof. The marginal is \(p_t(x) = \int p_t(x \mid x_0)\, p_0(dx_0)\), and the kernel is smooth in \(x\) with derivatives dominated by an integrable function, so differentiation under the integral is permitted:
\[\nabla p_t(x) = \int \nabla_x p_t(x \mid x_0)\, p_0(dx_0) = \int p_t(x \mid x_0)\, \nabla_x \log p_t(x \mid x_0)\, p_0(dx_0).\]Dividing by \(p_t(x)\) turns \(p_t(x \mid x_0)\, p_0(dx_0)/p_t(x)\) into the posterior law of \(x_0\) given \(x_t = x\), by Bayes' rule, and the integral into the conditional expectation. The second form substitutes Example 2.1.3 and \(z = (x - \alpha_t x_0)/\sigma_t\). ∎
This identity turns score estimation into regression. Least squares learns a conditional expectation, as Lemma 2.3.6 will show. During training we know the target noise \(z\) because we generated it ourselves. We never need to evaluate the unknown data density.
left 1,000 pairs (x_0, z) of R1 at sigma = 0.8, seed 1, drawn as z against x_0: two copies of one cloud,
the 510 draws with x_0 = +1 averaging +0.003 and the 490 with x_0 = -1 averaging -0.023, and
corr(z, x_0) = 0.013
right the same pairs as z against x_t = x_0 + sigma z, on the lines z = (x -+ 1)/sigma, corr(z, x_t) = 0.634
(exact 0.625); through them E[z | x_t = x] = -sigma s_sigma(x), and the means of the draws in 16 bins
of x_t, within 0.033 of the curve on average and 0.111 at worst; at x_t = 0.5 the noise was -0.625
with probability 0.827 or 1.875 with probability 0.173, mean -0.192
(the live figure needs JavaScript)
Example 2.1.5 (Theorem 2.1.4 on R1). For the two points \(\pm 1\) under a variance-exploding process, \(\alpha = 1\), the posterior of the sign given \(x_\sigma = x\) is, by Bayes' rule with equal priors and the kernels \(\mathcal N(\pm 1, \sigma^2)\),
\[P(+ \mid x) = \frac{e^{-(x - 1)^2/2\sigma^2}}{e^{-(x - 1)^2/2\sigma^2} + e^{-(x + 1)^2/2\sigma^2}} = \frac{1}{1 + e^{-2x/\sigma^2}},\]a logistic function of \(x\) with slope \(2/\sigma^2\). The conditional scores are \(-(x \mp 1)/\sigma^2\), and their posterior average is
\[\begin{aligned} s_\sigma(x) &= -\frac{x - 1}{\sigma^2} P(+ \mid x) - \frac{x + 1}{\sigma^2} P(- \mid x) \\ &= \frac{(2P(+ \mid x) - 1) - x}{\sigma^2} = \frac{\tanh(x/\sigma^2) - x}{\sigma^2}, \end{aligned}\]since \(2P(+ \mid x) - 1 = \tanh(x/\sigma^2)\). At \(\sigma = 1.6\) and \(x = 0.5\), \(P(+ \mid x) = 0.596\) and \(s_\sigma(x) = -0.120\); at \(\sigma = 0.8\) the same point has \(P(+ \mid x) = 0.827\) and \(s_\sigma(x) = +0.240\). The score changes sign between the two noise levels: at the larger one the hump of \(p_\sigma\) is centred at the origin and \(x = 0.5\) sits on its right flank, so the score points left; at the smaller one \(p_\sigma\) has a mode near \(+0.88\) and the score points toward it. ◁
Example 2.1.6 (mixtures). For \(p_0 = \sum_k w_k\, \mathcal N(m_k, s_k^2 I)\), the noised law is the mixture \(\sum_k w_k\, \mathcal N(\alpha_t m_k, v_{k,t} I)\) with \(v_{k,t} = \alpha_t^2 s_k^2 + \sigma_t^2\) (Example 1.2.3), the posterior responsibilities are \(r_k(x) \propto w_k\, \mathcal N(x; \alpha_t m_k, v_{k,t} I)\), and
\[s_t(x) = \sum_k r_k(x)\, \frac{\alpha_t m_k - x}{v_{k,t}} .\]This is what every figure of the series evaluates for R2, R3 and the 64-centre law of the opening figure: a softmax over the centres, then a weighted sum of pulls toward them. With \(s_k = 0\) and the training images as centres this is the posterior mean over a training set, the memorizing minimizer of Section 2.5, which the worked examples of Sections 2.2, 3.3 and 4.4 run on \(L\).
Tweedie's formula
Estimating the score and estimating the clean point are closely related tasks. Tweedie's formula gives the exact conversion between them: use the score to correct the noisy observation, then undo the signal scaling.
Theorem 2.2.1 (Tweedie's formula). Under the assumptions of Theorem 2.1.4,
\[D_t(x) \;:=\; \mathbb E[x_0 \mid x_t = x] \;=\; \frac{x + \sigma_t^2\, s_t(x)}{\alpha_t}, \qquad \hat\varepsilon_t(x) \;:=\; \mathbb E[z \mid x_t = x] \;=\; -\sigma_t\, s_t(x). \tag{2.2}\]Proof. The second identity is (2.1). For the first, take conditional expectations in \(x_0 = (x_t - \sigma_t z)/\alpha_t\). ∎
Robbins attributed the formula to Tweedie; Miyasawa established the Gaussian case, and Efron later used it to study selection bias.H. Robbins, "An empirical Bayes approach to statistics", Proceedings of the Third Berkeley Symposium on Mathematical Statistics and Probability I (1956), 157–163; K. Miyasawa, "An empirical Bayes estimator of the mean of a normal population", Bulletin of the International Statistical Institute 38 (1961), 181–188; B. Efron, "Tweedie's formula and selection bias", Journal of the American Statistical Association 106 (2011), 1602–1614. Its use for denoising, with the score estimated from noisy data alone, is M. Raphan and E. P. Simoncelli, "Least squares estimation without priors or supervision", Neural Computation 23 (2011), 374–420. Its practical meaning is simple. Move the noisy point by \(\sigma^2\) times the score, then divide by the signal coefficient. The result is the best mean-square estimate of the clean point.
We call this estimate \(D_t\), the denoiser. The clean estimate, noise estimate, and score contain the same information, expressed in different units.
A best estimate need not be a possible clean sample. For the two-point law at \(-1\) and \(+1\), an observation at zero leaves both signs equally likely. Their conditional mean is zero, which minimizes squared error but belongs to neither clean state. This explains why applying a denoiser once can produce an average of incompatible possibilities. A sampler uses denoising estimates to move through intermediate distributions, updating the estimate as the observation changes.
Example 2.2.2 (Tweedie on R1, and the splitting of the law). With the score of Example 2.1.5, the denoiser of the two-point law is \(D_\sigma(x) = x + \sigma^2 s_\sigma(x) = \tanh(x/\sigma^2) = 2P(+ \mid x) - 1\): the posterior mean of a sign is the difference of its two probabilities. At \(x = 0.5\) it is \(0.193\) when \(\sigma = 1.6\) and \(0.653\) when \(\sigma = 0.8\). Near the origin \(D_\sigma(x) \approx x/\sigma^2\), a line of slope \(1/\sigma^2\); far from it \(D_\sigma\) saturates at \(\pm 1\).
Whether the noised law has one mode or two is decided by the sign of the score's slope at the origin, \(s_\sigma'(0) = (1/\sigma^2 - 1)/\sigma^2\): the density \(p_\sigma(x) \propto e^{-(x^2 + 1)/2\sigma^2} \cosh(x/\sigma^2)\) has a minimum at \(0\) and two modes at \(\pm x_\sigma\), with \(x_\sigma = \tanh(x_\sigma/\sigma^2)\), exactly when \(\sigma < 1\), and one mode at \(0\) otherwise. At \(\sigma = 0.8\) the modes are at \(\pm 0.880\), where the density is \(0.262\) against \(0.228\) at the origin. The figure draws all of this, and the flow lines of the probability-flow ODE of Section 4.3, \(dx/d\sigma = -\sigma s_\sigma(x)\), run through it: sixty-four lines started at \(\sigma = 3\), thirty-two of which end at \(+1\), fan out as they cross \(\sigma = 1\), the noise level at which the law splits. Section 10.3 calls this the speciation of the two points and gives the general criterion. ◁
left p_sigma (shaded), the denoiser D_sigma(x) = tanh(x/sigma^2) (line), the posterior variance
1 - D^2 (dashed) and the diagonal; a probe at x = 0.5
right 64 flow lines of dx/dsigma = -sigma grad log p_sigma from sigma = 3 down to 0.1, over the
density with its contours, and the flow line through the probe; the modes +-x_sigma
(dashed) split at sigma = 1
strip the floors of the denoising loss against sigma: 0.709 at 1.6, 0.307 at 0.8, 0.0013 at 0.3
(the live figure needs JavaScript)
The sensitivity of the denoiser also has a statistical meaning. If the noisy point identifies its source confidently, a small change in the observation changes the estimate little. If several sources remain plausible, the same small change can shift their relative weights and move the estimate substantially. The next identity makes that connection between sensitivity and uncertainty exact.
Proposition 2.2.3 (the posterior covariance). Under the assumptions of Theorem 2.1.4, with \(\nabla s_t\) the Jacobian of the score, the posterior covariance and the Jacobian of the denoiser are
\[\begin{aligned} \mathrm{Cov}(x_0 \mid x_t = x) &= \frac{\sigma_t^2}{\alpha_t^2}\, \bigl(I + \sigma_t^2\, \nabla s_t(x)\bigr), \\ \nabla D_t(x) &= \frac{\alpha_t}{\sigma_t^2}\, \mathrm{Cov}(x_0 \mid x_t = x). \end{aligned}\]Proof. Differentiating \(p_t(x) D_t(x) = \int x_0\, p_t(x \mid x_0)\, p_0(dx_0)\) in \(x\) and using \(\nabla_x p_t(x \mid x_0) = p_t(x \mid x_0)(\alpha_t x_0 - x)/\sigma_t^2\) gives \(\nabla(p_t D_t) = p_t\, (\alpha_t \mathbb E[x_0 x_0^\top \mid x_t] - \mathbb E[x_0 \mid x_t]\, x^\top)/\sigma_t^2\); expanding the left side as \(p_t \nabla D_t + D_t (\nabla p_t)^\top\) with \(\nabla p_t = p_t s_t = p_t (\alpha_t D_t - x)/\sigma_t^2\) and simplifying leaves \(\sigma_t^2 \nabla D_t = \alpha_t\, (\mathbb E[x_0 x_0^\top \mid x_t] - D_t D_t^\top) = \alpha_t\, \mathrm{Cov}(x_0 \mid x_t)\). The first form follows by differentiating (2.2). ∎
The Jacobian of the denoiser is a scaled posterior covariance, hence symmetric and positive semidefinite: the ideal denoiser is the gradient of a convex function, and a learned denoiser whose Jacobian is far from symmetric is far from ideal. Section 4.4 uses the trace of this Jacobian, the divergence of the velocity, to compute likelihoods. The figure below draws the ellipse at a probe; the worked example of Section 4.4 measures the same Jacobian on a trained network.
Example 2.2.4 (Tweedie on R3). The figure below draws the four fields of this section on the two moons at the noise level \(\sigma\), with \(\alpha = 1\), and as a fifth option the output \(v\) of Section 2.4: the score, the displacement \(D(x) - x = \sigma^2 s\), the noise estimate \(\hat\varepsilon = -\sigma s\), and the velocity of the linear path of Section 6 at the same signal-to-noise ratio. In the default view, the probe at \((0.90, 0.60)\) with \(\sigma = 0.5\) has a score of length \(2.16\) and is moved \(0.541\) by the denoiser, toward the near end of the upper moon; its posterior over the 24 centres has an entropy whose exponential, the effective number of centres, is \(11.5\), and no centre carries more than \(0.16\) of the weight. The Jacobian of the denoiser there has eigenvalues \(1.386\) and \(0.344\), its larger eigenvector \(19.8^\circ\) below the \(x_1\) axis, and it equals \(\mathrm{Cov}(x_0 \mid x)/\sigma^2\) to the last digit, as Proposition 2.2.3 says; the posterior covariance has standard deviations \(0.589\) and \(0.293\) along the same axes, the ellipse of the figure.
The larger eigenvalue exceeds \(1\): the denoiser is expansive along that direction, a small move of the noisy point moving the clean estimate more, which is the signature of a point between centres whose weights are comparable. Along \(\sigma\) the larger eigenvalue at this point is \(3.52\) at \(\sigma = 0.2\), peaks where the effective number of centres is near four, and falls to \(0.562\) at \(\sigma = 1\). Clicking elsewhere shows the rule of the field: far from both moons the posterior is spread over a whole moon or both and the denoiser points at their centroid; close to a moon a single centre takes the weight and the arrow is short and local. ◁
the two moons R3 (24 centres) and, over the exact density p_sigma, the chosen field at noise level sigma: the score, the Tweedie displacement D(x) - x, the noise estimate, the output v = eps-hat - sigma D, or the velocity, as arrows or as streamlines; at the probe (0.90, 0.60) its streamline along the field, the posterior over the centres, the arrow to D(x) and the posterior covariance as an ellipse; right: the eigenvalues of the denoiser's Jacobian against sigma (the live figure needs JavaScript)
Remark 2.2.5 (takeaway: the denoiser averages, it does not project). \(D_t(x)\) is a conditional mean. When the posterior given \(x\) spreads over several modes of the data, \(D_t(x)\) lies between them, at a point of low density: on R1 at \(\sigma = 1.6\) the denoiser of \(x = 0.5\) is \(0.193\), a point where \(p_0\) has no mass at all, and on R3 the arrows from the gap between the moons land in the gap. A single application of the ideal denoiser returns the mean of what the data could have been, which is not a data point; sampling has to be iterative because each step may only remove the noise whose posterior is unimodal. The figures of Sections 3 to 5 show how small that noise is per step, and Section 8 shows what it takes to learn the whole path at once.
Example 2.2.6 (the denoiser averages, in practice). We train a one-step denoiser on FashionMNIST at \(28 \times 28\). It is a U-Net of widths \((16, 32, 64)\) with the architecture of Section 11, trained for six epochs, \(2{,}808\) steps of 128 images by Adam at \(10^{-3}\), to return the clean image from \(0.5\, x_0 + 0.5\, z\) at that single noise level, with no chain. We compare its answer at sixteen fixed draws of pure noise and at sixteen held-out images corrupted the same way, against the exact posterior mean over the 60,000 training images at the same inputs.
Fed pure noise, the network returns grey blots, and the figure shows what the blots are. The exact minimizer of the same objective is the posterior mean over the training set, \(\sum_i r_i(u)\, x^{(i)}\) with \(r_i(u) \propto \exp(-\|u - 0.5\, x^{(i)}\|^2/0.5)\): at a pure-noise input it concentrates on about \(1.0\) training images (the exponential of the posterior's entropy, median over the sixteen inputs) with a top weight of \(0.994\), so the minimizer returns one training image, while the network returns a blot of standard deviation \(0.493\) against the data's \(0.671\), at a root-mean-square distance of \(0.537\) from the exact answer.
On held-out images corrupted as in training the network's error is \(0.0404\) per pixel against the exact posterior mean's \(0.0538\), and the posterior there averages \(2.1\) images: in \(784\) dimensions the posterior over the training set collapses onto one or two training images, so the exact minimizer returns a near neighbour of the corrupted image rather than an average of many, and on an image it has never seen it does worse than the network. Here the trained network denoises held-out images better than the exact empirical minimizer. Section 10.3 asks what prevents it from becoming a lookup table of the training set. ◁
sixteen draws of pure noise (row 1), the one-step denoiser applied to them (row 2: grey blots), the exact posterior mean over the 60,000 training images for the same inputs (row 3: training images), and the difference network - exact as a change map, full at +-1.39 (row 4); in the second view sixteen held-out images, their corruption 0.5 x_0 + 0.5 z, the network's answer, the exact one and their difference, full at +-0.72; in the third the denoiser iterated four times (the live figure needs JavaScript and loads about 0.21 MB of arrays)
Score matching and its denoising form
We now have a target, the score, and need a loss we can evaluate from samples. Hyvärinen's score matching removes the unknown density from the loss by integration by parts. Vincent's denoising version goes further: once we add Gaussian noise ourselves, ordinary regression is enough.
Definition 2.3.1 (explicit score matching). For a family of vector fields \(s_\theta\) and a density \(p\), the explicit score-matching objective is \(J_{\mathrm{E}}(\theta) = \tfrac12\, \mathbb E_{x \sim p}\, \|s_\theta(x) - \nabla \log p(x)\|^2\).
Theorem 2.3.2 (Hyvärinen). If \(p\) and \(s_\theta\) are differentiable and \(p(x)\, s_\theta(x) \to 0\) as \(\|x\| \to \infty\), then
\[J_{\mathrm{E}}(\theta) \;=\; \mathbb E_{x \sim p}\Bigl[\tfrac12\, \|s_\theta(x)\|^2 + \nabla \cdot s_\theta(x)\Bigr] + \text{a constant independent of } \theta .\]Proof. Expand the square: \(\tfrac12\|s_\theta\|^2 - s_\theta \cdot \nabla \log p + \tfrac12 \|\nabla \log p\|^2\). The last term is the constant. In the middle term, \(\int p\, s_\theta \cdot \nabla \log p\, dx = \int s_\theta \cdot \nabla p\, dx = -\int p\, \nabla \cdot s_\theta\, dx\) by integration by parts, the boundary term vanishing by hypothesis. ∎A. Hyvärinen, "Estimation of non-normalized statistical models by score matching", Journal of Machine Learning Research 6 (2005), 695–709, Theorem 1.
The theorem removes the unknown score from the objective at the price of a divergence, which for a neural network in \(d\) dimensions costs \(d\) backward passes to evaluate exactly, or a random projection to estimate.Y. Song, S. Garg, J. Shi and S. Ermon, "Sliced score matching: a scalable approach to density and score estimation", UAI 2019 (arXiv 1905.07088), for the projected version; the Hutchinson estimator of the divergence is used again in Section 4.4.
Remark 2.3.3 (warning: the objective needs noise). Hyvärinen's objective is stated at a single density. Applied to the data law itself, it meets two difficulties that Song and Ermon identified and that the forward process cures. If the data lie on a lower-dimensional set, the score of \(p_0\) is undefined off it and the objective is ill-posed; and where \(p_0\) is small the objective gives no signal, so the score between modes, which a sampler must cross, is never learned. Adding Gaussian noise at a level \(\sigma\) gives a density everywhere, and a whole range of levels gives a usable score at every scale; the sampler of Section 4 then needs exactly the family \(s_t\), \(t \in (0, 1]\), and nothing at \(t = 0\).Y. Song and S. Ermon, "Generative modeling by estimating gradients of the data distribution", NeurIPS 2019 (arXiv 1907.05600), Section 3, for the two difficulties and the noise-conditional score network; the annealed Langevin sampler of that paper is the ancestor of the reverse SDE.
Theorem 2.3.4 (Vincent, denoising score matching). Fix \(t\) with \(\sigma_t > 0\) and let \(x_t = \alpha_t x_0 + \sigma_t z\) as in Theorem 2.1.4. Then for every vector field \(s_\theta\) with \(\mathbb E\|s_\theta(x_t)\|^2 < \infty\),
\[\tfrac12\, \mathbb E_{x_0, z}\Bigl\|s_\theta(x_t) + \frac{z}{\sigma_t}\Bigr\|^2 \;=\; \tfrac12\, \mathbb E_{x_t \sim p_t}\bigl\|s_\theta(x_t) - s_t(x_t)\bigr\|^2 + C_t, \tag{2.3}\]where \(C_t = \tfrac12\mathbb E\|s_t(x_t) + z/\sigma_t\|^2\) does not depend on \(\theta\).
Proof. Write \(s_\theta + z/\sigma_t = (s_\theta - s_t) + (s_t + z/\sigma_t)\) and expand. The cross term is \(\mathbb E[(s_\theta(x_t) - s_t(x_t)) \cdot (s_t(x_t) + z/\sigma_t)]\). Conditioning on \(x_t\), the first factor is fixed and the second has conditional mean \(s_t(x_t) + \mathbb E[z \mid x_t]/\sigma_t = 0\) by (2.1). So the cross term vanishes and the two remaining squares are the two sides. ∎P. Vincent, "A connection between score matching and denoising autoencoders", Neural Computation 23 (2011), 1661–1674, where the identity is proved for a general smoothing kernel. Theorem 2.1.4 is the mechanism, and the proof above is shorter than the original because it uses it.
The left side of (2.3) gives a training procedure: draw \(x_0\) and \(z\), form \(x_t\), and train \(s_\theta(x_t)\) to predict \(-z/\sigma_t\). Every quantity in that loss is available. We need neither a density evaluation nor a divergence calculation.
The right side explains why the procedure works. Its minimizer is the true score, and any loss above the constant floor is exactly the model's squared score error under \(p_t\). Section 2.4 shows how changing the output units and weights yields the familiar diffusion objectives.
Two different effects keep a training loss above zero. The network can be inaccurate, and the target can remain uncertain even when the network is optimal. Equation (2.3) separates them. Improving the model reduces the first part; the second is set by the data and the noise level. Comparing raw losses across noise levels without this distinction can make a good predictor look bad merely because its task is more ambiguous.
left against sigma from 0.1 to 3, the loss of the best linear denoiser a x of R1, a = 1/(1 + sigma^2), in
the units of the clean point: the floor sigma^4 (1/sigma^2 - I(p_sigma)) as a wash and the excess
above it, 0.083 = 0.0013 + 0.081 at sigma = 0.3, 0.390 = 0.307 + 0.084 at 0.8 and 0.719 = 0.709 +
0.010 at 1.6 (in the units of the noise 0.917 = 0.015 + 0.903, 0.610 = 0.479 + 0.130 and
0.281 = 0.277 + 0.004); the excess peaks at 0.131 near sigma = 0.50
right at sigma = 0.8, p_sigma, the exact denoiser tanh(x/sigma^2), the line a x and the local excess
p_sigma (a x - D_sigma)^2, which peaks at 0.0357 near |x| = 0.69; p_sigma puts 0.236 of its mass
in |x| < 0.5
(the live figure needs JavaScript)
Proposition 2.3.5 (the loss floor). The constant in Theorem 2.3.4 is \(C_t = \tfrac12\bigl(d/\sigma_t^2 - I(p_t)\bigr)\), where \(I(p_t)\) is the Fisher information of Definition 1.4.1.
Proof. \(\mathbb E\|s_t + z/\sigma_t\|^2 = \mathbb E\|s_t\|^2 + 2\, \mathbb E[s_t(x_t) \cdot z]/\sigma_t + \mathbb E\|z\|^2/\sigma_t^2\). The first term is \(I(p_t)\), the last is \(d/\sigma_t^2\), and in the middle, conditioning on \(x_t\), \(\mathbb E[s_t(x_t) \cdot z] = \mathbb E[s_t(x_t) \cdot \mathbb E[z \mid x_t]] = -\sigma_t\, I(p_t)\) by (2.1). ∎
The floor is positive because \(I(p_t) \le d/\sigma_t^2\), with equality only when \(p_t\) is the kernel of a point mass; it is the irreducible error of guessing the noise from the noisy point, and it is largest, in units of the data, at the noise levels where the posterior is most spread.
Lemma 2.3.6 (least squares computes conditional expectations). For square-integrable \(Y\) and \(X\) and the class of all measurable functions \(h\), \(\mathbb E\|h(X) - Y\|^2\) is minimized by \(h(x) = \mathbb E[Y \mid X = x]\), and \(\mathbb E\|h(X) - Y\|^2 = \mathbb E\|h(X) - \mathbb E[Y \mid X]\|^2 + \mathbb E\|Y - \mathbb E[Y \mid X]\|^2\) for every \(h\).
Proof. The cross term \(\mathbb E[(h(X) - \mathbb E[Y \mid X]) \cdot (\mathbb E[Y \mid X] - Y)]\) vanishes by conditioning on \(X\). ∎
The lemma applies to any regression target. Predict the clean point \(x_0\) from \(x_t\), and the ideal answer is the denoiser \(D_t\). Predict the noise \(z\), and it is \(\hat\varepsilon_t\). Predict \(\dot\alpha_t x_0 + \dot\sigma_t z\), and it is the velocity \(u_t\) of Section 4.3. Tweedie's formula connects these answers to the score. The algorithm below uses the clean point as its target.
Algorithm 2.3.7 DENOISING-REGRESSION
(training; the clean point as the target)
Picture: one training example. A data point x_0 is
pushed to x_t by the kernel, the network sees x_t
and t, and is pulled toward x_0 by the loss.
x_0 --(alpha_t x_0 + sigma_t z)--> x_t
^ |
| loss w(t) |D_theta(x_t, t) - x_0|^2
'----------------------------------'
D_theta
Input samples of p_0; a schedule (alpha_t, sigma_t); a
weight w(t) > 0; a law for t on (0, 1]; a network
D_theta(x, t) in R^d.
Output theta, with D_theta ~ D_t = E[x_0 | x_t] at every t.
1. repeat until converged:
2. draw x_0 from the data, t from its law, z ~ N(0, I)
3. x_t <- alpha_t x_0 + sigma_t z
4. L <- w(t) |D_theta(x_t, t) - x_0|^2
5. take a gradient step on L in theta
Example 2.3.8 (the floor on R1). On the two-point law under a variance-exploding process, Proposition 2.3.5 gives the irreducible error of the noise estimate, \(1/\sigma^2 - I(p_\sigma)\), and multiplying by \(\sigma^4\) gives the irreducible error of the denoiser, \(\mathbb E\|D_\sigma(x_\sigma) - x_0\|^2\), which is the posterior variance of the sign averaged over \(x_\sigma\). Numerically, \(I(p_\sigma) = 0.282\) at \(\sigma = 1.6\), so the denoiser's floor is \(0.709\), close to the variance \(1\) of a sign that is barely known; at \(\sigma = 0.8\) the floor is \(0.307\), and at \(\sigma = 0.3\) it is \(0.0013\), the sign being decided. A network trained with Algorithm 2.3.7 on this law cannot do better than these numbers, and a loss curve that reaches them has learned the score exactly. The strip under the first figure of Section 2.2 draws these floors against \(\sigma\), the first figure of Section 2.3 separates one model's loss into this floor and its excess, and the second figure of Section 2.3 trains a network on R3 against the same floor. ◁
The program below prints the posteriors, denoisers and scores of Examples 2.1.5 and 2.2.2 and checks Tweedie's formula on them digit by digit.
# Tweedie's formula on R1, the two points at -1 and +1, at noise level sigma.
#
# p_sigma(x) = (1/2) N(x; -1, sigma^2) + (1/2) N(x; +1, sigma^2)
# P(+ | x) = 1 / (1 + exp(-2x / sigma^2))
# D(x) = E[x_0 | x_t = x] = 2 P(+ | x) - 1 = tanh(x / sigma^2)
# score = (tanh(x / sigma^2) - x) / sigma^2 = (D(x) - x) / sigma^2
#
# Picture: at sigma = 0.8 the law has two modes and D is a
# step; at sigma = 1.6 it has one and D is nearly a line.
#
# D(x) +1 ..............----- sigma = 0.8
# 0 ---'
# -1 -----'
# -2 -1 0 1 2 x
import math
def posterior_plus(x, s): return 1 / (1 + math.exp(-2 * x / (s * s)))
def denoiser(x, s): return math.tanh(x / (s * s))
def score(x, s): return (math.tanh(x / (s * s)) - x) / (s * s)
def density(x, s):
n = lambda m: math.exp(-0.5 * (x - m) ** 2 / (s * s)) / (s * math.sqrt(2 * math.pi))
return 0.5 * (n(-1) + n(1))
for s in (1.6, 0.8):
print(f"sigma = {s}")
for x in (0.5, 1.5):
D, sc = denoiser(x, s), score(x, s)
print(f" x = {x}: P(+|x) = {posterior_plus(x, s):.4f} D(x) = {D:.4f} "
f"x + sigma^2 score = {x + s * s * sc:.4f} score = {sc:.4f}")
# the modes of p_sigma: zeros of the score, by bisection on x > 0
lo, hi = 1e-6, 3.0
if score(lo, s) * score(hi, s) < 0:
for _ in range(60):
mid = 0.5 * (lo + hi)
lo, hi = (lo, mid) if score(lo, s) * score(mid, s) < 0 else (mid, hi)
print(f" two modes, at x = +-{0.5 * (lo + hi):.4f}; p at the mode {density(0.5 * (lo + hi), s):.4f}, at 0 {density(0, s):.4f}")
else:
print(f" one mode, at 0; slope of D at 0 = 1/sigma^2 = {1 / (s * s):.4f}")
sigma = 1.6
x = 0.5: P(+|x) = 0.5964 D(x) = 0.1929 x + sigma^2 score = 0.1929 score = -0.1200
x = 1.5: P(+|x) = 0.7635 D(x) = 0.5270 x + sigma^2 score = 0.5270 score = -0.3801
one mode, at 0; slope of D at 0 = 1/sigma^2 = 0.3906
sigma = 0.8
x = 0.5: P(+|x) = 0.8267 D(x) = 0.6534 x + sigma^2 score = 0.6534 score = 0.2397
x = 1.5: P(+|x) = 0.9909 D(x) = 0.9817 x + sigma^2 score = 0.9817 score = -0.8098
two modes, at x = +-0.8797; p at the mode 0.2623, at 0 0.2283
Example 2.3.9 (a denoiser trained in the page). This example trains on R3, the two moons, with a noise level \(\sigma\) drawn log-uniform on \([0.02, 3]\). A multilayer perceptron has layers \(3 \to 64 \to 64 \to 2\) and SiLU activations, with 4,546 parameters. It uses the preconditioning of Remark 2.4.3 and trains in the page by Algorithm 2.3.7 with Adam at rate \(3 \times 10^{-3}\) on batches of 256, on fresh draws of the law or on \(n\) fixed training points, for 3,000 steps from fixed draws. At every step, we measure its distance to the exact denoiser.
With fresh samples at every batch, the law itself as training data, the held-out loss in the units of the noise after 3,000 steps is \(0.65\) against the floor \(0.61\) of Proposition 2.3.5 averaged over the training law of \(\sigma\), an excess of \(0.044\), and the error view on the left shows where it lives: away from the moons, growing with the distance from them.
With 64 training points the training loss after 3,000 steps is \(0.64\), still above the \(0.60\) that the exact denoiser of R3 scores on the same points, the held-out loss is \(0.66\), and \(0.029\) of 512 samples land within 0.02 of a training point, as do \(0.027\) of fresh draws of R3: after a minute of training the network is still a denoiser of the law, not of its sample, and its score stays closer to the population score than to the empirical score of the 64 points.
The onset of memorization, the first step from which the training loss stays below the exact denoiser's for 200 steps, comes at about step 11,000, and the third figure of Section 10.3 continues the run to it. The numbers are given to two digits because browsers differ in the last bits of their exponentials. ◁
left the field of a 4,546-parameter network trained in the page on R3 (arrows) over the exact
field at the chosen sigma; or its error over the plane; or 512 of its samples
right its held-out loss by noise level against the floor of Proposition 2.3.5 (Monte Carlo
with the exact score), and the loss, the excess over the floor and, for n training
points, the distances to the population and to the empirical score and the loss of the
exact denoiser on the points, against the training step
(the live figure needs JavaScript; it trains in a Web Worker)
Four parameterizations and their weights
A network can predict the clean point, the noise, the score, or the velocity. We can convert any one prediction into the others, but that does not make their training losses identical. The conversion changes how much each noise level contributes to the loss. Those weights help determine where a finite network spends its capacity.
Proposition 2.4.1 (the dictionary). At a point \(x\) and time \(t\) with \(\sigma_t > 0\), the following are affine functions of one another, and a network that outputs any one of them determines the others:
\[\begin{aligned} D_t(x) &= \frac{x - \sigma_t \hat\varepsilon_t(x)}{\alpha_t}, \qquad \hat\varepsilon_t(x) = -\sigma_t s_t(x), \\ v_t(x) &:= \alpha_t\hat\varepsilon_t(x) - \sigma_t D_t(x), \\ u_t(x) &= \dot\alpha_t D_t(x) + \dot\sigma_t \hat\varepsilon_t(x) \\ &= \frac{\dot\alpha_t}{\alpha_t}\, x + \Bigl(\dot\sigma_t - \frac{\dot\alpha_t \sigma_t}{\alpha_t}\Bigr)\hat\varepsilon_t(x) = f_t\, x - \tfrac12 g_t^2\, s_t(x). \end{aligned}\]The quantity \(v_t\) is the velocity parameterization of Salimans and Ho; \(u_t\) is the probability-flow velocity of Section 4.3, and the last equality is Proposition 1.1.5's \(g_t^2 = 2\sigma_t(\dot\sigma_t - \sigma_t\dot\alpha_t/\alpha_t)\).
Proof. The first line is (2.2) and a definition. For the second, substitute \(D_t\) from the first line into \(\dot\alpha_t D_t + \dot\sigma_t\hat\varepsilon_t\); then substitute \(\hat\varepsilon_t = -\sigma_t s_t\). ∎
Changing the output changes the units in which we charge for a mistake. From \(\hat\varepsilon_t = (x-\alpha_t D_t)/\sigma_t\), a clean-estimate error is multiplied by \(\alpha_t/\sigma_t\) when expressed as a noise-estimate error. That factor changes with time. Equal weight on noise errors therefore assigns unequal weight to the same-sized clean-estimate errors at different noise levels. The following formulas keep track of that rescaling, and the first figure of Section 2.3 shows it on one model: switch its units and the floor and the excess rescale together.
Proposition 2.4.2 (the weights). For a network with output \(D_\theta\) and the corresponding \(\varepsilon_\theta = (x - \alpha_t D_\theta)/\sigma_t\), \(v_\theta = \alpha_t \varepsilon_\theta - \sigma_t D_\theta\) and \(u_\theta = \dot\alpha_t D_\theta + \dot\sigma_t \varepsilon_\theta\), the squared errors against the targets \(x_0\), \(z\), \(v = \alpha_t z - \sigma_t x_0\) and \(\dot x_t = \dot\alpha_t x_0 + \dot\sigma_t z\) are proportional:
\[\begin{aligned} \|\varepsilon_\theta - z\|^2 &= \frac{\alpha_t^2}{\sigma_t^2}\|D_\theta - x_0\|^2, \\ \|v_\theta - v\|^2 &= \frac{1}{\sigma_t^2}\|D_\theta - x_0\|^2, \\ \|u_\theta - \dot x_t\|^2 &= \Bigl(\frac{\alpha_t g_t^2}{2\sigma_t^2}\Bigr)^2\|D_\theta - x_0\|^2 . \end{aligned}\]Proof. Each target and each output is the same affine function of \((x_0, D_\theta)\) respectively, with \(x_t\) fixed: \(\varepsilon_\theta - z = -(\alpha_t/\sigma_t)(D_\theta - x_0)\);
\[\begin{aligned} v_\theta - v &= \alpha_t(\varepsilon_\theta - z) - \sigma_t(D_\theta - x_0) = -\Bigl(\frac{\alpha_t^2}{\sigma_t} + \sigma_t\Bigr)(D_\theta - x_0) \\ &= -\frac{\alpha_t^2 + \sigma_t^2}{\sigma_t}(D_\theta - x_0), \end{aligned}\]which is \(-(D_\theta - x_0)/\sigma_t\) for a variance-preserving schedule, the case in which \(v\) is used; and \(u_\theta - \dot x_t = (\dot\alpha_t - \dot\sigma_t\alpha_t/\sigma_t)(D_\theta - x_0) = -\frac{\alpha_t g_t^2}{2\sigma_t^2}(D_\theta - x_0)\), since \(g_t^2/(2\sigma_t^2) = \dot\sigma_t/\sigma_t - \dot\alpha_t/\alpha_t\) by the expression for \(g_t^2\). ∎
So the loss of DDPM, which regresses onto the noise with unit weight, is the denoiser loss weighted by the signal-to-noise ratio \(\alpha_t^2/\sigma_t^2\); the velocity loss of flow matching on the linear path, where \(\alpha_t g_t^2/(2\sigma_t^2) = 1/t\), is the denoiser loss weighted by \(1/t^2\), largest at the data end; and the \(v\)-prediction loss on a variance-preserving schedule is the denoiser loss weighted by \(1/\sigma_t^2\). None of these weights changes the minimizer, which is \(D_t\) in every case by Lemma 2.3.6.
What they change is where, along the noise axis, the finite capacity of a network and the finite length of training are spent, and Kingma and Gao showed what the choice means in likelihood terms: every weighted denoising objective in use is a weighted integral of evidence lower bounds over noise levels, and when the weight is monotone in the noise level it is exactly an evidence lower bound for data augmented by Gaussian noise.T. Salimans and J. Ho, "Progressive distillation for fast sampling of diffusion models", ICLR 2022 (arXiv 2202.00512), Section 4 and Appendix D, for \(v \equiv \alpha_t\varepsilon - \sigma_t x\), "the most stable option" for distillation; D. P. Kingma and R. Gao, "Understanding diffusion objectives as the ELBO with simple data augmentation", NeurIPS 2023 (arXiv 2303.00848): "all commonly used diffusion model objectives equate to a weighted integral of ELBOs over different noise levels, where the weighting depends on the specific objective used. Under the condition of monotonic weighting, the connection is even closer: the diffusion objective then equals the ELBO, combined with simple data augmentation, namely Gaussian noise perturbation." Section 3.4 states the bound. The figure of Section 3.1, in its third view, draws these weights on the \(\lambda\) axis, and the worked example of Section 3.3 shows the loss of the \(\varepsilon\)-model by noise level in both units.
Remark 2.4.3 (trick: conditioning the regression). A shared network receives inputs of very different sizes across noise levels. It also faces different tasks: near the clean end, much of the answer can pass through unchanged, while at large noise it must infer the clean point from weak evidence. Converting predictions can magnify errors too: in \(D_\theta=(x-\sigma_t\varepsilon_\theta)/\alpha_t\), a noise-prediction error is multiplied by \(\sigma_t/\alpha_t\), which is large near the noise end.
Karras et al. use the variance-exploding coordinates, \(\alpha=1\), and separate a direct contribution from the input from a learned correction:
\[D_\theta(x;\sigma)=c_{\mathrm{skip}}(\sigma)\,x+c_{\mathrm{out}}(\sigma)\,F_\theta(c_{\mathrm{in}}(\sigma)\,x;c_{\mathrm{noise}}(\sigma)).\]The coefficients are \(c_{\mathrm{skip}}=\sigma_{\mathrm{data}}^2/(\sigma^2+\sigma_{\mathrm{data}}^2)\), \(c_{\mathrm{out}}=\sigma\sigma_{\mathrm{data}}/\sqrt{\sigma^2+\sigma_{\mathrm{data}}^2}\), \(c_{\mathrm{in}}=1/\sqrt{\sigma^2+\sigma_{\mathrm{data}}^2}\) and \(c_{\mathrm{noise}}=\tfrac14\ln\sigma\). The direct contribution approaches the identity at small noise and vanishes at large noise. The input and output scales keep the regression comparable across noise levels, using \(\sigma_{\mathrm{data}}\) to set the data scale. The loss weight \(1/c_{\mathrm{out}}^2\) measures error in the network's output units; the training noise levels are drawn with \(\ln\sigma\sim\mathcal N(-1.2,1.2^2)\).Karras, Aittala, Aila and Laine, "Elucidating the design space of diffusion-based generative models" (2022), Section 5 and Table 1; \(\sigma_{\mathrm{data}} = 0.5\) for their image data. At \(\sigma \to 0\), \(c_{\mathrm{skip}} \to 1\) and the network predicts the residual; at \(\sigma \to \infty\), \(c_{\mathrm{skip}} \to 0\) and it predicts the clean point scaled. The opening figure of the introduction runs exactly this parameterization with the exact denoiser in place of \(F_\theta\). The second figure of Section 2.3 trains its network with exactly this parameterization, and the readout of the first figure of that section prints \(c_{\mathrm{skip}}\) and \(c_{\mathrm{out}}^2\) at its noise level.
To read the training figure, first locate the irreducible floor, then inspect the excess loss by noise level. A falling average can hide a poorly learned interval that the sampler must traverse. Finally compare training and held-out examples: they answer different questions about which distribution the network has learned.
What the regression converges to
The ideal regression in Algorithm 2.3.7 averages over the data law. A real training set gives us only \(n\) points, so the objective instead averages over the empirical law \(\hat p_0 = \tfrac1n\sum_i \delta_{x^{(i)}}\). Lemma 2.3.6 still applies. The exact minimizer is the empirical denoiser, \(\hat D_t(x) = \sum_i r_i(x)\, x^{(i)}\) with \(r_i(x) \propto \mathcal N(x; \alpha_t x^{(i)}, \sigma_t^2 I)\): a weighted lookup over the training set.
Example 2.1.6 computes this denoiser by setting \(s_k = 0\). Section 10.3 runs its reverse flow and shows the consequence: samples end at training points. Perfectly minimizing the empirical denoising objective leads to a memorizing model.
The worked examples on \(L\) compare that ideal empirical solution with trained networks. Section 2.2 compares their outputs at one noise level, Section 3.3 measures held-out error across noise levels, and Section 4.4 compares their Jacobians. Where the empirical posterior has collapsed onto one training image, the network can still behave differently.
Section 10.3 asks why a trained network can avoid that memorizing solution. Two features of the problem help frame the question.
First, learning a smooth score from \(n\) samples at a fixed noise level is a nonparametric estimation problem. For the Lipschitz-score class discussed later, the minimax squared \(L^2\) error is of order \(n^{-2/(d+4)}\). Without additional structure, high dimension makes that rate impractically slow.A. Wibisono, Y. Wu and K. Y. Yang, "Optimal score estimation via empirical Bayes smoothing", COLT 2024 (arXiv 2402.07747): "we establish the optimal rate of \(\tilde\Theta(n^{-2/(d+4)})\) for this estimation problem under the loss function \(\|\hat s - s^*\|^2_{L^2(\rho^*)}\)", attained by a regularized score estimator based on a Gaussian kernel. Section 10.2 collects the estimation theory.
Second, noise hides the differences between a population and a finite sample of it. At scales larger than the spacing of the training points, the empirical and population scores can be close. As the noise falls, individual observations become distinguishable and the empirical sampler can commit to one of them. Section 10.3 measures where that happens as the training set grows.
The network's inductive bias matters in this low-noise region: it determines how the learned function departs from the empirical denoiser. Section 10.3 gives a partial account for convolutional architectures; a general explanation remains open.