Stochastic Differential Equation (SDE)

Forward and Backward Updates in Gradient Descent

Recall the equation of a step in gradient descent

$${\mathbf{x}}_i-{\mathbf{x}}_{i-1}=-{\beta }_{i-1}\nabla f({\mathbf{x}}_{i-1})$$

And we can rewrite this into

$$\Delta \mathbf{x}=-\beta \Delta t\cdot \nabla f(\mathbf{x})$$

Take $\Delta t\to 0$, we will have

$$\mathrm{d}\mathbf{x}=-\beta \nabla f(\mathbf{x})\mathrm{d}t$$

Consider the iteration of the descent procedure. For a quite small step, we have the forward equation, it is a ODE

$${\mathbf{x}}_i\approx {\mathbf{x}}_{i-1}-\beta \nabla f({\mathbf{x}}_{i-1})\mathrm{d}t$$

Similarly, we can write the reverse equation (See gradient ascent)

$${\mathbf{x}}_{i-1}\approx {\mathbf{x}}_{i-1}+\beta \nabla f({\mathbf{x}}_i)\mathrm{d}t$$

Stochastic Differential Equation

If we introduce a noise term ${\mathbf{z}}_t\sim \mathcal{N}(0,\text{I})$ to the gradient descent algorithm, then the ODE will become a stochastic differential equation (SDE). If denote the noise term by $\mathbf{z}(t)$, we will have (treating $\mathbf{x}(t)$ as a continuous function)

$$\mathbf{x}(t+\Delta t)=\mathbf{x}(t)-\tau \nabla f(\mathbf{x}(t))+\mathbf{z}(t)$$

And we define a random process $\mathbf{w}(t)$ such that $\mathbf{z}(t)=\mathbf{w}(t+\Delta t)-\mathbf{w}(t)\approx \frac{\mathrm{d}\mathbf{w}(t)}{\mathrm{d}t}\Delta t$ for a very small $\Delta t$. In computation, we can generate $\mathbf{w}(t)$ by integrating $\mathbf{z}(t)$. We can therefore write

$$\begin{array}{rl}\mathbf{x}(t+\Delta )-\mathbf{x}(t)& =-\tau \nabla f(\mathbf{x}(t))+\mathbf{w}(t+\Delta t)-\mathbf{w}(t)\\ ⟹\mathrm{d}\mathbf{x}& =-\tau \nabla f(\mathbf{x})\mathrm{d}t+\mathrm{d}\mathbf{w}\end{array}$$

This equation reveals a generic form of the SDE, that is

$$\boxed{\mathrm{d}\mathbf{x}=\underset{\text{drift}}{\underbrace{\mathbf{f}(\mathbf{x},t)}}\mathrm{d}t+\underset{\text{diffusion}}{\underbrace{g(t)}}\mathrm{d}\mathbf{w}}$$

where in terms of physics

• The drift coefficient defining how molecules in a closed system would move in the absence of random effects. For the gradient descent algorithm, the drift is defined by the negative gradient of the objective function. That is, we want the solution trajectory to follow the gradient of the objective.
• The diffusion coefficient is a scalar function describing how the molecules would randomly walk from one position to another.

Stochastic Differential Equation for DDPM

We consider the discrete-time DDPM iteration, for $i=1,2,\dots ,N$ (check this equation for details, and here we define ${\beta }_i={\bar{\alpha }}_i$)

$${\mathbf{x}}_i=\sqrt{1-{\beta }_i}{\mathbf{x}}_{i-1}+\sqrt{{\beta }_i}{\mathbf{z}}_{i-1},{\mathbf{z}}_{i-1}\sim \mathcal{N}(0,\text{I})$$

And this sampling equation can be written as an SDE via

$$\mathrm{d}\mathbf{x}=\underset{\text{drift}}{\underbrace{-\frac{\beta (t)}{2}\mathbf{x}}}\mathrm{d}t+\underset{\text{diffusion}}{\underbrace{\sqrt{\beta (t)}}}\mathrm{d}\mathbf{w}$$

This is because if we define a step size $\Delta t=\frac{1}{N}$, and rewrite ${\beta }_i$ as

$${\beta }_i=\beta \left(\frac{i}{N}\right)\cdot \frac{1}{N}=\beta (t+\Delta t)\Delta t$$

where we start by $t=0$ and $i=1$. Similarly, we can define

$$\begin{array}{rl}{\mathbf{x}}_i& =\mathbf{x}\left(\frac{i}{N}\right)=\mathbf{x}(t+\Delta t)\\ {\mathbf{z}}_i& =\mathbf{z}\left(\frac{i}{N}\right)=\mathbf{z}(t+\Delta t)\end{array}$$

Hence, we have

$$\begin{array}{rl}{\mathbf{x}}_i& =\sqrt{1-{\beta }_i}{\mathbf{x}}_{i-1}+\sqrt{{\beta }_i}{\mathbf{z}}_{i-1}\\ ⟹\mathbf{x}(t+\Delta t)& =\sqrt{1-\beta (t+\Delta t)\cdot \Delta t}\mathbf{x}(t)+\sqrt{\beta (t+\Delta t)\cdot \Delta t}\mathbf{z}(t)\\ ⟹\mathbf{x}(t+\Delta t)& \approx \left(1-\frac{1}{2}\beta (t+\Delta t)\cdot \Delta t\right)\mathbf{x}(t)+\sqrt{\beta (t+\Delta t)\cdot \Delta t}\mathbf{z}(t)\\ ⟹\mathbf{x}(t+\Delta t)& \approx \mathbf{x}(t)-\frac{1}{2}\beta (t)\Delta t\mathbf{x}(t)+\sqrt{\beta (t)\Delta t}\mathbf{z}(t)\end{array}$$

As $\Delta \to 0$, we have

$$\mathrm{d}\mathbf{x}=-\frac{1}{2}\beta (t)\mathbf{x}\mathrm{d}t+\sqrt{\beta (t)}\mathrm{d}\mathbf{w}$$

Being able to write the DDPM forward update iteration as an SDE means that the DDPM estimates can be determined by solving the SDE. In other words, for an appropriately defined SDE solver, we can throw the SDE into the solver. The solution returned by an appropriately chosen solver will be the DDPM estimate.