arXiv:2412.13795cs.LGcs.AI2024-12ICLR被引 36

混合层归一化让深层网络更有效,提升大模型训练质量。

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

  • 用Post-LN保浅层梯度,Pre-LN强深层梯度,混合使用
  • 70M到7B模型均表现更好,梯度分布更均衡
  • 适合追求更高训练效率与模型性能的研究者

大语言模型虽取得显著进展,但深层网络常贡献有限,可被剪枝而不影响性能。我们发现这源于广泛使用的Pre-LN导致深层梯度范数衰减。相比之下,Post-LN在深层保持较大梯度,但浅层易出现梯度消失。为此提出Mix-LN:前层用Post-LN,深层用Pre-LN,实现梯度均匀分布。在70M至7B不同规模模型上实验表明,Mix-LN优于Pre-LN与Post-LN,整体梯度更健康,预训练质量更高。经监督微调(SFT)和人类反馈强化学习(RLHF)验证,其学习能力更强。该方法无需增加模型大小即可释放深层潜力。代码已开源。

原文摘要 · Abstract (English)

Large Language Models (LLMs) have achieved remarkable success, yet recent findings reveal that their deeper layers often contribute minimally and can be pruned without affecting overall performance. While some view this as an opportunity for model compression, we identify it as a training shortfall rooted in the widespread use of Pre-Layer Normalization (Pre-LN). We demonstrate that Pre-LN, commonly employed in models like GPT and LLaMA, leads to diminished gradient norms in its deeper layers, reducing their effectiveness. In contrast, Post-Layer Normalization (Post-LN) preserves larger gradient norms in deeper layers but suffers from vanishing gradients in earlier layers. To address this, we introduce Mix-LN, a novel normalization technique that combines the strengths of Pre-LN and Post-LN within the same model. Mix-LN applies Post-LN to the earlier layers and Pre-LN to the deeper layers, ensuring more uniform gradients across layers. This allows all parts of the network--both shallow and deep layers--to contribute effectively to training. Extensive experiments with various model sizes from 70M to 7B demonstrate that Mix-LN consistently outperforms both Pre-LN and Post-LN, promoting more balanced, healthier gradient norms throughout the network, and enhancing the overall quality of LLM pre-training. Furthermore, we demonstrate that models pre-trained with Mix-LN learn better compared to those using Pre-LN or Post-LN during supervised fine-tuning (SFT) and reinforcement learning from human feedback (RLHF), highlighting the critical importance of high-quality deep layers. By effectively addressing the inefficiencies of deep layers in current LLMs, Mix-LN unlocks their potential, enhancing model capacity without increasing model size. Our code is available at https://github.com/pixeli99/MixLN.

大模型层归一化梯度优化

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