FreeFlow: A Unified Gradient Flow Framework

• FreeFlow is a gradient-flow framework for diffusion probabilistic models that recasts probability evolution as a steepest descent in the 2-Wasserstein space.
• It unifies stochastic (SDE-based) and deterministic (ODE-based) approaches under a common formulation rooted in optimal transport and free-energy minimization.
• The framework employs both Lagrangian and Eulerian perspectives, offering analytical insights into DPM dynamics, pathologies, and potential model improvements.

FreeFlow refers to a gradient-flow-based theoretical framework for diffusion probabilistic models (DPMs), rooted in the principles of optimal transport and Riemannian geometry. This framework recasts the time evolution of probability densities in DPMs as a unique gradient flow in the 2-Wasserstein space, presenting a mathematically unified view that bridges stochastic and deterministic models. FreeFlow offers a comprehensive explanation of DPM dynamics, elucidates the underlying geometrical structures, and provides insight into both algorithmic mechanisms and potential model pathologies (Sun et al., 2023).

1. Mathematical Formulation and Core Principle

The central insight of FreeFlow is that the evolution of a time-dependent probability density $\rho_t(x)$ on $\mathbb{R}^n$ can be expressed as a gradient flow of a free-energy functional $\mathcal{F}[\rho]$ in the 2-Wasserstein space $\mathcal{P}(\mathbb{R}^n)$. This is formalized as:

$$\partial_t \rho_t = -\mathrm{grad}_{\rho_t}\mathcal{F} = \nabla \cdot [D_t \nabla \rho_t + \rho_t \nabla \Psi]$$

where $D_t$ is a (possibly time-varying) diffusion coefficient, $\Psi$ a potential, and the functional

$$\mathcal{F}[\rho] = \int_{\mathbb{R}^n} D_t\, \rho(x) \log \rho(x) + \rho(x) \Psi(x) \,dx$$

serves as a free energy landscape. The kinetic energy for a curve of densities is

$$A(\rho, v) = \int_0^1 \int_{\mathbb{R}^n} \rho_t(x) |v_t(x)|^2 \, dx\,dt$$

subject to the continuity equation constraint

$\partial_t \rho_t + \nabla \cdot (\rho_t v_t) = 0.$

The Benamou–Brenier theorem yields that the infimum of $A(\rho, v)$ equals the squared 2-Wasserstein distance between the endpoints, establishing the metric structure.

2. Lagrangian and Eulerian Descriptions

FreeFlow maintains a duality between Lagrangian (particle-based) and Eulerian (density-based) perspectives. The Lagrangian approach tracks the evolution of individual sample trajectories via:

$$\frac{d}{dt} \gamma_x(t) = v_t(\gamma_x(t)),\qquad \gamma_x(0)=x$$

corresponding to SDE or ODE representations:

$$dX_t = \xi_t(X_t)dt + \sigma_t dW_t \qquad \text{(SDE)}$$

$dX_t = f_t(X_t)dt \qquad \text{(ODE)}$

The Eulerian perspective instead describes the entire probability density evolution through (forward) PDEs such as the Fokker–Planck equation:

$$\partial_t \rho + \nabla \cdot (\rho v) = D_t \Delta \rho.$$

FreeFlow identifies both perspectives as instantiations of the same infinite-dimensional Riemannian gradient flow, allowing rigorous transition between pointwise and distributional dynamics.

3. Unified Representation of Stochastic and Deterministic DPMs

The FreeFlow PDE subsumes both stochastic and deterministic DPMs via choices of $D_t$ and $\Psi$:

• DDPM / VP-SDE: $D_t = \beta_t$, $\Psi(x) = \frac{1}{2}\beta_t |x|^2$. The resulting gradient flow coincides with the Fokker–Planck equation and with the SDE

$dX_t = -\beta_t X_t dt + \sqrt{2\beta_t} dW_t.$

• VE-SDE: $D_t = \dot\alpha_t$, $\Psi \equiv 0$. The stochastic process is purely diffusive,

$$dX_t = \sqrt{2\dot\alpha_t} dW_t,\qquad \partial_t\rho = \dot\alpha_t \Delta \rho.$$

• ODE-based samplers (DDIM, DPM-Solver, PFGM, GenPhys): $D_t = 0$ and $\nabla \Psi(x) = -f_t(x)$, yielding the deterministic continuity equation

$$\partial_t \rho + \nabla \cdot (\rho f_t) = 0,\qquad dX_t = f_t(X_t)dt.$$

This unification demonstrates that every major DPM variant is a manifestation of steepest descent in probability space with respect to a suitably chosen energy functional (Sun et al., 2023).

4. Riemannian Geometry and the Wasserstein Metric

The space $\mathcal{P}(\mathbb{R}^n)$ of probability densities is equipped with a 2-Wasserstein Riemannian metric, endowing it with an infinite-dimensional differential structure. The tangent space at a density $\rho$ comprises zero-mean scalar fields, and the infinitesimal distance is defined by

$$\left\| \partial_t \rho \right\|_\rho = \inf_{\partial_t \rho + \nabla \cdot (\rho v) = 0} \int \rho |v|^2$$

with inner product

$$\langle s_1, s_2 \rangle_\rho = \int_{\mathbb{R}^n} \rho(x) (\nabla \varphi_1(x) \cdot \nabla \varphi_2(x)) dx$$

when $s_i = \nabla \cdot (\rho \nabla \varphi_i)$. This geometric structure allows for precise formulation and analysis of DPM dynamics as energy-dissipating flows.

5. Analytical and Practical Implications

Formulating DPMs as Wasserstein gradient flows provides a common foundation for analysis, revealing both methodological parallels and unique failure modes. For instance, the Lagrangian description can develop shock waves upon intersecting particle paths (non-invertibility), while the Eulerian PDE remains well-posed when $v_t$ is Lipschitz. Monitoring $\nabla \cdot v$ enables diagnosis of such pathologies. All DPMs can be seen as energy minimization processes, and their learned drift/score terms can be interpreted as approximations to the steepest descent in $\mathcal{P}(\mathbb{R}^n)$. This connection offers avenues for principled model improvement and generalization.

6. Extensions and Future Research Directions

The FreeFlow formalism is extensible: the quadratic transport cost $c(x,y)=\frac{1}{2}|x-y|^2$ can be generalized to more complex Riemannian metrics or higher-order flows (potentially fourth-order), and additional curvature terms can be incorporated into $\mathcal{F}[\rho]$. This enables the design of non-Euclidean score-based models and information-geometric diffusion processes directly in parameter space. Such extensions have the potential to yield more expressive and adaptive generative models, integrating advanced geometric and analytic techniques into DPM theory and practice (Sun et al., 2023).

Table: Main Unification in FreeFlow

Model Class	$D_t$	Potential $\Psi(x)$
DDPM / VP-SDE	$\beta_t$	$\frac{1}{2}\beta_t|x|^2$
VE-SDE	$\dot\alpha_t$	$0$
ODE-based Sampler	$0$	determined by $f_t(x)$ via $\nabla\Psi = -f_t$

The unified perspective offered by FreeFlow advances both the theoretical foundation and practical development of diffusion probabilistic models, clarifying existing algorithms and signaling directions for geometric and functional generalization.