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

Fuente: arXiv
Saved in:
Bibliographic Details
Main Authors: Lee, Hyungu, Kim, Taehyeong, Choi, Hayoung
Format: Preprint
Published: 2025
Subjects:
Online Access:
Tags: Add Tag
No Tags, Be the first to tag this record!
_version_ 1866916926027988992
author Lee, Hyungu
Kim, Taehyeong
Choi, Hayoung
author_facet Lee, Hyungu
Kim, Taehyeong
Choi, Hayoung
contents 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.
format Preprint
id arxiv_https___arxiv_org_abs_2509_00362
institution arXiv
publishDate 2025
record_format arxiv
spellingShingle Optimized Weight Initialization on the Stiefel Manifold for Deep ReLU Neural Networks
Lee, Hyungu
Kim, Taehyeong
Choi, Hayoung
Machine Learning
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.
title Optimized Weight Initialization on the Stiefel Manifold for Deep ReLU Neural Networks
topic Machine Learning
url https://arxiv.org/abs/2509.00362