BatchNormalization反向传播原理详解与实现避坑指南

1次阅读

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

BatchNormalization（BN）是深度学习中广泛使用的技术，主要用于解决内部协变量偏移（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$$

反向传播的核心是计算损失函数对输入 $x_i$、缩放参数 $\gamma$ 和偏移参数 $\beta$ 的梯度。以下是详细的推导过程：

损失函数对 $y_i$ 的梯度：
$$\frac{\partial L}{\partial y_i}$$
对 $\beta$ 的梯度：
$$\frac{\partial L}{\partial \beta} = \sum_{i=1}^m \frac{\partial L}{\partial y_i}$$
对 $\gamma$ 的梯度：
$$\frac{\partial L}{\partial \gamma} = \sum_{i=1}^m \frac{\partial L}{\partial y_i} \hat{x}_i$$
对 $\hat{x}_i$ 的梯度：
$$\frac{\partial L}{\partial \hat{x}_i} = \frac{\partial L}{\partial y_i} \gamma$$
对 $\sigma_B^2$ 的梯度：
$$\frac{\partial L}{\partial \sigma_B^2} = \sum_{i=1}^m \frac{\partial L}{\partial \hat{x}_i} (x_i – \mu_B) \left(-\frac{1}{2} (\sigma_B^2 + \epsilon)^{-3/2}\right)$$
对 $\mu_B$ 的梯度：
$$\frac{\partial L}{\partial \mu_B} = \left(\sum_{i=1}^m \frac{\partial L}{\partial \hat{x}i} \frac{-1}{\sqrt{\sigma_B^2 + \epsilon}}\right) + \frac{\partial L}{\partial \sigma_B^2} \frac{-2}{m} \sum^m (x_i – \mu_B)$$
对 $x_i$ 的梯度：
$$\frac{\partial L}{\partial x_i} = \frac{\partial L}{\partial \hat{x}_i} \frac{1}{\sqrt{\sigma_B^2 + \epsilon}} + \frac{\partial L}{\partial \sigma_B^2} \frac{2(x_i – \mu_B)}{m} + \frac{\partial L}{\partial \mu_B} \frac{1}{m}$$

以下是 PyTorch 中自定义 BN 层的完整实现：

import torch
import torch.nn as nn

class BatchNorm1dCustom(nn.Module):
    def __init__(self, num_features, eps=1e-5, momentum=0.1):
        super().__init__()
        self.gamma = nn.Parameter(torch.ones(num_features))
        self.beta = nn.Parameter(torch.zeros(num_features))
        self.eps = eps
        self.momentum = momentum
        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)
        out = self.gamma * x_hat + self.beta
        return out