SDE Diffusion

• Score-based Generative Model through SDE
  • 1. Notations
  • 2. Score-based Diffusion through SDE (Unconditional)
  • 3. Training (Unconditional)
  • 4. Score-based Diffusion through SDE (Conditional)
  • 5. Special Case
  • 6. Denoising Score Matching
  • 7. VPSDE（Continuous DDPM） (Unconditional)
  • 8. VESDE

Score-based Generative Model through SDE

1. Notations

The meaning of $\nabla \cdot [\cdot ]$:

For scalar function $\varphi (\mathbf{x}):{\mathbb{R}}^d\to \mathbb{R}$, we know $\nabla \varphi (\mathbf{x})$ is the gradient.

For vector function $\varphi (\mathbf{x}):{\mathbb{R}}^d\to {\mathbb{R}}^d$, $\nabla \cdot \varphi (\mathbf{x})\in \mathbb{R}$ is the divergence, defined as the sum of the partial derivatives of each component:

$$\nabla \cdot \varphi (\mathbf{x})=\sum _{i=1}^d\frac{\partial {\varphi }_i}{\partial x_i}$$

For matrix function $\varphi (\mathbf{x}):{\mathbb{R}}^d\to {\mathbb{R}}^{d\times d}$, its divergence is a vector, defined as:

$$\nabla \cdot \varphi (\mathbf{x})={\left(\sum _{j=1}^d\frac{\partial {\varphi }_{1j}}{\partial x_j},\dots ,\sum _{j=1}^d\frac{\partial {\varphi }_{dj}}{\partial x_j}\right)}^{\top }$$

or written in a more general form:

$$\nabla \cdot \varphi (\mathbf{x}):={\left(\nabla \cdot {\varphi }^1(\mathbf{x}),\dots ,\nabla \cdot {\varphi }^d(\mathbf{x})\right)}^{\top }$$

where ${\varphi }^i(\mathbf{x})$ is the $i$th row of $\varphi$.

2. Score-based Diffusion through SDE (Unconditional)

SDE of forward process(noise):

$$\begin{gathered} d\mathbf{x}=f(\mathbf{x},t)dt+g(\mathbf{x},t)dw \\ f:{\mathbb{R}}^d\times [0,T]\to {\mathbb{R}}^d,g:{\mathbb{R}}^d\times [0,T]\to {\mathbb{R}}^{d\times d} \end{gathered}$$

SDE of reverse process(denoise):

$$d\mathbf{x}=\left(f(\mathbf{x},t)-\nabla \cdot [g(\mathbf{x},t)g(\mathbf{x},t{)}^T]-g(\mathbf{x},t)g(\mathbf{x},t{)}^T{\nabla }_{\mathbf{x}}\ln p_t(\mathbf{x})\right)dt+g(\mathbf{x},t)dw$$

probability flow ODE:

$$d\mathbf{x}=\left(f(\mathbf{x},t)-\frac{1}{2}\nabla \cdot [g(\mathbf{x},t)g(\mathbf{x},t{)}^T]-\frac{1}{2}g(\mathbf{x},t)g(\mathbf{x},t{)}^T{\nabla }_{\mathbf{x}}\ln p_t(\mathbf{x})\right)dt$$

3. Training (Unconditional)

• Sliced Score Matching:

$${\theta }^{\ast }=\arg \underset{\theta }{\min }{\mathbb{E}}_t\left\{\lambda (t){\mathbb{E}}_{\mathbf{x}(0)}{\mathbb{E}}_{\mathbf{x}(t)}{\mathbb{E}}_{\mathbf{v}\sim p_{\mathbf{v}}}\left[{\mathbf{v}}^{\top }{\mathbf{s}}_{\theta }(\mathbf{x}(t),t)\mathbf{v}+\frac{1}{2}({\mathbf{v}}^T{\mathbf{s}}_{\theta }(\mathbf{x}(t),t){)}^2\right]\right\}$$

• Denoising Score Matching: $${\theta }^{\ast }=\arg \underset{\theta }{\min }{\mathbb{E}}_{t\sim U(0,T)}\lambda (t){\mathbb{E}}_{x_0\sim p_0}\left[{\mathbb{E}}_{x\sim p_{0t}}\left[{\Vert s_{\theta }(x,t)-{\nabla }_{\mathbf{x}}\ln p_{0t}(\mathbf{x}|{\mathbf{x}}_0)\Vert }_2^2\right]\right]$$

4. Score-based Diffusion through SDE (Conditional)

$$d\mathbf{x}=\left(f(\mathbf{x},t)-\nabla \cdot [g(\mathbf{x},t)g(\mathbf{x},t{)}^T]-g(\mathbf{x},t)g(\mathbf{x},t{)}^T{\nabla }_{\mathbf{x}}\ln p_t(\mathbf{x}|\mathbf{y})\right)dt+g(\mathbf{x},t)dw$$ $${\nabla }_{\mathbf{x}}\ln p_t(\mathbf{x}|\mathbf{y})={\nabla }_{\mathbf{x}}\ln \left(\frac{p_t(\mathbf{y}|\mathbf{x})p_t(\mathbf{x})}{p_t(\mathbf{y})}\right)={\nabla }_{\mathbf{x}}\ln p_t(\mathbf{y}|\mathbf{x})+{\nabla }_{\mathbf{x}}\ln p_t(\mathbf{x})$$ $$d\mathbf{x}=\left(f(\mathbf{x},t)-\nabla \cdot [g(\mathbf{x},t)g(\mathbf{x},t{)}^T]-g(\mathbf{x},t)g(\mathbf{x},t{)}^T\left({\nabla }_{\mathbf{x}}\ln p_t(\mathbf{y}|\mathbf{x})+{\nabla }_{\mathbf{x}}\ln p_t(\mathbf{x})\right)\right)dt+g(\mathbf{x},t)dw$$

We can train a neural network $c(\mathbf{x},t)$ to learn $\ln p_t(\mathbf{y}|\mathbf{x})$

We can also use some prior knowledge to directly determine ${\nabla }_{\mathbf{x}}\ln p_t(\mathbf{y}|\mathbf{x})$

5. Special Case

If $\exists \tilde{g}(t)\in \mathbb{R}$ s.t. $g(\mathbf{x},t)=\tilde{g}(t)\cdot I$, then the reverse process can be written as:

$$d\mathbf{x}=\left(f(\mathbf{x},t)-\tilde{g}(t{)}^2{\nabla }_{\mathbf{x}}\ln p_t(\mathbf{x})\right)dt+\tilde{g}(t)dw$$

6. Denoising Score Matching

Assume forward diffusion process can be written as ${\mathbf{x}}_t=a_t{\mathbf{x}}_0+b_tϵ,ϵ\sim \mathcal{N}(0,\mathbf{I})$, then

$${\mathbf{x}}_t\sim \mathcal{N}(a_t{\mathbf{x}}_0,b_t^2\mathbf{I})$$

minimize

$$\underset{t,{\mathbf{x}}_0,{\mathbf{x}}_t\sim p_{0t}({\mathbf{x}}_t|{\mathbf{x}}_0)}{\mathbb{E}}\Vert s_{\theta }({\mathbf{x}}_t,t)-{\nabla }_{{\mathbf{x}}_t}\ln p_{0t}({\mathbf{x}}_t|{\mathbf{x}}_0){\Vert }_2^2$$

• Equivalence of Epsilon Model and Score Model

$$\begin{array}{rl}\text{score}& ={\nabla }_{{\mathbf{x}}_t}\ln p_{0t}({\mathbf{x}}_t|{\mathbf{x}}_0)\\ & =-\frac{1}{b_t^2}({\mathbf{x}}_t-a_t{\mathbf{x}}_0)\\ & =-\frac{ϵ}{b_t}\end{array}$$

• Unconditional Score Matching

$\mathrm{argmin}\Vert \frac{-1}{b_t}{ϵ}_{\theta }({\mathbf{x}}_t,t)-\frac{-1}{b_t}ϵ{\Vert }_2^2⟺\mathrm{argmin}\Vert {ϵ}_{\theta }({\mathbf{x}}_t,t)-ϵ{\Vert }_2^2$

• Conditional (Loss Guidance) Score Matching

$\mathrm{argmin}\Vert \frac{-1}{b_t}{ϵ}_{\theta }({\mathbf{x}}_t,t)-\frac{-1}{b_t}ϵ-{\nabla }_x-l({\mathbf{x}}_t,{\mathbf{x}}_0){\Vert }_2^2⟺\mathrm{argmin}\Vert {ϵ}_{\theta }({\mathbf{x}}_t,t)-ϵ-{\nabla }_x{b}_{t}l({\mathbf{x}}_t,{\mathbf{x}}_0){\Vert }_2^2$

7. VPSDE（Continuous DDPM） (Unconditional)

$$P({\mathbf{x}}_t|{\mathbf{x}}_{t-1})=\mathcal{N}({\mathbf{x}}_t;\sqrt{1-{\beta }_t}\cdot {\mathbf{x}}_{t-1},{\beta }_t\cdot I)$$ $${\mathbf{x}}_t=\sqrt{1-{\beta }_t}\cdot {\mathbf{x}}_{t-1}+\sqrt{{\beta }_t}\cdot ϵ,ϵ\sim \mathcal{N}(0,I)$$ $${\mathbf{x}}_{t+\Delta t}-{\mathbf{x}}_t=\sqrt{1-{\beta }_{t+\Delta t}}\cdot {\mathbf{x}}_t+\sqrt{{\beta }_{t+\Delta t}}\cdot ϵ-{\mathbf{x}}_t$$

Because $\sqrt{1-x}=1-\frac{x}{2}+o(x)$, we have:

$${\mathbf{x}}_{t+\Delta t}-{\mathbf{x}}_t=-\frac{{\beta }_{t+\Delta t}}{2}\cdot {\mathbf{x}}_t+\sqrt{{\beta }_{t+\Delta t}}\cdot ϵ+o({\beta }_{t+\Delta t})\cdot {\mathbf{x}}_t$$

SDE of forward process(noise):

$$d\mathbf{x}=-\frac{{\beta }_t}{2}\cdot \mathbf{x}dt+\sqrt{{\beta }_t}dw$$

SDE of reverse process(denoise):

$$d\mathbf{x}=\left(-\frac{{\beta }_t}{2}\cdot \mathbf{x}-{\beta }_t\cdot {\nabla }_{\mathbf{x}}\ln p_t(\mathbf{x})\right)dt+\sqrt{{\beta }_t}dw$$

参数对比

model	beta	n_step
DDPM	0.0001-0.02	1000
VPSDE	0.1-20	1000

8. VESDE

$${\mathbf{x}}_t\sim \mathcal{N}({\mathbf{x}}_0,{\sigma }_t^{2}I),{\sigma }_t={\sigma }_{min}({\sigma }_{max}/{\sigma }_{min}{)}^t$$ $$\begin{array}{rl}dx& =g(t)dw\\ x_t& =x_0+{\int }_0^{t}g(s)ds\\ Var[x_t]& ={\int }_0^{t}g(s{)}^{2}ds=:{\sigma }_t^2\\ g(t)& =\sqrt{\frac{d}{dt}{\sigma }_t^2}={\sigma }_t\sqrt{2\ln ({\sigma }_{max}/{\sigma }_{min})}\end{array}$$

离散情况下：

$$\begin{array}{rl}{x}_t& =x_{t-1}+g(t){ϵ}_t,{ϵ}_t\sim \mathcal{N}(0,I)\\ {\sigma }_t^2& ={\sigma }_{t-1}^2+g(t{)}^2\\ g(t)& =\sqrt{{\sigma }_t^2-{\sigma }_{t-1}^2}\end{array}$$