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

Hiroshi Horii; Sothea Has

stat.MLAug 18, 2025

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

Authors: Hiroshi Horii, Sothea Has

Organizations: SU · KHM

Abstract

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. Within this framework, we study the relationship between the quasi-stationary distribution derived from this equation and the initial distribution through the Kullback-Leibler (KL) divergence. As the quasi-steady-state distribution depends on the expected cost function, the KL divergence eventually reveals the connection between the expected cost function and the initialization distribution. 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. Then, by minimizing this bound, we obtain an optimal condition of the initialization variance in the Gaussian case. This result provides a concrete mathematical criterion, rather than a heuristic approach, to select the scale of weight initialization in DNNs. In addition, we experimentally confirm our theoretical results by using the classical SGD to train fully connected neural networks on the MNIST and Fashion-MNIST datasets. 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. 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.

Figures & tables

Panel of Figure 1 — Fig. 1: (a) Time evolution of the averaged weight variance $V(W,t)$ over $N=10$ runs of a three-layer deep neural network trained with SGD under different initial variances $\sigma_{0}$ . Results are shown for MNIST. (b) The figure shows the accuracy for various $\sigma_{0}$ values on MNIST. When $\sigma_{0}=0.15$ , the validation accuracy takes the highest value. The black dashed line represents the result using He-Normal initialization.

Panel of Figure 2 — Fig. 2: The relationship between the expected training loss $\mathbb{E}[L(W)]$ and the steady - state variance $\bar{V}(W)$ in Equation ( 18 ) for a DNN network trained using SGD under different initial variances $\sigma_{0}^{2}$ . The red linear function is the most relaxed bound for small $\sigma_{0}$ shown in Equation ( 18 ), and a star marker represents the result using He-Normal initialization. For each $\sigma_{0}$ , weights are independently sampled from the same Gaussian distribution $\mathscr{N}(0,\sigma_{0}^{2})$ , and the average variance $\bar{V}(W)$ over 10 runs is plotted on the x - axis with the MNIST dataset. The figures show (a) the overall variance for this network, (b) the variance of the output layer. Each point is colored according to the ratio $\bar{V}(W)/\sigma_{0}^{2}$ .

Panel of Figure 3 — Fig. 3: The relationship between the expected training loss $\mathbb{E}[L(W)]$ and the steady - state variance $\bar{V}(W)$ in Equation ( 18 ) for a modern complex network trained using SGD under different initial variances $\sigma_{0}^{2}$ . The red linear function is the most relaxed bound for small $\sigma_{0}$ shown in Equation ( 18 ), and a star marker represents the result using He-Normal initialization. For each $\sigma_{0}$ , weights are independently sampled from the same Gaussian distribution $\mathscr{N}(0,\sigma_{0}^{2})$ , and the average variance $\bar{V}(W)$ over 10 runs is plotted on the x - axis with the CIFAR-10 dataset. The figures show (a) the CNN two-layer model and (b) the VGG Net model. Each point is colored according to the ratio $\bar{V}(W)/\sigma_{0}^{2}$ .

Fokker–Planck Equation

See all 16 papers

Nov 10, 2025

Can Stationary Distributions of Scale-Invariant Neural Networks Be Described by the Thermodynamics of an Ideal Gas?

Understanding the training dynamics of deep neural networks remains a major open problem, with physics-inspired approaches offering promising insights. Building on this perspective, we develop a thermodynamic framework to describe the stationary distributions of stochastic gradient descent (SGD) with weight decay for scale-invariant neural networks, a setting that both reflects practical architectures with normalization layers and permits theoretical analysis. We establish analogies between training hyperparameters (e.g., learning rate, weight decay) and thermodynamic variables such as temperature, pressure, and volume. Starting with a simplified isotropic noise model, we uncover a close correspondence between SGD dynamics and ideal gas behavior, validated through theory and simulation. Extending to training of neural networks, we show that key predictions of the framework, including the behavior of stationary entropy, align closely with experimental observations. This framework provides a principled foundation for interpreting training dynamics and may guide future work on hyperparameter tuning and the design of learning rate schedulers.

Ildus Sadrtdinov, Ekaterina Lobacheva, Ivan Klimov +3

Constructor University · Mila – Quebec AI Institute · Université de Montréal +2

Statistical Physics of Learning Neural Network Training Dynamics Stochastic Gradient Descent

Ask PDF

Oct 17, 2026

Global Convergence in Deep Networks via the Second Law of Thermodynamics

Existing convergence analyses of finite-width overparameterized deep neural networks under standard parameterization often rely on parameters remaining close to their initialization and typically neglect the role of noise, despite empirical evidence that stochasticity can improve optimization and generalization. To bridge this gap, we study

L_2

-regularized continuous-time Langevin dynamics (CLD) with isotropic noise, for training feedforward networks with smooth activations. We introduce a change-of-measure inequality, inspired by the second law of thermodynamics, that enables control of key properties of the training dynamics under the time-evolving parameter distribution. Leveraging a sharpened lower bound on the minimum eigenvalue of the Neural Tangent Kernel (NTK) at initialization, we show that injected noise helps preserve NTK stability with high probability without confining the weights to a small neighborhood of their initialization. As a result, we derive a non-asymptotic upper bound on the expected training loss. Our analysis yields explicit conditions on width and noise temperature that guarantee convergence to a finite loss floor. Finally, we experimentally show the network Jacobian can change from initialization, departing the commonly analyzed lazy regime.

Matan Tsipory, Ofir Gaash, Itamar Harel +2

Technion - Israel Institute of Technology, The Andrew and Erna Viterbi Faculty of Electrical & Computer Engineering, Technion City, Haifa 3200003, Israel. · Technion - Israel Institute of Technology, Faculty of Data and Decision Sciences, Technion City, Haifa 3200003, Israel. · Toyota Technological Institute at Chicago, Chicago, IL, USA.

Neural Network Training Dynamics Stochastic Optimization Convergence Stochastic Gradient Langevin Dynamics Neural Network Stability Neural Tangent Kernel

Ask PDF

May 21, 2026

Why SGD is not Brownian Motion: A New Perspective on Stochastic Dynamics

Stochastic Gradient Descent (SGD) is commonly modeled as a Langevin process, assuming that minibatch noise acts as Brownian motion. However, this approximation relies on a continuous-time limit and a sqrt(eta) noise scaling that does not match the discrete SGD update at finite learning rate. In this work, we propose an alternative formulation of SGD as deterministic dynamics in a fluctuating loss landscape induced by minibatch sampling. Starting directly from the discrete update, we derive a master equation for the parameter distribution and obtain a discrete Fokker--Planck equation that differs from the standard Langevin form at order eta^2. Using this framework, we analyze SGD dynamics near critical points of the loss. We show that the behavior decomposes along the eigenbasis of the mean Hessian into qualitatively distinct regimes. In particular, nearly-flat directions do not admit a stationary distribution: the variance grows over time, corresponding to effective diffusion along valleys with a coefficient proportional to the learning rate. We provide empirical evidence supporting these predictions on neural network models in computer vision and natural language processing, observing a clear qualitative separation between confined and diffusive modes.

Igor Ignashin, Anna Radovskaya, Andrew Semenov +7

Basic Research of Artificial Intelligence Laboratory (BRAIn Lab) · P.N. Lebedev Physical Institute of the Russian Academy of Sciences · Innopolis University

Fokker–Planck Equation Neural Network Training Dynamics Gradient Descent Dynamics Stochastic Gradient Descent

Ask PDF

Abstract

Fokker–Planck Equation

Explore similar work

Can Stationary Distributions of Scale-Invariant Neural Networks Be Described by the Thermodynamics of an Ideal Gas?

Global Convergence in Deep Networks via the Second Law of Thermodynamics

Why SGD is not Brownian Motion: A New Perspective on Stochastic Dynamics