Reversible Architectures for Arbitrarily Deep Residual Neural Networks

TL;DR

The paper addresses the memory and stability challenges of training very deep ResNets by reinterpreting ResNets as discretizations of ordinary differential equations (ODEs) and deriving three stable, reversible architectures (Hamiltonian, Midpoint, Leapfrog). The authors prove that the Jacobians of these dynamical systems have (approximately) purely imaginary eigenvalues, ensuring forward and backward stability, and demonstrate memory-efficient backpropagation through reversibility. They validate the approach on CIFAR-10, CIFAR-100, and STL-10, achieving competitive or superior accuracy with significantly reduced memory usage, and show robust performance even with limited training data, including a 1202-layer Hamiltonian network. The work enables arbitrarily deep networks with modest computational resources, offering practical benefits for data-efficient training and scalable deep learning.

Abstract

Recently, deep residual networks have been successfully applied in many computer vision and natural language processing tasks, pushing the state-of-the-art performance with deeper and wider architectures. In this work, we interpret deep residual networks as ordinary differential equations (ODEs), which have long been studied in mathematics and physics with rich theoretical and empirical success. From this interpretation, we develop a theoretical framework on stability and reversibility of deep neural networks, and derive three reversible neural network architectures that can go arbitrarily deep in theory. The reversibility property allows a memory-efficient implementation, which does not need to store the activations for most hidden layers. Together with the stability of our architectures, this enables training deeper networks using only modest computational resources. We provide both theoretical analyses and empirical results. Experimental results demonstrate the efficacy of our architectures against several strong baselines on CIFAR-10, CIFAR-100 and STL-10 with superior or on-par state-of-the-art performance. Furthermore, we show our architectures yield superior results when trained using fewer training data.

Figures (4)

• Figure 1: The Hamiltonian Reversible Block. First, the input feature map is equally channel-wise split to $\mathbf{Y}_j$ and $\mathbf{Z}_j$. Then the operations described in Eq. \ref{['eqn:hamiltonian-discretized']} are performed, resulting in $\mathbf{Y}_{j+1}$ and $\mathbf{Z}_{j+1}$. Finally, $\mathbf{Y}_{j+1}$ and $\mathbf{Z}_{j+1}$ are concatenated as the output of the block.
• Figure 2: The Hamiltonian Reversible Neural Network. It is the simple stacking of several Hamiltonian reversible blocks as shown in Fig. \ref{['fig:hamiltonian-reversible-block']} and pooling layer.
• Figure 3: Hamiltonian vs ResNet test accuracy on CIFAR10 with a small subset of training data.
• Figure 4: Hamiltonian vs ResNet test accuracy for STL10 with a small subset of training data.