【知识】关于参数初始化

Dymay
Dymay
发布于 2024-09-27 / 30 阅读
0
0

【知识】关于参数初始化

简单整理了一下参数初始化相关的知识。

PyTorch 中调整神经网络的参数初始化策略,可以使用 torch.nn.init 模块中的各种初始化方法。

示例代码

import torch
import torch.nn as nn
import torch.nn.init as init
​
class CustomNet(nn.Module):
    def __init__(self):
        super(CustomNet, self).__init__()
        self.conv1 = nn.Conv2d(in_channels=3, out_channels=16, kernel_size=3, stride=1, padding=1)
        self.fc1 = nn.Linear(in_features=16*32*32, out_features=128)
        self.fc2 = nn.Linear(in_features=128, out_features=10)
        
        # 调用初始化方法
        self._initialize_weights()
​
    def _initialize_weights(self):
        # 使用 Xavier 均匀分布初始化卷积层,gain=1.0
        init.xavier_uniform_(self.conv1.weight, gain=1.0)
        init.constant_(self.conv1.bias, 0)
        
        # 使用 Xavier 正态分布初始化全连接层,gain=sqrt(2) 适用于 ReLU 激活函数
        init.xavier_normal_(self.fc1.weight, gain=torch.nn.init.calculate_gain('relu'))
        init.constant_(self.fc1.bias, 0)
        
        init.xavier_normal_(self.fc2.weight, gain=torch.nn.init.calculate_gain('relu'))
        init.constant_(self.fc2.bias, 0)
​
    def forward(self, x):
        x = self.conv1(x)
        x = torch.relu(x)
        x = x.view(x.size(0), -1)  # 展平
        x = self.fc1(x)
        x = torch.relu(x)
        x = self.fc2(x)
        return x
​
# 创建网络实例
net = CustomNet()
print(net(torch.randn(4, 3, 32, 32)))

这个示例中定义了一个简单的卷积神经网络,并在 _initialize_weights 方法中使用了不同的初始化策略。

一些常用的初始化方法:

基于固定方差的参数初始化

一种最简单的随机初始化方法是从一个固定均值(通常为 0)和方差 𝜎^2 的分布中采样来生成参数的初始值。

  1. 正态分布初始化 (init.normal_):

    torch.nn.init.normal_(tensor, mean=0.0, std=0.02)

    即使用一个高斯分布 \mathcal{N}(0, 𝜎^2) 对每个参数进行随机初始化

  2. 均匀分布初始化 (init.uniform_):

    torch.nn.init.uniform_(tensor, a=0.0, b=1.0)

    即在一个给定的区间 [a, b] 内采用均匀分布来初始化参数。随机变量 𝑥 在区间 [𝑎, 𝑏] 内均匀分布则其方差为

    \mathrm{var}(x)=\frac{(b-a)^2}{12}

    因此,若使用区间为 [−𝑟, 𝑟] 的均匀分布来采样,并满足 \text{var}(𝑥) = 𝜎^2 时,则 𝑟 的取值为

    r=\sqrt{3\sigma^2}

    在基于固定方差的随机初始化方法中,比较关键的是如何设置方差 𝜎^2,如果参数范围取的太小,一是会导致神经元的输出过小,经过多层之后信号就慢慢消失了;二是还会使得 Sigmoid 型激活函数丢失非线性的能力。

    如果参数范围取的太大,会导致输入状态过。对于 Sigmoid 型激活函数来说,激活值变得饱和,梯度接近于 0,从而导致梯度消失问题。 为了降低固定方差对网络性能以及优化效率的影响,基于固定方差的随机初始化方法一般需要配合逐层归一化(Layer-wise Normalization)来使用。

基于方差缩放的参数初始化

初始化一个深度网络时,为了缓解梯度消失或爆炸问题,我们尽可能保持每个神经元的输入和输出的方差一致,根据神经元的连接数量来自适应地调整初始化分布的方差,这类方法称为方差缩放(Variance Scaling)

  1. Xavier 均匀分布初始化 (init.xavier_uniform_):

    torch.nn.init.xavier_uniform_(tensor, gain=1.0)
  2. Xavier 正态分布初始化 (init.xavier_normal_):

    torch.nn.init.xavier_normal_(tensor, gain=1.0)

    假设在一个神经网络中,第 𝑙 层的一个神经元 𝑎^{(𝑙)},其接收前一层的 𝑀_{𝑙−1} 个神经元的输出 𝑎^{(𝑙-1)},1\le i\le M_{l-1}

    a^{(l)}=f\Big(\sum_{i=1}^{M_{l-1}}w_i^{(l)}a_i^{(l-1)}\Big)

    假设激活函数 f(\cdot) 为恒等函数,w_i^{(l)}a_i^{(l-1)} 的均值都为0并相互独立,则

    \mathbb{E}[a^{(l)}]=\mathbb{E}\Big[\sum_{i=1}^{M_{l-1}}w_i^{(l)}a_i^{(l-1)}\Big]=\sum_{i=1}^{M_{l-1}}\mathbb{E}[w_i^{(l)}]\mathbb{E}[a_i^{(l-1)}]=0\\

    \begin{aligned} \operatorname{var}(a^{(l)})& =\mathrm{var}\Big(\sum_{i=1}^{M_{l-1}}w_i^{(l)}a_i^{(l-1)}\Big) \\ &=\sum_{i=1}^{M_{l-1}}\operatorname{var}(w_{i}^{(l)})\operatorname{var}(a_{i}^{(l-1)}) \\ &=M_{l-1}\operatorname{var}(w_{i}^{(l)})\operatorname{var}(a_{i}^{(l-1)}). \end{aligned}

    也就是说,输入信号在经过该神经元后均值不变,方差被放大或缩小了 𝑀_{𝑙−1} \operatorname{var}(𝑤_i^{(𝑙)}) 倍。为了使得在经过多层网络后,信号不被过分放大或过分减弱,要尽可能保持每个神经元的输入和输出的方差一致,即:

    𝑀_{𝑙−1} \operatorname{var}(𝑤_i^{(𝑙)})=1\Rightarrow\operatorname{var}(𝑤_i^{(𝑙)})=\frac{1}{M_{l-1}}

    同理,为了使得在反向传播中,误差信号也不被放大或缩小:

    𝑀_{𝑙} \operatorname{var}(𝑤_i^{(𝑙)})=1\Rightarrow\operatorname{var}(𝑤_i^{(𝑙)})=\frac{1}{M_{l}}

    折中一下,取调和均值:

    \mathrm{var}(w_i^{(l)})=\frac{2}{M_{l-1}+M_l}

    在计算出参数的理想方差后,可以通过高斯分布或均匀分布来随机初始化参数。 这种根据每层的神经元数量来自动计算初始化参数方差的方法称为Xavier 初始化(也称Glorot 初始化)。
    此处的推导为 神经网络与深度学习 (nndl.github.io) 书中过程,更一般情况的推导可以参考 深度学习中的初始化方法(Initialization) - 郑之杰的个人网站 (0809zheng.github.io)

    虽然在 Xavier 初始化中假设激活函数为恒等函数,但是 Xavier 初始化也适用于 Logistic 函数和 Tanh 函数。这是因为神经元的参数和输入的绝对值通常比较小,处于激活函数的线性区间。这时 Logistic 函数和 Tanh 函数可以近似为线性函数,由于 Logistic 函数在线性区间的斜率约为 0.25,因此其参数初始化的方差约为 \displaystyle 16 × \frac{2}{M_{l-1}+M_l},在实际应用中,使用 Logistic 函数或 Tanh 函数的神经层通常将方差 \displaystyle \frac{2}{M_{l-1}+M_l} 乘以一个缩放因子 𝜌,即对应于代码中的 gain 参数。

    可以注意到代码对应的实现是标准差,即:

    \text{std} = \text{gain} \times \sqrt{\frac{2}{\text{fan\_in} + \text{fan\_out}}}

    使用 torch.nn.init.calculate_gain 计算给定非线性激活函数的增益因子 gain 即可。使用方法如下:

    torch.nn.init.xavier_normal_(tensor, gain=torch.nn.init.calculate_gain(nonlinearity, param=None))
    • nonlinearity:字符串,指定所使用的非线性激活函数类型,例如 'relu''leaky_relu''tanh' 等。

    • param:可选参数,某些非线性激活函数(如 Leaky ReLU)需要额外的参数,例如负斜率。

    Xavier 初始化在理论设计时是针对 Sigmoid、Tanh 等激活函数,对于 ReLU 及其变体激活函数则更适用下面的 Kaiming 初始化(也称 He 初始化)。

  3. Kaiming 均匀分布初始化 (init.kaiming_uniform_):

    torch.nn.init.kaiming_uniform_(tensor, a=0, mode='fan_in', nonlinearity='leaky_relu')
  4. Kaiming 正态分布初始化 (init.kaiming_normal_):

    torch.nn.init.kaiming_normal_(tensor, a=0, mode='fan_in', nonlinearity='leaky_relu')

    当第 𝑙 层神经元使用 ReLU 激活函数时,通常有一半的神经元输出为 0,因此其分布的方差也近似为使用恒等函数时的一半。这样,只考虑前向传播时,参数 𝑤_i^{(𝑙)} 的理想方差为

    \operatorname{var}(𝑤_i^{(𝑙)})=\frac{2}{M_{l-1}}
    • a:这是负斜率参数,通常用于 Leaky ReLU 激活函数。如果使用普通的 ReLU 激活函数,a 通常设为 0。如果使用 Leaky ReLU 激活函数,a 应设为 Leaky ReLU 的负斜率值。

    • mode有两个选项:

      • 'fan_in':保持前向传播过程中输入的方差一致,适用于 ReLU 激活函数。

      • 'fan_out':保持反向传播过程中梯度的方差一致,适用于输出层。

    • nonlinearity:指定所使用的激活函数类型,通常为 'relu''leaky_relu'

    Kaiming 初始化总体上相较于Xavier 初始化效果更好,是 Pytorch 默认的参数初始化函数,基本不用改。

正交初始化

上面介绍的两种基于方差的初始化方法都是对权重矩阵中的每个参数进行独立采样。由于采样的随机性,采样出来的权重矩阵依然可能存在梯度消失或梯度爆炸问题。

  1. 正交初始化 (init.orthogonal_):

    torch.nn.init.orthogonal_(tensor, gain=1)

    假设一个 𝐿 层的等宽线性网络(激活函数为恒等函数,偏置初始化为 0)为

    y=W^{(L)}W^{(L-1)}\cdots W^{(1)}x

    其中 𝑾^{(𝑙)} ∈ ℝ^{𝑀×𝑀} (1 ≤ 𝑙 ≤ 𝐿) 为神经网络的第 𝑙 层权重矩阵。在反向传播中,误差项 𝛿^{(l)}(第 𝑙 层神经元对最终损失 的影响,也反映了最终损失对第 𝑙 层神经元的敏感程度)的反向传播公式为 𝛿^{(𝑙−1)} = (𝑾^{(𝑙)})^{\top}𝛿^{(𝑙)},为了避免梯度消失或梯度爆炸问题,我们希望误差项在反向传播中具有范数保持性(Norm-Preserving),即

    \|\delta^{(l-1)}\|^{2}=\|\delta^{(l)}\|^{2}=\|(W^{(l)})^{\top}\delta^{(l)}\|^{2}

    一种直接的方式是将 𝑾^{(𝑙)} 初始化为正交矩阵,即 𝑾^{(𝑙)}(𝑾^{(𝑙)})^{\top} =𝐼,这种方法称为正交初始化(Orthogonal Initialization)。正交初始化通常用在 RNN 中循环边的权重矩阵上。

    正交初始化的具体实现过程可以分为两步:

    • 用均值为 0、方差为 1 的高斯分布初始化一个矩阵;

    • 将这个矩阵用奇异值分解得到两个正交矩阵,并使用其中之一作为权重矩阵

    当在非线性神经网络中应用正交初始化时,通常需要将正交矩阵乘以一个缩放系数 𝜌,即对应于代码中的 gain 参数。比如当激活函数为 ReLU 时,激活函数在 0 附近的平均梯度可以近似为 0.5,为了保持范数不变,缩放系数 𝜌 可以设置为 \sqrt{2}。同样代码可以使用 torch.nn.init.calculate_gain 计算

其他初始化

  1. 稀疏初始化 (init.sparse_):

    torch.nn.init.sparse_(tensor, sparsity=0.1, std=0.01)

    sparsity 参数表示非零元素的比例,std 参数表示非零元素从正态分布中采样的标准差。

  2. 常数初始化 (init.constant_):

    torch.nn.init.constant_(tensor, val=0)
  3. 恒等初始化 (init.eye_init.dirac_):

    恒等初始化是指通过把神经网络的权重层初始化为一个单位矩阵(恒等变换),使得网络层的输出值与输出值相等。

    对于二维参数(如全连接层的权重参数),通过单位矩阵进行恒等初始化:

    torch.nn.init.eye__(tensor)

    对于高维参数(如卷积层的权重矩阵),通过Dirac-delta函数进行恒等初始化:

    torch.nn.init.dirac_(tensor, groups=1)

    恒等初始化具有动力等距(Dynamical Isometry)性质,使得神经网络具有稳定的信号传播以及梯度下降的行为。然而恒等初始化建立在各层的维度是相等的假设之上,在实际中这种假设有些过强。当各层的输入输出维度不相等时,可以把参数矩阵(非方阵)初始化为部分单位矩阵,对于行列中“超出”的部分补零即可。然而当使用部分单位矩阵在训练神经网络时,会出现训练衰减(Training Degeneracy)现象,即无论隐藏层维度有多高,部分的输入在激活函数阶段无法生效,导致神经网络的维度仅仅依赖于输入数据的维度,极大的限制了神经网络的表达能力。

  4. Zero初始化

    此部分参考:忘掉Xavier初始化吧!最强初始化方法ZerO来了 (aiorang.com)
    Code: GitHub - jiaweizzhao/ZerO-initialization

关于 Dynamical Isometry

dynamical isometry 的概念是用来表征雅可比矩阵与误差信号的乘积与等值算子的接近程度。也就是说雅可比矩阵(的乘积)的奇异值中越多接近1的奇异值越好。使用类似正交初始化的技术(奇异值接近1)可以使得权值结构有更好的 dynamical isometry。

具体的理论分析可以参考:

【L】a note on learning of deep architectures

【L】on dynamical isometry



评论