Mix-LN: Liberare la Potenza dei Livelli Più Profondi Combinando Pre-LN e Post-LN

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

December 18, 2024

Autori: Pengxiang Li, Lu Yin, Shiwei Liu

cs.AI

Abstract

I Large Language Models (LLM) hanno ottenuto un notevole successo, tuttavia recenti scoperte rivelano che i loro strati più profondi spesso contribuiscono minimamente e possono essere potati senza influenzare le prestazioni complessive. Mentre alcuni vedono ciò come un'opportunità per la compressione del modello, identifichiamo ciò come una carenza nella formazione radicata nell'ampio utilizzo della Normalizzazione Pre-Strato (Pre-LN). Dimostriamo che il Pre-LN, comunemente impiegato in modelli come GPT e LLaMA, porta a norme di gradiente ridotte nei suoi strati più profondi, riducendone l'efficacia. Al contrario, la Normalizzazione Post-Strato (Post-LN) preserva norme di gradiente più elevate negli strati più profondi ma soffre di gradienti che svaniscono negli strati precedenti. Per affrontare ciò, introduciamo Mix-LN, una nuova tecnica di normalizzazione che combina i punti di forza del Pre-LN e del Post-LN all'interno dello stesso modello. Mix-LN applica il Post-LN agli strati più precoci e il Pre-LN agli strati più profondi, garantendo gradienti più uniformi tra gli strati. Ciò consente a tutte le parti della rete - sia gli strati superficiali che quelli profondi - di contribuire efficacemente alla formazione. Estesi esperimenti con varie dimensioni di modello da 70M a 7B dimostrano che Mix-LN supera costantemente sia il Pre-LN che il Post-LN, promuovendo norme di gradiente più bilanciate e salutari in tutta la rete e migliorando la qualità complessiva della preformazione LLM. Inoltre, dimostriamo che i modelli pre-addestrati con Mix-LN apprendono meglio rispetto a quelli che utilizzano Pre-LN o Post-LN durante il fine-tuning supervisionato (SFT) e il reinforcement learning dal feedback umano (RLHF), sottolineando l'importanza critica degli strati profondi di alta qualità. Affrontando efficacemente le inefficienze degli strati profondi nei LLM attuali, Mix-LN sblocca il loro potenziale, potenziando la capacità del modello senza aumentarne le dimensioni. Il nostro codice è disponibile su 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.