BranchNorm: Robuust schalen van extreem diepe Transformers

BranchNorm: Robustly Scaling Extremely Deep Transformers

May 4, 2023

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

cs.AI

Samenvatting

Onlangs heeft DeepNorm Transformers geschaald naar extreem diepe architecturen (d.w.z. 1000 lagen) en het veelbelovende potentieel van diepe schaling onthuld. Om de training van diepe modellen te stabiliseren, probeert DeepNorm (Wang et al., 2022) de modelupdate te beperken tot een constante waarde. Hoewel het toepassen van een dergelijke beperking voordelig kan zijn in de vroege fase van de modeltraining, kan het leiden tot onvoldoende getrainde modellen gedurende het hele trainingsproces. In dit artikel stellen we BranchNorm voor, dat dynamisch de niet-residuele tak van de Transformer herschaalt in overeenstemming met de trainingsperiode. BranchNorm stabiliseert niet alleen theoretisch de training met vloeiende gradientnormen in de vroege fase, maar bevordert ook een betere convergentie in de latere trainingsfase. Experimentele resultaten op meerdere vertaaltaaken tonen aan dat BranchNorm een betere balans bereikt tussen trainingsstabiliteit en convergentieprestaties.

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.