一文搞懂深度网络初始化（Xavier and Kaiming i

最近时不时就有网友发私信让我帮忙debug程序，对于这些来信，我的回复通常都是一句话：“先跑通论文作者的开源代码，在此基础上再逐步修改数据集和模型。”

你或许会觉得我在摆谱，净说些高大上的政治正确的空话，但这还真不是。

和其他的软件程序不同，神经网络是个系统工程，数据、参数、模型内部结构、训练策略、学习率等等，这些因素不管哪一部分出错，它都不会报错，只是会输出一些不是你想要的结果而已。

参数初始化就是这么一个容易被忽视的重要因素，因为不仅使用者对其重要性缺乏概念，而且这些操作都被TF、pytorch这些框架封装了，你可能不知道的是，糟糕的参数初始化是会阻碍复杂非线性系统的训练的。

Xavier Initialization

早期的参数初始化方法普遍是将数据和参数normalize为高斯分布（均值0方差1），但随着神经网络深度的增加，这方法并不能解决梯度消失问题。

因此，他提出了Xavier初始化：bias初始化为0，为Normalize后的参数乘以一个rescale系数：1/，n是输入参数的个数。

公式的推导过程大致如下：

• 因为E（期望）等于均值，而输入数据（x）和参数（W）的均值都是0，因此，
• 又因为x和W恒等分布（方差都是1），因此，
• 我们的目标是，因此，

如果上述这段公式你看晕了，也没关系，只要记住结果就好。

接下来，我们要做实验来验证Xavier的洞见。

def linear(x, w, b): return x @ w + b

def relu(x): return x.clamp_min(0.)

nh = 50
W1 = torch.randn(784, nh)
b1 = torch.zeros(nh)
W2 = torch.randn(nh, 1)
b2 = torch.zeros(1)

z1 = linear(x_train, W1, b1)
print(z1.mean(), z1.std())

tensor(-0.8809) tensor(26.9281)

这是个简单的线性回归模型：，(W1, b1)和(W2, b2)分别是隐层和输出层的参数，W1/W2初始化为高斯分布，b1/b2初始为0。果然，第一个linear层的输出值（z1）的均值和标准差就已经发生了很大的变化。如果后续使用sigmoid作为激活函数，那梯度消失就会很明显。

现在我们按照Xavier的方法来初始化参数：

W1 = torch.randn(784, nh) * math.sqrt(1 / 784)
b1 = torch.zeros(nh)
W2 = torch.randn(nh, 1) * math.sqrt(1 / nh)
b2 = torch.zeros(1)

z1 = linear(x_train, W1, b1)
print(z1.mean(), z1.std())

tensor(0.1031) tensor(0.9458)

a1 = relu(z1)
a1.mean(), a1.std()

(tensor(0.4272), tensor(0.5915))

参数经过Xavier初始化后，linear层的输出值的分布没有大的变化（），依旧接近高斯分布，但是好景不长，relu的激活值分布就开始跑偏了（）。

Kaiming Initialization

Xavier初始化的问题在于，它只适用于线性激活函数，但实际上，对于深层神经网络来说，线性激活函数是没有价值，神经网络需要非线性激活函数来构建复杂的非线性系统。今天的神经网络普遍使用relu激活函数。

W1 = torch.randn(784, nh) * math.sqrt(2 / 784)
b1 = torch.zeros(nh)
W2 = torch.randn(nh, 1) * math.sqrt(2 / nh)
b2 = torch.zeros(1)

z1 = linear(x_train, W1, b1)
a1 = relu(z1)
a1.mean(), a1.std()

(tensor(0.4553), tensor(0.7339))

可以看到，Kaiming初始化的表现要优于Xavier初始化，relu之后的输出值标准差还有0.7339（浮动可以达到0.8+）。

实际上，Kaiming初始化已经被Pytorch用作默认的参数初始化函数。

import torch.nn.init as init

W1 = torch.zeros(784, nh)
b1 = torch.zeros(nh)
W2 = torch.zeros(nh, 1)
b2 = torch.zeros(1)

init.kaiming_normal_(W1, mode='fan_out', nonlinearity='relu')
init.kaiming_normal_(W2, mode='fan_out')
z1 = linear(x_train, W1, b1)
a1 = relu(z1)
print("layer1: ", a1.mean(), a1.std())
z2 = linear(a1, W2, b2)

layer1:  tensor(0.5583) tensor(0.8157)
tensor(1.1784) tensor(1.3209)

现在，方差的问题已经解决了，接下来就是均值不为0的问题。因为在x轴上平移data并不会影响data的方差，因此，如果把relu的激活值左移5，结果会如何？

def linear(x, w, b):
  return x @ w + b

def relu(x):
  return x.clamp_min(0.) - 0.5

def model(x):
  x = relu(linear(x, W1, b1))
  print("layer1: ", x.mean(), x.std())
  x = relu(linear(x, W2, b2))
  print("layer2: ", x.mean(), x.std())
  x = linear(x, W3, b3)
  print("layer3: ", x.mean(), x.std())
  return x

nh = [100, 50]
W1 = torch.zeros(784, nh[0])
b1 = torch.zeros(nh[0])
W2 = torch.zeros(nh[0], nh[1])
b2 = torch.zeros(nh[1])
W3 = torch.zeros(nh[1], 1)
b3 = torch.zeros(1)

init.kaiming_normal_(W1, mode='fan_out')
init.kaiming_normal_(W2, mode='fan_out')
init.kaiming_normal_(W3, mode='fan_out')
_ = model(x_train)

layer1:  tensor(0.0383) tensor(0.7993)
layer2:  tensor(0.0075) tensor(0.7048)
layer3:  tensor(-0.2149) tensor(0.4493)

结果出乎意料的好，这个三层的模型在没有添加batchnorm的情况下，每层的输入值和输出值都接近高斯分布，虽然数据方差是会逐层递减，但相比normalize初始化和Xavier初始化要好很多。

最后，因为Kaiming初始化是pytorch的默认初始化函数，因此我又用pytorch提供的nn.Linear()和nn.Relu()来构建相同的模型对比测试，结果是大跌眼镜。

class Model(nn.Module):
  def __init__(self):
    super().__init__()
    self.lin1 = nn.Linear(784, nh[0])
    self.lin2 = nn.Linear(nh[0], nh[1])
    self.lin3 = nn.Linear(nh[1], 1)
    self.relu = nn.ReLU()
  
  def forward(self, x):
    x = self.relu(self.lin1(x))
    print("layer 1: ", x.mean().item(), x.std().item())
    x = self.relu(self.lin2(x))
    print("layer 2: ", x.mean().item(), x.std().item())
    x = self.relu(self.lin3(x))
    print("layer 3: ", x.mean().item(), x.std().item())
    return x

m = Model()
_ = m(x_train)

layer 1:  0.2270725518465042 0.32707411050796
layer 2:  0.033514849841594696 0.23475737869739532
layer 3:  0.013271240517497063 0.09185370802879333

可以看到，第三层的输出已经均值为0、方差为0。去看nn.Linear()类的代码时会看到，它在做初始化时会传入参数a=math.sqrt(5)。我们知道，当输入为负数时，leaky relu的梯度为，，参数a就是这个。虽然kaiming_uniform_()的默认网络要使用的激活函数是leaky relu，但a默认值为0，此时leaky relu就等于relu。但现在数据存在负数，因此，mean相比relu模型更接近于0，甚至E(x) > 0的假设都不成立了，因此，rescale系数就不准确了，nn.Linear()才会有这样的表现。

    def reset_parameters(self):
        init.kaiming_uniform_(self.weight, a=math.sqrt(5))

END

本文通过Xavier和Kaiming初始化来展现了参数初始化的重要性，因为糟糕的初始化容易让神经网络陷入梯度消失的陷阱中。

References