Denoising Score Matching and the Noise-Prediction Objective

A likelihood model must normalise. The density p(x) is proportional to an unnormalised function, and the constant that makes it integrate to one is an integral over the whole space, which for images is a space no quadrature can cover. Score-based modelling avoids the constant by learning a different object: the gradient of the log density with respect to the data, ∇_x log p(x). That gradient does not depend on the normalising constant, because the constant is additive in the log and its gradient is zero.

This is the reason a diffusion model’s training loss can be a plain squared error against a known target with no likelihood term anywhere. The target is the noise that was added, and the claim connecting the two is that predicting the added noise is equivalent to estimating the score of the noisy data distribution. The equivalence is exact under the right scaling, and it is worth deriving because the scaling constants are where implementations go wrong.

Why the score sidesteps the partition function

Write p(x) = f(x) / Z with Z unknown. Then ∇_x log p(x) = ∇_x log f(x), and Z has vanished. A model s_θ(x) trained to approximate this vector field can generate samples by following it, either through Langevin dynamics or through the reverse of a diffusion process, without ever evaluating a density.

The catch is that the score of the data distribution is not directly available as a training target, because the data distribution is known only through samples. Fitting s_θ to ∇_x log p(x) requires an objective that can be estimated from samples alone. Hyvärinen’s score matching gives one: integrate by parts, and the objective becomes an expectation over the data of a combination of the model’s output and its Jacobian, with no reference to the true score. That objective is exact but requires the full Jacobian of the model’s output with respect to its input, which for a network with d outputs costs d backward passes per example. Denoising score matching replaces it with something cheaper and lower variance.

The denoising identity

Corrupt a clean sample x₀ with Gaussian noise of known scale σ:

x = x₀ + σ ε,    ε ~ N(0, I)

The conditional density of x given x₀ is Gaussian, so its score with respect to x is available in closed form:

∇_x log q(x | x₀) = −(x − x₀) / σ² = −ε / σ

Denoising score matching states that matching the model to the conditional score, averaged over the corruption process, has the same minimiser as matching it to the score of the marginal. The objective is

E_{x₀, ε} [ ‖ s_θ(x₀ + σ ε) + ε / σ ‖² ]

and its minimiser over all functions is the score of the noisy marginal, ∇_x log q(x). The proof is an integration by parts plus the observation that the conditional expectation of the marginal score given x equals the conditional score; the important part for implementation is that the target is computed from the noise draw, not from any density.

The noise-prediction form follows by substitution. If the network predicts ε rather than the score, define ε_θ(x) = −σ s_θ(x). Then the objective becomes

E [ ‖ ε_θ(x₀ + σ ε) − ε ‖² ]

which is a regression on the noise that was added, weighted by nothing. The two forms are the same objective up to a constant factor of σ², which matters when the loss is summed over timesteps with different σ: without the weighting, low-noise steps dominate the gradient because their target has unit scale and their error is measured in the same units as the signal.

What the network is actually being asked to compute

The minimiser of the squared error is the conditional expectation of the target given the input. For the noise-prediction objective, that is E[ε | x_t], the expected noise that produced the observed noisy sample. Because the corruption is x_t = √(A_t) x₀ + √(1 − A_t) ε, the posterior mean of ε given x_t is a linear function of the posterior mean of x₀. Predicting noise and predicting the clean sample are therefore equivalent parameterisations of the same estimate, related by

E[x₀ | x_t] = ( x_t − √(1 − A_t) · E[ε | x_t] ) / √(A_t)

This is why the choice between noise prediction, clean-sample prediction, and score prediction is mostly a matter of conditioning rather than of modelling. The three targets differ in their scale across timesteps: near t = 0 the clean sample is almost determined by the input, so its prediction has small error and a well-conditioned target; near t = T the clean sample is nearly independent of the input, so its target has the full data variance and the loss is dominated by irreducible uncertainty. The noise target has the opposite profile: it is nearly determined near t = T, where the input is almost pure noise, and uncertain near t = 0. Training on the noise target therefore puts more weight on the high-noise region, which is where the reverse process needs the most guidance to get the global structure right.

Choosing the weighting across noise levels

The unweighted sum over timesteps is not the only option, and the choice of weighting changes what the model learns. The likelihood bound derived from the diffusion process gives a specific weighting per timestep, which is small at both extremes and largest in the middle of the chain. Training with that weighting optimises the bound; training with uniform weighting optimises a different objective that tends to produce better perceptual samples at the cost of a worse likelihood.

The reason is that the bound’s weighting is derived from the KL between the reverse step and its true posterior, and that KL is small when the reverse step is nearly deterministic, which happens at both ends of the chain. Uniform weighting spends equal capacity on the ends, where the model’s task is either trivial or hopeless, and less on the middle, where the structure of the data is actually being decided. Neither is wrong; they answer different questions.

A practical middle ground rescales the loss so that its magnitude is comparable across timesteps. Since the noise target has unit variance, the raw error is already comparable, but the model’s output magnitude is not: near t = T the optimal output is close to the input’s noise, while near t = 0 the optimal output is close to zero. Networks handle this by conditioning on the timestep, usually through a sinusoidal embedding of t added to the activations, which lets one set of weights express very different functions at different noise levels.

import numpy as np

def train_step(model, x0, alpha_bar, rng):
    t = rng.integers(0, len(alpha_bar))
    eps = rng.standard_normal(x0.shape)
    a = alpha_bar[t]
    x_t = np.sqrt(a) * x0 + np.sqrt(1.0 - a) * eps
    # the target is the noise itself; no density, no normalising constant
    pred = model(x_t, t)
    loss = np.mean((pred - eps) ** 2)
    return loss

def predict_x0(x_t, eps_pred, a):
    # invert the corruption using the predicted noise
    return (x_t - np.sqrt(1.0 - a) * eps_pred) / np.sqrt(a)

The second function is the whole of the sampling arithmetic at one step. It also shows the numerical hazard: as a approaches zero, the division by √a amplifies any error in the noise prediction. At the last timesteps of a long chain, the estimate of the clean sample is a difference of two nearly equal large quantities, and the cancellation is severe. This is the reason samplers clamp the predicted clean sample to a valid range at high noise, and the reason the final steps are usually taken with reduced stochasticity.

The equivalence is exact only for the marginal

The statement that the denoising objective’s minimiser is the score of the noisy marginal is exact, but the marginal is the data convolved with Gaussian noise at that scale. It is not the score of the clean data distribution. For a single fixed noise level the difference is a bias: the model learns a smoothed version of the data distribution, and sampling from it produces samples from that smoothed distribution, not from the original.

The multi-scale construction fixes this by training at a range of noise levels and annealing from large to small. At large noise the smoothed distribution has usable gradients almost everywhere, because the smoothing fills in the low-density regions where the clean score is undefined or enormous. At small noise the gradients become sharp and the samples converge to the data. The sequence of noise levels is the same schedule that appears in the forward process, and its design is covered in signal-to-noise ratio and schedule shape.

The bias has a visible symptom: sampling at a fixed moderate noise level produces images that are recognisable but lack fine detail, and the detail returns as the final noise levels are annealed. This is not a defect of the training; it is the smoothing doing exactly what it was asked to do, and the annealed sampler is the mechanism that removes it.

Failure modes in the objective

Three things break the equivalence in practice.

If the noise scale used at training differs from the scale assumed at sampling, the model’s score estimate is off by a constant factor and the sampler’s step sizes are wrong. The symptom is a sampler that either barely moves or overshoots into saturation, and the cause is usually a mismatch between the schedule table and the timestep indexing used by the sampler.

If the data are discretised and the model is trained on continuous values, the score near the data manifold is a sum of delta functions rather than a smooth field, and the denoising target at very small noise is dominated by the discretisation rather than by the data. Training at noise levels below the quantisation step therefore teaches the model to reproduce the discretisation pattern, which is why the smallest noise level is usually set above the quantisation scale.

If the network’s output is not scaled to match the noise level, the loss at high noise is large and the gradient dominates. Predicting a unit-variance target at every timestep is the normalisation that keeps the loss balanced; a parameterisation that predicts the clean sample directly has a target whose variance grows with the noise level, and the resulting loss profile is dominated by the noisiest steps.

The relationship between this objective and the reverse-time process it is used to drive is developed in the reverse SDE and probability-flow ODE, where the score appears as the drift correction that makes the reverse dynamics well defined.