How Reversible Residual Networks Eliminate Activation Storage

Reversible Residual Networks (RevNets) solve one of the biggest bottlenecks in artificial intelligence training: memory consumption. During standard neural network training, intermediate activations from every layer must be stored in memory during the forward pass so they can be used to calculate gradients during the backward pass. RevNets eliminate this storage requirement by using a clever mathematical architecture that allows the network to recompute any layer's input activations on the fly during backpropagation, using only the outputs from the subsequent layer. This article breaks down the exact mathematical and architectural mechanics of how RevNets achieve this backward recomputation.

The Memory Bottleneck in Standard Backpropagation

To train a deep neural network, the backpropagation algorithm applies the chain rule of calculus to compute gradients. For any given layer, calculating the gradient of the weights requires two things: the gradient coming from the next layer, and the original input activation to the current layer.

In traditional architectures like standard ResNets, the forward pass transforms an input \(x\) into an output \(y\) via a residual function \(F\):

\[y = x + F(x)\]

Because \(F(x)\) is a non-invertible operation (often involving ReLU activations, batch normalization, and convolutions), it is mathematically impossible to determine \(x\) if you only know \(y\). Therefore, the system must save \(x\) in the GPU's RAM during the forward pass. As networks grow deeper, storing these activations consumes gigabytes of memory, severely limiting model size and batch capacity.

The Reversible Block Architecture

RevNets bypass this limitation by redesigning the residual block to be bijective (one-to-one and invertible). Instead of passing a single tensor through a layer, a RevNet splits the input activation tensor into two equal parts along the channel dimension, denoted as \(x_1\) and \(x_2\).

The reversible block uses two residual functions, \(F\) and \(G\) (which are standard neural network layers like convolutions), to produce two outputs, \(y_1\) and \(y_2\). The forward pass is computed in an interleaved manner:

  1. \[y_1 = x_1 + F(x_2)\]
  2. \[y_2 = x_2 + G(y_1)\]

Because of this coupling structure—inspired by additive coupling layers in normalizing flows—the forward operation becomes completely invertible, regardless of whether the underlying functions \(F\) and \(G\) are invertible themselves.

Recomputing Activations Backward

During the forward pass, a RevNet discards \(x_1\) and \(x_2\) immediately after computing \(y_1\) and \(y_2\). Only the final output activations of the last reversible sequence are kept in memory.

When backpropagation begins, the network moves backward layer by layer. To compute the gradients for a specific layer, it first reconstructs that layer's inputs (\(x_1\) and \(x_2\)) using the already known outputs (\(y_1\) and \(y_2\)). This is done by reversing the forward equations through simple subtraction:

  1. Recover \(x_2\): Because \(y_1\) is already known from the next layer, the network runs \(y_1\) forward through the function \(G\). It then subtracts the result from \(y_2\): \[x_2 = y_2 - G(y_1)\]

  2. Recover \(x_1\): Now that \(x_2\) has been successfully reconstructed, the network runs \(x_2\) forward through the function \(F\). It then subtracts this result from \(y_1\): \[x_1 = y_1 - F(x_2)\]

Note that the functions \(F\) and \(G\) themselves are never inverted. During the backward pass, \(F\) and \(G\) are actually executed as standard forward operations, and their outputs are simply subtracted instead of added.

Simultaneous Gradient Calculation

Once \(x_1\) and \(x_2\) are recomputed for the current block, the network temporarily holds them in memory just long enough to compute the local gradients for the weights inside functions \(F\) and \(G\), as well as the activation gradients needed to continue backpropagation.

After the gradients for the current block are calculated, \(y_1\) and \(y_2\) are deleted from memory, and the newly reconstructed \(x_1\) and \(x_2\) take their place as the "outputs" of the preceding layer. The process then repeats for the previous block.

The Trade-off: Compute for Memory

By recomputing activations backward from the next layer's outputs, RevNets reduce the memory complexity of storing activations from \(O(L)\)—where \(L\) is the number of layers—down to \(O(1)\). The memory required for activations becomes completely independent of the depth of the network.

The cost of this memory elimination is one additional forward computation per layer during the backward pass (evaluating \(F\) and \(G\) to reconstruct the inputs). However, because modern AI hardware is significantly more constrained by memory bandwidth and capacity than by raw floating-point compute power, this small computational overhead (roughly 33% more training time per step) is a massive net positive, enabling the training of vastly deeper and wider architectures on limited hardware.