Mix-LN: Die Kraft tieferer Schichten entfesseln durch die Kombination von Pre-LN und Post-LN.

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

December 18, 2024

Autoren: Pengxiang Li, Lu Yin, Shiwei Liu

cs.AI

Zusammenfassung

Große Sprachmodelle (LLMs) haben bemerkenswerte Erfolge erzielt, doch aktuelle Erkenntnisse zeigen, dass ihre tieferen Schichten oft nur minimalen Beitrag leisten und ohne Auswirkungen auf die Gesamtleistung beschnitten werden können. Während einige dies als Chance für die Modellkompression betrachten, identifizieren wir es als ein Trainingsdefizit, das in der weit verbreiteten Verwendung der Vor-Schichten-Normalisierung (Pre-LN) verwurzelt ist. Wir zeigen auf, dass Pre-LN, das üblicherweise in Modellen wie GPT und LLaMA verwendet wird, zu verringerten Gradientennormen in den tieferen Schichten führt, was ihre Wirksamkeit reduziert. Im Gegensatz dazu bewahrt die Nach-Schichten-Normalisierung (Post-LN) größere Gradientennormen in den tieferen Schichten, leidet jedoch unter verschwindenden Gradienten in den früheren Schichten. Um dies anzugehen, führen wir Mix-LN ein, eine neuartige Normalisierungstechnik, die die Stärken von Pre-LN und Post-LN innerhalb desselben Modells kombiniert. Mix-LN wendet Post-LN auf die früheren Schichten und Pre-LN auf die tieferen Schichten an, um gleichmäßigere Gradienten über die Schichten hinweg zu gewährleisten. Dies ermöglicht es allen Teilen des Netzwerks - sowohl den oberflächlichen als auch den tiefen Schichten -, effektiv zum Training beizutragen. Umfangreiche Experimente mit verschiedenen Modellgrößen von 70M bis 7B zeigen, dass Mix-LN durchweg sowohl Pre-LN als auch Post-LN übertrifft, indem es ausgewogenere, gesündere Gradientennormen im gesamten Netzwerk fördert und die Gesamtqualität des LLM-Vor-Trainings verbessert. Darüber hinaus zeigen wir, dass Modelle, die mit Mix-LN vorab trainiert wurden, im Vergleich zu denen, die Pre-LN oder Post-LN während des überwachten Feintunings (SFT) und des Lernens durch menschliches Feedback (RLHF) verwenden, besser lernen, was die entscheidende Bedeutung hochwertiger tiefer Schichten hervorhebt. Indem Mix-LN die Ineffizienzen der tiefen Schichten in aktuellen LLMs effektiv angeht, hebt es ihr Potenzial und verbessert die Modellkapazität, ohne die Modellgröße zu erhöhen. Unser Code ist unter https://github.com/pixeli99/MixLN verfügbar.

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.