Learning to Sample From Diffusion Models via Inverse Reinforcement Learning¶

会议: NeurIPS 2026
arXiv: 2602.08689
OpenReview: https://openreview.net/forum?id=9RGG0tZyPh
领域: 图像生成
关键词: 扩散模型, 逆向强化学习, 采样优化, 状态占用度量, f-散度

一句话总结¶

将冻结去噪器的扩散采样过程建模为离散时间有限时域 MDP，利用逆向强化学习直接最小化采样轨迹与真实数据边缘分布间的 f-散度状态占用度量，以单次训练替代昂贵的采样超参数网格搜索。

研究背景与动机¶

扩散模型通过参数化神经网络（去噪器）引导的反向时间去噪过程，将高斯噪声逐步转化为复杂的高维数据分布。然而，当去噪网络训练完毕并固定后，推理阶段的采样算法仍包含大量高度敏感的动态超参数，例如无分类器引导尺度（CFG scale）、随机性注入调度（如 EDM 的 \(\gamma_{\text{EDM}}\)）以及重启采样（restart/renoise）步长。现有方法几乎完全依赖经验性的网格搜索来确定这些控制参数，这一过程需要针对每组候选超参数生成数万张样本计算评测指标，导致计算成本随参数维度呈指数级爆炸。以最先进的 EDM 框架为例，在 ImageNet-64 上仅针对 3 个随机性超参数做细粒度网格搜索就需要消耗多达 1437 EFLOPs，几乎等同于重新训练整个去噪网络所需的 1645 EFLOPs。

这种高昂调优开销的核心矛盾在于：扩散采样本质上是一个多步序贯轨迹控制问题，但每个去噪单步并没有天然的显式即时奖励函数；同时，在去噪轨迹优化中也无法获得所谓的“专家行为动作”或完整的专家推理轨迹，我们能唯一观测到的监督信号只有真实数据在不同噪声尺度下的边际分布。现有的强化学习与最优控制方法往往侧重于微调去噪网络权重本身，或者需要人工构造能量函数与启发式奖励，不仅增加了模型退化风险，也无法解决保持去噪网络冻结时的采样器通用控制问题。

本文的切入角度是采用仅依赖状态的逆向强化学习（State-only IRL），直接将扩散采样的超参数调度建模为离散时间有限时域的马尔可夫决策过程。核心 idea：在保持去噪器网络权重完全冻结的前提下，通过判别器估计各噪声层级的状态密度比，利用 PPO 裁剪策略梯度直接最小化生成轨迹与专家数据边际分布之间的 f-散度状态占用度量，实现端到端的样本级自适应采样控制并消除网格搜索开销。

方法详解¶

整体框架¶

本文将离散化去噪采样过程形式化为一个有限时域 MDP \((\mathcal{S}, \mathcal{A}, \mathcal{P}, \rho, T)\)。在每个时间步 \(t\)，系统状态定义为当前样本与其对应的噪声水平 \(s_t = (x_t, \sigma_t)\)，其中 \(x_t \in \mathbb{R}^d\)，\(\sigma_t \in \{\Sigma_0, \dots, \Sigma_N\}\)，终端 \(\sigma_t = \Sigma_0\) 设为吸收态；动作 \(a_t \in \mathcal{A}\) 编码离散的采样控制决策（例如离散引导尺度选择、噪声注入比例或回退步数）；环境状态转移 \(\mathcal{P}\) 由选定动作控制下的数值求解器单步更新 \(H(x_t, \sigma_t, a_t)\) 确定。整个架构由冻结的预训练去噪网络、特征提取骨干网络（冻结的 ADM Encoder）、输出动作分布的策略网络 \(\pi_\theta(a_t \mid s_t)\) 以及估计状态密度比的判别器 \(D_\phi(x, \sigma)\) 协同构成。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["初始高斯噪声 s_0=(x_0, Σ_N)"] --> B["阶段1：状态感知与动作调度<br/>冻结ADM编码特征 + 策略网络输出离散控制动作"]
    B --> C["阶段2：冻结去噪器数值更新<br/>执行CFG引导/随机性注入/重加噪转移至下一状态"]
    C --> D["阶段3：判别器密度比估计<br/>多尺度判别器区分真实加噪样本与生成状态"]
    D --> E["阶段4：f-散度占用度量优化<br/>累积优势函数计算 + PPO裁剪策略梯度更新"]
    E -->|未达到吸收态 Σ_0| B
    E -->|到达终止态| F["输出高质量生成样本 x_T"]

关键设计¶

1. 状态占用度量与多尺度边缘分布匹配：无需专家动作的无奖励对齐

在缺少即时奖励与专家轨迹的前提下，直接优化采样轨迹需要匹配策略诱导的状态占用度量 \(\mu_\theta(s)\) 与真实数据的专家占用度量 \(\mu_E(s)\)。状态空间连续且包含噪声尺度，占用度量可正交分解为噪声能级权重与条件特征分布之积：\(\mu_\theta(x, \sigma) = w_\theta(\sigma) \times p_\theta(x \mid \sigma)\)。专家占用度量 \(\mu_E(s) = w_E(\sigma) \times p_E(x \mid \sigma)\) 中，每个能级上的专家条件分布 \(p_E(x \mid \sigma)\) 可直接通过向无噪声真实数据注入对应方差高斯噪声解析获得，而权重 \(w_E(\sigma)\) 集中在终端态以稠密化最终生成目标的监督信号。通过最小化二者间的 f-散度 \(\mathcal{D}_f(\mu_E \parallel \mu_\theta)\)，直接引导采样轨迹流向高质量数据流形。

2. 占用度量 f-散度策略梯度推导：从密度比到无奖励策略更新

为了避开繁琐的奖励函数逆向求解，论文推导了 f-散度目标关于策略参数 \(\theta\) 的解析策略梯度定理。利用链式法则与占用度量关于轨迹概率的展开，梯度可精确表示为轨迹上动作对数似然梯度与后向时间累积密度比变换函数的加权期望：

\[\nabla_\theta \mathcal{D}_f(\mu_E \parallel \mu_\theta) = \mathbb{E}_{\bm{s}_{1:T} \sim p_\theta} \left[ \frac{1}{T} \sum_{t=1}^T \nabla_\theta \log \pi_\theta(a_t \mid s_t) \sum_{t' \ge t} h_f \left( \frac{\mu_E(s_{t'})}{\mu_\theta(s_{t'})} \right) \right]\]

其中生成函数变换定义为 \(h_f(u) = f(u) - u f'(u)\)。该公式表明，未来时间步的状态密度比偏差能够直接作为即时优势信号反向传递，无需定义显式奖励函数即可端到端驱动策略网络更新。

3. 判别器辅助的密度比估计与 PPO 裁剪稳定优化：克服多步采样梯度方差

计算状态占用度量比例 \(\frac{\mu_E(x, \sigma)}{\mu_\theta(x, \sigma)}\) 需要精确评估高维分布密度比。论文引入以噪声层级为条件的二值分类判别器 \(D_\phi(x, \sigma)\)，通过交叉熵损失区分真实加噪样本与当前策略生成的中间状态。根据贝叶斯最优分类器原理，条件概率比值满足 \(\frac{p_E(x \mid \sigma)}{p_\theta(x \mid \sigma)} = \exp(\tilde{D}_\phi(x, \sigma))\)（\(\tilde{D}_\phi\) 为 logits 输出），再乘以重要性权重修正项 \(\frac{w_E(\sigma)}{w_\theta(\sigma)}\) 即可得到全局占用度量比。在优化策略时，针对长轨迹重采样导致的梯度方差发散问题，采用 PPO 风格的概率比率裁剪目标与全局批次均值基线中心化，确保策略更新的平滑收敛与数值稳定。

损失函数 / 训练策略¶

判别器 \(D_\phi\) 在每个噪声能级均匀采样的条件下，通过最小化标准二值交叉熵进行多能级判别训练：

\[\mathcal{L}_D(\phi) = \mathbb{E}_{\sigma \sim \mathbb{U}[\Sigma_0, \dots, \Sigma_N]} \left[ -\mathbb{E}_{x \sim p_E(\cdot \mid \sigma)} [\log D_\phi(x, \sigma)] - \mathbb{E}_{x \sim p_\theta(\cdot \mid \sigma)} [\log (1 - D_\phi(x, \sigma))] \right]\]

策略更新则采用 PPO 裁剪优势损失，在重采样轨迹缓冲区中最小化：

\[\hat{\mathcal{L}}_f^{\text{PPO}}(\theta) = \frac{1}{N_{\text{batch}}} \sum_{i=1}^{N_{\text{batch}}} \max \left( r_i(\theta) \hat{A}_i, \, \text{clip}(r_i(\theta), 1-\epsilon, 1+\epsilon) \hat{A}_i \right)\]

其中优势函数 \(\hat{A}_t = \frac{1}{T}(A_t - \bar{A})\) 为累积学习信号 \(A_t = \sum_{t' \ge t} h_f \left( \frac{\mu_E(s_{t'})}{\mu_\theta(s_{t'})} \right)\) 经批次均值减法中心化后的稳定项，重要性采样比率为 \(r_i(\theta) = \frac{\pi_\theta(a_i \mid s_i)}{\pi_{\theta_0}(a_i \mid s_i)}\)。

实验关键数据¶

主实验¶

在 CIFAR-10、FFHQ-64 和 ImageNet-64 三个基准数据集上，分别针对自适应 CFG 引导、随机性注入 \(\gamma_{\text{EDM}}\) 以及重加噪重启（Renoise）三种主流采样策略进行测试，对比预训练 EDM 的最佳网格搜索超参数基线。

数据集与策略	指标	本文 (KL)	本文 (rKL)	之前SOTA (EDM基线)	表现分析
CIFAR-10 (CFG)	FID \(\downarrow\) / Prec \(\uparrow\) / Rec \(\uparrow\)	2.02 / 79.8 / 71.9	2.38 / 81.2 / 70.6	2.07 / 80.6 / 70.9	KL散度提升覆盖度，rKL更偏向高保真模式
CIFAR-10 (\(\gamma_{\text{EDM}}\))	FID \(\downarrow\) / Prec \(\uparrow\) / Rec \(\uparrow\)	2.01 / 78.8 / 73.1	2.04 / 79.4 / 72.3	1.98 / 78.7 / 72.9	匹配精细微调的随机性水平，召回率提升0.2%
FFHQ-64 (CFG)	FID \(\downarrow\) / Prec \(\uparrow\) / Rec \(\uparrow\)	3.03 / 90.6 / 88.6	3.04 / 90.8 / 88.3	3.11 / 90.7 / 87.8	相比最佳固定引导（\(\omega=0.1\)），FID改善0.08
FFHQ-64 (\(\gamma_{\text{EDM}}\))	FID \(\downarrow\) / Prec \(\uparrow\) / Rec \(\uparrow\)	2.53 / 90.4 / 88.2	2.49 / 90.9 / 87.9	2.56 / 90.1 / 89.4	rKL在人脸生成上取得更佳FID与精确率
ImageNet-64 (Renoise)	FID \(\downarrow\) / NFE \(\downarrow\)	2.96 / 51.8	2.92 / 52.1	3.01 / 350.0	NFE骤降85%（从350步降至52步），FID反超

消融实验¶

在 CIFAR-10 自适应 CFG 任务上，对状态条件输入形式与特征骨干进行消融验证：

配置	FID \(\downarrow\)	Precision \(\uparrow\)	Recall \(\uparrow\)	说明
全状态条件 \(\pi(x, \sigma)\) (KL)	2.02	79.8	71.9	完整模型：每张图像样本独立自适应引导
仅噪声阶条件 \(\pi(\sigma)\) (KL)	2.26	80.3	71.1	仅依赖时间/噪声步长，退化为全局调度
显式类别条件 \(\pi(x, \sigma, c)\) (KL)	2.12	81.1	70.6	额外提供类别标签，与全状态表现相近
去除 ADM Encoder 骨干 (KL)	3.04	82.6	67.5	仅从原始像素输入学习，FID大幅劣化近1.0
动作空间基数 \(\lvert\mathcal{A}\rvert=5\)	2.20	79.8	72.4	动作离散粒度太粗，无法精细匹配
动作空间基数 \(\lvert\mathcal{A}\rvert=20\)	2.06	78.8	72.5	表现已接近 \(\lvert\mathcal{A}\rvert=100\) 的上限

关键发现¶

样本级自适应性至关重要：全状态条件 \(\pi(x, \sigma)\) 显著超越仅含噪声阶的策略 \(\pi(\sigma)\)。在中间噪声阶段（即图像拓扑结构定型的物种形成相/speciation phase），策略选择性地对部分模糊样本施加高引导强度，而保持其余样本低引导，从而在避免模式坍塌的同时兼顾多样性。
f-散度精准调控精度与召回率平衡：正向 KL 散度天然具备“模式覆盖（mode-covering）”特性，引导采样策略生成更多样化的边缘样本，使 Recall 显著提升；而反向 rKL 表现为“模式寻求（mode-seeking）”，集中于高概率密度区域，从而带来更高的 Precision。
算力开销大幅压缩：在 ImageNet-64 任务中，传统的 3 超参数细粒度网格搜索（\(8^3=512\) 次评估）总耗时高达 1437.4 EFLOPs，而本文仅需 160.6 EFLOPs（降低近 9 倍），且在推理阶段仅引入 16% 的计算开销。

亮点与洞察¶

无奖励的直接占用度量匹配：摆脱了传统逆向强化学习中“先学奖励函数再做 RL”的两阶段繁琐管线，数学上证明了占用度量散度直接可转化为解析策略梯度。
冻结去噪器的即插即用控制器：不改动大模型权重即可完成推理期轨迹控制，使大规模预训练扩散模型无需昂贵的重新微调即可适配不同的下游采样目标。
预训练特征表征的强力使能：冻结的 ADM Encoder 为策略网络和判别器提供了极高鲁棒性的语义隐空间，有效解决了高噪声状态下高维像素直接判别时信噪比极低导致的梯度发散问题。

局限与展望¶

依赖离散动作空间划分：当前控制决策均建立在预先离散化的超参数动作集合上，未能实现连续平滑的动态控制输出。
训练对策略初始启发式较敏感：非凸优化景观容易陷入局部极小，需要针对不同任务设定合理的初试策略偏置（例如在高噪声区域优先选择继续去噪）。
未来方向：探索连续动作空间的高阶强化学习算法，以及将该状态占用度量控制框架拓展到多模态扩散 Transformer（DiT）与流匹配（Flow Matching）生成模型中。

评分¶

新颖性: ⭐⭐⭐⭐⭐ 首次将状态边缘占用度量逆向强化学习引入冻结去噪器的采样控制，数学推导完备。
实验充分度: ⭐⭐⭐⭐⭐ 涵盖三种主流采样策略、三个经典图像数据集，提供详尽的算力、消融与理论验证。
写作质量: ⭐⭐⭐⭐⭐ 结构清晰，理论推导严谨，实验图表信息量极其丰富。
价值: ⭐⭐⭐⭐⭐ 彻底解决了扩散模型采样超参数调优的算力黑洞，具有显著的工程实践与学术指导价值。