Reversible Fixup Networks for Memory-Efficient Training

Deep residual networks (ResNets) have been used successfully for many computer vision tasks, but are difficult to scale to high resolution images or 3D volumetric datasets. Activation memory requirements, which scale with network depth and mini-batch size, quickly become prohibitive. Recently, invertible neural networks (INNs) have been successfully applied to classification problems to reduce memory requirements during neural network training, enabling constant memory complexity in depth. This is accomplished by removing the need to store the input activations computed in the forward pass and instead reconstruct them on-the-fly during the backward pass. But existing approaches require additive coupling layers, which perform channel-wise splitting and add additional parameters to the standard ResNet. In addition, they require a large enough mini-batch size to perform effec-tive Batch Normalization. We propose RevUp, an invertible ResNet architecture with Fixup initialization that is memory efficient, achieves high generalization, removes mini-batch dependence and does not require modifications via coupling layers as seen in previous methods. We show that RevUp achieves competitive test accuracy against comparable baselines on CIFAR-10 and ILSVRC (ImageNet).