[Paper Review] ReZero is All You Need: Fast Convergence at Large Depth

Thomas Bachlechner, Bodhisattwa Prasad Majumder+3 more|arXiv (Cornell University)|Mar 10, 2020

Domain Adaptation and Few-Shot Learning39 references150 citations

TL;DR

This paper proposes ReZero, a zero-initialized residual gate that enables initial dynamical isometry, allowing training of extremely deep networks and faster convergence across FC nets, ResNets, and Transformers.

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.

Motivation & Objective

• Motivate the difficulty of training very deep networks due to vanishing/exploding gradients.
• Propose a simple, architecture-agnostic modification to enforce initial dynamical isometry.
• Demonstrate large-depth training capability and faster convergence across multiple architectures (FCN, CNN/ResNet, Transformers).
• Show empirical gains on CIFAR-10 and language modeling benchmarks such as enwiki8.

Proposed method

• Add a residual gate alpha to each layer and initialize alpha to zero: x_{i+1}=x_i+alpha_i F(x_i).
• Demonstrate that this initialization makes the network behave like the identity at start, achieving dynamical isometry.
• Apply ReZero to fully connected networks, convolutional ResNets, and Transformer architectures.
• Compare convergence speed and final performance against standard residuals and normalization-based approaches.
• Analyze Jacobian spectra to illustrate maintained signal propagation during training.

Experimental results

Research questions

• RQ1 Can a minimal, zero-initialized residual gate enable stable training of extremely deep networks without normalization?
• RQ2 How does ReZero affect convergence speed and test performance across FCNs, CNNs/ResNets, and Transformers?
• RQ3 To what extent does initializing alpha to zero influence training dynamics versus initializing to one?
• RQ4 What are the practical limits on depth and depth-related performance gains when applying ReZero to modern architectures?
• RQ5 Does ReZero reduce or eliminate the need for warm-up or heavy normalization in deep models?

Key findings

• ReZero enables training of networks with thousands of layers (e.g., 10,000-layer FC networks).
• Transformers with over 100 layers can be trained without LayerNorm or warm-up under ReZero.
• On enwiki8, a 12-layer Transformer with ReZero converges 56% faster than the vanilla Transformer to 1.2 BPB.
• On CIFAR-10, ResNet variants with ReZero achieve faster convergence and better test accuracy (e.g., ResNet-110 with ReZero shows improved metrics).
• ReZero accelerates CIFAR-10 training by up to ~32% in iterations to reach target accuracy across several ResNet depths.
• In deeper Transformers (64, 128 layers), ReZero maintains dynamical isometry better than Post-Norm/GPT2-Norm placements and enables training where vanilla Transformers diverge.