How Controlling the Variance can Improve Training Stability of Sparsely Activated DNNs and CNNs

Technical Summary: Controlling the Variance to Improve Training Stability of Sparsely Activated DNNs and CNNs

Problem Statement
Training Deep Neural Networks (DNNs) and Convolutional Neural Networks (CNNs) with activation functions that induce high degrees of sparsity (e.g., 85% to 90%) presents significant initialization challenges. Traditional Edge-of-Chaos (EoC) theory, which characterizes intermediate layers as Gaussian processes to ensure stable information propagation, typically assumes activations are approximately linear around the origin. In this regime, EoC theory often advocates for the fixed-point variance of the pre-activation distribution, $q^*$, to converge to zero or a small non-zero value to achieve dynamic isometry.

However, for sparsity-inducing activations where $\phi(z) = 0$ for $|z| \leq \tau$ (a region near the origin), the linearity assumption fails. Previous work [1] demonstrated that standard EoC initialization strategies often lead to training instability in these settings. The instability arises because the empirical variance of hidden layers, $\tilde{q}^{(\ell)}$, fluctuates stochastically around the theoretical fixed point $q^*$. In the absence of linearity near the origin, these fluctuations can cause the multiplicative perturbation of errors (governed by $\chi_1$) to accumulate exponentially with depth, leading to exploding gradients or a collapse of the network's expressivity.

Methodology
The authors extend EoC theory to the setting of sparsity-inducing activations, specifically analyzing the exemplar function $CReLU_{\tau,m}$ (and related functions like $CST_{\tau,m}$, $HardShrink$, and $Sawtooth$). The core methodology involves treating the fixed-point variance $q^*$ not merely as a derived parameter but as a tunable hyperparameter.

The theoretical analysis focuses on three control metrics that determine training stability in the large-width limit:

1. Symmetry of the Variance Map: The authors analyze the second derivative of the variance mapping function, $V''(q^*)$. A large $V''(q^*)$ indicates a loss of symmetry around the fixed point, causing the empirical variance $\tilde{q}^{(\ell)}$ to drift asymmetrically.
2. Finite Dimensional Correction: Using an infinite series expansion for finite-width networks, the authors derive bounds for the next-leading-order correction term, $\tilde{q}^{(1)}$. This term represents the deviation from the infinite-width Gaussian limit.
3. Sensitivity of $\chi_1$: The authors examine the derivative of the multiplicative perturbation factor, $\chi'_{1,\phi}(q^*)$, with respect to the variance $q$. High sensitivity implies that small stochastic fluctuations in layer variance lead to large changes in error propagation.

The proposed initialization strategy involves explicitly increasing $q^*$ (e.g., from 1 to 2 or 3) and adjusting the activation parameters ($\tau, m$) and weight/bias variances ($\sigma_w^2, \sigma_b^2$) to satisfy the EoC conditions ($\chi_1(q^*) = 1$ and $q^* = V(q^*)$) at this new, higher fixed point.

Key Contributions

1. New Phenomenon Identification: The paper identifies a counter-intuitive phenomenon where, for activations zeroed near the origin, increasing the fixed-point variance $q^*$ improves training stability. This contrasts with traditional EoC approaches that seek $q^* \to 0$.
2. Theoretical Characterization: The authors prove that increasing $q^*$ simultaneously:
   • Improves the symmetry of the variance map $V(q)$ around the fixed point (reducing $V''(q^*)$).
   • Tightens the bound on the finite-dimensional correction term $\tilde{q}^{(1)}$, making the finite-width behavior closer to the infinite-width Gaussian limit.
   • Reduces the sensitivity of the error propagation factor $\chi_1$ to fluctuations in layer variance (reducing $\chi'_{1,\phi}(q^*)$).
3. Practical Initialization Scheme: A concrete initialization procedure is provided for sparsifying activations. This scheme allows for the training of networks with up to 90% sparsity in hidden layers, a regime previously difficult to stabilize.

Experimental Results
The authors validate their theory through experiments on DNNs (100 layers, width 300) and CNNs (50 layers, 300 channels) trained on MNIST, CIFAR10, and Tiny ImageNet using $CReLU_{\tau,m}$ activations.

• Stability and Accuracy: For high sparsity levels ($s=0.85$ and $s=0.9$), networks initialized with $q^*=1$ often failed to train or exhibited high sensitivity to the choice of the slope parameter $V'(q^*)$. Increasing $q^*$ to 2 or 3 significantly reduced this sensitivity and restored training stability.
• Performance: In cases where $q^*=1$ resulted in low accuracy (e.g., 11% for DNNs at $s=0.9$), increasing $q^*$ to 3 recovered accuracy to near-full levels (e.g., 89%).
• Convergence: Training loss and validation accuracy curves showed faster convergence for higher $q^*$ values, particularly in high-sparsity regimes.
• Efficiency: The paper includes a custom CUDA kernel implementation demonstrating that the induced sparsity yields computational speedups (up to ~2.9x) and energy reductions, validating the practical utility of the sparsity.

Significance and Claims
The paper claims that the fixed-point variance $q^*$ is a critical, yet previously underexplored, parameter for controlling the training dynamics of sparsely activated networks. By revisiting the EoC framework for non-linear activations that are zero near the origin, the authors demonstrate that a simple adjustment—increasing $q^*$—can overcome the instability inherent in high-sparsity regimes.

The authors modestly note that their results rely on the infinite-width limit assumption and are primarily demonstrated on deep proxy architectures (DNNs and CNNs) rather than optimized production models. They explicitly state that extensions to Transformers and State-Space models are not covered due to architectural differences (e.g., skip connections negating activation sparsity) and separate lines of analysis. However, the work provides a theoretical foundation and empirical proof-of-concept for enabling stable training of networks with extreme activation sparsity (up to 90%), potentially unlocking significant computational efficiencies without sacrificing model performance.