ReZero is All You Need:Fast Convergence at Large Depth
Table 3: Comparison of various 12 layer Transformers normalization variants against ReZero and the training iterations required to reach 1.2 BPB on enwiki8 validation set.
['Model', 'Layers', 'Parameters', 'BPB']
To verify this, we re-ran the Post-Norm setup with 100 steps of learning rate warm-up and find that the model is able to converge to 1.2 BPB in 13,690 iterations. Under this setting, we compared other LayerNorm placement schemes against Post-Norm. We find that the other placements led to initially faster convergence, but ultimately Post-Norm catches up in performance, resulting in relatively slower convergence for Pre-Norm and GPT2-Norm. However, other LayerNorm placements have an advantage over Post-Norm in that they do not require learning rate warm-up, thus have fewer hyperparameters to tune. ReZero with α=1 does not show an improvement over the vanilla Transformer, indicating the importance of initializing α=0. With our proposed initialization of α=0, ReZero converges 56% faster than the vanilla Transformer. This cost comes from an increase in memory requirements and poor signal propagation. We find that trying to train deep vanilla Transformers leads to convergence difficulties. When scaled to 64 layers, the vanilla Transformer fails to converge even with a warm-up schedule. A ReZero Transformer with initialization of α=1 diverges, supporting our theoretically motivated initialization at α=0. The deeper ReZero Transformers are able to attain better performance than the shallower Transformers.