Contrôle de la variance par renormalisation des poids dans le pré-entraînement des LLM

Variance Control via Weight Rescaling in LLM Pre-training

March 21, 2025

Auteurs: Louis Owen, Abhay Kumar, Nilabhra Roy Chowdhury, Fabian Güra

cs.AI

Résumé

Le résultat du pré-entraînement des grands modèles de langage (LLM) dépend fortement des stratégies d'initialisation des poids et de contrôle de la variance. Bien que l'importance du contrôle initial de la variance ait été bien documentée pour les réseaux de neurones en général, la littérature sur l'initialisation et la gestion de sa croissance pendant le pré-entraînement des LLM, en particulier, est relativement limitée. Dans cet article, nous introduisons le schéma d'initialisation des poids par redimensionnement basé sur l'indice de couche (Layer Index Rescaling, LIR) et la stratégie de contrôle de la variance par redimensionnement vers une variance cible (Target Variance Rescaling, TVR). Des expériences sur un modèle LLaMA de 1 milliard de paramètres démontrent qu'une meilleure gestion de la variance grâce à ces techniques entraîne des améliorations substantielles des performances sur les tâches en aval (jusqu'à 4,6 % sur les benchmarks courants de pré-entraînement) et réduit les valeurs d'activation extrêmes, atténuant ainsi les défis associés à la quantification et à l'entraînement en faible précision. Notre code est disponible à l'adresse suivante : https://github.com/bluorion-com/weight_rescaling.

English

The outcome of Large Language Model (LLM) pre-training strongly depends on weight initialization and variance control strategies. Although the importance of initial variance control has been well documented in neural networks in general, the literature on initialization and management of its growth during LLM pre-training, specifically, is somewhat sparse. In this paper, we introduce the Layer Index Rescaling (LIR) weight initialization scheme, and the Target Variance Rescaling (TVR) variance control strategy. Experiments on a 1B parameter LLaMA model demonstrate that better variance management using these techniques yields substantial improvements in downstream task performance (up to 4.6% on common pre-training benchmarks) and reduces extreme activation values, thus mitigating challenges associated with quantization and low-precision training. Our code is available at: https://github.com/bluorion-com/weight_rescaling.