Mix-LN: Pre-LN과 Post-LN을 결합하여 더 깊은 층의 파워를 발휘하기

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

December 18, 2024

저자: Pengxiang Li, Lu Yin, Shiwei Liu

cs.AI

초록

대형 언어 모델 (LLM)은 놀라운 성과를 이루었지만 최근 연구 결과에 따르면 그들의 깊은 층은 종종 미미하게 기여하며 전체 성능에 영향을 미치지 않고 가지치기할 수 있다는 것을 밝혀냈다. 일부는 이를 모델 압축의 기회로 보지만, 우리는 이를 Pre-Layer Normalization (Pre-LN)의 널리 사용된 훈련 결함으로 식별한다. 우리는 GPT 및 LLaMA와 같은 모델에서 일반적으로 사용되는 Pre-LN이 그 깊은 층에서 그래디언트 노름을 감소시켜 효과를 줄인다는 것을 입증한다. 이에 반해 Post-Layer Normalization (Post-LN)은 깊은 층에서 더 큰 그래디언트 노름을 유지하지만 초기 층에서 사그라들어버리는 문제가 있다. 이를 해결하기 위해 우리는 Pre-LN과 Post-LN의 장점을 결합한 혁신적인 정규화 기술인 Mix-LN을 소개한다. Mix-LN은 초기 층에 Post-LN을 적용하고 깊은 층에 Pre-LN을 적용하여 더 균일한 그래디언트를 각 층에서 보장한다. 이를 통해 네트워크의 모든 부분 - 얕은 층과 깊은 층 모두 -이 효과적으로 훈련에 기여할 수 있게 한다. 70M에서 7B까지 다양한 모델 크기로 수행된 광범위한 실험에서 Mix-LN이 일관되게 Pre-LN과 Post-LN을 능가하여 네트워크 전체에 걸쳐 균형 잡히고 건강한 그래디언트 노름을 촉진하고 LLM 사전 훈련의 전반적인 품질을 향상시킨다는 것을 입증한다. 더 나아가 Mix-LN으로 사전 훈련된 모델이 Pre-LN 또는 Post-LN을 사용하는 모델에 비해 지도 미세 조정 (SFT) 및 인간 피드백으로부터의 강화 학습 (RLHF) 중에 더 나은 학습을 한다는 것을 입증하여 깊은 층의 고품질이 중요함을 강조한다. 현재 LLM의 깊은 층의 비효율성을 효과적으로 다루면서, Mix-LN은 그들의 잠재력을 발휘하여 모델 용량을 향상시키고 모델 크기를 증가시키지 않는다. 우리의 코드는 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.