Methods of improving LLM training stability

大语言模型（LLMs）的训练稳定性是一个重要的研究课题。重现训练不稳定现象可能会耗费大量资源，因此我们使用了一个拥有8.3亿参数的小型语言模型，并通过提高学习率来迫使模型发散。训练不稳定的其中一个来源是注意力层中logits的增长。我们扩展了前人工作的关注点，不仅考察logits的大小，还考察Transformer块中所有线性层的输出。我们观察到，在高学习率下，所有线性层输出的L2范数会随着每次训练步骤而增长，最终导致模型发散。具体来说，我们发现QKV、Proj和FC2层的输出幅度增长最大。这促使我们探索几种方案：1）不仅在QK层之后应用层归一化，还在Proj和FC2层之后也应用；2）在QKV层之后应用层归一化（并移除预归一化）；3）结合softmax上限一起应用QK层归一化。我们表明，与仅基于QK层归一化的方法相比，最后两种方法可以将学习率提高1.5倍（而不导致模型发散）。此外，我们观察到，与基线模型相比，这三种方法都显著改善了困惑度。

该论文旨在解决大型语言模型（LLMs）训练过程中出现的稳定性问题，特别是高学习率导致的模型发散问题。这是一个在深度学习社区内广泛关注的问题，尤其是在训练大规模模型时。

论文的关键思路是不仅关注注意力层中的logits增长，还扩展到Transformer块中所有线性层的输出。通过观察QKV、Proj和FC2层的输出L2范数的增长，提出几种方法来提高模型的稳定性，包括在Proj和FC2层后应用层归一化，以及在QKV层后应用层归一化并移除预归一化。这些方法可以显著提高学习率而不导致模型发散，并且改善了模型的困惑度。

论文通过实验验证了提出的几种方法的有效性，特别是在增加学习率和提高模型性能方面。实验使用了一个8.3亿参数的小型语言模型，并通过高学习率强制模型发散来观察不同方法的效果。此外，论文还开源了实验代码，方便其他研究者复现和进一步探索。未来的工作可以进一步优化这些方法，并将其应用于更大规模的模型。

最近在这个领域内的相关研究包括： 1. "Understanding and Improving Layer Normalization" - 提出了一种改进的层归一化方法，以提高模型的训练稳定性。 2. "On Layer Normalization in the Transformer Architecture" - 深入分析了Transformer架构中层归一化的作用，并提出了新的归一化策略。 3. "Stable Training of Large Language Models" - 探讨了多种技术来提高大模型的训练稳定性，包括梯度裁剪和自适应学习率调整。