Mix-LN: Liberando o Poder de Camadas Mais Profundas ao Combinar Pré-LN e Pós-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

Resumo

Os Modelos de Linguagem de Grande Escala (LLMs) alcançaram um sucesso notável, no entanto, descobertas recentes revelam que suas camadas mais profundas frequentemente contribuem minimamente e podem ser podadas sem afetar o desempenho geral. Enquanto alguns veem isso como uma oportunidade para compressão de modelos, identificamos como uma deficiência no treinamento enraizada no amplo uso da Normalização Pré-Camada (Pre-LN). Demonstramos que o Pre-LN, comumente empregado em modelos como GPT e LLaMA, leva a normas de gradiente diminuídas em suas camadas mais profundas, reduzindo sua eficácia. Em contraste, a Normalização Pós-Camada (Post-LN) preserva normas de gradiente maiores em camadas mais profundas, mas sofre com gradientes desaparecidos em camadas anteriores. Para lidar com isso, introduzimos o Mix-LN, uma técnica de normalização inovadora que combina os pontos fortes do Pre-LN e do Post-LN dentro do mesmo modelo. O Mix-LN aplica o Post-LN às camadas mais antigas e o Pre-LN às camadas mais profundas, garantindo gradientes mais uniformes entre as camadas. Isso permite que todas as partes da rede - tanto camadas rasas quanto profundas - contribuam efetivamente para o treinamento. Experimentos extensivos com vários tamanhos de modelo de 70M a 7B demonstram que o Mix-LN consistentemente supera tanto o Pre-LN quanto o Post-LN, promovendo normas de gradiente mais equilibradas e saudáveis em toda a rede, e aprimorando a qualidade geral do pré-treinamento de LLM. Além disso, demonstramos que modelos pré-treinados com Mix-LN aprendem melhor em comparação com aqueles que usam Pre-LN ou Post-LN durante o ajuste fino supervisionado (SFT) e aprendizado por reforço a partir do feedback humano (RLHF), destacando a importância crítica de camadas profundas de alta qualidade. Ao abordar efetivamente as ineficiências das camadas profundas nos atuais LLMs, o Mix-LN desbloqueia seu potencial, aprimorando a capacidade do modelo sem aumentar o tamanho do modelo. Nosso código está disponível em 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.