Predict-Project-Renoise: Sampling Diffusion Models under Hard Constraints

会议: NeurIPS 2026
arXiv: 2601.21033
OpenReview: https://openreview.net/forum?id=0RZexLnRZZ
领域: 图像生成
关键词: 扩散模型、硬约束采样、预测-投影-再加噪、逆问题、物理约束生成

一句话总结

针对物理科学与工程中预训练扩散模型难以严格满足守恒律等硬约束的根本缺陷，本文从受约束前向过程的逆向动力学出发证明了以受约束边缘分布为唯一平稳分布的校正核，提出了免微调的 Predict-Project-Renoise（PPR）算法，通过去噪器隐式流形投影与前向再加噪的迭代，兼顾了近 80%+ 的可行解率与高质量分布保真度。

研究背景与动机

扩散概率模型在逼近高维数据分布方面展现出卓越的生成能力，但在流体力学、天气气候预报、医学成像等科学计算领域，生成样本往往必须严格满足硬约束 \(\mathcal{C}\)，例如不可压缩流体的无散度约束、质量与能量守恒定律、严格观测一致性或非线性动力系统边界条件。现有的约束生成范式主要分为微调训练与免训练后验采样两类：前者缺乏通用性且对不同约束需反复重训，同时仅能实现软约束；后者在推理阶段引导预训练无条件扩散模型，但现有方案存在严重的理论与实践瓶颈。

在免训练引导中，一类方法将硬约束建模为极端观测并采用贝叶斯分解近似后验得分，当约束似然退化为狄拉克 \(\delta\) 函数时，得分梯度出现数值奇异与爆炸，仅在可行集极小邻域内提供有效引导；另一类投影方法直接在反向扩散过程中执行几何投影，但往往面临三重困境：破坏样本的真实数据流形支撑集导致间断假影、破坏受约束后验分布导致多样性坍缩与模式崩溃、或随着逆向步推进再度违反约束。其本质矛盾在于：投影算子通常缺乏闭式解且只在干净数据空间有物理定义，若直接在带噪变量上投影会导致病态，若仅在去噪后点估计上投影则容易脱离数据流形，且确定性迭代投影会将采样器推向最近可行点的局部吸引子。

为克服现有方法在流形偏离、模式坍缩与约束违背之间的两难困境，本文从扩散边缘分布的测度几何出发，严格推导了由可行集诱导的受约束前向与反向扩散 SDE，并构造了一个以目标受约束边际分布为唯一平稳分布的校正马尔可夫转移核。核心 idea：将受约束采样形式化为在每个反向加噪层级逼近受约束边际分布，通过「预测器推进一步、通过去噪器向可行集隐式投影以维持流形支撑、再沿前向高斯核注入噪声」的交替迭代，以极小计算开销严格逼近真实受约束后验。

方法详解

整体框架

PPR（Predict-Project-Renoise）的输入为任意预训练无条件扩散去噪器 \(d_\theta(x_t)\)、标准无条件扩散预测器（如 DDPM 或 DDIM 步进）以及非负硬约束函数 \(c(x) \ge 0\)（其零水平集定义可行域 \(\mathcal{C} = \{x : c(x) = 0\}\)）；输出为严格满足约束且位于真实数据流形支撑集 \(\mathcal{S}\) 内的物理可行样本 \(x_0 \in \mathcal{S} \cap \mathcal{C}\)。

在反向扩散的每个离散时间步 \(t\)，标准预测器首先将上一时刻的状态推进一步得到先验带噪样本 \(x_t \sim p_t\)；随后算法进入一个循环迭代 \(M\) 次的校正循环：第一步利用优化器求解穿透去噪器投影问题 \(\min_x c(d_\theta(x))\) 得到最优带噪控制量 \(x_t^*\)，第二步利用前向扰动核以干净预测均值 \(d_\theta(x_t^*)\) 为中心重新注入标准方差的高斯噪声以完成再加噪。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    In["初始噪声状态<br/>x_T ~ N(0, sigma_T^2 I)"] --> StepT["反向时间步遍历<br/>t = T-1 down to 0"]
    StepT --> Pred["无条件扩散单步预测<br/>x_t <- Predict(x_t+1, d_theta(x_t+1), t+1)"]
    Pred --> LoopHead["校正循环迭代<br/>循环 M 次 (i = 1 ... M)"]
    LoopHead --> Proj["穿透去噪器的流形约束投影<br/>x_t* <- argmin_x c(d_theta(x))"]
    Proj --> Renoise["高斯前向核重采样再加噪<br/>x_t ~ p(x_t | x_0 = d_theta(x_t*))"]
    Renoise --> LoopCheck{"达到迭代上限 M ?"}
    LoopCheck -->|否| LoopHead
    LoopCheck -->|是| TimeCheck{"完成所有时间步 (t=0) ?"}
    TimeCheck -->|否| StepT
    TimeCheck -->|是| Out["物理可行高质量样本<br/>x_0 in S \cap C"]

关键设计

1. 受约束边际分布的平稳校正核构建：从理论上保证向真实约束测度收敛

传统几何投影算法仅将样本机械拉回可行超曲面，缺乏向受约束目标分布 \(p_t^{\mathcal{C}}(x_t) \propto \int p(x_t \mid x_0) p(x_0) \mathbb{1}_{\mathcal{C}}(x_0) \mathrm{d}x_0\) 混合的随机机制，容易诱发严重的模式坍缩。本文从理论上证明，理想的双步条件重采样核 $$ K_t(x_t' \mid x_t) = \int_{\mathcal{C}} p(x_t' \mid x_0) \frac{p(x_0 \mid x_t) \mathbb{1}_{\mathcal{C}}(x_0)}{\mathbb{P}(x_0 \in \mathcal{C} \mid x_t)} \mathrm{d}x_0 $$ 以受约束边际分布 \(p_t^{\mathcal{C}}\) 为其唯一的平稳分布。由于高斯重加噪转移核在全空间处处为正，马尔可夫链具备全空间不可约性与非周期性，因而在任意初值分布下重复作用 \(M\) 次必在全变差范数下依几何速率指数收敛到 \(p_t^{\mathcal{C}}\)。这为在每个去噪时间步注入高斯随机噪声提供了第一性原理依据，证明了再加噪步骤不仅不会抵消投影效果，反而是驱动采样遍历受约束后验测度空间的必要动力源。

2. 穿透去噪器的流形约束投影：通过隐式先验消除边界非物理间断

硬物理约束通常只对无噪干净物理场具备良好数学定义，直接在带噪变量 \(x_t\) 上最小化 \(c(x_t)\) 会因 \(x_t \notin \mathcal{S}\) 导致优化病态；而直接在单步去噪估计 \(x_0 = d_\theta(x_t)\) 上优化 \(c(x_0)\) 则缺乏正则约束，优化器倾向于在局部高频域寻找捷径，产生破坏物理连续性的伪影和剧烈跳变。PPR 提出了穿透去噪器的投影目标函数： $$ \textsc{Project}(x_t) \coloneqq \arg\min_x c(d_\theta(x)) $$ 该设计以当前带噪状态 \(x_t\) 初始化优化器（如 L-BFGS 或 Adam），通过反向传播将约束误差的梯度反传穿过神经网络去噪器 \(d_\theta\)。由于去噪器网络在训练集上学习了数据的全局先验流形与时空关联性，反传梯度会沿着学习到的流形切空间传播，引导输出 \(d_\theta(x_t^*)\) 严格落在可行集与真实数据流形交集 \(\mathcal{C} \cap \mathcal{S}\) 内，彻底根除了高维物理场生成中的空间间断缺陷。

3. 压缩映射误差分解与自适应计算权衡：收缩率与有界偏置的双重界定

由于精确计算受约束后验分布的积分不可行，PPR 采用狄拉克点估计 \(p^{\mathcal{C}}(x_0 \mid x_t) \approx \delta(x_0 - d_\theta(\textsc{Project}(x_t)))\)，将积分核退化为单点高斯扰动核 \(K_t^{\mathrm{PPR}}(x_t' \mid x_t) = p(x_t' \mid d_\theta(\textsc{Project}(x_t)))\)。本文推导了该近似核在 Wasserstein-2 距离下的严格误差界： $$ W_2\left(q_t^{(M)}, p_t^{\mathcal{C}}\right) \le (\alpha_t L_t)^M W_2\left(q_t^{(0)}, \tilde{p}_t^{\mathcal{C}}\right) + \frac{\Delta_t}{1 - \alpha_t L_t} $$ 其中 \(L_t\) 为穿透投影映射的 Lipschitz 常数，满足收缩条件 \(\alpha_t L_t < 1\)；\(\Delta_t\) 为点估计带来的单步不可约近似偏置。该误差上界清晰地表明：第一项混合误差随校正迭代步数 \(M\) 呈指数几何衰减，解释了多轮项目-再加噪循环对于消除预测器离散误差的必要性；第二项偏置由后验集中度决定。该机制使用户能够通过调节单步投影轮数与校正次数 \(M\)，在推理算力与物理可行率之间建立线性且单调可控的权衡杠杆。

损失函数 / 训练策略

PPR 为完全免微调（Training-free）的测试期推理采样算法，其依赖的无条件基础模型采用标准 EDM（Elucidating the Design Space of Diffusion Models）框架训练，优化加权去噪得分匹配目标函数： $$ \mathbb{E}{x_0 \sim p(x_0), \epsilon \sim \mathcal{N}(0, I)} \left[ | d\theta(x_0 + \sigma_t \epsilon) - x_0 |_2^2 \right] $$ 在推理阶段执行投影时，采用带线搜索的准牛顿 L-BFGS 优化器或 Adam，对于 2D 问题判定容差设为 \(c(x) \le 4 \times 10^{-6}\)，高维 KS 动力系统判定容差设为 \(c(x) \le 4 \times 10^{-4}\)。

实验关键数据

主实验

论文在三个代表性尺度上进行了严格评估：合成 2D 几何分布（包含 Uniform Checkerboard 与 Banana-GMM 两种先验，搭配 12 组高斯随机场非线性约束，共 24 种实验配置）、高维一维混沌动力系统 Kuramoto-Sivashinsky（KS）方程（多时刻非线性正弦观测硬约束，样本维度 \(128 \times 128\)），以及大规模全求气象模型 Appa（基于 ERA5 训练的潜在扩散模型，单步大气物理状态规模达 \(\mathcal{O}(10^8)\)，含 1.7 亿参数解码器）。

下表展示了在 Kuramoto-Sivashinsky（KS）混沌动力系统测试集上的关键定量对比结果（容差阈值 \(c(x) < 4 \times 10^{-4}\)，Max-Norm 越接近 1.0 代表样本平滑度越符合物理真值，偏离过大表明存在高频非连续间断）：

方法	可行率 (%) \(\uparrow\)	RMSE \(\downarrow\)	CRPS \(\downarrow\)	Skill \(\downarrow\)	Spread \(\downarrow\)	Max-Norm (理想值 1.0)	单次采样耗时 (s) \(\downarrow\)
PPR (本文)	82.7	0.503	0.177	0.365	0.385	1.80	80.2
DPS [Chung et al.]	0.0	1.089	0.352	0.666	0.865	17.11	25.5
MMPS [Rozet et al.]	1.0	0.830	0.265	0.519	0.672	19.28	307.0
TS (Trust Sampling)	0.0	0.614	0.206	0.427	0.487	2.01	398.0
CPS [Narasimhan et al.]	0.0	1.179	0.436	0.805	0.890	4.50	50.8
PDM [Christopher et al.]	0.8	1.137	0.463	0.831	0.798	4.48	296.0
HDC [Song et al.]	0.0	1.152	0.425	0.787	0.870	4.36	49.8

在二维合成数据分布（Data2D，24 个任务全平均，阈值 \(4 \times 10^{-6}\)，评估分布重合度的 \(k\)-NN 跨边缘率理想值为 0.5）的对比结果如下：

方法	可行解比例 (%) \(\uparrow\)	\(k\)-NN 交叉边缘率 (目标 0.5)	平均耗时 (s)
PPR (本文)	79.0%	0.375 \(\pm\) 0.026	14.9 \(\pm\) 0.6
DPS	1.0%	0.279 \(\pm\) 0.046	2.6 \(\pm\) 0.2
MMPS	3.6%	0.340 \(\pm\) 0.034	11.7 \(\pm\) 0.4
Trust Sampling (TS)	0.1%	0.124 \(\pm\) 0.027	3.8 \(\pm\) 0.2
CPS	31.6%	0.373 \(\pm\) 0.030	5.9 \(\pm\) 0.4
PDM	1.4%	0.002 \(\pm\) 0.000	10.8 \(\pm\) 0.1
HDC	30.3%	0.375 \(\pm\) 0.030	3.5 \(\pm\) 0.1

在千万维级全球气候模型 Appa 上的海平面气压、地表 2m 温度与风场分量上，PPR 同样以一致的计算预算显著降低了连续 24 小时预报的 CRPS 与 Skill 偏差（如 2m 气温 Skill 从 MMPS 的 1.536 K 降至 1.399 K，海平面气压 Skill 从 141.3 Pa 降至 111.2 Pa）。

消融实验

论文针对内部投影迭代步数（Projection Steps）以及是否经过去噪器反传进行了深入消融，证明了各核心模块的独立必要性（以 KS 与 Data2D 上的投影步数消融为例）：

投影步数	Data2D 可行率 (%) \(\uparrow\)	Data2D \(k\)-NN	KS 可行率 (%) \(\uparrow\)	KS RMSE \(\downarrow\)	KS CRPS \(\downarrow\)	KS Max-Norm	KS 耗时 (s) \(\downarrow\)
2 步	62.0	0.336 \(\pm\) 0.036	33.3	1.140 \(\pm\) 0.171	0.445 \(\pm\) 0.093	1.91 \(\pm\) 0.44	39.1 \(\pm\) 2.1
4 步	76.0	0.371 \(\pm\) 0.027	73.2	1.069 \(\pm\) 0.218	0.403 \(\pm\) 0.111	1.79 \(\pm\) 0.31	53.0 \(\pm\) 2.2
8 步 (默认)	78.5	0.374 \(\pm\) 0.027	87.2	1.012 \(\pm\) 0.605	0.361 \(\pm\) 0.134	1.65 \(\pm\) 0.29	83.7 \(\pm\) 6.6
16 步	78.2	0.373 \(\pm\) 0.027	94.0	0.925 \(\pm\) 0.340	0.337 \(\pm\) 0.150	1.60 \(\pm\) 0.21	150 \(\pm\) 25
32 步	78.0	0.370 \(\pm\) 0.029	95.9	0.904 \(\pm\) 0.440	0.330 \(\pm\) 0.162	1.61 \(\pm\) 0.41	305 \(\pm\) 97

关键发现

• 去噪器流形投影的决定性作用：若移除去噪器引导，直接在 \(x_0\) 空间投影（变体 \(x_0^M\)），KS 动力系统生成轨迹的可行率直接暴跌归零（0%），且轨迹导数范数 Max-Norm 飙升，在时空连续面上出现剧烈的跳跃撕裂；而经由 \(d_\theta\) 优化的 \(x_t^M\) 能保持 Max-Norm 在 1.6 左右，保持严密物理连续性。
• 再加噪（Renoise）是避免模式崩溃的生命线：完全移除再加噪步（\(M=0\)）会导致数值发散或崩溃为确定性局部极小值，无法对后验分布进行有效覆盖；增加重采样次数 \(M > 1\) 能单调降低约束偏离，使高维状态空间混合更充分。
• 优化器选型的鲁棒性差异：投影内部子问题采用准牛顿法 L-BFGS 的稳定度与收敛性显著优于一阶 Adam 与 SGD，一阶方法对步长极度敏感，容易跳出局部吸引子导致流形偏离。

亮点与洞察

• 严格将校正步落地为受约束平稳测度逼近：跳脱出直觉式「启发式几何拉回」的局限，建立了扩散采样校正子马尔可夫核收敛于受约束边缘分布的理论保证。
• 反向传播穿透预训练去噪网络充当隐式流形正则：巧妙利用预训练去噪器作为流形几何的隐式表示，使得对干净解的非线性约束优化自动获得流形结构保护，避免生成破碎间断解。
• 通用性极强且支持计算换精度：无需重训模型，对任意可微硬约束均可即插即用，通过调节内部优化轮数与重采样步数，可平滑控制计算耗时与物理守恒精度的折中。

局限与展望

• 依赖约束函数的一阶可微性：由于需要反向传播穿透去噪器网络计算梯度，当前框架无法直接适用于离散组合约束、逻辑规则或黑盒模拟器输出的硬约束。
• 计算代价随反向传播增加：穿透大型去噪主干网络（如大规模 Transformer/UNet）的高阶自动微分使得显存开销和推理耗时高于简单前向投影，未来需探索分步选择性投影（跳步校正）以加速推理。
• 不等式约束处理机制尚待拓展：当前实验主要聚焦于等式型物理守恒与精确观测，对于区间不等式硬约束（如非负浓度、能级上限），需进一步结合动量投影或屏障罚函数法设计专属算子。

相关工作与启发

• vs DPS [Chung et al., 2024] / MMPS [Rozet et al., 2025]：观测引导类方案假设似然具有连续支撑，在硬约束退化为 Dirac 极限时梯度出现病态爆炸与饱和；PPR 将约束视为几何可行集，利用确定性投影配合受控高斯再加噪规避了伪似然梯度的数值奇异性。
• vs CPS [Narasimhan et al., 2025] / PDM [Christopher et al., 2024]：现有投影扩散方法直接在干净样本 \(x_0\) 或带噪样本 \(x_t\) 上执行投影，前者脱离先验网络流形导致物理场时空出现撕裂间断，后者在带噪域中缺乏物理语义；PPR 创新采用穿透去噪器架构 \(\min c(d_\theta(x))\)，兼顾了物理定义域与数据流形自愈能力。
• vs Trust Sampling (TS) [Huang et al., 2024]：TS 基于噪声预测模长构建启发式置信域约束，容易由于过早截断导致投影不充分；PPR 直接对接通用二阶优化器并辅以理论证明的平稳重加噪核，在可行率与流形真实度上均大幅占优。

评分

• 新颖性: ⭐⭐⭐⭐⭐ 首次从受约束边缘扩散 SDE 视角严格推导并证明了保持真实约束边际测度的平稳校正核及其误差收敛上界。
• 实验充分度: ⭐⭐⭐⭐⭐ 涵盖 2D 几何多任务基准、一维高维偏微分方程系统以及高达 \(10^8\) 维度的全球气象真实模型，消融与基线极其详实。
• 写作质量: ⭐⭐⭐⭐⭐ 理论推导清晰严密，动机叙述直击痛点，算法流程与对比实验逻辑闭环完备。
• 价值: ⭐⭐⭐⭐⭐ 为科学机器学习与物理引导生成建模提供了坚实可靠的通用无监督硬约束后验采样基石。

相关论文

• [NeurIPS 2026] Temporal Alignment Guidance: On-manifold Sampling in Diffusion Models
• [ECCV 2026] Adaptive Noise Covariance Scheduling under Riemannian Metrics for Diffusion Models
• [ICLR 2026] A Probabilistic Hard Concept Bottleneck for Steerable Generative Models
• [NeurIPS 2026] CAB: Accelerating Flow and Diffusion Sampling via Rectification and Corrected Adams-Bashforth
• [NeurIPS 2026] Discrete Flow Matching: Convergence Guarantees Under Minimal Assumptions