NoiseRater：用于扩散模型训练的元学习噪声评估

# NoiseRater：用于扩散模型训练的元学习噪声估值

来源：https://arxiv.org/html/2605.08144

Fang Wu, Stanford University & Haokai Zhao, UNSW & Da Xing, UCL & Tinson Xu, The University of Chicago & Hanqun Cao, CUHK & Yanchao Li, Nanjing University & Zeqi Zhou, Brown University & Xiangru Tang, Yale University & Hanchen Wang, Stanford University & Hongbin Lin, CUHK & Zehong Wang, University of Notre Dame & Kuan Pang, Stanford University & Xia Peng, UNC–Chapel Hill & Yinxi Li, University of Waterloo & Aaron Tu, UCB & Molei Tao, Georgia Tech & Li Erran Li, Amazon & Aditya Joshi, UNSW & Jure Leskovec, Stanford University & Yejin Choi, Stanford University

###### 摘要

扩散模型在众多生成任务中取得了显著成功，但其训练范式通常将注入的噪声视为具有同等信息价值。在本工作中，我们挑战了这一假设，并引入了 NoiseRater，这是一种用于扩散模型训练中实例级噪声估值的元学习框架。我们提出了一种参数化的噪声评分器（noise rater），它根据数据和时间步长为单个噪声实现分配重要性分数，从而实现对训练目标的自适应重加权。该评分器通过双水平优化进行训练，以在内部循环扩散更新后提升下游验证性能。为了实现高效部署，我们进一步设计了一个解耦的两阶段流水线，从元训练期间的软加权过渡到标准训练期间的硬噪声选择。在 FFHQ 和 ImageNet 上的大量实验表明，并非所有噪声样本的贡献都是相等的，而优先处理具有信息量的噪声可以提高训练效率和生成质量。我们的结果确立了噪声估值作为改善扩散模型训练的一个互补且此前未被充分探索的维度。我们的代码可用地址为：https://anonymous.4open.science/r/NoiseRater-DEB116。

## 1 引言

扩散模型 [42, 44, 17] 已成为生成建模的主导范式，在图像 [34, 39]、视频 [32]、生物学 [51] 和多模态生成任务 [38, 52] 中均实现了最先进（SOTA）的性能。其成功背后的一个关键因素是迭代去噪过程，该过程通过一系列细化步骤将随机噪声转换为结构化数据。

最近，人们越来越关注扩散模型的*测试时计算优化* [15, 37, 24, 53, 29, 45]。特别是，噪声在推理阶段日益成为控制的核心对象 [11, 26]。诸如引导缩放 [2]、自适应步长调度 [8] 和噪声重采样策略 [1, 41] 等技术明确地操纵噪声轨迹，以在生成质量、多样性和计算成本之间进行权衡。这些方法表明，噪声的选择——包括其幅度、结构和演化——在塑造最终输出方面起着关键作用。

现有的方法主要将噪声视为*测试时的控制变量*。然而，在训练期间，噪声通常从固定的高斯分布中采样，并以 largely uniform、样本无关的方式并入目标函数。虽然先前的工作表明不同的噪声水平对学习 [49, 33, 35, 46, 13] 的贡献并不相等，但主要方法 largely 侧重于时间步级的重加权或调度设计，而*实例级*噪声变化的作用则未被充分探索。

这提出了一个基本问题：*所有的噪声实现对扩散模型的学习都具有同等效用吗？* 更具体地说，即使在相同的时间步长下，不同的噪声实例可能携带不同级别的学习信号。一些噪声可能为去噪提供清晰的监督，而另一些可能模糊、冗余或对优化而言信息量较少。如果是这种情况，那么在训练期间统一对待噪声可能导致次优的学习动力学。

在本工作中，我们将焦点从测试时的噪声操纵转移到*训练时的噪声估值*。我们提议显式建模单个噪声样本的重要性，并利用这些信息来指导训练过程。具体来说，我们引入了一种*元学习噪声评分器*，这是一种参数化函数，它为每个噪声实例分配一个分数，该分数取决于数据样本和时间步长。这些分数用于重加权扩散损失，使模型能够关注更具信息量的噪声，同时降低不太有用噪声的权重。为了学习噪声评分器，我们将问题表述为双水平优化框架 [30, 9]。扩散模型在内部循环中使用加权噪声样本进行训练，而评分器在外部循环中通过最小化验证损失进行优化。这种元学习设置使评分器能够直接捕捉噪声样本对泛化性能的贡献，而不是依赖启发式标准。

我们的方法简单、灵活，并与现代扩散框架 [43] 兼容。由于纯粹在训练目标层面操作，它不需要修改模型架构或推理过程。方法概览见图 1。

> **图 1：** 我们用于扩散训练的元学习噪声估值方法的插图。我们首先使用双水平优化训练噪声评分器，然后利用评分器的分数指导扩散训练。

我们总结如下我们的贡献：

- • 我们识别并形式化了扩散中*训练时噪声估值*的问题，这是一个未充分探索的维度，补充了时间步级噪声设计和测试时控制。
- • 我们提出了一种元学习噪声评分器，通过双水平优化框架在训练期间自适应地加权噪声样本。
- • 我们引入了一种解耦的两阶段训练流水线，将噪声评估与模型训练分离，从而实现元学习噪声策略的高效部署。
- • 我们证明了并非所有噪声实现都具有同等效用，有选择地强调具有信息量的噪声可以提高学习效率和模型性能。

## 2 相关工作

#### 推理时的噪声优化。

最近的工作表明，优化推理时的初始噪声可以显著提高生成质量。一类研究通过跨去噪步骤的反向传播直接优化噪声隐变量。这些方法迭代地更新噪声以最大化人类偏好奖励 [48]，满足特定的运动约束 [22]，或与内部注意力分数对齐 [12]。虽然有效且无需训练，但由于重复的去噪循环，它们在推理时产生了巨大的计算开销。为了分摊这一成本，其他方法训练辅助模型以直接预测最佳噪声。例如，FIND [5] 使用强化学习来调整初始高斯分布的均值和方差，而 GoldenNoise [54] 训练网络以预测高质量的噪声扰动。至关重要的是，所有这些都将噪声视为测试时的控制变量。相比之下，我们探索了评估和加权训练期间噪声的未被开发的潜力。

#### 训练时的损失重加权和噪声选择。

各种损失重加权和噪声调度策略提高了扩散训练的效率。大多数方法在*时间步级*操作，调整不同噪声幅度的相对重要性或采样频率。对于损失重加权，P2 [6] 优先考虑高度损坏的时间步长以学习全局概念，Min-SNR [14] 平衡跨时间步的优化冲突，Sun 和 Shi [46] 根据对数信噪比（log-SNR）级别上损失分布的方差动态调整权重。同时，噪声调度方法优化采样时间步的分布：Han 等人 [13] 引入了一种固定调度，专注于 $\log\text{SNR} \approx 0$ 区域。Kim 等人 [25] 使用课程学习逐步引入更困难的时间步长，Raya 等人 [35] 提出了一种基于熵减的动态、信息引导的分配。Wang 等人 [49] 从理论上分析了不同的总加权、时间和噪声调度。此外，EDM [20] 提供了对扩散设计的统一视角，表明噪声参数化、预处理和损失缩放的选择不自觉地定义了跨噪声水平的有效加权。然而，这些方法将给定时间步下的所有噪声实例视为具有同等信息量。最近的工作也开始质疑关于扩散模型中噪声的基本假设。例如，Sun 等人 [47] 表明，显式的噪声条件可能并非严格必要，模型即使无法访问时间步信息，往往也能保持有竞争力的性能。虽然这类工作研究了是否需要噪声信息，但它并未解决噪声*实例*的变化如何影响学习。相比之下，我们假设标准的条件设定，并关注一个互补的问题：*给定噪声，哪些特定的实现对训练最有用？* 在实例级噪声控制方面最接近的尝试是 Immiscible Diffusion [27, 28]，它使用最优传输将特定的噪声向量分配给数据样本，以最小化其预扩散欧几里得距离。虽然这加速了训练，但分配依赖于忽略模型实际学习动力学的静态指标。*我们的方法通过引入直接从模型泛化性能中学习的动态实例级噪声估值，超越了静态分配。*

#### 用于训练优化的元学习。

元学习技术已被探索用于自适应数据加权 [3]。早期方法 [36, 19] 使用在线元梯度动态加权训练样本。为了提高可扩展性和泛化能力，后续工作使用通过双水平优化训练的神经网络对这些加权机制进行参数化，从动态损失到权重映射 [40, 50] 演变为大规模数据集整理模型 [4]。我们首次将这种双水平优化范式引入扩散模型。至关重要的是，我们改变了估值的基本对象：我们的元网络旨在评估注入*噪声*实例的学习效用，而不是对训练*数据*的质量进行评分。

## 3 预备知识：扩散模型

扩散模型通过学习逆转数据逐渐腐化为噪声的过程来定义生成过程。在标准表述中，前向加噪过程逐步扰动干净的数据样本 $\mathbf{x}_0 \sim q(\mathbf{x}_0)$ 进入一系列越来越嘈杂的隐变量 $\{\mathbf{x}_t\}_{t=1}^T$。

#### 前向过程。

前向扩散过程是一个马尔可夫链，它在每个时间步长添加高斯噪声 $q(\mathbf{x}_t \mid \mathbf{x}_{t-1}) = \mathcal{N}(\mathbf{x}_t; \sqrt{1-\beta_t}\, \mathbf{x}_{t-1}, \beta_t \mathbf{I})$，其中 $t=1, \dots, T$，$\{\beta_t\}_{t=1}^T$ 是预定义的方差调度，控制每一步的噪声幅度。直观地说，这个过程逐渐破坏 $\mathbf{x}_0$ 中的结构，最终产生接近各向同性的高斯分布。

这一构造的一个关键特性是，任何时间步 $t$ 的边际分布 $q(\mathbf{x}_t \mid \mathbf{x}_0) = \mathcal{N}(\mathbf{x}_t; \sqrt{\bar{\alpha}_t}\mathbf{x}_0, (1-\bar{\alpha}_t)\mathbf{I})$ 都有闭式表达式，其中 $\alpha_t = 1-\beta_t$ 且 $\bar{\alpha}_t = \prod_{s=1}^t \alpha_s$。该表达式允许我们直接从 $\mathbf{x}_0$ 采样 $\mathbf{x}_t$，而无需模拟完整的马尔可夫链，这对高效训练至关重要。

#### 反向过程和训练目标。

生成模型被训练以逆转此加噪过程。与其直接参数化反向转移 $p_\theta(\mathbf{x}_{t-1} \mid \mathbf{x}_t)$，不如训练神经网络 $\epsilon_\theta(\mathbf{x}_t, t)$ 以预测为产生 $\mathbf{x}_t$ 而添加的噪声。具体而言，给定噪声样本 $\mathbf{x}_t = \sqrt{\bar{\alpha}_t}\mathbf{x}_0 + \sqrt{1-\bar{\alpha}_t}\, \epsilon$，其中 $\epsilon \sim \mathcal{N}(0, I)$，模型使用简单的均方误差（MSE）目标进行训练：

$$
\mathcal{L}_{\text{diff}}(\theta) = \mathbb{E}_{t, \mathbf{x}_0, \epsilon} \left[ \|\epsilon - \epsilon_\theta(\mathbf{x}_t, t)\|^2 \right]. \quad (1)
$$

训练过程通过采样三元组 $(\mathbf{x}_0, t, \epsilon)$，使用上述闭式表达式构造 $\mathbf{x}_t$，并最小化真实噪声 $\epsilon$ 与预测噪声 $\epsilon_\theta(\mathbf{x}_t, t)$ 之间的差异来进行。

#### 条件训练和无分类器引导。

在条件扩散模型中，去噪网络额外接收条件 $c \sim p(c)$（例如，文本提示）作为输入，并