【深度神经网络】优化深度神经网络

EMA（Exponential Moving Average）

EMA（指数移动平均）是一种 平滑技术，在深度学习中，常用于 存储模型可学习参数的局部平均值。

可以把它看作一个“影子模型”，它的参数不是简单地复制原模型，而是随着训练 以指数衰减的方式 逐步向原模型靠拢。

为什么要使用 EMA？

• 提高模型稳定性：在训练过程中，模型参数可能会剧烈波动，EMA 平均化了参数，使其更稳定。
• 提升泛化能力：直接使用 EMA 计算的参数进行推理，通常比原始参数表现更好，尤其是在 少量训练步数 下。
• 适用于生成模型（如 Diffusion Models）：Diffusers 库中的 Stable Diffusion 训练时 使用 EMA 来平滑 UNet 权重，使生成的图像更加稳定。
• 在半监督学习中常用：如 Mean Teacher 方法，使用 EMA 计算的模型作为“教师”模型指导学生模型学习。

EMA 在累积历史信息的同时，更关注最近的更新，从而对新数据变化更敏感，而不会受太早的参数扰动。

假设：

• ${\theta }_t$ 是第 $t$ 轮训练的模型参数
• ${\theta }_{\text{EMA},t}$ 是第 $t$ 轮的 EMA 计算的影子参数
• $\alpha$ 是 EMA 衰减系数（通常取 0.99 ~ 0.999）

EMA 参数的更新方式：
$${\theta }_{\text{EMA},t}=\alpha \cdot {\theta }_{\text{EMA},t-1}+(1-\alpha )\cdot {\theta }_t$$
这意味着：

• 较早的参数影响力逐渐减弱（因为乘以了 $\alpha$）。
• 最近的参数更新权重更大（乘以 $1-\alpha$）。
• 选择 较大的 $\alpha$（如 0.999），EMA 更新较慢，适用于平滑长时间的变化。

为什么较早的参数影响力逐渐减弱？

我们可以将 EMA 当前参数展开，看看它是如何由历史所有参数的加权平均组成的：

$${\theta }_{\text{EMA},t}=(1-\alpha )\cdot {\theta }_t+\alpha (1-\alpha )\cdot {\theta }_{t-1}+{\alpha }^2(1-\alpha )\cdot {\theta }_{t-2}+{\alpha }^3(1-\alpha )\cdot {\theta }_{t-3}+\dots$$

这说明：

• 最近的参数 ${\theta }_t$ 乘以 $1-\alpha$（即 0.01），虽然数值小，但它是最新的更新，影响直接而强烈。
• 较早的参数 ${\theta }_{t-1},{\theta }_{t-2}$ 乘以 $\alpha ,{\alpha }^2$ 等次幂，影响力随着时间推移呈指数级衰减。
• 老的参数贡献依然存在，但比重越来越小，这使得 EMA 更关注近期变化，而不会被早期的不稳定训练步骤影响太多。

💡直觉理解 EMA 的本质是一种带有“记忆衰减”的平滑机制：

• 老的参数不会立刻丢失，但它的影响会随着时间逐步减弱，让新数据有更大的话语权。
• 虽然最近参数的权重（1 - α = 0.01）看似小，但它不会被 EMA 继续削弱，因此它的相对影响力更大。
• 较早的参数影响力会随着 ${\alpha }^t$ 指数级减少，长期来看其贡献会趋近于 0。

如果 $\alpha =0.99$，那么过去 5 个时间步的参数贡献依次为：
$$\begin{array}{rl}\text{Step }t:& (1-\alpha )=0.01\\ \text{Step }t-1:& 0.99\times 0.01=0.0099\\ \text{Step }t-2:& 0.99^2\times 0.01=0.009801\\ \text{Step }t-3:& 0.99^3\times 0.01=0.00970299\\ \text{Step }t-4:& 0.99^4\times 0.01=0.0096059601\end{array}$$

下面是一个简单的 PyTorch EMA 代码示例，展示如何在训练过程中维护一个 EMA 版本的模型参数。

import torch
import torch.nn as nnclass EMA:"""指数移动平均（EMA），用于平滑模型参数"""def __init__(self, model, decay=0.999):self.model = modelself.decay = decay  # EMA 影子参数衰减系数self.shadow = {name: param.clone().detach() for name, param in model.named_parameters()}def update(self):"""更新 EMA 影子模型参数"""for name, param in self.model.named_parameters():if param.requires_grad:self.shadow[name] = self.decay * self.shadow[name] + (1 - self.decay) * param.detach()def apply_shadow(self):"""使用 EMA 参数更新原模型（推理时调用）"""for name, param in self.model.named_parameters():if param.requires_grad:param.data.copy_(self.shadow[name])# 创建简单的神经网络
class SimpleModel(nn.Module):def __init__(self):super().__init__()self.fc = nn.Linear(10, 1)def forward(self, x):return self.fc(x)# 初始化模型和 EMA 影子模型
model = SimpleModel()
ema = EMA(model, decay=0.99)# 模拟训练过程
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
for step in range(100):# 训练步骤（假设 x 是输入数据）x = torch.randn(16, 10)loss = model(x).mean()optimizer.zero_grad()loss.backward()optimizer.step()# 更新 EMA 影子模型ema.update()if step % 10 == 0:print(f"Step {step}: loss={loss.item():.4f}")# 在推理时应用 EMA 参数
ema.apply_shadow()

1. EMA 类

   • 维护了 shadow（影子模型参数）。
   • 通过 update() 逐步更新 EMA 版本的参数。
   • apply_shadow() 用于推理时将 EMA 参数应用到原模型上。

2. 训练过程中

   • 每次模型参数更新后，调用 ema.update()，让影子模型参数缓慢跟随原模型更新。

3. 推理时

   • ema.apply_shadow() 把 EMA 版本的参数复制到模型，通常能获得 更好的性能。

动量梯度下降

RMSProp 优化器

Adam 优化器