Optimal Condition for Initialization Variance in Deep Neural Networks: An SGD Dynamics Perspective

The gist: This work derives an analytic optimal condition for initialization variance in Gaussian weight initialization within Stochastic Gradient Descent (SGD) by analyzing the relationship between quasi-stationary distribution and initial distribution through Kullback-Leibler divergence, providing a concrete mathematical criterion to select weight initialization scale.

Theoretical Framework

The study recasts SGD as a continuous-time approximation of Langevin dynamics, which is then modeled using the Fokker-Planck equation to describe the evolution of the probability distribution function over time<ref:2508.128341, "Stochastic gradient descent (SGD), one of the most fundamental optimization algorithms in machine learning (ML) can be recast through a continuous-time approximation as a Fokker-Planck equation for Langevin dynamics, a viewpoint that has motivated many theoretical studies." The analysis establishes that the expected loss function bound is derived by considering the KL divergence between the stationary distribution and the initial Gaussian density<ref:2508.128343, By applying this to deep neural network models (DNNs), we can express the bounds of the expected loss function explicitly in terms of the initialization parameters. The paper assumes a quadratic form for the loss function around a deep local minimum, allowing for the description of Gibbs-like quasi-steady state distribution<ref:2508.1283410, We first assume that around the deep local minimum, the loss function is well approximated by a quadratic form, which is given by L(W) = 1/2 W T A W.

Derivation of Optimal Condition

The analysis leads to an inequality bounding the expected loss function based on the initial variance and steady-state variance<ref:2508.128343, Then, by minimizing this bound, we obtain an optimal condition of the initialization variance in the Gaussian case. The paper shows that for small initialization variances (σ 2 0 ≪ 1), the expected loss scales as E[L(W)] = O(V(W)], where V(W) is the variance of the parameter W at quasi-steady state<ref:2508.128343, We show that for small initialization variances (σ 2 0 ≪ 1), the expected loss scales as E[L(W)] = O(V(W)], where V(W) is the variance of the parameter W at quasi-steady state. The optimal condition is obtained by minimizing this upper bound with respect to σ 2 0<ref:2508.128343, By minimizing this bound, we obtain an optimal condition of the initialization variance in the Gaussian case.

Numerical Validation and Findings

The theoretical results are validated through numerical experiments conducted on MNIST and Fashion-MNIST datasets using a three-layer fully connected network trained with SGD<ref:2508.128344, In this study, we perform the simulation on the deep neural network (DNN) model. The simulations confirm that if the variance of the initialization distribution satisfies our theoretical optimal condition, then the corresponding DNN model always achieves lower final training loss and higher test accuracy than conventional He-normal initialization<ref:2508.128344, "The result shows that if the variance of the initialization distribution satisfies our theoretical optimal condition, then the corresponding DNN model always achieves lower final training loss and higher test accuracy than the conventional He-normal initialization." Figure 4 illustrates that at σ 2 0 = V¯ (W)/σ 2 0 = 1, the optimal condition of Theorem III.3 is satisfied, and the learning process also achieved the lowest loss<ref:2508.128347, At V¯ (W)/σ 2 0 = 1, the optimal condition of Theorem III.3 is satisfied, and the learning process also achieved the lowest loss.

Conclusion on Initialization Strategy

The paper concludes that choosing the initialization distribution is equivalent to selecting the exploration region in SGD<ref:2508.128349, Physically, choosing the initialization distribution is equivalent to choosing the exploration region in the learning process. The optimal condition implies that SGD begins its search in exactly the region of parameter space that will eventually contain the minimum loss solution<ref:2508.128349, This optimal condition means that the exploration region of SGD is already matched to the region where the final state is located. This approach yields more reliable optimization in landscapes riddled with multiple local minima<ref:2508.128349, Every weight is updated inside a productive basin from the first step, yielding more reliable optimization in landscapes riddled with multiple local minima. The optimal Gaussian parameter σ 2 0 = V¯ (W) is shown to improve both training loss and test accuracy compared to He-normal initialization<ref:2508.128349, In practice, DNN models are capable of achieving this optimal condition if they start from an ideal scale of the initialization distribution, especially in the Gaussian case. The work provides a mathematically grounded indicator that guides the choice of initialization variance and clarifies its physical meaning of the dynamics of parameters in DNNs<ref:2508.128341, Our work thus supplies a mathematically grounded indicator that guides the choice of initialization variance and clarifies its physical meaning of the dynamics of parameters in DNNs. The paper suggests future work could involve replacing KL divergence with metrics such as Wasserstein distance to recast the problem within optimal transport<ref:2508.128349, For future perspective, discovering the algorithms that actively search for the optimal initialization variance could further improve the efficiency of deep learning optimization. The analysis also notes that heavy-tailed initializations do not deliver uniformly high accuracy across trials<ref:2508.128349, Indeed, their divergent second moment often makes performance less stable than in the Gaussian case.

--- Page 1 ---

1

Optimal Condition for Initialization

Variance in Deep Neural Networks:

An SGD Dynamics Perspective

Hiroshi Horii1 and Sothea Has2<ref:2508.128347, Optimal Condition for Initialization Variance Variance in Deep Neural Networks: An SGD Dynamics Perspective.

--- Page 2 ---

2

strategy also has a significant impact on both convergence behavior and generalization performance of the model as well (see, for example, [10], [11])<ref:2508.128347, strategy also has a significant impact on both convergence behavior and generalization performance of the model as well (see, for example, [10], [11])

--- Page 3 ---

3

In addition, V (t) follows V(t) ∼ N(0, 1). Note that the coefficient of the stochastic term epsilon sqrt is obtained by considering lim∆t→0 sqrt(∆t η = dV sqrt.<ref:2508.128343, In addition, V (t) follows V(t) ∼ N(0, 1). Note that the coefficient of the stochastic term epsilon sqrt is obtained by considering lim∆t→0 sqrt(∆t η = dV sqrt.

--- Page 4 ---

4

From these assumptions, the main results of this study are presented in the following section.<ref:2508.128347, From these assumptions, the main results of this study are presented in the following section.

--- Page 5 ---

5

In summary, our analysis shows that the expected loss is highly sensitive to the scale of the initial weight distribution.<ref:2508.128347, In summary, our analysis shows that the expected loss is highly sensitive to the scale of the initial weight distribution.

--- Page 6 ---

6

In addition, for any input xi, the transformed value is defined by zi = xiW = (Pd j=1 xijW j1,...,Pd j=1 xijW jM) ∈ RM.<ref:2508.128344, In addition, for any input xi, the transformed value is defined by zi = xiW = (Pd j=1 xijW j1,...,Pd j=1 xijW jM) ∈ RM.

--- Page 7 ---

7

Fig. 3: The relationship between the expected training loss E[L(W)] and the steady-state variance V¯ (W) in Equation (18) for a DNN network trained using SGD under different initial variances σ 2 0.<ref:2508.128344, Fig. 3: The relationship between the expected training loss E[L(W)] and the steady-state variance V¯ (W) in Equation (18) for a DNN network trained using SGD under different initial variances σ 2 0.

--- Page 8 ---

8

In practice, DNN models are capable of achieving this optimal condition if they start from an ideal scale of the initialization distribution, especially in the Gaussian case.<ref:2508.128349, In practice, DNN models are capable of achieving this optimal condition if they start from an ideal scale of the initialization distribution, especially in the Gaussian case.

--- Page 9 ---

9

References

[1] A. Rakhlin, O. Shamir, and K.

1. Bold initialization variance optimization based on theoretical bounds: The AI system can be explicitly tuned to select an optimal initialization variance by minimizing an upper bound for expected loss, as derived in Theorem III.3: Ess[L(W)] ≤ ασ2/4b log(V¯ (W)) + log(2πeC1,loc), provided that the optimal standard deviation of the initial distribution is given by σ0 = r Ess[∥W∥2]K = q V¯ (W).

2. Guaranteed lower training loss through initialization matching: The system can achieve a lower training loss compared to standard He-normal initialization if the optimal condition is met, as demonstrated in Figure 4 where the point attaining the minimum expected loss occurs when σ20 = V¯ (W).

3. Dynamic exploration region matching: The AI system can be designed to dynamically modulate its weight distribution spread during training, aligning it with the theoretically optimal scale: the algorithm dynamically modulates the spread of the weight distribution by adapting the scale of the parameters around the optimal σ0 to align itself with the theoretically optimal size.

4. Improved generalization and accuracy: By setting initialization variance to σ0 = 0.15, which satisfies Theorem III.3, the model can achieve the highest accuracy in the final epoch, surpassing He-normal initialization trajectories in Figure 5 for both MNIST and Fashion-MNIST datasets.