Hugging Face Diffusers 的 DPMSolverSDEScheduler：基于随机微分方程的随机采样器完全指南

导读

DPMSolverSDEScheduler 是 Hugging Face Diffusers 中一款特殊的采样调度器：它不是像 DPMSolver 系列那样的确定性求解器，而是把扩散模型的去噪过程建模为一条随机微分方程（SDE）的反向求解，在每一步显式地注入随机噪声，从而同时保留"采样多样性"与"快速收敛"两种特性。本篇指南以 docs/source/en/api/schedulers/dpm_sde.md 为核心骨架，结合仓库内 scheduling_dpmsolver_sde.py 的完整实现与其测试用例，系统讲解该调度器的设计来源、全部构造参数、sigma 噪声调度体系、二阶求解流程，以及如何把它接入自定义去噪循环。读完本文，你将能够独立配置并运行基于该调度器的随机采样，也能理解其与确定性调度器在原理层面的本质区别。

1. 设计来源与定位

DPMSolverSDEScheduler 的实现灵感来自论文 Elucidating the Design Space of Diffusion-Based Generative Models（即常说的 Karras 等人 EDM 工作），采用其中提出的**随机采样器（stochastic sampler）**思想。在源码层面，该调度器由 Katherine Crowson（katherinecrowson 社区生态 的作者，也是 Stable Diffusion 社区广泛使用的 k-diffusion 采样器的作者）移植并创建，仓库内保留其版权声明（见 scheduling_dpmsolver_sde.py）。

从类层级看，DPMSolverSDEScheduler 同时继承自 SchedulerMixin 与 ConfigMixin（见 scheduling_dpmsolver_sde.py），因此它具备所有 Diffusers 调度器共有的能力：

• 配置保存/加载：通过 ConfigMixin，构造参数会被 @register_to_config 记录，可整体 save_pretrained / from_pretrained；
• 调度器互换：_compatibles 声明为 KarrasDiffusionSchedulers 的全部成员（见 scheduling_dpmsolver_sde.py），意味着在 DiffusionPipeline 中它可以与其他同族调度器（如 Euler、DPM、UniPC 等，枚举定义见 scheduling_utils.py）无缝替换。

一句话概括它的定位：在保留 SDE 随机采样质量的同时，通过基于中点（midpoint）的二阶预测-校正机制逼近确定性 ODE 求解器的效率。order = 2 的类属性（见 scheduling_dpmsolver_sde.py）正是这一二阶性质的直接体现。

2. 构造参数详解

调度器所有参数都通过 __init__ 中的 @register_to_config 注册（见 scheduling_dpmsolver_sde.py）。下表汇总了全部参数、默认值及其作用：

参数	默认值	可选值	作用
num_train_timesteps	1000	int	模型训练的扩散步数，决定 beta 序列长度
beta_start	0.00085	float	推理用 beta 序列的起始值
beta_end	0.012	float	推理用 beta 序列的终止值
beta_schedule	"linear"	"linear"、"scaled_linear"、"squaredcos_cap_v2"	从 beta 区间映射到 beta 序列的方式
trained_betas	None	np.ndarray / list[float]	直接传入训练好的 beta 数组，绕过 beta_start / beta_end
prediction_type	"epsilon"	"epsilon"、"sample"、"v_prediction"	模型输出的预测类型
use_karras_sigmas	False	bool	是否使用 Karras（EDM）sigma 噪声调度
use_exponential_sigmas	False	bool	是否使用指数 sigma 噪声调度
use_beta_sigmas	False	bool	是否使用 Beta 分布 sigma 噪声调度（需安装 scipy）
noise_sampler_seed	None	int	布朗噪声采样器的随机种子
timestep_spacing	"linspace"	"linspace"、"leading"、"trailing"	timestep 的缩放方式
steps_offset	0	int	推理步的偏移量，部分模型族需要

2.1 三种 beta schedule 的实现差异

从 scheduling_dpmsolver_sde.py 可以看到：

• "linear"：直接对 [beta_start, beta_end] 做 torch.linspace；
• "scaled_linear"：先对 beta_start**0.5 到 beta_end**0.5 做线性插值再平方，注释明确指出"该调度是潜在扩散模型（LDM）特有的"；
• "squaredcos_cap_v2"：调用 betas_for_alpha_bar（GLIDE 风格余弦调度，见 scheduling_dpmsolver_sde.py），该辅助函数支持 cosine、laplace、exp 三种 alpha_bar 变换并夹紧到 max_beta=0.999。

随后调度器计算 alphas = 1 - betas 与 alphas_cumprod，并据此得到 sigma 基础曲线 ((1 - alpha_cumprod) / alpha_cumprod) ** 0.5（见 scheduling_dpmsolver_sde.py 与 set_timesteps 中 L451）。

2.2 构造期的校验逻辑

构造函数内置了两条防御性校验（见 scheduling_dpmsolver_sde.py）：

1. 启用 use_beta_sigmas 但未安装 scipy 时直接抛 ImportError；
2. use_beta_sigmas、use_exponential_sigmas、use_karras_sigmas 三者最多只能启用一个，否则抛 ValueError。

此外，构造完成后 self.sigmas 会被主动移到 CPU（L301），以避免训练/推理过程中 CPU 与 GPU 之间频繁通信，这一细节与 set_timesteps 末尾的行为保持一致。

3. 噪声调度体系：Karras / Exponential / Beta sigmas

set_timesteps(num_inference_steps, device=None, num_train_timesteps=None) 是推理前的关键入口（见 scheduling_dpmsolver_sde.py）。其工作流程分三步：

第一步：生成离散 timesteps。 依据 timestep_spacing 三种策略生成（L432-L449），对应论文 Common Diffusion Noise Schedules and Sample Steps are Flawed 中 Table 2 的记号。

第二步：插值得到基础 sigma。 将 timesteps 经 np.interp 映射到基础 sigma 曲线上。

第三步（可选）：重排 sigma。 依据三个互斥开关之一转换 sigma，并把 sigma 反算回 timesteps：

• Karras sigmas（_convert_to_karras，见 L543-L566）：采用论文提出的公式，rho = 7.0 为论文原值，把 [sigma_min, sigma_max] 按 ramp 在倒数空间线性插值后再取 rho 次幂；
• Exponential sigmas（_convert_to_exponential，见 L568-L600）：在 log(sigma_max) 与 log(sigma_min) 之间均匀采样再取指数；
• Beta sigmas（_convert_to_beta，见 L602-L649）：基于论文 Beta Sampling is All You Need，默认取 Beta 分布 alpha=0.6, beta=0.6，通过 scipy.stats.beta.ppf 分位数函数生成 sigma。

3.1 关键实现细节：二阶步长的"膨胀"

set_timesteps 末尾有一段容易被忽略但至关重要的逻辑（L465-L474）：

second_order_timesteps = self._second_order_timesteps(sigmas, log_sigmas)
sigmas = np.concatenate([sigmas, [0.0]]).astype(np.float32)
self.sigmas = torch.cat([sigmas[:1], sigmas[1:-1].repeat_interleave(2), sigmas[-1:]])
timesteps = torch.cat([timesteps[:1], timesteps[1:].repeat_interleave(2)])
timesteps[1::2] = second_order_timesteps

含义是：除首尾外，每个 sigma 与 timestep 都被重复一次，且重复出来的奇数位置 timestep 被替换为"中点时间"。也就是说，用户请求 num_inference_steps 步时，实际调度表长度接近 2 * num_inference_steps：每个逻辑步由一次一阶（预测）与一次二阶（校正）子步交替构成。_second_order_timesteps（L491-L504）用 t = -log(sigma) 变换后在时间域取 midpoint_ratio = 0.5 的中点，再映射回 timestep。

3.2 init_noise_sigma 与首步噪声

init_noise_sigma 属性（L350-L356）决定初始潜变量的噪声缩放：timestep_spacing 为 "linspace" 或 "trailing" 时直接返回 sigmas.max()；"leading" 时返回 (sigmas.max() ** 2 + 1) ** 0.5（因为在 "leading" 下首步 sigma 对应的实际是 x0 的缩放域）。测试 test_scheduler_dpm_sde.py 中即用 sample = dummy_sample * scheduler.init_noise_sigma 初始化完整循环。

4. 随机噪声采样器：BrownianTreeNoiseSampler 与 BatchedBrownianTree

随机采样区别于确定性求解的核心在于每次迭代注入新的随机噪声，这部分由噪声采样器负责（见 scheduling_dpmsolver_sde.py）：

• BatchedBrownianTree 包装 torchsde.BrownianInterval，支持**单条或多条（按 batch 逐条分配种子）**布朗桥。构造时对时间区间排序并记录符号；支持传入 seed 列表，此时每条样本各用一棵独立的布朗树，w0 取第一维（L54-L108）；
• BrownianTreeNoiseSampler 是对前者的一层封装：接受 x（决定形状/设备/dtype）、sigma_min、sigma_max 与可选的 transform（把 sigma 映射到内部时间）；__call__(sigma, sigma_next) 返回两时间点之间的布朗增量，并除以 (t1 - t0).abs().sqrt() 归一化（L111-L139）。

在 step 内部，噪声采样器是惰性创建的（L689-L693）：首次调用时以当前 sample 的形状、sigmas 中大于 0 的最小值与最大值、以及 noise_sampler_seed 构建。因此：

• noise_sampler_seed 传入固定整数即可获得可复现的随机采样；
• 不传则每次运行随机，天然保留多样性。

5. step()：二阶随机采样单步实现剖析

step(model_output, timestep, sample, return_dict=True, s_noise=1.0)（见 scheduling_dpmsolver_sde.py）是每个去噪循环的核心。其内部通过 state_in_first_order（即 self.sample is None，L651-L653）区分一阶/二阶子步，以 midpoint_ratio = 0.5 计算时间域中点（L710-L714）。

第 1 步：由模型输出还原 x0 预测。 依据 prediction_type 分支处理（L716-L730）：

• "epsilon"：pred_original_sample = sample - sigma_input * model_output；
• "v_prediction"：按 v 参数化公式结合 sigma_input 还原 x0；
• "sample"：源码中直接抛出 NotImplementedError（见 L725-L726），这一点与构造参数文档中宣称支持三种类型存在出入，实际可用的是 epsilon 与 v_prediction；测试也仅覆盖这两种（见 test_scheduler_dpm_sde.py）。

第 2 步：末步（sigma_next == 0）走确定性收尾。 直接计算导数 (sample - pred_original_sample) / sigma 并用 prev_sample = sample + derivative * dt 完成最后一步（L732-L735），保证最终输出没有额外噪声。

第 3 步：非末步的随机祖先采样。 按 EDM 的拆分方式计算：

• sigma_up：本步要注入的随机噪声量，取 sigma_to 与公式 (sigma_to**2 * (sigma_from**2 - sigma_to**2) / sigma_from**2) ** 0.5 的较小值；
• sigma_down：确定性部分对应的 sigma；
• 更新公式为 prev_sample = (sigma_fn(ancestral_t)/sigma_fn(t)) * sample - (t - ancestral_t).expm1() * pred_original_sample + noise_sampler(...) * s_noise * sigma_up（L743-L753），其中 s_noise 是噪声放大系数，默认 1.0。

第 4 步：一阶/二阶状态切换。 一阶子步结束后保存 sample 与 mid_point_sigma 供下一个二阶子步复用（L755-L758）；二阶子步结束时清除状态回到一阶模式。每次 step 后 _step_index 自增（L765）。

返回值。 return_dict=True 时返回 DPMSolverSDESchedulerOutput(prev_sample, pred_original_sample)；False 时返回 (prev_sample, pred_original_sample) 元组。pred_original_sample 可用于进度预览或 guidance（定义见 scheduling_dpmsolver_sde.py）。

6. 其他重要接口

6.1 scale_model_input 与输入归一化

scale_model_input(sample, timestep)（L383-L408）确保与需要按当前 timestep 缩放输入的调度器保持互换性：取当前步 sigma（二阶子步取 mid_point_sigma），执行 sample / ((sigma_input**2 + 1) ** 0.5)。测试的完整循环中每次都先调用它再送入模型（见 test_scheduler_dpm_sde.py）。

6.2 add_noise 与 img2img / inpainting

add_noise(original_samples, noise, timesteps)（L775-L822）按 timestep 从 sigma 表取对应噪声水平加噪，支持三种场景的 begin_index / step_index 语义（训练、img2img 初始加噪、inpainting 中途加噪），与 set_begin_index（L372-L381）配合使用。

6.3 与调度器通用的 SchedulerOutput

文档末尾引用的 SchedulerOutput 定义于 scheduling_utils.py，是所有调度器 step 输出的基类，只含 prev_sample 字段；DPMSolverSDESchedulerOutput 在其基础上扩展了可选的 pred_original_sample。

7. 接入自定义去噪循环的完整示例

以下代码展示了在 Diffusers 生态中手动驱动该调度器的标准流程（结构参照 test_scheduler_dpm_sde.py 的 test_full_loop_no_noise 与 test_full_loop_device_karras_sigmas 编写，可直接运行于 CPU/GPU）：

import torch
from diffusers import DPMSolverSDEScheduler

# 1. 构造调度器（使用与测试一致的配置，含固定噪声种子保证可复现）
scheduler = DPMSolverSDEScheduler(
    num_train_timesteps=1000,
    beta_start=0.00085,
    beta_end=0.012,
    beta_schedule="linear",
    prediction_type="epsilon",      # 或 "v_prediction"
    use_karras_sigmas=False,        # 三选一：karras / exponential / beta
    use_exponential_sigmas=False,
    use_beta_sigmas=False,
    noise_sampler_seed=0,           # 固定种子 → 结果可复现
)

# 2. 设定推理步数并放到目标设备
num_inference_steps = 25
scheduler.set_timesteps(num_inference_steps, device="cuda")

# 3. 用 init_noise_sigma 生成初始潜变量（shape 依模型而定）
sample = torch.randn(1, 4, 64, 64, device="cuda") * scheduler.init_noise_sigma

# 4. 去噪循环
for t in scheduler.timesteps:
    # 4.1 按当前 timestep 归一化输入
    sample = scheduler.scale_model_input(sample, t)
    # 4.2 模型前向（此处替换为你的 UNet/DiT）
    model_output = your_model(sample, t)
    # 4.3 单步反演 SDE
    output = scheduler.step(model_output, t, sample)
    sample = output.prev_sample  # 也可读取 output.pred_original_sample 预览 x0

# 5. 得到的 sample 再经 VAE decoder 还原为像素

实用建议：

• 追求随机多样性时保持 noise_sampler_seed=None；需要对比实验可复现时固定种子；
• 追求更快的低步数收敛可尝试 use_karras_sigmas=True（测试 test_full_loop_device_karras_sigmas 给出了其在 CPU/CUDA/MPS 上的数值基准）；
• 若启用 use_beta_sigmas=True，请确保环境已安装 scipy，否则构造时会直接抛 ImportError（见 scheduling_dpmsolver_sde.py）。

8. 测试覆盖与数值稳定性参考

仓库为 DPMSolverSDEScheduler 提供了专门的测试套件 tests/schedulers/test_scheduler_dpm_sde.py，全部用例被 @require_torchsde 装饰（L9），意味着运行前必须先安装 torchsde。覆盖点包括：

• test_timesteps：验证 10/50/100/1000 四种步数；
• test_betas / test_schedules：验证 beta 区间与 linear、scaled_linear 两种 schedule；
• test_prediction_type：验证 epsilon 与 v_prediction；
• test_full_loop_no_noise / test_full_loop_with_v_prediction / test_full_loop_device：分别在 CPU、MPS、CUDA/XPU 上断言完整去噪循环的数值结果（如 CPU 下 result_sum ≈ 162.5238，见 L71-L72）；
• test_full_loop_device_karras_sigmas：Karras sigma 模式下的完整循环基准；
• test_beta_sigmas / test_exponential_sigmas：两种 sigma 调度的配置一致性校验。

这些测试既验证了实现正确性，也为你本地复现或二次开发提供了可直接套用的最小运行骨架。在 API 层面，DPMSolverSDEScheduler 已通过 src/diffusers/schedulers/init.py 的懒加载机制对外导出（from diffusers import DPMSolverSDEScheduler 即可使用），并在 docs/source/en/_toctree.yml 中登记为独立 API 文档页。

结语

DPMSolverSDEScheduler 在 Diffusers 调度器家族中的独特价值，在于把"随机噪声注入"与"二阶求解"合二为一：前者保证了采样结果的多样性和高保真度，后者让它在相同步数下逼近确定性求解器的效率。理解其 sigma 调度体系、布朗噪声采样器与一阶/二阶交替机制，是正确调参（如切换 Karras/Exponential/Beta sigmas、控制 s_noise 与 noise_sampler_seed）的前提。建议读者结合本文示例与 test_scheduler_dpm_sde.py 中的完整循环，在实际模型上对比不同 sigma 调度与种子策略的生成效果，从而找到最适合自身任务的质量-多样性平衡点。