Mix-LN: Разблокирование мощи более глубоких слоев путем объединения предшествующего LN и последующего 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-LN). Мы демонстрируем, что Pre-LN, обычно применяемая в моделях, таких как GPT и LLaMA, приводит к уменьшению норм градиента в более глубоких слоях, снижая их эффективность. В отличие от этого, постнормализация слоев (Post-LN) сохраняет более крупные нормы градиента в более глубоких слоях, но страдает от затухания градиентов в более ранних слоях. Для решения этой проблемы мы представляем Mix-LN, новую технику нормализации, которая объединяет преимущества Pre-LN и Post-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.