Weight Initialization › Initialization for Different Activations

Michael BrenndoerferApril 25, 20255 min read
Language AI HandbookAll chapters
Part 10 of 67Neural Network Foundations
Chapter 70Weight Initialization › Initialization for Different Activations
Reading section 8 of 19 · Initialization for Different Activations · 5 minSection 8 of 19 · 5 min

Book contents

Section 8 of 19 · 5 min read

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:

Recommended initialization schemes for common activation functions. The variance formulas assume sampling from a normal distribution; multiply variance by 3 if using a uniform distribution instead.
ActivationRecommended InitWeight Variance
LinearXavier2nin+nout\frac{2}{n_{\text{in}} + n_{\text{out}}}
TanhXavier2nin+nout\frac{2}{n_{\text{in}} + n_{\text{out}}}
SigmoidXavier2nin+nout\frac{2}{n_{\text{in}} + n_{\text{out}}}
ReLUHe2nin\frac{2}{n_{\text{in}}}
Leaky ReLU (α\alpha)He (adjusted)2(1+α2)nin\frac{2}{(1 + \alpha^2) n_{\text{in}}}
SELULeCun1nin\frac{1}{n_{\text{in}}}
GELUHe (approximate)2nin\frac{2}{n_{\text{in}}}

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, Var(w)=1/nin\text{Var}(w) = 1/n_{\text{in}}. 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 α\alpha (typically 0.01 or 0.2), negative inputs are scaled by α\alpha rather than zeroed. This preserves some variance from the negative half of the distribution. The corrected variance formula is:

Var(w)=2(1+α2)⋅nin\text{Var}(w) = \frac{2}{(1 + \alpha^2) \cdot n_{\text{in}}}

where:

  • α\alpha: the negative slope of Leaky ReLU (the multiplier applied to negative pre-activations)
  • (1+α2)(1 + \alpha^2): a correction factor accounting for variance from both positive inputs (coefficient 1, so contribution is 12=11^2 = 1) and negative inputs (coefficient α\alpha, so contribution is α2\alpha^2)
  • ninn_{\text{in}}: fan-in (number of input connections)

When α=0\alpha = 0 (standard ReLU), the factor (1+0)=1(1 + 0) = 1 recovers He initialization. When α=1\alpha = 1 (linear activation), the factor (1+1)=2(1 + 1) = 2 gives Var(w)=1/nin\text{Var}(w) = 1/n_{\text{in}}, approaching LeCun initialization for the forward-only case. For typical Leaky ReLU with α=0.01\alpha = 0.01, 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 2\sqrt{2}, 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 2/nin2/n_{\text{in}} variance formula to compensate, while activations that are closer to linear around zero (tanh, sigmoid) do not need this correction factor.

In[17]:
Code
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)
Out[18]:
Console
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 2\sqrt{2} 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 1/nin1/n_{\text{in}}.