Variance Control via Weight Rescaling in LLM Pre-training

2025/03/21 by Louis Owen, Abhay Kumar, Owen, Louis +5 · 1 citation
Computer Science · #Computation and Language (cs.CL) #Domain Adaptation and Few-Shot Learning #FOS: Computer and information sciences #Machine Learning (cs.LG) #Machine Learning (stat.ML) #Natural Language Processing Techniques #Topic Modeling

paper · pdf · doi:10.48550/arxiv.2503.17500

openalex publication_date 2025/03/21 · openalex created_date 2025/10/10 · openalex updated_date 2026/07/28

Abstract

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.

Cited by

• The Spike, the Sparse and the Sink: Anatomy of Massive Activations and Attention Sinks

Related

• Crown, Frame, Reverse: Layer-Wise Scaling Variants for LLM Pre-Training
• WISCA: A Lightweight Model Transition Method to Improve LLM Training via Weight Scaling
• Weight Initialization and Variance Dynamics in Deep Neural Networks and Large Language Models
• Methods of improving LLM training stability
• Take the Bull by the Horns: Hard Sample-Reweighted Continual Training Improves LLM Generalization