Source-linked AI summary
The Reversible Residual Network: Backpropagation Without Storing Activations
Aidan N. Gomez, Mengye Ren, Raquel Urtasun, Roger B. Grosse
TL;DR
Backpropagation makes activation storage a bottleneck as networks grow, motivating a method that preserves deep-network performance with lower memory use. RevNets reconstruct most activations during backward computation rather than storing them, and achieve nearly identical accuracy to standard ResNets across three image-classification benchmarks. The trade-off is that non-reversible layers still require storage, while reversible backpropagation adds computational overhead and can accumulate numerical error.
Problem
As networks grow deeper and wider, storing activations for backpropagation creates a substantial memory burden that can exceed GPU capacity.
Method
RevNets partition channels into reversible blocks whose activations can be reconstructed online during backpropagation, avoiding storage for most layers.
Results
RevNets achieve nearly identical classification accuracy to standard ResNets on CIFAR-10, CIFAR-100, and ImageNet.
Takeaways & Limitations
RevNets provide considerable memory-efficiency gains with little or no performance cost, supporting larger networks under limited computational resources.
Takeaways & Limitations
Stride-1 reversible blocks cannot replace information-discarding strided layers, so activations must still be stored for non-reversible layers.
Abstract
from arXiv · showhide
Deep residual networks (ResNets) have significantly pushed forward the state-of-the-art on image classification, increasing in performance as networks grow both deeper and wider. However, memory consumption becomes a bottleneck, as one needs to store the activations in order to calculate gradients using backpropagation. We present the Reversible Residual Network (RevNet), a variant of ResNets where each layer's activations can be reconstructed exactly from the next layer's. Therefore, the activations for most layers need not be stored in memory during backpropagation. We demonstrate the effectiveness of RevNets on CIFAR-10, CIFAR-100, and ImageNet, establishing nearly identical classification accuracy to equally-sized ResNets, even though the activation storage requirements are independent of depth.
1 Introduction
Deepening and widening networks improve visual-processing performance but make activation storage during backpropagation a major memory bottleneck. RevNets address this by reconstructing activations from subsequent layers while retaining nearly identical classification accuracy to standard ResNets.
- Motivation: Growing network depth and width increases the memory burden of storing activations for backpropagation.GPU memory limits can force expensive, complicated multi-GPU training for very deep architectures.
- Contribution: RevNets reconstruct each layer’s activations from the next layer’s activations, avoiding activation storage for most layers.The method is a reversible variant of ResNets.
- Contribution: Activation storage requirements become independent of network depth, typically using at least an order of magnitude less memory than equally sized ResNets.Only a handful of non-reversible layers require explicit activation storage.
- Results: RevNets achieve nearly identical classification accuracy to standard ResNets on CIFAR-10, CIFAR-100, and ImageNet.The reported result comes with only a modest increase in training time.
2 Background
Backpropagation computes network gradients in reverse through a computation graph, while residual blocks stabilize deep networks by passing signals directly through. RevNets build on these ideas alongside reversible transformations such as additive coupling.
- 2.1 Backpropagation: Backpropagation computes total derivatives of the cost with respect to nodes in a network’s computation graph.It processes nodes in reverse topological order and accounts for indirect effects through descendants.
- 2.2 Deep Residual Networks: Residual networks use skip-connected residual blocks to pass information directly and reduce exploding or vanishing gradient problems.Residual functions typically stack batch normalization, ReLU, and convolution layers.
- 2.2 Deep Residual Networks: ResNets use basic and bottleneck residual functions, with bottlenecks reducing feature dimensionality before a 3×3 convolution and output projection.The bottleneck contains three convolutions.
- 2.3 Reversible Architectures: NICE introduced reversible blocks based on nonlinear bijective transformations between data and latent spaces.The architecture’s invertibility and unit-determinant Jacobian make log-likelihoods and gradients tractable, while constraining mappings to be volume-preserving.
- 2.3 Reversible Architectures: Dinh et al. addressed NICE’s volume-preserving limitation with a reversible transformation whose Jacobian determinant is non-unit because of element-wise multiplication by exp(F(x1)).The Hadamard product denotes element-wise multiplication.
3 Methods
RevNets use reversible blocks to reconstruct activations during backpropagation, reducing storage while preserving compatibility with residual functions and handling non-reversible layers explicitly.
- Backpropagation Without Storing Activations: The backpropagation algorithm reconstructs activations online and uses ordinary automatic differentiation for the residual functions’ Jacobians.This preserves modularity across different choices of F and G.
- Reversible Residual Networks: RevNets partition each layer’s channels into two groups and use additive coupling with residual functions F and G.Channel-wise partitioning performed best in the reported experiments.
- Reversible Residual Networks: The reverse computation reconstructs x2 from y2 and y1, then reconstructs x1 from y1 and F(x2).These computations invert the forward block in exact arithmetic.
- Reversible Residual Networks: Reversible blocks require stride 1 because larger strides discard information; non-reversible layers therefore require explicit activation storage.Typical architectures contain only a handful of such non-reversible layers.
- Backpropagation Without Storing Activations: Reversible backpropagation costs approximately 4N operations versus 3N for ordinary backpropagation, or roughly 33% more in the stated model.In practice, the authors report overhead closer to 50% when forward and backward passes are equally expensive on GPUs.
4 Related Work
Prior memory-saving approaches trade storage against recomputation or replace backpropagation; RevNets instead achieve constant activation storage through reversibility.
- Checkpointing: Checkpointing stores selected activations and recomputes the others, trading spatial complexity for additional computation.Recursive checkpointing reduces memory to O(log L) while increasing computation to O(L log L).
- Complexity Comparisons: For L identical layers, ordinary backpropagation has computational cost 2L and activation storage cost L.The comparison assumes unit forward, backward, and single-layer activation-storage costs.
- Complexity Comparisons: RevNet backpropagation uses O(1) storage and computational cost 3L in the simplified comparison.Only a single block must be stored under the stated assumptions.
- Alternative Gradient Methods: Synthetic-gradient methods achieve O(1) activation storage only after their estimators are trained, whereas estimator training requires all activations to be stored.These methods replace global backpropagation with gradient approximations based on node activations.
5 Experiments
The experiments compare RevNets with similarly sized ResNets on CIFAR and ImageNet, finding closely matched classification performance despite reversibility and modest computational overhead.
- 5 Experiments: RevNets were evaluated on CIFAR-10, CIFAR-100, and ImageNet against ResNets matched closely for computational depth and parameter count.Each reversible block has the computation depth of two original residual blocks, so the experiments approximately halved residual blocks and doubled channels per block.
- 5 Experiments: RevNets roughly matched traditional ResNets’ classification error across CIFAR-10, CIFAR-100, and ImageNet.In no condition did RevNets underperform ResNets by more than 0.5%, and some results were slightly better.
- 5 Experiments: Reversibility caused no noticeable per-iteration training slowdown, although each RevNet update was about 1.5-2× more expensive depending on implementation.The reported computational overhead reflects additional recomputation during reversible backpropagation.
- 5 Experiments: Gradient reconstruction remained numerically stable: the gradient angle increased during training but stayed small, with no instability after four additional epochs.Training curves using stored and reconstructed activations showed no meaningful difference in training efficiency or final performance.
6 Conclusion and Future Work
The conclusion presents RevNets as memory-efficient architectures that retain classification performance and points to applications where activation storage is a critical constraint.
- 6 Conclusion and Future Work: RevNets avoid storing activations for most layers and provide considerable memory efficiency at little or no performance cost.The conclusion identifies semantic segmentation and reversible recurrent networks as future applications.
- 6 Conclusion and Future Work: Semantic segmentation is a target application because high-resolution image processing requires large input patches while effective batch normalization requires sufficiently large batches.The authors describe this combination as a critical memory bottleneck.
- 6 Conclusion and Future Work: RevNets are intended to enable training larger and more powerful networks with limited computational resources.This envisioned use follows from the reversible block’s reduced activation-storage requirements.
7 Appendix
The appendix describes training settings and implementation strategies for realizing RevNet memory savings, including TensorFlow graph management and while-loop execution.
- 7 Appendix: CIFAR-10/100 training used batch size 100, learning rate 0.1, decays at 40K and 60K steps, and 80K total steps.Weight decay was 2 × 10^-4, momentum was 0.9, and data augmentation used random cropping and horizontal flipping.
- 7 Appendix: ImageNet training used batch size 256 across four Titan X GPUs with synchronous SGD, momentum 0.9, and 600K training steps.Learning-rate decays occurred at 160K, 320K, and 480K steps, with weight decay 1 × 10^-4.
- 7 Appendix: Realizing RevNet’s theoretical memory savings requires precise low-level GPU memory management.The authors experimented with TensorFlow implementations using Tensor Handles and segmented computation graphs.
- 7 Appendix: A tf.while_loop implementation with back_prop=False discards activations after their utility expires during forward and backward passes.This implementation trained a 600-layer RevNet on ImageNet using a single GPU.
- 7 Appendix: During reversible backpropagation, the implementation reconstructs block inputs, stops gradients through those reconstructions, recomputes outputs, and then differentiates to obtain input and parameter gradients.The reversible layer recovers x2 from y2 − G(z1) and x1 from y1 − F(x2).