我们知道，神经网络大多靠随机梯度下降训练，参数从一组初值出发一步步更新。初值选得好不好，直接决定网络能不能训起来、收敛快不快。初值偏大，每层输出越乘越大，几层之后数值冲到溢出，这就是梯度爆炸。初值偏小，信号一层层衰减，传到后面几乎不动，梯度也跟着消失，这就是梯度消失。最极端的是把所有权重都设成同一个值，比如全 0，同一层每个神经元算出一样的输出，反向传回一样的梯度，更新完还是一模一样，网络永远学不出差别，等于白搭。

自然的想法是随机初始化，从高斯分布或均匀分布里抽权重。可随便定个方差又会引出新麻烦，每层输出的方差会跟着这一层输入神经元的个数一起放大或缩小，网络一深就失控。Xavier 初始化要做的，就是给权重方差定一个和层尺寸挂钩的公式，让信号和梯度在前向、反向传播时都别跑太远。接下来，我们先把方差在层间怎么传播讲清楚，再推导出高斯分布和均匀分布两种 Xavier 的具体取值，然后拿一个具体尺寸的网络把数字算一遍，最后落到 PyTorch 的现成接口。读之前你需要知道反向传播和方差的基本概念，手写过全连接层会更顺。

1. 方差传播

我们先看清楚，为什么权重随便选个方差，深层网络就会失控。这得回到一层里信号是怎么算的。

一层全连接，输出 y 是 n_in 个输入 x_i 乘权重 w_i 再加起来，即：y = Σ w_i x_i，偏置我们先忽略不看。这里要关心的是 y 的方差。假设权重和输入互相独立、各自均值为 0，把乘积方差展开，每一项 w_i x_i 的方差等于 E[w_i²]E[x_i²]，均值为 0 时就是 Var(w_i)·Var(x_i)。n_in 项加起来，就得到 Var(y) = n_in·Var(w)·Var(x)。

这个式子正好解释了失控。如果权重固定用某个方差，每过一层，输出方差就比输入多乘一个 n_in 倍的因子。n_in 大的层，比如第一层接 784 个像素，乘出来就是几百倍，方差瞬间被撑大。层数一多，要么层层放大溢出，要么层层缩小趋零。所以初始化合适的目标，就是让 Var(y) 和 Var(x) 大致相等，信号在每层之间不放大也不缩小。

2. 前后向折中

目标定成 Var(y)=Var(x)，代入上式就能解出权重该取多大方差。可这只照顾了前向，反向传播还有另一头要顾及。

令 Var(y)=Var(x)，从 Var(y)=n_in·Var(w)·Var(x) 两边约掉，得到 n_in·Var(w)=1，也就是前向要求 Var(w)=1/n_in。再看反向，传给前一层的梯度要对这一层 n_out 个输出节点各自求导再求和，推导和前向对称，得到反向要求 Var(w)=1/n_out。前向要 1/n_in，反向要 1/n_out，这两个数通常不一样，没法同时满足。Xavier 的折中办法是取两者平均，令 N=(n_in+n_out)/2，让 Var(w)=1/N，也就是 Var(w)=2/(n_in+n_out)。这样前向和反向的方差都不会离目标太远。

2.1 高斯分布取值

高斯版本就是让权重服从均值 0、方差 2/(n_in+n_out) 的正态分布，标准差是它的开方。我们拿一个具体网络算一遍，假设某层输入 784 个神经元、输出 256 个，那么 std = sqrt(2/(784+256)) = sqrt(2/1040) ≈ 0.044，权重就从 N(0, 0.044²) 里抽。对比一下，如果图省事直接用 N(0,1)，这一层输出方差会被放大 784 倍，数字立刻失控，这就是随便初始化的下场。

3. 均匀区间

高斯分布好推导，可实际框架里更常用均匀分布版本。区间端点怎么定，思路完全一样。

让 w 服从对称均匀分布 U(-a, a)。我们知道，均匀分布在长度为 L 的区间上的方差是 L²/12，这里区间长度是 2a，所以 Var(w)=(2a)²/12=a²/3。把它和 Xavier 要求的方差 2/(n_in+n_out) 划等号，即：a²/3=2/(n_in+n_out)，解出 a=sqrt(6/(n_in+n_out))，权重就在这个对称区间内均匀抽样。那个常被问起的除以 12 就来自这里，它是均匀分布方差积分出来的固定系数，不是随便定的。

回到刚才那层 784 进 256 出，端点 a=sqrt(6/1040)≈0.076，权重从 U(-0.076, 0.076) 里抽。两种分布的取值汇总在这张表里：

分布	抽样方式	关键参数
高斯 Xavier	正态分布 N(0, 方差)	std = sqrt(2/(n_in+n_out))
均匀 Xavier	对称均匀分布 U(-a, a)	a = sqrt(6/(n_in+n_out))

4. 接口用法

推导归推导，真写代码的时候不用自己开方，PyTorch 把这些都包好了。下面是 xavier_uniform_ 的实现，你对照上面的公式看，std 就是 sqrt(2/(fan_in+fan_out))，端点 a 再乘一个 sqrt(3)。

def xavier_uniform_(tensor: Tensor, gain: float = 1.) -> Tensor:
    fan_in, fan_out = _calculate_fan_in_and_fan_out(tensor)
    std = gain * math.sqrt(2.0 / float(fan_in + fan_out))
    a = math.sqrt(3.0) * std

    return _no_grad_uniform_(tensor, -a, a)

代码里多出来的 gain 是个补偿项。前面推导假设激活函数近似线性、关于原点对称，可不同激活函数会改变输出分布，PyTorch 就用一个增益系数 gain 来修正。它具体怎么算，写在 calculate_gain 里：

def calculate_gain(nonlinearity, param=None):
    """
    Return the recommended gain value for the given nonlinearity function.
    The values are as follows:
    ================= ====================================================
    nonlinearity      gain
    ================= ====================================================
    Linear / Identity :math:`1`
    Conv{1,2,3}D      :math:`1`
    Sigmoid           :math:`1`
    Tanh              :math:`\frac{5}{3}`
    ReLU              :math:`\sqrt{2}`
    Leaky Relu        :math:`\sqrt{\frac{2}{1 + \text{negative\_slope}^2}}`
    SELU              :math:`\frac{3}{4}`
    ================= ====================================================
    """

    linear_fns = ['linear', 'conv1d', 'conv2d', 'conv3d', 'conv_transpose1d', 'conv_transpose2d', 'conv_transpose3d']
    if nonlinearity in linear_fns or nonlinearity == 'sigmoid':
        return 1
    elif nonlinearity == 'tanh':
        return 5.0 / 3
    elif nonlinearity == 'relu':
        return math.sqrt(2.0)
    elif nonlinearity == 'leaky_relu':
        if param is None:
            negative_slope = 0.01
        elif not isinstance(param, bool) and isinstance(param, int) or isinstance(param, float):
            # True/False are instances of int, hence check above
            negative_slope = param
        else:
            raise ValueError("negative_slope {} not a valid number".format(param))
        return math.sqrt(2.0 / (1 + negative_slope ** 2))
    elif nonlinearity == 'selu':
        return 3.0 / 4  # Value found empirically (https://github.com/pytorch/pytorch/pull/50664)
    else:
        raise ValueError("Unsupported nonlinearity {}".format(nonlinearity))

把返回值列成表更直观：

激活函数	gain 推荐值
Linear / Identity / Conv	1
Sigmoid	1
Tanh	5/3 ≈ 1.667
ReLU	sqrt(2) ≈ 1.414
Leaky ReLU	sqrt(2/(1+negative_slope²))
SELU	3/4 = 0.75

实际调用时，传入激活函数的名字就能拿到对应的增益：

import torch.nn as nn


if __name__ == '__main__':

    print('relu:', nn.init.calculate_gain('relu'))
    print('sigmoid:', nn.init.calculate_gain('sigmoid'))
    print('tanh:', nn.init.calculate_gain('tanh'))
    print('leaky_relu:', nn.init.calculate_gain('leaky_relu'))
    print('selu:', nn.init.calculate_gain('selu'))

程序输出结果：

relu: 1.4142135623730951
sigmoid: 1
tanh: 1.6666666666666667
leaky_relu: 1.4141428569978354
selu: 0.75

5. 局限与延伸

不过 Xavier 不是万能的，它有一个明确前提，碰到 ReLU 这类激活函数就会露短板。

整个推导建立在一个假设上，激活函数近似线性、关于原点对称，输出均值为 0。tanh 接近这个条件，所以 Xavier 配 tanh 效果不错。但 ReLU 会把所有负值压成 0，经过它之后一半信号没了，方差直接腰斩，前面算好的平衡就被打破。2015 年 He 等人针对这个问题提出改进，把 ReLU 造成的方差损失算进去，把权重方差改成 2/n_in，这就是后来的 He 初始化，也叫 Kaiming 初始化。现在网络里用 ReLU 居多，选初始化时先看激活函数，tanh 配 Xavier，ReLU 配 He，别一套初始化用到所有网络上。