Methods: Reversible Residual Layers

Reversible Residual layers were presented in Gomez et al Ref. The idea is to make layers "reversible", i.e. each layer’s activations can be computed from the next layer’s activations. Adding Reversible Residial Layers enables us to perform backpropagation without storing the activations in memory. The memory requirements of architectures based on Reversible Residual blocks are independent of the number of layers in the model. Hence it allows one to train arbitrarily deep models, at the cost of increased compute.

The forward pass for a Reversible Residual Layer is given as

$y_1=x_1+\mathcal{F}(x_2)$

$y_2=x_2+\mathcal{G}(y_1)$

When doing a backward pass, the intermediate activations can be reconstructed as follows:

$x_2=y_2-\mathcal{G}(y_1)$

$x_1=y_1-\mathcal{F}(x_2)$

Applying this to Transformer architecture we set $\mathcal{F}=\mathbf{M}\mathbf{u}\mathbf{l}\mathbf{t}\mathbf{i}\mathbf{H}\mathbf{e}\mathbf{a}\mathbf{d}\mathbf{A}\mathbf{t}\mathbf{t}\mathbf{e}\mathbf{n}\mathbf{t}\mathbf{i}\mathbf{o}\mathbf{n}$ and $\mathcal{G}=\mathbf{F}\mathbf{e}\mathbf{e}\mathbf{d}\mathbf{F}\mathbf{o}\mathbf{r}\mathbf{w}\mathbf{a}\mathbf{r}\mathbf{d}$. $x_1$ and $x_2$ are initialized as copies of original input.

Our implementation of Reversible Residual Layers is largely inspired by https://github.com/lucidrains/reformer-pytorch and can be found here.