Proximal Implicit ODE Solvers for Accelerating Learning Neural ODEs

The computational bottlenecks of solving the neural ODE (1) and its adjoint ODE (3) come from the numerical stability and numerical accuracy guarantees. Since ${\bm{f}}({\bm{h}}(t),t,\theta)$ is usually high dimensional; direct application of implicit ODE solvers requires solving a system of high dimensional nonlinear equations, which is computationally inefficient, see Section 4.1 for an illustration. As such, explicit solvers — especially the explicit adaptive solvers, e.g., the Dormand-Prince method [15] — are the current default numerical solvers for neural ODE training, test, and inference; see Appendix A for a brief review of explicit adaptive neural ODE solvers. On the one hand, to guarantee numerical accuracy, we have to employ very small step sizes to discretize the neural ODE and its adjoint ODE. On the other hand, both ODEs involved in learning neural ODEs are usually very stiff, which further constrains the step size. Intuitively, a stiff ODE system has very different eigenvalues that reveal dramatically different scales of dynamics. To resolve these dynamics accurately, we have to accommodate for dynamics at significantly different scales, resulting in the use of very small step sizes. To demonstrate this issue, we consider training a recently proposed neural ODE-based GNN, named graph neural diffusion (GRAND) [8], for the benchmark CoauthorCS graph node classification; the detailed experimental settings can be found in Section 4.3. GRAND parametrizes the right-hand side of (1) with the Laplacian operator, resulting in a class of diffusion models. The diffusion model is very stiff from the numerical ODE viewpoint since the ratio between the magnitude of the largest and smallest eigenvalues is infinite. Figure 1 plots the error tolerance of the adaptive solver vs. forward and backward NFEs of different adaptive solvers — including adaptive Heun [49] and two Dormand-Prince methods [15] (DOPRI5 and DOPRI8) — for training GRAND for CoauthorCS node classification. We see that 1) all three adaptive solvers require significant NFEs in solving both neural ODE and its adjoint ODE; as the error tolerance reduces both forward and backward NFEs increase very rapidly. And 2) the higher-order scheme, e.g., DOPRI8 requires more NFEs than the lower-order scheme, indicating that the stability is a dominating factor in choosing the step size. Moreover, it is worth mentioning that both DOPRI5 and DOPRI8 often fail to solve the adjoint ODE when a large error tolerance is used, based on our experiments.

The fundamental issues of learning stiff neural ODEs using explicit solvers due to stability concerns and the high computational cost of direct application of implicit solvers motivate us to study the following question:

Can we modify or relax the implicit solvers to make them suitable for learning neural ODEs with significantly reduced computational costs than existing benchmark solvers?

We answer the above question affirmatively by considering learning neural ODE-style models using the proxy of a few celebrated implicit ODE solvers — including backward Euler, backward differentiation formulas (BDFs), and Crank-Nicolson — leveraging the proximal operator [38]. These proximal solvers reformulate implicit ODE solvers as variational problems. Each solver contains inner-outer iterations, where the inner iterations approximate a one-step update of the implicit solver, and the outer iterations solve the ODE over time. Leveraging fast optimization algorithms for solving the inner optimization problem, these proximal solvers are remarkably faster than explicit adaptive ODE solvers in learning stiff neural ODEs. We summarize the major benefits of proximal solvers below:

• •

  Due to the implicit nature of proximal algorithms, they are much less encumbered by the numerical stability issue than explicit solvers. Therefore, the proximal algorithms allow the use of very large step sizes for neural ODEs training, test, and inference.

• •

  To achieve the same numerical accuracy in solving stiff neural ODEs, proximal algorithms can save significant NFEs in both forward and backward propagation.

• •

  Training neural ODE-style models using proximal algorithms maintains the generalization accuracy of the model compared to using benchmark adaptive solvers.

In this part, we discuss several related works in three directions: reducing NFEs in learning neural ODEs, advances of proximal algorithms, and the development of algorithms for learning neural ODEs.

We denote scalars by lower or upper case letters; vectors and matrices by lower and upper case boldface letters, respectively. We denote the magnitude of a complex number $z$ as $|z|$. For a vector ${\bm{x}}=(x_{1},\ldots,x_{d})^{\top}\in\mathbb{R}^{d}$, we use $\|{\bm{x}}\|:={(\sum_{i=1}^{d}|x_{i}|^{2})^{1/2}}$ to denote its $\ell_{2}$-norm. We denote the vector whose entries are all 0s as $\mathbf{0}$. For two vectors ${\bm{a}}$ and ${\bm{b}}$, we denote their inner product as $\langle{\bm{a}},{\bm{b}}\rangle$. For a matrix ${\bm{A}}$, we use ${\bm{A}}^{\top}$ and ${\bm{A}}^{-1}$ to denote its transpose and inverse, respectively. We denote the identity matrix as ${\bm{I}}$. For a function $f({\bm{x}}):\mathbb{R}^{d}\rightarrow\mathbb{R}$, we denote $\nabla f({\bm{x}})$ and $\nabla^{2}f({\bm{x}})$ as its gradient and Hessian, respectively. We denote $a=\mathcal{O}(b)$, if there is a constant $C$ such that $a\leq Cb$.

We organize the paper as follows: In Section 2, we present the proximal form of a few celebrated implicit ODE solvers for learning neural ODEs. We analyze the numerical stability and convergence of proximal solvers and contrast them with explicit solvers in Section 3. We verify the efficacy of proximal solvers in Section 4, followed by concluding remarks. Technical proofs and more experimental details are provided in the appendix.

We first consider solving a neural ODE (1) using the proximal backward Euler solver. Directly discretizing (1) using the backward Euler scheme with a constant step size $s$ gives

Equation (4) is a system of nonlinear equations, which can be solved efficiently by Newton’s method. But when the dimension of ${\bm{f}}$ is very high, solving (4) can be computationally infeasible. Nevertheless, if ${\bm{f}}({\bm{z}})$ is the gradient of a function $-F({{\bm{z}}})$, we can reformulate the update in (4) using the proximal operator as follows

where ${\bm{f}}({\bm{z}})=-\nabla F({\bm{z}})$, and $\mathcal{A}$ is the admissible domain of ${\bm{z}}$, which is considered to be the whole space in this work; we omit $\mathcal{A}$ in the following discussion. To see why (5) is equivalent to (4), we note that ${\bm{h}}_{k+1}$ is a stationary point of $G({\bm{z}})=\frac{1}{2}\|{\bm{z}}-{\bm{h}}_{k}\|_{2}^{2}+sF({\bm{z}})$, resulting in

replacing ${\bm{z}}$ with ${\bm{h}}_{k+1}$ in the above equation, we get the backward Euler solver (4) (note $\nabla F({\bm{z}})=-{\bm{f}}({\bm{z}})$). Based on the proximal formulation of the backward Euler method (4), we can solve the neural ODE (1) using an inner-outer iteration scheme. The outer iterations perform backward Euler iterations over time, and the inner iterations solve the optimization problem formulated in (5). We summarize the inner-outer iteration scheme for the proximal backward Euler solver in Algorithm 1.

A natural extension of the backward Euler scheme is the single-step, multi-stage schemes whose proximal form are proposed in [60], and we formulated those schemes in Algorithm 2.

Clearly, the 1-stage scheme is the backward Euler scheme. The authors of [60] give conditions on the coefficient $\gamma_{m,i}$. However, it is quite tedious to get those coefficients. For example, a second-order unconditionally stable scheme can be constructed by taking the matrix $\bm{\Lambda}$ — whose entries are $\{\gamma_{m,i}\}_{m,i=1}^{3}$ that appeared in Algorithm 2 — as follows

Furthermore, the matrix $\bm{\Lambda}$ of a third-order scheme is

The main drawback of these multi-stage algorithms for training neural ODEs lies in the computational inefficiency. Note that for an $M$-stage scheme, one need to solve $M$ optimization problems in each iteration.

Another single-step, second-order extension for the backward Euler is the Crank-Nicolson scheme, given by

The Crank-Nicolson scheme is also known as an implicit Runge–Kutta method or trapezoidal rule, which can be formulated in the following proximal form [16]

The backward Euler scheme has first-order accuracy, i.e., discretizing the ODE (1) using (4) with step size $s$ has error $\mathcal{O}(s)$. Backward differentiation formulas (BDFs) are higher-order multi-step methods that can be formulated as follows

where $a_{s}$ is a constant, and ${\bm{A}}_{s}$ is an operator acts on ${\bm{h}}_{k},{\bm{h}}_{k-1},\ldots$. The backward Euler method corresponds to the case when $s=1$, and $a_{s}=1$ and $A_{1}({\bm{h}}_{k})={\bm{h}}_{k}$. The second-order BDF2 scheme can be constructed by taking $a_{s}=\frac{3}{2}$ and ${\bm{A}}_{2}({\bm{h}}_{k},{\bm{h}}_{k-1})=2{\bm{h}}_{k}-\frac{1}{2}{\bm{h}}_{k-1}$. Furthermore, the BDF2 scheme can be written in the following proximal form [34, 16]

if there exists $F({\bm{z}})$ such that ${\bm{f}}({\bm{z}})=-\nabla F({\bm{z}})$.

We can further construct high-order BDFs, for instance the third- and fourth-order BDFs are given below

• •

  BDF3:

  $$\small a_{3}=\frac{11}{6}\ \ \mbox{and}\ \ {\bm{A}}_{3}(\cdot)=3{\bm{h}}_{k}-\frac{3}{2}{\bm{h}}_{k-1}+\frac{1}{3}{\bm{h}}_{k-2}.$$

• •

  BDF4:

  $$\small a_{4}=\frac{25}{12}\ \ \mbox{and}\ \ {\bm{A}}_{4}(\cdot)=4{\bm{h}}_{k}-3{\bm{h}}_{k-1}+\frac{4}{3}{\bm{h}}_{k-2}-\frac{1}{4}{\bm{h}}_{k-3}.$$

The corresponding proximal form of BDF3 and BDF4 are given as follows

• •

  BDF3:

  $\displaystyle{\bm{h}}_{k+1}=$

  $\displaystyle\mathop{\arg\min}_{{\bm{z}}}\left\{\frac{3\|{{\bm{z}}}-{\bm{h}}_{k}\|^{2}}{2s}-\frac{3\|{\bm{z}}-{\bm{h}}_{k-1}\|^{2}}{4s}\right.\left.+\frac{\|{\bm{z}}-{\bm{h}}_{k-2}\|^{2}}{6s}+F({{\bm{z}}})\right\},$

• •

  BDF4:

  $\displaystyle{\bm{h}}_{k+1}$

  $\displaystyle=\mathop{\arg\min}_{{\bm{z}}}\left\{\frac{2\|{{\bm{z}}}-{\bm{h}}_{k}\|^{2}}{s}-\frac{3\|{\bm{z}}-{\bm{h}}_{k-1}\|^{2}}{2s}\right.\left.+\frac{2\|{\bm{z}}-{\bm{h}}_{k-2}\|^{2}}{3s}-\frac{\|{\bm{z}}-{\bm{h}}_{k-2}\|^{2}}{8s}+F({{\bm{z}}})\right\},$

Only one optimization problem needs to be solved in BDFs, which is the advantage of BDF methods over many other high-order implicit schemes.

It is well known that explicit solvers often require a small step size for numerical stability and convergence guarantees, especially for solving stiff ODEs. Consider the linear ODE

Assume ${\bm{A}}$ has spectrum $\sigma({\bm{A}})=\{\lambda_{j}\}_{j=1}^{d}$ and ${\rm Re}(\lambda_{j})<0$ for $j=1,\ldots d$, one can define the stiffness ratio as

We say the ODE is stiff if $Q\gg 1$. The linear stability domain $\mathcal{D}$ of the underlying numerical method is the set of all numbers $z:=s\lambda_{j}$ for $j=1,\ldots,d$, such that $\lim_{k\rightarrow\infty}{\bm{h}}_{k}={\bf 0}$, where ${\bm{h}}_{k}$ is the numerical solution of (13) at the $k$-th step. Table. 1 lists the linear stability domain of some single-step numerical schemes described above.

We call an ODE solver $A$-stable if $\mathbb{C}^{-}:=\{z\in\mathbb{C}~{}|~{}{\rm Re}(z)<0\}\in\mathcal{D}$, where $\mathcal{D}$ is the linear stability domain of the solver. From Table 1, it is clear that the backward Euler and Crank-Nicolson methods are both $A$-stable. For multi-step methods, e.g., BDFs, there is no closed-form expression for the linear stability domain; see [24] for a detailed discussion. Among all BDFs, only BDF2 is $A$-stable. The $A$-stability of ODE solvers is related to the acceleration of the first-order optimization method [30].

Compared to the forward methods, backward methods obtain their next step by solving an implicit equation. Dense linear equations can be solved efficiently using exact methods like QR decomposition. In general, nonlinear equations, or even sparse linear equations, do not have the luxury. At best, an approximation to a solution bounded by a specific error tolerance is attained. Fortunately, a bound of linear growth can be established on the total effect caused by the inner solver error against per step tolerance. For example, we consider a one-dimensional equation $y^{\prime}=f(y)$, and the backward Euler scheme with error can be written as

where $\epsilon_{k}\leq\epsilon$ is the error for the inner loop with $\epsilon>0$ being the tolerance. Let $\hat{y}^{k}$ be the exact backward Euler scheme solution, Then we have

Where $C_{k}=\min_{\xi\in I(y^{k+1},\hat{y}^{k+1})}\lvert 1-sf^{\prime}(\xi)\rvert$, and $I(a,b)$ denote the interval $[\min(a,b),\max(a,b)]$. Hence

which provides a linear upper bound on the error. Similar arguments can be generalized to other solvers and for high-dimensional ODEs.

Besides the linear stability, there is another notion of stability, known as energy stability. We say a numerical method is energy stable if $F({\bm{h}}_{k+1})\leq F({\bm{h}}_{k})$ for all $k$ [48]. As an advantage of the proximal formulation, one can show that the backward Euler is unconditionally energy stable [57] for arbitrary $s$ even for non-convex $F({\bm{z}})$ (here we assume that there exists $F({\bm{z}})$ such that ${\bm{f}}({\bm{z}})=-\nabla F({\bm{z}})$). The energy stability guarantees that the numerical solver takes each step toward minimizing $F({\bm{z}})$, which further lead to the convergence of the numerical scheme [54]. More precisely, we have the following result (see Appendix C for the detailed proof).

Similar energy stability analysis holds for other type of proximal solvers. For instance, for BDF2, one can show there exists ${\bm{h}}_{k+1}$ such that

For Crank-Nicolson, the proximal formulation (12) leads to

These energy stability result can further lead to the convergence of these scheme in the discrete level [34].

We consider solving the following 1D diffusion equation

for $t\in[0,1]$ and $x\in[0,1]$ with a periodic boundary condition. We initialize $u(x,0)$ with the standard Gaussian. To solve the diffusion equation, we first partition $[0,1]$ into uniform grids $\{x_{i}\}_{i=1}^{128}$ and discretize $\partial^{2}h/\partial x^{2}$ using the central finite difference scheme, resulting in the following coupled ODEs

where ${\bm{L}}\in{\mathbb{R}}^{128\times 128}$ is the Laplacian matrix of a cyclic graph scaled by $1/\Delta x^{2}$ with $\Delta x=1/127$ being the spatial resolution, and ${\bm{h}}=[h(x_{1}),\ldots,h(x_{128})]$. Then we apply different ODE solvers to solve (15). It is worth noting that the diffusion equation has an infinite domain of dependence, thus in method of lines discretization, the time step to be taken for explicit methods has to be much finer compared to the spatial discretization to guarantee numerical stability.

We compare our proximal methods against the predominantly used Adaptive Heun and DOPRI5 in solving (15). As the matrix ${\bm{L}}$ is circulant, we can compute the exact solution of (15) using the discrete Fourier transform for comparison. Figure 3 compares the NFEs of proximal methods and two benchmark adaptive solvers with different errors at the final time. For each final error, we use a fixed step size for each proximal solver with FR for solving the inner minimization problem. A list of the configuration of each ODE solver can be found in Table 2 in the appendix. Figure 3 shows that adaptive solvers require many more NFEs than proximal solvers. NFEs required by explicit solvers are almost independent of the error, revealing that the step sizes are constrained by numerical stability. Each iteration of higher-order BDF schemes requires more NFEs than lower-order BDF schemes, showing a tradeoff between controlling numerical error and using high-order schemes.

We train CNFs for MNIST [28] generation using proximal solvers and contrast them to adaptive solvers. In particular, we use the FFJORD approach outlined in [23]. We use an architecture of a multi-scale encoder and two CNF blocks, containing convolutional layers of dimension $64\rightarrow 64\rightarrow 64\rightarrow 64$ with a uniform stride length of 1 and a softplus activation function. We use two scale blocks and train for 10 epochs using Adam optimizer [27] with an initial learning rate of $0.001$. We run the model using the proximal BDFs with FR and DOPRI5 with relative tolerance of $10^{-5}$ and $10^{-7}$, respectively. All the other experimental settings are adapted from [23]. The proximal BDFs are tuned to have an error profile which was consistent with the DOPRI5 as shown in Figure 6 (a). We choose $t_{0}=0$ and $T=1$ and the step sizes for numerical integration and the inner optimizer are 0.2 and 0.3, respectively. We terminate the inner optimization solver once $\|{\bm{z}}^{i+1}-{\bm{z}}^{i}\|\leq 2\times 10^{-3}$. A comparison of the forward and backward NFEs are depicted in Figure 6 (b) and (c), respectively. We see that proximal BDF2, BDF3, and BDF4 require approximately half of the NFEs as DOPRI5 to reach similar training curves.

GRAND is a continuous-depth graph neural network proposed in [8]. It consists of a dense embedding layer, a diffusion layer, and a dense output layer, with ReLU activation function in between. The diffusion layer takes ${\bm{u}}(t_{0})$ as input and solves (16) below to time $T$

where ${\bm{A}}$ is the graph attention function [51] with learnable parameter set $\theta$. This ODE to be solved in GRAND is nonlinear and diffusive, which is particularly challenging to solve numerically.

We compare the performance of proximal solvers with three major adaptive ODE solver options (DOPRI5, DOPRI8, Adaptive Heun) on training the GRAND model for CoauthorCS node classification. The CoauthorCS graph contains 18333 nodes and 81894 edges, and each node is represented by a 6805 dimensional vector. The graph nodes contains 15 classes, and we aim to classify the unlabelled nodes. We choose $t_{0}=0$ and $T=10$ as our starting and ending times of the ODE (16). For each solver, we run the experiment for 50 epochs and record the average NFEs per epoch and accuracy of trained models. For adaptive solvers we choose tolerance $10^{-4}$, and for proximal solvers we choose the integration step size 1.25 (8 steps in total) so that for the first iteration with same input admits output with error within $10^{-3}$. For optimization method we use Fletcher-Reeves with learning rate 0.3, and stop when error is below $10^{-4}$ tolerance. All the other experimental settings are adapted from [8]. Figure 7 shows the advantage of proximal methods over adaptive solvers in training efficiency and maintaining the convergence of the training.