A Self-Contained Tutorial on Ito Stochastic Differential Equations for Diffusion Models

In the machine learning community in the context of Diffusion Models, there has been a big interest in Stochastic Differential Equations (SDEs) recently. SDEs are a powerful concept driving powerful generative AI tools from image generation (see DALL-E-2 or Stable Diffusion) to protein generation (see RF-Diffusion). However, the accompanying papers assume understanding of complex results from the theory of SDEs - such as the time-reversal formula. As a result, the works are more and more inaccessible for non-SDE experts. This blog post aims to help out here and give a self-contained introduction to SDEs for diffusion models. All mathematical results for SDEs that are used for diffusion models are proven here. More specifically, in this tutorial, you will know:

1. What SDEs are.
2. How SDEs are implemented.
3. The distribution and convergence properties of SDEs with affine drift coefficients.
4. An overview of famous diffusion models and how they can be framed as SDEs with affine drift coefficients.
5. How SDEs can be time-reversed and how that corresponds to novel generation.
6. How to train a diffusion model.

First time hear about SDEs?: If you have never heard about SDEs, I would recommend you to first read my previous tutorial where I explore the Langevin SDE that is fundamental for machine learning.

More interested in practical implementations? in my next tutorial, we are going to put the theory into practice.

Required background: I assume that you have background knowledge in probability theory.

Acknowledgement: the content of this tutorial has partially grown out of several nights trying to understand the famous diffusion model SDE paper by Song et al..

import numpy as np
import matplotlib.pyplot as plt
from typing import Callable, List
from itertools import product
from tqdm.notebook import tqdm
#!pip install ipywidgets

from celluloid import Camera # getting the camera
import matplotlib.pyplot as plt
import numpy as np
from IPython.display import HTML
import seaborn as sns
from IPython.display import Video
import os
import pandas as pd

1. Introduction: What are Stochastic Differential Equations?

The solutions of a stochastic differential equation (SDE) is a stochastic process. A stochastic process is a collection of random variables $X_t\in\mathbb{R}^d$ that evolve over the time $t\geq 0$. While each $X_t$ is random for each individual $t$, the beauty of stochastic processes comes in studying how $X_{t+s}$ is related to $X_{t}$ - i.e. how values at different time points depend on each other.

1.1. Stochastic Processes and Brownian Motion

A fundamental stochastic process is a Brownian motion $W_t$. A Brownian motion is characterized by:

1. Normal increments: Increments are normally distributed with variance scaling proportional to the time difference: $W_{t+s}-W_{t}\sim\mathcal{N}(0,s\mathbf{I}_d)$ for all $t,s\geq 0$.
2. Independent increments: $W_{t_1}-W_{t_2}$ is independent of $W_{t_2}-W_{t_3}$ for $t_1>t_2>t_3$.

We can sample trajectories of a Brownian motion by discretizing the stochastic process by: $$W_{t+s} = W_{t} +\sqrt{s} \epsilon\quad\text{where }\epsilon\sim\mathcal{N}(0,1)$$

Let's plot 5 trajectories:

def run_brownian_motion(n_traj: int = 5, n_grid_points: int = 10000, t_start: float = 0.0, t_end: float = 10.0):
    time_grid = np.linspace(t_start,t_end,n_grid_points)
    grid_size = (t_end-t_start)/n_grid_points
    x_0 = np.array([n_traj*[0.0]])
    gaussian_noise = np.random.normal(size=(len(time_grid)-1,n_traj))
    brownian_motion = np.sqrt(grid_size)*gaussian_noise.cumsum(axis=0)
    brownian_motion = np.concatenate([x_0,brownian_motion],axis=0)
    return time_grid,brownian_motion

fig,ax = plt.subplots(figsize=(9,9))
time,X_t = run_brownian_motion()
sample_labels = [f"Trajectory {i}" for i in range(1,6)]
ax.plot(time,X_t,label=sample_labels)
ax.set_xlabel("time t")
ax.set_ylabel("W_t")
ax.legend();

If you feel reminded of a graph describing a stock market index, good guess! Brownian motion is widely used in Mathematical Finance to model stock prices.

1.2 From Ordinary Differential Equations to Stochastic Differential Equations

To understand SDE, let's start by reviewing the classical ordinary differential equations (ODEs). An ODE is an equation with the following characteristics:

1. Its solution $\mathbf{x}(t)$ is a deterministic function: $x:\mathbb{R}\to\mathbb{R}^d$.
2. The rate of change $$\frac{d}{dt}\mathbf{x}(t) = f(\mathbf{x}(t),t)$$ is time- and location-dependent and given by a function $$f:\mathbb{R}^d\times\mathbb{R}\to\mathbb{R}^d,(\mathbf{x},t)\to f(\mathbf{x},t)$$.
3. A fixed start value: $x(0)=v\in\mathbb{R}^d$.

A Stochastic Differential Equation "is an ODE but just random". In other words, an SDE is an equation with the following characteristics:

1. Its solution $X_t$ is a stochastic process. I.e. a sample of a stochastic process is a function $t\to X_t\in\mathbb{R}^d$.
2. The rate of change $$dX(t) = f(X_t,t)dt+g(t)dW_t$$ is given by
   • A. A deterministic drift specified by a function $f$.
   • B. A random drift specified by a function $g\geq 0$ and a Brownian motion $W_t$.
3. A start distribution $X_0\sim p_0$.

Intuitively, the above equation means that for (infinitisimally) small $s>0$, we have

$$X_{t+s}\approx X_{t}+s f(X_t,t)+g(t)\sqrt{s}(W_{t+s}-W_{t})$$

In other words, the conditional distribution for the next time step $X_{t+s}$ given the current time step $X_{t}$ is given by $$X_{t+s}|X_t\sim\mathcal{N}(X_{t}+s f(X_t,t),sg^2(t))$$

So the deterministic drift $f$ describes the infinitesimal change in mean and the volatility function $g$ describes the infinitesimal standard deviation.

2. Implementing SDEs with the Euler-Maruyama method

In this section, we simulate SDEs and look at a few example SDEs.

To simulate SDEs, we discretize them in time and then sampling the next time step with the above update rule:

• Input: number of steps $n_{\text{steps}}$, step size $s>0$
• Sample $X_0\sim p_0$
• Set $t=0$

• For $i=1,...n_{\text{steps}}$:

  Sample $\epsilon\sim\mathcal{N}(0,I)$

  Set $X_{t+s} = X_{t} + sf(X_{t},t)+g(t)\sqrt{s}\epsilon$

• Return: $[X_{0},X_{s},X_{2s},X_{3s},\dots,X_{sn_{\text{steps}}}]$

The above method is called the Euler-Maruyama method.

def run_sde(f_determ_drift: Callable,
            g_random_drift: Callable,
            x_start: np.array,
            t_start: float = 0.0, 
            t_end: float = 1.0, 
            n_steps: int = 10000,
            **kwargs):
    """Function to run stochastic differential equation. We assume a deterministic initial distribution p_0."""
    
    #Number of trajectories, dimension of data:
    n_traj,dim_x = x_start.shape

    #Compute time grid for discretization and step size:
    time_grid = np.linspace(t_start,t_end,n_steps)
    step_size = time_grid[1]-time_grid[0]

    #Compute the random drift at every time point:
    random_drift_grid = g_random_drift(time_grid)
    
    #Sample random drift at every time point:
    noise = np.random.normal(size=(n_steps,n_traj,dim_x))
    random_drift_grid_sample = np.sqrt(step_size)*random_drift_grid[:,None,None]*noise
    
    #Initialize list of trajectory:
    x_traj = [x_start]
    
    
    for idx,time in tqdm(enumerate(time_grid)):
        
        #Get last location and time
        x = x_traj[idx]
        t = float(time_grid[idx])
        
        #Get deterministic drift and random drift sample
        determ_drift = step_size*f_determ_drift(x,t)
        random_drift_sample = random_drift_grid_sample[idx]
        
        #Compute next step:
        next_step = x + determ_drift + random_drift_sample
        
        #Save step:
        x_traj.append(next_step)

    return np.stack(x_traj),time_grid    

Note: There are fancier (but much less intuitive) methods to compute solutions other than the Euler-Maryuama method. This is beyond the scope of this tutorial.

2.1. Define Example Random Drift Functions g

Let's define two example drift functions and plot them:

• Linear: $g(t)=0.1+0.05t$
• Periodic: $g(t)=\sin(2\pi t)+1.10$

def linear_g(t,constant=0.2):
    """t - 1d np.ndarray or float
    Returns: 1d np.ndarray or float"""
    return 0.1+constant*t

def periodic_g(t):
    """t - 1d np.ndarray or float
    Returns: 1d np.ndarray or float"""
    return np.sqrt(np.sin(2*np.pi*t)+1.10)

def plot_g_random_drift(drift_func_list: List[Callable], min_t: float = 0.0, max_t: float = 5.0,n_grid_points: int = 10000):
    """Function to plot a random drift function g."""
    fig, axs = plt.subplots(1,len(drift_func_list),figsize=(len(drift_func_list)*6,6))
    time = np.linspace(min_t,max_t,n_grid_points)
    for idx,drift_func in enumerate(drift_func_list):
        axs[idx].plot(time,drift_func(time))
        axs[idx].set_title(drift_func.__name__)
        axs[idx].set_xlabel("time")
        axs[idx].set_ylabel("g")

plot_g_random_drift([linear_g,periodic_g])

2.2. Define Example Deterministic Drift Functions f

Let's define example drift functions $f$:

• Shift: $f(x,t)=-(x-4)$
• Periodic: $f(x,t)=\cos(2\pi t^2)\cos(2\pi x^2)$

def shift_f(x,t,shift=4.0):
    """x - shape (n,k): a space coefficients 
       t - shape (n) or float: time coefficients"""
    return -(x-shift)

def periodic_f(x,t):
    """x - shape (n,k): space coefficients 
       t - shape (n) or float: time coefficients"""

    if isinstance(t,float):
        return 2*np.cos(2*np.pi*t**2)*(np.cos(2*np.pi*x)) #*normalizer[:,None])
    else:
        return 2*np.cos(2*np.pi*t**2)[:,None]*(np.cos(2*np.pi*x**2)) #*normalizer[:,None])

def plot_f_determ_drift(determ_func_list: List[Callable],
                min_t: float = 0.0, 
                max_t: float = 2.5, 
                min_x: float = 0.0,
                max_x: float = 2.5,
                n_grid_points: int = 1000,
                plot_contours=False):
    """Function to plot the function f(x,t) in an Ito-SDE."""
    time = np.linspace(min_t,max_t,n_grid_points)
    oned_grid = np.linspace(min_x, max_x, n_grid_points)
    twod_grid = np.array([[x,t] for t,x in product(time,oned_grid)])
    extent = [min_x, max_x, min_t, max_t]
    
    fig,axs = plt.subplots(1,len(determ_func_list),figsize=(6*len(determ_func_list),6))
    
    for idx, determ_func in enumerate(determ_func_list):
        derivative = determ_func(twod_grid[:,0].reshape(-1,1),twod_grid[:,1])
        axs[idx].imshow(derivative.reshape(n_grid_points,n_grid_points).transpose(),interpolation='bilinear',origin='lower', extent = extent, cmap=plt.get_cmap('YlOrRd'))
        if plot_contours:
            axs[idx].contour(derivative.reshape(n_grid_points,n_grid_points).transpose(),interpolation='bilinear',origin='lower', extent = extent,color='black')
        axs[idx].set_title(determ_func.__name__)
        axs[idx].set_xlabel("time t")
        axs[idx].set_ylabel("location x")

plot_f_determ_drift([shift_f,periodic_f])

2.3 Animation of SDEs

Let's write a function that animates SDEs over time

def animate_ito_sde(f_drift: Callable, g_drift: Callable, fpath_anim: str = None, n_steps: int = 100000, n_samples: int = 400,n_grid_points: int = 100, initial_dist:str ="normal", t_end=4.0):
    """Function to animate and plot a 1d SDE over time."""
    if initial_dist == "normal":
        x_start = np.random.normal(size=n_samples).reshape(-1,1)
    elif initial_dist == "uniform":
        x_start = np.random.uniform(size=n_samples).reshape(-1,1)
    else:
        raise ValueError
        
    x_traj,time_grid = run_sde(f_drift,g_drift,x_start=x_start,t_end=t_end,n_steps=n_steps)
    x_traj = x_traj.squeeze()
    
    fig, axs = plt.subplots(3,1,figsize=(24,36))
    camera = Camera(fig)
    
    oned_grid = np.linspace(x_traj.min(), x_traj.max(), len(time_grid))
    twod_grid = np.array([[x,t] for t,x in product(time_grid,oned_grid)])
    extent = [time_grid.min(), time_grid.max(),x_traj.min(), x_traj.max()]
    derivative = f_drift(twod_grid[:,0].reshape(-1,1),twod_grid[:,1])
    derivative_plot = derivative.reshape(len(time_grid),len(time_grid)).transpose()
    
    g_drift_per_time = g_drift(time_grid)

    for idx in tqdm(range(1,len(x_traj),max(int(len(x_traj)/200),1))):
    
        #Plot evolution over time over deterministic drift:
        axs[0].imshow(derivative_plot,interpolation='bilinear',origin='lower', extent = extent, cmap=plt.get_cmap('YlOrRd'),aspect='auto')
        if idx>5:
            for idy in range(10):
                axs[0].plot(time_grid[:idx],x_traj[:idx,idy].squeeze()) #,label=sample_labels)
            axs[0].plot(time_grid[:idx],x_traj[:idx].squeeze().mean(axis=1),color='black',linewidth=10)
        axs[0].set_title("Example trajectories (color=f)")

        #Plot distribution:
        sns.kdeplot(x_traj[idx],ax=axs[1])
        axs[1].set_title("Distribution of X_t")

        #Plot random drift:
        axs[2].plot(time_grid[:idx],g_drift_per_time[:idx])
        axs[2].set_title("Random drift g")
        camera.snap()
        
    animation = camera.animate() # animation ready
    plt.close()
    return animation

Set hyperparameters.

RUN_ANIMATION = False
N_STEPS = 1000
N_SAMPLES = 10000

2.4. Run first simple SDE

Let's set:

• $X_0\sim\mathcal{N}(0,1)$
• $f(x,t)=-(x-4)$
• $g(t)=\sin(2\pi t)+1.10$

and run the SDE.

The below animation shows:

• Top figure: 5 example trajectories (plot over the function $f(x,t)$)
• Middle figure: the change of distribution over time
• Bottom figure: the change in drift $g(t)$ over time.

fpath = f"simple_sde.mp4"
animation = animate_ito_sde(shift_f,periodic_g,n_steps=N_STEPS,n_samples=N_SAMPLES)
animation.save(fpath)
HTML(animation.to_html5_video())

0it [00:00, ?it/s]

  0%|          | 0/200 [00:00<?, ?it/s]

We can observe that:

• The mean $\mathbb{E}[X_t]$ converges towards $4$ over time (black line in animation in 1st figure). The reason for that is that $f(x,t)=-(x-4)$ pushes $x$ towards $4$.
• The curves becomes "rougher" or noisier periodically due to the periodic $g(t)=\sin(2\pi t)+1.10$ (see plot 2 and 3 above).
• In combination of both effects, the distribution moves from a Gaussian distribution $\mathcal{N}(0,1)$ towards a Gaussian distribution $\mathcal{N}(4,\sigma^2(t))$ with mean 4 and variance $\sigma^2(t)$ shifting periodically (see plot 1 above).

2.5. Run complex SDE

Let's run a more complex SDE with:

• $X_0\sim\mathcal{N}(0,1)$
• $f(x,t)=\cos(2\pi t^2)\cos(2\pi x^2)$
• $g(t)=0.1+0.1t$

fpath = f"complex_sde.mp4"
animation = animate_ito_sde(periodic_f,linear_g,n_steps=N_STEPS,n_samples=N_SAMPLES)
animation.save(fpath)
HTML(animation.to_html5_video())
# HTML(f"""
# <video muted autoplay width="1280" height="960" controls>
#   <source src="complex_sde.mp4" type="video/mp4">
# </video>
# """)

0it [00:00, ?it/s]

  0%|          | 0/200 [00:00<?, ?it/s]