Chenggang Liu

Sampling with the Score

Once you have \(s_\theta(x) \approx \nabla_x \log p(x)\) , sample via Langevin dynamics:

\[ x_{t+1} = x_t + \frac{\eta}{2} s_\theta(x_t) + \sqrt{\eta}\, \xi_t, \quad \xi_t \sim \mathcal{N}(0, I) \]

As \(\eta \to 0\) and \(T \to \infty\) , \(x_T \sim p_\theta(x)\) .

For multi-scale: use annealed Langevin dynamics — start with high noise score, gradually decrease.

How to Sample

  1. Start from random noise: \(x_0 \sim \mathcal{N}(0, I)\)
  2. Iteratively update: \[ x_{t+1} = x_t + \frac{\eta}{2} s_\theta(x_t) + \sqrt{\eta}\, \xi_t, \quad \xi_t \sim \mathcal{N}(0, I) \]
    • The score \(s_\theta(x_t)\) pushes \(x\) toward high-density regions
    • The noise \(\xi_t\) ensures exploration of the full distribution
  3. As \(\eta \to 0\) and \(T \to \infty\), samples converge to \(p(x)\)

Annealed Langevin Dynamics (practical version)

  • Use decreasing noise levels \(\sigma_1 > \sigma_2 > \cdots > \sigma_L\)
  • At each level, run a few Langevin steps using the score trained at that noise level
  • High-noise scores provide a coarse landscape to avoid getting stuck in low-density regions early on

Test: Sample from a Known Gaussian

import torch

def langevin_sample(score_fn, n_steps=1000, eta=0.01, dim=2, n_samples=64):
    x = torch.randn(n_samples, dim)
    for _ in range(n_steps):
        x = x + (eta / 2) * score_fn(x) + torch.sqrt(torch.tensor(eta)) * torch.randn_like(x)
    return x

# Score of N(mu, I) is -(x - mu)
mu = torch.tensor([3.0, -2.0])
score_fn = lambda x: -(x - mu)

samples = langevin_sample(score_fn, n_steps=2000, eta=0.01, dim=2, n_samples=1000)
print(f"Mean: {samples.mean(0)}")  # should be close to [3, -2]
print(f"Std:  {samples.std(0)}")   # should be close to [1, 1]

Mean: tensor([ 2.9596, -1.9918])
Std:  tensor([1.0213, 0.9915])