BranchNorm : Mise à l'échelle robuste des Transformers extrêmement profonds

BranchNorm: Robustly Scaling Extremely Deep Transformers

May 4, 2023

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

cs.AI

Résumé

Récemment, DeepNorm a permis de mettre à l'échelle les Transformers à des profondeurs extrêmes (c'est-à-dire 1000 couches) et a révélé le potentiel prometteur de la mise à l'échelle en profondeur. Pour stabiliser l'entraînement des modèles profonds, DeepNorm (Wang et al., 2022) tente de contraindre la mise à jour du modèle à une valeur constante. Bien que l'application d'une telle contrainte puisse bénéficier aux premières étapes de l'entraînement du modèle, elle peut conduire à des modèles sous-entraînés tout au long de la procédure d'entraînement. Dans cet article, nous proposons BranchNorm, qui redimensionne dynamiquement la branche non résiduelle du Transformer en fonction de la période d'entraînement. BranchNorm stabilise théoriquement l'entraînement avec des normes de gradient lisses dès les premières étapes, tout en favorisant une meilleure convergence lors des étapes ultérieures de l'entraînement. Les résultats expérimentaux sur plusieurs tâches de traduction démontrent que BranchNorm atteint un meilleur compromis entre stabilité de l'entraînement et performance de convergence.

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.