ReZero is All You Need: Fast Convergence at Large Depth

TL;DR

This paper introduces a simple gating mechanism in residual networks, called ReZero, which achieves fast convergence and improved performance in deep neural networks by maintaining dynamical isometry.

Contribution

The paper demonstrates that a single zero-initialized gating parameter in residual connections ensures dynamical isometry, enabling training of much deeper networks with faster convergence.

Findings

ReZero enables training thousands of layers in fully connected networks.

ReZero accelerates convergence by 56% in 12-layer Transformers.

ReZero improves test performance on CIFAR-10.

Abstract

Deep networks often suffer from vanishing or exploding gradients due to inefficient signal propagation, leading to long training times or convergence difficulties. Various architecture designs, sophisticated residual-style networks, and initialization schemes have been shown to improve deep signal propagation. Recently, Pennington et al. used free probability theory to show that dynamical isometry plays an integral role in efficient deep learning. We show that the simplest architecture change of gating each residual connection using a single zero-initialized parameter satisfies initial dynamical isometry and outperforms more complex approaches. Although much simpler than its predecessors, this gate enables training thousands of fully connected layers with fast convergence and better test performance for ResNets trained on CIFAR-10. We apply this technique to language modeling and find that we can easily train 120-layer Transformers. When applied to 12 layer Transformers, it converges 56% faster on enwiki8.