Neural ODEs and continuous normalizing flows

Let $d,q\ge 1$, let $I\subseteq \mathbb{R}$ and $\mathrm{\Theta }\subseteq {\mathbb{R}}^q$ be open, and let $v:I\times {\mathbb{R}}^d\times \mathrm{\Theta }\to {\mathbb{R}}^d$. For a parameter $\theta \in \mathrm{\Theta }$, write $v_{\theta }(t,z):=v(t,z,\theta )$. The parameterized ordinary differential equation on $I$ is the initial-value problem

$\dot{z}(t)=v_{\theta }(t,z(t)),z(s)=z_s,$

where $s\in I$ and $z_s\in {\mathbb{R}}^d$. A trajectory on an interval $J\subseteq I$ containing $s$ is a differentiable function $z:J\to {\mathbb{R}}^d$ satisfying both displayed conditions. Here $v_{\theta }$ is the state-velocity field.

It is a parameterized neural ordinary differential equation, or neural ODE, when $v_{\theta }(t,z)$ is specified by a neural network with input $(t,z)$ and $d$ outputs. Here a network is a fixed finite directed acyclic graph evaluated in an order in which every edge points to a later vertex: input vertices hold the coordinates of $(t,z)$, each other vertex applies a fixed scalar activation to an affine combination of earlier vertex values, and the output vertices use the identity activation. The weights and biases are specified functions of $\theta$ and may be shared or fixed. No particular depth, width, or hidden activation is prescribed; affine networks with no hidden vertices are allowed.

The general parameterized ODE remains defined for fields without such a network specification. This definition alone asserts neither solution existence beyond its stated interval nor numerical accuracy.

Chen et al., equation (2) and continuous normalizing flows

Papamakarios et al., §4.1

Let $T>0$, let $d,q\ge 1$, let $\mathrm{\Theta }\subseteq {\mathbb{R}}^q$ be open, let $I\subseteq \mathbb{R}$ be an open interval containing $[0,T]$, and let $v:I\times {\mathbb{R}}^d\times \mathrm{\Theta }\to {\mathbb{R}}^d$ be jointly $C^{\mathrm{\infty }}$. For every compact $K\subseteq \mathrm{\Theta }$, suppose there are constants $L_K,M_K\ge 0$ such that, for all $t\in [0,T]$, $\theta \in K$, and $z,w\in {\mathbb{R}}^d$,

$\mathrm{\parallel }v(t,z,\theta )-v(t,w,\theta ){\mathrm{\parallel }}_2\le L_K\mathrm{\parallel }z-w{\mathrm{\parallel }}_2,\mathrm{\parallel }v(t,0,\theta ){\mathrm{\parallel }}_2\le M_K\mathrm{.}$

Then for each $\theta \in \mathrm{\Theta }$, $s,t\in [0,T]$, and $z_s\in {\mathbb{R}}^d$, the parameterized ODE has a unique solution on all of $[0,T]$. Write ${\mathrm{\Phi }}_{t,s}^{\theta }(z_s)$ for its value at time $t$. For each fixed $s,t$, the map $(z_s,\theta )\mapsto {\mathrm{\Phi }}_{t,s}^{\theta }(z_s)$ is locally $C^{\mathrm{\infty }}$. For each fixed $\theta$, ${\mathrm{\Phi }}_{t,s}^{\theta }:{\mathbb{R}}^d\to {\mathbb{R}}^d$ is a $C^{\mathrm{\infty }}$ diffeomorphism with inverse ${\mathrm{\Phi }}_{s,t}^{\theta }$. Its state Jacobian $J_{t,s}:=D_{z_s}{\mathrm{\Phi }}_{t,s}^{\theta }(z_s)$ satisfies

$\frac{d}{dt}{J}_{t,s}=D_{z}v(t,{\mathrm{\Phi }}_{t,s}^{\theta }(z_s),\theta )J_{t,s},J_{s,s}=I_d\mathrm{.}$

Time derivatives at $0$ and $T$ are one-sided.

Teschl, Ordinary Differential Equations and Dynamical Systems, Chapter 2, §§2.2–2.5

Assume Countable Choice. Let $T>0$, let $d,q\ge 1$, let $\mathrm{\Theta }\subseteq {\mathbb{R}}^q$ be open, let $I\subseteq \mathbb{R}$ be an open interval containing $[0,T]$, and let $v:I\times {\mathbb{R}}^d\times \mathrm{\Theta }\to {\mathbb{R}}^d$ be jointly $C^{\mathrm{\infty }}$. For every compact $K\subseteq \mathrm{\Theta }$, assume there are constants $L_K,M_K\ge 0$ such that, for all $t\in [0,T]$, $\eta \in K$, and $z,w\in {\mathbb{R}}^d$,

$\mathrm{\parallel }v(t,z,\eta )-v(t,w,\eta ){\mathrm{\parallel }}_2\le L_K\mathrm{\parallel }z-w{\mathrm{\parallel }}_2,\mathrm{\parallel }v(t,0,\eta ){\mathrm{\parallel }}_2\le M_K\mathrm{.}$

Fix $\theta \in \mathrm{\Theta }$ and let ${\mathrm{\Phi }}_t:={\mathrm{\Phi }}_{t,0}^{\theta }:{\mathbb{R}}^d\to {\mathbb{R}}^d$ be its flow. Let $p_0:{\mathbb{R}}^d\to (0,\mathrm{\infty })$ be a continuous probability density with respect to Lebesgue measure, define ${\mu }_0(E):={\int }_E{p}_0(z)dz$ for Borel $E$, and define the pushforward measure ${\mu }_t(E):={\mu }_0({\mathrm{\Phi }}_t^{-1}(E))$. For $y\in {\mathbb{R}}^d$, set

$p_t(y):=\frac{p_0({\mathrm{\Phi }}_t^{-1}(y))}{\mathrm{\mid }\det D{\mathrm{\Phi }}_t({\mathrm{\Phi }}_t^{-1}(y))\mathrm{\mid }}\mathrm{.}$

Then $p_t$ is a density of ${\mu }_t$. For every $z_0\in {\mathbb{R}}^d$, write $z_t:={\mathrm{\Phi }}_t(z_0)$ and $J_t:=D{\mathrm{\Phi }}_t(z_0)$; the determinant of $J_t$ is positive and

$p_t({\mathrm{\Phi }}_t(z_0))=\frac{p_0(z_0)}{\det J_t}>0,\frac{d}{dt}\log p_t(z_t)=-\mathrm{tr}(D_{z}v(t,z_t,\theta ))=-{\sum }_{i=1}^d{\mathrm{\partial }}_{z_i}{v}_i(t,z_t,\theta ),$

The derivative identity holds for $0<t<T$ and has one-sided meanings at $0$ and $T$.

Chen et al., §4 and Appendix A

Papamakarios et al., §4.1

Let $T>0$, let $d,q\ge 1$, let $\mathrm{\Theta }\subseteq {\mathbb{R}}^q$ be open, and let $v:I\times {\mathbb{R}}^d\times \mathrm{\Theta }\to {\mathbb{R}}^d$ be jointly $C^{\mathrm{\infty }}$ on an open time interval $I$ containing $[0,T]$. For each compact $K\subseteq \mathrm{\Theta }$, suppose there are constants $L_K,M_K\ge 0$ such that for all $t\in [0,T]$, $\theta \in K$, and $z,w\in {\mathbb{R}}^d$,

$\mathrm{\parallel }v(t,z,\theta )-v(t,w,\theta ){\mathrm{\parallel }}_2\le L_K\mathrm{\parallel }z-w{\mathrm{\parallel }}_2,\mathrm{\parallel }v(t,0,\theta ){\mathrm{\parallel }}_2\le M_K\mathrm{.}$

Let $\mathrm{\Omega }\subseteq \mathrm{\Theta }$ be open with compact closure contained in $\mathrm{\Theta }$, let $U_0\subseteq \mathrm{\Theta }$ be an open neighborhood of $\bar{\mathrm{\Omega }}$, and let $u_0:U_0\to {\mathbb{R}}^d$ be $C^1$. Write $u_{\theta }$ for the exact solution of ${\dot{u}}_{\theta }(t)=v(t,u_{\theta }(t),\theta )$, $u_{\theta }(0)=u_0(\theta )$, and let $\mathcal{T}:=\{(t,u_{\theta }(t),\theta ):t\in [0,T],\theta \in \bar{\mathrm{\Omega }}\}$ be its common compact trajectory tube. Let

${\mathcal{T}}_T:=\{(u_{\theta }(T),\theta ):\theta \in \bar{\mathrm{\Omega }}\}$

be the terminal state-parameter slice. Let $U_{\varphi }\subseteq {\mathbb{R}}^d\times \mathrm{\Theta }$ and $U_{\mathrm{\ell }}\subseteq I\times {\mathbb{R}}^d\times \mathrm{\Theta }$ be open sets containing ${\mathcal{T}}_T$ and $\mathcal{T}$, respectively, and let $\varphi :U_{\varphi }\to \mathbb{R}$ and $\mathrm{\ell }:U_{\mathrm{\ell }}\to \mathbb{R}$ be $C^1$. Assume their first derivatives, measured in Euclidean operator norms, are bounded and uniformly continuous on ${\mathcal{T}}_T$ and $\mathcal{T}$, respectively. Define

$\mathcal{J}(\theta )=\varphi (u_{\theta }(T),\theta )+{\int }_0^T\mathrm{\ell }(t,u_{\theta }(t),\theta )dt\mathrm{.}$

For each $\theta \in \mathrm{\Omega }$, the exact continuous adjoint sensitivity procedure is:

1. Solve the state initial-value problem forward on $[0,T]$.
2. Solve the terminal adjoint problem backward, $-p^{\mathrm{\prime }}(t)=D_{u}v(t,u_{\theta }(t),\theta {)}^{T}p(t)+D_u\mathrm{\ell }(t,u_{\theta }(t),\theta {)}^T,p(T)=D_u\varphi (u_{\theta }(T),\theta {)}^T\mathrm{.}$
3. Return the vector $G(\theta )=D_{\theta }\varphi (u_{\theta }(T),\theta {)}^T+D_{\theta }{u}_0(\theta {)}^{T}p(0)+{\int }_0^T(D_{\theta }\mathrm{\ell }(t,u_{\theta }(t),\theta {)}^T+D_{\theta }v(t,u_{\theta }(t),\theta {)}^{T}p(t))dt\mathrm{.}$

Here $D_u$ and $D_{\theta }$ denote Jacobian matrices in the displayed variables, $p(t)\in {\mathbb{R}}^d$, and $G(\theta )\in {\mathbb{R}}^q$. All state and adjoint equations and the integral are exact continuous mathematical objects. Replacing them by numerical solver outputs defines a finite computation and is outside this algorithm's accuracy guarantees.

Chen et al., §2 and Appendix B

Papamakarios et al., §4.2.2