Variance Control via Weight Rescaling in LLM Pre-training
Journal:
arXiv
Published Date:
Mar 21, 2025
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.
