Mix-LN: Het Ontketenen van de Kracht van Diepere Lagen door Pre-LN en Post-LN te Combineren

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

Samenvatting

Grote Taalmodellen (LLM's) hebben opmerkelijk succes behaald, maar recente bevindingen tonen aan dat hun diepere lagen vaak minimaal bijdragen en kunnen worden weggesnoeid zonder de algehele prestaties te beïnvloeden. Terwijl sommigen dit zien als een kans voor modelcompressie, identificeren wij dit als een trainingsgebrek dat geworteld is in het wijdverbreide gebruik van Pre-Laagnormalisatie (Pre-LN). We tonen aan dat Pre-LN, veelgebruikt in modellen zoals GPT en LLaMA, leidt tot verminderde gradiëntnormen in de diepere lagen, waardoor hun effectiviteit afneemt. In tegenstelling hiermee behoudt Post-Laagnormalisatie (Post-LN) grotere gradiëntnormen in diepere lagen, maar heeft te maken met verdwijnende gradiënten in eerdere lagen. Om dit aan te pakken, introduceren we Mix-LN, een nieuwe normalisatietechniek die de sterke punten van Pre-LN en Post-LN combineert binnen hetzelfde model. Mix-LN past Post-LN toe op de eerdere lagen en Pre-LN op de diepere lagen, waardoor meer uniforme gradiënten over de lagen worden gegarandeerd. Dit maakt het mogelijk dat alle delen van het netwerk - zowel oppervlakkige als diepe lagen - effectief bijdragen aan de training. Uitgebreide experimenten met verschillende modelgroottes van 70M tot 7B tonen aan dat Mix-LN consequent beter presteert dan zowel Pre-LN als Post-LN, wat zorgt voor meer gebalanceerde, gezondere gradiëntnormen door het hele netwerk en het verbeteren van de algehele kwaliteit van LLM-voortraining. Bovendien tonen we aan dat modellen die zijn voorgetraind met Mix-LN beter leren in vergelijking met die welke Pre-LN of Post-LN gebruiken tijdens begeleid finetunen (SFT) en versterkend leren van menselijke feedback (RLHF), waarbij het cruciale belang van hoogwaardige diepe lagen wordt benadrukt. Door de inefficiënties van diepe lagen in huidige LLM's effectief aan te pakken, ontsluit Mix-LN hun potentieel, waardoor de modelcapaciteit wordt verbeterd zonder de modelgrootte te vergroten. Onze code is beschikbaar op 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.