Momentum Residual Neural Networks

Abstract

The training of deep residual neural networks (ResNets) with backpropagation has a memory cost that increases linearly with respect to the depth of the network. A way to circumvent this issue is to use reversible architectures. In this paper, we propose to change the forward rule of a ResNet by adding a momentum term. The resulting networks, momentum residual neural networks (Momentum ResNets), are invertible. Unlike previous invertible architectures, they can be used as a drop-in replacement for any existing ResNet block. We show that Momentum ResNets can be interpreted in the infinitesimal step size regime as second-order ordinary differential equations (ODEs) and exactly characterize how adding momentum progressively increases the representation capabilities of Momentum ResNets. Our analysis reveals that Momentum ResNets can learn any linear mapping up to a multiplicative factor, while ResNets cannot. In a learning to optimize setting, where convergence to a fixed point is required, we show theoretically and empirically that our method succeeds while existing invertible architectures fail. We show on CIFAR and ImageNet that Momentum ResNets have the same accuracy as ResNets, while having a much smaller memory footprint, and show that pre-trained Momentum ResNets are promising for fine-tuning models.

arXiv comment24 pages

References in corpus (10)

• TensorFlow: Large-Scale Machine Learning on Heterogeneous Distributed Systems
• An overview of gradient descent optimization algorithms
• Convolutional Neural Networks for Medical Image Analysis: Full Training or Fine Tuning?
• Dissecting Neural ODEs
• ANODE: Unconditionally Accurate Memory-Efficient Gradients for Neural ODEs
• ANODEV2: A Coupled Neural ODE Evolution Framework
• Universal Approximation Property of Neural Ordinary Differential Equations
• Stochastic Training of Residual Networks: a Differential Equation Viewpoint
• Towards Understanding Normalization in Neural ODEs
• MomentumRNN: Integrating Momentum into Recurrent Neural Networks

Cited by in corpus (6)

• Heavy Ball Neural Ordinary Differential Equations
• Large-time asymptotics in deep learning
• Can convolutional ResNets approximately preserve input distances? A frequency analysis perspective
• Efficient and Accurate Gradients for Neural SDEs
• Neural Architecture Search via Bregman Iterations
• HeunNet: Extending ResNet using Heun's Methods