Noise-Conditional Score Networks and Annealed Langevin Dynamics

Dive into Deep Learning · §17.6

Noise-conditional score networks and annealed Langevin dynamics
a ladder of noise levels · one network for all of them · sample coarse to fine

A Ladder of Noise Levels Removes the Trade-off

Coarse levels have no valleys and reliable scores wherever the noisy data reach; fine levels match the data. Use all of them, in order.

One Network, Conditioned on the Level

\mathcal{L}(\boldsymbol{\theta}) = \frac{1}{L} \sum_{i=1}^{L} \sigma_i^2\; \mathbb{E}\Big[ \tfrac{1}{2} \big\| \mathbf{s}_{\boldsymbol{\theta}}(\mathbf{x} + \sigma_i \boldsymbol{\epsilon}, \sigma_i) + \boldsymbol{\epsilon} / \sigma_i \big\|^2 \Big]

the \sigma^2 weight gives every level a term of comparable size;
with \mathbf{s}_{\boldsymbol{\theta}} = \mathbf{f}_{\boldsymbol{\theta}} / \sigma, the network predicts -\boldsymbol{\epsilon};
the fit is best at the coarse levels and worst at the finest.

with torch.no_grad():
    print(f'{"sigma":>6} {"Fisher divergence":>18} {"0.5 E||score||^2":>17}')
    for sigma in sigmas[::3]:
        x_noisy = test + sigma * torch.randn_like(test)
        true = mix.score(x_noisy, sigma=sigma.item())
        s = score(x_noisy, torch.full((len(test),), sigma.item()))
        print(f'{sigma:>6.3f} {0.5 * ((s - true) ** 2).sum(1).mean():>18.4f} '
              f'{0.5 * (true ** 2).sum(1).mean():>17.3f}')

 sigma  Fisher divergence  0.5 E||score||^2
 3.000             0.0011             0.079
 0.766             0.0412             1.133
 0.196             0.1244             3.587
 0.050             1.1162             4.086

Annealed Langevin Dynamics

\alpha_i = \epsilon\, \frac{\sigma_i^2}{\sigma_L^2}, \qquad \mathbf{x} \leftarrow \mathbf{x} + \frac{\alpha_i}{2}\, \mathbf{s}_{\boldsymbol{\theta}}(\mathbf{x}, \sigma_i) + \sqrt{\alpha_i}\, \boldsymbol{\xi}

For i = 1, \ldots, L, a fixed number of steps each, every level starting where the last one stopped. The step size tracks the score’s scale; the injected noise shrinks with the level.

The Ladder Recovers What One Level Could Not

def annealed_langevin(score, x, sigmas, steps_per_level=100, scale=0.1):
    """Langevin dynamics at each level of the ladder, coarse to fine."""
    with torch.no_grad():
        for sigma in sigmas:
            alpha = scale * sigma ** 2  # alpha_i = eps * sigma_i^2 / sigma_L^2
            sigma_vec = torch.full((len(x),), sigma.item())
            for _ in range(steps_per_level):
                x = (x + 0.5 * alpha * score(x, sigma_vec)
                     + alpha ** 0.5 * torch.randn_like(x))
    return x

def summarize(x):
    """Mode weights, within-mode spread, and the fraction of stray chains."""
    d2 = ((x[:, None, :] - mix.means[None]) ** 2).sum(-1)
    weights = [round(w, 3) for w in
               (torch.bincount(d2.argmin(1), minlength=3) / len(x)).tolist()]
    nearest = d2.min(1).values
    inside = nearest < (3 * mix.std) ** 2
    return (f'mode weights {weights}, spread {nearest[inside].mean():.2f}, '
            f'stray {1 - inside.float().mean():.3f}')

torch.manual_seed(2)
x_init = 3 * torch.randn(2000, 2)
print(f'target (2000 samples of the mixture): {summarize(data)}')
for steps in (20, 100):
    x = annealed_langevin(score, x_init, sigmas, steps_per_level=steps)
    print(f'annealed, {steps:>3d} steps per level: {summarize(x)}')
x_single = annealed_langevin(score, x_init, sigmas[-1:], steps_per_level=1000)
print(f'finest level only, 1000 steps:  {summarize(x_single)}')

target (2000 samples of the mixture): mode weights [0.5, 0.317, 0.183], spread 0.49, stray 0.012
annealed,  20 steps per level: mode weights [0.428, 0.314, 0.257], spread 0.75, stray 0.133
annealed, 100 steps per level: mode weights [0.497, 0.298, 0.206], spread 0.54, stray 0.019
finest level only, 1000 steps:  mode weights [0.343, 0.317, 0.34], spread 1.04, stray 0.469

Same budget of a thousand evaluations: the annealed chain recovers the weights with a slightly wider spread and few strays; the finest level alone moves the chains only part of the way, because its steps are 3600 times smaller than those of the coarsest level.

A U-Net Approximates the Scores at Every Level

One Network Denoises at Three Levels

sigma = 0.09: per-pixel squared error noisy 0.009, denoised 0.002
sigma = 0.43: per-pixel squared error noisy 0.186, denoised 0.022
sigma = 2.00: per-pixel squared error noisy 4.004, denoised 0.132

Fine level: the original. Middle: the posterior mean. Coarse: a blurred silhouette, the layout without the detail. Layout, shape, detail.

Recap

A noise-conditional network learns \nabla \log p_{\sigma_i} for a whole ladder, by noise prediction.
Annealed Langevin dynamics samples coarse to fine; each level starts from the previous level’s samples, likely where its score is well estimated.
Conditioning: add a classifier gradient to the score (Bayes’ rule).
Remaining costs: a long chain, no tractable likelihood, a ladder rather than a process. Next: diffusion models.