BranchNorm: 極めて深いTransformerのロバストなスケーリング

BranchNorm: Robustly Scaling Extremely Deep Transformers

May 4, 2023

著者: Yijin Liu, Xianfeng Zeng, Fandong Meng, Jie Zhou

cs.AI

要旨

近年、DeepNormはTransformerを極めて深い層（例えば1000層）にスケールアップし、深層スケーリングの有望な可能性を明らかにしました。深層モデルの学習を安定化させるため、DeepNorm（Wang et al., 2022）はモデルの更新を一定値に制約することを試みています。このような制約を適用することは、モデル学習の初期段階では有益であるものの、学習プロセス全体を通じてモデルが十分に学習されない可能性があります。本論文では、学習期間に応じてTransformerの非残差ブランチを動的に再スケーリングするBranchNormを提案します。BranchNormは、理論的に初期段階で滑らかな勾配ノルムによる学習の安定化を実現するだけでなく、その後の学習段階でより良い収束を促進します。複数の翻訳タスクにおける実験結果は、BranchNormが学習の安定性と収束性能の間でより良いトレードオフを達成することを示しています。

English

Recently, DeepNorm scales Transformers into extremely deep (i.e., 1000 layers) and reveals the promising potential of deep scaling. To stabilize the training of deep models, DeepNorm (Wang et al., 2022) attempts to constrain the model update to a constant value. Although applying such a constraint can benefit the early stage of model training, it may lead to undertrained models during the whole training procedure. In this paper, we propose BranchNorm, which dynamically rescales the non-residual branch of Transformer in accordance with the training period. BranchNorm not only theoretically stabilizes the training with smooth gradient norms at the early stage, but also encourages better convergence in the subsequent training stage. Experiment results on multiple translation tasks demonstrate that BranchNorm achieves a better trade-off between training stability and converge performance.