提升大语言模型训练稳定性的方法
计算与语言
2024-10-23 v1 机器学习
摘要
大语言模型 (LLM) 的训练稳定性是一个重要的研究课题。重现训练不稳定性成本高昂，因此我们使用一个具有 8.3 亿参数的小型语言模型，并通过实验采用更高的学习率来强制模型发散。训练不稳定的一个来源是注意力层中 logits 的增长。我们扩展了先前工作的关注点，不仅关注 logits 的幅度，还关注 Transformer 模块中所有线性层的输出。我们观察到，在高学习率下，所有线性层输出的 L2 范数会随着每个训练步骤而增长，导致模型发散。具体来说，我们观察到 QKV、Proj 和 FC2 层的输出幅度增长最大。这促使我们探索几种方案：1) 不仅对 QK 层应用层归一化，也对 Proj 和 FC2 层应用层归一化；2) 在 QKV 层之后应用层归一化（并移除预归一化）；3) 将 QK 层归一化与 softmax 上限结合使用。我们证明，与仅基于 QK 层归一化的方法相比，使用后两种方法可以将学习率提高 1.5 倍（而模型不发散）。同时，我们观察到与基线模型相比，所有三种方法的困惑度都有显著改善。
引用
@article{arxiv.2410.16682,
title = {Methods of improving LLM training stability},
author = {Oleg Rybakov and Mike Chrzanowski and Peter Dykas and Jinze Xue and Ben Lanir},
journal= {arXiv preprint arXiv:2410.16682},
year = {2024}
}