0. TL; DR

本文提出了一种新的归一化技术Mix-LN，旨在解决大型语言模型（LLMs）中深度层训练不足的问题。研究表明，常用的Pre-LN会导致深度层的梯度减小，降低其有效性，而Post-LN虽然能保留深度层的梯度，但会导致早期层梯度消失。Mix-LN通过在模型的早期层应用Post-LN，在深度层应用Pre-LN，实现了更均匀的梯度分布，使网络的浅层和深层都能有效地参与训练。

1. 背景介绍

大型语言模型（LLMs）在理解和生成类似人类文本方面取得了巨大成功，推动了多个行业和学术领域的进步。然而，最近的研究发现，LLMs的深度层往往贡献较小，甚至可以被剪枝而不影响整体性能。这一现象被许多人视为模型压缩的机会，但本文作者认为这反映了训练过程中的不足，主要是由于广泛使用的Pre-LN导致的。Pre-LN在模型如GPT和LLaMA中被广泛使用，但会导致深度层的梯度减小，降低其有效性。相比之下，Post-LN虽然在深度层保留了较大的梯度，但会导致早期层的梯度消失。为了解决这一问题，本文提出了Mix-LN，一种结合了Pre-LN和Post-LN优势的新归一化技术。

2. Mix-LN技术

层归一化（LN）是一种在现代语言模型中广泛使用的归一化技术，它通过估计隐藏层中神经元输入的分布来确保输入分布的稳定性。原始Transformer架构中使用的是Post-LN，即在残差连接后应用LN。然而，后续研究表明，将LN放在残差连接前（Pre-LN）可以实现更稳定的性能，尤其是在大型语言模型中。本文通过实验验证了Pre-LN和Post-LN对LLMs深度层有效性的影响，并提出了Mix-LN技术。

Mix-LN的核心思想是结合Pre-LN和Post-LN的优势。具体来说，Mix-LN在模型的早期层（前$aL$层）应用Post-LN，在深度层（后$(1-a)L$层）应用Pre-LN。这样做的目的是利用Post-LN在深度层增强梯度流动的优势，同时利用Pre-LN在早期层稳定梯度的优势。通过这种方式，Mix-LN在中间和深度层实现了更健康的梯度范数，促进了整个网络的平衡训练，从而提高了模型的整体性能。

超参数$a$控制应用Post-LN的层的比例。通过在LLaMA-1B模型上进行超参数搜索，发现$a=0.25$时性能最佳，因此在所有模型大小上均使用了这一值。

3. 实验分析

实验使用了基于LLaMA架构的不同大小模型（从1.1B到5B参数），并比较了Post-LN、DeepNorm、Pre-LN和Mix-LN的性能。结果表明，Mix-LN在所有模型大小上均优于其他归一化技术，显著降低了困惑度，表明Mix-LN在预训练阶段就展现出了更好的性能。

为了验证Mix-LN在更大模型上的有效性，实验还使用了LLaMa-7B架构。尽管由于计算限制，训练只进行了13,000步，但Mix-LN在早期训练阶段就显示出了一致的性能提升，表明其优势在更大规模模型上也成立。

在监督微调阶段，Mix-LN预训练的模型在多个下游任务上的表现优于使用Pre-LN和Post-LN的模型。这表明Mix-LN训练的深度层能够捕捉到更丰富和多样化的特征，从而在复杂任务中实现更好的泛化能力。