Mix-LN: Unleashing the Power of Deeper Layers by Combining Pre-LN and Post-LN

大型语言模型（LLMs）取得了显著的成功，但最近的研究发现表明，其较深层的贡献通常很小，甚至可以在不影响整体性能的情况下进行剪枝。虽然有些人认为这是一个模型压缩的机会，但我们将其视为一个由广泛使用的预层归一化（Pre-LN）导致的训练不足问题。我们展示了预层归一化（Pre-LN），这种在GPT和LLaMA等模型中常用的归一化方法，会导致其较深层的梯度范数减小，从而降低其有效性。相比之下，后层归一化（Post-LN）在较深层保留了较大的梯度范数，但在早期层却遭受梯度消失的问题。为了解决这一问题，我们引入了混合层归一化（Mix-LN），这是一种新的归一化技术，它在同一模型中结合了Pre-LN和Post-LN的优点。Mix-LN对早期层应用Post-LN，对较深层应用Pre-LN，确保各层之间的梯度更加均匀。这使得网络的所有部分——无论是浅层还是深层——都能有效地参与到训练中。通过从70M到7B不同规模模型的大量实验表明，Mix-LN在所有情况下都优于Pre-LN和Post-LN，促进了更均衡、更健康的梯度范数分布，提升了LLM预训练的整体质量。此外，我们还证明了使用Mix-LN预训练的模型在监督微调（SFT）和基于人类反馈的强化学习（RLHF）过程中比使用Pre-LN或Post-LN的模型学习效果更好，突显了高质量深层的重要性。通过有效解决当前LLMs中深层效率低下的问题，Mix-LN释放了它们的潜力，在不增加模型大小的情况下增强了模型容量。我们的代码可在https://github.com/pixeli99/MixLN获取。

论文试图解决大型语言模型（LLMs）中深层网络贡献度低的问题，这一现象被广泛认为是训练过程中的不足，特别是由于Pre-Layer Normalization (Pre-LN) 导致的梯度范数在深层网络中减弱。这并不是一个全新的问题，但之前的研究大多集中在模型压缩上，而这篇论文则从改进训练方法的角度出发。

关键思路在于引入了一种新的归一化技术——Mix-LN，它结合了Pre-LN和Post-Layer Normalization (Post-LN) 的优势。具体来说，Mix-LN在模型的浅层应用Post-LN，在深层应用Pre-LN，从而确保了更均匀的梯度分布，使所有层次都能有效地参与训练。相比现有的研究，这种方法不仅解决了梯度消失的问题，还提升了模型的整体性能。

该论文通过一系列广泛的实验验证了Mix-LN的有效性，涵盖了不同规模的模型（从70M到7B参数）。此外，论文还展示了Mix-LN在监督微调（SFT）和基于人类反馈的强化学习（RLHF）中的优越表现。值得注意的是，作者提供了开源代码，使得其他研究人员可以复现实验结果并进一步探索。未来的工作可以深入研究如何将Mix-LN应用于更大规模的模型或不同的任务类型。

最近在这个领域，相关的研究还包括：1.《Understanding the Difficulty of Training Deep Feedforward Neural Networks》探讨了初始化对深度网络训练的影响；2.《Layer Normalization》首次提出了层归一化的概念；3.《On Layer Normalization in the Transformer Architecture》分析了Transformer架构中层归一化的位置选择。这些研究都为理解梯度行为和改进训练方法提供了重要的背景知识。