Unified Diffusion Theory

Module: Generative AI | Difficulty: Advanced

Forward SDE

Reverse SDE (Anderson, 1982)

Fokker-Planck Equation

Probability Flow ODE

Connection Between All Three

class UnifiedDiffusion:
    def __init__(self, score_model, f, g, T=1000):
        self.score = score_model
        self.f, self.g, self.T = f, g, T
    def reverse_sde(self, x, t, dt):
        score = self.score(x, t)
        drift = self.f(x, t) - self.g(t)**2 * score
        diffusion = self.g(t)
        return drift * dt + diffusion * torch.randn_like(x) * (dt**0.5)
    def probability_flow(self, x, t, dt):
        score = self.score(x, t)
        drift = self.f(x, t) - 0.5 * self.g(t)**2 * score
        return x + drift * dt
    def sample_sde(self, x_init):
        x = x_init
        dt = -1.0 / self.T
        for i in range(self.T):
            t = torch.full((x.size(0),), i/self.T)
            x = self.reverse_sde(x, t, dt)
        return x
    def sample_ode(self, x_init):
        x = x_init
        dt = -1.0 / self.T
        for i in range(self.T):
            t = torch.full((x.size(0),), i/self.T)
            x = self.probability_flow(x, t, dt)
        return x

VP SDE (Variance Preserving)

VE SDE (Variance Exploding)

Sub-VP SDE

Research Insight: The Fokker-Planck equation provides a deterministic characterization of the entire probability flow, enabling analysis of mode connectivity, barrier heights, and convergence rates without sampling.