BatchNormalization反向传播推导：从数学原理到PyTorch实现

1次阅读

没有评论

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

BatchNormalization（BN）是现代深度神经网络中的重要组件，它能有效加速训练并提升模型泛化能力。前向传播过程可分为三个主要步骤：

计算当前 batch 数据的均值 μ 和方差 σ²
对数据进行归一化：x̂ = (x – μ) / √(σ² + ε)
缩放和平移：y = γx̂ + β

其中 γ 和 β 是可学习的参数，ε 是为数值稳定性添加的小常数。这个过程使得每层的输入保持相似的分布，缓解了内部协变量偏移问题。

反向传播的目标是计算以下梯度：∂L/∂x, ∂L/∂γ, ∂L/∂β。我们需要使用链式法则逐步推导。

这两个梯度的计算相对简单，因为 γ 和 β 是直接作用于归一化后的数据：

∂L/∂γ = Σ(∂L/∂y * x̂)
∂L/∂β = Σ(∂L/∂y)

从 y = γx̂ + β 可得：

∂L/∂x̂ = ∂L/∂y * γ

这是最复杂的部分，因为 x 参与了均值、方差和归一化三个计算。我们需要考虑所有路径的梯度贡献：

∂L/∂x = ∂L/∂x̂ * ∂x̂/∂x + ∂L/∂σ² * ∂σ²/∂x + ∂L/∂μ * ∂μ/∂x

经过推导（详细步骤见下文），最终得到：

∂L/∂x = (∂L/∂x̂ – mean(∂L/∂x̂) – x̂ * mean(∂L/∂x̂ * x̂)) / √(σ² + ε)

∂x̂/∂x = 1/√(σ² + ε)
∂σ²/∂x = 2(x – μ)/m（m 是 batch 大小）
∂μ/∂x = 1/m

将这些部分组合起来，经过化简后得到上述最终表达式。

下面我们实现一个自定义的 BatchNorm 层，包含完整的前向和反向传播：

import torch
import torch.nn as nn

class CustomBatchNorm1d:
    def __init__(self, num_features, eps=1e-5, momentum=0.1):
        self.gamma = nn.Parameter(torch.ones(num_features))
        self.beta = nn.Parameter(torch.zeros(num_features))
        self.eps = eps
        self.momentum = momentum
        self.running_mean = torch.zeros(num_features)
        self.running_var = torch.ones(num_features)

    def forward(self, x, training=True):
        if training:
            # 计算当前 batch 的统计量
            mean = x.mean(dim=0)
            var = x.var(dim=0, unbiased=False)

            # 更新 running 统计量
            self.running_mean = (1 - self.momentum) * self.running_mean + self.momentum * mean
            self.running_var = (1 - self.momentum) * self.running_var + self.momentum * var

            # 归一化
            x_hat = (x - mean) / torch.sqrt(var + self.eps)
        else:
            # 推理阶段使用 running 统计量
            x_hat = (x - self.running_mean) / torch.sqrt(self.running_var + self.eps)

        # 缩放和平移
        y = self.gamma * x_hat + self.beta

        # 保存中间结果供反向传播使用
        if training:
            self.cache = (x, mean, var, x_hat)

        return y

    def backward(self, dy):
        x, mean, var, x_hat = self.cache
        m = x.shape[0]  # batch size

        # 计算参数梯度
        dgamma = (dy * x_hat).sum(dim=0)
        dbeta = dy.sum(dim=0)

        # 计算输入梯度
        dx_hat = dy * self.gamma
        dvar = (dx_hat * (x - mean) * (-0.5) * (var + self.eps)**(-1.5)).sum(dim=0)
        dmean = (dx_hat * (-1) / torch.sqrt(var + self.eps)).sum(dim=0) + dvar * (-2) * (x - mean).sum(dim=0) / m
        dx = dx_hat / torch.sqrt(var + self.eps) + dvar * 2 * (x - mean) / m + dmean / m

        return dx, dgamma, dbeta

我们可以使用 PyTorch 的自动求导功能来验证我们的实现是否正确：

# 创建测试数据
batch_size = 32
features = 64
x = torch.randn(batch_size, features, requires_grad=True)

# 自定义 BN 层
custom_bn = CustomBatchNorm1d(features)

# 前向传播
y_custom = custom_bn(x)

# 计算梯度
loss = y_custom.sum()
loss.backward()

# 保存自定义 BN 的梯度
dx_custom = x.grad.clone()
dgamma_custom = custom_bn.gamma.grad.clone()
dbeta_custom = custom_bn.beta.grad.clone()

# 重置
x.grad = None

# 官方 BN 层
official_bn = nn.BatchNorm1d(features)
official_bn.gamma.data.copy_(custom_bn.gamma.data)
official_bn.beta.data.copy_(custom_bn.beta.data)

# 前向传播
y_official = official_bn(x)

# 计算梯度
loss = y_official.sum()
loss.backward()

# 比较梯度
print("输入梯度差异:", torch.allclose(dx_custom, x.grad, atol=1e-5))
print("gamma 梯度差异:", torch.allclose(dgamma_custom, official_bn.weight.grad, atol=1e-5))
print("beta 梯度差异:", torch.allclose(dbeta_custom, official_bn.bias.grad, atol=1e-5))

在实际实现中，有几个关键点需要注意：

running_mean 和 running_var 的更新时机：只在训练阶段更新，推理阶段保持不变
方差的校正：PyTorch 默认使用有偏估计（除以 m 而不是 m -1）
ε 的作用：防止除以零，通常设为 1e-5
动量参数：控制 running 统计量更新的速度

当 batch size 较小时，BN 的效果会下降，因为统计量估计不准确。可以考虑以下替代方案：

Group Normalization：不依赖 batch 维度，将通道分组计算统计量
Layer Normalization：在特征维度上归一化
Instance Normalization：对每个样本的每个通道单独归一化

BN 能缓解梯度消失问题，主要体现在：

保持各层输入的分布稳定，避免某些层的梯度变得过小
归一化后的数据通常在 0 附近，激活函数的梯度处于较大部分
可以支持更大的学习率，因为参数更新更稳定

BatchNormalization 的反向传播推导虽然复杂，但理解它对实现自定义网络层和调试模型非常重要。通过手动实现 BN 层，我们不仅能更深入理解其工作原理，也能在需要时灵活调整其行为。在实际应用中，建议优先使用框架提供的 BN 实现，但在特殊需求下（如研究新型归一化方法），这些知识将非常有用。

正文完