训练深度神经网络,特别是现代LLM中使用的巨型Transformer架构,需要仔细关注诸多细节。在影响训练过程的诸多因素中,模型参数的初始状态,特别是其权重,尤为重要。简单地将权重初始化为零或从朴素分布中抽取的微小随机数,常导致严重的训练难题。恰当的初始化不只是一种启发式方法;它是使深度网络能够有效学习的基本要求。核心问题源于信号(包括前向传播时的激活和反向传播时的梯度)如何在网络层间传递。设想单层内的计算,通常涉及矩阵乘法后接非线性激活函数。当多层这样堆叠时,这些操作会依序重复。梯度消失在反向传播过程中,损失函数对早期层中某个权重的梯度,是使用链式法则计算的,即乘以所有后续层的梯度。如果这些梯度(或层变换的雅可比矩阵)的平均大小小于1,梯度信号在反向传播时会指数级减小。$$ \frac{\partial L}{\partial W_l} \propto \left( \prod_{k=l+1}^{N} J_k \right) \frac{\partial L}{\partial a_N} $$这里,$J_k$ 表示第 $k$ 层的雅可比矩阵。如果平均而言 $|J_k| < 1$,乘积项会随着层数 $N-l$ 的增加而迅速减小。这种现象被称为梯度消失问题,意味着早期层中的权重得到极其微小的更新。因此,这些层学习非常缓慢,有时甚至完全不学。这实际上阻止了网络学习依赖于其全部深度参数配合调整的复杂表示。历史上,这是训练深度网络的一大障碍,尤其是那些使用如sigmoid或tanh等激活函数的网络,这些函数在饱和区域的导数小于1。梯度爆炸反之,如果这些梯度(或雅可比矩阵)的平均大小大于1,梯度信号在反向通过各层时会指数级增长。$$ | \frac{\partial L}{\partial W_l} | \rightarrow \infty \quad \text{as } N-l \rightarrow \infty \quad \text{if } |J_k| > 1 $$这种梯度爆炸问题导致权重更新过大。过大的更新会使优化过程变得不稳定,可能导致越过最佳点或剧烈震荡。在极端情况下,梯度变得如此之大,以至于导致数值溢出(表示为NaN或Inf值),使训练过程停滞。尽管梯度裁剪等方法(将在第17章讨论)可以在训练期间缓解梯度爆炸,但恰当的初始化旨在从一开始就防止它们频繁出现。让我们看一个简化的模拟。设想一个具有10层的非常简单的线性网络。我们将权重初始化为略小或略大的值,并观察一个模拟梯度反向传播后的大小。import torch import torch.nn as nn import math # 模拟信号传播(简化) def check_grad_magnitude(init_scale, num_layers=10): """ 模拟通过线性层的反向梯度大小。 """ # 这里输入维度不那么重要,侧重于尺度 layer_dim = 100 layers = [] for _ in range(num_layers): layer = nn.Linear(layer_dim, layer_dim, bias=False) # 使用特定的标准差初始化权重 nn.init.normal_(layer.weight, mean=0.0, std=init_scale) layers.append(layer) network = nn.Sequential(*layers) # 模拟输入和输出梯度 x = torch.randn(1, layer_dim) # 假设输出梯度大小为1 output_grad = torch.ones(1, layer_dim) # 计算相对于输入的梯度,使用链式法则模拟 current_grad = output_grad # 手动反向传播以观察大小变化 with torch.no_grad(): for i in range(num_layers - 1, -1, -1): # 层输入梯度 = 层输出梯度 @ W的转置 current_grad = current_grad @ layers[i].weight # 为了演示,进行归一化以避免实际的爆炸/消失 # 在实际情况中,这种归一化不会逐层发生 # current_grad /= math.sqrt(layer_dim) # 示例归一化因子 # 计算相对于第一层输入的梯度大小 # 这模拟了梯度到达最早参数的情况 input_grad_norm = torch.norm(current_grad) return input_grad_norm.item() # 检查大小 num_layers = 10 small_init_scale = 0.05 # 可能导致梯度消失 large_init_scale = 0.5 # 可能导致梯度爆炸 ideal_init_scale = math.sqrt(1.0 / 100) # 类似于Xavier/He的缩放提示 grad_norm_small = check_grad_magnitude(small_init_scale, num_layers) grad_norm_large = check_grad_magnitude(large_init_scale, num_layers) grad_norm_ideal = check_grad_magnitude(ideal_init_scale, num_layers) print( f"Initialization Scale: {small_init_scale:.3f}, " f"Final Gradient Norm: {grad_norm_small:.4e}" ) print( f"Initialization Scale: {ideal_init_scale:.3f}, " f"Final Gradient Norm: {grad_norm_ideal:.4e}" ) print( f"Initialization Scale: {large_init_scale:.3f}, " f"Final Gradient Norm: {grad_norm_large:.4e}" ) # 预期(近似)输出: # Initialization Scale: 0.050, Final Gradient Norm: 9.5367e-08 # Initialization Scale: 0.100, Final Gradient Norm: 1.0000e+00 # Initialization Scale: 0.500, Final Gradient Norm: 9.7656e+06上述简单模拟(不含归一化和非线性,它们会使情况更复杂)说明了权重尺度如何直接影响经过多层反向传播后的梯度大小。较小的初始尺度会导致梯度范数消失,而较大的尺度则会导致其爆炸。一个“理想”的尺度(例如Xavier/Glorot初始化中提示的 $1/\sqrt{\text{fan_in}}$)有助于使梯度范数保持在其原始大小附近。前向传播的稳定性初始化也影响前向传播。如果权重过大,线性层的输出会大幅增长,可能将激活函数的输入推到饱和区域(例如,对于sigmoid或tanh),使得梯度接近零。这再次阻碍了学习。反之,如果权重过小,激活值可能会逐层减弱,导致表示趋近于零,并降低网络的有效容量。因此,我们将要讨论的Xavier和Kaiming初始化等原则性权重初始化策略的目标,是根据层维度谨慎设置权重的初始尺度。目的是确保前向传播中的激活和反向传播中的梯度在整个网络中保持合理的方差,从而防止信号消失或爆炸,进而促进更快、更稳定的训练收敛。这对本课程中非常重要的深度Transformer模型尤为重要。