Conv2D 默认初始化(`Normal(0, √(2/fan_in))`)在无归一化深层/共享权重结构下激活指数放大,导致训练发散与 NaN 梯度
bug描述 Describe the Bug
paddle 的 Conv2D 默认权重初始化为 Normal(0, std),其中 std = sqrt(2/filter_elem_num)(He Normal,见 python/paddle/nn/layer/conv.py 的 _get_default_param_initializer;英文文档已注明,相关 issue #24993)。该默认增益(每层方差 ×2)是针对 ReLU 网络设计的;对于 SiLU 等平滑激活、无归一化的深层结构,每层增益 >1,当同一个卷积模块被反复应用(权重复用/动态深度 U-Net、递归结构等)时,激活值逐层指数放大,几步之内溢出到 ±10^4 量级:训练 loss 在第二步爆发到数百,反向传播同时产生 NaN 梯度。
对比:PyTorch Conv2d 默认初始化为 kaiming_uniform_(a=sqrt(5)),数学上等价于 U(±1/√fan_in)。同形状(16→16, 3x3)卷积实测 torch 默认 std=0.048,paddle 默认 std=0.117,相差约 2.45×(方差约 6×)。从 PyTorch 迁移"同构"网络到 paddle 时,若不显式指定初始化,会出现 torch 稳定收敛、paddle 第二步即发散的隐蔽差异。
最小复现(CPU 即可,无外部数据):
import paddle
import paddle.nn.functional as F
from paddle import nn
class ResBlock(nn.Layer):
def __init__(self, channels, conv):
super().__init__()
self.conv1 = conv(channels, channels)
self.conv2 = conv(channels, channels)
def forward(self, x):
return x + self.conv2(F.silu(self.conv1(x)))
def default_conv(cin, cout):
return nn.Conv2D(cin, cout, 3, padding=1) # paddle 默认: Normal(0, sqrt(2/fan_in))
def uniform_conv(cin, cout):
b = (cin * 9) ** -0.5
attr = paddle.ParamAttr(initializer=nn.initializer.Uniform(-b, b))
return nn.Conv2D(cin, cout, 3, padding=1, weight_attr=attr, bias_attr=attr)
def trial(tag, conv):
paddle.seed(0)
paddle.device.set_device("cpu")
blocks = [ResBlock(16, conv) for _ in range(2)]
x = paddle.rand([1, 16, 64, 64])
y = x
peaks = []
for i in range(8): # 共享 2 个 ResBlock 反复应用,模拟动态深度/递归复用
for blk in blocks:
y = blk(y)
peaks.append(round(float(y.abs().max()), 1))
print(tag, "每轮 max|y| =", peaks)
y2 = paddle.rand([1, 16, 64, 64])
for blk in blocks:
for _ in range(4):
y2 = blk(y2)
y2.sum().backward()
n_nan = sum(1 for blk in blocks for p in blk.parameters()
if p.grad is not None and bool(paddle.isnan(p.grad).any()))
print(tag, f"4 轮后 max|y|={float(y2.abs().max()):.1f}, NaN 梯度参数数={n_nan}")
trial("默认初始化 ", default_conv)
trial("U(±1/√fan_in)", uniform_conv)实际输出:
默认初始化 每轮 max|y| = [5.0, 15.9, 49.7, 162.4, 530.1, 1678.3, 6065.0, 20140.9]
默认初始化 4 轮后 max|y|=511.4, NaN 梯度参数数=8
U(±1/√fan_in) 每轮 max|y| = [1.4, 1.7, 2.3, 2.9, 3.7, 4.5, 5.5, 7.2]
U(±1/√fan_in) 4 轮后 max|y|=3.3, NaN 梯度参数数=0期望行为(任一即可):1) Conv2D 文档显著位置标注默认初始化与 PyTorch 默认的差异(约 2.45× std),说明其适用前提(ReLU 系网络);2) 提供与 PyTorch 对齐的默认初始化选项,降低迁移成本;3) 若社区认为合适,评估默认增益在非 ReLU 激活下的鲁棒性。
环境:paddlepaddle-gpu 3.5.0.dev20260819(本地源码编译),Windows 11,Python 3.12。CPU 与 GPU 均可复现(初始化数值问题,与设备无关)。
其他补充信息 Additional Supplementary Information
- 实测对照:
torch.nn.Conv2d(16,16,3,padding=1)默认 weight std=0.0475(max 0.083);paddle.nn.Conv2D同形状 std=0.1164(max 0.412)。 - 真实场景:我们在一个共享权重的动态深度 VSR 网络(同一批卷积模块按金字塔层级反复调用约 46-69 次)上遇到:GPU 上训练第二步 L1 loss 从 0.0006 爆发到 255;CPU 上反向传播出现 28/44 个参数的 NaN 梯度。将全部卷积显式改为
U(±1/√fan_in)后,paddle 与 torch 两个同构实现的训练曲线几乎完全重合(150 步 total loss 0.1137→0.0766 vs 0.1131→0.0740),网络其他部分未做任何修改。 - 补充观察:激活放大后,fused
F.silu的 backward 路径上出现 NaN(把 silu 改写为等价的x * F.sigmoid(x)分解形式后 NaN 消失);但对大幅值输入(|x| 直至 1000)单独测试F.silu的 backward 是稳定的,因此这更可能是激活放大后的下游现象,不构成独立的 silu bug,仅作线索供参考。 - 相关 issue:#24993(Conv2D 初始化的中英文文档不一致;英文文档描述的即当前实现)。
Source: PaddlePaddle/Paddle