GPAS : Accélération de la convergence du pré-entraînement des LLM via une mise à l'échelle des activations préservant le gradient

GPAS: Accelerating Convergence of LLM Pretraining via Gradient-Preserving Activation Scaling

June 27, 2025

Auteurs: Tianhao Chen, Xin Xu, Zijing Liu, Pengxiang Li, Xinyuan Song, Ajay Kumar Jaiswal, Fan Zhang, Jishan Hu, Yang Wang, Hao Chen, Shizhe Diao, Shiwei Liu, Yu Li, Yin Lu, Can Yang

cs.AI

Résumé

Les modèles de langage modernes de grande taille, tels que les séries LLaMA, Qwen et DeepSeek, adoptent principalement l'architecture Transformer avec Pre-LayerNorm (Pre-LN). Bien que stable lors du pré-entraînement et adaptable à des tailles de modèles importantes, Pre-LN souffre d'une croissance exponentielle de la variance des activations à travers les couches, ce qui entraîne une domination du chemin résiduel sur les sorties des sous-couches et limite la capacité d'apprentissage des couches plus profondes. Pour atténuer ce problème, nous proposons le Gradient-Preserving Activation Scaling (GPAS), une technique simple pouvant être combinée avec des approches existantes. GPAS fonctionne en réduisant l'échelle des activations intermédiaires tout en conservant leurs gradients inchangés. Cela préserve l'information contenue dans les activations et évite le problème de disparition des gradients associé à la réduction d'échelle des gradients. Des expériences approfondies sur diverses tailles de modèles, allant de 71M à 1B, montrent que GPAS permet des gains de performance constants. Au-delà de l'amélioration des Transformers Pre-LN, GPAS montre également un potentiel pour améliorer des architectures alternatives telles que Sandwich-LN et DeepNorm, démontrant ainsi sa polyvalence et son potentiel pour optimiser la dynamique d'entraînement dans un large éventail de configurations.

English

Modern Large Language Models, such as the LLaMA, Qwen and DeepSeek series, predominantly adopt the Pre-LayerNorm (Pre-LN) Transformer architecture. While being stable during pretraining and scalable to large model sizes, Pre-LN suffers from an exponential growth in activation variance across layers, causing the residual path to dominate over sub-layer outputs and limiting the learning capacity of deeper layers. To mitigate this issue, we propose Gradient-Preserving Activation Scaling (GPAS), a simple technique that can be used in combination with existing approaches. GPAS works by scaling down the intermediate activations while keeping their gradients unchanged. This leaves information in the activations intact, and avoids the gradient vanishing problem associated with gradient downscaling. Extensive experiments across various model sizes from 71M to 1B show that GPAS achieves consistent performance gains. Beyond enhancing Pre-LN Transformers, GPAS also shows promise in improving alternative architectures such as Sandwich-LN and DeepNorm, demonstrating its versatility and potential for improving training dynamics in a wide range of settings.