Momentum Residual Neural Networks

Michael E. Sander, Pierre Ablin, Mathieu Blondel, Gabriel Peyré

Introduction

As a particular instance of deep learning (LeCun et al., 2015; Goodfellow et al., 2016), residual neural networks (He et al., 2016, ResNets) have achieved great empirical successes due to extremely deep representations and their extensions keep on outperforming state of the art on real data sets (Kolesnikov et al., 2019; Touvron et al., 2019). Most of deep learning tasks involve graphics processing units (GPUs), where memory is a practical bottleneck in several situations (Wang et al., 2018; Peng et al., 2017; Zhu et al., 2017). Indeed, backpropagation, used for optimizing deep architectures, requires to store values (activations) at each layer during the evaluation of the network (forward pass). Thus, the depth of deep architectures is constrained by the amount of available memory. The main goal of this paper is to explore the properties of a new model, Momentum ResNets, that circumvent these memory issues by being invertible: the activations at layer nn is recovered exactly from activations at layer n+1n+1. This network relies on a modification of the ResNet’s forward rule which makes it exactly invertible in practice. Instead of considering the feedforward relation for a ResNet (residual building block)

we define its momentum counterpart, which iterates

where ff is a parameterized function, vv is a velocity term and γ∈\gamma\in is a momentum term. This radically changes the dynamics of the network, as shown in the following figure.

In contrast with existing reversible models, Momentum ResNets can be integrated seamlessly in any deep architecture which uses residual blocks as building blocks (cf. in Section 3).

We introduce momentum residual neural networks (Momentum ResNets), a new deep model that relies on a simple modification of the ResNet forward rule and which, without any constraint on its architecture, is perfectly invertible. We show that the memory requirement of Momentum ResNets is arbitrarily reduced by changing the momentum term γ\gamma (Section 3.2), and show that they can be used as a drop-in replacement for traditional ResNets.

On the theoretical side, we show that Momentum ResNets are easily used in the learning to optimize setting, where other reversible models fail to converge (Section 3.3). We also investigate the approximation capabilities of Momentum ResNets, seen in the continuous limit as second-order ODEs (Section 4). We first show in Proposition 3 that Momentum ResNets can represent a strictly larger class of functions than first-order neural ODEs. Then, we give more detailed insights by studying the linear case, where we formally prove in Theorem 1 that Momentum ResNets with linear residual functions have universal approximation capabilities, and precisely quantify how the set of representable mappings for such models grows as the momentum term γ\gamma increases. This theoretical result is a first step towards a theoretical analysis of representation capabilities of Momentum ResNets.

Our last contribution is the experimental validation of Momentum ResNets on various learning tasks. We first show that Momentum ResNets separate point clouds that ResNets fail to separate (Section 5.1). We also show on image datasets (CIFAR-10, CIFAR-100, ImageNet) that Momentum ResNets have similar accuracy as ResNets, with a smaller memory cost (Section 5.2). We also show that parameters of a pre-trained model are easily transferred to a Momentum ResNet which achieves comparable accuracy in only few epochs of training. We argue that this way to obtain pre-trained Momentum ResNets is of major importance for fine-tuning a network on new data for which memory storage is a bottleneck. We provide a Pytorch package with a method that takes a torchvision ResNet model and returns its Momentum counterpart that achieves similar accuracy with very little refit. We also experimentally validate our theoretical findings in the learning to optimize setting, by confirming that Momentum ResNets perform better than RevNets (Gomez et al., 2017). Our code is available at https://github.com/michaelsdr/momentumnet.

Background and previous works.

Backpropagation is the method of choice to compute the gradient of a scalar-valued function. It operates using the chain rule with a backward traversal of the computational graph (Bauer, 1974). It is also known as reverse-mode automatic differentiation (Baydin et al., 2018; Rumelhart et al., 1986; Verma, 2000; Griewank & Walther, 2008). The computational cost is similar to the one of evaluating the function itself. The only way to back-propagate gradients through a neural architecture without further assumptions is to store all the intermediate activations during the forward pass. This is the method used in common deep learning libraries such as Pytorch (Paszke et al., 2017), Tensorflow (Abadi et al., 2016) and JAX (Jacobsen et al., 2018). A common way to reduce this memory storage is to use checkpointing: activations are only stored at some steps and the others are recomputed between these check-points as they become needed in the backward pass (e.g., Martens & Sutskever (2012)).

However, models that allow backpropagation without storing any activations have recently been developed. They are based on two kinds of approaches. The first is discrete and relies on finding ways to easily invert the rule linking activation nn to activation n+1n+1 (Gomez et al., 2017; Chang et al., 2018; Haber & Ruthotto, 2017; Jacobsen et al., 2018; Behrmann et al., 2019). In this way, it is possible to recompute the activations on the fly during the backward pass: activations do not have to be stored. However, these methods either rely on restricted architectures where there is no straightforward way to transfer a well performing non-reversible model into a reversible one, or do not offer a fast inversion scheme when recomputing activations backward. In contrast, our proposal can be applied to any existing ResNet and is easily inverted. The second kind of approach is continuous and relies on ordinary differential equations (ODEs), where ResNets are interpreted as continuous dynamical systems (Weinan, 2017; Chen et al., 2018; Teh et al., 2019; Sun et al., 2018; Weinan et al., 2019; Lu et al., 2018; Ruthotto & Haber, 2019). This allows one to import theoretical and numerical advances from ODEs to deep learning. These models are often called neural ODEs (Chen et al., 2018) and can be trained by using an adjoint sensitivity method (Pontryagin, 2018), solving ODEs backward in time. This strategy avoids performing reverse-mode automatic differentiation through the operations of the ODE solver and leads to a O(1)O(1) memory footprint. However, defining the neural ODE counterpart of an existing residual architecture is not straightforward: optimizing ODE blocks is an infinite dimensional problem requiring a non-trivial time discretization, and the performances of neural ODEs depend on the numerical integrator for the ODE (Gusak et al., 2020). In addition, ODEs cannot always be numerically reversed, because of stability issues: numerical errors can occur and accumulate when a system is run backwards (Gholami et al., 2019; Teh et al., 2019). Thus, in practice, neural ODEs are seldom used in standard deep learning settings. Nevertheless, recent works (Zhang et al., 2019; Queiruga et al., 2020) incorporate ODE blocks in neural architectures to achieve comparable accuracies to ResNets on CIFAR.

Studying the representation capabilities of such models is also important, as it gives insights regarding their performance on real world data. It is well-known that a single residual block has universal approximation capabilities (Cybenko, 1989), meaning that on a compact set any continuous function can be uniformly approximated with a one-layer feedforward fully-connected neural network. However, neural ODEs have limited representation capabilities. Teh et al. (2019) propose to lift points in higher dimensions by concatenating vector fields of data with zeros in an extra-dimensional space, and show that the resulting augmented neural ODEs (ANODEs) achieve lower loss and better generalization on image classification and toy experiments. Li et al. (2019) show that, if the output of the ODE-Net is composed with elements of a terminal family, then universal approximation capabilities are obtained for the convergence in LpL^{p} norm for p<+∞p<+\infty, which is insufficient (Teshima et al., 2020). In this work, we consider the representation capabilities in L∞L^{\infty} norm of the ODEs derived from the forward iterations of a ResNet. Furthermore, Zhang et al. (2020) proved that doubling the dimension of the ODE leads to universal approximators, although this result has no application in deep learning to our knowledge. In this work, we show that in the continuous limit, our architecture has better representation capabilities than Neural ODEs. We also prove its universality in the linear case.

Some recent works (He et al., 2020; Chun et al., 2020; Nguyen et al., 2020; Li et al., 2018) have explored momentum in deep architectures. However, these methods differ from ours in their architecture and purpose. Chun et al. (2020) introduce a momentum to solve an optimization problem for which the iterations do not correspond to a ResNet. Nguyen et al. (2020) (resp. He et al. (2020)) add momentum in the case of RNNs (different from ResNets) where the weights are tied to alleviate the vanishing gradient issue (resp. link the key and query encoder layers). Li et al. (2018) consider a particular case where the linear layer is tied and is a symmetric definite matrix. In particular, none of the mentioned architectures are invertible, which is one of the main assets of our method.

We show that adding a momentum term corresponds to an Euler integration scheme for integrating a second-order ODE. Some recently proposed architectures (Norcliffe et al., 2020; Rusch & Mishra, 2021; Lu et al., 2018; Massaroli et al., 2020) are also motivated by second-order differential equations. Norcliffe et al. (2020) introduce second-order dynamics to model second-order dynamical systems, whereas our model corresponds to a discrete set of equations in the continuous limit. Also, in our method, the neural network only acts on xx, so that although momentum increases the dimension to 2d2d, the computational burden of a forward pass is the same as a ResNet of dimension dd. Rusch & Mishra (2021) propose second-order RNNs, whereas our method deals with ResNets. Finally, the formulation of LM-ResNet in Lu et al. (2018) differs from our forward pass (xn+1=xn+γvn+(1−γ)f(xn,θn)x_{n+1}=x_{n}+\gamma v_{n}+(1-\gamma)f(x_{n},\theta_{n})), even though they both lead to second-order ODEs. Importantly, none of these second-order formulations are invertible.

Momentum Residual Neural Networks

We now introduce Momentum ResNet, a simple transformation of any ResNet into a model with a small memory requirement, and that can be seen in the continuous limit as a second-order ODE.

In this paper, we consider initial speeds v0v_{0} that depend on x0x_{0} through a simple relation. The simplest options are to set v0=0v_{0}=0 or v0=f(x0,θ0)v_{0}=f(x_{0},\theta_{0}). We prove in Section 4 that this dependency between v0v_{0} and x0x_{0} has an influence on the set of mappings that Momentum ResNets can represent. The parameter γ\gamma controls how much a Momentum ResNet diverges from a ResNet, and also the amount of memory saving. The closer γ\gamma is to , the closer Momentum ResNets are to ResNets, but the less memory is saved. In our experiments, we use γ=0.9\gamma=0.9, which we find to work well in various applications.

so that activations can be reconstructed on the fly during the backward pass in a Momentum ResNet. In practice, in order to exactly reverse the dynamics, the information lost by the finite-precision multiplication by γ\gamma in (2) has to be efficiently stored. We used the algorithm from Maclaurin et al. (2015) to perform this reversible multiplication. It consists in maintaining an information buffer, that is, an integer that stores the bits that are lost at each iteration, so that multiplication becomes reversible. We further describe the procedure in Appendix C. Note that there is always a small loss of floating point precision due to the addition of the learnable mapping ff. In practice, we never found it to be a problem: this loss in precision can be neglected compared to the one due to the multiplication by γ\gamma.

Our approach makes it possible to turn any existing ResNet into a reversible one. In other words, a ResNet can be transformed into its Momentum counterpart without changing the structure of each layer. For instance, consider a ResNet-152 (He et al., 2016). It is made of 44 layers (of depth 33, 88, 3636 and 33) and can easily be turned into its Momentum ResNet counterpart by changing the forward equations (1) into (2) in the 44 layers. No further change is needed and Momentum ResNets take the exact same parameters as inputs: they are a drop-in replacement. This is not the case of other reversible models. Neural ODEs (Chen et al., 2018) take continuous parameters as inputs. i-ResNets (Behrmann et al., 2019) cannot be trained by plain SGD since the spectral norm of the weights requires constrained optimization. i-RevNets (Jacobsen et al., 2018) and RevNets (Gomez et al., 2017) require to train two networks with their own parameters for each residual block, split the inputs across convolutional channels, and are half as deep as ResNets: they do not take the same parameters as inputs. Table 1 summarizes the properties of reversible residual architectures. We discuss in further details the differences between RevNets and Momentum ResNets in sections 3.3 and 5.3.

2 Memory cost

As another example, consider a ResNet-152 (He et al., 2016) which can be used for ImageNet classification (Deng et al., 2009). Its layer named “conv4_x” has a depth of 3636: it has 4040 M parameters, whereas storing the activations would require storing 5050 times more parameters. Since storing the activations is here the main obstruction, the memory requirement for this layer can be arbitrarily reduced by taking γ\gamma close to 11.

3 The role of momentum

When γ\gamma is set to in (2), we recover a ResNet. Therefore, Momentum ResNets are a generalization of ResNets. When γ→1\gamma\xrightarrow[]{}1, one can scale f→11−γff\to\frac{1}{1-\gamma}f to get in (2) a symplectic scheme (Hairer et al., 2006) that recovers a special case of other popular invertible neural network: RevNets (Gomez et al., 2017) and Hamiltonian Networks (Chang et al., 2018). A RevNet iterates

where φ\varphi and ψ\psi are two learnable functions.

The usefulness of such architecture depends on the task. RevNets have encountered success for classification and regression. However, we argue that RevNets cannot work in some settings. For instance, under mild assumptions, the RevNet iterations do not have attractive fixed points when the parameters are the same at each layer: θn=θ\theta_{n}=\theta, θn′=θ′\theta_{n}^{\prime}=\theta^{\prime}. We rewrite (4) as (vn+1,xn+1)=Ψ(vn,xn)(v_{n+1},x_{n+1})=\Psi(v_{n},x_{n}) with Ψ(v,x)=(v+φ(x,θ),x+ψ(v+φ(x,θ),θ′))\Psi(v,x)=(v+\varphi(x,\theta),x+\psi(v+\varphi(x,\theta),\theta^{\prime})).

This shows that (v∗,x∗)(v^{*},x^{*}) cannot be a stable fixed point. As a consequence, in practice, a RevNet cannot have converging iterations: according to (4), if xnx_{n} converges then vnv_{n} must also converge, and their limit must be a fixed point. The previous proposition shows that it is impossible.

This result suggests that RevNets should perform poorly in problems where one expects the iterations of the network to converge. For instance, as shown in the experiments in Section 5.3, this happens when we use reverible dynamics in order to learn to optimize (Maclaurin et al., 2015). In contrast, the proposed method can converge to a fixed point as long as the momentum term γ\gamma is strictly less than 11.

Proposition 1 has a continuous counterpart. Indeed, in the continuous limit, (4) writes v˙=φ(x,θ),x˙=ψ(v,θ′)\dot{v}=\varphi(x,\theta),\quad\dot{x}=\psi(v,\theta^{\prime}). The corresponding Jacobian in (v∗,x∗)(v^{*},x^{*}) is \big{(}\begin{smallmatrix}0&A\\ B&0\end{smallmatrix}\big{)}. The eigenvalues of this matrix are the square roots of those of ABAB: they cannot all have a real part <0<0 (same stability issue in the continuous case).

4 Momentum ResNets as continuous models

The ResNets equation (1) with initial condition x0x_{0} (the input of the ResNet) can be seen as a discretized Euler scheme of the ODE x˙=f(x,θ)\dot{x}=f(x,\theta) with x(0)=x0x(0)=x_{0}. Denoting TT a time horizon, the neural ODE maps the input x(0)x(0) to the output x(T)x(T), and, as in Chen et al. (2018), is trained by minimizing a loss L(x(T),θ)L(x(T),\theta).

Let ε=11−γ\varepsilon=\frac{1}{1-\gamma}. We can then rewrite (2) as

which corresponds to a Verlet integration scheme (Hairer et al., 2006) with step size 11 of the differential equation εx¨+x˙=f(x,θ).\varepsilon\ddot{x}+\dot{x}=f(x,\theta). Thus, in the same way that ResNets can be seen as discretization of first-order ODEs, Momentum ResNets can be seen as discretization of second-order ones. Figure 3 sums up these ideas.

Representation capabilities

We now turn to the analysis of the representation capabilities of Momentum ResNets in the continuous setting. In particular, we precisely characterize the set of mappings representable by Momentum ResNets with linear residual functions.

We denote by φt(x0)\varphi_{t}(x_{0}) the solution at time tt starting at initial condition x(0)=x0x(0)=x_{0}. It is called the flow of the ODE. For all t∈[0,T]t\in[0,T], where TT is a time horizon, φt\varphi_{t} is a homeomorphism: it is continuous, bijective with continuous inverse.

2 Representation capabilities of second-order ODEs

We consider the second-order model for which we recall that Momentum ResNets are a discretization:

In Section 3.3, we showed that Momentum ResNets generalize existing models when setting γ=0\gamma=0 or 11. We now state the continuous counterparts of these results. Recall that 11−γ=ε\frac{1}{1-\gamma}=\varepsilon. When ε→0\varepsilon\xrightarrow{}0, we recover the first-order model.

We let x∗x^{*} (resp. xεx_{\varepsilon}) be the solution of (5) (resp. (6)) on [0,T][0,T], with initial conditions x∗(0)=xε(0)=x0x^{*}(0)=x_{\varepsilon}(0)=x_{0} and x˙ε(0)=v0\dot{x}_{\varepsilon}(0)=v_{0}. Then ∥xε−x∗∥∞→0\|x_{\varepsilon}-x^{*}\|_{\infty}\xrightarrow[]{}0 as ε→0\varepsilon\xrightarrow{}0.

The proof of this result relies on the implicit function theorem and can be found in Appendix A.1. Note that Proposition 2 is true whatever the initial speed v0v_{0}. When ε→+∞\varepsilon\xrightarrow{}+\infty, one needs to rescale ff to study the asymptotics: the solution of x¨+1εx˙=f(x,θ)\ddot{x}+\frac{1}{\varepsilon}\dot{x}=f(x,\theta) converges to the solution of x¨=f(x,θ)\ddot{x}=f(x,\theta) (see details in Appendix B.1). These results show that in the continuous regime, Momentum ResNets also interpolate between x˙=f(x,θ)\dot{x}=f(x,\theta) and x¨=f(x,θ)\ddot{x}=f(x,\theta).

There exists a function f^\hat{f} such that for all xx solution of (5), xx is also solution of the second-order model εx¨+x˙=f^(x,θ)\varepsilon\ddot{x}+\dot{x}=\hat{f}(x,\theta) with (x(0),x˙(0))=(x0,f(x0,θ0))(x(0),\dot{x}(0))=(x_{0},f(x_{0},\theta_{0})).

Furthermore, even with the restrictive initial condition v0=0v_{0}=0, x↦λxx\mapsto\lambda x for λ>−1\lambda>-1 can always be represented by a second-order model (6) (see details in Appendix B.4). This supports the claim that the set of representable mappings increases with ε\varepsilon.

3 Universality of Momentum ResNets with linear residual functions

As a first step towards a theoretical analysis of the universal representation capabilities of Momentum ResNets, we now investigate the linear residual function case. Consider the second-order linear ODE

At time 11, (7) defines the linear mapping x0↦φ1(x0)=Ψε(θ)x0x_{0}\mapsto\varphi_{1}(x_{0})=\Psi_{\varepsilon}(\theta)x_{0} where

When ε→0\varepsilon\xrightarrow[]{}0, Proposition 2 shows that Ψε(θ)→Ψ0(θ)=exp⁡θ\Psi_{\varepsilon}(\theta)\xrightarrow{}\Psi_{0}(\theta)=\exp\theta. The range of the matrix exponential is indeed the set of representable mappings of a first order linear model

Experiments

We now demonstrate the applicability of Momentum ResNets through experiments. We used Pytorch and Nvidia Tesla V100 GPUs.

2 Image experiments

We also compare the accuracy of ResNets and Momentum ResNets on real data sets: CIFAR-10, CIFAR-100 (Krizhevsky et al., 2010) and ImageNet (Deng et al., 2009). We used existing ResNets architectures. We recall that Momentum ResNets can be used as a drop-in replacement and that it is sufficient to replace every residual building block with a momentum residual forward iteration. We set γ=0.9\gamma=0.9 in the experiments. More details about the experimental setup are given in Appendix D.

For these data sets, we used a ResNet-101 (He et al., 2016) and a Momentum ResNet-101 and compared the evolution of the test error and test loss. Two kinds of Momentum ResNets were used: one with an initial speed v0=0v_{0}=0 and the other one where the initial speed v0v_{0} was learned: v0=f(x0)v_{0}=f(x_{0}). These experiments show that Momentum ResNets perform similarly to ResNets. Results are summarized in Table 2.

For this data set, we used a ResNet-101, a Momentum ResNet-101, and a RevNet-101. For the latter, we used the procedure from Gomez et al. (2017) and adjusted the depth of each layer for the model to have approximately the same number of parameters as the original ResNet-101. Evolution of test errors are shown in Figure 6 (lower row), where comparable performances are achieved.

We compare the memory (using a memory profiler) for performing one epoch as a function of the batch size for two datasets: ImageNet (depth of 152) and CIFAR-10 (depth of 1201). Results are shown in Figure 7 and illustrate how Momentum ResNets can benefit from increased batch size, especially for very deep models. We also show in Figure 7 the final test accuracy for a full training of Momentum ResNets on CIFAR-10 as a function of the memory used (directly linked to γ\gamma (section 3.2)).

It has been shown (Tajbakhsh et al., 2016) that in various medical imaging applications the use of a pre-trained model on ImageNet with adapted fine-tuning outperformed a model trained from scratch. In order to easily obtain pre-trained Momentum ResNets for applications where memory could be a bottleneck, we transferred the learned parameters of a ResNet-152 pre-trained on ImageNet to a Momentum ResNet-152 with γ=0.9\gamma=0.9. In only 11 epoch of additional training we reached a top-1 error of 26.5%26.5\% and in 55 additional epochs a top-1 error of 23.5%23.5\%. We then empirically compared the accuracy of these pre-trained models by fine-tuning them on new images: the hymenopterahttps://www.kaggle.com/ajayrana/hymenoptera-data data set.

As a proof of concept, suppose we have a GPU with 33 Go of RAM. The images have a resolution of 500×500500\times 500 pixels so that the maximum batch size that can be taken for fine-tuning the ResNet-152 is 22, against 44 for the Momentum ResNet-152. As suggested in Tajbakhsh et al. (2016) (“if the distance between the source and target applications is significant, one may need to fine-tune the early layers as well”), we fine-tune the whole network in this proof of concept experiment. In this setting the Momentum ResNet leads to faster convergence when fine-tuning, as shown in Figure 8: Momentum ResNets can be twice as fast as ResNets to train when samples are so big that only few of them can be processed at a time. In contrast, RevNets (Gomez et al., 2017) cannot as easily be used for fine-tuning since, as shown in (4), they require to train two distinct networks.

We also compare accuracy when using first-order ODE blocks (Chen et al., 2018) and second-order ones on CIFAR-10. In order to emphasize the influence of the ODE, we considered a neural architecture which down-sampled the input to have a certain number of channels, and then applied 1010 successive ODE blocks. Two types of blocks were considered: one corresponded to the first-order ODE (5) and the other one to the second-order ODE (6). Training was based on the odeint function implemented by Chen et al. (2018). Figure 9 shows the final test accuracy for both models as a function of the number of channels used. As a baseline, we also include the final accuracy when there are no ODE blocks. We see that an ODE Net with momentum significantly outperforms an original ODE Net when the number of channels is small. Training took the same time for both models.

3 Learning to optimize

We consider a Momentum ResNet and a RevNet variant of LISTA which use the residual function ff. For the RevNet, the activations xnx_{n} are first duplicated: the network has twice as many parameters at each layer. The matrix DD is generated with i.i.d. Gaussian entries with p=32p=32, d=16d=16, and its columns are then normalized to unit variance. Training and testing samples yy are generated as normalized Gaussian i.i.d. entries. More details on the experimental setup are added in Appendix D. The next Figure 10 shows the test loss of the different methods, when the depth of the networks varies.

As predicted by Proposition 1, the RevNet architecture fails on this task: it cannot have converging iterations, which is exactly what is expected here. In contrast, the Momentum ResNet works well, and even outperforms the LISTA baseline. This is not surprising: it is known that momentum can accelerate convergence of first order optimization methods.

Conclusion

This paper introduces Momentum ResNets, new invertible residual neural networks operating with a significantly reduced memory footprint compared to ResNets. In sharp contrast with existing invertible architectures, they are made possible by a simple modification of the ResNet forward rule. This simplicity offers both theoretical advantages (better representation capabilities, tractable analysis of linear dynamics) and practical ones (drop-in replacement, speed and memory improvements for model fine-tuning). Momentum ResNets interpolate between ResNets (γ=0\gamma=0) and RevNets (γ=1\gamma=1), and are a natural second-order extension of neural ODEs. As such, they can capture non-homeomorphic dynamics and converging iterations. As shown in this paper, the latter is not possible with existing invertible residual networks, although crucial in the learning to optimize setting.

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). The authors would like to thank David Duvenaud and Dougal Maclaurin for their helpful feedbacks. M. S. thanks Pierre Rizkallah and Pierre Roussillon for fruitful discussions.

References

Appendix A Proofs

If f:U×V→Wf:U\times V\to W is a function, we denote by ∂uf\partial_{u}f, when it exists, the partial derivative of ff with respect to u∈Uu\in U.

A.0 Instability of fixed points – Proof of Proposition 1

Since (x∗,v∗)(x^{*},v^{*}) is a fixed point of the RevNet iteration, we have

Then, a first order expansion, writing x=x∗+εx=x^{*}+\varepsilon and v=v∗+δv=v^{*}+\delta gives at order one

We start by showing that λ≠1\lambda\neq 1 by contradiction. Indeed, if λ=1\lambda=1, then (10) gives Aw=0Aw=0, which implies w=0w=0 since AA is invertible. Then, (11) gives Bu=0Bu=0, which also implies u=0u=0. This contradicts the fact that (u,v)(u,v) is an eigenvector (which is non-zero by definition).

Then, the first equation (10) gives Aw=(λ−1)uAw=(\lambda-1)u, and multiplying (11) by AA on the left gives

We also cannot have λ=0\lambda=0, since it would imply u=0u=0. Then, dividing (12) by λ\lambda shows that (λ−1)2λ\frac{(\lambda-1)^{2}}{\lambda} is an eigenvalue of ABAB.

Next, we let μ≠0\mu\neq 0 the eigenvalue of ABAB such that μ=(λ−1)2λ\mu=\frac{(\lambda-1)^{2}}{\lambda}. The equation can be rewritten as the second order equation

This equation has two solutions λ1(μ)\lambda_{1}(\mu), λ2(μ)\lambda_{2}(\mu), and since the constant term is 11, we have λ1(μ)λ2(μ)=1\lambda_{1}(\mu)\lambda_{2}(\mu)=1. Taking modulus, we get ∣λ1(μ)∣∣λ2(μ)∣=1|\lambda_{1}(\mu)||\lambda_{2}(\mu)|=1, which shows that necessarily, either ∣λ1(μ)∣≥1|\lambda_{1}(\mu)|\geq 1 or ∣λ1(μ)∣≥1|\lambda_{1}(\mu)|\geq 1.

Leveraging the fact that uu is an eigenvector of ABAB, we have λABu=λμu\lambda ABu=\lambda\mu u, and finally:

Which recovers exactly (11): λ\lambda is indeed an eigenvalue of J(A,B)J(A,B). ∎

A.1 Momentum ResNets in the limit ε→0absent→𝜀0\varepsilon\xrightarrow[]{}0 – Proof of Proposition 2

We take T=1T=1 without loss of generality. We are going to use the implicit function theorem. Note that xεx_{\varepsilon} is solution of (6) if and only if (xε,vε=xε˙)(x_{\varepsilon},v_{\varepsilon}=\dot{x_{\varepsilon}}) is solution of

so that xεx_{\varepsilon} is solution of (6) if and only if uε=(xε,vε=xε˙)u_{\varepsilon}=(x_{\varepsilon},v_{\varepsilon}=\dot{x_{\varepsilon}}) satisfies Ψ(uε,ε)=0.\Psi(u_{\varepsilon},\varepsilon)=0. Let u∗=(x∗,x∗˙)u^{*}=(x^{*},\dot{x^{*}}). One has Ψ(u∗,0)=0.\Psi(u^{*},0)=0. Ψ\Psi is differentiable everywhere, and at (u∗,0)(u^{*},0) we have

∂uΨ(u∗,0)\partial_{u}\Psi(u^{*},0) is continuous, and it is invertible with continuous inverse because it is linear and continuous, and because ∂uΨ(u∗,0)(x,v)=0\partial_{u}\Psi(u^{*},0)(x,v)=0 if and only if

This in particular ensures that xεx_{\varepsilon} converges uniformly to x∗x^{*} as ε\varepsilon goes to ∎

A.2 Momentum ResNets are more general than neural ODEs – Proof of Proposition 3

If xx satisfies (5) we get by derivation that

Then, if we define f^(x,θ)=ε[∂xf(x,θ)f(x,θ)+∂θf(x,θ)θ˙]+f(x,θ)\hat{f}(x,\theta)=\varepsilon[\partial_{x}f(x,\theta)f(x,\theta)+\partial_{\theta}f(x,\theta)\dot{\theta}]+f(x,\theta), we get that xx is also solution of the second-order model εx¨+x˙=f^(x,θ)\varepsilon\ddot{x}+\dot{x}=\hat{f}(x,\theta) with (x(0),x˙(0))=(x0,f(x0,θ0))(x(0),\dot{x}(0))=(x_{0},f(x_{0},\theta_{0})). ∎

A.3 Solution of (7) – Proof of Proposition 4

For which the solution at time tt writes

The calculation of this exponential gives

Note that it can be checked directly that this expression satisfies (7) by derivations. At time 11 this effectively gives x(1)=Ψε(θ)x0x(1)=\Psi_{\varepsilon}(\theta)x_{0}.

A.4 Representable mappings for a Momentum ResNet with linear residual functions – Proof of Theorem 1

In what follows, we denote by fεf_{\varepsilon} the function of matrices defined by

We first prove Theorem 1 in the diagonalizable case.

A.4.2 Theorem 1 in the diagonalizable case

Necessity Suppose that DD can be represented by a second-order model (7). This means that there exists a real matrix XX such that D=fε(X)D=f_{\varepsilon}(X) with XX real and

is such that fε(X)=Df_{\varepsilon}(X)=D, and DD is represented by a second-order model (7). ∎

We now state and demonstrate the general version of Theorem 1.

First, we need to demonstrate properties of the complex derivatives of the entire function fεf_{\varepsilon}.

so that Gε′(z)=−2zfε′(−z2)G_{\varepsilon}^{\prime}(z)=-2zf_{\varepsilon}^{\prime}(-z^{2}) and it is sufficient to prove that the zeros of Gε′G_{\varepsilon}^{\prime} are all real.

The zeros of GεG_{\varepsilon} are all real.

This class being stable under differentiation, we get that Gε′G_{\varepsilon}^{\prime} also belongs to the Laguerre-Pólya class. So that the roots of Gε′G_{\varepsilon}^{\prime} are all real, and hence those of fεf_{\varepsilon} as well.

A.4.4 Theorem 1 in the general case

When ε=0\varepsilon=0, we have in the general case the following from Culver (1966):

We now state and demonstrate the equivalent of this result for second order models (7).

If AA can be represented by a second-order model (7), then each Jordan block of A corresponding to an eigen value λ<λε\lambda<\lambda_{\varepsilon} occurs an even number of time.

Reciprocally, if each Jordan block of A corresponding to an eigen value λ≤λε\lambda\leq\lambda_{\varepsilon} occurs an even number of time, then AA can be represented by a second-order model.

We refer to the arguments from Culver (1966) and use results from Gantmacher (1959) for the proof.

Since fε(zkˉ)=fε(zk)=λkf_{\varepsilon}(\bar{z_{k}})=f_{\varepsilon}(z_{k})=\lambda_{k}, we can conclude that the Jordan blocks of A corresponding λk<λε\lambda_{k}<\lambda_{\varepsilon} occur an even number of time.

Now, suppose that each Jordan block of AA corresponding to an eigen value λ≤λε\lambda\leq\lambda_{\varepsilon} occurs an even number of times. Let λk\lambda_{k} be an eigenvalue of AA.

This shows that the Jordan blocks of AA are necessarily of the form

Appendix B Additional theoretical results

\varepsilon\xrightarrow{}+\infty). We let x∗x^{*} (resp. xεx_{\varepsilon}) be the solution of x¨=f(x,θ)\ddot{x}=f(x,\theta) (resp. x¨+1εx˙=f(x,θ)\ddot{x}+\frac{1}{\varepsilon}\dot{x}=f(x,\theta)) on [0,T][0,T], with initial conditions x∗(0)=xε(0)=x0x^{*}(0)=x_{\varepsilon}(0)=x_{0} and x˙∗(0)=x˙ε(0)=v0\dot{x}^{*}(0)=\dot{x}_{\varepsilon}(0)=v_{0}. Then xεx_{\varepsilon} converges uniformly to x∗x^{*} as ε→+∞\varepsilon\xrightarrow{}+\infty.

The equation x¨+1εx˙=f(x,θ)\ddot{x}+\frac{1}{\varepsilon}\dot{x}=f(x,\theta) with xε(0)=x0x_{\varepsilon}(0)=x_{0}, x˙ε(0)=v0\dot{x}_{\varepsilon}(0)=v_{0} writes in phase space (x,v)(x,v)

It then follows from the Cauchy-Lipschitz Theorem with parameters (Perko, 2013, Theorem 2, Chapter 2) that the solutions of this system are continuous in the parameter 1ε\frac{1}{\varepsilon}. That is xεx_{\varepsilon} converges uniformly to x∗x^{*} as ε→+∞\varepsilon\xrightarrow{}+\infty.

B.2 Universality of Momentum ResNets

This is because the solution is φt(x0)=x0−v0(e−t−1)\varphi_{t}(x_{0})=x_{0}-v_{0}(e^{-t}-1). ∎

This in particular proves that x↦λxx\mapsto\lambda x for λ≤−1\lambda\leq-1 cannot be represented by this ODE with initial conditions (x0,0)(x_{0},0).

Consider such an x0x_{0} and hh. Since φ1(x0)=h(x0)≤−x0\varphi_{1}(x_{0})=h(x_{0})\leq-x_{0}, that φ0(x0)=x0\varphi_{0}(x_{0})=x_{0} and that t↦φt(x0)t\mapsto\varphi_{t}(x_{0}) is continuous, we know that there exists t0∈t_{0}\in such that φt0(x0)=−x0\varphi_{t_{0}}(x_{0})=-x_{0}. We denote x(t)=φt(x0)x(t)=\varphi_{t}(x_{0}), solution of

Since d=1d=1, one can write ff as a derivative: f=−E′f=-E^{\prime}. The energy Em=12x˙2+EE_{m}=\frac{1}{2}\dot{x}^{2}+E satisfies:

So that E(−x0)≤E(x0)E(-x_{0})\leq E(x_{0}) We now apply the exact same argument to the solution starting at x1=−x0x_{1}=-x_{0}. Since x0≤h(−x0)=h(x1)x_{0}\leq h(-x_{0})=h(x_{1}) there exists t1∈t_{1}\in such that φt1(x1)=x0\varphi_{t_{1}}(x_{1})=x_{0}. So that:

So that E(x0)≤E(−x0)E(x_{0})\leq E(-x_{0}). We get that

There exits ff such that the solution of

with initial condition (x0,0)(x_{0},0) at time 11 is

with initial condition (x0,0)(x_{0},0) The solution of this ODE is

B.5 Orientation preservation of first-order ODEs

where kk is ff’s Lipschitz constant. So that ∥Φtn(xn)−Φtn(x)∥→0\|\Phi_{t_{n}}(x_{n})-\Phi_{t_{n}}(x)\|\xrightarrow{}0 as n→∞n\xrightarrow{}\infty. In addition, it is obvious that ∥Φtn(x)−Φt(x)∥→0\|\Phi_{t_{n}}(x)-\Phi_{t}(x)\|\xrightarrow{}0 as n→∞n\xrightarrow{}\infty. We conclude that

Let’s denote by HH the set of homeomorphisms defined on KK. The application

where MfM_{f} bounds the continuous function ff on CC defined in lemma 3. Since MfM_{f} does not depend on x0x_{0}, we have that

as ε→0\varepsilon\xrightarrow{}0, which proves that Ψ\Psi is continuous. Since Ψ(0)=IdK\Psi(0)=Id_{K}, we get that ∀t∈\forall t\in, Φt\Phi_{t} is connected to IdKId_{K}. ∎

B.6 On the linear mappings represented by autonomous first order ODEs in dimension 111

Suppose d=1d=1. If (15) represents a linear mapping x↦axx\mapsto ax at time 11, we have that ff is linear.

Consider FF a primitive of 1f\frac{1}{f}. Integrating (16), we get

This proves that f(0)=0f(0)=0. We now suppose that a>1a>1. We also have that

But when n→∞n\xrightarrow{}\infty, f(xan)=xanf′(0)+o(1an)f(\frac{x}{a^{n}})=\frac{x}{a^{n}}f^{\prime}(0)+o(\frac{1}{a^{n}}) so that

and ff is linear. The case a<1a<1 treats similarly by changing ana^{n} to a−na^{-n}. ∎

B.7 There are mappings that are connected to the identity that cannot be represented by a first order autonomous ODE

Appendix C Exact multiplication

Appendix D Experiment details

Appendix E Backpropagation for Momentum ResNets

In order to backpropagate the gradient of some loss in a Momentum ResNet, we need to formulate an explicit version of (2). Indeed, (2) writes explicitly

Writing z=(x,v)z=(x,v), the backpropagation for Momentum ResNets then writes, for some loss LL

We implement these formula to obtain a custom Jacobian-vector product in Pytorch.

Appendix F Additional figures

We here show the learning curves when training a ResNet-101 and a Momentum ResNet-101 on CIFAR-10.