Do Residual Neural Networks discretize Neural Ordinary Differential Equations?

Introduction

On the other hand, a Neural ODE (Chen et al., 2018) uses a neural network $\varphi_{\Theta}(x,s)$, that takes time $s$ into account, to parameterise a vector field (Kidger, 2022) in a differential equation, as follows,

Neural ODEs also provide a theoretical framework to study deep learning models from the continuous viewpoint, using the arsenal of ODE theory (Teh et al., 2019; Li et al., 2019; Teshima et al., 2020). Importantly, they can also be seen as the continuous analog of ResNets. Indeed, consider for $N$ an integer, the Euler scheme for solving Eq. (2) with time step $\frac{1}{N}$ starting from $x_{0}$ and iterating $x_{n+1}=x_{n}+\frac{1}{N}\varphi_{\Theta}(x_{n},\frac{n}{N})$. Under mild assumptions on $\varphi_{\Theta}$, this scheme is known to converge to the true solution of Eq. (2) as $N$ goes to $+\infty$. Also, if $\Theta=(\theta_{n}^{N})_{i\in[N-1]}$ and $\varphi_{\Theta}(.,\frac{n}{N})=f(.,\theta_{n}^{N})$, then the ResNet equation Eq. (1) corresponds to a Euler discretization with time step $\frac{1}{N}$ of Eq. (2). However, for a given ResNet with fixed depth $N$ and weights, the activations in Eq. (1) can be far from the solution of Eq. (2). This is illustrated in Figure 1 where we show that a deep ResNet can easily break the topology of the input space, which is impossible for a Neural ODE. In this paper, we study the link between ResNets and Neural ODEs. We make the following contributions:

In Section 3, we propose a framework to define a set of associated Neural ODEs for a given ResNet. We control the error between the discrete and the continuous trajectory. We show that without additional assumptions on the smoothness with depth of the residual functions, this error does not go to as $N\to\infty$ (Prop. 1). However, we show that under some assumptions on the weight initialization, the trained parameters of a deep linear ResNet uniformly (with respect to both depth and training time) approach a Lipschitz function as the depth $N$ of the network goes to infinity, at speed $\frac{1}{N}$ (Prop. 2 and Th. 1). This result highlights an implicit regularization towards a limit Neural ODE.

In Section 4, we investigate a simple technique to train ResNets without storing activations. Inspired by the adjoint method, we propose to recover the approximated activations during the backward pass by using a reverse-time Euler scheme. We control the error for recovering the activations and gradients with this method. We show that if the residuals of the ResNet are bounded and Lipschitz continuous, with constants independent of $N$, then this error scales in $O(\frac{1}{N})$ (Prop. 3). Hence, the adjoint method needs a large number of layers to lead to correct gradients (Prop. 4). We then consider a smoothness-dependent reconstruction with Heun’s method to bound the error between the true and approximated gradient by a term that depends on $\frac{1}{N}$ times the smoothness in depth of the residual functions, hence guaranteeing a better approximation when successive weights are close one to another (Prop. 5 and 6).

In Section 5, on the experimental side, we show that the adjoint method fails when training a ResNet 101 on ImageNet. Nevertheless, we empirically show that very deep ResNets pretrained with tied weights (constant weights: $\theta^{N}_{n}=\theta$ $\forall n$) can be refined -using our adjoint method- on CIFAR-10 and ImageNet by untying their weights, leading to a better test accuracy. Last, but not least, we show using a ResNet architecture with heavy downsampling in the first layer that our adjoint method succeeds at large depth and that Heun’s method leads to a better behaved training, hence confirming our theoretical results.

Background and related work

Neural ODEs are a class of implicit deep learning models defined by an ODE where a neural network parameterises the vector field (Weinan, 2017; Chen et al., 2018; Teh et al., 2019; Sun et al., 2018; Weinan et al., 2019; Lu et al., 2018; Ruthotto and Haber, 2019; Kidger, 2022). Given an input $x_{0}$, the output of the model is the solution of the ODE (2) at time $1$. From a theoretical viewpoint, the expression capabilities of Neural ODEs have been investigated in (Cuchiero et al., 2020; Teshima et al., 2020; Li et al., 2019) and the Neural ODE framework has been used to better understand the dynamics of more general architectures that include residual connections such as Transformers (Sander et al., 2022; Lu et al., 2019). Experimentaly, Neural ODEs have been successful in a various range of applications, among which physical modelling (Greydanus et al., 2019; Cranmer et al., 2019) and generative modeling (Chen et al., 2018; Grathwohl et al., 2018). However, there are many areas where Neural ODEs have failed to replace ResNets, for instance for building computer vision classification models. Neural ODEs fail to compete with ResNets on ImageNet, and to the best of our knowledge, previous works using Neural ODEs on ImageNet consider weight-tied architectures and only achieves the same accuracy as a ResNet18 (Zhuang et al., 2021).

Implicit Regularization of ResNets towards ODEs.

Recent works have studied the link between ResNets and Neural ODEs. In (Cohen et al., 2021), the authors carry experiments to better understand the scaling behavior of weights in ResNets as a function of the depth. They show that under the assumption that there exists a scaling limit $\theta(s)=N^{\beta}\lim\theta^{N}_{\lfloor Ns\rfloor}$ for the weights of the ResNets (with $0<\beta<1$) and if the scale of the ResNet is $\frac{1}{N^{\alpha}}$ with $0<\alpha<1$ and $\alpha+\beta=1$, then the hidden state of the ResNet converges to a solution of a linear ODE. In this paper, we are interested in the case where $\alpha=1$, which seems more natural since it is the scaling that appears in Euler’s method with step $\frac{1}{N}$. In addition, we do not assume the existence of a scaling limit $\theta(s)=\lim\theta^{N}_{\lfloor Ns\rfloor}$. In subsection 3.2, we demonstrate the existence of this scaling limit in the linear setting, under some assumptions. The recent work (Cont et al., 2022) shows results regarding linear convergence of gradient descent in ResNets and prove the existence of an $\frac{1}{2}$-Hölder continuous scaling limit as $N\to\infty$ with a scaling factor for the residuals in $\frac{1}{\sqrt{N}}$ which is different from ours. In contrast, we show that our limit function is Lipschitz continuous, which is a stronger regularity. We also show that our convergence is uniform in depth and optimization time. More generally, recent works have proved the convergence of gradient descent training of ResNet when the initial loss is small enough. This include ResNet with finite width but arbitrary large depth (Du et al., 2019; Liu et al., 2020) and ResNet with both infinite width and depth (Lu et al., 2020; Barboni et al., 2021). These convergence proofs leverage an implicit bias toward weights with small amplitudes. They however leave open the question of convergence of individual weights as depth increases, which we tackle in this work in the linear case. This requires showing an extra bias toward weights with small variations across depth.

Memory bottleneck in ResNets.

Training deep learning models involve graphics processing units (GPUs) where memory is a practical bottleneck (Wang et al., 2018; Peng et al., 2017; Zhu et al., 2017). Indeed, backpropagation requires to store activations at each layer during the forward pass. Since samples are processed using mini batches, this storage can be important. For instance, with batches of size 128, the memory needed to compute gradients for a ResNet 152 on ImageNet is about 22 GiB. Note that the memory needed to store the parameters of the model is only 220 MiB, which is negligible compared to the memory needed to store the activations. Thus, designing deep invertible architectures where one can recover the activations on the fly during the backpropagation iterations has been an active field in recent years (Gomez et al., 2017; Sander et al., 2021a; Jacobsen et al., 2018). In this work, we propose to approximate activations using a reverse-time Euler scheme, as we detail in the next subsection.

Adjoint Method.

Consider a loss function $L(x_{N})$ for the ResNet (1). The backpropagation equations (Baydin et al., 2018) are

Now, consider a loss function $L(x(1))$ for the Neural ODE (2). The adjoint state method (Pontryagin, 1987; Chen et al., 2018) gives

Note that if $\Theta=(\theta_{n}^{N})_{n\in[N-1]}$ and $\varphi_{\Theta}(.,\frac{n}{N})=f(.,\theta_{n}^{N})$, then Eq. (3) corresponds to a Euler discretization with time step $\frac{1}{N}$ of Eq. (4). The key advantage of using Eq. (4) is that one can recover $x(s)$ on the fly by solving the Neural ODE (2) backward in time starting from $x(1)$. This strategy avoids storing the forward trajectory $(x(s))_{s\in}$ and leads to a $O(1)$ memory footprint (Chen et al., 2018). In this work, we propose to use a discrete adjoint method by using a reverse-time Euler scheme for approximately recovering the activations in a ResNet (Section 4). Contrarily to other models such as RevNets (Gomez et al., 2017) (architecture change) or Momentum ResNets (Sander et al., 2021b) (forward rule modification) which rely on an exactly invertible forward rule, the proposed method requires no change at all in the network, but gives approximate gradients.

Notations.

ResNets as discretization of Neural ODEs

In this section we first show that without further assumptions, the distance between the discrete trajectory and the solution of associated ODEs can be constant with respect to the depth of the network if the residual functions lack smoothness with depth. We then present a positive result by studying the linear case where we show that, under some hypothesis (small loss initialization and initial smoothness with depth), the ResNet converges to a Neural ODE as the number of layers goes to infinity. We show that this convergence is uniform with depth and optimization time.

We first define associated Neural ODEs for a given ResNet.

Note that we omit the dependency of $\Theta$ in $N$ to simplify notations. For example, for a given ResNet, there are two natural ways to interpolate it with a Neural ODE, either by interpolating the residuals, or by interpolating the weights. Indeed, one can interpolate the residuals with $\varphi_{\Theta}(\cdot,s)=(n+1-Ns)f(.,\theta_{n}^{N})+(Ns-n)f(\cdot,\theta_{n+1}^{N})$ when $s\in[\frac{n}{N},\frac{n+1}{N}]$, or interpolate the weights with $\varphi_{\Theta}(\cdot,s)=f(\cdot,(n+1-Ns)\theta_{n}^{N}+(Ns-n)\theta_{n+1}^{N})$ for $s\in[\frac{n}{N},\frac{n+1}{N}]$. If $\theta_{n}^{N}=\theta^{N}$ does not depend on $n$, then both interpolations are identical and one can simply consider $\varphi_{\Theta}(x,s)=f(x,\theta^{N})$, $\forall(x,s)$.

We now consider any smooth interpolation $\varphi_{\Theta}$ for the ResNet (1) and a Euler scheme for the Neural ODE (2) with time step $\frac{1}{N}$.

2 Linear Case

Gradient.

Two continuous variables involved.

Assumption 1 is the classical assumption in the literature (Zou et al., 2020; Barboni et al., 2021) to prove linear convergence of our loss and that the $\theta^{N}_{n}(t)$’s stay bounded with $t$. Note that this bounded norm assumption implies that $\frac{1}{N}\theta_{n}(0)=O(\frac{1}{N})$. This is in contrast with classical initialization scales in the feedforward case where the initialization only depends on width (He et al., 2015). However this initialization scale is coherent with those of ResNets for which the scale has to depend on depth (Yang and Schoenholz, 2017). In addition, the experimental findings in Cohen et al. (2021) suggest that the weights in ResNets scale in $\frac{1}{N^{\beta}}$ with $\beta>0$.

We now prove an implicit regularization result showing that if at initialization, in addition to assumption 1, the weights are close from one another ($O(\frac{1}{N})$), they will stay at distance $O(\frac{1}{N})$: the discrete derivative stay in $O(\frac{1}{N})$, which is a central result to consider the infinite depth limit in our Th. 1.

Lemma 1 is proved in appendix A.3, and gives us the existence of a Lipschitz continuous accumulation point, but not the uniqueness nor the convergence speed. For the uniqueness, we show in appendix A.5 that, under the assumptions of Th. 1, one has that any accumulation point of $\psi_{\sigma}$ satisfies the limit Neural ODE

and show that $F$ satisfies the hypothesis of the Picard–Lindelöf theorem, hence showing the uniqueness of $\psi$. We finally show that, as intuitively expected, trajectories of the weights of our linear ResNets of depth $N$ and $2N$ remain close one to each other. This gives the convergence speed in Th. 1. See appendix A.4 for a proof.

Adjoint Method in Residual Networks

The approximate recovery of the activations in Eq. (6) is implementable for any ResNet: there is no need for particular architecture or forward rule modification. The drawback is that the recovery is only approximate. We devote the remainder of the section to the study of the corresponding errors and to error reduction using second order Heun’s method. We first show that, if $f(.,\theta^{N}_{n})$ and its derivative are bounded by a constant independent of $N$, then the error for reconstructing the activations in the backward scheme (6) is $O(\frac{1}{N})$. Proofs of the theoretical results are in appendix A.

Then the error made by reconstructing the activations is in $O(\frac{1}{N})$.

Prop. 3 shows a slow convergence of the error for recovering activations. This bound does not depend on the discrete derivative $f(.,\theta_{n+1}^{N})-f(.,\theta_{n}^{N})$, contrarily to the errors between the ResNet activations and the trajectory of the interpolating Neural ODE in Prop 1. In summary, even though regularity in depth is necessary to imply closeness to a Neural ODE, it is not necessary to recover activations, and neither gradients, as we now show.

Error in gradients when using the adjoint method.

We use the result obtained in Prop. 3 to derive a bound in $O(\frac{1}{N^{2}})$ on the error made for computing gradients using formulas (7).

For a proof, see appendix A.7, where we give the dependency of our upper bound as a function of $\Delta,L_{f},C_{f},\Omega$ and $L_{df}$.

Smoothness-dependent reconstruction with Heun’s method.

The bounds in Prop. 3 and 4 do not depend on the smoothness with respect to the weights of the $f(.,\theta_{n}^{N})$. Only the magnitude of the residuals plays a role in the correct recovery of the activations and estimation of the gradient. Hence, there is no apparent benefit of having such a network behave like a Neural ODE. We now turn to Heun’s method, a second order integration scheme, and show that in this case smoothness in depth of the network improves activation recovery. A HeunNet (Maleki et al., 2021) of depth $N$ with parameters $\theta_{1}^{N},\dots,\theta_{N}^{N}$ iterates for $n=0,\dots,N-1$:

These forward iterations can once again be approximately reversed by doing for $n=N-1,\dots,0$:

which also enables approximated backpropagation without storing activations. When discretizing an ODE, Heun’s method has a better $O(\frac{1}{N^{2}})$ error, hence we expect a better recovery than in Prop. 3. Indeed, we have:

Just like with activation, we see that Heun’s method allows for a better gradient estimation when the weights are smooth with depth. Equivalently, for a fixed depth, this proposition indicates that HeunNets have a better estimation of the gradient with the adjoint method than ResNets which ultimately leads to better training and overall better performances by such memory-free model.

Experiments

We now present experiments to investigate the applicability of the results presented in this paper. We use Pytorch (Paszke et al., 2017) and Nvidia Tesla V100 GPUs. Our code will be open-sourced. All the experimental details are given in appendix B, and we provide a recap on ResNet architectures in appendix C.

The ResNet model (1) is different from the classical ResNet because of the $\frac{1}{N}$ term. This makes the model depth aware, and we want to study the impact of this modification on the accuracy on CIFAR and ImageNet.

We first train a ResNet-101 (He et al., 2016a) on CIFAR-10 and ImageNet using the same hyper-parameters. Experimental details are in appendix B and results are summarized in table 1, showing that the explicit addition of the step size $\frac{1}{N}$ does not affect accuracy. In strike contrast, the classical ResNet rule without the scaling $\frac{1}{N}$ makes the network behave badly at large depth, while it still works well with our scaling $\frac{1}{N}$, as shown in Figure 3 (a). On ImageNet, the scaling $\frac{1}{N}$ also leads to similar test accuracy in the weight tied setting: $72.5\%$ with $4$ blocks per layer, $73.2\%$ with $8$ blocks per layer and $72\%$ with $16$ blocks per layer (mean over $2$ runs).

2 Adjoint method

Our results in Prop. 3 and 4 assume uniform bounds in $N$ on our residual functions and their derivatives. We also formally proved in the linear setting that these assumptions hold during the whole learning process if the initial loss is small. A natural idea to start from a small loss is to consider a pretrained model.

In addition, we also want our pretrained model to verify assumption 2 so we consider the following setup. On CIFAR (resp. ImageNet) we train a ResNet with 4 (resp. 8) blocks in each layer, where weights are tied within each layer. A first observation is that one can transfer these weights to deeper ResNets without significantly affecting the test accuracy of the model: it remains above $94.5\%$ on CIFAR-10 and $72\%$ on ImageNet. We then untie the weights of our models and refine them. More precisely, for CIFAR, we then transfer the weights of our model to a ResNet with $4$, $4$, $64$ and $4$ blocks within each layer and fine-tune it only by refining the third layer, using our adjoint method. We display in table 2 the median of the new test accuracy, over $5$ runs for the initial pretraining of the model. For ImageNet, we transfer the weights to a ResNet with $100$ blocks per layer and fine-tune the whole model with our adjoint method for the residual layers. Results are summarized in table 2. To the best of our knowledge, this is the first time a Neural-ODE like ResNet achieves a test-accuracy of $75.1\%$ on ImageNet.

Failure in usual settings.

In Prop. 3 we showed under assumption 2, that is if the residuals are bounded and Lipschitz continuous with constant independent of the depth $N$, then the error for computing the activations backward would scale in $\frac{1}{N}$ as well as the error for the gradients (Prop. 4). First, this results shows that the architecture needs to be deep enough, because it scales in $\frac{1}{N}$: for instance, we fail to train a ResNet-101 (He et al., 2016a) on the ImageNet dataset using the adjoint method on its third layer (depth $23$), as shown in Figure 3 (b).

Success at large depth.

To further investigate the applicability of the adjoint method for training deeper ResNets, we train a simple ResNet model on the CIFAR data set. First, the input is processed by a $5\times 5$ convolution with $16$ out channels, and the image is down-sampled to a size $10\times 10$.

We then apply a batch norm, a ReLU and iterate relation (1) where $f$ is a pre-activation basic block (He et al., 2016b). We consider the zero residual initialisation: the last batch norm of each basic block is initialized to zero. We consider different values for the depth $N$ and notice that in this setup, the deeper our model is, the better it performs in term of test accuracy. We then compare the performance of our model using a ResNet (forward rule (1)) or a HeunNet (forward rule (8)). We train our networks using either the classical backpropagation or our corresponding proxys using the adjoint method (formulas (6) and (9)). We display the final test accuracy (median over $5$ runs) for different values of the depth $N$ in Figure 4. The true backpropagation gives the same curves for the ResNet and the HeunNet. Approximated gradients, however, lead to a large test error at small depth, but give the same performance at large depth, hence confirming our results in Prop. 4 and 6. In addition, at fixed depth, the accuracy when training a HeunNet with the adjoint method is better (or similar at depths $2$, $32$ and $64$) than for the ResNet with the adjoint method. This is to be linked with the two different bounds in Prop. 4 and 6: for the HeunNet, smoothness with depth, which is expected at large depth, according to the theoretical results for the linear case (Prop. 2), implies a faster convergence to the true gradients for the HeunNet than for the ResNet. We finally validate this convergence in Figure 3 (c): the deeper the architecture, the better the approximation on the gradients. In addition, the HeunNet approximates the true gradient better than the ResNet.

Conclusion, limitations and future works

We propose a methodology to analyze how well a ResNet discretizes a Neural ODE. The positive results predicted by our theory in the linear case are also observed in practice with real architectures: one can successfully use the adjoint method to train ResNets (or even more effectively HeunNets) using very deep architectures on CIFAR, or fine-tune them on ImageNet, without memory cost in the residual layers. However, we also show that for large scale problems such as ImageNet classification from scratch, the adjoint method fails at usual depths.

Our work provides a theoretical guarantee for the convergence to a Neural ODE in the linear setting under a small loss initialization. A natural extension would be to study the non-linear case. In addition, the adjoint method is time consuming, and an improvement would be to propose a cheaper method than a reverse mode traversal of the architecture for approximating the activations.

Acknowledgments

This work was granted access to the HPC resources of IDRIS under the allocation 2020-[AD011012073] made by GENCI. This work was supported in part by the French government under management of Agence Nationale de la Recherche as part of the “Investissements d’avenir” program, reference ANR19-P3IA-0001 (PRAIRIE 3IA Institute). This work was supported in part by the European Research Council (ERC project NORIA). M. S. thanks Mathieu Blondel and Zaccharie Ramzi for helpful discussions.

References

APPENDIX

In Section A we give the proofs of all the propositions, lemmas and the theorem presented in this work.

Section B gives details for the experiments in the paper.

We also give a recap on ResNet architectures in Section C.

Appendix A Proofs

Our proof is inspired by [Demailly, 2016].

We denote $h=\frac{1}{N}$ and $s_{n}=nh$. We define

We have that $\varphi_{\Theta}(x(s_{n}),s_{n})=\dot{x}(s_{n})$.

with $\|R_{1}(h)\|\leq\frac{1}{2}h^{2}\|\ddot{x}\|_{\infty}$. This implies that

The true error we are interested in is the global error $e_{n}=x(s_{n})-x_{n}$. One has

Because $\varphi_{\Theta}$ is $L$-Lipschitz, this gives $\|e_{n+1}-e_{n}\|\leq\|\varepsilon_{n}\|+hL\|e_{n}\|$ and hence

this implies from the discrete Gronwall lemma, since $e_{0}=0$ that

Note that we have $\ddot{x}=\partial_{s}\varphi_{\Theta}+\partial_{x}\varphi_{\Theta}[\varphi_{\Theta}].$ This gives the desired result. ∎

A.2 Proof of Prop. 2

Recall that we denote $\Pi^{N}=\prod_{n=1}^{N}(I_{d}+\frac{\theta_{n}^{N}}{N})$, $\Pi^{N}_{:n}=(I_{d}+\frac{\theta_{N}^{N}}{N})\dots(I_{d}+\frac{\theta_{n+1}^{N}}{N})$ and $\Pi^{N}_{n:}=(I_{d}+\frac{\theta_{n-1}^{N}}{N})\dots(I_{d}+\frac{\theta_{1}^{N}}{N})$. We denote $\nabla^{N}_{n}=\nabla_{\theta_{n}^{N}}L$. One has

One has that $\forall t\in[0,t^{*}]$, $\sigma^{2}_{min}(\Pi^{N}_{:n})\geq(1-\frac{1}{2N})^{2(N-n)}$ and $\sigma^{2}_{min}(\Pi^{N}_{n:})\geq(1-\frac{1}{2N})^{2(n-1)}$ which implies that

We now show our main result. Note that we have the relationship $(I+\frac{\theta_{n+1}^{N}}{N})^{\top}\nabla_{n+1}=\nabla_{n}(I+\frac{\theta_{n}^{N}}{N})^{\top}$ so that

Because $\|(I+A)^{-1}\|\leq 2$ if $\|A\|\leq\frac{1}{2}$ this gives $\|\nabla^{N}_{n+1}-\nabla^{N}_{n}\|\leq\frac{2}{N}\|\nabla^{N}_{n}\|$. Integrating we get

A.3 Proof of lemma 1

We adapt a variant of the Ascoli–Arzelà theorem [Brezis and Brézis, 2011]. We showed in Prop. 2 that there exists $C>0$ that only depends on the initialization such that, $\forall t\geq 0,\forall i\in[N-1]$,

Its follows that $\|\psi_{N}(s_{1},t_{1})-\psi_{N}(s_{2},t_{2})\|\leq\|\psi_{N}(s_{1},t_{1})-\psi_{N}(s_{1},t_{2})\|+\|\psi_{N}(s_{1},t_{2})-\psi_{N}(s_{2},t_{2})\|$ and thus

These two properties are essential to prove our lemma. We proceed as follows.

(we denote the limit $\psi(s_{j},t_{j})$).

$\|\psi_{\sigma(N)}(s,t)-\psi_{\sigma(M)}(s,t)\|\leq\|\psi_{\sigma(N)}(s,t)-\psi_{\sigma(N)}(s_{k},t_{k})\|+\|\psi_{\sigma(N)}(s_{k},t_{k})-\psi_{\sigma(M)}(s_{k},t_{k})\|+\|\psi_{\sigma(M)}(s_{k},t_{k})-\psi_{\sigma(M)}(s,t)\|$

and $\|\psi(s,t)-\psi(u,t)\|\leq\varepsilon$. There exists a finite set of $\{s_{j}\}_{j=1}^{k}$ such that

For our $s$, there exists $j\in\{1,\dots,k\}$ such that $\|s-s_{j}\|\leq\delta$.

There also exists $t_{0}\geq 0$ such that if $t\geq t_{0}$,

$\|\psi_{\sigma(N)}(s,t_{0})-\psi(s,t_{0})\|\leq\|\psi_{\sigma(N)}(s,t_{0})-\psi_{\sigma(N)}(s_{j},t_{0})\|+\|\psi_{\sigma(N)}(s_{j},t_{0})-\psi(s_{j},t_{0})\|+\|\psi(s_{j},t_{0})-\psi(s,t_{0})\|$

$\|\psi_{\sigma(N)}(s,t_{0})-\psi(s,t_{0})\|\leq 2\varepsilon+\frac{C}{\sigma(N)}+\max_{j\in\{1,\dots,k\}}\|\psi_{\sigma(N)}(s_{j},t_{0})-\psi(s_{j},t_{0})\|\leq 4\varepsilon$ for $N$ big enough.

Finally, $\|\psi_{\sigma(N)}(s,t)-\psi(s,t)\|\leq\|\psi_{\sigma(N)}(s,t)-\psi_{\sigma(N)}(s,t_{0})\|+\|\psi_{\sigma(N)}(s,t_{0})-\psi(s,t_{0})\|+\|\psi(s,t_{0})-\psi(s,t)\|\leq 6\varepsilon$

for $N$ big enough, independently of $t$ and $s$. This concludes the proof. ∎

A.4 Proof of lemma 2

Note also that since the Jacobian of $(\theta_{1},..,\theta_{N})\to\Pi^{N}$ is

for some constants $\alpha$, $\beta$. Finally, we have

Our (PL) conditions precisely write $-\Delta^{\top}H(\Delta)\leq-\lambda\|\Delta\|^{2}$ for some $\lambda>0$. Let $\varphi^{N}=\frac{1}{2}\|\Delta^{N}\|^{2}.$ One has

Since $\|\Delta^{N}\|=\sqrt{2\varphi^{N}}$ we get

Let $n(t)$ be such that $D^{N}_{n(t)}(t)=\max_{i\in[1,N]}D^{N}_{i}(t)$. We have

A.5 Proof of Th. 1

We first prove the following lemma 3 before proving Th. 1.

and the Euler scheme with time step $\frac{1}{\sigma(N)}$ for its discretization

We know by Prop. 1, since $x_{0}$ has unit norm that

Since $\|\theta^{N}_{n}\|\leq\frac{1}{2}$ and $x_{0}$ has unit norm, there exists $M>0$ independent of $x_{0}$ such that, $\forall n$ and $N$, $\|x_{n}\|\leq M$. Thus

as $N\to\infty$. We obtain the uniform convergence with $t$.

Consider $(\psi_{\sigma(N)})_{N}$ a sub-sequence of $(\psi_{N})_{N}$ as in lemma 1 that converges to some $\psi_{\sigma}$.

1) We first prove the uniqueness of the limit.

We want to show that $\psi_{\sigma}$ does not depend on $\sigma$. This will imply the uniqueness of any accumulation point of the relatively compact sequence $(\psi_{N})_{N}$ and thus its convergence.

As $N\to\infty$, we have thanks to lemma 3 that the right hand term converges uniformly to

This uniform convergence makes it possible to consider the limit ODE as $N\to\infty$:

Let $\psi_{1}$, $\psi_{2}$ with $\|\psi_{1}(s,t)\|\leq\frac{1}{2}$ and $\|\psi_{2}(s,t)\|\leq\frac{1}{2}$ and $\Pi_{1}(t)$, $\Pi_{2}(t)$ the corresponding flows.

One has $\Pi_{1}(t)x_{0}=x_{1}(1)$ and $\Pi_{2}(t)x_{0}=x_{2}(1)$. One has $\dot{y}=\psi_{1}x_{1}-\psi_{2}x_{2}=\psi_{2}y+(\psi_{1}-\psi_{2})x_{1}$. Hence, since $y(0)=0$, $\|y(s)\|\leq\int_{0}^{s}\|\psi_{2}\|\|y\|+\|\psi_{1}-\psi_{2}\|_{\infty}|\|x_{1}\|_{\infty}$, we have

for some $\alpha>0$. The same arguments go for $\Pi_{:s}$ and $\Pi_{s:}$.

Since we only consider maps $\psi_{\sigma}$ such that $\|\psi_{\sigma}(s,t)\|\leq\frac{1}{2}$, this implies that the product is also Lipschitz and thus $F$ is Lipschitz. This guarantees the uniqueness of a solution $\psi$ to the Cauchy problem and we have that $\psi_{N}\to\psi$ uniformly.

Letting $k\to\infty$ finally gives $\|\psi-\psi_{N}\|\leq\frac{2D}{N}.$

A.6 Proof of Prop. 3

and since $r_{N}=0$, the discrete Gronwall lemma leads to $\|r_{n}\|\leq\frac{e^{L_{f}}-1}{{L_{f}}N}K_{N}+O(\frac{1}{N^{2}}).$ In addition, one has $K_{N}\leq L_{f}C_{f}$ so that

A.7 Proof of Prop. 4

1) We first control the error made in the gradient with respect to activations.

2) We can now control the gradients with respect to the parameters $\theta^{N}_{n}$’s.

Using our bound on $\|g_{n}\|$ and Prop. 3 we get

A.8 Proof of Prop. 5

In the following, we let for short $f_{n}(x)=f(x,\theta_{n}^{N})$, and we define

so that Heun’s forward and backward equations are

We have the following lemma that quantifies the reconstruction error over one iteration:

where $J_{n}=\partial_{x}f_{n}(x)$ is the Jacobian of $f_{n}$.

As $N$ goes to infinity, we have the following expansions of (12):

Putting everything together, we find that the zero-th order in $\psi_{n}(x+\frac{1}{N}\varphi_{n}(x))-\varphi_{n}(x)$ cancels, and that the first order simplifies to $\frac{1}{4N}\left(J_{n+1}(x)-J_{n}(x)\right)[f_{n+1}(x)-f_{n}(x)]$. ∎

We now turn the the proof of the main proposition:

Using the triangle inequality, and the $L_{f}^{\prime}-$Lispchitz continuity of $\psi_{n}$, we get

The last term is controlled with the previous Lemma 4:

A.9 Proof of Prop. 6

1) We first control the error made in the gradient with respect to activations. We have the following recursions:

The last term is controled with the previous proposition, and we find

2) We can now control the gradients with respect to parameters. Since Heun’s method involves parameters $\theta_{n}^{N}$ both for the computation of $x_{n}$ and $x_{n+1}$, the gradient formula is slightly more complicated than for the classical ResNet. It is the sum of two terms, the first one $\nabla^{1}_{\theta^{N}_{n}}L$ corresponding to iteration $n$ and the second one $\nabla^{2}_{\theta^{N}_{n}}L$ corresponding to iteration $n-1$.

The gradient $\nabla_{\theta^{N}_{n}}L$ is finally

Overall, these equations map the activations $x_{n}$ and $x_{n-1}$, and the gradients $\nabla_{x_{n-1}}L$ and $\nabla_{x_{n}}L$ to the gradient $\nabla_{\theta_{n}^{N}}$, which we rewrite as

where the function $\Psi$ is explicitly defined by the above equations. With the memory-free backward pass, the gradient is rather estimated as

The function $\Psi$ is Lispchitz-continuous since all functions involved in its composition are Lipschitz-continuous and the activations belong to a compact set, and its Lipschitz constant scales as $\frac{1}{N}$. We write its Lipschitz constant as $\frac{L_{\Psi}}{N}$, and we get:

Appendix B Experimental details

In all our experiments, we use Nvidia Tesla V100 GPUs.

For our experiments on CIFAR-10 (training from scratch), we used a batch-size of $128$ and we employed SGD with a momentum of $0.9$. The training was done over $200$ epochs. The initial learning rate was $0.1$ and we used a cosine learning rate scheduler. A constant weight decay was set to $5\times 10^{-4}$. Standard inputs preprocessing as proposed in Pytorch [Paszke et al., 2017] was performed.

For our finetuning experiment on CIFAR-10, we used a batch-size of $128$ and we employed SGD with a momentum of $0.9$. The training was done over $5$ epochs. The learning rate was kept constant to $10^{-3}$. A constant weight decay was set to $5\times 10^{-4}$. Standard inputs preprocessing as proposed in Pytorch was also performed.

For our experiment with our simple ResNet model that processes the input by a $5\times 5$ convolution with $16$ out channels, we used a batch-size of $256$ and we employed SGD with a momentum of $0.9$. The training was done over $90$ epochs. The learning rate was set to $10^{-1}$ and was decayed by a factor $10$ every $30$ epochs. A constant weight decay was set to $5\times 10^{-4}$. Standard inputs preprocessing as proposed in Pytorch was also performed.

B.2 ImageNet

For our experiments on ImageNet (training from scratch), we used a batch-size of $256$ and we employed SGD with a momentum of $0.9$. The training was done over $100$ epochs. The initial learning rate was $0.1$ and was decayed by a factor $10$ every $30$ epochs. A constant weight decay was set to $10^{-4}$. Standard inputs preprocessing as proposed in Pytorch was performed: normalization, random croping of size $224\times 224$ pixels, random horizontal flip.

For our finetuning experiment on ImageNet, we used a batch-size of $256$ and we employed SGD with a momentum of $0.9$. The training was done over $3$ epochs. The learning rate was kept constant to $5\times 10^{-4}$. A constant weight decay was set to $10^{-4}$. Standard inputs preprocessing as proposed in Pytorch was performed: normalization, random croping of size $224\times 224$ pixels, random horizontal flip.

Appendix C Architecture details

In computer vision, the ResNet as presented in [He et al., 2016a] first applies non residual transformations to the input image: a feature extension convolution that goes to $3$ channels to 64, a batch norm, a non-linearity (ReLU) and optionally a maxpooling.

Finally, there is a classification module: average pooling followed by a fully connected layer.