BranchNorm: Escalonamento Robusto de Transformadores Extremamente Profundos

BranchNorm: Robustly Scaling Extremely Deep Transformers

May 4, 2023

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

cs.AI

Resumo

Recentemente, o DeepNorm escalona Transformers para profundidades extremas (ou seja, 1000 camadas) e revela o potencial promissor do escalonamento profundo. Para estabilizar o treinamento de modelos profundos, o DeepNorm (Wang et al., 2022) tenta restringir a atualização do modelo a um valor constante. Embora a aplicação de tal restrição possa beneficiar o estágio inicial do treinamento do modelo, ela pode levar a modelos subtreinados durante todo o procedimento de treinamento. Neste artigo, propomos o BranchNorm, que redimensiona dinamicamente o ramo não residual do Transformer de acordo com o período de treinamento. O BranchNorm não apenas estabiliza teoricamente o treinamento com normas de gradiente suaves no estágio inicial, mas também incentiva uma melhor convergência no estágio subsequente de treinamento. Resultados experimentais em múltiplas tarefas de tradução demonstram que o BranchNorm alcança um melhor equilíbrio entre estabilidade de treinamento e desempenho de convergência.

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.