归一化

总结

深度学习中的归一化（Normalization）技术是训练高性能模型的关键环节之一。它们主要解决两个问题：

  1. 内部协变量偏移（Internal Covariate Shift）：在训练过程中，网络前面层的参数不断更新，导致后面层接收到的数据分布一直变化，这使得后面层需要不断适应新的分布，减慢了训练速度。
  2. 梯度消失/爆炸：不稳定的数据分布可能导致梯度在反向传播时变得过大或过小，使模型难以收敛。

归一化层通过将数据拉回一个稳定的、标准的分布（通常是均值为0、方差为1），从而加速收敛、提高模型泛化能力。

第一类：对激活值的归一化

这类方法作用于层的输出（激活值）上。它们的核心区别在于计算均值和方差的数据范围不同。 为了方便理解，我们假设一个思维的特征图张量（Tensor）[N, C, H, W]，其中N：batch size，C：Channels，H：Height（特征图高度），W：Width（特征图宽度）。

快速记忆

为了方便对比，下表总结了各种对激活值进行归一化的方法（以[N, C, H, W]四维张量为例）：

方法 (Method)归一化维度 (Normalization Dimensions)可学习参数 (Learnable Params)参数尺寸 (Parameter Size)是否依赖Batch (Batch Dependent?)主要应用 (Primary Use Case)
BatchNorm[N, H, W] (对每个通道 C 计算)γ, βC是CNN (大批量)
InstanceNorm[H, W] (对每个样本 N、每个通道 C 计算)γ, βC否风格迁移, GANs
GroupNorm[C/G, H, W] (对每个样本 N、每个通道组 G 计算)γ, βC否CNN (小批量)
LayerNorm[C, H, W] (对每个样本 N 计算)γ, βNLP: d_model, CV: [C,H,W]否Transformers, RNNs
RMSNorm[C, H, W] (对每个样本 N 计算)g (或 γ)NLP: d_model, CV: [C,H,W]否Transformers (计算高效)

BatchNorm（BN）

  • 核心思想：对每个通道，在整个批次内进行归一化。
  • 如何工作：
    • 对于每一个通道c，BN会收集这一个批次中，所有样本(N)在该通道上的所有激活值(H*W个)，然后计算这些值的均值和方差。
    • 归一化维度B：[N, H, W]。它的C个通道中，每一个通道都计算一个独立的均值和方差。
    • 公式：$y=\gamma \frac{x - \mu_{B}}{\sigma_{B}} + \beta$
      • $\mu_B$和$\sigma_B$是在一个批次B上计算得到的均值和标准差。
      • $\gamma$和$\beta$的是可学习的缩放和平移参数，让网络可以自己学习恢复数据的原始分布，增加了模型的表达能力。
      • $\mu_B$、$\sigma_B$、$\gamma$、$\beta$的维度只和通道数C相关。
      • $\mu = \frac{1}{N} \sum_{i=1}^{N} x_i$，$\sigma^2 = \frac{1}{N} \sum_{i=1}^{N} (x_i - \mu)^2$，为了防止除零：$\sigma=\sqrt{sigma^{2}+\epsilon}$
      • 这里计算的时候采用的是有偏样本方差，即除数是N而不是N-1，这是因为工程实现上方便求导；且一般N会很大，所以差别不会特别大。

BatchNorm的局限性：

  • 对Batch Size敏感：BN严重依赖于足够大的batch size。如果batch size很小（比如1或2），计算出的均值和方差将会有很大的噪声，无法代表整个数据集的分布，从而导致模型的性能下降。
  • 不适用于序列数据(RNN/Transformer)：在处理文本等变长序列数据时，一个批次的每个样本可能有不同的长度。如果对每个时间步的特征进行归一化，强行凑成的矩形批次会引入大量无效padding，干扰均值和方差的计算。
  • 训练和推理不一致：会给部署带来一定的复杂性。

Sync BatchNorm

TODO

InstanceNorm

Instance Normalization（实例归一化，简称IN）是另一种重要的归一化技术，它尤其在风格迁移和图像生成任务中大放异彩。

  • 核心思想：对每个样本的每个通道独立进行归一化。
  • 如何工作：
    • 假设特征图张量为[N, C, H, W]。
    • IN会为N个样本中的每一个，C个通道中的每一个，都计算一个独立的均值和方差。
    • 归一化维度B：[H, W]。这意味着它在空间维度上进行归一化。总共会计算N * C组均值和方差。
    • 它完全独立于Batch Size。
  • 为什么在风格迁移中有效？
    • 风格迁移的目标是将一张内容图的内容和一张风格图的风格结合起来。
    • 图像的风格（如色调、对比度、纹理）通常被认为体现在特征图通道的统计数据（均值和方差）中。
    • IN对每个通道、每个样本进行归一化，相当于“抹掉”了原始图像的对比度和颜色分布等风格信息，使得模型可以更专注于内容本身，然后再将新的风格（通过AdaIN等机制）施加上去。

优点：

  • 与Batch Size无关：计算完全在单个样本内部完成。
  • 强大的风格“清洗”能力：在生成任务中，能有效去除实例特定的风格信息。

缺点：

  • 不适用于所有任务：在分类等判别性任务中，图像的原始对比度等信息可能是有用的特征。IN会破坏这些信息，因此在这些任务上效果通常不如BN。

AdaIN

Adaptive Instance Normalization（自适应实例归一化，简称AdaIN）是风格迁移领域的里程碑式工作。它无需训练新的网络即可实现任意风格的实时迁移。

  • 核心思想：将内容图像的特征分布（均值和方差）对齐到风格图像的特征分布上。
  • 公式：$\text{AdaIN}(x, y) = \sigma(y) \left( \frac{x - \mu(x)}{\sigma(x)} \right) + \mu(y)$
    • $x$：内容图像的特征图（Content Input）。
    • $y$：风格图像的特征图（Style Input）。
    • $\mu(x), \sigma(x)$：内容特征的通道均值和标准差。
    • $\mu(y), \sigma(y)$：风格特征的通道均值和标准差。
  • 与BN/IN的区别：
    • BN/IN/LN 等使用可学习的仿射参数 $\gamma$ 和 $\beta$。
    • AdaIN 直接使用风格输入 $y$ 的统计量作为 $\gamma$ 和 $\beta$。即 $\gamma = \sigma(y)$，$\beta = \mu(y)$。
  • 代码实现：
def calc_mean_std(feat, eps=1e-5):
    # feat: [N, C, H, W]
    size = feat.size()
    assert (len(size) == 4)
    N, C = size[:2]
    feat_var = feat.view(N, C, -1).var(dim=2) + eps
    feat_std = feat_var.sqrt().view(N, C, 1, 1)
    feat_mean = feat.view(N, C, -1).mean(dim=2).view(N, C, 1, 1)
    return feat_mean, feat_std

def adaptive_instance_normalization(content_feat, style_feat):
    assert (content_feat.size()[:2] == style_feat.size()[:2])
    size = content_feat.size()
    style_mean, style_std = calc_mean_std(style_feat)
    content_mean, content_std = calc_mean_std(content_feat)

    normalized_feat = (content_feat - content_mean.expand(size)) / content_std.expand(size)
    return normalized_feat * style_std.expand(size) + style_mean.expand(size)

SPADE

  • 细节见：SPADE
  • 与AdaIN的区别：SPADE计算的参数$\beta_{c,y,x}$ $\gamma_{c,y,x}$是和通道、像素相关的，也就是每个位置都分别变化，控制粒度比AdaIN更细。

GroupNorm

Group Normalization（组归一化，简称GN）是BatchNorm在小批量（small batch size）场景下的有力替代品。

  • 核心思想：将通道（Channels）分成若干组（Groups），然后在每个样本的每个组内进行归一化。
  • 如何工作：
    • 假设特征图张量为[N, C, H, W]，并设定一个超参数G（组数）。
    • 首先，将C个通道分成G个组，每个组有C/G个通道。
    • 然后，对于每个样本n和每个组g，计算这C/G个通道内的所有元素（(C/G) * H * W个值）的均值和方差，并进行归一化。
    • 总共会计算N * G组均值和方差。
  • 与其它归一化的关系：
    • GN可以看作是LayerNorm和InstanceNorm的折中。
    • 当G=1时，GN等价于LayerNorm（所有通道在一个组内归一化）。 不等价，因为可学习参数的尺寸不相同。
    • 当G=C时，GN等价于InstanceNorm（每个通道自成一个组进行归一化）。

优点：

  • 与Batch Size无关：计算独立于批次大小，因此在batch size很小（如1, 2, 4）时性能依然稳定，完美解决了BN的痛点。
  • 性能稳定：在从物体检测、分割到视频分类等需要使用大模型但受限于显存只能用小批量的任务中，GN是首选。
  • 灵活性：通过调整组数G，可以在LN和IN之间进行插值，为模型提供更灵活的归纳偏置。

缺点：

  • 引入超参数：需要手动设置组数G，给调参带来一些额外工作。
  • 大批量下可能不如BN：在batch size足够大的情况下，BN由于利用了整个批次的信息，估计的统计量更准，性能通常会略优于GN。

LayerNorm

LayerNorm，对单个样本的所有特征进行归一化，与BatchNorm不同的是，它的计算完全独立于batch中的其他样本，这个特性使得其在处理序列数据（如自然语言处理）的循环神经网络（RNN）和Transformer模型中取得巨大成功。 样本均值（Mean）：$\bar{x}=\frac{1}{n}\sum\limits_{i=1}^{n}x_{i}$ 样本方差（Variance）：$S_{n}^{2}=\frac{1}{n}\sum\limits_{i=1}^{n}(x_{i} - \bar{x})^{2}$，这里方便梯度回传用的是有偏估计。 归一化（Normalization）：$\hat{x}_{i}=\frac{x_{i}-\bar{x}}{\sqrt{S_{n}^{2} + \epsilon}}$ 仿射变换（Affine Transformation）：$y_{i}=\gamma\hat{x}_{i}+\beta$

  • $\gamma$（gamma，gain/增益）和$\beta$（beta，bias/偏量）是两个可学习的参数，维度与特征维度相同（eg：对于[N, C, H, W]的激活值，可学习参数的维度就是[C, H, W]）。
  • 这一步非常重要，它让网络可以根据需要“撤销”或调整归一化操作。如果网络发现原始的特征分布更有用，它可以通过学习让$\gamma$等于原始标准差，$\beta$等于原始均值，从而恢复原始激活值，这给模型了更大的灵活性。

LayerNorm的优点：

  • 与BatchSize无关：这是他最大的优势，使其在小批量、在线学习或样本量变化较大的场景中表现优异。
  • 确定性：训练和推理时行为一致，不需要保存全局统计量。
  • 非常适合序列模型：在RNNs和Transformers中效果显著，已成为标配。

缺点：

  • 不适用于CNN：卷积神经网络中，BatchNorm的效果通常优于LayerNorm。因为CNN假设图像的不同通道（特征）具有独立的、不同的含义。LayerNorm强制将这些不同含义的通道拉到同一个尺度上进行归一化，可能会破坏有效的特征表示。而BatchNorm保持了每个通道的独立性，更符合CNN的归纳偏置。
  • 计算开销：相比于后面RMSNorm，LayerNorm需要计算均值，多了一步计算，速度稍慢。

代码实现：

import torch
import torch.nn as nn
class LayerNorm(nn.Module):
	def __init__(self, d_model: int, eps: float = 1e-6):
		super().__init__()
		self.d_model = d_model
		self.eps = eps
		self.gamma = nn.Parameter(torch.ones(d_model))
		self.beta = nn.Parameter(torch.zeros(d_model))
	
	def forward(self, x):
		# NC
		mean = x.mean(dim=-1, keepdim=True)
		var = x.var(dim=-1, unbiased=False, keepdim=True)
		x_normalized = (x - mean) / torch.sqrt(var + self.eps)
		
		output = x_normalized * self.gamma + self.beta
		return output

RMSNorm

RMSNorm(Root Mean Square Layer Normalization)： RMSNorm是一种比Layer Normalization(LayerNorm)更简单、计算效率更高的归一化技术。它的核心思想是：保留LayerNorm的重新缩放（re-scale）特性，但去除了重新中心化(re-centering)的步骤。实践证明，这种简化在保持模型性能的同时，显著提升了计算速度。

$$\begin{aligned} \mathbf{y} = \frac{\mathbf{x}}{\text{RMS}(\mathbf{x})} \odot \mathbf{g} \\ \end{aligned}$$
  • 输入向量：$x=(x_{1},x_{2},x_{3},...,x_{n})$
  • 计算均方根（Root Mean Square，RMS）：$RMS(x) = \sqrt{\frac{1}{n} \sum\limits_{i=1}^{n} x_i^2 + \epsilon}$，这里$RMS$没有减去均值，它直接计算了向量元素的平方和的均值，然后开根号。
  • 归一化（Normalization）：这一步只进行重新缩放。$\bar{x}_{i} = \frac{x}{RMS(x)}$。
  • 仿射变换（Affine Transformation）：$y_{i}=g_{i}\bar{x}_{i}$，这里只有一个可学习的增益参数$g_{i}$（gain）。由于没有了中心化的步骤，偏置参数$\beta$（bias）也被认为是不必要的，因此被移除了（尽管有些实现为了灵活性而保留了它）。

RMSNorm的优势：

  • 计算效率高：
    • 在GPU等硬件上，计算均值和方差需要两次遍历数据或更复杂的并行算法。
    • RMSNorm只需要计算平方和，这是一个更简单的操作，可以更快完成。
    • 根据论文实验，RMSNorm在GPU上的速度比LayerNorm快**7%到64%。
  • 性能相当或更优：
    • 在多种任务和模型架构上，RMSNorm的性能与LayerNorm相当，有时甚至略有提升。
    • 这证明了当初的假设：对于Transformer模型来说，归一化中的重新缩放比重新中心化更重要。
  • 隐式正则化效果：
    • 论文作者推测，不减去均值可能会给模型带来一种隐式的正则化效果，因为梯度的变化会受到输入均值的影响，这可能有助于模型的泛化。

代码实现：

import torch
import torch.nn as nn

class RMSNorm(nn.Module):
	def __init__(self, d_model: int, eps: float = 1e-6):
		"""
		初始化 RMSNorm 模块。

		参数: 
		d_model (int): 模型的维度，即输入张量最后一个维度的大小。 
		eps (float): 一个很小的数，用于防止除以零。
		"""
		super().__init__()
		self.d_model = d_model
		self.eps = eps
		self.gamma = nn.Parameter(torch.ones(d_model))
	def _norm(self, x):
		# 计算输入的均方根
		# torch.rsqrt是 平方根的倒数 1/sqrt
		return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
	
	def forward(self, x):
		# NC
		output = self._norm(x.float()).type_as(x)
		return output * self.gamma

第二类：对权重的归一化

SpectralNorm

只在训练中生效，推理前模型导出可以把其烘焙到参数中。

Spectral Normalization（谱归一化，简称SpectralNorm或SN）是一种权重归一化技术，主要用于稳定GAN的训练。它的核心思想是通过限制网络中每一层权重矩阵的谱范数（Spectral Norm）来控制判别器（Discriminator）的Lipschitz常数，从而使判别器函数更加平滑，提供给生成器。（Gemerator）的梯度更加稳定有效。

动机：GAN训练的困境

GAN的训练实际上是一个“猫鼠游戏”：

  • 生成器（G）试图生成以假乱真的数据。
  • 判别器（D）试图区分真实数据和生成数据。

理想情况下，两者同步优化，最终G能够生成高质量的数据。但实际训练中，这个过程非常脆弱，经常出现问题：

  1. 梯度消失：如果判别器D太强大，能轻易地区分真假样本，它返回给生成器G地梯度会非常小，G几乎学不到任何东西。
  2. 梯度爆炸：如果判别器D的梯度过大且不稳定，G的更新会剧烈震荡，无法收敛。
  3. 模式崩溃：G 发现了一个或几个能轻易骗过 D 的样本，然后就只生成这些样本，不再探索数据的其他多样性。

这些问题的根源是判别器D的函数行为过于“剧烈”、“陡峭”。我们需要一种方法来稳定判别器。 为了解决这个问题，研究者们引入了Lipschitz约束。

Lipschitz连续性

一个函数$f$被称为$K$-Lipschitz连续的，如果对于其定义域内的任意两个点$x_1$和$x_2$，都满足：

$$ ||f(x_1) - f(x_2)|| \le K \cdot ||x_1 - x_2|| $$
  • $K$是Lipschitz常数，它限制了函数输出变化的剧烈程度。
  • 当$K=1$时，称为1-Lipschitz连续。这意味着函数的梯度范数（gradient norm）在任何地方都不会超过1。
  • 对于判别器D来说，如果它是1-Lipschitz连续的，那么它的输出就不会对输入的微小变化产生剧烈响应，梯度会变得平滑且有界，从而为生成器G提供稳定、可靠的学习信号。

谱范数与Lipschitz常数

那么如何将判别器D约束为1-Lipschitz连续呢？

  • 判别器D通常是由一系列神经网络层（如全连接层、卷积层）堆叠而成。
  • 对于一个线性层$f(x) = Wx$，它的Lipschitz常数就是其权重矩阵$W$的谱范数（Spectral Norm），记为$\sigma(W)$。
  • 谱范数$\sigma(W)$ 被定义为$W$最大的奇异值（Singular Value）。
  • 如果我们能让网络中每一层的权重矩阵$W_i$的谱范数$\sigma(W_i) \le 1$，那么整个函数（判别器）的Lipschitz常数也将被约束在1附近。

如何实现谱归一化

谱归一化的做法非常直接：

  1. 计算谱范数：对于每一层的权重矩阵$W$，计算其谱范数$\sigma(W)$。直接计算奇异值分解（SVD）的开销很大，因此实践中通常使用一种叫做幂迭代（Power Iteration） 的算法来快速地近似估算出最大奇异值。
  2. 归一化权重：将原始权重$W$替换为归一化后的权重$W_{SN}$： $$ W_{SN} = \frac{W}{\sigma(W)} $$ 这样，归一化后的权重矩阵$W_{SN}$的谱范数就等于1了。这个操作会在每次训练迭代（或每次前向传播）时执行，以确保约束持续有效。 优点：
  • 训练稳定：是稳定GAN训练最有效的方法之一，效果优于梯度惩罚（Gradient Penalty）等早期技术。
  • 计算高效：相比其他正则化方法，幂迭代的计算开销相对较小。
  • 改善生成质量：通过提供平滑的梯度，帮助生成器学习更真实、更多样化的数据分布。

python example

在 PyTorch 中，可以使用 torch.nn.utils.spectral_norm 来轻松实现谱归一化。它通常作为装饰器包裹卷积层或全连接层。

import torch
import torch.nn as nn
from torch.nn.utils import spectral_norm

# 1. 基本用法：包裹一个层
# 比如在一个全连接层上使用谱归一化
# 这会自动将层的 weight 参数替换为 weight_orig，并注册 weight_u 和 weight_v 两个 buffer 用于幂迭代计算
layer = spectral_norm(nn.Linear(20, 40))
print(f"Layer: {layer}")
# 查看 registered buffers and parameters
print(f"Keys: {list(layer.state_dict().keys())}")
# 输出 Keys: ['bias', 'weight_orig', 'weight_u', 'weight_v']
# weight_orig 是原始的可学习权重，weight 是 forward 过程中根据 spectral norm 公式计算出来的临时权重

# 2. 在 GAN 判别器中的典型应用
class Discriminator(nn.Module):
    def __init__(self, in_channels=3, features=64):
        super(Discriminator, self).__init__()
        self.model = nn.Sequential(
            # 第一层
            spectral_norm(nn.Conv2d(in_channels, features, 4, 2, 1)),
            nn.LeakyReLU(0.2, inplace=True),
            
            # 中间层
            spectral_norm(nn.Conv2d(features, features * 2, 4, 2, 1)),
            nn.LeakyReLU(0.2, inplace=True),
            
            spectral_norm(nn.Conv2d(features * 2, features * 4, 4, 2, 1)),
            nn.LeakyReLU(0.2, inplace=True),
            
            # 最后一层输出 1 (或者 logits)
            spectral_norm(nn.Conv2d(features * 4, 1, 4, 1, 0)),
        )

    def forward(self, x):
        return self.model(x)

# 验证
D = Discriminator()
x = torch.randn(2, 3, 64, 64)
out = D(x)
print(f"Output shape: {out.shape}")

关键点：

  • spectral_norm 会在 Module 中注册 weight_orig（原始权重）和 weight_u、weight_v（用于幂迭代的向量）。
  • 每次 forward 时，它会根据 weight_orig、weight_u 和 weight_v 计算出归一化后的 weight 并使用，这个过程是在 C++ 层面高效实现的。
  • 默认 n_power_iterations=1，即每次前向传播迭代一次，这在训练中足够高效且能收敛。