Score-based generative models

What is the Score Function?

Why the gradient of the log probability?

Let's derive this from first principles. Suppose we have $p(x) = \frac{\tilde{p}(x)}{Z}$ where $Z = \int \tilde{p}(x) dx$ is the partition function.

$\nabla_x \log p(x) = \nabla_x \log \left(\frac{\tilde{p}(x)}{Z}\right)$

$= \nabla_x [\log \tilde{p}(x) - \log Z]$

Since $Z$ doesn't depend on $x$ : $= \nabla_x \log \tilde{p}(x)$

The magic: The intractable normalization constant $Z$ disappears! This is why score-based models are powerful — we never need to compute or estimate $Z$ .

Score Matching: Learning the Score

The problem: We don't know $\nabla_x \log p_{\text{data}}(x)$ ! That's what we're trying to learn.

Denoising Score Matching

The breakthrough: We can estimate scores by adding noise and learning to denoise.

Derivation from scratch:

Add Gaussian noise: $\tilde{x} = x + \sigma \epsilon$ where $\epsilon \sim \mathcal{N}(0, I)$
The perturbed distribution is $q_\sigma(\tilde{x}|x) = \mathcal{N}(\tilde{x}; x, \sigma^2 I)$
The marginal is $q_\sigma(\tilde{x}) = \int p_{\text{data}}(x) q_\sigma(\tilde{x}|x) dx$

Key insight: For Gaussian noise, the score of $q_\sigma(\tilde{x}|x)$ is: $\nabla_{\tilde{x}} \log q_\sigma(\tilde{x}|x) = \nabla_{\tilde{x}} \log \mathcal{N}(\tilde{x}; x, \sigma^2 I)$

For a Gaussian $\mathcal{N}(y; \mu, \Sigma)$ , we have: $\log \mathcal{N}(y; \mu, \Sigma) = -\frac{1}{2}(y - \mu)^T \Sigma^{-1}(y - \mu) + \text{const}$

Taking the gradient with respect to $y$ : $\nabla_y \log \mathcal{N}(y; \mu, \Sigma) = -\Sigma^{-1}(y - \mu)$

Applying to our case ( $y = \tilde{x}$ , $\mu = x$ , $\Sigma = \sigma^2 I$ ): $\nabla_{\tilde{x}} \log q_\sigma(\tilde{x}|x) = -\frac{1}{\sigma^2}(\tilde{x} - x) = -\frac{\epsilon}{\sigma}$

Why this works: By learning to denoise, the network implicitly learns the score function at noise level $\sigma$ .

Multi-Scale Score Matching

Real data lives on low-dimensional manifolds in high-dimensional space. A single noise level $\sigma$ won't cover the whole distribution well.

Formula Noise-Conditional Score Network

Use multiple noise levels $\{\sigma_1 > \sigma_2 > \cdots > \sigma_L\}$ where $\sigma_1$ is large (covers all space) and $\sigma_L$ is small (close to data).

Train a noise-conditional score network $s_\theta(x, \sigma)$ with objective: $J(\theta) = \sum_{i=1}^L \lambda(\sigma_i) \mathbb{E}_{x \sim p_{\text{data}}} \mathbb{E}_{\epsilon \sim \mathcal{N}(0,I)} \left[\frac{1}{2} \left\|s_\theta(x + \sigma_i\epsilon, \sigma_i) + \frac{\epsilon}{\sigma_i}\right\|^2\right]$ where $\lambda(\sigma_i) = \sigma_i^2$ is a common weighting choice (makes the objective scale-invariant).

Why multiple scales?

Large $\sigma$ : Ensures connectivity, no isolated modes
Small $\sigma$ : Captures fine data structure
Interpolation: Smooth transition between scales

Sampling via Langevin Dynamics

Once we have the score function, we generate samples using Langevin MCMC:

Derivation: This comes from the Langevin SDE: $dx = \nabla_x \log p(x) dt + \sqrt{2} dW$ where $dW$ is a Wiener process. Discretizing with step size $\alpha$ : $x_{t+1} = x_t + \alpha \nabla_x \log p(x_t) + \sqrt{2\alpha} \epsilon$

Why the $\sqrt{2\alpha}$ noise term? It's the discretization of the diffusion term $\sqrt{2}dW$ . The variance of $dW$ over time $dt$ is $dt$ , so over step $\alpha$ it's $\sqrt{2\alpha}$ .

Annealing schedule: Start at $\sigma_1$ (coarse), progressively decrease to $\sigma_L$ (fine). This ensures:

Initial exploration in high-noise regime
Refinement in low-noise regime
Avoidance of local minima

Worked example Sampling Algorithm

def sample(score_model, sigmas, n_steps_each=100):
    x = torch.randn(batch_size, data_dim) * sigmas[0]
    for sigma in sigmas:
        alpha = 0.1 * (sigma / sigmas[-1]) ** 2  # Adaptive step size
        for t in range(n_steps_each):
            score = score_model(x, sigma)
            noise = torch.randn_like(x)
            x = x + alpha * score + sqrt(2* alpha) * noise
    return x

Why this step?

x = torch.randn(...) * sigmas[0]: Initialize from high-variance Gaussian matching the largest noise scale
alpha = ... * (sigma / sigmas[-1]) ** 2: Smaller steps for smaller noise levels (more precision needed)
x = x + alpha * score: Move toward higher probability (follow the score)
+ sqrt(2 * alpha) * noise: Add stochasticity to explore and satisfy detailed balance

Connection to Diffusion Models

Derivation: Given $x_t = \alpha_t x_0 + \sigma_t \epsilon$ , the conditional distribution is: $q(x_t | x_0) = \mathcal{N}(x_t; \alpha_t x_0, \sigma_t^2 I)$

The score is: $\nabla_{x_t} \log q(x_t | x_0) = -\frac{x_t - \alpha_t x_0}{\sigma_t^2} = -\frac{\sigma_t \epsilon}{\sigma_t^2} = -\frac{\epsilon}{\sigma_t}$

Therefore, learning to predict $\epsilon$ is equivalent to learning the score $s_\theta(x_t, t) \approx -\epsilon/\sigma_t$ .

Worked example Training Loop

def train_score_model(model, dataloader, sigmas):
    optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)
    
    for x_batch in dataloader:  # x_batch: clean data
        # Sample noise level
        sigma_idx = torch.randint(0, len(sigmas), (x_batch.size(0),))
        sigma = sigmas[sigma_idx]  # shape: (batch,)
        
        # Add noise
        epsilon = torch.randn_like(x_batch)
        x_noisy = x_batch + sigma.view(-1, 1, 1, 1) * epsilon
        # Predict score
        score_pred = model(x_noisy, sigma)
        
        # Target score: -epsilon / sigma
        score_target = -epsilon / sigma.view(-1, 1, 1)
        # Loss (weighted by sigma^2 for scale invariance)
        loss = (sigma**2).view(-1, 1, 1) * (score_pred - score_target)**2
        loss = loss.mean()
        
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

Why this step?

sigma_idx = torch.randint(...): Randomly sample noise levels to train across all scales
x_noisy = x_batch + sigma * epsilon: Create training pairs (noisy input, target score)
score_target = -epsilon / sigma: The true score of the noised distribution
(sigma**2) * (...)^2: Weight by $\sigma^2$ so gradients have similar magnitude across noise levels

Common mistake Common Mistake: Forgetting the Noise Variance Weighting

Wrong idea: Use loss $\mathcal{L} = \|s_\theta(x, \sigma) + \epsilon/\sigma\|^2$ directly.

Why it feels right: This is the squared error between predicted and true score, seems natural.

The problem: At large $\sigma$ , scores are small ( $\approx -\epsilon/\sigma$ where $\sigma$ is large). At small $\sigma$ , scores are large. Without weighting, the model focuses only on large $\sigma$ (large gradients) and ignores fine details.

The fix: Weight by $\sigma^2$ : $\mathcal{L} = \sigma^2 \|s_\theta(x, \sigma) + \epsilon/\sigma\|^2$ . This makes the loss scale-invariant and balances learning across noise levels.

Recall Explain to a 12-year-old

Imagine you're lost in the mountains in thick fog. You can't see the valley where you need to go, but you have a magic compass that always points downhill.

Score-based models are like that compass. Instead of trying to memorize the exact shape of every mountain (which would be impossible), we learn a "which way is down" function everywhere. Then to get to the valley (generate realistic data), we just start anywhere random and follow the compass downhill, taking small steps.

The trick is that we learn this compass by playing a game: we take real valley locations, throw them randomly up the mountain (add noise), then train our compass to point back toward where they came from. By doing this thousands of times at different heights (noise levels), our compass learns the right direction everywhere.

Connections

Diffusion Models: Score-based models are the continuous-time limit of discrete diffusion
VAE: Both learn latent structures, but scores model $p(x)$ directly vs. encoder-decoder
GAN: Both generate samples, but scores use gradients vs. adversarial training
Energy-Based Models: Score is the gradient of the energy function $E(x) = -\log p(x)$
Langevin Dynamics: The MCMC sampling procedure used for generation
Stochastic Differential Equations: Modern formulation treats diffusion as reverse-time SDE

#flashcards/ai-ml

What is the score function in score-based generative models? :: The gradient of the log-probability: $s(x) = \nabla_x \log p(x)$ . It points in the direction of increasing probability density.

Why do we model the score instead of the density directly?

(1) The intractable normalization constant

Z

p(x) = \tilde{p}(x)/Z

cancels out when taking the gradient of log

p

. (2) Scores are often easier to estimate in high dimensions.

What is the key insight of denoising score matching?

Add Gaussian noise

\tilde{x} = x + \sigma\epsilon

to data. The score of the noised distribution is

\nabla_{\tilde{x}} \log q(\tilde{x}|x) = -\epsilon/\sigma

, which we can compute directly and use as training target.

Write the denoising score matching objective :: $J(\theta) = \mathbb{E}_{x, \epsilon}[\frac{1}{2} \|s_\theta(x + \sigma\epsilon) + \epsilon/\sigma\|^2]$ where $\epsilon \sim \mathcal{N}(0, I)$ .

Why do we use multiple noise scales in score-based models?

(1) Large

\sigma

ensures coverage of the entire space. (2) Small

\sigma

captures fine data details. (3) Intermediate scales bridge the gap, avoiding mode collapse and poor mixing.

What is annealed Langevin dynamics?

An MCMC sampling procedure that starts at high noise

\sigma_1

and progressively decreases to low noise

\sigma_L

, at each level running:

x_{t+1} = x_t + \alpha s_\theta(x_t, \sigma) + \sqrt{2\alpha} z_t

where

z_t \sim \mathcal{N}(0, I)

Why weight the score matching loss by

\sigma^2

Without weighting, small

\sigma

(large scores) dominate the gradient. Weighting by

\sigma^2

makes the loss scale-invariant and balances learning across all noise levels.

How are score-based models related to diffusion models?

In diffusion,

x_t = \alpha_t x_0 + \sigma_t \epsilon

. The score is

\nabla \log p(x_t) = -\epsilon/\sigma_t

. So predicting noise

\epsilon

(diffusion objective) is equivalent to predicting the score (up to a scale factor).

What is the Fisher divergence minimized in score matching?

\mathbb{E}_{p_{\text{data}}}[\|s_\theta(x) - \nabla_x \log p_{\text{data}}(x)\|^2]

. It measures the squared difference between the model score and the true score, averaged over the data distribution.

Why does the

\sqrt{2\alpha}

noise term appear in Langevin dynamics?

It comes from discretizing the Langevin SDE

dx = \nabla \log p(x) dt + \sqrt{2} dW

. The diffusion term

\sqrt{2}dW

has variance

2dt

over time step

dt

, so over discrete step

\alpha

it becomes

\sqrt{2\alpha}

Concept Map

Hinglish (regional understanding)

Intuition Hinglish mein samjho

Hinglish (regional understanding)

Intuition Hinglish mein samjho

Chalo, ise simple tarike se samajhte hain. Normally generative models mein hum $p(x)$ yaani exact probability density seekhne ki koshish karte hain, lekin ye high dimensions mein bahut mushkil hota hai, kyunki usme ek normalizing constant $Z$ aa jaata hai jise calculate karna almost impossible hai. Score-based models ka jugaad ye hai ki hum density ke bajaye uska score function seekhte hain, matlab $\nabla_x \log p(x)$ . Ise aise socho jaise pahaad par khade ho aur tumhe har jagah ki exact height nahi pata, par tumhe ye pata hai ki "upar" kis direction mein hai. Bas isi slope ko follow karte hue tum random noise se shuru karke dheere-dheere realistic data tak pahunch sakte ho.

Ab magic ye hai — jab tum $\log p(x)$ ka gradient lete ho, to woh painful normalizing constant $Z$ apne aap gayab ho jaata hai, kyunki $Z$ ka $x$ se koi lena-dena nahi hota. Isliye hume kabhi $Z$ calculate karne ki zaroorat hi nahi padti, jo ki ek badi jeet hai. Lekin ek problem hai: true score to hume pata hi nahi hai, kyunki wahi to seekhna hai. Iska solution hai denoising score matching — hum data mein thoda Gaussian noise daalte hain aur network ko sikhaate hain ki us noise ko kaise hataana hai. Gaussian ke liye score ka formula seedha $-\epsilon/\sigma$ nikalta hai, to network basically noise predict karna seekh jaata hai aur isi process mein woh score function bhi seekh leta hai.

Ye baat isliye important hai kyunki real-world data usually high-dimensional space mein ek chhoti si "manifold" par hota hai, aur ek hi noise level se poori distribution cover nahi hoti. Isliye hum multiple noise levels use karte hain — bade $\sigma$ se poora space cover hota hai aur chhote $\sigma$ se data ke paas ki fine details aati hain. Yahi core idea aage jaake diffusion models ka foundation banta hai, jo aaj Stable Diffusion aur DALL-E jaise powerful image generators chalate hain. To ye topic samajhna future ke advanced generative AI ke liye ekdum zaroori hai.

What is the Score Function?

Score Matching: Learning the Score

Denoising Score Matching

Multi-Scale Score Matching

Sampling via Langevin Dynamics

Connection to Diffusion Models

Connections

Concept Map

Hinglish (regional understanding)

Hinglish (regional understanding)

Go deeper — visual, from zero

Test yourself — Generative Models

Connections