[Paper Review] m-RevNet: Deep Reversible Neural Networks with Momentum

Duo Li, Shanghua Gao|arXiv (Cornell University)|Aug 12, 2021

Advanced Neural Network Applications39 references4 citations

TL;DR

This paper proposes m-RevNet, a novel reversible neural network that integrates momentum into residual blocks by modeling the architecture through second-order ordinary differential equations (ODEs). By leveraging reversible dynamics and momentum, m-RevNet achieves superior memory efficiency and performance over ResNet, especially in challenging learning scenarios where standard ResNet fails.

ABSTRACT

In recent years, the connections between deep residual networks and first-order Ordinary Differential Equations (ODEs) have been disclosed. In this work, we further bridge the deep neural architecture design with the second-order ODEs and propose a novel reversible neural network, termed as m-RevNet, that is characterized by inserting momentum update to residual blocks. The reversible property allows us to perform backward pass without access to activation values of the forward pass, greatly relieving the storage burden during training. Furthermore, the theoretical foundation based on second-order ODEs grants m-RevNet with stronger representational power than vanilla residual networks, which potentially explains its performance gains. For certain learning scenarios, we analytically and empirically reveal that our m-RevNet succeeds while standard ResNet fails. Comprehensive experiments on various image classification and semantic segmentation benchmarks demonstrate the superiority of our m-RevNet over ResNet, concerning both memory efficiency and recognition performance.

Motivation & Objective

• To bridge deep neural network design with second-order ODEs for improved representational capacity.
• To address the high memory cost of training deep networks by enabling gradient computation without storing activations.
• To develop a reversible architecture that maintains strong performance even in learning scenarios where standard ResNet fails.
• To theoretically and empirically validate the advantages of momentum integration in residual blocks.

Proposed method

• Modeling the residual block dynamics using second-order ODEs to incorporate momentum into the network's hidden state evolution.
• Designing a reversible architecture where the forward pass is invertible, eliminating the need to store intermediate activations during backpropagation.
• Introducing a momentum term in the residual block update rule, derived from the second-order ODE formulation.
• Ensuring the invertibility of the transformation by structuring the network such that each block's output can be exactly recovered from its input using a learned inverse mapping.
• Using a discretization scheme compatible with reversible ODE solvers to train the network end-to-end.
• Formulating the network such that the hidden state evolution is governed by a second-order dynamical system, enabling richer representational dynamics than first-order ODEs in ResNets.

Experimental results

Research questions

• RQ1 Can modeling residual blocks via second-order ODEs with momentum lead to improved representational power compared to first-order ODEs in ResNets?
• RQ2 Does the integration of momentum in a reversible network architecture enhance performance on challenging learning tasks where standard ResNet fails?
• RQ3 To what extent does the reversible property of m-RevNet reduce memory consumption during training compared to standard ResNet?
• RQ4 How does the theoretical foundation based on second-order ODEs explain the empirical performance gains of m-RevNet?

Key findings

• m-RevNet achieves higher accuracy than ResNet on multiple image classification benchmarks, including CIFAR-10, CIFAR-100, and ImageNet.
• The model demonstrates superior performance on semantic segmentation tasks, outperforming ResNet while maintaining low memory usage.
• In certain learning scenarios, m-RevNet successfully trains and converges, whereas standard ResNet fails to learn effectively.
• The reversible architecture of m-RevNet enables memory-efficient training by eliminating the need to store intermediate activations.
• Theoretical analysis confirms that the second-order ODE formulation endows m-RevNet with stronger representational capacity than first-order ODE-based ResNets.
• Empirical results validate that momentum integration enhances optimization dynamics, particularly in deep architectures.