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

Hyungu Lee; Taehyeong Kim; Hayoung Choi

cs.LGAug 30, 2025

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

Authors: Hyungu Lee, Taehyeong Kim, Hayoung Choi

Organizations: Department of Mathematics & Nonlinear Dynamics and Mathematical Application Center, Kyungpook National University, Daegu 41566, Republic of Korea · Nonlinear Dynamics and Mathematical Application Center, Kyungpook National University, Daegu 41566, Republic of Korea

Abstract

Stable and efficient training of ReLU networks with large depth is highly sensitive to weight initialization. Improper initialization can cause permanent neuron inactivation dying ReLU and exacerbate gradient instability as network depth increases. Methods such as He, Xavier, and orthogonal initialization preserve variance or promote approximate isometry. However, they do not necessarily regulate the pre-activation mean or control activation sparsity, and their effectiveness often diminishes in very deep architectures. This work introduces a semi-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. A family of closed-form solutions and an efficient sampling scheme are derived. Theoretical analysis at initialization shows that the proposed initialization mitigates dying ReLU, slows the decay of activation variance, and alleviates gradient vanishing, thereby stabilizing signal and gradient flow in deep architectures. Empirically, across MNIST, Fashion-MNIST, multiple tabular datasets, few-shot settings, and ReLU-family activations, the proposed initialization maintains stable performance as network depth increases, with particularly clear advantages in few-shot learning and large-depth settings.

Figures & tables

Panel of Figure 1 — Figure 1: A two-dimensional UMAP projection of vectors $\bm{v}_{k}=\bm{W}_{(k)}^{T}\bm{1}$ .

Figure 2: Convergence of $\bm{W}\bm{x}$ to normal distribution as dimension $n$ increases.

Dataset	Depth	Proposed	Lee	ZerO	IDInit	He	Xavier	Orth.	Rand.
MNIST	10	96.96 $\pm$ 0.14	95.84 $\pm$ 0.26	93.25 $\pm$ 0.65	93.50 $\pm$ 0.60	96.42 $\pm$ 0.19	96.27 $\pm$ 0.15	96.46 $\pm$ 0.22	95.64 $\pm$ 0.22
	20	96.86 $\pm$ 0.25	95.79 $\pm$ 0.21	91.96 $\pm$ 1.08	91.85 $\pm$ 1.24	95.65 $\pm$ 0.22	94.19 $\pm$ 0.80	94.45 $\pm$ 0.59	93.04 $\pm$ 0.47
	30	96.78 $\pm$ 0.20	95.97 $\pm$ 0.22	90.70 $\pm$ 2.05	91.09 $\pm$ 1.25	95.12 $\pm$ 0.24	93.41 $\pm$ 0.35	92.64 $\pm$ 0.61	92.55 $\pm$ 0.63
	40	96.69 $\pm$ 0.34	96.16 $\pm$ 0.23	90.86 $\pm$ 1.28	90.83 $\pm$ 1.20	94.42 $\pm$ 0.36	92.82 $\pm$ 0.46	92.83 $\pm$ 0.40	91.00 $\pm$ 0.54
	50	96.76 $\pm$ 0.16	96.44 $\pm$ 0.23	90.69 $\pm$ 1.00	90.84 $\pm$ 1.06	93.94 $\pm$ 0.35	92.02 $\pm$ 0.52	88.07 $\pm$ 18.07	90.57 $\pm$ 0.78
	60	96.59 $\pm$ 0.25	96.59 $\pm$ 0.18	91.36 $\pm$ 1.05	91.11 $\pm$ 0.73	93.62 $\pm$ 0.45	91.86 $\pm$ 0.45	92.00 $\pm$ 0.57	89.86 $\pm$ 0.88

Table 1: Test accuracy (%) by depth and initialization method. Values are reported as mean

\pm

standard deviation over twenty random seeds.

Dataset	Depth	Proposed	Mean-only	Orthogonal
MNIST	50	96.76 $\pm$ 0.16	95.79 $\pm$ 0.18	88.07 $\pm$ 18.07
MNIST	100	96.57 $\pm$ 0.29	95.21 $\pm$ 0.40	63.15 $\pm$ 39.01
Fashion-MNIST	50	87.60 $\pm$ 0.35	86.62 $\pm$ 0.22	21.07 $\pm$ 27.03
Fashion-MNIST	100	87.52 $\pm$ 0.27	86.06 $\pm$ 0.29	10.00 $\pm$ 0.00

Table 2: Test accuracy (%) for the ablation of mean calibration and semi-orthogonality at depths 50 and 100. Values are reported as mean

\pm

standard deviation over twenty random seeds.

Dataset	Activation	Depth	Proposed	Lee	ZerO	IDInit	He	Xavier	Orth.	Rand.
MNIST	LeakyReLU	50	97.30 $\pm$ 0.29	97.70 $\pm$ 0.15	97.26 $\pm$ 0.20	97.24 $\pm$ 0.23	96.73 $\pm$ 0.26	76.16 $\pm$ 33.57	72.53 $\pm$ 32.62	11.35 $\pm$ 0.00
	LeakyReLU	100	94.56 $\pm$ 5.51	97.66 $\pm$ 0.18	97.19 $\pm$ 0.19	97.17 $\pm$ 0.20	91.16 $\pm$ 5.58	11.35 $\pm$ 0.00	11.35 $\pm$ 0.00	11.35 $\pm$ 0.00
	PReLU	50	97.36 $\pm$ 0.20	97.75 $\pm$ 0.22	97.34 $\pm$ 0.22	97.31 $\pm$ 0.21	96.93 $\pm$ 0.26	95.28 $\pm$ 1.80	95.68 $\pm$ 3.92	35.16 $\pm$ 37.37
	PReLU	100	96.53 $\pm$ 1.42	97.73 $\pm$ 0.18	97.22 $\pm$ 0.21	97.24 $\pm$ 0.15	91.17 $\pm$ 5.98	11.35 $\pm$ 0.00	11.35 $\pm$ 0.00	11.35 $\pm$ 0.00
	ELU	50	97.63 $\pm$ 0.18	97.38 $\pm$ 0.17	97.37 $\pm$ 0.19	97.47 $\pm$ 0.14	97.18 $\pm$ 0.24	97.46 $\pm$ 0.18	97.73 $\pm$ 0.21	95.97 $\pm$ 0.84
	ELU	100	97.56 $\pm$ 0.21	97.38 $\pm$ 0.15	97.41 $\pm$ 0.29	97.43 $\pm$ 0.15	91.58 $\pm$ 4.24	65.86 $\pm$ 30.12	97.10 $\pm$ 1.64	11.35 $\pm$ 0.00

Table 3: Accuracy (%) of different initialization methods and various activation functions on MNIST and Fashion-MNIST at depths

50

and

100

. Values are reported as mean

\pm

standard deviation over twenty random seeds.

Dataset ( $k$ )	Depth	Proposed	Lee	ZerO	IDInit	He	Xavier	Orthogonal	Random
MNIST (1)	10	43.33 $\pm$ 2.89	34.54 $\pm$ 4.53	9.94 $\pm$ 1.09	9.94 $\pm$ 1.10	27.42 $\pm$ 2.93	28.09 $\pm$ 4.50	32.54 $\pm$ 3.58	23.79 $\pm$ 4.87
	50	39.95 $\pm$ 4.70	29.24 $\pm$ 4.64	9.94 $\pm$ 1.08	9.93 $\pm$ 1.01	13.80 $\pm$ 2.62	10.78 $\pm$ 1.41	10.00 $\pm$ 1.45	10.08 $\pm$ 0.63
	100	37.33 $\pm$ 4.32	27.43 $\pm$ 4.86	9.93 $\pm$ 1.03	9.94 $\pm$ 1.09	12.85 $\pm$ 2.92	9.93 $\pm$ 0.44	9.92 $\pm$ 0.40	9.98 $\pm$ 0.69
MNIST (2)	10	52.18 $\pm$ 4.02	41.35 $\pm$ 4.49	10.51 $\pm$ 1.81	10.77 $\pm$ 2.39	34.35 $\pm$ 5.01	35.85 $\pm$ 5.66	42.22 $\pm$ 4.20	28.67 $\pm$ 5.05
	50	47.83 $\pm$ 3.61	34.13 $\pm$ 4.07	10.57 $\pm$ 1.92	10.57 $\pm$ 1.93	15.52 $\pm$ 2.90	10.64 $\pm$ 1.52	9.91 $\pm$ 0.53	9.93 $\pm$ 0.41
	100	45.75 $\pm$ 4.27	34.17 $\pm$ 4.18	10.30 $\pm$ 1.70	10.31 $\pm$ 1.71	13.57 $\pm$ 2.47	9.86 $\pm$ 0.67	10.17 $\pm$ 0.59	10.10 $\pm$ 0.81

Table 4: Few-shot classification accuracy (%) for various initialization methods and network depths. The number following each dataset name indicates the number of shots (

k

). Values are reported as mean

\pm

standard deviation. The highest mean accuracy in each row is shown in bold.

Panel of Figure 3 — Figure 3: 4-shot classification accuracy for various initialization methods and network depths.

Dataset	Samples	Features $N_{0}$	Output $N_{D}$	Task	Architecture
Adult	48,842	14	2	Classification	$\,14\rightarrow 8\rightarrow\cdots\rightarrow 8\rightarrow 2\,$
Cancer	569	30	2	Classification	$\,30\rightarrow 16\rightarrow\cdots\rightarrow 16\rightarrow 2\,$
Pima	768	8	2	Classification	$\,8\rightarrow 4\rightarrow\cdots\rightarrow 4\rightarrow 2\,$
Ionosphere	351	34	2	Classification	$\,34\rightarrow 16\rightarrow\cdots\rightarrow 16\rightarrow 2\,$
Wine	178	13	3	Classification	$\,13\rightarrow 8\rightarrow\cdots\rightarrow 8\rightarrow 3\,$
Diabetes	442	10	1	Regression	$\,10\rightarrow 8\rightarrow\cdots\rightarrow 8\rightarrow 1\,$

Table 5: Description of tabular datasets and architectures.

Dataset	$\alpha_{0}$	Depth	Proposed	Lee	ZerO	IDInit	He	Xavier	Orthogonal	Random
Adult	-2	50	76.07 $\pm$ 0.00	85.14 $\pm$ 0.53	85.74 $\pm$ 0.43	85.86 $\pm$ 0.44	79.93 $\pm$ 5.29	79.85 $\pm$ 5.17	83.56 $\pm$ 4.21	76.07 $\pm$ 0.00
		100	76.07 $\pm$ 0.00	80.06 $\pm$ 4.78	85.97 $\pm$ 0.46	85.48 $\pm$ 0.88	79.88 $\pm$ 5.22	76.07 $\pm$ 0.00	76.07 $\pm$ 0.00	76.07 $\pm$ 0.00
	0	50	85.40 $\pm$ 0.47	85.90 $\pm$ 0.29	86.19 $\pm$ 0.22	86.10 $\pm$ 0.24	80.00 $\pm$ 5.38	79.89 $\pm$ 5.23	83.67 $\pm$ 4.27	76.07 $\pm$ 0.00
		100	85.41 $\pm$ 0.31	85.85 $\pm$ 0.73	86.03 $\pm$ 0.54	86.02 $\pm$ 0.45	79.78 $\pm$ 5.10	76.07 $\pm$ 0.00	77.98 $\pm$ 4.27	76.07 $\pm$ 0.00
	2	50	85.31 $\pm$ 0.29	85.54 $\pm$ 0.50	85.79 $\pm$ 0.47	85.75 $\pm$ 0.43	81.85 $\pm$ 5.28	76.07 $\pm$ 0.00	83.58 $\pm$ 4.20	76.07 $\pm$ 0.00
		100	85.57 $\pm$ 0.36	85.79 $\pm$ 0.53	85.81 $\pm$ 0.39	85.86 $\pm$ 0.47	81.49 $\pm$ 4.95	76.07 $\pm$ 0.00	77.84 $\pm$ 3.96	76.07 $\pm$ 0.00

Table 6: Classification accuracy (%) for the classification datasets and RMSE for Diabetes regression. Values are reported as mean

\pm

standard deviation. Boldface denotes the best performance for each

(\text{dataset},\alpha_{0},\text{depth})

configuration, while underlined boldface additionally indicates the best

\alpha_{0}

for the corresponding initialization method at the same dataset and depth.

Neural Network Initialization

Neural network initialization studies how starting weights shape signal propagation, optimization, generalization, and reproducibility.

See all 51 papers

Path-conditioned training: a principled way to rescale ReLU neural networks
Feb 23, 2026Arthur Lebeurrier, Titouan Vayer, Rémi GribonvalReLU Neural Networks Neural Network Training Dynamics
ENS de Lyon, CNRS, Inria, Université Claude Bernard Lyon 1, LIP, UMR 5668, 69342, Lyon cedex 07, France · Inria, Rennes, France · Inria, CNRS, ENS de Lyon, Université Claude Bernard Lyon 1, LIP, UMR 5668, 69342, Lyon cedex 07, France
Optimal Condition for Initialization Variance in Deep Neural Networks: An SGD Dynamics Perspective
Aug 18, 2025Hiroshi Horii, Sothea HasFokker–Planck Equation Neural Network Initialization
SU · KHM
Favorability of Loss Landscape with Weight Decay Requires Both Large Overparametrization and Initialization
May 28, 2025Etienne Boursier, Matthew Bowditch, Matthias Englert +1Loss Landscape Overparameterized Neural Network
INRIA LMO, Université Paris-Saclay Orsay, France · Mathematics Institute University of Warwick Coventry, UK · Department of Computer Science University of Warwick Coventry, UK

Feb 23, 2026J/K move · Enter open · S save

Path-conditioned training: a principled way to rescale ReLU neural networks

Arthur Lebeurrier, Titouan Vayer, Rémi Gribonval

ENS de Lyon, CNRS, Inria, Université Claude Bernard Lyon 1, LIP, UMR 5668, 69342, Lyon cedex 07, France · Inria, Rennes, France · Inria, CNRS, ENS de Lyon, Université Claude Bernard Lyon 1, LIP, UMR 5668, 69342, Lyon cedex 07, France

ReLU Neural Networks Neural Network Training Dynamics Neural Network Initialization

Open paper AskPDF arXiv

Despite recent algorithmic advances, we still lack principled ways to leverage the well-documented rescaling symmetries in ReLU neural network parameters. While two properly rescaled weights implement the same function, the training dynamics can be dramatically different. To offer a fresh perspective on exploiting this phenomenon, we build on the recent path-lifting framework, which provides a compact factorization of ReLU networks. We introduce a geometrically motivated criterion to rescale neural network parameters which minimization leads to a conditioning strategy that aligns a kernel in the path-lifting space with a chosen reference. We derive an efficient algorithm to perform this alignment. In the context of random network initialization, we analyze how the architecture and the initialization scale jointly impact the output of the proposed method. Numerical experiments illustrate its potential to speed up training.

Abstract

Neural Network Initialization

Explore similar work

Path-conditioned training: a principled way to rescale ReLU neural networks

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

Favorability of Loss Landscape with Weight Decay Requires Both Large Overparametrization and Initialization