Source-linked AI summary
On weight initialization in deep neural networks
Siddharth Krishna Kumar
TL;DR
Deep-network training is sensitive to weight initialization, while existing theory largely concerns linear activations. The paper develops initialization theory for nonlinear activations, deriving a general strategy for functions differentiable at 0 and a RELU-specific strategy, and explains why Xavier performs poorly with RELU.
Problem
Weight initialization theory has primarily focused on linear activations, despite initialization errors causing shrinking variances in deeper layers and slower convergence.
Method
The paper derives a general initialization formula for activation functions differentiable at 0 and rigorously analyzes RELU, including the He initialization.
Results
The analysis provides theoretical insights into why Xavier initialization fails with RELU and demonstrates that nonlinearities produce different activation dynamics.
Takeaways & Limitations
Proper weight initialization must account for the activation function’s nonlinearities rather than relying universally on linear-activation results.
Takeaways & Limitations
The analysis of activation functions not differentiable at 0 is largely incomplete beyond the RELU case.
Abstract
from arXiv · showhide
A proper initialization of the weights in a neural network is critical to its convergence. Current insights into weight initialization come primarily from linear activation functions. In this paper, I develop a theory for weight initializations with non-linear activations. First, I derive a general weight initialization strategy for any neural network using activation functions differentiable at 0. Next, I derive the weight initialization strategy for the Rectified Linear Unit (RELU), and provide theoretical insights into why the Xavier initialization is a poor choice with RELU activations. My analysis provides a clear demonstration of the role of non-linearities in determining the proper weight initializations.
1 Introduction
The paper addresses how to initialize weights so layerwise output variances remain stable, extending initialization theory beyond linear activations to nonlinear functions and explaining RELU-specific behavior.
- Arbitrary initialization can shrink variances in deeper layers, slowing backpropagation and overall convergence.
- The paper asks how to choose v2 for weights sampled from N(0, v2) so output variances remain approximately equal across layers.
- For linear activations, Xavier initialization sets v2 = 1/N, where N is the number of inputs to a layer.
- For RELU activations, He and colleagues propose v2 = 2/N after showing a 30-layer network converges with He initialization but not Xavier initialization.
- The paper generalizes initialization theory to nonlinear activations, proves the He initialization for RELU, and provides insights into the contrasting convergence behavior.
2 The setup
The setup models layerwise propagation with independently initialized Gaussian weights and seeks a variance choice that preserves first-pass activation variance across layers. It establishes independence and approximate normality results before analyzing activation-specific nonlinearities.
- The setup: The network uses Gaussian weights with mean 0 and variance v2, while first-layer inputs are independent and identically distributed with mean 0 and variance 1.
- The setup: The target is to choose v2 so the first forward pass maintains s1^2 ≈ s2^2 ... ≈ sM^2 = 1 across layers.
- The setup: During the first iteration, each layer’s weights are independent of all inputs to that layer, including products of input variables.
- The setup: The resulting variance expression holds for arbitrary covariance structures among inputs to the layer.
- The setup: Assuming a large hidden-layer width N, the pre-activation ym(i) is approximately normally distributed for all layers and units.
- The setup: This approximate normality follows from summing independent or weakly dependent contributions and applying the Central Limit Theorem when N is large.
- The setup: The later analysis splits according to whether the activation function is differentiable at 0.
3 Activation functions differentiable at 0
The paper derives a general initialization strategy for differentiable nonlinear activations by approximating layerwise moments, then applies it to tanh and sigmoid networks. For tanh, the analysis recovers Xavier initialization and explains changing saturation across depth.
- General strategy: A Taylor expansion around zero approximates the activation’s behavior when g(x) is differentiable at 0, enabling moment-based initialization analysis.The derivation assumes higher-order terms can be ignored.
- General strategy: The expected input value of successive layers has little dependence on the preceding layer’s input moments, and recursively follows the activation value at zero.This result applies to layers beyond the first.
- General strategy: Equation (12) provides a general weight initialization strategy for arbitrary differentiable activation functions.The paper applies this strategy to hyperbolic tangent and sigmoid activations.
- Hyperbolic tangent: For hyperbolic tangent, g(0) = 0 and g′(0) = 1, so the derived initialization becomes the Xavier initialization.The paper identifies the resulting expression as precisely Xavier initialization.
- Hyperbolic tangent: Under hyperbolic tangent activation, shallower layers are more saturated while saturation decreases toward deeper layers.The analysis associates saturation with activations concentrated near -1 and +1, and reduced saturation with deeper-layer distributions.
- Sigmoid: The paper compares the derived sigmoid initialization with Xavier initialization using a 10-layer network on CIFAR-10 and top-5 validation accuracy.The dataset contains 50,000 training examples and 10,000 validation examples.
4 Activation functions not differentiable at 0
Because RELU is not differentiable at 0, the general differentiable-activation analysis does not apply, so the paper focuses on deriving RELU-specific initialization results.
- 4 Activation functions not differentiable at 0: RELU requires a separate analysis because it is not differentiable at 0.The paper states that the preceding results cannot be used to compute the optimal variance for RELU.
- 4 Activation functions not differentiable at 0: The RELU analysis targets the mean and variance of the next-layer activations, xm+1(i) = max(0, ym(i)).
- 4 Activation functions not differentiable at 0: Maintaining variance across iterations is imposed as the condition for selecting the RELU initialization.
- 4 Activation functions not differentiable at 0: The derived mean is approximately 0.7 across layers during the first forward pass, consistent with the result reported by He and colleagues.
To converge or not to converge, that is the question.
The paper analyzes why Xavier initialization can fail with RELU networks by deriving moment recursions and comparing their predictions with simulations. It finds that deeper layers receive exponentially smaller input variances, providing a possible explanation for depth-dependent convergence differences.
- To converge or not to converge, that is the question.: The analysis computes central moments layer by layer under Xavier initialization to explain differing convergence behavior in RELU networks.
- To converge or not to converge, that is the question.: The predicted moment approximations are remarkably accurate compared with the simulation experiments in Figure 4.
- To converge or not to converge, that is the question.: Xavier initialization produces exponentially smaller input variances in deeper layers, worsening as network depth increases.
- To converge or not to converge, that is the question.: The variance at the 30th layer is 3 × 10^-3 times the variance at the 22nd layer, possibly explaining their different convergence outcomes.
5 Conclusion
The paper presents a general framework for initializing weights with nonlinear activations, gives a rigorous RELU-specific result, and explains why Xavier initialization can fail with RELU. The conclusion also acknowledges that the non-differentiable-function analysis remains incomplete.
- 5 Conclusion: The paper provides a general formula for ideal weight initialization when activation functions are differentiable at 0.
- 5 Conclusion: For non-differentiable activations, the paper focuses on RELU and provides a rigorous proof of the He initialization.
- 5 Conclusion: The paper shows why Xavier initialization fails to work with RELU activation.
- 5 Conclusion: The analysis illustrates that introducing nonlinearities can produce drastically different system dynamics.
- 5 Conclusion: The treatment of non-differentiable activation functions is described as largely incomplete, motivating a more general analysis.