MALI: A memory efficient and reverse accurate integrator for Neural ODEs

Notations   We summarize the notations following Zhuang et al. (2020).

• •

  $z_{i}(t_{i})/\overline{z}(\tau_{i})$: hidden state in forward/reverse time trajectory at time $t_{i}/\tau_{i}$.

• •

  $\psi_{h}(t_{i},z_{i})$: the numerical solution at time $t_{i}+h$, starting from $(t_{i},z_{i})$ with a stepsize $h$.

• •

  $N_{f},N_{z}$: $N_{f}$ is the number of layers in $f$ in Eq. 1, $N_{z}$ is the dimension of $z$.

• •

  $N_{t}/N_{r}$: number of discretized points (outer iterations in Algo. 1) in forward / reverse integration.

• •

  $m$: average number of inner iterations in Algo. 1 to find an acceptable stepsize.

Numerical Integration   The algorithm for general adaptive-stepsize numerical ODE solvers is summarized in Algo. 1 (Wanner & Hairer, 1996). The solver repeatedly advances in time by a step, which is the outer loop in Algo. 1 (blue curve in Fig. 1). For each step, the solver decreases the stepsize until the estimate of error is lower than the tolerance, which is the inner loop in Algo. 1 (green curve in Fig. 1). For fixed-stepsize solvers, the inner loop is replaced with a single evaluation of $\psi_{h}(t,z)$ using predefined stepsize $h$. Different methods typically use different $\psi$, for example different orders of the Runge-Kutta method (Runge, 1895).

We compare different numerical implementations of the analytical form in this section. The forward-pass and backward-pass of different methods are demonstrated in Fig. 1 and Fig. 2 respectively. Forward-pass is similar for different methods. The comparison of backward-pass among different methods are summarized in Table. 1. We explain methods in the literature below.

Naive method  The naive method saves all of the computation graph (including search for optimal stepsize, green curve in Fig. 2) in memory, and backpropagates through it. Hence the memory cost is $N_{z}N_{f}\times N_{t}\times m$ and depth of computation graph are $N_{f}\times N_{t}\times m$, and the computation is doubled considering both forward and backward passes. Besides the large memory and computation, the deep computation graph might cause vanishing or exploding gradient (Pascanu et al., 2013).

Theorem 2.1 is stated as Theorem 3.2 in Zhuang et al. (2020); please see reference paper for detailed proof. To summarize, due to inevitable errors with numerical ODE solvers, the reverse-time trajectory (red curve, $\overline{z}(\tau)$) cannot match the forward-time trajectory (blue curve, $z(t)$) accurately. The error in $\overline{z}$ propagates to $\frac{\mathrm{d}L}{\mathrm{d}\theta}$ by Eq. 2, hence affects the accuracy in gradient estimation.

Adaptive checkpoint adjoint (ACA)  To solve the inaccuracy of adjoint method, Zhuang et al. (2020) proposed ACA: ACA stores forward-time trajectory in memory for backward-pass, hence guarantees accuracy; ACA deletes the search process (green curve in Fig. 2), and only back-propagates through the accepted step (blue curve in Fig. 2), hence has a shallower computation graph ($N_{f}\times N_{t}$ for ACA vs $N_{f}\times N_{t}\times m$ for naive method). ACA only stores $\{z(t_{i})\}_{i=1}^{N_{t}}$, and deletes the computation graph for $\{f\big{(}z(t_{i}),t_{i}\big{)}\}_{i=1}^{N_{t}}$, hence the memory cost is $N_{z}(N_{f}+N_{t})$. Though the memory cost is much smaller than the naive method, it grows linearly with $N_{t}$, and can not handle very high dimensional models. In the following sections, we propose a method to overcome all these disadvantages of existing methods.

In this section we give a brief introduction to the asynchronous leapfrog (ALF) method (Mutze, 2013), and we provide theoretical analysis which is missing in Mutze (2013). For general first-order ODEs in the form of Eq. 1, the tuple $(z,t)$ is sufficient for most ODE solvers to take a step numerically. For ALF, the required tuple is $(z,v,t)$, where $v$ is the “approximated derivative”. Most numerical ODE solvers such as the Runge-Kutta method (Runge, 1895) track state $z$ evolving with time, while ALF tracks the “augmented state” $(z,v)$. We explain the details of ALF as below. Input $(z_{in},v_{in},s_{in},h)$ where $s_{in}$ is current time, $z_{in}$ and $v_{in}$ are correponding values at time $s_{in}$, $h$ is stepsize. Forward    $s_{1}=s_{in}+h/2$                $k_{1}=z_{in}+v_{in}\times h/2$                $u_{1}=f(k_{1},s_{1})$                $v_{out}=v_{in}+2(u_{1}-v_{in})$                $z_{out}=k_{1}+v_{out}\times h/2$                $s_{out}=s_{1}+h/2$ Output      $(z_{out},v_{out},s_{out},h)$ Algorithm 2 Forward of $\psi$ in ALF Input $(z_{out},v_{out},s_{out},h)$ where $s_{out}$ is current time, $z_{out}$ and $v_{out}$ are corresponding values at $s_{out}$, $h$ is stepsize. Inverse      $s_{1}=s_{out}-h/2$                $k_{1}=z_{out}-v_{out}\times h/2$                $u_{1}=f(k_{1},s_{1})$                $v_{in}=2u_{1}-v_{out}$                $z_{in}=k_{1}-v_{in}\times h/2$                $s_{in}=s_{1}-h/2$ Output      $(z_{in},v_{in},s_{in},h)$ Algorithm 3 $\psi^{-1}$ (Inverse of $\psi$) in ALF

Procedure of ALF  Different ODE solvers have different $\psi$ in Algo. 1, hence we only summarize $\psi$ for ALF in Algo. 2. Note that for a complete algorithm of integration for ALF, we need to plug Algo. 2 into Algo. 1. The forward-pass is summarized in Algo. 2. Given stepsize $h$, with input $(z_{in},v_{in},s_{in})$, a single step of ALF outputs $(z_{out},v_{out},s_{out})$.

As in Fig. 3, given $(z_{0},v_{0},t_{0})$, the numerical forward-time integration calls Algo. 2 iteratively:

Invertibility of ALF  An interesting property of ALF is that $\psi$ defines a bijective mapping; therefore, we can reconstruct $(z_{in},v_{in},s_{in},h)$ from $(z_{out},v_{out},s_{out},h)$, as demonstrated in Algo. 7. As in Fig. 3, we can reconstruct the entire trajectory given the state $(z_{j},v_{j})$ at time $t_{j}$, and the discretized time points $\{t_{0},...t_{N_{t}}\}$. For example, given $(z_{N_{t}},v_{N_{t}})$ and $\{t_{i}\}_{i=0}^{N_{t}}$, the trajectory for Eq. 3.1 is reconstructed:

In the following sections, we will show the invertibility of ALF is the key to maintain accuracy at a constant memory cost to train Neural ODEs. Note that “inverse” refers to reconstructing the input from the output without computing the gradient, hence is different from “back-propagation”.

Initial value   For an initial value problem (IVP) such as Eq. 1, typically $z_{0}=z(t_{0})$ is given while $v_{0}$ is undetermined. We can construct $v_{0}=f(z(t_{0}),t_{0})$, so the initial augmented state is $(z_{0},v_{0})$.

Difference from midpoint integrator   The midpoint integrator (Süli & Mayers, 2003) is similar to Algo. 2, except that it recomputes $v_{in}=f(z_{in},s_{in})$ for every step, while ALF directly uses the input $v_{in}$. Therefore, the midpoint method does not have an explicit form of inverse.

Proof is in Appendix A.4 and A.5. Theorem 3.2 implies the following: when $\eta=1$, the damped ALF reduces to ALF, and the stability region is empty; when $0<\eta<1$, the stability region is non-empty. However, stability describes the behaviour when $T$ goes to infinity; in practice we always use a bounded $T$ and ALF performs well. Inverse of damped ALF is in Appendix A.5.

Procedure of MALI  Details of MALI are summarized in Algo. 4. For the forward-pass, we only keep the end-time state $(z_{N_{t}},v_{N_{t}})$ and the accepted discretized time points (blue curves in Fig. 1 and  2). We ignore the search process for optimal stepsize (green curve in Fig. 1 and 2), and delete other variables to save memory. During the backward pass, we can reconstruct the forward-time trajectory as in Eq. 5, then calculate the gradient by numerical discretization of Eq. 2 and Eq. 3.

Constant memory cost w.r.t number of solver steps in integration  We delete the computation graph and only keep the end-time state to save memory. The memory cost is $N_{z}(N_{f}+1)$, where $N_{z}N_{f}$ is due to evaluating $f(z,t)$ and is irreducible for all methods. Compared with the adjoint method, MALI only requires extra $N_{z}$ memory to record $v_{N_{t}}$, and also has a constant memory cost w.r.t time step $N_{t}$. The memory cost is $N_{z}(N_{f}+1)$.

Accuracy   Our method guarantees the accuracy of reverse-time trajectory (e.g. blue curve in Fig. 2 matches the blue curve in Fig. 1), because ALF is explicitly invertible for free-form $f$ (see Algo. 7). Therefore, the gradient estimation in MALI is more accurate compared to the adjoint method.

Shallow computation graph  Similar to ACA, MALI only backpropagates through the accepted step (blue curve in Fig. 2) and ignores the search process (green curve in Fig. 2), hence the depth of computation graph is $N_{f}\times N_{t}$. The computation graph of MALI is much shallower than the naive method, hence is more robust to vanishing and exploding gradients (Pascanu et al., 2013).

Summary  The adjoint method suffers from inaccuracy in reverse-time trajectory, the naive method suffers from exploding or vanishing gradient caused by deep computation graph, and ACA finds a balance but the memory grows linearly with integration time. MALI achieves accuracy in reverse-time trajectory, constant memory w.r.t integration time, and a shallow computation graph.

We validate MALI on image recognition tasks using Cifar10 and ImageNet datasets. Similar to Zhuang et al. (2020), we modify a ResNet18 into its corresponding Neural ODE: the forward function is $y=x+f_{\theta}(x)$ and $y=x+\int_{0}^{T}f_{\theta}(z)\mathrm{d}t$ for the residual block and Neural ODE respectively, where the same $f_{\theta}$ is shared. We compare MALI with the naive method, adjoint method and ACA.

Accuracy on ImageNet  Due to the heavy memory burden caused by large images, the naive method and ACA are unable to train a Neural ODE on ImageNet with 4 GPUs; only MALI and the adjoint method are feasible due to the constant memory. We also compare the Neural ODE to a standard ResNet. As shown in Fig. 6, the accuracy of the Neural ODE trained with MALI closely follows ResNet, and significantly outperforms the adjoint method (top-1 validation: 70% v.s. 63%).

Invariance to discretization scheme  A continuous model should be invariant to discretization schemes (e.g. different types of ODE solvers) as long as the discretization is sufficiently accurate. We test the Neural ODE using different solvers without re-training; since ResNet is often viewed as a one-step Euler discretization of an ODE (Haber & Ruthotto, 2017), we perform similar experiments. As shown in Table 2, Neural ODE consistently achieves high accuracy ($\sim$70%), while ResNet drops to random guessing ($\sim$0.1%) because ResNet as a one-step Euler discretization fails to be a meaningful dynamical system (Queiruga et al., 2020).

Robustness to adversarial attack  Hanshu et al. (2019) demonstrated that Neural ODE is more robust to adversarial attack than ResNet on small-scale datasets such as Cifar10. We validate this result on the large-scale ImageNet dataset. The top-1 accuracy of Neural ODE and ResNet under FGSM attack (Goodfellow et al., 2014) are summarized in Table 3. For Neural ODE, due to its invariance to discretization scheme, we derive the gradient for attack using a certain solver (row in Table 3), and inference on the perturbed images using various solvers. For different combinations of solvers and perturbation amplitudes, Neural ODE consistently outperforms ResNet.

Summary   In image recognition tasks, we demonstrate Neural ODE is accurate, invariant to discretization scheme, and more robust to adversarial attack than ResNet. Note that detailed explanation on the robustness of Neural ODE is out of the scope for this paper, but to our knowledge, MALI is the first method to enable training of Neural ODE on large datasets due to constant memory cost.

We apply MALI to latent-ODE (Rubanova et al., 2019) and Neural Controlled Differential Equation (Neural CDE) (Kidger et al., 2020a; b). Our experiment is based on the official implementation from the literature. We report the mean squared error (MSE) on the Mujoco test set in Table 4, which is generated from the “Hopper” model using DeepMind control suite (Tassa et al., 2018); for all experiments with different ratios of training data, MALI achieves similar MSE to ACA, and both outperform the adjoint and naive method. We report the test accuracy on the Speech Command dataset for Neural CDE in Table 5; MALI achieves a higher accuracy than competing methods.

We apply MALI on FFJORD (Grathwohl et al., 2018), a free-from continuous generative model, and compare with several variants in the literature (Finlay et al., 2020; Kidger et al., 2020a). Our experiment is based on the official implementaion of Finlay et al. (2020); for a fair comparison, we train with MALI, and test with the same solver as in the literature (Grathwohl et al., 2018; Finlay et al., 2020), the Dopri5 solver with $\text{rtol}=\text{atol}=10^{-5}$ from the torchdiffeq package (Chen et al., 2018). Bits per dim (BPD, lower is better) on validation set for various datasets are reported in Table 6. For continuous models, MALI consistently generates the lowest BPD, and outperforms the Vanilla FFJORD (trained with adjoint), RNODE (regularized FFJORD) and the SemiNorm Adjoint (Kidger et al., 2020a). Furthermore, FFJORD trained with MALI achieves comparable BPD to state-of-the-art discrete-layer flow models in the literature. Please see Sec. B.3 for generated samples.