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

TL;DR

This paper proposes an orthogonal weight initialization method on the Stiefel manifold tailored for deep ReLU networks, improving stability and performance by better controlling activation and gradient flow.

Contribution

It introduces a novel optimization-based orthogonal initialization on the Stiefel manifold, specifically designed for ReLU networks, with theoretical analysis and practical benefits.

Findings

Outperforms previous initializations on multiple datasets

Enables stable training of very deep ReLU networks

Reduces dying ReLU and gradient vanishing issues

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 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. A family of closed-form solutions and an efficient sampling scheme are derived. Theoretical analysis at initialization shows that prevention of the dying ReLU problem, slower decay of activation variance, and mitigation of gradient vanishing, which together stabilize signal and gradient flow in deep architectures. Empirically, 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.

Figures

Tables

Table 1. (a) MNIST

Depth	Proposed	Lee	He	Xavier	Orth.	Rand.
10	97.59	97.87	97.55	97.57	97.56	97.30
20	97.57	97.53	97.43	97.23	97.00	96.31
30	97.17	96.14	97.21	97.10	97.14	96.44
40	97.19	95.83	97.06	96.04	97.07	95.13
50	97.28	96.02	96.52	82.74	11.35	11.35
60	97.17	94.74	96.44	82.07	11.35	11.35
70	97.25	96.00	93.70	11.35	11.35	11.35
80	96.64	84.82	95.95	11.35	11.35	11.35
90	97.04	88.46	95.91	11.35	11.35	11.35
100	96.98	95.17	93.69	11.35	11.35	11.35

Table 2. (a) MNIST

Depth	Proposed	Lee	He	Xavier	Orth.	Rand.
10	97.59	97.87	97.55	97.57	97.56	97.30
20	97.57	97.53	97.43	97.23	97.00	96.31
30	97.17	96.14	97.21	97.10	97.14	96.44
40	97.19	95.83	97.06	96.04	97.07	95.13
50	97.28	96.02	96.52	82.74	11.35	11.35
60	97.17	94.74	96.44	82.07	11.35	11.35
70	97.25	96.00	93.70	11.35	11.35	11.35
80	96.64	84.82	95.95	11.35	11.35	11.35
90	97.04	88.46	95.91	11.35	11.35	11.35
100	96.98	95.17	93.69	11.35	11.35	11.35

Table 3. (b) Fashion-MNIST

Depth	Proposed	Lee	He	Xavier	Orth.	Rand.
10	88.15	87.66	88.17	87.52	87.96	87.58
20	88.29	87.65	88.17	87.44	87.67	87.04
30	87.92	88.65	87.80	87.00	87.62	86.93
40	88.03	87.73	87.53	87.64	87.45	10.00
50	87.77	88.68	87.23	10.00	10.00	10.00
60	87.98	85.94	87.59	10.00	10.00	10.00
70	87.83	86.59	86.97	10.00	10.00	10.00
80	87.78	77.89	84.92	10.00	10.00	10.00
90	87.88	81.78	85.50	10.00	10.00	10.00
100	87.70	79.09	83.33	10.00	10.00	10.00

Table 4. TABLE II: Accuracy (%) of different initialization methods and various activation functions on MNIST and Fashion-MNIST at depths 50 50 and 100 100 .

Dataset	Activation	Proposed		Lee		He		Xavier		Orthogonal		Random
		50	100	50	100	50	100	50	100	50	100	50	100
MNIST	ReLU	97.12	96.33	97.25	84.15	96.39	95.24	11.35	11.35	93.94	11.35	11.35	11.35
	LeakyReLU	97.08	96.29	97.16	11.35	96.62	95.39	86.52	11.35	87.36	11.35	11.35	11.35
	PReLU	96.93	96.52	97.46	11.35	95.99	94.60	95.02	11.35	11.35	11.35	11.35	11.35
	ELU	97.67	97.59	97.79	74.25	96.62	96.34	97.17	11.35	97.73	88.75	86.56	11.35
	SELU	97.39	96.45	97.62	96.40	96.57	95.33	96.59	96.61	97.22	97.07	96.52	96.31
Fashion-MNIST	ReLU	88.15	87.81	87.82	86.54	87.63	86.97	10.00	10.00	85.30	10.00	10.00	10.00
	LeakyReLU	88.36	87.89	86.87	74.37	87.10	86.16	86.22	10.00	85.81	10.00	10.00	10.00
	PReLU	87.75	87.77	88.28	75.94	86.81	86.91	79.44	10.00	86.80	10.00	10.00	10.00
	ELU	88.05	88.29	88.43	78.15	87.29	85.30	87.90	10.00	88.22	70.18	83.02	10.00
	SELU	87.60	87.25	84.55	71.13	87.01	85.45	87.64	85.97	87.92	87.50	87.22	83.85

Table 5. TABLE III: Few-shot classification accuracy (%) for various initialization methods and network depths. The number following each dataset name indicates the number of shots ( k k ). Values are reported as mean ± \pm standard deviation. The highest mean accuracy in each row is shown in bold.

Dataset ($k$)	Depth	Proposed	Lee	He	Xavier	Orthogonal	Random
MNIST (1)	10	42.48 $\pm$ 3.93	37.75 $\pm$ 4.56	27.42 $\pm$ 2.93	28.09 $\pm$ 4.50	32.54 $\pm$ 3.58	23.79 $\pm$ 4.87
	50	39.57 $\pm$ 3.80	36.74 $\pm$ 4.47	13.80 $\pm$ 2.62	10.78 $\pm$ 1.41	10.00 $\pm$ 1.45	10.08 $\pm$ 0.63
	100	37.23 $\pm$ 4.92	28.23 $\pm$ 3.97	12.85 $\pm$ 2.92	9.93 $\pm$ 0.44	9.92 $\pm$ 0.40	9.98 $\pm$ 0.69
MNIST (2)	10	52.86 $\pm$ 3.89	44.30 $\pm$ 4.15	34.35 $\pm$ 5.01	35.85 $\pm$ 5.66	42.22 $\pm$ 4.20	28.67 $\pm$ 5.05
	50	48.64 $\pm$ 3.89	43.85 $\pm$ 4.35	15.52 $\pm$ 2.90	10.64 $\pm$ 1.52	9.91 $\pm$ 0.53	9.93 $\pm$ 0.41
	100	45.72 $\pm$ 4.20	34.05 $\pm$ 4.20	13.57 $\pm$ 2.47	9.86 $\pm$ 0.67	10.17 $\pm$ 0.59	10.10 $\pm$ 0.81
MNIST (4)	10	63.79 $\pm$ 2.73	53.00 $\pm$ 3.99	43.67 $\pm$ 4.90	44.73 $\pm$ 4.84	52.28 $\pm$ 3.52	35.73 $\pm$ 5.96
	50	58.86 $\pm$ 3.73	54.73 $\pm$ 3.38	16.41 $\pm$ 2.95	11.62 $\pm$ 2.11	9.96 $\pm$ 0.47	9.98 $\pm$ 0.65
	100	56.70 $\pm$ 3.11	44.81 $\pm$ 3.85	13.90 $\pm$ 2.17	10.08 $\pm$ 0.45	9.76 $\pm$ 0.49	10.15 $\pm$ 0.70
MNIST (8)	10	71.73 $\pm$ 2.21	62.36 $\pm$ 2.92	53.92 $\pm$ 4.76	52.90 $\pm$ 3.76	61.97 $\pm$ 3.80	44.67 $\pm$ 6.13
	50	67.67 $\pm$ 2.58	63.33 $\pm$ 2.77	17.96 $\pm$ 3.45	10.94 $\pm$ 2.46	9.99 $\pm$ 0.67	9.84 $\pm$ 0.50
	100	66.35 $\pm$ 2.73	52.45 $\pm$ 4.01	14.45 $\pm$ 2.44	9.89 $\pm$ 0.60	9.93 $\pm$ 0.44	10.05 $\pm$ 0.53
Fashion-MNIST (1)	10	46.78 $\pm$ 4.27	43.48 $\pm$ 4.94	37.62 $\pm$ 5.08	36.80 $\pm$ 5.11	41.50 $\pm$ 6.28	31.55 $\pm$ 6.74
	50	46.19 $\pm$ 4.08	37.19 $\pm$ 5.42	20.29 $\pm$ 3.97	11.81 $\pm$ 3.10	10.14 $\pm$ 1.67	10.00 $\pm$ 0.00
	100	45.07 $\pm$ 4.28	28.88 $\pm$ 8.82	14.81 $\pm$ 4.18	10.00 $\pm$ 0.00	10.00 $\pm$ 0.00	10.00 $\pm$ 0.00
Fashion-MNIST (2)	10	55.99 $\pm$ 4.15	52.44 $\pm$ 4.05	46.37 $\pm$ 4.62	45.29 $\pm$ 5.25	49.68 $\pm$ 5.68	42.05 $\pm$ 5.27
	50	54.94 $\pm$ 4.15	47.05 $\pm$ 5.48	23.61 $\pm$ 5.48	12.36 $\pm$ 2.96	10.24 $\pm$ 0.86	10.00 $\pm$ 0.00
	100	55.21 $\pm$ 4.19	40.40 $\pm$ 8.90	16.20 $\pm$ 3.62	10.00 $\pm$ 0.00	10.00 $\pm$ 0.00	10.00 $\pm$ 0.00
Fashion-MNIST (4)	10	62.32 $\pm$ 3.09	58.53 $\pm$ 3.17	53.33 $\pm$ 4.18	52.62 $\pm$ 4.60	56.22 $\pm$ 3.34	49.17 $\pm$ 3.58
	50	60.85 $\pm$ 2.97	57.05 $\pm$ 3.68	25.12 $\pm$ 4.78	13.12 $\pm$ 3.35	10.00 $\pm$ 0.00	10.00 $\pm$ 0.00
	100	60.40 $\pm$ 3.11	51.14 $\pm$ 8.45	17.15 $\pm$ 4.17	10.00 $\pm$ 0.00	10.00 $\pm$ 0.00	10.00 $\pm$ 0.00
Fashion-MNIST (8)	10	67.65 $\pm$ 1.66	64.49 $\pm$ 2.32	60.68 $\pm$ 2.91	58.83 $\pm$ 3.57	62.65 $\pm$ 2.27	54.36 $\pm$ 5.38
	50	66.75 $\pm$ 1.96	63.65 $\pm$ 2.97	30.00 $\pm$ 6.98	13.69 $\pm$ 3.98	10.33 $\pm$ 1.45	10.00 $\pm$ 0.00
	100	66.05 $\pm$ 2.25	58.66 $\pm$ 9.23	20.98 $\pm$ 4.25	10.00 $\pm$ 0.00	10.00 $\pm$ 0.00	10.00 $\pm$ 0.00

Table 6. TABLE IV: Description of tabular datasets.

Dataset	Samples	Features $N_0$	Output $N_D$	Task	Architecture
Adult	48,842	14	2	Classification	$\mathrm{\hspace{0.17em}14}\to 8\to \cdots \to 8\to 2$
Cancer	569	30	2	Classification	$\mathrm{\hspace{0.17em}30}\to 16\to \cdots \to 16\to 2$
Pima	768	8	2	Classification	$\mathrm{\hspace{0.17em}8}\to 4\to \cdots \to 4\to 2$
Ionosphere	351	34	2	Classification	$\mathrm{\hspace{0.17em}34}\to 16\to \cdots \to 16\to 2$
Wine	178	13	3	Classification	$\mathrm{\hspace{0.17em}13}\to 8\to \cdots \to 8\to 3$
Diabetes	442	10	1	Regression	$\mathrm{\hspace{0.17em}10}\to 8\to \cdots \to 8\to 1$

Table 7. TABLE V: Classification accuracy (%) for tabular datasets and RMSE for regression. Numbers shown under each initializer indicate network depth. Boldface denotes the best score among initializations for each ( dataset , α 0 ) (\text{dataset},\alpha_{0}) , and underlined boldface marks the best α \alpha within the same initializer.

Dataset	${\alpha }_0$	Proposed		Lee		He		Xavier		Orthogonal		Random
		50	100	50	100	50	100	50	100	50	100	50	100
Adult	-2	$76.07$	$76.07$	$85.33$	76.64	85.58	$76.07$	$76.07$	$76.07$	$76.07$	$76.07$	$76.07$	$76.07$
	0	85.91	$85.39$	$85.33$	85.64	$76.07$	$85.52$	$76.07$	$76.07$	$76.07$	$76.07$	$76.07$	$76.07$
	2	85.79	85.86	$85.59$	$85.64$	$76.07$	$76.07$	$76.07$	$76.07$	$76.07$	$76.07$	$76.07$	$76.07$
	50	$76.07$	$76.07$	$76.07$	$76.07$	$76.07$	$76.07$	$76.07$	$76.07$	$76.07$	$76.07$	$76.07$	$76.07$
Cancer	-2	$83.33$	$86.84$	96.49	93.86	$89.47$	$88.60$	$63.16$	$63.16$	96.49	$63.16$	$63.16$	$63.16$
	0	$94.74$	$92.11$	$94.74$	$93.86$	$89.47$	$81.58$	97.37	95.61	$96.49$	$63.16$	$63.16$	$63.16$
	2	97.37	97.37	$95.61$	$94.74$	$90.35$	$83.33$	$93.86$	$63.16$	$96.49$	$63.16$	$63.16$	$63.16$
	50	$71.05$	$42.11$	86.84	78.07	$63.16$	$63.16$	$63.16$	$63.16$	$63.16$	$63.16$	$63.16$	$63.16$
Pima	-2	$64.29$	65.58	68.18	65.58	$64.94$	$64.94$	$64.94$	$64.94$	$64.94$	$64.94$	$64.94$	$64.94$
	0	74.03	72.73	$73.38$	$71.43$	$64.94$	$64.94$	$64.94$	$64.94$	$64.94$	$64.94$	$64.94$	$64.94$
	2	$71.43$	73.38	72.73	$72.73$	$64.94$	$64.94$	$64.94$	$64.94$	$64.94$	$64.94$	$64.94$	$64.94$
	50	$56.49$	$56.49$	64.94	64.94	64.94	64.94	64.94	64.94	64.94	64.94	64.94	64.94
Ionosphere	-2	$70.42$	$69.01$	$88.73$	91.55	$84.51$	$87.32$	92.96	$64.79$	$87.32$	$64.79$	$64.79$	$64.79$
	0	94.37	$87.32$	94.37	91.55	$88.73$	$78.87$	$91.55$	$64.79$	$92.96$	$64.79$	$64.79$	$64.79$
	2	$91.55$	$90.14$	92.96	94.37	$78.87$	$70.42$	$87.32$	$64.79$	$91.55$	$64.79$	$64.79$	$64.79$
	50	$64.79$	80.28	92.96	$73.24$	$64.79$	$64.79$	$64.79$	$64.79$	$64.79$	$64.79$	$64.79$	$64.79$
Wine	-2	80.56	$47.22$	80.56	77.78	$38.89$	$38.89$	$38.89$	$38.89$	$38.89$	$38.89$	$38.89$	$38.89$
	0	86.11	94.44	$83.33$	$88.89$	$38.89$	$38.89$	$38.89$	$38.89$	$38.89$	$38.89$	$38.89$	$38.89$
	2	$91.67$	$88.89$	94.44	97.22	$38.89$	$38.89$	$38.89$	$38.89$	$38.89$	$38.89$	$38.89$	$38.89$
	50	$44.44$	$58.33$	83.33	83.33	$38.89$	$38.89$	$38.89$	$38.89$	$38.89$	$38.89$	$38.89$	$38.89$
Diabetes	-2	$70.53$	71.46	$161.88$	$161.88$	62.63	$162.78$	$73.33$	$73.87$	$70.68$	$72.84$	$72.79$	$162.83$
	0	$61.59$	59.62	$161.88$	$161.88$	$66.42$	$72.84$	56.46	$162.83$	$73.32$	$160.03$	$162.41$	$162.83$
	2	53.75	54.09	$161.88$	$161.88$	$64.64$	$73.18$	$63.65$	$72.97$	$60.77$	$162.83$	$162.79$	$162.83$
	50	72.15	71.09	$161.90$	$161.90$	$73.08$	$71.93$	$73.46$	$143.81$	$73.83$	$72.98$	$156.69$	$162.83$

Equations

$${\mathbf{y}}^{\ell }={\mathbf{W}}^{\ell }{\mathbf{x}}^{\ell -1}+{\mathbf{b}}^{\ell },{\mathbf{x}}^{\ell }=\varphi ({\mathbf{y}}^{\ell }),$$

$$\mathrm{ReLU}(\mathbf{z})=\max (0,\mathbf{z}).$$

$$W_{ij}\sim\mathcal{U}\Bigl{(}-\sqrt{\tfrac{6}{m+n}},\,\sqrt{\tfrac{6}{m+n}}\Bigr{)}.$$

$$W_{ij}\sim\mathcal{N}\Bigl{(}0,\tfrac{2}{m}\Bigr{)},$$

$$\frac{⟨\mathbf{x},1⟩}{\Vert \mathbf{x}\Vert \Vert 1\Vert }\approx \frac{⟨\mathbf{W}\mathbf{x},1⟩}{\Vert \mathbf{W}\mathbf{x}\Vert \Vert 1\Vert },$$

$$⟨\mathbf{W}\mathbf{x},1⟩=⟨\mathbf{x},{\mathbf{W}}^{T}1⟩,$$

$${\mathcal{O}}_{m,n}:=\left\{\mathbf{W}\in {\mathbb{R}}^{m\times n}:\mathbf{W}{\mathbf{W}}^T={\mathbf{I}}_m\text{ or }{\mathbf{W}}^T\mathbf{W}={\mathbf{I}}_n\right\}.$$

$$\mathbf{W}{\mathbf{W}}^T={\mathbf{I}}_m,$$

$${\mathbf{W}}^T\mathbf{W}={\mathbf{I}}_n.$$

$${\left\{{\mathbf{W}}_{(k)}\right\}}_{k=1}^{5000}\subset {\mathcal{O}}_{64,64},$$

$$s_k=⟨{\mathbf{v}}_k,1⟩=1^T{\mathbf{W}}_{(k)}1=\sum _{i=1}^{64}\sum _{j=1}^{64}({\mathbf{W}}_{(k)}{)}_{ij},$$

$$\underset{\mathbf{W}\in {\mathcal{O}}_{m,n}}{\max }\mathrm{tr}({\mathbf{J}}_{m\times n}^T\mathbf{W}).$$

$$\underset{\mathbf{W}\in {\mathcal{O}}_{m,n}}{\min }\Vert {\mathbf{J}}_{m\times n}-\mathbf{W}{\Vert }_F,$$

$$\tilde{\mathbf{W}}=\mathbf{U}{\mathbf{V}}^T,$$

$${\mathbf{J}}_{m\times n}=\hat{\mathbf{U}}\mathbf{\Sigma }{\hat{\mathbf{V}}}^T,$$

$$\|\bm{J}_{m\times n}-\bm{W}\|_{F}=\bigl{\|}\widehat{\bm{U}}^{T}(\bm{J}_{m\times n}-\bm{W})\,\widehat{\bm{V}}\bigr{\|}_{F}=\bigl{\|}\bm{\Sigma}-\widehat{\bm{U}}^{T}\bm{W}\,\widehat{\bm{V}}\bigr{\|}_{F}.$$

$$\underset{\mathbf{Q}\in {\mathcal{O}}_{m,m}}{\min }\Vert \mathbf{\Sigma }-\mathbf{Q}{\Vert }_F^2.$$

$$\Vert \mathbf{\Sigma }-\mathbf{Q}{\Vert }_F^2=mn+m-2\sqrt{mn}{Q}_{11},$$

$$\tilde{\mathbf{Q}}=\left[\begin{array}{cc}1& 0\\ 0& {\mathbf{Q}}^{\prime }\end{array}\right],\text{where}{\mathbf{Q}}^{\prime }\in {\mathcal{O}}_{m-1,m-1}.$$

$$\mathbf{U}=\hat{\mathbf{U}}\tilde{\mathbf{Q}}\in {\mathcal{O}}_{m,m},\mathbf{V}=\hat{\mathbf{V}}\in {\mathcal{O}}_{n,m}.$$

$$1_m^T\mathbf{W}{1}_n={\left({\mathbf{U}}^T{1}_m\right)}^T\left({\mathbf{V}}^T{1}_n\right)=\sqrt{mn}.$$

$$\widetilde{\mathcal{O}}_{m,n}:=\left\{\bm{W}\in\mathcal{O}_{m,n}\Bigl{|}\frac{1}{\sqrt{mn}}\bm{1}_{m}^{T}\bm{W}\bm{1}_{n}=\bm{\xi}_{m}^{T}\bm{W}\bm{\xi}_{n}=1\right\}.$$

$$\mathbf{W}=\mathbf{U}{\mathbf{V}}^T=\left[\begin{array}{ccc}\frac{\sqrt{6}-\sqrt{3}}{6}& \frac{\sqrt{6}-\sqrt{3}}{6}& \frac{\sqrt{6}+2\sqrt{3}}{6}\\ \frac{\sqrt{6}+\sqrt{3}}{6}& \frac{\sqrt{6}+\sqrt{3}}{6}& \frac{\sqrt{6}-2\sqrt{3}}{3},\end{array}\right],$$

$$\mathbf{U}=\left[\begin{array}{cc}\frac{1}{\sqrt{2}}& \frac{1}{\sqrt{2}}\\ \frac{1}{\sqrt{2}}& -\frac{1}{\sqrt{2}}\end{array}\right]\in {\mathcal{O}}_{2,2},\mathbf{V}=\left[\begin{array}{cc}\frac{1}{\sqrt{3}}& \frac{1}{\sqrt{6}}\\ \frac{1}{\sqrt{3}}& \frac{1}{\sqrt{6}}\\ \frac{1}{\sqrt{3}}& -\frac{2}{\sqrt{6}}\end{array}\right]\in {\mathcal{O}}_{3,2}.$$

$$\mathbf{W}=\mathbf{U}{\mathbf{V}}^T=\left[\begin{array}{cccc}|& |& & |\\ {\mathbf{u}}_1& {\mathbf{u}}_2& \cdots & {\mathbf{u}}_m\\ |& |& & |\end{array}\right]\left[\begin{array}{c}\text{—}{\mathbf{v}}_1^T\text{—}\\ \text{—}{\mathbf{v}}_2^T\text{—}\\ ⋮\\ \text{—}{\mathbf{v}}_m^T\text{—}\end{array}\right],$$

$$\Vert {\mathbf{W}}^T\mathbf{x}\Vert =\Vert \mathbf{x}\Vert \text{for all }\mathbf{x}\in {\mathbb{R}}^m.$$

$$\mathbf{W}=\left[\begin{array}{c}{\mathbf{w}}_1^T\\ ⋮\\ {\mathbf{w}}_m^T\end{array}\right]\in {\tilde{\mathcal{O}}}_{m,n},{\mathbf{W}}^{\perp }=\left[\begin{array}{c}{\mathbf{w}}_{m+1}^T\\ ⋮\\ {\mathbf{w}}_n^T\end{array}\right]\in {\mathcal{O}}_{n-m,n},$$

$${\mathbf{W}}^{\perp }{\mathbf{\xi }}_n=0.$$

$$1=\Vert {\mathbf{\xi }}_n{\Vert }^2={\Vert \left[\begin{array}{c}\mathbf{W}\\ {\mathbf{W}}^{\perp }\end{array}\right]{\mathbf{\xi }}_n\Vert }^2=1+\Vert {\mathbf{W}}^{\perp }{\mathbf{\xi }}_n{\Vert }^2,$$

$${\mathbf{h}}_i:={\mathbf{w}}_i-{\mathrm{proj}}_{{\mathbf{\xi }}_n}({\mathbf{w}}_i)\text{for all }i=1,\dots ,m.$$

Full text

I Introduction

Deep learning has achieved remarkable success across various domains, driven by its ability to learn effective representations from data through deep neural networks [1]. These architectures come in various forms, each designed to address specific data processing challenges. A representative example is the feed-forward neural network (FFNN), which serves as a fundamental building block in deep learning, where information propagates unidirectionally from the input layer, through multiple hidden layers, to the output layer [2]. The depth of the network determines representational capacity, since additional layers enable the progressive extraction of more abstract and complex features [3, 4]. On the other hand, increasing depth may lead to the vanishing and exploding gradient phenomenon that hinders practical training and delays convergence [5].

The Rectified Linear Unit (ReLU) activation function is widely adopted in modern deep learning models due to its simple form, computational efficiency, and sparsity-inducing properties in neuron activity [2, 5, 6]. However, the ReLU activation suffers from the dying ReLU problem, when a pre-activation remains negative, its gradient vanishes and the corresponding parameters stop updating [7, 8]. This problem is one of the vanishing gradient [5], and it becomes more pronounced in deeper networks.

To address the dying ReLU problem, various studies have been proposed. The main approaches can be broadly classified into three categories. First, some studies employ activation functions that permit small gradients for negative inputs, such as Leaky ReLU [7] or the exponential linear unit (ELU) [9], or modify the network structure itself by introducing residual connections [10]. Second, there are widely used methods that apply normalization techniques, such as batch normalization [11], to each layer to stabilize the distribution of activation values. Lastly, there are studies on initialization techniques that carefully set weights and biases to prevent neuron deactivation early in training [5, 8]. Among these three approaches, this paper focuses on weight initialization, which is involved in the most fundamental stage of training. It proposes a new method to enhance the stability and performance of ReLU networks.

Early methods, such as Xavier initialization [5], were designed under the assumption of symmetric, zero-centered activations (e.g., tanh, sigmoid), and often underperform with ReLU-based models due to the asymmetry of the function. In response, He initialization [8] scales weight variances by $2/m$, where $m$ is the number of input units (fan-in) of the layer to compensate for ReLU’s half-activation. More recent approaches exploit geometric constraints and dynamical-isometry conditions to better tailor initialization to ReLU’s properties [12, 13]. In a recent study, Lee et al. [14] proposed a specific constructive method for orthogonal weight matrices by applying QR decomposition to an all-ones matrix perturbed by a small term $\epsilon>0$. While it is effective, this constructive method treats a key property—approximate angle preservation with the all-ones vector—as a byproduct of its procedure rather than a direct optimization objective. Consequently, these desirable characteristics are only approximately guaranteed.

To address these limitations, this paper introduces a novel weight initialization method derived from a principled optimization framework. We formulate 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. This approach is motivated by our empirical and theoretical findings that such alignment is critical for preventing neuron inactivation in ReLU networks. The resulting initialization scheme, obtained as the exact solution to this optimization problem, is designed to inherently preserve signal propagation and stabilize training dynamics from the very first epoch, thereby mitigating the dying ReLU and vanishing gradient problems in deep and narrow architectures.

To validate the practical benefits of our approach, the proposed initialization is evaluated on several benchmark tasks. Deep ReLU neural networks trained on MNIST and Fashion-MNIST exhibit faster convergence and improved generalization compared to standard initializations [15]. Robustness to activation choice is assessed through experiments with various ReLU variants, confirming consistent performance gains. In few-shot learning scenarios—where only a handful of labeled examples are available—the initialization maintains stability and accuracy. Finally, the application to diverse tabular datasets demonstrates its versatility across different data modalities. Across these diverse scenarios, our proposed scheme consistently demonstrates superior performance compared to existing initialization methods, underscoring its effectiveness in addressing the specific challenges of ReLU activation.

The main contributions of this paper are summarized as follows.

We solve an optimization problem on the Stiefel manifold to derive novel weight initialization and analyze its key mathematical properties. 2. 2.

We design an efficient algorithm that enables practical weight initialization. 3. 3.

We theoretically and empirically demonstrate that the proposed initialization alleviates early-stage issues in ReLU networks such as dying neurons and unstable gradients. 4. 4.

We show superior performance over previous initializations and stability of training in deep ReLU networks across MNIST, Fashion-MNIST, and various tabular benchmarks, including few-shot regimes and variants of ReLU activations.

Notations

Throughout this paper, we adopt the following notation conventions. Let $\mathbb{R}$ (resp. $\mathbb{R}_{+}$) be the set of real numbers (resp. nonnegative real numbers). The standard inner product of two vectors $\bm{u}$ and $\bm{v}$ is denoted by $\langle\bm{u},\bm{v}\rangle$, and $\|\bm{v}\|_{2}$ denotes the Euclidean norm. The vector max norm is denoted by $\|\bm{x}\|_{\infty}$. The superscript $T$ denotes the transpose operator. Denote the all-ones vector of size $m$ by $\bm{1}_{m}=[1,1,\dots,1]^{T}\in\mathbb{R}^{m}$ and its normalized version by $\bm{\xi}_{m}=\tfrac{1}{\sqrt{m}}\bm{1}_{m}$. The all-ones matrix of size $m\times n$ is denoted by $\bm{J}_{m\times n}$ where the subscript is shortened to $\bm{J}_{m}$ for a square matrix of size $m$. And the $m\times m$ identity matrix is denoted by $\bm{I}_{m}$. Denote the zero matrix of size $m\times n$ by $\bm{0}_{m\times n}$. When the dimension is clear from context, we omit subscripts on $\bm{1}$, $\bm{\xi}$, $\bm{J}$, $\bm{I}$, and $\bm{0}$.

We use the notation $x\sim\mathcal{N}(\mu,\sigma^{2})$ to indicate that the scalar random variable $x$ follows a univariate normal distribution with mean $\mu\in\mathbb{R}$ and variance $\sigma^{2}>0$. For the multivariate case, $\bm{x}\sim\mathcal{N}(\bm{\mu},\Sigma)$ denotes that the random vector $\bm{x}\in\mathbb{R}^{d}$ follows a multivariate normal distribution with mean vector $\bm{\mu}\in\mathbb{R}^{d}$ and covariance matrix $\bm{\Sigma}\in\mathbb{R}^{d\times d}$, which is symmetric and positive semidefinite. We write $x\sim\mathcal{U}(a,b)$ to denote that $x$ is distributed uniformly over the interval $[a,b]\subset\mathbb{R}$, where $a<b$.

II Preliminaries

Before introducing our proposed weight initialization method, we provide a brief overview of the basic concepts and review prior works.

II-A Basic Concepts

FFNN is a composition of affine maps and pointwise nonlinearities. Given an input $\bm{x}^{0}\in\mathbb{R}^{N_{0}}$, the forward recursion for layers $\ell=1,\ldots,L$ is given by

[TABLE]

where $N_{\ell}$ denotes the number of units in the $\ell$-th layer. Here, $\bm{W}^{\ell}\in\mathbb{R}^{N_{\ell}\times N_{\ell-1}}$ is the weight matrix, $\bm{b}^{\ell}\in\mathbb{R}^{N_{\ell}}$ is its bias vector, $\phi:\mathbb{R}\to\mathbb{R}$ is an activation function.

For FFNN architectures based on fully connected layers, non-expanding configurations are widely employed, where the number of neurons either remains constant or decreases across layers. Such architectures, commonly found in standard designs [2, 16], support hierarchical compression of features and often lead to improved generalization and interpretability. Accordingly, in this work, a non-expanding architecture is considered, specified by the condition $N_{0}\geq N_{1}\geq\cdots\geq N_{L}$.

An activation function $\phi$ is a key component in neural networks, introducing non-linearity that allows the model to learn complex, non-linear mappings. Common choices include sigmoid [17], hyperbolic tangent (tanh) [18], and rectified linear unit (ReLU) functions [6], each with distinct characteristics in terms of saturation, gradient flow, and computational cost. In this paper, the activation function $\phi$ is chosen to be the ReLU, defined element-wise as

[TABLE]

The ReLU activation is widely used due to its computational simplicity and beneficial properties for optimization. It introduces non-linearity while preserving gradient flow for positive inputs, thereby mitigating the vanishing gradient problem. Furthermore, ReLU induces sparsity in the activations, as it outputs zero for all non-positive values. While ReLU activations offer several optimization advantages, their effectiveness is closely tied to the choice of weight initialization [19, 20]. Inappropriate initialization may lead to unbalanced activation statistics across layers, undermining training dynamics [21]. Consequently, considerable research has been devoted to developing initialization schemes, as reviewed below.

II-B Prior Works

Before the widespread adoption of ReLU, activation functions such as sigmoid and tanh were more commonly used. One of the most influential methods in this category is Xavier initialization [5], which draws a weight matrix $\bm{W}=[W_{ij}]\in\mathbb{R}^{m\times n}$ from a uniform distribution, i.e., for all $i,j$

[TABLE]

Xavier initialization works well for activation functions with symmetry around zero, effectively preserving variance in shallow networks. However, because ReLU activations are asymmetric, this method often struggles to avert the vanishing gradient problem in deeper ReLU-based networks.

To address the reduction in activation caused by ReLU zeroing out approximately half of its inputs, He et al. [8] proposed to increase the variance more aggressively:

[TABLE]

where weights are sampled from a zero-mean normal distribution with variance $\frac{2}{m}$. This strategy compensates for the sparsity introduced by ReLU, helping to preserve the variance of activations in deeper networks. It has become standard in many ReLU-based architectures, including various residual network designs, although it does not entirely prevent gradient degradation in extremely deep architectures.

Other researchers investigated matrix structures and their impact on signal propagation. For instance, Saxe et al. [13] advocated orthogonal initialization, leveraging the property $\bm{W}\bm{W}^{T}=\bm{I}$ to preserve signal norms through forward and backward passes in linear regimes. Orthogonal initialization has been shown to maintain dynamical isometry, thereby stabilizing gradient flow in deep networks. Building upon this intuition, Hu et al. [22] provided a theoretical justification for orthogonal initialization in deep linear networks. Their analysis demonstrates that orthogonal weights enable depth-independent convergence, in contrast to Gaussian initialization, which requires network width to increase linearly with depth for efficient training. This result offers a provable advantage for orthogonal initialization in specific settings.

Despite these theoretical benefits, orthogonal initialization remains challenging to apply in practice. Generating exact orthogonal matrices is computationally expensive, particularly for high-dimensional or convolutional layers. Moreover, its effectiveness diminishes in highly nonlinear regimes, where the orthogonality of weight matrices does not guarantee stable signal propagation due to the nonlinear distortions introduced by activation functions.

While these methods have advanced the field, deep and narrow ReLU networks pose additional challenges. In such architecture, the vanishing gradient and the dying ReLU become especially pronounced. Recently, Lee et al. [14] addressed these difficulties by proposing a deterministic initialization method that constructs an orthogonal weight matrix via a perturbed all-one matrix $\bm{1}\bm{1}^{T}+\epsilon\bm{I}$ for sufficiently small $\epsilon>0$. This method preserves orthogonality in a manner particularly advantageous for narrow networks. Empirically, this weight initialization maintains more positive entries than either He initialization or standard orthogonal matrices, thereby reducing the risk of the dying ReLU. Furthermore, Lee et al. establish that their construction approximately preserves the angle between any input vector $\bm{x}$ and the all-ones vector $\bm{1}$. Specifically, the proposed initialization $\bm{W}$ satisfies

[TABLE]

so the angle between any positive input $\bm{x}$ and the all-ones vector $\bm{1}$ is preserved by $\bm{W}$. Since

[TABLE]

a strong positive alignment of $\bm{W}^{T}\bm{1}$ with $\bm{x}$ immediately implies $\langle\bm{W}\bm{x},\bm{1}\rangle>0$. In other words, all pre-activations remain strictly positive, thus preventing the dying ReLU phenomenon in deep ReLU networks.

III Methodology

Motivated by the angle-preserving behavior observed in (2), we hypothesize that the alignment between the weight matrix and the all-ones vector $\bm{1}$ serves as a key factor in determining ReLU activation. Denote the set of all semi-orthogonal matrices as

[TABLE]

If $m\leq n$, the rows of $\bm{W}$ are orthonormal, so that

[TABLE]

whereas if $m\geq n$, the columns of $\bm{W}$ are orthonormal, so that

[TABLE]

III-A Empirical Motivation via Alignment Optimization

We introduce an optimization-based framework for constructing semi-orthogonal weight matrices that are specifically aligned with directions favorable to ReLU activation. The guiding principle stems from empirical observations indicating that the alignment between the semi-orthogonal weight matrix and the all-ones vector, representing uniform positive inputs, correlates positively with network performance.

To understand how semi-orthogonal weight matrices align with the all-ones vector, we first observe how randomly sampled semi-orthogonal matrices align with the all-ones vector to understand their effect on ReLU activation behavior. Consider a collection of random semi-orthogonal matrices,

[TABLE]

where each $\bm{W}_{(k)}$ is randomly generated using the orthogonal initialization scheme of Saxe et al. [13], involving QR decomposition of a Gaussian matrix. For each $k$, define the 64-dimensional vector $\bm{v}_{k}=\bm{W}_{(k)}^{T}\bm{1}$. Since $\{\bm{v}_{k}\}_{k=1}^{5000}$ lie in a 64-dimensional space, it is challenging to visualize clustering or distribution patterns directly. There are several dimension reduction methods to visualize them in a low-dimensional space. Among them, we adopt Uniform Manifold Approximation and Projection (UMAP) [23] because it supports user-defined distance metrics, which allows us to employ the 1-Wasserstein distance. This choice provides invariance to coordinate permutations and enables a meaningful comparison of the alignment vectors [24].

In Fig. 1 (a) and (b), an identical 2D UMAP embedding is used for both plots, i.e., each point occupies the same location in both subfigures and corresponds to the same $\bm{W}_{(k)}$. In Fig. 1 (a), points are colored by their alignment score

[TABLE]

which quantifies how strongly $\bm{v}_{k}$ aligns with the uniform direction $\bm{1}=(1,1,\dots,1)$; that is, $s_{k}=\langle\bm{v}_{k},\bm{1}\rangle$ measures the extent to which $\bm{v}_{k}$ projects onto the subspace spanned by $\bm{1}$, or equivalently, the sum of all entries in $\bm{W}_{(k)}$. This reflects how uniformly the components of $\bm{v}_{k}$ are distributed across all coordinates, indicating a globally consistent activation pattern. In Fig. 1 (b), the points are colored instead by the average one-shot classification accuracy $\alpha_{k}$, obtained when $\bm{W}_{(k)}$ is used to initialize the middle layer of a simple fully connected ReLU network trained on MNIST [15]. Specifically, the network architecture is $784\to 64\to 64\to 10$. The input-to-hidden ($64\times 784$) and hidden-to-output ($10\times 64$) weight matrices are initialized using orthogonal initialization. The middle $64\times 64$ weight matrix is set to the specific sample $\bm{W}_{(k)}$ under evaluation. For each $\bm{W}_{(k)}$, we run 100 independent one-shot trials: in each trial, one random image per digit class is sampled from the MNIST training set, training proceeds for 100 epochs using vanilla SGD (learning rate 0.01, batch size 64, no weight decay) under cross-entropy loss, and we measure test accuracy on the full MNIST test set. The color in Fig. 1(b) represents the average test accuracy $\alpha_{k}$ across these 100 trials.

Although Fig. 1 (a) and (b) are colored by two different scalar metrics—alignment score $s_{k}$ and average classification accuracy $\alpha_{k}$, respectively—the resulting spatial patterns are remarkably similar. This visual correspondence suggests a positive relationship between the alignment of $\bm{W}_{(k)}$ with the direction of the all-ones vector and the effectiveness of the corresponding initialization. A Pearson correlation analysis on the dataset $\left\{(s_{k},\alpha_{k})\right\}_{k=1}^{5000}$ revealed a strong, statistically significant linear correlation between $s_{k}$ and $\alpha_{k}$ ($r=0.8178$). This finding supports the hypothesis that alignment with the all-one direction positively contributes to performance in ReLU networks. Motivated by this result, we construct a new class of weight matrices $\bm{W}\in\mathbb{R}^{m\times n}$ that satisfy the following two criteria:

• (i)

$\bm{W}$ is semi-orthogonal,

• (ii)

The value of $\bm{1}_{m}^{T}\bm{W}\bm{1}_{n}$ is maximized.

Since no more than $n$ orthonormal vectors can exist in $\mathbb{R}^{n}$, the requirement of row orthonormality implies $m\leq n$. Throughout the paper, we adopt the assumption $m\leq n$, i.e., $\bm{W}\bm{W}^{T}=\bm{I}$.

III-B Optimization on the Stiefel Manifold

Note that $\mathcal{O}_{m,n}$ in (4) is a real Stiefel manifold, which mathematically represents the collection of all ordered orthonormal $m$-frames in $\mathbb{R}^{n}$. Each element in $\mathcal{O}_{m,n}$ can be viewed as a point on a nonlinear manifold embedded in Euclidean space, constrained by $\frac{1}{2}m(m+1)$ nonlinear equations due to the orthonormality condition. Notably, $\mathcal{O}_{m,n}$ is a smooth Riemannian manifold, and it generalizes both the unit sphere $\mathbb{S}^{n-1}=\mathcal{O}_{1,n}$ and the orthogonal group $\mathcal{O}_{n,n}$.

An optimization problem is now formulated over the Stiefel manifold $\mathcal{O}_{m,n}$ to satisfy the above criteria. The scalar $\bm{1}_{m}^{T}\bm{W}\bm{1}_{n}$, representing the sum of all entries of $\bm{W}$, is equivalent to $\operatorname*{tr}(\bm{J}_{m\times n}^{T}\bm{W})$, where $\bm{J}_{m\times n}$ is the all-ones matrix. Thus, the optimization can be posed as

[TABLE]

Equivalently, it can be expressed as

[TABLE]

which interprets the problem as finding the closest point on the Stiefel manifold to $\bm{J}_{m\times n}$ in terms of the Frobenius norm.

Theorem 1**.**

Every optimal solution of (6) can be explicitly expressed as follows:

[TABLE]

where $\bm{U}\in\mathcal{O}_{m,m}$ has its first column equal to $\bm{\xi}_{m}$, and $\bm{V}\in\mathcal{O}_{n,m}$ has its first column equal to $\bm{\xi}_{n}$.

Proof.

Note that the thin SVD of $\bm{J}_{m\times n}$ is

[TABLE]

where $\bm{\Sigma}$ is the $m\times m$ diagonal matrix with $\sqrt{mn}$ in the first diagonal element and [math] elsewhere, $\widehat{\bm{U}}\in\mathcal{O}_{m,m}$ and $\widehat{\bm{V}}\in\mathcal{O}_{n,m}$ have their first columns equal to $\bm{\xi}_{m}$ and $\bm{\xi}_{n}$, respectively. By the orthogonal invariance of the Frobenius norm, it holds

[TABLE]

By setting $\bm{Q}=\widehat{\bm{U}}^{T}\bm{W}\,\widehat{\bm{V}}\in\mathcal{O}_{m,m}$, the problem (6) can be simply rewritten as

[TABLE]

Note that

[TABLE]

where $Q_{11}$ is $(1,1)$ entry of $\bm{Q}$. So, solving (8) is equivalent to maximizing $Q_{11}$. Since the first column vector of $\bm{Q}$ has norm $1$, it follows that the maximum possible value of $Q_{11}$ is $1$. Thus, the optimal solution of (8) is

[TABLE]

Since $\widehat{\bm{U}}\,\widetilde{\bm{Q}}$ is a semi-orthogonal matrix and $\widehat{\bm{U}}\,\widetilde{\bm{Q}}$ has the first column $\bm{\xi}_{m}$, $\widehat{\bm{U}}\,\widetilde{\bm{Q}}$ can be considered as an element in $\mathcal{O}_{m,m}$ whose the first column equals $\bm{\xi}_{m}$. Finally, set

[TABLE]

By construction, the first columns of $\bm{U}$ and $\bm{V}$ are $\bm{\xi}_{m}$ and $\bm{\xi}_{n}$, respectively, and $\widetilde{\bm{W}}=\bm{U}\,\bm{V}^{T}$. ∎

The optimal value of the objective function in (6) can be explicitly computed by evaluating the solution characterized in Theorem 1. Substituting the optimal solution $\widetilde{\bm{W}}=\bm{U}\bm{V}^{T}$ yields

[TABLE]

Denote the set of all optimal solutions to the maximization problem in (6) as $\widetilde{\mathcal{O}}_{m,n}$. Equivalently,

[TABLE]

The following is one example from the set $\mathcal{O}_{m,n}$for small values $m,n$.

Example 1**.**

Let $m=2$, $n=3$, and consider the matrix

[TABLE]

where $\bm{U}$ and $\bm{U}$ are semi-orthogonal matrices given by

[TABLE]

Then $\bm{W}\in\mathcal{O}_{2,3}$ and $\bm{\xi}_{m}^{T}\bm{W}\bm{\xi}_{n}=1$. Thus, $\bm{W}\in\widetilde{\mathcal{O}}_{2,3}$.

This article proposes to initialize the weight matrix of each layer by independently sampling from $\widetilde{\mathcal{O}}_{m,n}$.

IV Theoretical properties

This section investigates several properties of matrices in $\widetilde{\mathcal{O}}_{m,n}$ whose elements are the optimal solutions to the maximization problem in (6). We first analyze its theoretical characteristics, highlighting key structural and alignment features that are preserved under the orthogonality constraint. Next, we present an efficient construction method for weight matrices in the solution set $\widetilde{\mathcal{O}}_{m,n}$ that exactly satisfy the prescribed optimization constraints.

IV-A Matrix Structure and Characterization

Each matrix $\bm{W}\in\widetilde{\mathcal{O}}_{m,n}$ admits the decomposition

[TABLE]

where $\bm{U}\in\mathcal{O}_{m,m}$, $\bm{V}\in\mathcal{O}_{n,m}$, with the leading singular vectors aligned to the normalized all-ones vectors: $\bm{u}_{1}=\xi_{m}$, $\bm{v}_{1}=\xi_{n}$. Since every $\bm{W}\in\widetilde{\mathcal{O}}_{m,n}$ is semi-orthogonal, it satisfies $\bm{W}\bm{W}^{T}=\bm{I}_{m}$, ensuring that the transformation preserves input norms:

[TABLE]

On the other hand, the right-side composition $\bm{W}^{T}\bm{W}=\sum_{i=1}^{m}\bm{v}_{i}\bm{v}_{i}^{T}$ does not yield the identity. Instead, it defines an orthogonal projection onto the span of $\{\bm{v}_{1},\dots,\bm{v}_{m}\}\subset\mathbb{R}^{n}$.

In addition to orthogonality and projection properties, the following lemma captures a key structural characterization of the matrices in $\widetilde{\mathcal{O}}_{m,n}$.

Lemma 1**.**

Suppose that $\bm{W}$ is a semi-orthogonal matrix. Then the following are equivalent:

1. $($i$)$

$\bm{W}\in\widetilde{\mathcal{O}}_{m,n}$, 2. $($ii$)$

$\bm{W}\bm{\xi}_{n}=\bm{\xi}_{m}$, 3. $($iii$)$

$\bm{W}^{T}\bm{\xi}_{m}=\bm{\xi}_{n}$.

Proof.

$[(i)\Rightarrow(ii)]$ If $\bm{W}\in\widetilde{\mathcal{O}}_{m,n}$, by Theorem 1, $\bm{W}$ has the decomposition $\bm{W}=\bm{U}\,\bm{V}^{T}$ for $\bm{U}$, $\bm{V}$ as in (11). Since $\bm{v}_{1}^{T}\,\bm{\xi}_{n}=1$ and $\bm{v}_{i}^{T}\,\bm{\xi}_{n}=0$ for all $i\geq 2$, it follows that $\bm{W}\bm{\xi}_{n}=\bm{U}\,\left(\bm{V}^{T}\bm{\xi}_{n}\right)=\bm{U}[1\,0\,\cdots\,0]^{T}=\,\bm{u}_{1}=\bm{\xi}_{m}.$

$[(ii)\Rightarrow(iii)]$ If $\bm{W}\bm{\xi}_{n}=\bm{\xi}_{m}$, then $1=\bm{\xi}_{m}^{T}\bm{\xi}_{m}=(\bm{W}\bm{\xi}_{n})^{T}\bm{\xi}_{m}=\langle\bm{\xi}_{n},~\bm{W}^{T}\bm{\xi}_{m}\rangle$, implying that $\bm{W}^{T}\bm{\xi}_{m}=\bm{\xi}_{n}$, since $\|\bm{W}^{T}\bm{\xi}_{m}\|=1$.

$[(iii)\Rightarrow(i)]$ $\bm{1}_{n}^{T}\bm{W}^{T}\bm{1}_{m}=\sqrt{mn}(\bm{\xi}_{n}^{T}\bm{W}^{T}\bm{\xi}_{m})=\sqrt{mn}(\bm{\xi}_{n}^{T}\bm{\xi}_{n})=\sqrt{mn}$, so $\bm{W}$ is an optimal solution of (5), i.e., $\bm{W}\in\widetilde{\mathcal{O}}_{m,n}$. ∎

Lemma 1 characterizes $\widetilde{\mathcal{O}}_{m,n}$ as the set of all semi-orthogonal matrices that preserve the normalized all-ones vector. Since the three conditions are equivalent, either (ii) or (iii) in Lemma 1 can fully characterize $\widetilde{\mathcal{O}}_{m,n}$. This lemma plays a central role in the derivation of subsequent properties.

It is well known that any semi-orthonormal matrix can be extended to an orthonormal matrix when the number of rows is strictly less than the number of columns [25]. Specifically, for a semi-orthonormal matrix $\bm{W}\in\widetilde{\mathcal{O}}_{m,n}$, where $m<n$, there exists a matrix $\bm{W}^{\perp}\in\mathbb{R}^{(n-m)\times n}$ whose rows form an orthonormal basis for the orthogonal complement of the row space of $\bm{W}$ such that $\begin{bmatrix}\bm{W}\\ \bm{W}^{\perp}\end{bmatrix}\in\mathcal{O}_{n,n}.$ Write

[TABLE]

where $\{\bm{w}_{i}\}_{i=1}^{n}\subset\mathbb{R}^{n}$ forms orthonormal basis for $\mathbb{R}^{n}$.

Proposition 1**.**

The rows of $\bm{W}^{\perp}$ in (12) are in the subspace orthogonal to $\bm{\xi}_{n}$, i.e.,

[TABLE]

Proof.

By Lemma 1 (ii), it follows that

[TABLE]

implying ${\bm{W}^{\perp}}\bm{\xi}_{n}=\bm{0}$. ∎

This proposition reveals a fundamental structural property of any full orthonormal basis extended from a matrix $\bm{W}\in\widetilde{O}_{m,n}$. It demonstrates that the row space of the complete $n\times n$ orthogonal matrix is partitioned with respect to the all-ones vector $\bm{\xi}_{n}$. The first $m$ basis vectors (the rows of $\bm{W}$) collectively map $\bm{\xi}_{n}$ to $\bm{\xi}_{m}$, while the remaining $n-m$ basis vectors (the rows of $\bm{W}^{\perp}$) lie entirely in the subspace orthogonal to $\bm{\xi}_{n}$. This complete structural understanding forms the theoretical basis for designing and verifying the construction methods that follow.

IV-B Construction of Proposed Weight Matrices

As stated in Theorem 1, when $m=1$, the set $\widetilde{\mathcal{O}}_{1,n}$ contains only a single vector $\bm{\xi}_{n}$, which leads to a fully deterministic construction. So we focus on the more general and nontrivial cases where $m\geq 2$, for which the proposed algorithm is applicable. Algorithm 1 presents a construction method to generate random weight matrices from $\widetilde{\mathcal{O}}_{m,n}$ using (11). Algorithm 1 first constructs two tall matrices $[\bm{\xi}_{m}\,\bm{A}]\in\mathbb{R}^{m\times m}$ and $[\bm{\xi}_{n}\,\bm{B}]\in\mathbb{R}^{n\times m}$, where $\bm{A}$ and $\bm{B}$ have i.i.d. $\mathcal{N}(0,1)$ entries. Then it performs thin QR decompositions on these blocks to obtain orthogonal factors $\bm{U}\in\mathbb{R}^{m\times m}$ and $\bm{V}\in\mathbb{R}^{n\times m}$, respectively. Next, diagonal sign-correction matrices $\bm{\Lambda}$ and $\bm{\Gamma}$ are applied so that the first columns of $\bm{U}$ and $\bm{V}$ match $\bm{\xi}_{m}$ and $\bm{\xi}_{n}$, respectively. Finally, $\bm{W}=\bm{U}\bm{V}^{T}$ yields a semi-orthogonal matrix in $\widetilde{\mathcal{O}}_{m,n}$. In particular, $\bm{U}$ and $\bm{V}$ retain their fixed first column. In contrast, their remaining columns are semi-orthogonal and uniformly sampled from the corresponding subspace according to the Haar distribution [26].

Note that Algorithm 1 requires QR decomposition for both $n\times n$ matrices and $m\times m$ matrices, respectively. The QR decomposition of an $m\times m$ matrix has a computational complexity of $O(m^{3})$. To reduce high complexity, we propose an improved algorithm to construct $\bm{W}\in\widetilde{\mathcal{O}}_{m,n}$ with just one QR decomposition through the following process.

For $\bm{W}\in\widetilde{\mathcal{O}}_{m,n}$, consider the set of orthonormal vectors $\{\bm{w}_{i}\}_{i=1}^{m}\subset\mathbb{R}^{n}$ such that $\bm{w}_{i}^{T}$ is the $i$-th row vector of $\bm{W}$ for all $i=1,\ldots,m$. Define

[TABLE]

Geometrically, each $\bm{h}_{i}$ is the orthogonal projection of $\bm{w}_{i}$ onto the $(m-1)$-dimensional subspace $\{\bm{x}\in\mathbb{R}^{n}:\bm{\xi}_{n}^{T}\bm{x}=0\}$.

Proposition 2**.**

Let $n\geq m\geq 2$. Consider $\bm{H}=[\bm{h}_{1}\,\bm{h}_{2}\,\cdots\,\bm{h}_{m}]$ where $\bm{h}_{i}$ is defined in (13), then Gram matrix $\bm{H}^{T}\bm{H}$ has following property:

[TABLE]

for all $i,j=1,\ldots,m$.

Proof.

By Lemma 1 it holds that for $i=1,\ldots,m$,

[TABLE]

Then it follows that

[TABLE]

for all $i,j=1,\ldots,m$. ∎

For $m\geq 2$, let

[TABLE]

Consider multiplying $\bm{P}_{m}$ by an $m$-dimensional vector $\bm{x}\in\mathbb{R}^{m}$,

[TABLE]

This subtracts the component of the original vector $\bm{x}$ in the direction of $\bm{\xi}_{m}$, which geometrically corresponds exactly to projecting $\bm{x}$ onto $\bm{\xi}_{m}^{\perp}$. Moreover, one can see that $\bm{P}_{m}$ is an $m\times m$ symmetric positive semidefinite matrix. So, $\bm{P}_{m}$ admits a Cholesky decomposition $\bm{P}_{m}=\bm{L}\bm{L}^{T}$ where $\bm{L}$ is a lower triangular matrix with nonnegative diagonal entries.

Lemma 2**.**

Let $m\geq 2$. Then $\bm{P}_{m}=\bm{I}_{m}-\frac{1}{m}\bm{J}_{m}\in\mathbb{R}^{m\times m}$ admits a unique Cholesky factorization

[TABLE]

where $\bm{L}\in\mathbb{R}^{m\times m}$ is a lower triangular matrix with nonnegative diagonal entries, explicitly given by

[TABLE]

Proof.

The uniqueness follows from the fact that all leading principal minors of $\bm{P}_{m}$ are positive, ensuring the existence of a unique Cholesky factor with nonnegative diagonal entries [27]. ∎

Since the Cholesky factor $\bm{L}$ of $\bm{P}_{m}$ is explicitly available, we can construct the matrix $\bm{H}^{T}=\bm{L}\bm{Q}$ for a certain semi-orthogonal matrix $\bm{Q}\in\mathcal{O}_{m,n}$. Using the relation in (14),

[TABLE]

In matrix form, it can be expressed as

[TABLE]

The following theorem provides the mathematical foundation for this decomposition and guarantees that all such matrices satisfy the defining conditions of $\widetilde{\mathcal{O}}_{m,n}$.

Theorem 2**.**

Let $m\geq 2$.

[TABLE]

where $\bm{L}$ is the $m\times m$ lower triangular matrix defined in Lemma 2.

Proof.

First, we prove the reverse inclusion ($\supseteq$). Let $\bm{W}$ be an arbitrary matrix from the set on the right-hand side, such that $\bm{W}=\bm{L}\bm{Q}+\frac{1}{\sqrt{mn}}\bm{J}_{m\times n}$, where $\bm{Q}\in\mathcal{O}_{m,n}$ is a semi-orthogonal matrix whose $m$-th row is $\bm{\xi}_{n}^{T}$. By Lemma 1 it is enough to show that $\bm{W}\bm{\xi}_{n}=\bm{\xi}_{m}$ and $\bm{W}\bm{W}^{T}=\bm{I}_{m}$. Since the rows of $\bm{Q}$ are orthonormal and its $m$-th row is $\bm{\xi}_{n}^{T}$, the vector $\bm{Q}\bm{\xi}_{n}$ is equal to $\bm{e}_{m}=[0,\dots,0,1]^{T}$. Then it follows that

[TABLE]

Next, to show that $\bm{W}$ is semi-orthogonal, we compute the product $\bm{W}\bm{W}^{T}$. Since $\bm{L}\bm{Q}\bm{J}^{T}=\sqrt{n}\bm{L}(\bm{Q}\bm{\xi}_{n})\bm{1}_{m}^{T}=\sqrt{n}\bm{L}\bm{e}_{m}\bm{1}_{m}^{T}=\bm{0}$,

[TABLE]

Having shown that $\bm{W}$ is a semi-orthogonal matrix satisfying the conditions of Lemma 1, it implies that $\bm{W}\in\widetilde{\mathcal{O}}_{m,n}$.

Next, we prove the forward inclusion ($\subseteq$). Let $\bm{W}$ be an arbitrary matrix in $\widetilde{\mathcal{O}}_{m,n}$. Define the matrix $\bm{M}=\bm{W}-\frac{1}{\sqrt{mn}}\bm{J}_{m\times n}$. By using the properties of $\bm{W}$ from Lemma 1, we compute the product $\bm{M}\bm{M}^{T}$ as following:

[TABLE]

From Lemma 2, we have $\bm{M}\bm{M}^{T}=\bm{P}_{m}=\bm{L}\bm{L}^{T}$. It follows from [25, Theorem 7.3.11] that there exists a semi-orthogonal matrix $\bm{Q}\in\mathcal{O}_{m,n}$ such that $\bm{M}=\bm{L}\bm{Q}$. To determine the properties of $\bm{Q}$, we use the property from Lemma 1 that $\bm{W}\bm{\xi}_{n}=\bm{\xi}_{m}$, which implies $\bm{M}\bm{\xi}_{n}=\bm{W}\bm{\xi}_{n}-\frac{1}{\sqrt{mn}}\bm{J}\bm{\xi}_{n}=\bm{\xi}_{m}-\bm{\xi}_{m}=\bm{0}$. Substituting the factorization gives $\bm{L}(\bm{Q}\bm{\xi}_{n})=\bm{0}$. From Lemma 2, the last column of $\bm{L}$ is zero. Thus, the product $\bm{M}=\bm{L}\bm{Q}$ implies that $\bm{M}$ is determined solely by the first $m-1$ rows of $\bm{Q}$. The property $\bm{W}\in\widetilde{\mathcal{O}}_{m,n}$ requires $\bm{M}\bm{\xi}_{n}=\bm{0}$. This yields a system of linear equations for the terms $x_{j}=\bm{q}_{j}^{T}\bm{\xi}_{n}$ for $j=1,2,\ldots,m-1$. Since the first $m-1$ columns of $\bm{L}$ are linearly independent, $\bm{q}_{j}^{T}\bm{\xi}_{n}=0$ for all $j=1,2,\ldots,m-1$. For $\bm{q}_{m}^{T}$, the last row of $\bm{Q}$, it must be a unit vector orthogonal to $\{\bm{q}_{1}^{T},\ldots,\bm{q}_{m-1}^{T}\}$. For the decomposition in (17) to be a unique characterization, we constrain the choice of the factorization such that $\bm{q}_{m}^{T}$ is selected from the remaining 1-dimensional space spanned by $\bm{\xi}_{n}$. This leads to $\bm{q}_{m}^{T}=\pm\bm{\xi}_{n}$. By convention, we choose the positive sign for the set. Thus, any $\bm{W}\in\widetilde{\mathcal{O}}_{m,n}$ can be decomposed into the desired form. ∎

The following example with $m=2$, $n=3$ demonstrates the construction of a matrix in $\widetilde{\mathcal{O}}_{m,n}$ as defined in (17).

Example 2**.**

Let $m=2$, $n=3$. Then, by Lemma 2, the Cholesky factor is given as

[TABLE]

Consider the semi-orthogonal matrix

[TABLE]

Note that the last row of $\bm{Q}$ is $\bm{\xi}^{T}_{n}$. Then, $\bm{L}\,\bm{Q}+\frac{1}{\sqrt{mn}}\bm{J}_{m\times n}$ is exactly the same as the matrix in Example 1, which belongs to $\widetilde{\mathcal{O}}_{m,n}$.

Computational Complexity

While both Algorithm 1 and Algorithm 2 provide methods for generating a matrix $\bm{W}\in\widetilde{\mathcal{O}}_{m,n}$, the approach in Algorithm 2, derived from Theorem 2, is significantly more efficient. Algorithm 1 requires two separate QR decompositions: one for an $m\times m$ matrix and another for an $n\times m$ matrix, resulting in a computational complexity of about $O(m^{3}+nm^{2})$. In contrast, Algorithm 2 first computes the fixed lower-triangular matrix $\bm{L}$ using a direct closed-form expression from Lemma 2, which costs only $O(m^{2})$. It then performs a single QR decomposition on an $n\times m$ matrix to build the semi-orthogonal matrix $\bm{Q}$. This approach reduces the total complexity to $O(m^{2}+nm^{2})$, removing the $O(m^{3})$ term entirely. When the ratio $m/n$ is close to $1$, this eliminated cost becomes a substantial portion of the total computation, potentially cutting the initialization time in half.

V Statistical properties

In this section, the statistical characteristics of matrices in $\widetilde{\mathcal{O}}_{m,n}$ are examined, together with their implications for signal propagation in ReLU‑based neural networks. The discussion begins with the linear transformation $f(\bm{x})=\bm{W}\bm{x}$ for a weight matrix $\bm{W}\in\widetilde{\mathcal{O}}_{m,n}$, with emphasis on its influence over key input-output statistics-most notably covariance structure and distributional behavior in high‑dimensional regimes. The scope is then extended to multi‑layer ReLU networks. A heuristic mean-field framework [12, 28] is adopted to trace the evolution of activation statistics across layers.

V-A Linear Transform Behavior

Before delving into the implications of the proposed initialization scheme for ReLU-based neural networks, we first examine a fundamental aspect of forward propagation: the statistical behavior of the linear transformation $f(\bm{x})=\bm{W}\bm{x}$ for the proposed weight matrix $\bm{W}\in\widetilde{\mathcal{O}}_{m,n}$. We aim to characterize the statistical properties of $\bm{W}$ and its effect on the image of $\bm{W}$, as it plays a central role in shaping the dynamics of signal propagation through fully connected layers.

We denote by $\mathbb{E}[\bm{x}]$ the expectation of the random vector $\bm{x}$, and by $\mathrm{Cov}[\bm{x}]$ its covariance matrix (for a scalar random variable $X$, $\mathrm{Cov}[X]$ simply becomes $\mathrm{Var}[X]$). We write $X\stackrel{{\scriptstyle d}}{{=}}Y$ to denote that $X$ and $Y$ are equal in distribution.

Proposition 3**.**

Let $\bm{W}\in\widetilde{\mathcal{O}}_{m,n}$ be given. For a random vector $\bm{x}\in\mathbb{R}^{n}$ with $\mathbb{E}[\bm{x}]=\mu\bm{1}_{n}$, $\operatorname{Cov}[\bm{x}]=\sigma^{2}\bm{I}_{n},$ ($\mu\in\mathbb{R}$, $\sigma>0$), it holds that

[TABLE]

Proof.

By the linearity of expectation and the decomposition $\bm{x}=(\bm{x}-\mu\bm{1}_{n})+\mu\bm{1}_{n}$, one can obtain that

[TABLE]

The last equality holds from Lemma 1. And by the linearity of expectation, it holds that

[TABLE]

Since $\bm{W}$ is a semi-orthogonal matrix, the last equality holds. ∎

Proposition 3 shows that when the input vector $\bm{x}\in\mathbb{R}^{n}$ has a constant mean $\mathbb{E}[\bm{x}]=\mu\bm{1}_{n}$ and isotropic covariance $\mathrm{Cov}[\bm{x}]=\sigma^{2}\bm{I}_{n}$, the output $\bm{W}\bm{x}$ under a semi-orthogonal transformation $\bm{W}\in\widetilde{\mathcal{O}}_{m,n}$ also has an explicit mean and covariance:

[TABLE]

This result indicates that the transformation preserves isotropy and uniformly rescales the mean, which is beneficial for stabilizing signal propagation in deep networks.

If we further assume that $\bm{x}$ follows a multivariate normal distribution $\mathcal{N}(\mu\bm{1}_{n},\sigma^{2}\bm{I}_{n})$, then the transformed vector $\bm{W}\bm{x}$ also follows a multivariate normal distribution with the same mean and covariance as derived above, due to the affine invariance of Gaussian distributions. However, this Gaussian assumption may not hold in real-world data. Empirical evidence shows that many practical datasets exhibit significant departures from normality, such as skewness, heavy tails, or outliers [29, 30]. Consequently, the exact Gaussian form of the output $\bm{y}$ is not guaranteed in general.

Let $\mathbb{P}(\mathcal{E})$ denote the probability of an event $\mathcal{E}$ under a given probability space. Recall that the cumulative distribution function (CDF) of the standard normal distribution is defined by

[TABLE]

We write $X_{n}\overset{d}{\longrightarrow}X$ to denote convergence in distribution of a sequence of random variables $\{X_{n}\}$ to a random variable $X$, that is,

[TABLE]

for all $z\in\mathbb{R}$ at which the cumulative distribution function of $X$ is continuous.

And we write $X_{n}\overset{p}{\longrightarrow}X$ to denote convergence in probability, for $\epsilon>0$,

[TABLE]

Recall that the Stiefel manifold is defined as $\mathcal{O}_{n,m}=\{\bm{Q}\in\mathbb{R}^{n\times m}|\;\bm{Q}^{T}\bm{Q}=\bm{I}_{m}\}$. The elements of the Stiefel manifold are sometimes called $m$–frames in $\mathbb{R}^{n}$. Let $\bm{A}\in\mathbb{R}^{n\times m}$ be a standard Gaussian random matrix. We perform the QR decomposition of $\bm{A}$, that is, $\bm{A}=\bm{Q}\bm{R}$, where $\bm{Q}\in O_{n,m}$ lies on the Stiefel manifold and $\bm{R}\in\mathbb{R}^{m\times m}$ is an upper triangular matrix with nonnegative diagonal entries. It is well known that, under this construction, $\bm{Q}$ is uniformly distributed on the Stiefel manifold $\mathcal{O}_{n,m}$ [31].

Using the result of [32], the following lemma can be obtained.

Lemma 3**.**

Let $\widehat{\bm{Q}}=[\widehat{Q}_{ij}]\in\mathcal{O}_{n,m-1}$ be a uniformly distributed random $(m-1)$–frame in the subspace $\bm{\xi}_{n}^{\perp}$. Then for any $C>2$,

[TABLE]

as $n\rightarrow\infty$.

Note that $\bm{Q}=[\widehat{\bm{Q}}\;\xi_{n}]^{T}\in\mathcal{O}_{m,n}$, where $\widehat{\bm{Q}}$ in Lemma 3 coincides with the matrix $\bm{Q}$ appearing in line 11 of Algorithm 2. Using Lemma 3, the following theorem holds.

Theorem 3**.**

Let $\bm{W}=[W_{ij}]$ be a random matrix as $\bm{W}=\bm{L}\bm{Q}+\frac{1}{\sqrt{mn}}\bm{J}_{m\times n}$, where $\bm{L}$ is the fixed matrix from Lemma 2 and $\bm{Q}=[\widehat{\bm{Q}}\;\xi_{n}]^{T}\in\mathcal{O}_{m,n}$ where $\widehat{\bm{Q}}$ is a uniformly distributed random $(m-1)$–frame in the subspace $\bm{\xi}_{n}^{\perp}$. Then for any $C>2$ and fixed $m$,

[TABLE]

as $n\rightarrow\infty$.

Proof.

Since $\widehat{\bm{Q}}$ is uniformly distributed as an $(m-1)$–frame in the subspace $\bm{\xi}_{n}^{\perp}$, it is invariant under left multiplication by any $\bm{R}\in\mathcal{O}_{m-1}$, i.e., $\bm{R}\widehat{\bm{Q}}^{T}\stackrel{{\scriptstyle d}}{{=}}\widehat{\bm{Q}}^{T}$. For any nonzero vector $\bm{v}\in\mathbb{R}^{m-1}$, let $\bm{R}\in\mathcal{O}_{m-1}$ be such that $\frac{\bm{v}}{\|\bm{v}\|}=\bm{R}^{T}\bm{e}_{1}$ , where $\bm{e}_{1}=[1~0~\cdots~0]^{T}$ is the first standard basis vector. Then

[TABLE]

where $\bm{q}_{1}$ is the first column of $\widehat{\bm{Q}}$. Now denote $\bm{L}=[\,\widehat{\bm{L}}\;\;\mathbf{0}\,]$, and denote by $\bm{l}_{i}^{T}$ the $i$-th row of $\widehat{\bm{L}}$. Using (20), $(\bm{L}\bm{Q})_{i}=\bm{l}_{i}^{T}\widehat{\bm{Q}}^{T}\stackrel{{\scriptstyle d}}{{=}}\ \|\bm{l}_{i}\|\bm{q}_{1}^{T},$ where $(\bm{L}\bm{Q})_{i}$ is the $i$-th row of $\bm{L}\bm{Q}$. Since $\|\bm{l}_{i}\|_{2}^{2}=(\bm{L}\bm{L}^{T})_{ii}=1-\tfrac{1}{m}$, it follows that

[TABLE]

For $i=1,\dots,m$,

[TABLE]

using Lemma 3. Then

[TABLE]

as $n\rightarrow\infty$. ∎

Note that $\bm{W}$ in Theorem 3 exactly matches $\bm{W}$ generated by Algorithm 2.

Lemma 4**.**

[33, Theorem 1.1]** Let $\bm{y}_{1},\ldots,\bm{y}_{n}$ be independent mean $\bm{0}$ random vectors in $\mathbb{R}^{m}$ such that $\sum_{i=1}^{n}\bm{y}_{i}$ has the identity covariance matrix. Let $\bm{g}$ be the standard Gaussian random vector in $\mathbb{R}^{m}$. Then for all measurable convex sets $\mathcal{A}\in\mathbb{R}^{m}$,

[TABLE]

Theorem 4**.**

Fix $m\in\mathbb{N}$. Let $\bm{x}=(x_{1},\dots,x_{n})\in\mathbb{R}^{n}$ be a random vector with independent components satisfying

[TABLE]

Let $\bm{W}=[W_{ij}]\in\widetilde{\mathcal{O}}_{m,n}$ be a random matrix as defined in Theorem 3, independent of $\bm{x}$. Then

[TABLE]

as $n\to\infty$.

Proof.

Let $\bm{z}=(z_{1},\ldots,z_{n})$ such that $z_{j}:=(x_{j}-\mu)/\sigma$ for all $j$. Then $\mathbb{E}[z_{j}]=0$, $\operatorname{Var}(z_{j})=1$, and $\beta_{3}:=\mathbb{E}[|z_{j}|^{3}]<\infty$. For each $j=1,\ldots,n$, let $\bm{w}_{j}\in\mathbb{R}^{m}$ be the $j$-th column vector of $\bm{W}$. For $j=1,\ldots,n$, define $\bm{y}_{j}:=\bm{w}_{j}z_{j}$. Since the $z_{j}$ are independent scalars, the vectors $\bm{y}_{1},\ldots,\bm{y}_{n}$ are independent and $\mathbb{E}[\bm{y}_{j}]=\bm{0}$ and

[TABLE]

Therefore, Lemma 4 applies to $\bm{y}_{j}$. For all measurable convex sets $\mathcal{A}\subset\mathbb{R}^{m}$,

[TABLE]

where $\bm{g}\sim\mathcal{N}(\bm{0},\bm{I}_{m})$.

Since $\|\bm{y}_{j}\|_{2}=|z_{j}|\|\bm{w}_{j}\|_{2}$, $\sum_{i=1}^{n}\mathbb{E}[\|\bm{y}_{i}\|_{2}^{3}]=\beta_{3}\sum_{i=1}^{n}\mathbb{E}[\|\bm{w}_{j}\|_{2}^{3}]$. And $\sum_{i=1}^{n}\|\bm{w}_{j}\|_{2}^{3}$ can be bounded as follows:

[TABLE]

It implies that $\sum_{i=1}^{n}\mathbb{E}[\|\bm{w}_{j}\|_{2}^{3}]\leq m^{3/2}\mathbb{E}[\max_{i,j}|W_{ij}|]$. Since $\max_{i,j}|W_{ij}|<1$ and $\max_{i,j}|W_{ij}|\overset{p}{\longrightarrow}0$, by uniform integrability and Theorem 2.20 in [34], we have $\mathbb{E}[\max_{i,j}|W_{ij}|]\rightarrow 0$. Let $K=(42m^{1/4}+16)\beta_{3}m^{3/2}$, then

[TABLE]

as $n\rightarrow\infty$. For any $\bm{a}\in\mathbb{R}^{m}$ and $t\in\mathbb{R}$, consider the closed half-space $\mathcal{H}_{\bm{a},t}:=\{\bm{x}\in\mathbb{R}^{m}|\ \bm{a}^{T}\bm{x}\leq t\}$. Since $\mathcal{H}_{\bm{a},t}$ is convex,

[TABLE]

By the Cramér-Wold theorem [34, Proposition 2.17], as $n\rightarrow\infty$,

[TABLE]

∎

To empirically validate this result, we consider a non-Gaussian input distribution. Fig. 2 shows the empirical distribution of the first component of $\bm{W}\bm{x}$, where $\bm{W}\in\widetilde{\mathcal{O}}_{m,n}$ and each entry of $\bm{x}\in\mathbb{R}^{n}$ is independently drawn from a chi-squared distribution with three degrees of freedom, $\chi^{2}(3)$, for $m=n\in\{5,10,100\}$. A total of 1,000,000 input vectors were sampled for each value of $n$ to estimate the output distribution. The original input distribution is plotted as a dashed line, and a Gaussian distribution with the same mean and variance is plotted as a dash-dotted black line for reference. As observed in the figure, the transformed distribution becomes increasingly Gaussian as $n$ increases.

V-B Propagation in ReLU Networks

Deep ReLU networks suffer from variance collapse and dying ReLU problems [35]. These problems are related to weight initialization and activation functions.

Consider a standard FFNN architecture in which each layer $\ell$ is recursively defined by:

[TABLE]

As stated in the preliminary section, we assume a non-increasing architecture where the layer widths $N_{\ell}$ for $\ell=0,\ldots,L$ satisfy $N_{0}\geq N_{1}\geq\cdots\geq N_{L}$.

Given the central limit behavior established in Theorem 4, it is reasonable to assume that the pre-activation vectors $\bm{y}^{(\ell)}$ in deep ReLU networks follow approximately Gaussian distributions in high-dimensional settings.

The single-variable rectified Gaussian distribution is defined as $\max(0,x)\sim\mathcal{N}^{\text{R}}(\mu,\sigma^{2})$ where $x\sim\mathcal{N}(\mu,\sigma^{2})$ [36]. This definition naturally extends to the multivariate. If $\bm{x}\sim\mathcal{N}(\bm{\mu},\bm{\Sigma})$ and $\bm{z}=\max(0,\bm{x})$ where the maximum is applied elementwise, then $\bm{z}$ is said to follow a multivariate rectified Gaussian distribution denotes $\bm{z}\sim\mathcal{N}^{\mathrm{R}}(\bm{\mu},\bm{\Sigma})$. Note that $\mathbb{E}[\bm{z}]$ and $\operatorname{Cov}(\bm{z})$ are not $\bm{\mu}$ and $\bm{\Sigma}$, respectively. Using the result of [36, 37], the following lemma provides the closed-form expressions of $\mathbb{E}[\bm{z}]$ and $\operatorname{Cov}(\bm{z})$ when $\bm{z}\sim\mathcal{N}^{R}(\mu\bm{1}_{n},\sigma^{2}\bm{I}_{n})$.

Lemma 5** ([36]).**

Let $\bm{z}\sim\mathcal{N}^{R}(\mu\bm{1}_{n},\sigma^{2}\bm{I}_{n})$. Then

[TABLE]

where $\alpha=\frac{\mu}{\sigma}$ and $\phi$ and $\Phi$ are the probability density function and cumulative distribution function of the Gaussian distribution, respectively.

However, since $\lim\limits_{\alpha\to\infty}\Phi(\alpha)=1$ and $\lim\limits_{\alpha\to\infty}\phi(\alpha)=0$, a sufficiently large $\alpha=\mu/\sigma$ implies that the statistical properties of the ReLU output (e.g. mean and variance) become increasingly close to those of the Gaussian input.

Remark 1**.**

Moreover, for each coordinate $z_{i}$ of $\bm{z}\sim\mathcal{N}(\mu\mathbf{1},\sigma^{2}I)$, \mathbb{P}(z_{i}>0)\;=\;\Phi\bigl{(}\tfrac{\mu}{\sigma}\bigr{)}\;=\;\Phi(\alpha).

Therefore, as $\alpha$ becomes large, most of the probability mass of the input lies in the positive region. In this regime, the ReLU function behaves nearly as the identity map, and its effect on the input distribution becomes negligible. Consequently, the rectified Gaussian output is nearly indistinguishable from the original Gaussian input in functional terms.

In our setting, the proposed initialization ensures that the pre-activation at layer $\ell$ follows $\mathcal{N}\left(\sqrt{\frac{N_{\ell-1}}{N_{\ell}}}\,\mu,\sigma^{2}\right)$, leading to the rectification parameter:

[TABLE]

In deep networks with non-increasing widths ($N_{\ell}<N_{\ell-1}$), the scaling factor $\sqrt{\frac{N_{\ell-1}}{N_{\ell}}}$ contributes to increasing $\alpha$. Furthermore, since ReLU increases the mean and decreases the variance of Gaussian inputs, these effects accumulate across layers. As a result, $\alpha$ grows with depth, pushing the activation distribution further into the positive domain. This behavior alleviates the dying ReLU problem and helps preserve stable signal propagation.

We formalize this observation in the following proposition, which characterizes the depth-wise behavior of activation statistics under our initialization.

Proposition 4**.**

Suppose that the input feature $\bm{x}^{(0)}=(x^{(0)}_{1},\dots,x^{(0)}_{N_{0}})\in\mathbb{R}^{N_{0}}$ satisfies the following condition:

[TABLE]

$\text{for all }j=1,\dots,N_{0},$* with input dimension $N_{0}$ is large. Assume that, for each layer $\ell$, the weight matrix $\bm{W}^{(\ell)}\in\widetilde{\mathcal{O}}_{N_{\ell},N_{\ell-1}}$ is a random matrix as defined in Theorem 3, and that the bias satisfies $\bm{b}^{(\ell)}=\bm{0}$. Then for each layer $\ell=1,\ldots,L$, the distribution of the post-activation vector $\bm{x}^{(\ell)}$ is approximately distributed as*

[TABLE]

where $\mu_{\ell}$ and $\sigma_{\ell}^{2}$ denote the mean and variance of the components of $\bm{x}^{(\ell)}$, respectively.

Proof.

We proceed by induction on the layer index $\ell$.

First, we prove for $l=1$. Since $N_{0}$ is large, Theorem 4 implies that the pre-activation vector $\bm{y}^{(1)}=\bm{W}^{(1)}\bm{x}^{(0)}$ is approximately Gaussian:

[TABLE]

Applying the ReLU activation yields $\bm{x}^{(1)}=\text{ReLU}(\bm{y}^{(1)})$, which is approximately distributed as a rectified Gaussian $\mathcal{N}^{R}(\mu_{1},\sigma_{1}^{2})$. This establishes the base case.

Now we prove for $\ell\geq 2$. Assume the proposition holds for layer $\ell-1$. We show that $\bm{y}^{(\ell)}=\bm{W}^{(\ell)}\bm{x}^{(\ell-1)}$ remains approximately Gaussian. There are two cases for the input vector $\bm{x}^{(\ell-1)}$.

(i) If $N_{\ell-1}$ is large, then by the multivariate central limit theorem, the linear transformation $\bm{y}^{(\ell)}=\bm{W}^{(\ell)}\bm{x}^{(\ell-1)}$ is approximately Gaussian, even though $\bm{x}^{(\ell-1)}$ follows a rectified Gaussian distribution.

(ii) If $N_{\ell-1}$ is not large, we exploit the fact that the rectification parameter $\alpha_{\ell-1}=\sqrt{\frac{N_{\ell-2}}{N_{\ell-1}}}\frac{\mu_{\ell-2}}{\sigma_{\ell-2}}$ grows with depth. By the inductive hypothesis, $\bm{x}^{(\ell-1)}$ follows a rectified Gaussian law. For a sufficiently large $\alpha_{\ell-1}$, the rectified Gaussian distribution of $x_{j}^{(\ell-1)}$ becomes almost indistinguishable from a true Gaussian distribution, as the mass at zero vanishes.

[TABLE]

Therefore, $\bm{y}^{(\ell)}=\bm{W}^{(\ell)}\bm{x}^{(\ell-1)}$ is a linear transformation of an approximately Gaussian vector, which results in $\bm{y}^{(\ell)}$ also being approximately Gaussian.

In both cases, the pre-activation $\bm{y}^{(\ell)}$ is approximately Gaussian. Thus, $\bm{x}^{(\ell)}=\text{ReLU}(\bm{y}^{(\ell)})$ is by definition distributed as $\mathcal{N}^{R}(\mu_{\ell},\sigma_{\ell}^{2})$. By the principle of induction, the proposition holds for all $\ell=1,\dots,L$. ∎

For a fixed network architecture, ensuring a large rectification parameter, $\alpha_{\ell}=\sqrt{\frac{N_{\ell-1}}{N_{\ell}}}\frac{\mu_{\ell-1}}{\sigma_{\ell-1}}$, involves either increasing the mean $\mu_{\ell-1}$ or decreasing the variance $\sigma_{\ell-1}$. However, reducing the variance is an undesirable strategy, as it would cause activations to become near-constant. Therefore, the viable approach is to control the mean. This mean shift is effective because as $\alpha$ increases, $\Phi(\alpha)$ approaches 1 while the PDF $\phi(\alpha)$ approaches 0. A key implication is that for inputs with a large $\alpha$, the statistical properties of the signal are almost perfectly preserved after passing through the ReLU activation. This mean shift mitigates the variance reduction by creating a large rectification parameter, and thus these stable properties can be preserved throughout all subsequent layers.

Building on this result, we now examine how our proposed method mitigates common challenges in deep ReLU networks during initialization. ReLU is known to suffer from three major issues [8, 38]: (a) the dying ReLU phenomenon, where neurons become inactive due to persistent negative pre-activations; (b) the variance reduction, where ReLU truncation reduces the dynamic range of signals; and (c) the gradient vanishing problem, especially in deeper networks with improper scaling.

Dying ReLU

Under the proposed initialization, Proposition 4 yields $\bm{x}^{(\ell)}\sim\mathcal{N}^{R}(\mu_{\ell}\mathbf{1},\,\sigma_{\ell}^{2}I)$ with nondecreasing $\mu_{\ell}$ and nonincreasing $\sigma_{\ell}$, so the rectification parameter $\alpha_{\ell}=\mu_{\ell}/\sigma_{\ell}$ increases with depth. By Remark 1, $\mathbb{P}(x_{j}^{(\ell)}>0)=\Phi(\alpha_{\ell})$ and hence the inactivity probability $1-\Phi(\alpha_{\ell})$ decays rapidly, mitigating the dying ReLU phenomenon in deep networks.

Variance reduction

For $\alpha_{\ell-1}=\mu_{\ell-1}/\sigma_{\ell-1}$, Lemma 5 gives

[TABLE]

Since $\Phi(a)\to 1$ and $\phi(a)\to 0$ as $a\to\infty$, it follows that $\frac{\sigma_{\ell}^{2}}{\sigma_{\ell-1}^{2}}\to 1$ as $\alpha_{\ell-1}\to\infty$. Hence, for any $\epsilon\in(0,1)$ there exists $\delta>0$ such that if $\alpha_{\ell-1}\geq\delta$ then $(1-\epsilon)\,\sigma_{\ell-1}^{2}\;\leq\;\sigma_{\ell}^{2}\;\leq\;\sigma_{\ell-1}^{2}$. In particular, once $\alpha_{\ell-1}\geq\delta$, the layerwise variance is maintained within a one-sided $(1-\epsilon)$ lower bound, preventing cumulative variance decay and preserving activation dynamic range.

Vanishing gradients

Let $\bm{D}^{(\ell)}$ be a diagonal matrix whose $i$th diagonal element is $1$ if $x^{(\ell)}_{i}>0$ and [math] otherwise. Backpropagation satisfies

[TABLE]

where $\bm{\delta}^{(\ell)}$ is the gradient of the loss with respect to the pre-activation at layer $\ell$. Under the proposed initialization, $\alpha_{\ell}=\mu_{\ell}/\sigma_{\ell}$ increases with depth, hence $\mathbb{P}(x^{(\ell)}_{j}>0)=\Phi(\alpha_{\ell})\to 1$ and $\mathbb{E}[\bm{D}^{(\ell)}]=\Phi(\alpha_{\ell})\bm{I}\to\bm{I}$. Together with the semi-orthogonality of $\bm{W}^{(\ell)}$ and the variance reduction above, this limits layerwise Jacobian contraction and prevents exponential gradient decay in deep ReLU networks.

VI Experimental Results

This section empirically demonstrates the effectiveness of the proposed initialization on deep ReLU networks. The proposed initialization is evaluated against several widely adopted initialization methods—Lee, Orthogonal, Xavier, He, and Random—on the MNIST and Fashion-MNIST datasets. All experiments employ an FFNN with ReLU activations in all hidden layers, implemented in PyTorch. The networks are trained using cross-entropy loss and optimized with the Adam algorithm. The learning rate is scaled with depth as $\text{lr}=0.001/\sqrt{\text{depth}}$, following established practice [39], except for Lee et al. [14], where a fixed learning rate of 0.001 is used to match prior work. All models are trained with a batch size of 256 for 100 epochs, and classification accuracy on the test split is reported as the evaluation metric.

VI-A Effect of Network Depth on Performance

To investigate how well the proposed initialization supports deep ReLU networks, experiments are conducted by varying the network depth. Specifically, a fully connected FFNN with architecture:

[TABLE]

are trained on the full MNIST and the full Fashion-MNIST datasets, respectively. The output layer uses softmax for classification. TABLE I (a) and (b) present the classification accuracy obtained on MNIST and Fashion-MNIST datasets across varying network depths and initialization methods.

First, for $D=10\sim 30$, all initialization schemes perform comparably well, achieving high accuracy on both datasets. However, as depth increases, significant differences emerge. Notably, the proposed initialization consistently maintains stable and high accuracy even at depths up to $D=100$, achieving $96.98\%$ on MNIST and $87.70\%$ on Fashion-MNIST. In contrast, commonly used initializations such as Xavier, Orthogonal, and Random exhibit severe performance degradation beyond $D=50$, with accuracy collapsing to near random levels. While Lee initialization and He initialization demonstrate better robustness due to their ReLU-aware design, they do not entirely prevent performance degradation in very deep networks. In particular, the Lee initialization exhibits noticeable fluctuations as the depth increases, indicating unstable training dynamics at large depths. In comparison, He initialization exhibits a more gradual decline in accuracy as depth increases.

This stability is a result of the fact that the proposed initialization method maintains the signal variance stably even in deep networks and effectively mitigates the dying ReLU phenomenon, as theoretically proven in the previous section. On the other hand, other initialization methods are analyzed to fail to prevent variance collapse as depth increases, leading to learning failure.

VI-B Effectiveness across ReLU Family Activation

To test the robustness of our initialization, we examine its interaction with other popular activation functions from the ReLU family. To examine the interaction between activation variants and initialization schemes, we replaced the ReLU activation in each hidden layer with LeakyReLU [7], PReLU [8], ELU [9], and SELU [40], respectively. The network architecture and training protocol were identical to those in the previous experiments, and the same set of initialization methods was evaluated. We tested depths of $D=50$ and $D=100$ to highlight failures due to poor initialization, as confirmed by the results in TABLE I.

The proposed initialization achieves high accuracy on MNIST and Fashion-MNIST for LeakyReLU, PReLU, ELU, SELU, and standard ReLU, with a drop of less than $1\%$ when transitioning from $D=50$ to $D=100$. This consistent performance across activation variants demonstrates the robustness and low variance of the proposed method. In contrast, Lee initialization performs competitively at $D=50$, often matching or slightly exceeding the proposed method. However, its stability degrades significantly at $D=100$, where it fails to train under several activation functions, such as LeakyReLU and PReLU, on MNIST and Fashion-MNIST. He initialization, by comparison, consistently supports learning across depths and activation types but exhibits lower overall accuracy than the proposed method, especially on MNIST at depth 100. Orthogonal, Xavier, and Random initializations already exhibited poor performance under standard ReLU and continue to fail at deeper depths across most activation functions. Nonetheless, a limited recovery is observed when combined with ELU or SELU, particularly on the Fashion-MNIST dataset.

VI-C Few-Shot Learning

To evaluate the proposed initialization method in data-scarce scenarios, we conducted a series of few-shot learning experiments. These scenarios are particularly challenging because the model must learn to generalize from a tiny number of training examples per class, making the initial state of the network weights a critical factor for successful training and convergence.

The experiments were conducted with the number of training data $k=1,2,4,8$ using the same network architecture and comparison methods as in the previous section. Each configuration was repeated 50 times to ensure statistical significance.

As shown in TABLE III, in deep networks (50 or 100 layers), most conventional initialization methods, such as Xavier and He, suffer from a sharp performance drop or fail to train. In contrast, the proposed method exhibits a significantly smaller drop in accuracy with increasing depth, demonstrating superior performance compared to any other method.

Fig. 3 provides a representation of these trends for the $4$-shot learning. The shaded area represents the standard deviation of the accuracy of each method. The plot of the proposed method shows a remarkably stable and slow decrease in accuracy as the network gets deeper, demonstrating its robustness. In contrast, the curves for Xavier, He, Orthogonal, and Random initialization, respectively, exhibit a sharp decline in performance, with accuracy collapsing after $30$-$40$ layers, and fail to learn in deeper architectures. Lee initialization also maintains a respectable performance but shows a steeper decline and a larger standard deviation compared to the proposed method.

VI-D Tabular Data

To assess generalization beyond image domains, several tabular datasets were used, covering both classification (Adult [41], Cancer [42], Pima [43], Ionosphere [44], Wine [45]) and regression (Diabetes [46]) tasks (TABLE IV). As tabular inputs require explicit scaling, all features were standardized to unit variance ($\sigma^{2}=1$), and the mean $\mu$ was shifted to control $\alpha=\mu/\sigma^{2}$. Experiments used $\alpha\in\{-2,\,0,\,2,\,50\}$. Models were fully connected FFNN(architecture details in TABLE IV). Training settings followed those of the prior experiments, and evaluation used accuracy for classification and root mean squared error (RMSE) for regression.

The tabular results are summarized in TABLE V. Across both classification and regression datasets, the proposed initialization consistently shows strong performance among the compared methods (TABLE V). Performance is sensitive to the rectification parameter: negative shifts ($\alpha_{0}=-2$) and excessively large shifts ($\alpha_{0}=50$) degrade accuracy and increase RMSE, whereas setting ($\alpha_{0}=2$) yields the best results. When $\alpha_{0}<0$, many first-layer pre-activations fall below zero, so ReLU turns those units off and gradients weaken. When $\alpha_{0}$ is very large, pre-activations stay strongly positive, so ReLU behaves nearly linearly, reducing the model’s nonlinearity. Hence, selecting a moderate $\alpha_{0}$ (here, $\alpha_{0}=2$) is critical for stable and accurate training on a deep ReLU network.

VII Conclusion

This paper provided a weight initialization for deep ReLU networks by formulating and solving an optimization problem on the Stiefel manifold. Unlike conventional approaches such as Xavier, He, and orthogonal initialization, the proposed method simultaneously preserves semi-orthogonality. It maximizes alignment with the all-ones vector, thereby directly regulating pre-activation statistics. We derived closed-form characterizations of the optimal solution set, developed efficient sampling algorithms, and established theoretical guarantees on variance preservation, mean calibration, and asymptotic distributional behavior.

Theoretical analysis demonstrated that 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. Empirical validation confirmed these advantages on MNIST, Fashion-MNIST, tabular classification/regression tasks, and few-shot learning scenarios. 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.

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. Future work may extend this framework to convolutional architectures, attention-based models, and adaptive optimization strategies on manifolds, further bridging rigorous mathematical design with practical advances in deep learning.