arXiv:2410.16682cs.CLcs.LG2024-10被引 15

通过优化线性层归一化提升大模型训练稳定性,显著放宽学习率限制。

Methods of improving LLM training stability

  • 在Transformer中对QKV、Proj和FC2层添加层归一化,抑制输出幅值增长。
  • 新方法使学习率可提升1.5倍,且不引发模型发散,困惑度显著降低。
  • 适合追求高效训练的大模型研究者,尤其关注稳定性的实践场景。

大语言模型训练稳定性是重要研究课题。为快速暴露训练不稳定性,我们使用830M参数的小型语言模型,在高学习率下强制模型发散。发现注意力层的logits增长是不稳定根源之一,但更广泛地观察到Transformer块中所有线性层的输出幅值(L2范数)在训练过程中持续增长,导致模型发散。其中,QKV、Proj和FC2层的输出增长最为显著。为此提出三种改进:1)不仅在QK层后加入层归一化,也在Proj和FC2层后增加;2)在QKV层后加归一化并移除预归一化;3)结合QK层归一化与softmax截断。结果显示,后两种方法相比仅使用QK层归一化的基准方法,学习率可提升1.5倍而不发散;三者均显著降低困惑度。

原文摘要 · Abstract (English)

Training stability of large language models(LLMs) is an important research topic. Reproducing training instabilities can be costly, so we use a small language model with 830M parameters and experiment with higher learning rates to force models to diverge. One of the sources of training instability is the growth of logits in attention layers. We extend the focus of the previous work and look not only at the magnitude of the logits but at all outputs of linear layers in the Transformer block. We observe that with a high learning rate the L2 norm of all linear layer outputs can grow with each training step and the model diverges. Specifically we observe that QKV, Proj and FC2 layers have the largest growth of the output magnitude. This prompts us to explore several options: 1) apply layer normalization not only after QK layers but also after Proj and FC2 layers too; 2) apply layer normalization after the QKV layer (and remove pre normalization). 3) apply QK layer normalization together with softmax capping. We show that with the last two methods we can increase learning rate by 1.5x (without model divergence) in comparison to an approach based on QK layer normalization only. Also we observe significant perplexity improvements for all three methods in comparison to the baseline model.

大模型训练归一化稳定性Transformer

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。