Taming Transformer Without Using Learning Rate Warmup

Scaling Transformer to a large scale without using some technical tricks such as learning rate warump and using an obviously lower learning rate is an extremely challenging task, and is increasingly gaining more attention. In this paper, we provide a theoretical analysis for the process of training Transformer and reveal the rationale behind the model crash phenomenon in the training process, termed textitspectral energy concentration of bWqtop bWk, which is the reason for a malignant entropy collapse, where bWq and bWk are the projection matrices for the query and the key in Transformer, respectively. To remedy this problem, motivated by textitWeyl's Inequality, we present a novel optimization strategy, ie, making the weight updating in successive steps smooth -- if the ratio fracσ1(nabla bWt)σ1(bWt-1) is larger than a threshold, we will automatically bound the learning rate to a weighted multiple of fracσ1(bWt-1)σ1(nabla bWt), where nabla bWt is the updating quantity in step t. Such an optimization strategy can prevent spectral energy concentration to only a few directions, and thus can avoid malignant entropy collapse which will trigger the model crash. We conduct extensive experiments using ViT, Swin-Transformer and GPT, showing that our optimization strategy can effectively and stably train these Transformers without using learning rate warmup.