Understanding the Difficulty of Training Transformers (2020)

Abstract: Transformers have proved effective in many NLP tasks. However, their training requires non-trivial efforts regarding carefully designing cutting-edge optimizers and learning rate schedulers (e.g., conventional SGD fails to train Transformers effectively). Our objective here is to understand __what complicates Transformer training__ from both empirical and theoretical perspectives. Our analysis reveals that unbalanced gradients are not the root cause of the instability of training. Instead, we identify an amplification effect that influences training substantially—for each layer in a multi-layer Transformer model, heavy dependency on its residual branch makes training unstable, since it amplifies small parameter perturbations (e.g., parameter updates) and results in significant disturbances in the model output. Yet we observe that a light dependency limits the model potential and leads to inferior trained models. Inspired by our analysis, we propose Admin (Adaptive model initialization) to stabilize the early stage’s training and unleash its full potential in the late stage. Extensive experiments show that Admin is more stable, converges faster, and leads to better performance

## introduction transformers @xcite have led to a series of breakthroughs in various deep learning tasks @xcite . they do not contain recurrent connections and can parallelize all computations in the same layer, thus improving effectiveness, efficiency, and scalability. training transformers, however, requires extra efforts. for example, although stochastic gradient descent (sgd) is the standard algorithm for conventional rnns and cnns, it converges to bad/suspicious local optima for transadmin not only stabilizes model training but unleashes the model potential for better performance. here, we conduct comprehensive analyses in empirical and theoretical manners to answer the question: what complicates transformer training. our analysis starts from the observation: the original transformer (referred to as post-ln) is less robust than its pre-ln variant foot_1 @xcite @xcite . we recognize that gradient vanishing issue is not the direct reason causing such difference, since fixing this issue alone cannot stabilize post-ln training. it implies that, besides unbalanced gradients, there exist other factors influencing model training greatly. with further analysis, we recognize that for each transformer residual block, the dependency table : intermediate output : residual output n: layer # d: hidden # h: head # ⇥n ⇥n layer norm in light of our analysis, we propose admin, an adaptive initialization method which retains the merits of pre-ln stability without hurting the performance. it restricts the layer dependency on its residual branches in the early stage and unleashes the model potential in the late stage. we conduct experiments on iwslt'14 de-en, wmt'14 en-de, and wmt'14 en-fr; admin is more stable, converges faster, and achieves better performance. for example, without introducing any additional hyper-parameters, admin successfully stabilizes 72-layer transformer training on wmt'14 en-fr and achieves a 43.80 bleu score. 3 for a residual block x + f (x), its shortcut output refers to x, its residual branch output refers to f (x), and the dependency on its residual branch refers to var[f (x)] var[x+f (x)] . ## preliminaries transformer architectures and notations. the transformer architecture contains two types of sub-layers, i.e., attention sub-layers and feedforward (ffn) sub-layers. they are composed of mainly three basic modules @xcite , i.e., layer norm (f ln ), multi-head attention (f att ), and feedforward network (f ffn ). as illustrated in figure 2 , the pre-ln transformer and the post-ln transformer organize these modules differently. for example, a pre-ln encoder organizes the self-attention sub-layer as layer norm. layer norm @xcite plays a vital role in transformer architecture. it is defined ## unbalanced gradients in this study, we strive to answer the question: what complicates transformer training. our analysis starts from the observation: pre-ln training is more robust than post-ln, while post-ln is more likely to reach a better performance than pre-ln. in a parameter grid search (as in figure 10 ), pre-ln converges in all 15 settings, and post-ln diverges in 7 out of 15 settings; when post-ln converges, it outperforms pre-ln in 7 out of 8 settings. we seek to reveal the underlying factor that destabilizes post-ln training and restricts the performance of pre-ln. in this section, we focus on the unbalanced gradients (e.g., gradient vanishing). we find that, although post-ln suffers from gradient vanishing and pre-ln does not, gradient vanishing is not the direct reason causing the instability of post-ln. specifically, we first theoretically and empirically establish that only post-ln decoders suffer from gradient vanishing and post-ln encoders do not. we then observe that fixing the gradient vanishing issue alone cannot stabilize training. ## instability from amplification effect we find that unbalanced gradients are not the root cause of the instability of post-ln, which implies the existence of other factors influencing model training. now, we go beyond gradient vanishing and introduce the amplification effect. specifically, we first examine the difference between pre-ln and post-ln, including their early-stage and latestage training. then, we show that post-ln's training instability is attributed to layer dependency's amplification effect, which intensifies gradient updates and destabilizes training. ## conclusion in this paper, we study the difficulties of training transformers in theoretical and empirical manners. our study in section 3 suggests that the gradient vanishing problem is not the root cause of unstable transformer training. also, the unbalanced gradient distribution issue is mostly addressed by adaptive optimizers. in section 4, we reveal the root cause of the instability to be the strong dependency on residual branches, which amplifies the fluctuation caused by parameter changes and destabilizes model training. in light of our analysis, we propose admin, an adaptive initialization method to stabilize transformers training. it controls the dependency at the beginning of training and maintains the flexibility to capture those dependencies once training stabilizes. extensive experiments verify our intuitions and show that, without introducing additional hyper-parameters, admin achieves more stable training, faster convergence, and better performance. our work opens up new possibilities to not only further push the state-of-the-art but understand deep network training better. it leads to many interesting future works, including generalizing theorem 2 to other models, designing new algorithms to automatically adapt deep networks to different training configurations, upgrading the transformer architecture, and applying our proposed admin to conduct training in a larger scale.