ReZero is All You Need: Fast Convergence at Large Depth

By Thomas C. Bachlechner, Bodhisattwa Prasad Majumder, Huanru Henry Mao, Garrison W. Cottrell, Julian McAuley. 2020

404 citations recorded by the source corpus.

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.

Subfields: Artificial Intelligence, Computer Vision and Pattern Recognition, Model Architectures, Optimization & Training

How researchers see this paper

The researchers whose trusted networks rank it highest, strongest advocate first. Every researcher gets a different order.

In Artificial Intelligence

In Computer Vision and Pattern Recognition

In Model Architectures

In Optimization & Training

Citations in the corpus

Flamingo: a Visual Language Model for Few-Shot Learning (2022, background): “This multiplies the output of a newly added layer by tanh(α) before adding it to the input representation from the residual connection, where α is a layer-specific learnable scalar initialized to 0 [4].”
Stable and low-precision training for large-scale vision-language models (2023, background): “Consequently, researchers have proposed initialization and scaling schemes which remedy this issue [1, 70, 4, 18].”
High-Performance Large-Scale Image Recognition Without Normalization (2021, methodology): “…a learnable zero-initialized scalar gain in addition to α which results in the residual block being initialized to the identity (except in transition layers), similar to Goyal et al. (2017); Zhang et al. (2019a); Bachlechner et al. (2020), which we find to improve stability for very deep networks.”
Generating Images with Sparse Representations (2021, methodology): “…of modiﬁcations to the original architecture that we found to improve stability, training speed, and memory consumption: ReZero We use ReZero (Bachlechner et al., 2020; De & Smith, 2020), multiplying each residual connection with a zero-initialized scalar value, which is optimized jointly…”
Do Transformer Modifications Transfer Across Implementations and Applications? (2021, methodology): “We explored “RMS (root-mean-square) norm” (Zhang and Sennrich, 2019) as an alternative to layer normalization as well as the Rezero (Bachlechner et al., 2020) initialization scheme, including combining Rezero with Layer Norm and RMS Norm.”
On the Transfer of Disentangled Representations in Realistic Settings (2020, methodology): “Each residual block consists of two 3x3 convolutions with a leaky ReLU nonlinearity, and a learnable scalar gating mechanism (Bachlechner et al., 2020).”
Going deeper with Image Transformers (2021, methodology): “As pointed in our paper, in that respect our LayerScale approach is more related to Rezero [2] and Skipinit [16], Fixup [76], and T-Fixup [34].”
MoViNets: Mobile Video Networks for Efficient Video Recognition (2021, methodology): “To improve convergence speed in searching and training we use ReZero [3] by applying zero-initialized learnable scalar weights that are multiplied with features before the ﬁnal sum in a residual block.”
Deep Transformers with Latent Depth (2020, background): “Depending on the position of layer norm sub-layer, backpropagating gradients through multiple layers may suffer from gradient vanishing [18, 30, 5].”
Characterizing signal propagation to close the performance gap in unnormalized ResNets (2021, background): “In unnormalized networks, one can preserve this beneﬁt by introducing a learnable scalar at the end of each residual branch, initialized to zero (Zhang et al., 2019; De & Smith, 2020; Bachlechner et al., 2020).”
NormFormer: Improved Transformer Pretraining with Extra Normalization (2021, methodology): “Similarly, some other approaches targeted initialization as well, in particular ReZero (Bachlechner et al., 2020), FixUp (Huang et al., 2020) and LookLinear (Balduzzi et al., 2017).”
A Survey of Transformers (2021, methodology): “ReZero [5] replace LN module with a learnable residual connection.”
VECO: Variable Encoder-decoder Pre-training for Cross-lingual Understanding and Generation (2020, result): “This phenomenon is consistent with the ﬁndings in recent researches about deep transformer (Liu et al., 2020a; Bachlechner et al., 2020).”
Transformers in Time-Series Analysis: A Tutorial (2022, methodology): “The ReZero approach replaces the layer normalization operationwith a trainable parameter, α, which is initialized to 0 in each residual layer [7].”
IDF++: Analyzing and Improving Integer Discrete Flows for Lossless Compression (2020, methodology): “(4) are adjusted to include a learnable scalar parameter that ensures initialization to the identity operator, similar to the rezero trick by Bachlechner et al. [3]: Here α is a learnable scalar parameter initialized to zero.”
A Survey on Vision Transformer (2020, neutral): “Some other normalization algorithms [249], [198], [5] have been proposed to improve training of transformer.”
DeepNet: Scaling Transformers to 1,000 Layers (2022, methodology): “In order to alleviate the above issue, there have been efforts on improving the optimization of deep Transformer by means of better initialization [16], [17], [18], or better architecture [15], [19], [20], [21].”
Image Synthesis from Layout with Locality-Aware Mask Adaption (2021, methodology): “Thirdly, we use ReZero [2] to stabilize the training.”
Understanding the Difficulty of Training Transformers (2020, methodology): “…that scaling residual outputs to smaller values helps to stabilize training (Hanin and Rolnick, 2018; Mishkin and Matas, 2015; Zhang et al., 2019a; Bachlechner et al., 2020; Goyal et al., 2017), there is no systematic analysis on what complicates Transformer training or its underlying connection…”
Cerebras-GPT: Open Compute-Optimal Language Models Trained on the Cerebras Wafer-Scale Cluster (2023, methodology): “Recent works, such as Maximal Update Parameterization (µP), also show techniques to improve training stability and performance as models scale up (e.g., Bachlechner et al. (2020); Yang et al. (2021)).”