Variance Control via Weight Rescaling in LLM Pre-training
This addresses variance management for LLM developers, offering incremental improvements in training stability and performance.
The paper tackles the problem of variance control in LLM pre-training by introducing Layer Index Rescaling (LIR) initialization and Target Variance Rescaling (TVR) strategy, resulting in up to 4.6% improvement on downstream tasks and reduced activation values for better quantization.
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.