深入解析Batch Norm反向传播：从数学原理到PyTorch实现

1次阅读

没有评论

共计 3187 个字符，预计需要花费 8 分钟才能阅读完成。

Batch Normalization（BN）在训练阶段需要维护运行时统计量（均值 / 方差），这使得其反向传播比普通层更复杂。关键在于：

梯度路径分叉 ：损失L 对输入 x_i 的梯度需同时考虑 x_i 对归一化结果 y_i 的直接影响，以及通过全局均值 μ 和方差 σ² 的间接影响
链式法则嵌套：计算∂L/∂x_i 时需展开三层链式法则：
归一化输出y_i = (x_i - μ)/√(σ² + ε)
均值μ = 1/m ∑x_j
方差σ² = 1/m ∑(x_j - μ)²
计算图膨胀 ：每个样本x_i 的梯度计算都依赖全体样本的统计量，导致显存访问模式复杂化

设 batch 大小为 m，输入x_i 的梯度公式推导如下（原始论文 [1] 附录 A）：

展开归一化变换：
$$y_i = \frac{x_i – μ}{\sqrt{σ^2 + ε}}$$
损失对 x_i 的总梯度：
$$\frac{∂L}{∂x_i} = \frac{∂L}{∂y_i} \cdot \frac{∂y_i}{∂x_i} + \frac{∂L}{∂μ} \cdot \frac{∂μ}{∂x_i} + \frac{∂L}{∂σ^2} \cdot \frac{∂σ^2}{∂x_i}$$
逐项计算（关键步骤）：
直接梯度项：
$$\frac{∂y_i}{∂x_i} = \frac{1}{\sqrt{σ^2 + ε}}$$
均值相关项：
$$\frac{∂μ}{∂x_i} = \frac{1}{m}$$
$$\frac{∂L}{∂μ} = \sum_{j=1}^m \frac{∂L}{∂y_j} \cdot \frac{-1}{\sqrt{σ^2 + ε}}$$
方差相关项：
$$\frac{∂σ^2}{∂x_i} = \frac{2(x_i – μ)}{m}$$
$$\frac{∂L}{∂σ^2} = \sum_{j=1}^m \frac{∂L}{∂y_j} \cdot (x_j – μ) \cdot \frac{-1}{2}(σ^2 + ε)^{-3/2}$$
最终合并形式：
$$\frac{∂L}{∂x_i} = \frac{1}{m\sqrt{σ^2 + ε}} \left[m\frac{∂L}{∂y_i} – \sum_{j=1}^m \frac{∂L}{∂y_j} – (x_i – μ) \cdot \sum_{j=1}^m \frac{∂L}{∂y_j}(x_j – μ) \cdot \frac{1}{σ^2 + ε} \right]$$

以 torch.nn.BatchNorm2d._backward() 为例（v1.12 源码）：

梯度预处理：

# 将上层梯度 (dL/dy) 与标准化系数相乘
grad_output = grad_output * self.weight.view(1, -1, 1, 1)

均值梯度计算：

# 对应公式中的 sum(dL/dy_j)项
grad_mean = torch.sum(grad_output, dim=(0, 2, 3), keepdim=True)

方差梯度计算：

# 计算点乘项 sum(dL/dy_j * (x_j - μ))
dot_p = torch.sum(grad_output * (input - self.running_mean.view(1, -1, 1, 1)),
    dim=(0, 2, 3),
    keepdim=True
)

最终梯度合成：

# 对应完整梯度公式
grad_input = (grad_output - grad_mean / N 
             - dot_p * (input - mean) / (var + self.eps) / N
            ) / torch.sqrt(var + self.eps)

完整 BatchNorm 层实现（训练模式）：

class MyBatchNorm2d:
    def __init__(self, num_features, eps=1e-5):
        self.gamma = torch.ones(num_features)
        self.beta = torch.zeros(num_features)
        self.eps = eps
        self.running_mean = torch.zeros(num_features)
        self.running_var = torch.ones(num_features)

    def forward(self, x):
        if self.training:
            dims = (0, 2, 3)
            mean = x.mean(dims, keepdim=True)
            var = x.var(dims, unbiased=False, keepdim=True)
            self.running_mean = 0.9 * self.running_mean + 0.1 * mean.squeeze()
            self.running_var = 0.9 * self.running_var + 0.1 * var.squeeze()
        else:
            mean, var = self.running_mean, self.running_var

        x_hat = (x - mean) / torch.sqrt(var + self.eps)
        return self.gamma * x_hat + self.beta

    def backward(self, grad_output, x):
        N = x.shape[0] * x.shape[2] * x.shape[3]
        mean = x.mean((0, 2, 3), keepdim=True)
        var = x.var((0, 2, 3), unbiased=False, keepdim=True)

        grad_gamma = (grad_output * x_hat).sum((0, 2, 3))
        grad_beta = grad_output.sum((0, 2, 3))

        dx_hat = grad_output * self.gamma.view(1, -1, 1, 1)
        dvar = torch.sum(dx_hat * (x - mean) * -0.5 * (var + self.eps)**(-1.5), 
                        (0, 2, 3))
        dmean = torch.sum(dx_hat * (-1 / torch.sqrt(var + self.eps)), 
                         (0, 2, 3)) \
               + dvar * torch.mean(-2 * (x - mean), (0, 2, 3))

        grad_input = (dx_hat / torch.sqrt(var + self.eps) +
                     dvar.view(1, -1, 1, 1) * 2 * (x - mean) / N +
                     dmean.view(1, -1, 1, 1) / N)
        return grad_input, grad_gamma, grad_beta