Weight Initialization and Variance Dynamics in Deep Neural Networks and Large Language Models
It provides incremental insights into initialization methods for practitioners training deep networks and transformers, improving robustness.
This paper investigates how weight initialization affects training stability and convergence in deep neural networks and large language models, finding that Kaiming initialization outperforms Xavier for ReLU networks and identifying stable initialization ranges (1e-2 to 1e-1) and depth-dependent variance patterns in transformers.
Weight initialization governs signal propagation and gradient flow at the start of training. This paper offers a theory-grounded and empirically validated study across two regimes: compact ReLU multilayer perceptrons and GPT-2-style transformers. First, a logarithmic sweep of the initial standard deviation maps vanishing and exploding regimes and identifies a broad stability band with standard deviations between 1e-2 and 1e-1. Second, a controlled comparison shows that Kaiming (fan-in) initialization converges faster and more stably than Xavier under ReLU, consistent with variance-preserving theory. Third, in a from-scratch 12-layer GPT-2-style model, this paper tracks layerwise Q/K/V weight variance through pretraining and observe depth-dependent equilibration into narrow bands: shallow layers expand rapidly while deeper layers change more gradually. Together, these results connect classic initialization principles with modern transformer behavior and yield simple, practical recipes for robust training.