Score function

Score function

score functions (gradients of log probability density functions)

learn score functions on a large number of noise-perturbed data distributions, then generate samples with Langevin-type sampling.

denote data distribution is $\boldsymbol{p(x)}$, the goal of generative modeling is to fit a model to the data distribution such that we can synthesize new data points at will by sampling from the distribution.

assume probability density function: $\boldsymbol{p_\theta(x) = \dfrac{e^{-f_\theta(x)}}{Z_\theta},Z_\theta > 0, \int p_\theta(x)dx = 1}$, it is unnormalized probabilistic model, or energy-based model: $\boldsymbol{f_\theta(x)}$,

train $\boldsymbol{p_\theta(x)}$ by maximizing the log-likelihood of the data: $\underset{\theta}{max}\sum_{i=1}^{N}{log\ \boldsymbol{p_\theta(x_i)}}$, requires $\boldsymbol{p_\theta(x)}$ to be a normalized probability density function $\Rightarrow$ evaluate the normalizing constant $Z_\theta\ \Rightarrow$ a typically intractable quantity for any general $\boldsymbol{f_\theta(x)}$

The score function of a distribution $\boldsymbol{p(x) = \triangledown_xlog\ p(x)}$, a model for the score function is called a score-based model $\boldsymbol{s_\theta(x) \approx \triangledown_xlog\ p(x)}$, it can be parameterized without worrying about the normalizing constant $\boldsymbol{Z_\theta} \Leftrightarrow$ the score-based model $\boldsymbol{s_\theta(x)}$ is independent of the normalizing constant $\boldsymbol{Z_\theta}$.

$\boldsymbol{s_\theta(x) \approx \triangledown_xlog\ p_\theta(x) = \triangledown_xlog \dfrac{e^{-f_\theta(x)}}{Z_\theta} = -\triangledown_xf_\theta(x) -\underset{=0}{\underbrace{ \triangledown_xlog\ Z_\theta}}} = -\triangledown_xf_\theta(x)$

unknown data score $\triangledown_xlog\ p(x)\ \Rightarrow$ score matching

Score matching objectives can directly be estimated on a dataset and optimized with stochastic gradient descent.

$1$-$D$ data mathematical derivation:

$\int \dfrac{1}{2} p_{data}(x)(\triangledown_xlog\ p_{\theta}(x))^2 dx = \dfrac{1}{2} \mathbb{E}{p{data}}[(\triangledown_xlog\ p_{\theta}(x))^2]$

By integration by parts:

muti-dimensional data score matching objective, $\triangledown_x^2$ denotes the Hessian with respect to $x$:

if $f_\theta(x)$ is parameterized by a deep neural network, $tr(\triangledown_x^2log\ p_{\theta}(x)) $ requires a number of backpropagation that is proportional to the data dimension $D$.

sliced score matching: one dimensional data distribution is much easier to estimate for score matching. project the scores onto random directions, such that the vector fields of scores of the data and model distribution become scalar fields. then compare the scalar fields to determine how far the model distribution is from the data distribution. It is clear to see that the two vector fields are equivalent if and only if their scalar fields corresponding to projections onto all directions are the same.

denote $\boldsymbol{v}$ as the random projection direction, The random projected version of Fisher divergence is sliced Fisher divergence:

$2 * \dfrac{1}{2}\mathbb{E}{p{data}}[(v^T\triangledown_xlog\ p_{data}(x) - v^T\triangledown_xlog\ p_{\theta}(x))^2] = 2 * (\mathbb{E}{p{data}}[v^T\triangledown_x^2log\ p_{\theta}(x) - \dfrac{1}{2}(v^T\triangledown_xlog\ p_{\theta}(x))^2] + const)$

$v^T\triangledown_x^2log\ p_{\theta}(x)$ is in the form of Hessian-vector products, which can be computed within $O(1)$ backpropagations

Langevin dynamics

given trained $\boldsymbol{s_\theta(x) \approx \triangledown_xlog\ p(x)}$, we can use an iterative procedure called Langevin dynamics to draw samples from it.

Langevin dynamics provides an MCMC procedure to sample from a distribution $\boldsymbol{p(x)}$ using only its score function $\triangledown_xlog\ p(x)$, Specifically, it initializes the chain from an arbitrary prior distribution $x_0 \sim \pi(x)$, and then iterates the following:

$x_{i+1} \leftarrow x_i + \epsilon \triangledown_xlog\ p(x) + \sqrt{2\epsilon} z_i, i = 0,1,…,K, z_i \sim \mathcal{N}(0,I) \ x_{i+1} \leftarrow x_i + \epsilon s_\theta(x) + \sqrt{2\epsilon} z_i$

When $ \rightarrow 0$ and $K \rightarrow \infty$, $x_K$ converges to a sample from $p(x)$.

Pitfalls

The key challenge is the fact that the estimated score functions are inaccurate in low density regions, where few data points are available for computing the score matching objective. This is expected as score matching minimizes the Fisher divergence.

score-based model are weighted by $p(x)$, they are largely ignored in low density regions where $p(x)$ is small.

Multiple noise perturbations

perturb data points with noise and train score-based models on the noisy data points instead. When the noise magnitude is sufficiently large, it can populate low data density regions to improve the accuracy of estimated scores.

choose an appropriate noise scale for the perturbation process: use multiple scales of noise perturbations simultaneously. Suppose we always perturb the data with isotropic Gaussian noise, and let there be a total of $L$ increasing standard deviations $\sigma_1 < … < \sigma_L$.

1. perturb the data distribution $p(x)$ with each of the Gaussian noise $\mathcal{N}(0, \sigma^2iI), i=1,…,L$ to obtain a noise-perturbed distribution: $p{\sigma_i}(x) = \int p(y)\mathcal{N}(x;y,\sigma_i^2I)dy \rightarrow$ sample $x \sim p(x)$ and $z \sim \mathcal{N}(0,I)$ and compute $x + \sigma_iz$.

2. estimate the score function of each noise-perturbed distribution, $\triangledown_xlog\ p_{\sigma_i}(x)$, by training a Noise Conditional Score-Based Model $s_\theta(x,i)$ with score matching, such that $s_\theta(x,i) \approx \triangledown_xlog\ p_{\sigma_i}(x), i=1,…,L$.

   $\begin{align}\mathbb{E}{p{\sigma_i}(x)}[

   \triangledown_xlog\ p_{\sigma_i}(x) - s_\theta(x, i)

   2^2] & = 2\ \mathbb{E}{p_{data, \sigma_i}}[tr(\triangledown_x^2log\ p_{\theta}(x, i)) + \dfrac{1}{2}

   \triangledown_xlog\ p_{\theta}(x, i)

   2^2] +const \ & = 2\ \mathbb{E}{p_{data, \sigma_i}}[tr(\triangledown_xs_{\theta}(x, i)) + \dfrac{1}{2}

   s_{\theta}(x, i)

   _2^2] + const \end{align}$

3. The training objective for $s_\theta(x,i)$ is a weighted sum of Fisher divergences for all noise scales: $\mathbb{E}_{p(x)}[

   \triangledown_xlog\ p(x) - s_\theta(x,i)

   2^2] = \sum{i=1}^L\lambda(i)\mathbb{E}{p{\sigma_i}(x)}[

   \triangledown_xlog\ p_{\sigma_i}(x) - s_\theta(x, i)

   2^2], \lambda(i) \in \mathbb{R}{>0}, \lambda(i) = \sigma_i^2$

4. annealed Langevin dynamics: produce samples from it by running Langevin dynamics for $i = L,L-1,…,1$ in sequence, since the noise scale $\sigma_i$ decreases (anneals) gradually over time.

   

Stochastic differential equations (SDEs)

By generalizing the number of noise scales to infinity, we obtain not only higher quality samples, but also, among others, exact log-likelihood computation, and controllable generation for inverse problem solving.

When the number of noise scales approaches infinity, we essentially perturb the data distribution with continuously growing levels of noise. In this case, the noise perturbation procedure is a continuous-time stochastic process.

Many stochastic processes (diffusion processes in particular) are solutions of stochastic differential equations (SDEs, hand designed):

$dx = f(x,t)dt + g(t)dw, f(·，t):\mathbb{R}^d \rightarrow \mathbb{R}^d, g(t) \in \mathbb{R} \Rightarrow dx = e^tdw$

$f(·，t)$ called drift coefficient, $g(t)$ called diffusion coefficient, $w$ denotes a standard Brownian motion. The solution of a stochastic differential equation is a continuous collection of random variables ${x(t)}_{t \in [0,T]}$. These random variables trace stochastic trajectories as the time index $t$ grows from the start time $0$ to the end time $T$. $p_t(x)$ denote the (marginal) probability density function of $x(t),\ p_0(x) = p(x), p_T(x) \approx \pi(x)$.

$\boldsymbol{dx = e^tdw}$: perturbs data with a Gaussian noise of mean zero and exponentially growing variance ($f(x,t) = 0, g(t) = e^t$).

reverse the perturbation process for sample generation by using the reverse SDE. any SDE has a corresponding reverse SDE:

$dx = [f(x,t) - g^2(t)\triangledown_xlog\ p_t(x)]dt + g(t)dw$

train a Time-Dependent Score-Based Model $s_\theta(x,t) \approx \triangledown_xlog\ p_t(x)$:

$dx = [f(x,t) - g^2(t)s_\theta(x,t)]dt + g(t)dw$

start with $x(T) \sim \pi$, and solve the above reverse SDE to obtain a sample $x(0), x(0) \sim p_\theta(x)$.

solving the estimated reverse SDE with numerical SDE solvers(Euler-Maruyama method), it discretizes the SDE using finite time steps and small Gaussian noise.

Probability flow ODE

a sampler based on ordinary differential equations (ODEs) that allow for exact likelihood computation.

it s possible to convert any SDE into an ordinary differential equation (ODE) without changing its marginal distributions ${p_t(x}_{t \in [0,T]}$. The corresponding ODE of an SDE is named probability flow ODE: $dx = [f(x,t) - \dfrac{1}{2}g^2(t)\triangledown_xlog\ p_t(x)]dt$.

trajectories obtained by solving the probability flow ODE have the same marginal distributions as the SDE trajectories. can leverage the instantaneous change-of-variable formula to compute the unknown data density $p_0$ from the known prior density $p_T$ with numerical ODE solvers.

Controllable generation

Connection with Diffusion Model

By generalizing the number of noise scales to infinity, we proved that score-based generative models and diffusion probabilistic models can both be viewed as discretizations to stochastic differential equations determined by score functions, bridges both score-based generative modeling and diffusion probabilistic modeling into a unified framework.