BranchNorm: Escalado Robusto de Transformers Extremadamente Profundos

BranchNorm: Robustly Scaling Extremely Deep Transformers

May 4, 2023

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

cs.AI

Resumen

Recientemente, DeepNorm ha escalado Transformers a profundidades extremas (es decir, 1000 capas) y ha revelado el potencial prometedor del escalado profundo. Para estabilizar el entrenamiento de modelos profundos, DeepNorm (Wang et al., 2022) intenta restringir la actualización del modelo a un valor constante. Aunque aplicar tal restricción puede beneficiar la etapa inicial del entrenamiento del modelo, puede resultar en modelos subentrenados durante todo el procedimiento de entrenamiento. En este artículo, proponemos BranchNorm, que reescala dinámicamente la rama no residual del Transformer de acuerdo con el período de entrenamiento. BranchNorm no solo estabiliza teóricamente el entrenamiento con normas de gradiente suaves en la etapa inicial, sino que también fomenta una mejor convergencia en la etapa posterior de entrenamiento. Los resultados experimentales en múltiples tareas de traducción demuestran que BranchNorm logra un mejor equilibrio entre la estabilidad del entrenamiento y el rendimiento de convergencia.

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.