Optimized Weight Initialization on the Stiefel Manifold for Deep ReLU Neural Networks

arXiv:2509.00362 · cs.LG · Submitted 2025-08-30 · Read on arXiv

Listen

Transcript

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

Tom: We're looking at the paper: "Optimized Weight Initialization on the Stiefel Manifold for Deep ReLU Neural Networks".

Jane: Here's the gist of the paper. The gist Stable and efficient training of ReLU networks with large depth is highly sensitive to weight initialization, and this work introduces an orthogonal initialization specifically optimized for ReLU by solving an optimization problem on the Stiefel manifold, thereby preserving scale and calibrating the pre-activation statistics from the outset<ref:2509.00362#pg2>.

Tom: Deep neural networks are highly sensitive to weight initialization because improper initialization can cause permanent neuron inactivation (“dying ReLU”) and exacerbate gradient instability as network depth increases.

Jane: The Rectified Linear Unit (ReLU) activation function suffers from the dying ReLU problem, when a pre-activation remains negative, its gradient vanishes and the corresponding parameters stop updating<ref:2509.00362#pg2>.

Tom: Early methods like Xavier initialization often underperform with ReLU-based models due to the asymmetry of the function<ref:2509.00362#pg2>.

Paper summary: Tom: So wrapping up the "Optimized Weight Initialization on the Stiefel Manifold for Deep ReLU Neural Networks" paper. The authors are Hyungu Lee, Taehyeong Kim, and Hayoung Choi.

Jane: It’s about using geometry—specifically the Stiefel manifold—to design initial weights so they inherently preserve scale and calibrate those pre-activation statistics from the very beginning one <ref:2509.00362#pg1,pre-activation statistics from the>.

Lu: The core idea is solving an optimization problem on that manifold to find a semi-orthogonal matrix that aligns with the all-ones vector, which prevents neuron inactivation and stabilizes gradients in deep ReLU networks two <ref:2509.00362#pg2>.

Meng: From an engineering standpoint, it means we have a method that addresses the vanishing gradient and dying ReLU issues by controlling the activation distribution dynamics across layers two <ref:2509.00362#pg2>.

Tom: And empirically, they showed this works well on MNIST up to a depth of one hundred layers, achieving ninety-six point nine eight percent, which is way better than what conventional initializations do beyond fifty layers two.

Jane: It suggests that for anyone building a very deep network using ReLU, the initialization choice isn't just an afterthought; it’s a critical structural component of the training process one <ref:2509.00362#pg1>.

Lu: The implication is that this geometric approach to initialization provides a robust framework for extending depth scalability and improving generalization performance across different data types two <ref:2509.00362#pg2>.

Conclusion: Tom: So we're wrapping up on this paper, "Optimized Weight Initialization on the Stiefel Manifold for Deep ReLU Neural Networks."

Jane: It basically says they found a specific geometric way to set those starting weights so deep networks with ReLU don't just break.

Tom: Right, it’s about using the Stiefel manifold to solve an optimization problem to find weights that preserve scale and pre-activation statistics from the very start.

Lu: The big idea is finding a semi-orthogonal matrix that aligns with the all-ones vector, which directly helps prevent those neurons from dying in ReLU networks.

Meng: So what does this actually mean for me when I'm trying to deploy something? Does it make my training run faster or more stable on real traffic?

Lalam: The paper shows that this initialization mitigates the dying ReLU phenomenon and slows down the decay of activation variance, which is key for keeping signals flowing through deep layers.

Jane: It means instead of guessing weights randomly or using older methods like Xavier, we have a structured way to set them up that keeps the network healthy as it gets deeper.

Tom: The empirical results are pretty strong; they show stable performance on MNIST and Fashion-MNIST up to a depth of one hundred layers, which is a huge jump compared to previous methods that failed past fifty.

Lu: They also showed how the rectification parameter alpha increases with depth, pushing activations into the positive domain, which helps keep things balanced.

Meng: If this initialization stabilizes training across such deep architectures without needing constant manual tuning of learning rates or regularization, that’s a big win for production deployment.

Jane: It suggests that geometric information about the space the weights live in can actually dictate how well an AI model learns and performs on complex tasks.

Tom: Yeah, so this moves initialization from being just a heuristic to something more principled based on the geometry of the problem itself.

Lalam: It’s a solid mathematical framework that translates directly into better signal and gradient flow in deep learning systems.

Hyungu Lee, Taehyeong Kim, Hayoung Choi

Kyungpook National University

cs.LG

Submitted: 2025-08-30

Updated: 2026-10-20

Comments: 18 pages, 3 figures, 6 tables

License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/

Importance score: 80/100

The gist: The gist Stable and efficient training of ReLU networks with large depth is highly sensitive to weight initialization, and this work introduces an orthogonal initialization specifically optimized for

Key concepts

Stiefel Manifold: This is a mathematical space where matrices with orthonormal columns live. The paper uses it to find a specific type of weight matrix that is 'semi-orthogonal,' meaning its columns are nearly orthogonal while preserving the direction of the all-ones vector. This geometric constraint helps control how the network's weights scale and interact with activations.
Dying ReLU Problem: This occurs when a neuron's input consistently results in a negative value under the ReLU function. Because its gradient is zero for negative inputs, the neuron stops learning and its associated parameters never update. The proposed initialization mitigates this by calibrating pre-activation statistics to keep more neurons in the active (positive) regime.
Semi-orthogonal Matrix: A matrix is semi-orthogonal if it satisfies a specific alignment condition related to the all-ones vector. By optimizing for this property on the Stiefel manifold, the initialization ensures that as you pass through layers, the activation distribution remains well-behaved. This prevents variance from decaying too quickly and stabilizes gradient flow.
Pre-activation Statistics Calibration: This refers to controlling the mean and covariance of the inputs to each neuron before they are passed through the ReLU function. The initialization method explicitly solves for a matrix that calibrates these statistics. This calibration is crucial because it ensures that as network depth increases, the signal and gradient flow remain stable, rather than vanishing or exploding.

Terminology

Stiefel manifold — 2× in this paper · explained in 3 episodes
Activation Sparsity — 1× in this paper · explained in 2 episodes
Gradient Flow — 1× in this paper · explained in 2 episodes
ReLU Network — 1× in this paper · explained in 2 episodes
Sparsity — 1× in this paper · explained in 5 episodes

Summary

The gist Stable and efficient training of ReLU networks with large depth is highly sensitive to weight initialization, and this work introduces an orthogonal initialization specifically optimized for ReLU by solving an optimization problem on the Stiefel manifold, thereby preserving scale and calibrating the pre-activation statistics from the outset<ref:2509.00362#pg2>.

Motivation and Problem Statement

Deep neural networks are highly sensitive to weight initialization because improper initialization can cause permanent neuron inactivation (“dying ReLU”) and exacerbate gradient instability as network depth increases. The Rectified Linear Unit (ReLU) activation function suffers from the dying ReLU problem, when a pre-activation remains negative, its gradient vanishes and the corresponding parameters stop updating<ref:2509.00362#pg2>. Early methods like Xavier initialization often underperform with ReLU-based models due to the asymmetry of the function<ref:2509.00362#pg2>.

Proposed Methodology

The paper introduces a novel weight initialization method derived from a principled optimization framework, formulating the design of the weight matrix as an optimization problem on the Stiefel manifold, which seeks a semi-orthogonal matrix that is maximally aligned with the all-ones vector<ref:2509.00362#pg2>. This approach is motivated by empirical and theoretical findings that such alignment is critical for preventing neuron inactivation in ReLU networks. The optimization problem is posed as max W∈Om,n tr (Jm×nW), which can be equivalently expressed as min W∈Om,n∥Jm×n − W∥F. The exact solution to this maximization problem is given by Wf = UV T, where U ∈ Om,m has its first column equal to ξm, and V ∈ On,m has its first column equal to ξn.

Theoretical Guarantees

The theoretical analysis reveals several key structural properties of the matrices in Oem,n. Lemma 1 characterizes Oem,n as the set of all semiorthogonal matrices that preserve the normalized all-ones vector. Furthermore, Proposition 3 shows that for a random vector x ∈ R n with E[x] = µ1n and Cov[x] = σ2In, the output W x under a semi-orthogonal transformation W ∈ Oem,n also has an explicit mean and covariance: E[W x] = µr/n m 1m, Cov[W x] = σ2Im. Theorem 4 guarantees that for a random vector x with independent components satisfying E[xj] = µ, Var(xj) = σ2, E[xj 3] < ∞ for all j = 1,..., n, and W ∈ Oem,n being a random matrix as defined in Theorem 3, then W x − µp n m 1m σd−→ N (0, Im) as n → ∞.

Empirical Validation and Performance

Empirically, the proposed initialization is evaluated on several benchmark tasks and consistently demonstrates superior performance compared to existing initialization methods. Across MNIST, Fashion-MNIST, multiple tabular datasets, few-shot settings, and ReLU-family activations, our method outperforms previous initializations and enables stable training in deep networks. Specifically, the proposed initialization maintains stable and high accuracy even at depths up to D = 100 for MNIST (96.98%) and Fashion-MNIST (87.70%), whereas conventional initializations exhibit severe performance degradation beyond D = 50<ref:2509.00362#pg2>.

Mitigation of Deep Network Issues

The proposed initialization systematically alleviates critical early-stage training issues—most notably the dying ReLU phenomenon, variance decay, and gradient vanishing—by promoting stable signal and gradient propagation across layers. Under the proposed initialization, Proposition 4 yields x(l) ∼ N R(µl1Nlµ, σ2lINl), with nondecreasing µl and nonincreasing σl, so the rectification parameter αl = µl/σl increases with depth. This increase in α pushes the activation distribution further into the positive domain, mitigating the dying ReLU phenomenon. Furthermore, for any layer l−1, Lemma 5 gives σ2l/σ2l−1=(1 + α2l−1) Φ(αl−1) + αl−1 ϕ(αl−1) − ϕ(αl-1) + αl−1 Φ(αell-2)2. Since Φ(a) → 1 and ϕ(a) → 0 as a → ∞, it follows that σ2l/σ2l−1→ 1 as αl−1 → ∞, preventing cumulative variance decay and preserving activation dynamic range.

Conclusion

Taken together, these results highlight the importance of geometrically informed initialization in the design of deep networks. By exploiting the structure of the Stiefel manifold, our work demonstrates that initialization can play a decisive role in stabilizing training, extending depth scalability, and improving generalization performance. The method consistently outperformed existing initialization schemes across depths of up to 100 layers and showed robustness across the ReLU family of activation functions, thereby establishing its broad applicability.

Improvements for AI systems

Weight initialization via Stiefel manifold optimization will stabilize training dynamics by solving an optimization problem on Om,n to derive a semi-orthogonal matrix that is maximally aligned with the all-ones vector. This directly addresses preventing neuron inactivation in ReLU networks and stabilize signal and gradient flow in deep architectures.
The proposed initialization scheme will prevent the dying ReLU phenomenon by ensuring that for each layer l, the rectification parameter alpha l grows with depth, leading to P(x(l)j > 0) approaching 1. This is achieved because as α becomes large, most of the probability mass of the input lies in the positive region.
The system will maintain stable and non-decreasing variance across layers by ensuring that for any epsilon in (0, 1), there exists delta such that if alpha l-1 >= delta then (1 - epsilon) sigma2 l-1 <= sigma2 l. This prevents cumulative variance decay and preserving activation dynamic range.
The AI system will exhibit superior performance on deep networks up to D = 100 layers across MNIST, Fashion-MNIST, and tabular datasets compared to standard initializations like Xavier or He. This is because the method consistently maintains stable and high accuracy even at depths up to D = 100.
Few-shot learning performance will be significantly enhanced because the proposed initialization exhibits a remarkably stable and slow decrease in accuracy as the network gets deeper, contrasting sharply with other methods that exhibit a sharp decline in performance, with accuracy collapsing after 30-40 layers.