Mix-LN : Libérer la puissance des couches plus profondes en combinant Pre-LN et Post-LN

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

December 18, 2024

Auteurs: Pengxiang Li, Lu Yin, Shiwei Liu

cs.AI

Résumé

Les grands modèles de langage (LLMs) ont connu un succès remarquable, cependant, des découvertes récentes révèlent que leurs couches plus profondes contribuent souvent de manière minimale et peuvent être élaguées sans affecter les performances globales. Alors que certains voient cela comme une opportunité pour la compression de modèle, nous l'identifions comme une lacune dans l'entraînement enracinée dans l'utilisation généralisée de la Normalisation Avant-Couche (Pre-LN). Nous démontrons que la Pre-LN, couramment utilisée dans des modèles tels que GPT et LLaMA, entraîne une diminution des normes de gradient dans ses couches plus profondes, réduisant ainsi leur efficacité. En revanche, la Normalisation Après-Couche (Post-LN) préserve des normes de gradient plus élevées dans les couches plus profondes mais souffre de gradients qui disparaissent dans les couches antérieures. Pour remédier à cela, nous introduisons Mix-LN, une nouvelle technique de normalisation qui combine les avantages de la Pre-LN et de la Post-LN au sein du même modèle. Mix-LN applique la Post-LN aux couches antérieures et la Pre-LN aux couches plus profondes, garantissant des gradients plus uniformes à travers les couches. Cela permet à toutes les parties du réseau - à la fois les couches peu profondes et profondes - de contribuer efficacement à l'entraînement. Des expériences approfondies avec différentes tailles de modèle de 70M à 7B démontrent que Mix-LN surpasse systématiquement à la fois la Pre-LN et la Post-LN, favorisant des normes de gradient plus équilibrées et plus saines dans l'ensemble du réseau, et améliorant la qualité globale de la pré-formation des LLM. De plus, nous démontrons que les modèles pré-entraînés avec Mix-LN apprennent mieux par rapport à ceux utilisant la Pre-LN ou la Post-LN lors du fine-tuning supervisé (SFT) et de l'apprentissage par renforcement à partir des retours humains (RLHF), soulignant l'importance critique des couches profondes de haute qualité. En traitant efficacement les inefficacités des couches profondes dans les LLM actuels, Mix-LN libère leur potentiel, améliorant la capacité du modèle sans augmenter sa taille. Notre code est disponible sur https://github.com/pixeli99/MixLN.

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.