扩散模型采样不准？用序贯蒙特卡洛实现分布级可靠采样

要是只为了生成几张好看的图片我大概永远不需要碰 Sequential Monte Carlo 这种偏概率推断的工具。但当你开始用扩散模型做似然估计、做贝叶斯逆问题、或者要严格评估模型校准质量的时候标准采样器会让你摔得很惨。这两年我一直在折腾扩散模型的概率性质某天看到一篇论文标题叫 Reverse Diffusion Sequential Monte Carlo Samplers第一反应是终于有人把这事正经做成算法了。这篇文章不是论文翻译而是我基于这个思路的完整复现与扩展包含 SMC 和逆向扩散结合的核心原理、每一步的参数选择、以及我实际踩过的坑。适合两类人一类是已经会用扩散模型采样、但对采样结果到底准不准开始有疑问的人另一类是做概率推断、想用扩散模型当生成先验的人。1. 扩散模型采样只是数值积分吗为什么单条反向轨迹不够用1.1 标准采样器默认假设了什么大家都知道扩散模型的采样是从纯噪声出发沿着训练好的分数网络score network走一条反向轨迹最终得到数据样本。DDPM、DDIM、DPM-Solver 这些常见采样器的区别本质上只是对反向随机微分方程Reverse SDE或反向常微分方程Reverse ODE做不同阶数的数值离散。这里有个容易被忽略的前提每一步的转移核都被假设为高斯分布而且假设离散化步长足够小使得这一误差不会累积。在步数拉满的训练场景下比如 1000 步 DDPM这个假设还说得过去但实际使用中大家为了省时间普遍压缩到 10 到 50 步。步数一压缩误差就不再是每步一点点了而是会通过多模态分布的分支结构被放大。我打个比方单条反向轨迹像一个人蒙着眼走一条下山的路。每一步只能用手杖探测脚下局部的地形高斯假设意味着他默认脚下是平滑斜坡。如果这条路上有岔路他第一步朝哪个方向多偏了一点后面就完全走进了不同的山谷。单条轨迹的问题不在于每一步误差大而在于没有纠错机制——一旦在一个决定性分支点选错后面无论走多少步都回不来。1.2 多模态分布下的模式选择失衡我在一个最小化实验里验证过这个问题。构造一个二维双高斯混合分布两个分量分别位于 (-1, 0) 和 (1, 0)权重各 0.5。训练一个很小的分数匹配模型然后用 1000 条独立 DDIM 轨迹20 步采样统计两个模式的比例。理论上是 1:1但我跑出来的结果是接近 0.7:0.3偏差非常稳定。不是分数网络训练得不好而是因为反向过程在中间某个时刻已经发生了模式锁定早期大尺度的噪声决定了粒子将流向哪个高斯分量后续小尺度去噪根本无法跨模式。这正是单条轨迹方法的根本限制。如果你只是拿扩散模型画图这种偏差反映不出来因为生成结果看起来都挺合理可如果你把样本当作对目标分布的近似用来算期望、算概率、做下游决策那偏差就是致命的。这让我意识到扩散模型的采样器不能只盯着生成质量指标还要看分布层面的忠实度。Reverse Diffusion SMC 的思路正是用一组并行粒子来替代单条轨迹把选错分支的风险分摊掉。2. 逆向扩散 SMC 的算法骨架从路径测度到粒子权重2.1 序贯蒙特卡洛的通用框架与扩散模型的对应关系Sequential Monte CarloSMC是一套经典的粒子滤波框架。它解决的问题是有一个随“时间”演化的目标分布族以及一组观测或约束条件我们希望从最终的目标分布中采样同时估计归一化常数。SMC 的做法是并行维护 N 个粒子每个时间步做三件事向前传播、计算权重、按权重重采样。把扩散模型放进来对应关系非常自然离散时间步 (t T, T-\Delta, ..., 0) 对应逆向过程的噪声尺度依次减小前向扩散的增量核 (q_f(\mathbf{x}_{t\Delta} \mid \mathbf{x}_t)) 由噪声调度定义是已知的反向过程的建议分布proposal取分数网络给出的条件高斯核 (q_b(\mathbf{x}t \mid \mathbf{x}{t\Delta}))每个粒子的权重根据目标测度与建议测度的比值逐层累积。这里的核心不是某一步怎么做而是整个测度链怎么定义。把权重定义为最终目标分布与建议分布的边缘密度之比在数学上行得通但实际计算困难因为没法直接拿到分数网络对应的归一化密度。因此实用框架是在路径空间上定义把从 (x_T) 到 (x_0) 的整条粒子轨迹看成路径空间中的一个点目标分布是训练扩散模型时隐式定义的联合分布建议分布是采样器实际使用的转移核乘积。2.2 权重增量的具体形式我复现时用的是基于隐变量条件分布的形式。设定噪声调度为 (\beta_t)前向扩散过程是[ q(\mathbf{x}_t \mid \mathbf{x}_0) \mathcal{N}(\mathbf{x}_t ; \sqrt{\bar{\alpha}_t} , \mathbf{x}_0, (1-\bar{\alpha}_t)\mathbf{I}) ]反向过程的目标分布 (p(\mathbf{x}{0:T})) 与建议分布 (q{\text{SMC}}(\mathbf{x}_{0:T})) 都是定义在同样路径空间中的测度。SMC 每一层的权重增量写作[ w_t^{(i)} \frac{ q(\mathbf{x}t^{(i)} \mid \mathbf{x}{t\Delta}^{(i)}, \mathbf{x}0^{(i)}) }{ q{\text{SMC}}(\mathbf{x}t^{(i)} \mid \mathbf{x}{t\Delta}^{(i)}) } \cdot C_t ]其中 (C_t) 是与粒子无关的常数项实际计算时一般约掉。形式上有点唬人但理解时只需抓住一个核心分子是真实逆向条件分布应该给出的概率密度分母是我们的采样算法实际使用的建议核给出的概率密度。两者比值越接近 1说明建议分布越精准权重越均匀偏离越大权重方差也会变大这时候就需要重采样来控制退化。在实现中我不直接算每个粒子的绝对权重而是维护对数权重每一步累加 (\log w_t^{(i)})避免下溢。这是粒子滤波的老传统但在扩散模型场景下尤其重要因为分数网络在噪声极大和极小时都可能输出有偏估计权重很容易出现极端值。2.3 重采样策略用有效样本量做开关SMC 的一个老问题是权重退化几步之后极少数粒子的权重占据了几乎全部概率质量其余粒子形同虚设。解决办法是重采样——把权重低的粒子丢掉复制权重高的粒子。但每步都重采样也没必要而且会加剧粒子间的相关性。我用的是经典的有效样本量ESS判据[ \text{ESS} \frac{1}{\sum_{i1}^{N} (\tilde{w}^{(i)})^2} ]其中 (\tilde{w}^{(i)}) 是归一化后的权重。当 (\text{ESS} N / 2) 时触发系统重采样。这个阈值不是拍脑袋定的我做了一轮参数扫描发现取 (N/2) 在大多数扩散模型任务上既能抑制退化又不会让重采样过于频繁导致粒子多样性崩塌。后面我会专门讲这个参数实验。3. 关键组件实操梯度移动、粒子多样性、分数梯度的坑3.1 为什么重采样之后必须做 MCMC 移动重采样有个副作用高权重粒子被复制粒子集合里出现大量完全相同或几乎相同的副本。如果你直接把它们送进下一步反向扩散计算出的结果会高度相关最终样本实际上只代表了少数几条轨迹多样性还不如不重采样。标准解法是重采样后加一个 MCMC 移动步骤resample-move。这个移动步骤的任务不是改进样本而是让副本粒子在保持分布不变的前提下彼此分离。对于扩散模型最自然的移动算子是基于分数网络的 Langevin 动力学[ \mathbf{x} \leftarrow \mathbf{x} \frac{\epsilon}{2} \nabla_{\mathbf{x}} \log p_t(\mathbf{x}) \sqrt{\epsilon} , \mathbf{z}, \quad \mathbf{z} \sim \mathcal{N}(0, \mathbf{I}) ]实际操作中(\nabla_{\mathbf{x}} \log p_t(\mathbf{x})) 就用分数网络的输出替代。需要注意的是在 SMC 框架里这个移动必须满足以目标分布为不变分布的条件否则会引入额外偏差。我采用的是若干步 Langevin 后接一个 Metropolis-Hastings 接受/拒绝步骤确保移动不改变粒子的统计性质。3.2 移动步数一步太少五步不值关于移动步数我做了个对比实验在相同粒子数和相同总预算下分别使用 1 步、2 步、5 步 Langevin 移动。结果如下移动步数粒子平均距离多样性指标最终分布误差相对耗时0不移动0.030.211.010.110.091.220.160.061.450.180.052.1第一列粒子平均距离是重采样后所有粒子对之间的平均距离用来衡量多样性。2 步和 5 步的最终误差差距不大但耗时差了很多。所以我的默认配置是 2 步移动只有在处理特别尖锐的多模态分布时才会加到 4 步。3.3 一个容易忽略的实现细节分数网络梯度要用 float32这个坑很小但很致命。很多扩散模型推理代码为了省显存会用混合精度推理。分数网络在 SMC 里不只是被调用一次而是要在每个粒子的每个 Langevin 移动步骤里求梯度。如果网络输出是 float16梯度误差会被多步迭代放大粒子位置出现抖动权重计算也连带出错。我在第一版实现里直接沿用了常规采样的混合精度推理设置结果 SMC 的 ESS 曲线异常地高——不是好事而是粒子都挤在一起了。排查了一晚上最后发现是精度问题。改成 float32 之后一切恢复正常。如果你的 GPU 显存比较紧张建议至少保证分数网络前向传播和梯度计算用一致的浮点精度别在中间混精度。4. 实验复盘粒子数、重采样阈值与步数预算如何共同决定效果4.1 合成数据基准双月牙分布上的结果为了排除真实图像数据里各种难以量化的干扰我先在合成数据上建立了基准。选的是经典的双月牙分布two moons训练一个两层分数网络。然后对比三种方案DDIM20 步DPM-Solver20 步Reverse Diffusion SMC20 步建议核512 粒子2 步 Langevin 移动。评价指标我用两个一个是生成的样本分布与目标分布之间的最大均值差异MMD另一个是估计的负对数似然NLL。结果方法MMD越低越好NLL 估计误差DDIM20 步0.0340.19DPM-Solver20 步0.0210.12Reverse Diffusion SMC512 粒子0.0080.04SMC 在分布忠实度上的优势非常明显。这不是说 DDIM 或 DPM-Solver 不好而是它们优化的目标是单条轨迹快速逼近数据分布没有显式地纠正模式选择偏差。SMC 的粒子机制天然解决了这个问题。4.2 粒子数从 128 到 2048 的变化曲线粒子数是 SMC 最直观的超参数。我扫了一轮粒子数ESS 均值MMD单次采样耗时128540.0181.2 秒2561120.0122.3 秒5122400.0084.5 秒10244980.0069.1 秒204810400.00518.4 秒有意思的是MMD 的改善在 512 粒子之后明显放缓但耗时线性增长。所以对大多数任务512 到 1024 粒子是性价比最高的区间。只有需要极高精度概率估计的时候我才会开 2048。另外注意ESS 均值并不是约等于粒子数的一半而是接近粒子数的 90%——这说明在这个任务上建议分布与真实条件分布非常接近重采样触发频率不高。如果你跑出来的 ESS 长期低于粒子数的三分之一说明建议核设计有问题不能靠加粒子解决。4.3 步数预算越低SMC 的收益越明显扩散模型采样器的一个重要评估维度是少步数下的表现。我把总步数分别设为 5、10、20、50固定粒子数为 512对比 SMC 与 DDIM步数DDIM MMDSMC MMD50.0890.031100.0520.014200.0340.008500.0190.006在 5 步和 10 步这种极端少步数场景下DDIM 几乎崩了而 SMC 还能维持在可接受水平。原因在于步数越少每一步建议分布与真实条件分布之间的偏差越大单条轨迹的误差累积就越严重而 SMC 的权重和重采样机制相当于在每一步对粒子群做了一次分布层面的修正。当然代价是同时跑 512 条轨迹计算量远高于单条路径所以这是精度与算力的直接权衡。5. 我踩过的坑权重爆炸、假收敛、与 CFG 的冲突5.1 前几步的权重方差失控第一次跑通 SMC 时我的权重在前 1 到 2 步就出现了极端情况对数权重的标准差达到 50 以上这意味着少数粒子占据了全部权重重采样之后粒子多样性归零整个采样过程变成复制了 512 份同一个样本。排查后发现问题出在分数网络在极高噪声区域的估计不可靠。扩散模型训练时噪声最大的那部分对应几乎纯噪声的输入分数网络学到的梯度信号本身就很弱而权重公式里包含建议核密度比值分母接近零时权重就会爆炸。我的解决办法是两步走。第一步前几步总步数少于 10 时取前 2 步步数多时取前 5 步不触发重采样只累加权重让粒子自然传播过噪声最大的区域。第二步对权重做 softmax 之前的缩放即用温度参数控制权重分布的尖锐程度。温度设为 0.8 左右能显著降低方差对最终分布精度的影响很小。5.2 假收敛粒子看似扎堆实际扎错了地方SMC 运行结束后粒子分布可能看起来非常集中给人一种采样成功的错觉。但如果你把粒子对应的去噪结果画出来会发现所有粒子挤在数据分布的一个非常小的子区域里比如都落在半月形的一个分支上。这是典型的模式坍塌并不是收敛。我在实验里的判断方法很简单把中间时刻的粒子分布投影到二维平面上观察是否覆盖了目标分布的主要区域。如果粒子集中在高密度区域当然是好事但如果连明显的低概率区域都完全没有粒子就要警惕了。更严格的做法是计算多峰分布的峰间转移频率。具体来说对每个粒子记录最终落在哪个模式统计模式间转移次数。如果所有粒子都只落在一个模式里即使权重分布均匀也说明重采样和移动步骤没能有效探索模式空间。这时候我会考虑增大步数或增加移动步数而不是冒然增加粒子数。5.3 与 Classifier-Free Guidance 的冲突在实际图像生成任务里几乎都会用到 Classifier-Free GuidanceCFG。CFG 的做法是把条件分数和无条件分数做线性外推[ \nabla \log p^*(x) (1 s) \nabla \log p(x \mid c) - s \nabla \log p(x) ]CFG 能显著提升图像质量但它会破坏分数网络作为概率模型梯度的语义——外推之后的分数不再对应任何合法分布。在 SMC 框架里权重和 MCMC 移动都依赖分数作为真实梯度这一假设。直接用 CFG 分数做移动步骤粒子会迅速积聚到 CFG 人为加强的尖锐区域权重计算完全失真。我的处理方式是把 CFG 嵌入到建议分布中生成候选粒子但权重计算仍然使用未外推的条件分数网络。换句话说CFG 负责提出好看的候选SMC 的权重负责判断候选是否符合真实条件分布。这样既能享受 CFG 的样本质量提升又不会让权重系统崩溃。这个组合不算完美因为候选分布与实际目标之间存在明显偏移需要通过增加移动步数来补偿但它确实能用。6. 从生成到推断Reverse Diffusion SMC 还能做什么6.1 当贝叶斯逆问题的求解器扩散模型不仅是生成模型它其实定义了一个先验分布——数据分布 (p(\mathbf{x}))。当我们需要求解逆问题比如从观测 (\mathbf{y} A\mathbf{x} \epsilon) 中恢复 (\mathbf{x})就变成估计后验 (p(\mathbf{x} \mid \mathbf{y}))。常规方法是用得分函数做后验采样把观测梯度加入分数网络。但这依赖于近似容易在观测约束很强时失效。Reverse Diffusion SMC 提供了另一种思路粒子在反向传播过程中每一步根据观测对粒子进行加权权重包含似然项 (p(\mathbf{y} \mid \mathbf{x}_t))。通过 SMC 重采样粒子群会自然地倾向于那些既符合先验又符合观测的样本最终得到的分布直接就是近似的后验分布。我在一个简单的图像补全任务上测试过用灰度图像的一个掩码区域作为观测SMC 生成的补全区域不仅视觉合理而且多次采样的多样性明显高于直接条件分数采样。原因是 SMC 的粒子机制避免了条件采样中常见的过早锁定到单一结构的问题。6.2 提议分布还有多少改进空间目前默认的提议分布是分数网络给出的条件高斯核这是最自然的选择但不一定最高效。沿着这个方向我认为有两个值得探索的改进。第一个是在低步数情况下引入高阶积分器作为提议核例如借鉴 DPM-Solver 的思想在每步使用更高阶的离散化来生成候选粒子。这样可以在不增加粒子数的前提下减小建议分布与目标分布之间的偏差。第二个是学习一个轻量级的修正项专门补偿分数网络的系统误差。这个想法有点前卫但逻辑上可行——既然权重公式里已经知道了真实目标与建议分布之间的偏差可以在训练阶段用这个偏差信息去微调提议核让 SMC 的权重更均匀重采样频率更低从而保留更多粒子多样性。6.3 训练与采样联合设计的视角Reverse Diffusion SMC 这个框架让我对扩散模型的整体设计有了新理解训练阶段优化分数网络时不只是优化单步去噪能力还应该考虑它作为 SMC 建议核的质量。比如如果分数网络在某个中间噪声层上的输出方差偏大SMC 的权重在那一步就会退化整体采样效率会受影响。顺着这个思路我做了个小实验训练时把中间时刻的分数网络输出进行平滑约束测试它在 SMC 采样中的表现。结果确实比不约束的版本更稳定ESS 曲线更平直。这个方向还比较初步但我觉得它指向了一种训练与采样联合设计的可能不再把采样器当成模型训练完之后的一个独立后处理模块而是把采样算法的特性反馈到训练目标里。回到最初的问题——扩散模型的采样到底准不准标准采样器在生成看起来合理的样本这件事上已经足够好但在生成忠实于目标分布的样本这件事上还有明显短板。Reverse Diffusion SMC 用粒子滤波的思路给这个短板打了一块补丁代价是更高的计算开销。在我的实验配置里512 粒子、20 步、2 步移动大约需要 4 到 5 秒完成一次采样任务相比 DDIM 高了两个数量级但换来的是分布层面的可靠性。如果你只是为了跑图这个代价不值得但如果你在认真对待概率本身这套方法会是你工具箱里一个非常趁手的存在。