Découplage Échelle-Distribution : Permettre un Entraînement Stable et Efficace des Modèles de Langage à Grande Échelle

Scale-Distribution Decoupling: Enabling Stable and Effective Training of Large Language Models

February 21, 2025

Auteurs: Ya Wang, Zhijian Zhuo, Yutao Zeng, Xun Zhou, Jian Yang, Xiaoqing Li

cs.AI

Résumé

La stabilité de l'entraînement constitue un défi persistant dans le pré-entraînement des grands modèles de langage (LLM), en particulier pour les architectures telles que les Transformers Post-Norm, qui sont sujettes à l'explosion et à la dissipation des gradients. Dans cet article, nous proposons le découplage échelle-distribution (Scale-Distribution Decoupling, SDD), une approche novatrice qui stabilise l'entraînement en découplant explicitement l'échelle et la distribution de la matrice de poids dans les couches entièrement connectées. SDD applique un mécanisme de normalisation pour réguler les activations et un vecteur d'échelle apprenable pour maintenir des gradients bien conditionnés, empêchant ainsi efficacement l'explosion et la dissipation des gradients. Cette séparation améliore l'efficacité de l'optimisation, en particulier dans les réseaux profonds, en assurant une propagation stable des gradients. Les résultats expérimentaux montrent que notre méthode stabilise l'entraînement pour diverses architectures de LLM et surpasse les techniques existantes dans différentes configurations de normalisation. De plus, la méthode proposée est légère et compatible avec les frameworks existants, ce qui en fait une solution pratique pour stabiliser l'entraînement des LLM. Le code est disponible à l'adresse suivante : https://github.com/kaihemo/SDD.

English

Training stability is a persistent challenge in the pre-training of large language models (LLMs), particularly for architectures such as Post-Norm Transformers, which are prone to gradient explosion and dissipation. In this paper, we propose Scale-Distribution Decoupling (SDD), a novel approach that stabilizes training by explicitly decoupling the scale and distribution of the weight matrix in fully-connected layers. SDD applies a normalization mechanism to regulate activations and a learnable scaling vector to maintain well-conditioned gradients, effectively preventing gradient explosion and dissipation. This separation improves optimization efficiency, particularly in deep networks, by ensuring stable gradient propagation. Experimental results demonstrate that our method stabilizes training across various LLM architectures and outperforms existing techniques in different normalization configurations. Furthermore, the proposed method is lightweight and compatible with existing frameworks, making it a practical solution for stabilizing LLM training. Code is available at https://github.com/kaihemo/SDD.