Score-based SDE: Unified score matching framework for generative modeling.

Implements the Score SDE framework from "Score-Based Generative Modeling through Stochastic Differential Equations" (Song et al., ICLR 2021). Unifies DDPM, SMLD, and other score-based methods under a single SDE perspective with two main variants:

SDE Variants

VP-SDE (Variance Preserving, generalizes DDPM):

dx = -0.5 * beta(t) * x dt + sqrt(beta(t)) dw
beta(t) = beta_min + t * (beta_max - beta_min)

VE-SDE (Variance Exploding, generalizes SMLD):

dx = sqrt(d[sigma^2(t)]/dt) dw
sigma(t) = sigma_min * (sigma_max/sigma_min)^t

Key Innovation: Score Function

The score function s(x, t) = grad_x log p_t(x) is learned via denoising score matching. Once learned, samples are generated by solving the reverse-time SDE:

dx = [-0.5 * beta(t) * x - beta(t) * s(x, t)] dt + sqrt(beta(t)) dw_rev

Or the probability flow ODE (deterministic):

dx = [-0.5 * beta(t) * x - 0.5 * beta(t) * s(x, t)] dt

Architecture

Input (x_t, t)
      |
      v
+-----------------------+
| Score Network         |
| s_theta(x_t, t)       |
|                       |
| x_proj + time_embed   |
| -> residual blocks    |
| -> score prediction   |
+-----------------------+
      |
      v
Score [batch, input_dim]  (gradient of log density)

Usage

model = ScoreSDE.build(
  input_dim: 64,
  hidden_size: 256,
  num_layers: 4,
  sde_type: :vp
)

# Training: denoising score matching
loss = ScoreSDE.dsm_loss(score_pred, noise, sigma)

Reference

• Paper: "Score-Based Generative Modeling through Stochastic Differential Equations"
• arXiv: https://arxiv.org/abs/2011.13456