面向奖励引导扩散的分层变分策略

# 面向奖励引导扩散的分层变分策略
来源：https://arxiv.org/html/2605.21661 Kushagra Pandey∗ 计算机科学系 加州大学尔湾分校 pandeyk1@uci\.edu&Farrin Marouf Sofian 计算机科学系 加州大学尔湾分校 fmaroufs@uci\.edu&Jan Niklas Groeneveld 计算机科学系 加州大学尔湾分校 jgroenev@uci\.edu&Felix Draxler 计算机科学系 加州大学尔湾分校 fdraxler@uci\.edu&Stephan Mandt 计算机科学系 加州大学尔湾分校 mandt@uci\.edu ###### 摘要 将预训练的扩散模型适配到下游目标（例如逆问题）通常需要昂贵的测试时引导或优化。我们提出一个原则性框架，用于在显著降低推理成本的情况下生成高质量的对齐奖励样本。我们的方法将测试时适配表述为分层变分模型，其中控制被摊销到一个轻量级但表达能力强的随机策略中。该表述天然支持少步扩散采样：大步长实现快速推理，而学习到的策略通过提供结构化的逐步控制来保持样本质量。由此产生的完全摊销采样器实现了强大的质量–速度权衡，在匹配或超越最近的测试时缩放基线的同时，所需计算量显著减少。例如，在 4×\\times 超分辨率任务上，我们的方法相比表现最佳的基线实现了更好的感知质量，并且推理速度快了 5 倍以上。我们进一步将方法扩展到半摊销机制，该机制将廉价的摊销提案与有限的测试时优化相结合，在多个具有挑战性的逆问题上实现了最先进的感知质量。 ## 1 引言 参考图 1 captionFigure 1:我们的方法（AHVP、SHVP）能够生成满足测量约束的高质量样本，且推理成本降低。基线方法常出现伪影（红色框），而我们的方法保留了精细细节（绿色框）。AHVP 在快速推理下提供强大的感知质量；SHVP 以适中的额外成本进一步提升质量。详见表格 1 (https://arxiv.org/html/2605.21661#S3.T1)。（放大查看最佳效果。） 扩散模型 (Ho et al., 2020 (https://arxiv.org/html/2605.21661#bib.bib23); Sohl-Dickstein et al., 2015 (https://arxiv.org/html/2605.21661#bib.bib21)）及相关家族 (Albergo et al., 2023 (https://arxiv.org/html/2605.21661#bib.bib121); Lipman et al., 2023 (https://arxiv.org/html/2605.21661#bib.bib63)）是下游生成任务的有力先验。扩散模型的测试时适配通过引导去噪过程来解决新颖任务，例如文本到图像对齐 (Ma et al., 2025 (https://arxiv.org/html/2605.21661#bib.bib202); Singhal et al., 2025 (https://arxiv.org/html/2605.21661#bib.bib203); Uehara et al., 2025 (https://arxiv.org/html/2605.21661#bib.bib199)）和逆问题 (Chung et al., 2022 (https://arxiv.org/html/2605.21661#bib.bib130); Mardani et al., 2024 (https://arxiv.org/html/2605.21661#bib.bib169); Zhang et al., 2025 (https://arxiv.org/html/2605.21661#bib.bib198); Pandey et al., 2025 (https://arxiv.org/html/2605.21661#bib.bib181); Zheng et al., 2025 (https://arxiv.org/html/2605.21661#bib.bib204); Geyfman et al., 2026 (https://arxiv.org/html/2605.21661#bib.bib201)），而无需昂贵地训练条件模型：只需指定任务，例如通过似然项或学习到的奖励 (Xu et al., 2023 (https://arxiv.org/html/2605.21661#bib.bib196); Wu et al., 2023 (https://arxiv.org/html/2605.21661#bib.bib3)），然后进行采样。然而，大多数现有方法依赖于重复的梯度评估或内部循环优化，这在高分辨率或实时场景中可能过于昂贵。一个自然的补救措施是将这种成本摊销：训练一个轻量级策略来学习引导去噪，将开销从推理转移到一次性训练阶段。然而，这种策略很难学习，特别是在病态逆问题或大步长扩散 (Salimans and Ho, 2022 (https://arxiv.org/html/2605.21661#bib.bib193); Song et al., 2023b (https://arxiv.org/html/2605.21661#bib.bib25); Yin et al., 2024 (https://arxiv.org/html/2605.21661#bib.bib192); Geng et al., 2025 (https://arxiv.org/html/2605.21661#bib.bib185)）中出现的多模态后验上。我们通过向策略中添加隐变量来解决这个问题，得到一个更灵活的分层变分模型 (Ranganath et al., 2016 (https://arxiv.org/html/2605.21661#bib.bib188)）。我们通过变分推断公式来训练这些策略 (Blei et al., 2017 (https://arxiv.org/html/2605.21661#bib.bib163); Zhang et al., 2018a (https://arxiv.org/html/2605.21661#bib.bib8)）：给定一个预训练的先验和一个下游奖励，我们近似去噪轨迹的后验。最近的几种方法 (Geyfman et al., 2026 (https://arxiv.org/html/2605.21661#bib.bib201); Pandey et al., 2025 (https://arxiv.org/html/2605.21661#bib.bib181); Mardani et al., 2024 (https://arxiv.org/html/2605.21661#bib.bib169)）也采用类似的观点，但仍在每个推理步骤中需要昂贵的优化。我们学习到的策略用单次前向传播取代了该成本，同时联合优化了所有采样步骤的控制信号。 我们的框架为构造摊销变分策略提供了一种通用且模块化的方法，适用于具有可微似然或奖励的广泛下游任务。作为一个具体实例，我们提出一个两阶段过程：首先学习一个最大化奖励的初始噪声分布，然后训练捕捉残差结构的逐步随机控制器。我们进一步将框架扩展到半摊销机制，该机制将廉价的摊销提案与额外的测试时精炼相结合。在具有挑战性的逆问题上，我们的方法在匹配或超越最先进的测试时缩放基线的同时，显著降低了推理成本。我们的贡献总结如下： - •我们引入了一个**统一框架**，将扩散模型中的测试时引导重新表述为分层变分策略上的推理，从而能够针对具有可微奖励的广泛下游任务进行原则性的摊销控制。 - •我们开发了**摊销 HVP (AHVP)**，这是一种新颖的两阶段方法，联合学习初始噪声分布和逐步随机策略，在单次前向展开中产生高质量的对齐奖励样本。 - •我们设计了**半摊销 HVP (SHVP)**，它将摊销提案与轻量级测试时精炼相结合，并**在几个具有挑战性的逆问题上实现了最先进的感知质量**，且额外成本适中。 - •我们在 FFHQ-256 和 ImageNet-256 上的四个逆问题上展示了**卓越的质量–速度权衡**，AHVP 在匹配或超越领先测试时方法的感知质量的同时，实现了**5 倍以上的推理加速**（图 1 (https://arxiv.org/html/2605.21661#S1.F1) 和图 2 (https://arxiv.org/html/2605.21661#S1.F2)）。 参考图 (a) SR (×4\\times 4) FFHQ-256 参考图 (b) SR (×4\\times 4) ImageNet-256 图 2：我们根据感知质量（LPIPS，越低越好）与每个样本的推理时间（对数尺度）比较了基于扩散的逆问题求解器。我们的方法（AHVP、SHVP）实现了良好的质量–速度权衡。结果展示了 ×4\\times 4 超分辨率任务。 ## 2 用于奖励引导的分层变分策略 直觉上，我们的方法通过**学习**如何推动无条件扩散轨迹来摊销引导。在每个去噪步骤，一个轻量级控制器观察当前状态、时间步和条件 y\{\\bm\{y\}\}，预测一个控制 u_t\{\\mathbf\{u\}\}\_\{t\}，然后将预训练的去噪器应用于受控状态。扩散模型保持固定；任务适配由学习到的初始噪声预测器和逐步控制器承载。固定去噪器和控制器共同起到任务条件采样器的作用，其中控制器注入来自 y\{\\bm\{y\}\} 的信息，同时保持预训练扩散模型不变。因此，推理变成了简单的条件展开，而不是内部循环优化。我们使用基于 Tweedie 估计的替代目标，通过奖励或似然目标来训练逐步控制器。下面我们将该过程形式化为受控去噪轨迹上的变分推断：首先作为**分层变分策略 (HVP)**（第 2.1 节 (https://arxiv.org/html/2605.21661#S2.SS1)），然后通过易于处理的学习目标（第 2.2 节 (https://arxiv.org/html/2605.21661#S2.SS2)），最后通过所得策略的摊销参数化（第 2.3 节 (https://arxiv.org/html/2605.21661#S2.SS3)）。 ### 2.1 分层变分策略 (HVP) 引导 我们现在通过分层变分近似的视角来呈现扩散模型中的引导。给定一个观察到的条件信号 y\{\\bm\{y\}\}，我们定义一个生成过程，逐步将噪声 x_T\{\\mathbf\{x\}\}\_\{T\} 转换为样本 x_0\{\\mathbf\{x\}\}\_\{0\}。更具体地说， p(\{\\mathbf\{x\}\}\_\{0:T\},\{\\bm\{y\}\}) = p(\{\\mathbf\{x\}\}\_\{T\}) \\left\[ \\prod\_{t=1}^T p(\{\\mathbf\{x\}\}\_{t-1} \\mid \{\\mathbf\{x\}\}_t) \\right\] p(\{\\bm\{y\}\} \\mid \{\\mathbf\{x\}\}\_0), (1) 其中 p(\{\\mathbf\{x\}\}\_T) 表示初始噪声上的先验（通常为标准高斯分布），p(\{\\mathbf\{x\}\}_{t-1} \\mid \{\\mathbf\{x\}\}_t) 直观上表示给定先前噪声状态 x_t 时推断去噪状态 x_{t-1} (Ho et al., 2020 (https://arxiv.org/html/2605.21661#bib.bib23); Albergo et al., 2023 (https://arxiv.org/html/2605.21661#bib.bib121); Lipman et al., 2023 (https://arxiv.org/html/2605.21661#bib.bib63)）。分布 p(\{\\bm\{y\}\} \\mid \{\\mathbf\{x\}\}_0) 表示在最终去噪状态 x_0 下观察到 y 的似然。在实践中，似然的形式是特定于应用的。例如，在逆问题（如修复）中，它可以定义为 p(\{\\bm\{y\}\} \\mid \{\\mathbf\{x\}\}_0) = \{\\mathcal\{N\}\}(\{\\bm\{y\}\}; \{\\mathcal\{A\}\}(\{\\mathbf\{x\}\}_0), \\sigma\_y^2 \{\\bm\{I\}\}\_d)，其中 \{\\mathcal\{A\}\} **破坏**去噪后的样本 x_0。在本文中，我们假设似然分布是**已知**且**可微**的。我们的目标是推断最大化 y 似然的中间状态 x_{0:T}。虽然主要感兴趣的对象是 p(\{\\mathbf\{x\}\}_0 \\mid \{\\bm\{y\}\})，但近似完整的轨迹后验 p(\{\\mathbf\{x\}\}_{0:T} \\mid \{\\bm\{y\}\}) 使我们能够将扩散生成动态重用于下游任务。由于这个后验难以处理，我们引入一个变分分布 q(\{\\mathbf\{x\}\}_{0:T} \\mid \{\\bm\{y\}\})。由此产生的证据下界 (ELBO) 为 \\log p(\{\\bm\{y\}\}) = \\log \\int q(\{\\mathbf\{x\}\}_{0:T} \\mid \{\\bm\{y\}\}) \\frac{p(\{\\mathbf\{x\}\}_{0:T},\{\\bm\{y\}\})}{q(\{\\mathbf\{x\}\}_{0:T} \\mid \{\\bm\{y\}\})} d\{\\mathbf\{x\}\}_{0:T} \\; \\geq \\mathbb{E}\_{q(\{\\mathbf\{x\}\}_{0:T} \\mid \{\\bm\{y\}\})} \\left[ \\log p(\{\\mathbf\{x\}\}_{0:T},\{\\bm\{y\}\}) - \\log q(\{\\mathbf\{x\}\}_{0:T} \\mid \{\\mathbf\{y\}\}) \\right]. (2) 变分分布 q(\{\\mathbf\{x\}\}_{0:T} \\mid \{\\bm\{y\}\}) 可以被解释为扩散轨迹上的**随机策略**。虽然式 (1) (https://arxiv.org/html/2605.21661#S2.E1) 定义了非引导的扩散动态，但优化式 (2) (https://arxiv.org/html/2605.21661#S2.E2) 中的 ELBO 相当于学习一个受控过程，将这些动态引导到与 y 一致的状态。 为了优化式 (2) (https://arxiv.org/html/2605.21661#S2.E2) 中的目标，我们将变分分布分解为 q(\{\\mathbf\{x\}\}_{0:T} \\mid \{\\bm\{y\}\}) = q(\{\\mathbf\{x\}\}_T \\mid \{\\bm\{y\}\}) \\prod_{t=1}^T q(\{\\mathbf\{x\}\}_{t-1} \\mid \{\\mathbf\{x\}\}_t, \{\\bm\{y\}\}). (3) 我们现在需要为每个 q(\{\\mathbf\{x\}\}_{t-1} \\mid \{\\mathbf\{x\}\}_t, \{\\bm\{y\}\}) 选择参数化。我们发现简单的选择（如单峰高斯分布）不够：例如，在修复这样的逆问题中，y 表示具有缺失像素的图像，给定 x_t 和 y 时，x_{t-1} 可能存在多个合理的补全，这会导致强多模态性。对于已知会导致更复杂且可能多模式的条件去噪分布 (Song et al., 2023b (https://arxiv.org/html/2605.21661#bib.bib25); Geng et al., 2025 (https://arxiv.org/html/2605.21661#bib.bib185); Boffi et al., 2025 (https://arxiv.org/html/2605.21661#bib.bib186); Zhou et al., 2025 (https://arxiv.org/html/2605.21661#bib.bib187)）的较大去噪步长，情况也是如此。因此，为了获得表达力强的变分策略，我们在每个时间步 (Ranganath et al., 2016 (https://arxiv.org/html/2605.21661#bib.bib188)）向变分分布添加额外的隐变量 u_t， q(\{\\mathbf\{x\}\}_{0:T} \\mid \{\\bm\{y\}\}) = \\int q(\{\\mathbf\{x\}\}_{0:T}, \{\\mathbf\{u\}\}_{1:T} \\mid \{\\bm\{y\}\}) \\, d\{\\mathbf\{u\}\}_{1:T}. (4) 通过最优控制的视角，这些隐变量 u_{1:T} 可以被解释为**随机控制**。由于这种具有边缘化隐控制的分层结构，我们将所得公式称为**分层变分策略 (HVP)**，定义为 q(\{\\mathbf\{x\}\}_{0:T}, \{\\mathbf\{u\}\}_{1:T} \\mid \{\\bm\{y\}\}) = q(\{\\mathbf\{x\}\}_T \\mid \{\\bm\{y\}\}) \\prod_{t=1}^T q(\{\\mathbf\{u\}\}_t \\mid \{\\mathbf\{x\}\}_t, \{\\mathbf\{u\}\}_{>t}, \{\\bm\{y\}\}) q(\{\\mathbf\{x\}\}_{t-1} \\mid \{\\mathbf\{x\}\}_t, \{\\mathbf\{u\}\}_t). (5) **初始噪声**分布 q(\{\\mathbf\{x\}\}_T \\mid \{\\bm\{y\}\}) 根据侧信息预测初始噪声。**逐步策略** q(\{\\mathbf\{u\}\}_t \\mid \{\\mathbf\{x\}\}_t, \{\\mathbf\{u\}\}_{>t}, \{\\bm\{y\}\}) 输出在时间 t 的条件于当前状态、侧信息和过去控制的控制；在我们早期的实验中，我们发现这种条件作用提高了训练稳定性和实证性能。定义变分分布的关键步骤是通过**条件状态转移**将其链接到预训练的去噪器， q(\{\\mathbf\{x\}\}_{t-1} \\mid \{\\mathbf\{x\}\}_t, \{\\mathbf\{u\}\}_t) = p(\{\\mathbf\{x\}\}_{t-1} \\mid \{\\mathbf\{x\}\}_t + \\gamma \{\\mathbf\{u\}\}_t), (6) 其中 p(\{\\mathbf\{x\}\}_{t-1} \\mid \{\\mathbf\{x\}\}_t) 就是生成模型的去噪分布，而 \\gamma 是一个调节参数。因此，控制 u_t 作为对状态 x_t 的加性校正，而 d