BranchNorm: Robustes Skalieren extrem tiefer Transformer

BranchNorm: Robustly Scaling Extremely Deep Transformers

May 4, 2023

Autoren: Yijin Liu, Xianfeng Zeng, Fandong Meng, Jie Zhou

cs.AI

Zusammenfassung

Kürzlich hat DeepNorm Transformers auf extrem tiefe Architekturen (d.h. 1000 Schichten) skaliert und das vielversprechende Potenzial des tiefen Skalierens aufgezeigt. Um das Training tiefer Modelle zu stabilisieren, versucht DeepNorm (Wang et al., 2022), die Modellaktualisierung auf einen konstanten Wert zu beschränken. Obwohl die Anwendung einer solchen Beschränkung in der frühen Phase des Modelltrainings von Vorteil sein kann, kann sie während des gesamten Trainingsprozesses zu unzureichend trainierten Modellen führen. In diesem Artikel schlagen wir BranchNorm vor, das den nicht-residualen Zweig des Transformers dynamisch in Abhängigkeit von der Trainingsphase neu skaliert. BranchNorm stabilisiert nicht nur theoretisch das Training mit glatten Gradientennormen in der Anfangsphase, sondern fördert auch eine bessere Konvergenz in der späteren Trainingsphase. Experimentelle Ergebnisse bei mehreren Übersetzungsaufgaben zeigen, dass BranchNorm eine bessere Balance zwischen Trainingsstabilität und Konvergenzleistung erreicht.

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.