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

summary

The gist

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

In short

The study models Stochastic Gradient Descent (SGD) as Langevin dynamics to find an optimal weight initialization variance for deep neural networks. By comparing the initial distribution to the quasi-stationary distribution using Kullback-Leibler divergence, the authors derive a mathematical criterion. This criterion suggests that initializing weights with a specific variance minimizes expected loss and improves training performance compared to standard methods like He initialization.

Key concepts

Langevin Dynamics: SGD is approximated as continuous-time Langevin dynamics, which describes the stochastic movement of parameters during optimization. This framework allows researchers to treat the learning process as a physical system evolving over time, helping to understand how parameters move towards an optimal state.
Fokker-Planck Equation: This equation is used to model the probability distribution function of the parameters evolving under Langevin dynamics. It helps describe how the probability density of weights changes as SGD progresses, allowing for a more rigorous mathematical description of the optimization process.
Kullback-Leibler (KL) Divergence: KL divergence measures how much one probability distribution differs from another. In this study, it is used to compare the initial weight distribution with the quasi-stationary (steady-state) distribution reached by SGD. Minimizing this divergence helps determine the best starting point for weight initialization.
Quasi-Stationary Distribution: This is the distribution of weights that SGD settles into after a long training period, assuming it has reached a steady state. The paper analyzes how this final state relates to the initial random distribution of weights, which is crucial for understanding the impact of initialization.

Terminology used across episodes

Reinforcement Learning (RL) — explained in 25 episodes
Sample Complexity — explained in 19 episodes
Federated Learning — explained in 17 episodes
Quantum Fisher Information (QFI) — explained in 16 episodes
Retrieval-Augmented Generation (RAG) — explained in 16 episodes
Interpretability — explained in 15 episodes

This episode discusses

The paper

Optimal Condition for Initialization Variance in Deep Neural Networks: An SGD Dynamics Perspective · Read on arXiv

Hiroshi Horii, Sothea Has

Laboratoire Jacques-Louis Lions · Research and Innovation Center, Institute of Technology of Cambodia

DOI: 10.1109/TNNLS.2026.3721153

Transcript

Introduction to the show: ident: AI Radio. Generated commentary on the latest Artificial Intelligence papers.

Tom: I'm Tom.

Jane: And I'm Jane, here with Lu, Meng and our in-house large language model, Lalam.

Tom: Today's paper: "Optimal Condition for Initialization Variance in Deep Neural Networks".

Jane: The gist:

Tom: First, who's behind it and why it matters.

Paper summary: Tom: So, to wrap up this paper, "Optimal Condition for Initialization Variance in Deep Neural Networks: An SGD Dynamics Perspective" by Horii and Has, they found a way to use the Fokker-Planck equation and KL divergence to derive an optimal condition for initial weight variance <ref:2508.12834#pg1>.

Jane: The main thing is that this condition provides a concrete mathematical criterion for selecting the scale of weight initialization in DNNs, moving beyond just standard rules like He or Xavier initialization <ref:2508.12834#pg2>. It suggests that starting the search in exactly the region where the final solution will be located is mathematically optimal <ref:2508.12834#pg6>.

Lu: What this means for the broader field is that we can use this approach to select initialization variance based on the expected final parameter variance, (W), rather than just picking a fixed value <ref:2508.12834#pg6>. It gives us a way to improve both training loss and test accuracy compared to He-normal initialization <ref:2508.12834#pg4>.

Meng: From an engineering standpoint, it means we have a mathematically grounded indicator that guides the choice of initialization variance <ref:2508.12834#pg6>. It’s less about magic and more about understanding the dynamics of the parameters in SGD <ref:2508.12834#pg1>.

Tom: Exactly. And they also noted a limitation, which is that heavy-tailed initializations don't deliver uniformly high accuracy across trials because their divergent second moment makes performance less stable than in the Gaussian case <ref:2508.12834#pg6>.

Jane: So, the paper gives us a tool to select initialization variance based on theory, and it shows that starting in the right spot can lead to better training loss and test accuracy when using SGD <ref:2508.12834#pg4>. It also points toward future work needing things like Wasserstein distance instead of KL divergence for more robust analysis <ref:2508.12834#pg6>.

Conclusion: Tom: So, to wrap up, this paper looks at how we set up the initial weights for neural networks using standard Stochastic Gradient Descent, looking at it through this Fokker-Planck equation stuff<ref:2508.12834#pg1>.

Jane: It’s about finding a mathematically precise way to pick the starting scale for those weights so the learning process actually works better than just guessing a random value<ref:2508.12834#pg6>.

Lu: The core idea is using Kullback-Leibler divergence between where things start and where they end to find an optimal starting variance, sigma two zero <ref:2508.12834#pg3>.

Meng: So, what does this actually mean for deployment? Does it give us a simple rule like "use He initialization"?

Tom: Not exactly. It gives us a condition based on the expected final variance of the weights, (W), which we can use to guide our initial setup<ref:2508.12834#pg7>.

Jane: So, if you know how big the final weight distribution will be after training, you can set your starting distribution to match it exactly<ref:2508.12834#pg6>. It’s about matching the exploration region to where the minimum loss solution is actually hiding<ref:2508.12834#pg9>.

Lu: That means we’re not just throwing weights randomly into a deep space; we're starting in a spot that already points toward the valley floor<ref:2508.12834#pg9>. This should make convergence much more reliable across different network architectures<ref:2508.12834#pg6>.

Meng: If it works better, then scaling up those models to production would be less risky because the initial training loss is lower and the final test accuracy stays higher<ref:2508.12834#pg4>.

Tom: It’s about making sure the dynamics of parameter updates are actually productive from step one, not just drifting around randomly<ref:2508.12834#pg9>. This moves us away from just heuristic guesses for initialization<ref:2508.12834#pg6>.

Jane: It’s a mathematical way to tell the AI how to start its journey through the loss landscape, so we know exactly where it should begin<ref:2508.12834#pg9>. This is a concrete indicator for choosing that initial variance<ref:2508.12834#pg9>.

Lu: And while this Gaussian case is great, the authors also noted that if we use heavy-tailed initializations, the performance can be less stable because their second moment diverges<ref:2508.12834#pg6>. We need to keep those tails controlled<ref:2508.12834#pg6>.

Meng: So, for production, we stick to Gaussian scaling because it’s more predictable than those heavy-tailed setups that might give us a good result sometimes but fail other times<ref:2508.12834#pg9>.

Tom: Exactly. It’s about moving from a "try and see" approach for initialization to something that actually has been derived mathematically<ref:2508.12834#pg9>. Next time, we'll talk about how this relates to other optimization methods like the Wasserstein distance<ref:2508.12834#pg9>.

More episodes

← Home