Conv2D 默认初始化(Normal(0, √(2/fan_in)))在无归一化深层/共享权重结构下激活指数放大，导致训练发散与 NaN 梯度 #79706

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)