Exact continuous adjoint sensitivity for a neural ODE

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