RNE: plug-and-play diffusion inference-time control and energy-based training¶

会议: ICLR 2026
arXiv: 2506.05668
代码: 无
领域: Image Generation / Diffusion Models
关键词: 扩散模型, 密度比估计, 推理时控制, 能量模型训练, Radon-Nikodym 导数

一句话总结¶

提出 Radon-Nikodym 估计器 (RNE)，基于路径分布间的密度比揭示边际密度与转移核的基本联系，提供统一的即插即用框架，同时实现扩散密度估计、推理时控制和能量扩散训练。

研究背景与动机¶

扩散模型通过逐步去噪生成数据，对应加噪过程的时间反转。在许多应用中，仅获取去噪核 (denoising kernels) 是不够的，我们需要知道生成轨迹上的边际密度 (marginal densities)。边际密度的知识可以支持：

密度估计：评估生成模型在任意点的概率密度

推理时控制 (inference-time control)：在生成过程中动态引导输出，如条件生成、组合多个模型

能量扩散训练：训练能量函数来参数化扩散模型

然而，获取扩散模型的边际密度是一个长期难题： - 直接计算需要积分所有可能的前向路径，计算上不可行 - 现有方法（如 ODE 概率流的似然估计）计算昂贵或精度不足 - 推理时控制方法通常需要特定假设（如 Tweedie 公式的近似），适用范围有限

核心洞察：利用 Radon-Nikodym 导数（密度比）的概念，可以建立边际密度与转移核之间的基本数学联系。这个联系无需训练额外模型，也不依赖特定的扩散模型架构。

方法详解¶

整体框架¶

RNE 的核心思想是在生成路径的分布层面工作：

定义两个路径分布：前向 (加噪) 路径分布 \(\mathbb{P}\) 和后向 (去噪) 路径分布 \(\mathbb{Q}\)
利用 Radon-Nikodym 定理，\(\mathbb{Q}\) 关于 \(\mathbb{P}\) 的密度比 \(\frac{d\mathbb{Q}}{d\mathbb{P}}\) 可以用转移核来表示
通过操纵这个密度比，实现对边际密度的估计和操控

关键设计¶

路径分布密度比：核心数学工具
- 考虑前向过程 \(q(x_0, x_1, ..., x_T)\) 和后向过程 \(p(x_T, x_{T-1}, ..., x_0)\)
- Radon-Nikodym 导数将两个路径分布联系起来：\(\frac{d\mathbb{Q}}{d\mathbb{P}}(x_{0:T}) = \frac{p(x_T) \prod_{t=1}^{T} p_\theta(x_{t-1}|x_t)}{q(x_0) \prod_{t=1}^{T} q(x_t|x_{t-1})}\)
- 这个比率可以分解为逐步的局部比率之积，每一步只涉及已知的转移核
- 设计动机：将全局密度问题转化为局部（逐步）的密度比估计
扩散密度估计：通过密度比获取边际密度
- 在任意中间时刻 \(t\)，状态 \(x_t\) 的边际密度可以通过路径密度比的部分乘积来估计
- 不需要额外训练密度模型，直接利用扩散模型自身的转移核
- 可以通过蒙特卡洛采样在路径空间中估计这个密度比
推理时控制：即插即用的条件生成
- 退火 (Annealing)：利用估计的边际密度作为能量函数的退火温度调度，实现更精确的条件采样
- 模型组合 (Model Composition)：将多个扩散模型的密度比相乘，实现多条件组合控制
- RNE 作为即插即用模块，不需要修改预训练扩散模型的权重
- 支持推理时缩放 (inference-time scaling)：更多的计算预算可以转化为更好的控制效果
能量扩散训练：将 RNE 用作训练正则化
- 传统能量扩散模型训练需要估计配分函数 (partition function)，计算困难
- RNE 提供了一种简单的正则化方式：利用密度比约束能量函数的训练
- 这避免了显式估计配分函数，简化了训练流程
模态无关性：不限于连续扩散
- RNE 的理论框架基于路径分布的一般概念
- 不仅适用于连续状态空间的扩散模型，也适用于离散扩散模型（如文本生成中的离散去噪）
- 这使得 RNE 成为一个跨模态通用的工具

损失函数 / 训练策略¶

RNE 在推理时控制场景下不需要额外训练（即插即用），在能量扩散训练场景下用作辅助正则化损失：

推理时控制：冻结预训练模型，仅通过密度比调节采样轨迹
能量训练正则化：在标准去噪损失基础上添加基于 RNE 的正则项，约束学到的能量函数与真实密度比一致

实验关键数据¶

主实验¶

任务	方法	关键指标	说明
退火采样	RNE	优于标准方法	更精确的条件采样
模型组合	RNE	多条件生成质量高	组合多个预训练模型
推理时缩放	RNE	性能随计算量提升	验证 scaling 特性
能量扩散训练	RNE 正则化	简单高效	无需估计配分函数

消融实验¶

配置	关键指标	说明
无 RNE 密度估计	密度估计不准确	缺少路径级别的密度比信息
有 RNE	密度估计精度提升	利用了转移核的完整信息
连续扩散	验证有效	标准场景
离散扩散	同样有效	验证模态无关性

关键发现¶

推理时控制的统一框架：RNE 将退火、模型组合等看似不同的推理时控制方法统一到密度比的视角下
推理时缩放：增加计算量（更多采样路径）可以持续提升控制精度，这与 inference-time compute scaling 的趋势一致
能量训练简化：RNE 正则化避免了传统能量模型训练中配分函数估计的困难
模态通用性：在连续和离散扩散模型上都验证了 RNE 的有效性

亮点与洞察¶

理论优美：利用 Radon-Nikodym 导数这一测度论基本工具，建立了扩散模型中看似独立的三个问题（密度估计、推理控制、能量训练）之间的统一联系
即插即用设计：不需要修改预训练模型，不需要训练额外的控制网络（如 ControlNet），大幅降低了使用门槛
路径分布视角的创新：不在单步转移层面工作，而是在完整轨迹的分布层面建立联系，这是一个更高层次的抽象
推理时缩放特性：呼应了当前 AI 社区对 test-time compute 和 inference-time scaling 的关注趋势
离散扩散的适用性：扩展了框架的适用范围，对文本和蛋白质等离散序列的扩散生成有潜在价值

局限与展望¶

蒙特卡洛估计的方差：路径空间中的密度比估计可能有较高方差，特别是在长扩散轨迹中
计算成本：虽然不需要额外训练，但推理时需要多次采样路径来估计密度比，增加了推理延迟
在大规模视觉生成上的验证不足：需要在如 Stable Diffusion、DALL-E 等大规模模型上验证
与已有推理控制方法的系统比较：如 Classifier Guidance、Classifier-Free Guidance、DPS 等的详细对比
理论与实践的差距：理论框架基于精确的前向/后向核，实际中使用的是学到的近似模型，近似误差的影响需要更深入分析

评分¶

新颖性: ⭐⭐⭐⭐⭐ — 用 Radon-Nikodym 导数统一三个独立问题是原创性强的理论贡献
实验充分度: ⭐⭐⭐ — 概念验证充分，但大规模验证不足
写作质量: ⭐⭐⭐⭐ — 理论清晰，框架统一
价值: ⭐⭐⭐⭐ — 即插即用特性和理论统一性有重要的实用和学术价值