Weight Initialization › Initialization for Different Activations
Language AI HandbookAll chapters Weight Initialization › Initialization for Different ActivationsOne section · Sections in this chapter8 / 19 ↓
Reading tools Definitions & inline links
Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.
Article links
Make inline references clickable
We've seen Xavier for symmetric activations and He for ReLU. The deep learning toolkit includes many more activations, each with its own variance characteristics. The general principle remains: analyze how the activation distorts variance and adjust the initialization accordingly. This generalization is what PyTorch's nn.init.calculate_gain() function systematizes. This provides gain factors for each activation that can be plugged into the LeCun base formula.
The following table summarizes the recommended initialization for common activations:
| Activation | Recommended Init | Weight Variance |
|---|---|---|
| Linear | Xavier | |
| Tanh | Xavier | |
| Sigmoid | Xavier | |
| ReLU | He | |
| Leaky ReLU () | He (adjusted) | |
| SELU | LeCun | |
| GELU | He (approximate) |
LeCun Initialization and SELU
The SELU activation (Scaled Exponential Linear Unit) is a special case that deserves its own mention. SELU is designed to be self-normalizing: under the right conditions, activations in a SELU network converge to a fixed-point distribution with mean zero and unit variance regardless of network depth. The activation function itself applies a specific scaling and shift to maintain this fixed-point property.
For SELU to work correctly, the weights must be initialized with LeCun initialization, . This is the simplest of all the initialization schemes and the theoretical precursor to Xavier. The SELU paper by Klambauer et al. (2017) provides rigorous proofs that the combination of LeCun initialization, SELU activations, and specific rescaling factors produces self-normalizing behavior, eliminating the need for batch normalization in networks designed around this activation function.
Leaky ReLU
For Leaky ReLU with negative slope (typically 0.01 or 0.2), negative inputs are scaled by rather than zeroed. This preserves some variance from the negative half of the distribution. The corrected variance formula is:
where:
- : the negative slope of Leaky ReLU (the multiplier applied to negative pre-activations)
- : a correction factor accounting for variance from both positive inputs (coefficient 1, so contribution is ) and negative inputs (coefficient , so contribution is )
- : fan-in (number of input connections)
When (standard ReLU), the factor recovers He initialization. When (linear activation), the factor gives , approaching LeCun initialization for the forward-only case. For typical Leaky ReLU with , the correction factor is barely distinguishable from He initialization, which is why many practitioners simply use He for Leaky ReLU without the explicit correction.
GELU and Modern Activations
GELU (Gaussian Error Linear Unit), used in BERT and GPT, applies a smooth, probabilistic gating that scales each input by the probability that a standard normal random variable is less than the input. Like ReLU, it zeros or nearly zeros a fraction of inputs, reducing variance by roughly half. He initialization works well for GELU networks in practice, as the variance reduction is similar in magnitude to standard ReLU.
Swish, another smooth activation used in some architectures, has a similar profile to GELU. Both can be approximated as having a gain factor of , identical to ReLU, making He initialization a reasonable default. The key insight is that any activation function that zeros or heavily suppresses roughly half its inputs requires the variance formula to compensate, while activations that are closer to linear around zero (tanh, sigmoid) do not need this correction factor.
def init_weights(
shape, activation="relu", mode="fan_in", distribution="normal"
):
"""
Initialize weights based on activation function.
Parameters:
- shape: (fan_in, fan_out)
- activation: 'linear', 'tanh', 'sigmoid', 'relu', 'leaky_relu', 'selu', 'gelu'
- mode: 'fan_in', 'fan_out', or 'fan_avg'
- distribution: 'normal' or 'uniform'
"""
fan_in, fan_out = shape
if mode == "fan_in":
fan = fan_in
elif mode == "fan_out":
fan = fan_out
else: # fan_avg
fan = (fan_in + fan_out) / 2
# Gain factors derived from each activation's variance behavior
gain = {
"linear": 1.0,
"tanh": 5 / 3, # Empirical correction for tanh nonlinearity
"sigmoid": 1.0,
"relu": np.sqrt(2), # Compensates for variance halving
"leaky_relu": np.sqrt(2 / (1 + 0.01**2)), # Negative slope alpha=0.01
"selu": 1.0, # LeCun initialization: variance = 1/n
"gelu": np.sqrt(2), # Similar to ReLU in practice
}.get(activation, 1.0)
std = gain / np.sqrt(fan)
if distribution == "normal":
return np.random.randn(*shape) * std
else:
limit = np.sqrt(3) * std
return np.random.uniform(-limit, limit, shape)Initialization statistics for different activations (shape 256x256): ------------------------------------------------------------ relu : std = 0.0883, range = [-0.354, 0.357] tanh : std = 0.1036, range = [-0.428, 0.423] leaky_relu : std = 0.0885, range = [-0.391, 0.404] selu : std = 0.0628, range = [-0.264, 0.260] gelu : std = 0.0883, range = [-0.339, 0.423]
ReLU and GELU initializations have the largest standard deviation because they use the gain factor to compensate for variance halving. Tanh uses a gain of approximately 5/3. SELU uses LeCun initialization with a gain of 1, giving variance .