Mix-LN: Liberando el Poder de Capas más Profundas al Combinar Pre-LN y Post-LN

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

December 18, 2024

Autores: Pengxiang Li, Lu Yin, Shiwei Liu

cs.AI

Resumen

Los Modelos de Lenguaje Grandes (LLMs) han logrado un éxito notable, sin embargo, hallazgos recientes revelan que sus capas más profundas a menudo contribuyen mínimamente y pueden ser podadas sin afectar el rendimiento general. Mientras algunos ven esto como una oportunidad para la compresión del modelo, nosotros lo identificamos como una deficiencia en el entrenamiento arraigada en el uso generalizado de la Normalización Pre-Capa (Pre-LN). Demostramos que la Pre-LN, comúnmente empleada en modelos como GPT y LLaMA, conduce a normas de gradiente disminuidas en sus capas más profundas, reduciendo su efectividad. En contraste, la Normalización Post-Capa (Post-LN) preserva normas de gradiente más grandes en las capas más profundas pero sufre de gradientes desvanecientes en las capas más tempranas. Para abordar esto, presentamos Mix-LN, una técnica de normalización novedosa que combina las fortalezas de la Pre-LN y la Post-LN dentro del mismo modelo. Mix-LN aplica Post-LN a las capas más tempranas y Pre-LN a las capas más profundas, asegurando gradientes más uniformes a lo largo de las capas. Esto permite que todas las partes de la red, tanto las capas superficiales como las profundas, contribuyan efectivamente al entrenamiento. Experimentos extensos con varios tamaños de modelo desde 70M hasta 7B demuestran que Mix-LN supera consistentemente tanto a Pre-LN como a Post-LN, promoviendo normas de gradiente más equilibradas y saludables en toda la red, y mejorando la calidad general del pre-entrenamiento de LLM. Además, demostramos que los modelos pre-entrenados con Mix-LN aprenden mejor en comparación con aquellos que utilizan Pre-LN o Post-LN durante el ajuste fino supervisado (SFT) y el aprendizaje por refuerzo a partir de retroalimentación humana (RLHF), resaltando la importancia crítica de las capas profundas de alta calidad. Al abordar eficazmente las ineficiencias de las capas profundas en los LLM actuales, Mix-LN desbloquea su potencial, mejorando la capacidad del modelo sin aumentar su tamaño. Nuestro código está disponible en 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.