Deriving the Continuity Equation and the Fokker-Planck Equation via Test Functions

测试函数法推导连续性方程和Fokker-Planck方程 苏剑林 · 2023-02-11 · Mathematics probability differential equation random diffusion

In the article Diffusion Models Explained (VI): General Framework — ODE, we derived the Fokker-Planck equation for SDEs; and in Diffusion Models Explained (XII): "Brute-Forcing" the Diffusion ODE, we separately derived the continuity equation for ODEs. Both describe how the distribution of a random variable changes as it evolves along an SDE/ODE — the continuity equation is a special case of the Fokker-Planck equation. When deriving the Fokker-Planck equation, we forced a Taylor expansion onto the Dirac delta function; the result was correct, but the approach felt a bit awkward. When deriving the continuity equation, we combined the Jacobian determinant with a Taylor expansion — a fairly conventional method in itself, but one that doesn't generalize to the Fokker-Planck equation.

In this article we introduce the "test function method," one of the standard techniques for deriving the continuity equation and the Fokker-Planck equation. Its analytical process is more rigorous, and it applies to a broader range of scenarios. more

Integration by parts

Before diving into the formal derivation, let's first introduce a key result we'll need later — the higher-dimensional generalization of integration by parts.

Most textbooks only cover the one-dimensional case of integration by parts, namely

\begin{equation}\int_a^b uv'dx = uv|_a^b - \int_a^b vu'dx\end{equation}

Here $u,v$ is a function of $x$, and $'$ denotes the derivative of the function with respect to $x$. Below we need a higher-dimensional version of this. To that end, let's first recall the derivation of the one-dimensional formula, which relies on the product rule for derivatives:

\begin{equation}(uv)' = uv' + vu'\end{equation}

Then integrating both sides with respect to $x$ and rearranging terms gives the integration-by-parts formula. For the higher-dimensional case, consider the analogous identity:

\begin{equation}\nabla\cdot(u\boldsymbol{v}) = \boldsymbol{v}\cdot\nabla u + u\nabla \cdot\boldsymbol{v}\end{equation}

where $u$ is a scalar function of $\boldsymbol{x}$, and $\boldsymbol{v}$ is a vector-valued function of $\boldsymbol{x}$ (matching the dimension of $\boldsymbol{v}$), and $\nabla$ denotes the gradient of the function with respect to $\boldsymbol{x}$. Now we integrate both sides over a region $\Omega$:

\begin{equation}\int_{\Omega}\nabla\cdot(u\boldsymbol{v})d\boldsymbol{x} = \int_{\Omega}\boldsymbol{v}\cdot\nabla u d\boldsymbol{x} + \int_{\Omega}u\nabla \cdot\boldsymbol{v} d\boldsymbol{x}\end{equation}

By the Gauss divergence theorem, the left-hand side equals $\int_{\partial\Omega}u\boldsymbol{v}\cdot\hat{\boldsymbol{n}}dS$, where $\partial\Omega$ is the boundary of $\Omega$, $\hat{\boldsymbol{n}}$ is the outward unit normal vector at the boundary, and $dS$ is the area element. So, after rearranging terms, we get

\begin{equation}\int_{\Omega}\boldsymbol{v}\cdot\nabla u d\boldsymbol{x} = \int_{\partial\Omega}u\boldsymbol{v}\cdot\hat{\boldsymbol{n}}dS - \int_{\Omega}u\nabla \cdot\boldsymbol{v} d\boldsymbol{x}\label{eq:int-by-parts}\end{equation}

This is the higher-dimensional integration-by-parts formula we wanted to derive. In particular, for a probability density function $p$, the constraints of non-negativity and integrating to one necessarily imply that $p\to 0$ and $\nabla p\to \boldsymbol{0}$ vanish at infinity. So if we take $\Omega$ to be the entire space (unless otherwise specified, the integration domain is assumed to be the whole space), and substitute $u=p$ and $\boldsymbol{v}=\nabla p$ into the above formula, we obtain

\begin{align}\int\boldsymbol{v}\cdot\nabla p d\boldsymbol{x} =&\, - \int p\nabla \cdot\boldsymbol{v} d\boldsymbol{x}\label{eq:int-by-parts-p} \\ \int u\nabla \cdot\nabla p d\boldsymbol{x} = &\,-\int\nabla p\cdot\nabla u d\boldsymbol{x}\label{eq:int-by-parts-gp}\end{align}

If one wants to make the above conclusion fully rigorous, one can assume that $p$ has compact support. But this is purely a matter of mathematical rigor — for a general understanding, it suffices to simply assume that $p\to 0$ and $\nabla p\to \boldsymbol{0}$ hold at infinity.

ODE Evolution

The principle behind the test function method is that if, for any function $\phi(\boldsymbol{x})$, the equation

\begin{equation}\int f(\boldsymbol{x})\phi(\boldsymbol{x})d\boldsymbol{x} = \int g(\boldsymbol{x})\phi(\boldsymbol{x})d\boldsymbol{x}\end{equation}

holds, then $f(\boldsymbol{x})=g(\boldsymbol{x})$ holds, where $\phi(\boldsymbol{x})$ is called the test function. A more rigorous definition would need to specify the space from which $\phi(\boldsymbol{x})$ is chosen, as well as the precise meaning of the equality sign (e.g., exact equality / equality almost everywhere / equality in probability). We won't get into these details here.

For the ODE

\begin{equation}\frac{d\boldsymbol{x}_t}{dt}=\boldsymbol{f}_t(\boldsymbol{x}_t)\label{eq:ode}\end{equation}

we discretize it as

\begin{equation}\boldsymbol{x}_{t+\Delta t} = \boldsymbol{x}_t + \boldsymbol{f}_t(\boldsymbol{x}_t)\Delta t\label{eq:ode-diff}\end{equation}

which gives us

\begin{equation}\phi(\boldsymbol{x}_{t+\Delta t}) = \phi(\boldsymbol{x}_t + \boldsymbol{f}_t(\boldsymbol{x}_t)\Delta t)\approx \phi(\boldsymbol{x}_t) + \Delta t\,\,\boldsymbol{f}_t(\boldsymbol{x}_t)\cdot\nabla_{\boldsymbol{x}_t}\phi(\boldsymbol{x}_t)\end{equation}

Taking the expectation of both sides, we get:

\begin{equation}\int p_{t+\Delta t}(\boldsymbol{x}_{t+\Delta t})\phi(\boldsymbol{x}_{t+\Delta t}) d\boldsymbol{x}_{t+\Delta t}\approx \int p_t(\boldsymbol{x}_t)\phi(\boldsymbol{x}_t)d\boldsymbol{x}_t + \Delta t\int p_t(\boldsymbol{x}_t)\boldsymbol{f}_t(\boldsymbol{x}_t)\cdot\nabla_{\boldsymbol{x}_t}\phi(\boldsymbol{x}_t)d\boldsymbol{x}_t\end{equation}

Since the result of an integral doesn't depend on the notation used for the integration variable, replacing $\boldsymbol{x}_{t+\Delta t}$ with $\boldsymbol{x}_t$ on the left-hand side gives an equivalent expression:

\begin{equation}\int p_{t+\Delta t}(\boldsymbol{x}_t)\phi(\boldsymbol{x}_t) d\boldsymbol{x}_t\approx \int p_t(\boldsymbol{x}_t)\phi(\boldsymbol{x}_t)d\boldsymbol{x}_t + \Delta t\int p_t(\boldsymbol{x}_t)\boldsymbol{f}_t(\boldsymbol{x}_t)\cdot\nabla_{\boldsymbol{x}_t}\phi(\boldsymbol{x}_t)d\boldsymbol{x}_t\label{eq:change-var}\end{equation}

Moving the first term on the right to the left-hand side, and then taking the limit as $\Delta t\to 0$, we get:

\begin{equation}\int \frac{\partial p_t(\boldsymbol{x}_t)}{\partial t}\phi(\boldsymbol{x}_t) d\boldsymbol{x}_t = \int p_t(\boldsymbol{x}_t)\boldsymbol{f}_t(\boldsymbol{x}_t)\cdot\nabla_{\boldsymbol{x}_t}\phi(\boldsymbol{x}_t)d\boldsymbol{x}_t\label{eq:dt-0}\end{equation}

Applying the integration-by-parts formula $\eqref{eq:int-by-parts-p}$ to the right-hand side gives

\begin{equation}\int \frac{\partial p_t(\boldsymbol{x}_t)}{\partial t}\phi(\boldsymbol{x}_t) d\boldsymbol{x}_t = -\int \Big[\nabla_{\boldsymbol{x}_t}\cdot\big(p_t(\boldsymbol{x}_t)\boldsymbol{f}_t(\boldsymbol{x}_t)\big)\Big]\phi(\boldsymbol{x}_t)d\boldsymbol{x}_t\end{equation}

By the equality principle of the test function method, we then have

\begin{equation}\frac{\partial p_t(\boldsymbol{x}_t)}{\partial t} = -\nabla_{\boldsymbol{x}_t}\cdot\big(p_t(\boldsymbol{x}_t)\boldsymbol{f}_t(\boldsymbol{x}_t)\big)\end{equation}

This is called the "continuity equation."

SDE Evolution

For the SDE

\begin{equation}d\boldsymbol{x}_t = \boldsymbol{f}_t(\boldsymbol{x}_t) dt + g_t d\boldsymbol{w}\label{eq:sde}\end{equation}

we discretize it as

\begin{equation}\boldsymbol{x}_{t+\Delta t} = \boldsymbol{x}_t + \boldsymbol{f}_t(\boldsymbol{x}_t) \Delta t + g_t \sqrt{\Delta t}\boldsymbol{\varepsilon},\quad \boldsymbol{\varepsilon}\sim \mathcal{N}(\boldsymbol{0}, \boldsymbol{I})\label{eq:sde-diff}\end{equation}

which gives us

\begin{equation}\begin{aligned} \phi(\boldsymbol{x}_{t+\Delta t}) =&\, \phi(\boldsymbol{x}_t + \boldsymbol{f}_t(\boldsymbol{x}_t) \Delta t + g_t \sqrt{\Delta t}\boldsymbol{\varepsilon}) \\ \approx&\, \phi(\boldsymbol{x}_t) + \left(\boldsymbol{f}_t(\boldsymbol{x}_t) \Delta t + g_t \sqrt{\Delta t}\boldsymbol{\varepsilon}\right)\cdot \nabla_{\boldsymbol{x}_t}\phi(\boldsymbol{x}_t) + \frac{1}{2} \left(g_t\sqrt{\Delta t}\boldsymbol{\varepsilon}\cdot \nabla_{\boldsymbol{x}_t}\right)^2\phi(\boldsymbol{x}_t) \end{aligned}\end{equation}

Taking the expectation of both sides — noting that on the right-hand side we must take the expectation over both $\boldsymbol{x}_t$ and $\boldsymbol{\varepsilon}$, where the expectation over $\boldsymbol{\varepsilon}$ can be computed in advance, giving the result

\begin{equation}\phi(\boldsymbol{x}_t) + \Delta t\,\,\boldsymbol{f}_t(\boldsymbol{x}_t)\cdot \nabla_{\boldsymbol{x}_t}\phi(\boldsymbol{x}_t) + \frac{1}{2} \Delta t\,g_t^2\nabla_{\boldsymbol{x}_t}\cdot\nabla_{\boldsymbol{x}_t}\phi(\boldsymbol{x}_t) \end{equation}

we obtain

\begin{equation}\begin{aligned} &\,\int p_{t+\Delta t}(\boldsymbol{x}_{t+\Delta t})\phi(\boldsymbol{x}_{t+\Delta t}) d\boldsymbol{x}_{t+\Delta t}\\ \approx&\, \int p_t(\boldsymbol{x}_t)\phi(\boldsymbol{x}_t)d\boldsymbol{x}_t + \Delta t\int p_t(\boldsymbol{x}_t)\boldsymbol{f}_t(\boldsymbol{x}_t)\cdot\nabla_{\boldsymbol{x}_t}\phi(\boldsymbol{x}_t)d\boldsymbol{x}_t + \int\frac{1}{2} \Delta t\,g_t^2 p_t(\boldsymbol{x}_t)\nabla_{\boldsymbol{x}_t}\cdot\nabla_{\boldsymbol{x}_t}\phi(\boldsymbol{x}_t) d\boldsymbol{x}_t \end{aligned}\end{equation}

Similar to equations $\eqref{eq:change-var}$ and $\eqref{eq:dt-0}$, taking the limit as $\Delta\to 0$, we get

\begin{equation}\int \frac{\partial p_t(\boldsymbol{x}_t)}{\partial t}\phi(\boldsymbol{x}_t) d\boldsymbol{x}_t = \int p_t(\boldsymbol{x}_t)\boldsymbol{f}_t(\boldsymbol{x}_t)\cdot\nabla_{\boldsymbol{x}_t}\phi(\boldsymbol{x}_t)d\boldsymbol{x}_t + \int\frac{1}{2} \,g_t^2 p_t(\boldsymbol{x}_t)\nabla_{\boldsymbol{x}_t}\cdot\nabla_{\boldsymbol{x}_t}\phi(\boldsymbol{x}_t) d\boldsymbol{x}_t\end{equation}

Applying formula $\eqref{eq:int-by-parts-p}$ to the first term on the right-hand side, and applying formula $\eqref{eq:int-by-parts-gp}$ followed by formula $\eqref{eq:int-by-parts-p}$ to the second term, we get

\begin{equation}\int \frac{\partial p_t(\boldsymbol{x}_t)}{\partial t}\phi(\boldsymbol{x}_t) d\boldsymbol{x}_t = \int \left[-\nabla_{\boldsymbol{x}_t}\cdot\big(p_t(\boldsymbol{x}_t)\boldsymbol{f}_t(\boldsymbol{x}_t)\big)+\frac{1}{2}g_t^2 \nabla_{\boldsymbol{x}}\cdot\nabla_{\boldsymbol{x}}p_t(\boldsymbol{x})\right]\phi(\boldsymbol{x}_t)d\boldsymbol{x}_t\end{equation}

By the equality principle of the test function method, we obtain

\begin{equation}\frac{\partial p_t(\boldsymbol{x}_t)}{\partial t} = -\nabla_{\boldsymbol{x}_t}\cdot\big(p_t(\boldsymbol{x}_t)\boldsymbol{f}_t(\boldsymbol{x}_t)\big)+\frac{1}{2}g_t^2 \nabla_{\boldsymbol{x}}\cdot\nabla_{\boldsymbol{x}}p_t(\boldsymbol{x})\end{equation}

This is the "Fokker-Planck equation."

Summary

This article introduced the test function method for deriving certain probability evolution equations. Its main content includes the higher-dimensional generalization of integration by parts, along with the derivations of the continuity equation for ODEs and the Fokker-Planck equation for SDEs.