SparseOpt: Addressing Normalization-induced Gradient Skew in Sparse Training

Dynamic Sparse Training (DST) methods train neural networks by maintaining sparsity while dynamically adapting the network topology. Despite the promise of reduced computation, DST methods converge significantly slower than dense training, often requiring comparable training time to achieve similar accuracy. We demonstrate both analytically and empirically that Batch Normalization (BN) adversely affects sparse training, and propose SparseOpt — a sparsity-aware optimizer — to address this. Experiments on ResNet models across CIFAR-100 and ImageNet demonstrate consistently faster convergence and improved generalization with our proposed method. Our work highlights the limitations of current normalization layers in sparse training and provides the first systematic study of the interaction between Batch Normalization, sparse layers, and DST, taking a significant step toward making DST practically competitive with dense training.

Sparse training aims to train neural networks using only a small fraction of the connections, instead of using every parameter in the network. Since fewer connections are active, sparse training has the potential to significantly reduce the computational and memory cost of training large models. However, in practice, sparse neural networks are much harder to train than standard dense networks. Existing sparse training methods often converge substantially slower and may require significantly longer training schedules to match the accuracy of dense models, limiting their practical efficiency benefits. To stabilize and accelerate the training of deep neural networks, modern architectures almost universally rely on normalization layers such as Batch Normalization (BN). Intuitively, BN keeps activations and gradients well-scaled during training, which helps optimization remain stable and enables faster convergence. While BN has been extensively studied in dense neural networks, prior work largely assumed that it behaves similarly in sparse networks. In this work, we show both theoretically and empirically that this assumption does not hold. Specifically, we demonstrate that BN interacts poorly with heterogeneous sparse connectivity, leading to neuron-dependent gradient scaling that distorts optimization dynamics and destabilizes Dynamic Sparse Training (DST). Motivated by this observation, we propose a simple sparsity-aware correction method, SparseOpt, that accounts for the gradient imbalance introduced by BN. Our method improves optimization stability, accelerates convergence, and consistently improves sparse training performance across multiple datasets and architectures.