深度学习基础：BN层反向传播原理与实现详解

1次阅读

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

Batch Normalization（BN）层是深度学习中的一项重要技术，由 Ioffe 和 Szegedy 在 2015 年提出。它的主要作用是解决内部协变量偏移（Internal Covariate Shift）问题，通过规范化每一层的输入，使得网络的训练更加稳定和高效。BN 层不仅加速了模型的收敛速度，还允许使用更大的学习率，提高了模型的泛化能力。

BN 层的前向传播过程可以分为以下几个步骤：

计算当前 batch 的均值：
$$\mu_B = \frac{1}{m}\sum_{i=1}^m x_i$$
计算当前 batch 的方差：
$$\sigma_B^2 = \frac{1}{m}\sum_{i=1}^m (x_i – \mu_B)^2$$
归一化：
$$\hat{x_i} = \frac{x_i – \mu_B}{\sqrt{\sigma_B^2 + \epsilon}}$$
缩放和偏移：
$$y_i = \gamma \hat{x_i} + \beta$$

其中，$\gamma$ 和 $\beta$ 是可学习的参数，$\epsilon$ 是一个很小的常数，用于防止除以零。

反向传播是 BN 层最复杂的部分。我们需要计算损失函数对输入 $x_i$、$\gamma$ 和 $\beta$ 的梯度。

对 $\gamma$ 和 $\beta$ 的梯度：
$$\frac{\partial L}{\partial \gamma} = \sum_{i=1}^m \frac{\partial L}{\partial y_i} \hat{x_i}$$
$$\frac{\partial L}{\partial \beta} = \sum_{i=1}^m \frac{\partial L}{\partial y_i}$$
对输入 $x_i$ 的梯度：
$$\frac{\partial L}{\partial x_i} = \frac{\gamma}{\sqrt{\sigma_B^2 + \epsilon}} \left(\frac{\partial L}{\partial y_i} – \frac{1}{m} \sum_{j=1}^m \frac{\partial L}{\partial y_j} – \frac{\hat{x_i}}{m} \sum_{j=1}^m \frac{\partial L}{\partial y_j} \hat{x_j} \right)$$

BN 层通过规范化输入，使得每一层的输入分布更加稳定，从而允许使用更大的学习率。此外，BN 层还有一定的正则化效果，可以减少对 Dropout 的依赖。

以下是一个简化的 BN 层实现：

import torch
import torch.nn as nn

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

    def forward(self, x):
        if self.training:
            mean = x.mean(dim=0)
            var = x.var(dim=0, unbiased=False)
            # Update running statistics
            self.running_mean = (1 - self.momentum) * self.running_mean + self.momentum * mean
            self.running_var = (1 - self.momentum) * self.running_var + self.momentum * var
        else:
            mean = self.running_mean
            var = self.running_var

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

与原生 PyTorch BN 层的对比：

我们的实现更加简单，便于理解
原生实现考虑了更多边界条件和优化
原生实现支持多种输入维度（1D, 2D, 3D）

忘记在训练和推理模式下切换统计量的计算
没有正确处理 $\epsilon$ 的值
忘记更新 running_mean 和 running_var

在训练模式下，使用当前 batch 的统计量
在推理模式下，使用 running_mean 和 running_var
使用 model.eval() 和 model.train() 来切换模式

当 batch size 较小时，batch 统计量可能不准确
可以考虑使用 Group Normalization 或 Layer Normalization 替代
或者使用更大的 batch size

我们可以设计一个简单的实验来验证 BN 层的效果：

import matplotlib.pyplot as plt

# 准备数据
x = torch.randn(100, 10)
# 应用 BN 层
bn = BatchNorm1d(10)
y = bn(x)

# 可视化
plt.figure(figsize=(10, 5))
plt.subplot(1, 2, 1)
plt.hist(x.numpy().flatten(), bins=50)
plt.title("Before BN")
plt.subplot(1, 2, 2)
plt.hist(y.detach().numpy().flatten(), bins=50)
plt.title("After BN")
plt.show()

实验结果会显示，BN 层成功将输入数据规范化为均值为 0，方差为 1 的分布。