Hierarchical Variational Policies for Reward-Guided Diffusion

会议: NeurIPS 2026
arXiv: 2605.21661
代码: https://github.com/czi-ai/oc-guidance
领域: 图像生成
关键词: 扩散模型, 奖励引导, 逆问题, 分层变分推断, 摊销控制

一句话总结

针对预训练扩散模型测试时引导优化计算开销巨大的瓶颈，提出将引导采样形式化为分层变分策略（HVP），通过两阶段摊销学习初始噪声分布与步进随机控制策略，以单次前向采样（AHVP）实现提速 5 倍以上并保持 SOTA 感知质量，同时支持轻量测试时微调（SHVP）。

研究背景与动机

预训练扩散模型作为强大的生成先验，已被广泛应用于下游条件生成与逆问题求解（如超分辨率、图像修复、HDR 复原以及文本对齐）。为了将无条件扩散模型引导至满足特定测量约束或下游奖励的目标分布，测试时引导（Test-Time Guidance）技术成为主流路线。这类方法无需为每个新任务重新训练庞大的条件扩散模型，而是直接利用已知的可微似然函数或显式奖励信号，在反向去噪采样过程中调整去噪轨迹。

然而，现有测试时引导方法普遍面临严重的计算效率瓶颈。无论是基于 Tweedie 估计的一阶近似方法，还是在反向去噪每一步嵌入内层优化循环的测试时扩展（Test-Time Scaling）算法，都需要在推理阶段重复计算网络雅可比矩阵、反向传播穿透长扩散链或执行密集的梯度更新。在少步数快速扩散采样或高分辨率场景下，多模态后验分布的复杂性进一步加剧，使得传统优化方法在极端步长下极易崩塌或产生严重伪影，无法满足实时部署与规模化落地的需求。

解决该计算开销的自然思路是将推断计算转移至离线训练阶段，即通过神经网络摊销（Amortization）去噪控制信号。然而，病态逆问题与大步长去噪天然诱导高度多模态的后验分布，简单的单峰高斯策略无法充分建模多解空间。核心 idea：将测试时引导重构为潜变量分层变分策略（HVP），在保持预训练去噪器完全冻结的前提下，先通过两阶段学习分别摊销初始全局噪声分布与步进随机控制，使条件逆问题求解简化为轻量前向推演（AHVP），并可无缝叠加轻量测试时梯度微调（SHVP）。

方法详解

整体框架

本文将观测条件 \(\bm{y}\) 下扩散生成反向轨迹的后验推断问题，形式化为在无条件扩散先验动力学之上的受控马尔可夫决策过程。整个框架在保持预训练无条件扩散模型参数 \(\bm{\mu}_{\text{pretrained}}\) 完全冻结的前提下，通过引入辅助潜变量控制向量 \(\mathbf{u}_t\) 扩展变分分布族，形成分层变分策略（Hierarchical Variational Policy, HVP）。

在推断阶段，算法通过已训练的轻量化 DiT 网络依次执行生成：首先由初始噪声网络根据测量 \(\bm{y}\) 预测出引导的初始噪声 \(\mathbf{x}_T\)，随后在每一步去噪过程中，步进随机控制器根据当前含噪状态 \(\mathbf{x}_t\)、历史控制累计及 \(\bm{y}\) 预测控制增量 \(\mathbf{u}_t\)。该控制向量作为加性修正项直接注入预训练去噪网络输入中，驱动状态完成一步条件反向去噪更新。整体采样完全摆脱测试时内层循环迭代，呈现单次前向推演的极高效率。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["观测条件 y<br/>降质图像 / 奖励约束"] --> B["初始噪声预测网络 E_phi<br/>生成全局对齐初始噪声 x_T"]
    B --> C["步进随机控制器 pi_psi<br/>状态 x_t 与历史控制采样 u_t"]
    C --> D["加性状态校正注入<br/>x_t + gamma * u_t"]
    D --> E["冻结预训练去噪器<br/>mu_pretrained 前向计算单步更新"]
    E -->|进入下一步 t-1| C
    E -->|完成 T 步去噪| F["最终高质量生成样本 x_0<br/>可选：SHVP 轻量测试时优化微调"]

关键设计

1. 分层潜变量变分形式与可处理熵下界：解决大步长采样下的多模态后验表达

病态逆问题由于测量不完备（例如 90% 随机掩码补全）存在大量视觉上自洽的解空间，在大步长反向去噪时状态转移后验极易呈现多峰分布。直接采用单峰高斯近似会导致严重的模式坍缩与伪影。为此，HVP 引入步进潜变量控制向量 \(\mathbf{u}_{1: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}\)。为了化解混合模型熵无法闭式求值的计算困境，本文利用辅助推断后验分布 \(\bar{r}(\mathbf{u}_{1:T}\mid\mathbf{x}_{0:T},\bm{y})=\prod_t \mathcal{N}(\mathbf{0},\mathbf{I}_d)\) 构造变分熵下界，推导出包含奖励项与三项正则约束的全局可处理证据下界（ELBO）：

\[\log p(\bm{y}) \ge \mathbb{E}_{\mathbf{x}_0 \sim q}[\log p(\bm{y}\mid\mathbf{x}_0)] - D_{\mathrm{KL}}(q(\mathbf{x}_T\mid\bm{y})\parallel p(\mathbf{x}_T)) - \sum_{t=1}^T \mathbb{E}[D_{\mathrm{KL}}(q(\mathbf{x}_{t-1}\mid\mathbf{x}_t,\mathbf{u}_t)\parallel p(\mathbf{x}_{t-1}\mid\mathbf{x}_t))] - \sum_{t=1}^T \mathbb{E}[D_{\mathrm{KL}}(q(\mathbf{u}_t\mid\mathbf{x}_t,\mathbf{u}_{>t},\bm{y})\parallel \bar{r}(\mathbf{u}_t))]\]

该目标由四部分组成：终末似然/奖励对齐项促使最终解满足条件测量；初始噪声 KL 正则约束初始潜分布贴近标准高斯先验；状态转移 KL 项强制修正后的去噪转移不偏离原扩散骨干的数据流形，防止奖励黑客（Reward Hacking）；控制先验正则项限制控制增量的能量范数。

2. 加性状态注入与轻量 DiT 架构：复用骨干强表征的轻量化参数化

为了最大程度保留预训练大模型的高阶表征能力而不产生重新微调庞大骨干的额外负担，状态转移分布被显式定义为经加性位移调制的高斯核：

\[q(\mathbf{x}_{t-1}\mid\mathbf{x}_t,\mathbf{u}_t,\bm{y}) = \mathcal{N}(\bm{\mu}_{\text{pretrained}}(\mathbf{x}_t + \gamma\mathbf{u}_t, t), \sigma_{t-1}^2\mathbf{I}_d)\]

其中 \(\gamma\) 为控制注入尺度因子。初始噪声预测器 \(\mathbf{x}_T = \epsilon + E_\phi(\bm{y},\epsilon)\) 与步进随机控制器 \(\mathbf{u}_t \sim \mathcal{N}(\bm{\pi}_\psi(\mathbf{x}_t,\mathbf{u}_{>t},\bm{y},t), \bar{\sigma}_t^2\mathbf{I}_d)\) 均基于参数量仅为 51M-92M 的轻量级 Diffusion Transformer（DiT）搭建。控制器方差随反向去噪过程的时间步噪声水平单调衰减（\(\bar{\sigma}_t = \kappa \sigma_t\)），从而在去噪初期赋予高随机性以探索多峰构型，在去噪后期收敛确定性以保留精细纹理。

3. 两阶段离线解耦训练：宏观拓扑捕获与微观残差并行精化

在训练策略上，若直接端到端通过全长扩散计算图反向传播训练所有参数，不仅显存开销巨大，且长轨迹梯度极易耗散或爆炸。本文提出分层解耦的两阶段训练方案： - 阶段 1（初始噪声策略学习）：固定步进控制器为零，通过简化的反向链直接训练初始噪声预测网络 \(E_\phi\)。该阶段以全局低频结构恢复为主，引导生成轨迹落入符合条件约束的高概率吸引盆地； - 阶段 2（步进控制器并行优化）：冻结阶段 1 的初始网络，训练步进控制器 \(\bm{\pi}_\psi\)。为了彻底摆脱链式反向传播，利用单步 Tweedie 估计 \(\hat{\mathbf{x}}_0\) 替代真值回传构建局部代理似然 \(\log p(\bm{y}\mid\hat{\mathbf{x}}_0)\)，使得所有时间步 \(t \in [1,T]\) 上的控制网络梯度可以并行独立反向更新，大幅压缩离线训练时间与显存占用。

4. 半摊销混合微调（SHVP）：消除有限网络容量的摊销误差

纯前向的完全摊销策略（AHVP）因策略网络参数容量有限和训练数据分布覆盖限制，天然存在一定的摊销误差（Amortization Error）。为追求极高感知保真度，作者进一步提出了半摊销变体 SHVP（Semi-Amortized HVP）：以 AHVP 预测的先验控制值作为极佳初始化点，在测试阶段允许对控制向量 \(\mathbf{u}_t\) 执行有限几步梯度上升优化：

\[\mathbf{u}_t^* = \arg\max_{\mathbf{u}_t} \mathbb{E}_q \left[ \log p(\bm{y}\mid\hat{\mathbf{x}}_0) - \lambda_2 \|\bm{\mu}(\mathbf{x}_t + \gamma\mathbf{u}_t, t) - \bm{\mu}(\mathbf{x}_t, t)\|_2^2 - \lambda_3 \|\mathbf{u}_t\|_2^2 \right]\]

相较于传统从随机初始点开始优化搜索的测试时方法，高质量的摊销初始化使优化过程能够在极少步数内快速收敛至高保真局部最优点，实现感知质量与推理速度的弹性可调。

损失函数 / 训练策略

阶段 1 针对噪声超网络采用终端测量均方误差加上隐空间能量正则化：\(\mathcal{L}_{\text{Stage-1}} = w_T \|\bm{y} - \mathcal{A}(\hat{\mathbf{x}}_0)\|_2^2 + w_{\text{control}} \|E_\phi(\bm{y},\epsilon)\|_2^2\)。阶段 2 在并行各时间步对控制器优化解耦目标：

\[\mathcal{L}_{\text{Stage-2}} = w_T \|\bm{y} - \mathcal{A}(\hat{\mathbf{x}}_0)\|_2^2 + w_{\text{score}} \|\bm{\mu}(\mathbf{x}_t + \gamma\mathbf{u}_t, t) - \bm{\mu}(\mathbf{x}_t, t)\|_2^2 + w_{\text{control}} \|\mathbf{u}_t\|_2^2\]

所有网络均使用 Adam 优化器，学习率设为 \(1\times 10^{-4}\)，在 FFHQ-256 和 ImageNet-256 训练集上分别仅需 1k 和 10k 样本即可充分收敛。

实验关键数据

主实验

在 FFHQ-256 与 ImageNet-256 数据集上，对四项经典图像逆问题（\(4\times\) 超分辨率、\(8\times\) 超分辨率、90% 随机掩码 Inpainting、HDR 复原）进行基准评测。为评估低推断预算下的性能，所有基线采样步数均压缩 5-6 倍（步数 \(\le 20\)，最少 8 步）。

任务	方法	FFHQ 时间(s) ↓	FFHQ LPIPS ↓	FFHQ FID ↓	ImageNet 时间(s) ↓	ImageNet LPIPS ↓	ImageNet FID ↓
SR (\(4\times\))	DAPS	18.9	0.174	50.45	30.9	0.249	80.35
	NDTM	4.1	0.064	40.62	5.8	0.304	84.09
	RED-Diff	0.3	0.224	86.98	0.7	0.378	125.08
	DDRM	0.3	0.136	71.49	0.7	0.338	101.69
	DPS	10.4	0.138	72.61	46.7	0.357	120.77
	HyperNoise	0.6	0.147	79.91	1.3	0.431	155.14
	AHVP (本文)	0.6	0.073	45.17	1.3	0.244	84.44
	SHVP (本文)	4.1	0.059	38.50	5.8	0.202	65.54
Inpainting (90%)	DAPS	18.9	0.266	84.43	30.9	0.397	222.05
	NDTM	4.1	0.105	63.92	5.8	0.287	141.81
	RED-Diff	0.3	0.828	306.64	0.7	0.931	321.10
	DDRM	0.3	0.749	246.97	0.7	0.901	355.76
	DPS	10.4	0.115	66.62	46.7	0.274	102.02
	HyperNoise	0.6	0.239	104.60	1.3	0.484	188.07
	AHVP (本文)	0.6	0.113	80.21	1.3	0.239	111.27
	SHVP (本文)	4.1	0.065	44.53	5.8	0.309	100.55
HDR	DAPS	18.9	0.145	36.22	30.9	0.186	47.95
	NDTM	4.1	0.089	50.85	5.8	0.141	56.04
	RED-Diff	0.3	0.163	74.83	0.7	0.202	84.47
	HyperNoise	0.6	0.193	102.90	1.3	0.468	196.98
	AHVP (本文)	0.6	0.089	45.76	1.3	0.138	61.01
	SHVP (本文)	4.1	0.073	43.35	5.8	0.135	50.63

消融实验

训练阶段有效性消融（SR \(4\times\) 任务）：

方法	阶段 1 (噪声策略)	阶段 2 (步进控制)	FFHQ LPIPS ↓	FFHQ FID ↓	ImageNet LPIPS ↓	ImageNet FID ↓	说明
AHVP	✓	✗	0.147	79.91	0.432	155.14	仅恢复全局粗糙轮廓
AHVP	✓	✓	0.073	45.17	0.244	84.44	补全细粒度高频纹理，LPIPS降低近50%
SHVP	✓	✗	0.066	40.95	0.225	66.42	仅微调无步进先验网络
SHVP	✓	✓	0.059	38.50	0.203	65.54	完整模型，达最优感知质量

控制策略分布随机性消融（SR \(4\times\) 任务）：

策略形式	FFHQ LPIPS ↓	FFHQ FID ↓	ImageNet LPIPS ↓	ImageNet FID ↓	说明
AHVP (确定性策略)	0.081	43.69	0.328	117.66	面对多样性高的复杂分布时表达力受限
AHVP (随机策略, 高斯分布)	0.073	45.17	0.244	84.44	ImageNet 上感知指标大幅领先

关键发现

• 推理加速与质量双优：完全摊销算法 AHVP 在单张 A6000 GPU 上的推理耗时仅为 0.6s（FFHQ）与 1.3s（ImageNet），相比 SOTA 测试时优化方法 NDTM 实现了超过 \(5\times\) 的端到端推理加速，而在复杂数据集 ImageNet 的超分辨率任务中，LPIPS 逆势领先（0.244 vs 0.304）。
• 两阶段解耦互补性明显：仅使用阶段 1 预测初始噪声（对应 HyperNoise 路线）只能重建粗粒度的面部或物体轮廓；阶段 2 引入步进控制器后，高频残差图清晰捕获毛发、布料和背景纹理，直接促成感知距离指标削减近半。
• 随机性对于多模态数据必不可少：在类别丰富、后验模式极其分散的 ImageNet 上，随机策略相对确定性策略实现 LPIPS 从 0.328 到 0.244 的断崖式优化，验证了分层潜变量对表达复杂条件流形的不可替代性。

亮点与洞察

• 提出了一套基于变分推断的扩散引导严格理论框架，巧妙地借助辅助推断后验分布推导出可处理的下界，将复杂的逆问题转化为控制能量与无条件扩散先验流形约束下的最优控制问题。
• 两阶段分离策略将「全局构型定位」与「局部细化驱动」成功解耦，并借助 Tweedie 单步伪似然实现控制器沿时间步的完全并行优化，兼具训练收敛速度与生成稳定性。
• 该框架建立了一座沟通完全摊销（AHVP）与测试时优化（SHVP）的桥梁，可根据实际工业场景对时延和画质的不同侧重实现性能连续调节。

局限与展望

• 当前推导与优化高度依赖终端奖励/似然函数的一阶可微性，尚无法直接原生应用于非可微度量（如黑盒分类置信度、离散规则惩罚）。
• 步进控制器采用简单的单峰高斯假定，未来若引入基于归一化流（Normalizing Flow）等更高阶的隐变量分布，有望进一步刻画极端多峰问题。
• 实验主要验证于图像超分和修复等视觉逆问题，后续可向文生图偏好对齐（如 ImageReward 对齐）及视频生成等高维连续动作空间任务做进一步迁移验证。

相关工作与启发

• vs DPS / DAPS: 后者属于典型的测试时优化算法，在去噪单步需要反复计算网络梯度甚至 MCMC 链式采样，推理成本高昂；本文通过神经网络摊销将单步优化收缩为前向推演，实现超 5 倍加速。
• vs NDTM: NDTM 同样提出了最优控制视角，但依赖测试时在线求解每步的最优控制向量 \(\mathbf{u}_t\)；本文将其进行变分参数化摊销为 AHVP，并以之为基底实现了收敛速度更快的半摊销采样 SHVP。
• vs HyperNoise: HyperNoise 仅在初始噪声处进行摊销预测，缺乏去噪过程中的步进精细纠偏；消融实验证实本文引入的两阶段步进潜变量控制能显著提升纹理清晰度并压低感知误差。

评分

• 新颖性: ⭐⭐⭐⭐⭐ 首次从分层变分推断与辅助分布熵下界推导出端到端摊销的扩散最优控制引导方案。
• 实验充分度: ⭐⭐⭐⭐⭐ 覆盖 4 类经典病态线性/非线性逆问题，跨双数据集提供详尽的速度-质量量化表与阶段消融。
• 写作质量: ⭐⭐⭐⭐⭐ 数学证明严密、结构层次分明，理论与实践落地映射清晰。
• 价值: ⭐⭐⭐⭐⭐ 为少步数扩散模型的实时受控生成与下游部署提供了极具实用价值的范式参考。

相关论文

• [ICML 2026] Recovering Hidden Reward in Diffusion-Based Policies
• [NeurIPS 2026] Reward Score Matching: Unifying Reward-based Fine-tuning for Flow and Diffusion Models
• [NeurIPS 2026] Are we really tilting? The mechanics of reward guidance in flow and diffusion models
• [NeurIPS 2026] ParetoSlider: Diffusion Models Post-Training for Continuous Reward Control
• [ICLR 2026] Improving Discrete Diffusion Unmasking Policies Beyond Explicit Reference Policies (UPO)