技术摘要：NoiseRater——用于扩散模型训练的元学习噪声估值

1. 问题陈述

扩散模型通过学习逆转渐进式加噪过程，在生成任务中实现了最先进（SOTA）的性能。然而，标准的训练范式将注入的噪声视为具有均匀的信息量。尽管近期研究集中于优化推理时的噪声（例如，引导缩放、噪声重采样）或在时间步级别重新加权噪声（例如，P2、Min-SNR），但训练过程中实例级噪声变化的作用仍未得到充分探索。

作者提出了一个根本性问题：所有噪声实现是否对学习同样有用？即使在同一时间步且针对同一数据样本，不同的噪声实例也可能携带不同级别的学习信号。有些可能提供清晰的去噪监督，而另一些则可能是模糊的、冗余的或有害的。均匀地对待这些变化可能导致次优的学习动态。这项工作通过将焦点从测试时的噪声操纵转移到训练时的噪声估值，填补了这一空白。

2. 方法论：NoiseRater

所提出的框架NoiseRater引入了一种元学习方法，为各个噪声实现分配重要性分数，从而实现对训练目标的自适应重加权。

2.1 参数化噪声评分器

核心组件是一个参数化函数 $\phi_\eta$，它充当“噪声评分器”。给定清晰图像 $x_0$、条件 $c$（可能被丢弃）、时间步 $t$ 以及特定的噪声实现 $\epsilon$，该评分器输出一个标量分数，代表该噪声实例对学习信号的相对重要性。

为了确保评分器评估的是噪声而非图像内容，该方法采用了分组噪声采样：

• 对于小批量中的每个数据样本 $(x_0, c, t)$，抽取 $K$ 个独立的噪声样本 $\{\epsilon^{(k)}\}_{k=1}^K$。
• 这些样本生成 $K$ 个噪声输入 $\{x_t^{(k)}\}_{k=1}^K$。
• 评分器的分数使用 softmax 函数进行组内归一化：
  $$w^{(k)}(x_0) = \frac{\exp(\phi_\eta(\epsilon^{(k)}, t, x_0, \tilde{c}))}{\sum_{k'=1}^K \exp(\phi_\eta(\epsilon^{(k')}, t, x_0, \tilde{c}))}$$
  这确保了权重反映固定条件下噪声实例的相对效用，防止模型利用图像内容的变化。

2.2 双层优化框架

噪声评分器通过双层优化框架进行训练，以直接优化下游泛化性能：

• 内循环（模型训练）： 扩散模型参数 $\theta$ 使用加权损失 $\mathcal{L}_{inner}$ 进行更新，其中标准 MSE 损失由评分器的分数 $w^{(k)}$ 重新加权。
  $$\mathcal{L}_{inner}(\theta; \eta) = \mathbb{E}_{x_0 \sim \mathcal{B}} \left[ \sum_{k=1}^K w^{(k)}(x_0) \cdot \|\epsilon^{(k)} - \epsilon_\theta(x_t^{(k)}, t, \tilde{c})\|^2 \right]$$
• 外循环（评分器训练）： 评分器参数 $\eta$ 被更新以最小化内循环更新后扩散模型的验证损失 $\mathcal{L}_{val}$。
  $$\eta \leftarrow \arg\min_\eta \mathcal{L}_{val}(\theta^*(\eta))$$
  其中 $\theta^*(\eta)$ 表示经过 $S$ 次内循环梯度步后的模型参数。这种元学习设置允许评分器学习哪些噪声实例能带来更好的泛化，而不是依赖启发式标准。

2.3 解耦的两阶段流程

为了实现高效部署并避免在标准训练期间维护多个噪声候选所带来的计算开销，作者提出了一个两阶段流程：

1. 元训练阶段： 使用软加权（利用所有 $K$ 个候选）和双层优化来训练噪声评分器。
2. 标准训练阶段： 训练好的评分器被冻结。在随后的扩散训练中，对于每个数据样本，采样 $K'$ 个候选噪声，并选择Top-1噪声（得分最高者）进行训练更新。这种“硬选择”策略将噪声评估与训练过程解耦，在保持标准训练效率的同时保留了实例级估值的优势。

3. 主要贡献

该论文概述了四项主要贡献：

1. 问题形式化： 识别并将训练时噪声估值形式化为现有时间步级噪声设计和测试时控制方法的补充维度。
2. 元学习评分器： 提出了一种参数化噪声评分器，通过双层优化进行训练，根据其泛化贡献自适应地加权噪声样本。
3. 高效流程： 引入了一种解耦的两阶段训练流程，从软加权（元训练）过渡到硬噪声选择（标准训练），实现了可扩展的部署。
4. 实证验证： 证明了并非所有噪声实现都同样有用，且有选择地强调信息丰富的噪声可以提高训练效率和生成质量。

4. 实验结果

该方法在 FFHQ（256×256）和 ImageNet（256×256）上使用扩散 Transformer（DiT）骨干网络进行了评估。

• 性能提升： 与标准训练（Vanilla）和朴素噪声选择策略（选择最小/最大范数噪声）相比，NoiseRater 始终实现了更低的 FID-50K 分数。在 ImageNet（无条件）上，该方法在 +80k 步时将 FID 从 61.56 降低到 59.95（相对降低 2.6%）。

• 消融研究：

  • 输入条件： 评分器需要所有输入（噪声、图像 $x_0$、时间步 $t$ 和类别标签）才能达到最佳性能。移除时间步嵌入导致了最大的性能下降，表明最优噪声在扩散步骤间存在显著差异。
  • 候选池大小（$k$）： $k=4$ 的池大小在质量提升和计算成本之间提供了最佳权衡。更大的池（$k=8$）收益递减，并增加了偏离高斯先验的偏差。
  • 泛化性： 在较小骨干网络（DiT-S/2）上训练的评分器成功迁移到了较大的骨干网络（DiT-B/2 和 DiT-L/2）而无需重新训练，表明评分器学习了与骨干网络无关的噪声质量属性。
  • 训练阶段： 评分器在训练的不同阶段（100k、400k、500k 步）应用时均提供了持续改进，其中在 400k 检查点附近应用时观察到的增益最为显著。

• 可解释性： 对评分器行为的分析显示：

  • 它不会退化为简单的噪声范数过滤器（初始训练后与噪声范数的相关性 $\approx 0$）。
  • 它在中间噪声水平（$t \in [0.6, 0.7]$）最具区分度，此时扩散模型置信度最低，学习信号最为关键。
  • 其区分能力在扩散模型训练约 60k 步后趋于稳定。

5. 意义与主张

作者声称，NoiseRater 确立了噪声估值作为改善扩散模型训练的一个此前未被充分探索的维度。通过显式建模各个噪声样本的重要性，该方法超越了噪声具有均匀信息量的假设。

这项工作的意义在于：

• 互补性： 它提供了一个新的优化方向，补充了现有的时间步级调度和测试时控制方法。
• 效率： 它提高了训练效率和生成质量，而无需对扩散模型进行架构更改或修改推理过程。
• 元学习应用： 它成功地将双层优化适应于评估注入噪声的具体问题，将估值对象从数据点转移到噪声实现。

该论文得出结论，优先处理信息丰富的噪声是增强扩散模型的一种可行策略，而所提出的框架为实现这一目标提供了一种灵活且有效的机制。