基于分数的生成模型：通过随机微分方程

本仓库包含论文《基于分数的生成模型：通过随机微分方程》的 PyTorch 实现，该论文由 Yang Song、Jascha Sohl-Dickstein、Diederik P. Kingma、Abhishek Kumar、Stefano Ermon 和 Ben Poole 共同撰写。

---

我们提出了一种统一的框架，通过随机微分方程（SDE）的视角，对先前基于分数的生成模型工作进行了泛化和改进。具体而言，我们可以利用由 SDE 描述的连续时间随机过程，将数据转换为简单的噪声分布。如果已知每个中间时间步长上边缘分布的分数函数，就可以逆转这一 SDE 过程以进行采样；而这些分数函数则可通过分数匹配方法来估计。其基本思想如图所示：

我们的工作不仅有助于更好地理解现有方法，还带来了新的采样算法、精确的似然计算、唯一可识别的编码、潜在空间代码操控等功能，并为基于分数的生成模型家族增添了新的条件生成能力（包括但不限于类别条件生成、图像修复和彩色化等）。综合来看，在 CIFAR-10 数据集上的无条件生成任务中，我们取得了 FID 为 2.20、Inception 分数为 9.89 的优异成绩，并能高质量地生成 1024px 尺寸的 Celeba-HQ 图像（见下方样本）。此外，在对 CIFAR-10 图像进行均匀去量化处理后，我们还获得了 2.99 比特/维度的似然值。

该代码的功能是什么？

除了我们论文中的 NCSN++ 和 DDPM++ 模型外，本代码库还在一处重新实现了许多先前的基于分数的模型，包括来自论文《通过估计数据分布的梯度进行生成建模》的 NCSN、来自论文《训练基于分数的生成模型的改进技术》的 NCSNv2，以及来自论文《去噪扩散概率模型》的 DDPM。

该代码支持训练新模型，也支持评估现有模型的样本质量和似然值。我们精心设计了代码结构，使其具有模块化特性，便于扩展到新的随机微分方程、预测器或校正器。

与 🤗 Diffusers 库的集成

目前，大多数模型也已在 🧨 Diffusers 中提供，并可通过 ScoreSdeVE 流水线 访问。

借助 Diffusers，您只需几行代码即可在 PyTorch 中测试基于分数 SDE 的模型。

安装 Diffusers 的命令如下：

pip install diffusers torch accelerate

然后只需几行代码即可试用这些模型：

from diffusers import DiffusionPipeline

model_id = "google/ncsnpp-ffhq-1024"

# 加载模型和调度器
sde_ve = DiffusionPipeline.from_pretrained(model_id)

# 在推理模式下运行流水线（生成随机噪声并逐步去噪）
image = sde_ve().images[0]

# 保存图像
image[0].save("sde_ve_generated_image.png")

更多模型可以直接在 Hub 上找到。

JAX 版本

请参阅 此处 的 JAX 实现，该版本额外支持使用预训练分类器进行类别条件生成，以及在中断后恢复评估过程。

JAX 与 PyTorch 的对比

总体而言，PyTorch 版本的内存消耗较低，但运行速度比 JAX 慢。以下是在 4 张 Nvidia Tesla V100 GPU（每张 32GB）硬件上训练 NCSN++ VE SDE 模型的基准测试结果：

框架	每步耗时（秒）	总内存占用（GB）
PyTorch	0.56	20.6
JAX (n_jitted_steps=1)	0.30	29.7
JAX (n_jitted_steps=5)	0.20	74.8

如何运行代码

依赖项

运行以下命令以安装本代码所需的部分 Python 包：

pip install -r requirements.txt

定量评估用统计文件

我们提供了 CIFAR-10 数据集的统计文件。您可以下载 cifar10_stats.npz，并将其保存到 assets/stats/ 目录下。有关如何为新数据集计算此类统计文件，请参阅 #5。

使用方法

通过 main.py 训练和评估我们的模型。

main.py:
  --config: 训练配置。
    (默认值：'None')
  --eval_folder: 存储评估结果的文件夹名称
    (默认值：'eval')
  --mode: <train|eval>：运行模式：训练或评估
  --workdir: 工作目录

• config 是配置文件的路径。我们提供的预设配置文件位于 configs/ 目录下。这些配置文件按照 ml_collections 的格式编写，应该非常直观易懂。

  配置文件命名规范：配置文件的路径由以下维度组合而成：

  • 数据集：cifar10、celeba、celebahq、celebahq_256、ffhq_256、celebahq、ffhq 中的一个。
  • 模型：ncsn、ncsnv2、ncsnpp、ddpm、ddpmpp 中的一个。
  • 连续性：使用连续采样的时间步长进行训练。

• workdir 是存储一次实验所有产物的路径，例如检查点、样本和评估结果。

• eval_folder 是 workdir 中用于存储评估过程所有产物的子文件夹名称，例如用于防止中断的元检查点、图像样本以及定量结果的 NumPy 文件。

• mode 可以是 "train" 或 "eval"。当设置为 "train" 时，将开始训练新模型；如果 workdir/checkpoints-meta 中存在该模型的元检查点（用于在云环境中被抢占后恢复运行），则会从中断处继续训练。当设置为 "eval" 时，可以执行以下任意组合的操作：

  • 在测试/验证数据集上评估损失函数。
  • 生成固定数量的样本，并计算其 Inception 分数、FID 或 KID。在评估之前，必须先下载或计算好统计文件，并将其存储在 assets/stats 目录中。
  • 在训练或测试数据集上计算对数似然。

这些功能可以通过配置文件进行配置，也可以更方便地通过 ml_collections 包的命令行支持来实现。例如，要生成样本并评估样本质量，可以使用 --config.eval.enable_sampling 标志；要计算对数似然，则使用 --config.eval.enable_bpd 标志，并通过 --config.eval.dataset=train/test 指定是在训练集还是测试集上计算似然。

如何扩展代码

• 新的随机微分方程 (SDE)：继承 sde_lib.SDE 抽象类，并实现所有抽象方法。discretize() 方法是可选的，默认采用欧拉-丸山离散化。现有的采样方法和似然计算将自动适用于这个新的 SDE。
• 新的预测器：继承 sampling.Predictor 抽象类，实现 update_fn 抽象方法，并使用 @register_predictor 注册其名称。新预测器可以直接用于 sampling.get_pc_sampler 中的预测-校正采样，以及 controllable_generation.py 中的所有其他可控生成方法。
• 新的校正器：继承 sampling.Corrector 抽象类，实现 update_fn 抽象方法，并使用 @register_corrector 注册其名称。新校正器可以直接用于 sampling.get_pc_sampler，以及 controllable_generation.py 中的所有其他可控生成方法。

预训练检查点

所有检查点均在此 Google Drive 中提供。

说明：对于某些模型，您可能会找到两个检查点。第一个检查点（编号较小）是我们论文表3中报告FID分数的版本（也对应下表中的FID和IS列）。第二个检查点（编号较大）是我们论文表2中报告黑箱ODE采样器的似然值和FID的版本（也对应下表中的FID(ODE)和NNL (bits/dim)列）。前者对应于训练过程中每5万次迭代时取得的最小FID，而后者则是训练结束时的最后一个检查点。

根据Google的政策，我们无法发布原始的CelebA和CelebA-HQ检查点。不过，我已利用个人资源在FFHQ 1024px、FFHQ 256px以及CelebA-HQ 256px数据集上重新训练了模型，其性能与我们的内部检查点相当。

以下是检查点及其在论文中报告结果的详细列表。FID (ODE) 对应于应用于概率流ODE的黑箱ODE求解器的样本质量。

检查点路径	FID	IS	FID (ODE)	NNL (bits/dim)
ve/cifar10_ncsnpp/	2.45	9.73	-	-
ve/cifar10_ncsnpp_continuous/	2.38	9.83	-	-
ve/cifar10_ncsnpp_deep_continuous/	2.20	9.89	-	-
vp/cifar10_ddpm/	3.24	-	3.37	3.28
vp/cifar10_ddpm_continuous	-	-	3.69	3.21
vp/cifar10_ddpmpp	2.78	9.64	-	-
vp/cifar10_ddpmpp_continuous	2.55	9.58	3.93	3.16
vp/cifar10_ddpmpp_deep_continuous	2.41	9.68	3.08	3.13
subvp/cifar10_ddpm_continuous	-	-	3.56	3.05
subvp/cifar10_ddpmpp_continuous	2.61	9.56	3.16	3.02
subvp/cifar10_ddpmpp_deep_continuous	2.41	9.57	2.92	2.99

检查点路径	样本
ve/bedroom_ncsnpp_continuous
ve/church_ncsnpp_continuous
ve/ffhq_1024_ncsnpp_continuous
ve/ffhq_256_ncsnpp_continuous
ve/celebahq_256_ncsnpp_continuous

演示与教程

链接	描述
	加载我们的预训练检查点，体验采样、似然计算和可控生成（JAX + FLAX）
	加载我们的预训练检查点，体验采样、似然计算和可控生成（PyTorch）
	JAX + FLAX中的基于分数的生成模型教程
	PyTorch中的基于分数的生成模型教程

小贴士

• 使用JAX代码库时，可以将多个训练步骤一起进行jit编译，以提高训练速度，但会增加内存消耗。这可以通过config.training.n_jitted_steps来设置。对于CIFAR-10数据集，如果您的GPU/TPU内存充足，建议使用config.training.n_jitted_steps=5；否则建议使用config.training.n_jitted_steps=1。目前的实现要求config.training.log_freq能够被n_jitted_steps整除，以便日志记录和检查点保存正常工作。
• LangevinCorrector的snr（信噪比）参数在某种程度上类似于温度参数。较大的snr通常会产生更平滑的样本，而较小的snr则会产生更多样但质量较低的样本。典型的snr取值为0.05 - 0.2，需要通过调优找到最佳平衡点。
• 对于VE SDE，建议将config.model.sigma_max设置为训练数据集中任意两个样本之间的最大距离。

参考文献

如果您发现该代码对您的研究有帮助，请考虑引用以下文献：

@inproceedings{
  song2021scorebased,
  title={基于分数的生成模型：通过随机微分方程},
  author={杨松、雅沙·索尔-迪克斯坦、迪德里克·P·金格玛、阿比谢克·库马尔、斯特凡诺·埃尔蒙、本·普尔},
  booktitle={国际表示学习大会},
  year={2021},
  url={https://openreview.net/forum?id=PxTIG12RRHS}
}

本工作建立在一些先前的论文基础上，这些论文也可能对您有所兴趣：

• 杨松和斯特凡诺·埃尔蒙. “通过估计数据分布的梯度进行生成建模.” 第33届神经信息处理系统年会论文集. 2019年.
• 杨松和斯特凡诺·埃尔蒙. “训练基于分数的生成模型的改进方法.” 第34届神经信息处理系统年会论文集. 2020年.
• 霍, 乔纳森、阿贾伊·贾因和皮特·阿贝尔. “去噪扩散概率模型.” 第34届神经信息处理系统年会论文集. 2020年.

score_sde_pytorch 快速上手指南

score_sde_pytorch 是基于随机微分方程（SDE）的分数生成模型的 PyTorch 实现。该工具支持训练和评估多种生成模型（如 NCSN++、DDPM++），并提供了高质量的图像生成、似然计算及条件生成能力。

环境准备

在开始之前，请确保您的开发环境满足以下要求：

• 操作系统: Linux (推荐) 或 macOS
• Python: 3.7 或更高版本
• GPU: 推荐使用 NVIDIA GPU (需安装对应的 CUDA 驱动)，虽然支持 CPU 运行，但训练和采样速度会显著较慢。
• 前置依赖:
  • PyTorch
  • torchvision
  • ml_collections
  • 其他 Python 科学计算库 (numpy, scipy, PIL 等)

注意：本项目代码库未直接提供国内镜像源配置，建议在安装 Python 依赖时使用清华或阿里镜像加速。

安装步骤

1. 克隆代码库

首先从 GitHub 克隆项目源码：

git clone https://github.com/yang-song/score_sde_pytorch.git
cd score_sde_pytorch

2. 安装依赖

使用 pip 安装必要的 Python 包。推荐使用国内镜像源以加快下载速度：

pip install -r requirements.txt -i https://pypi.tuna.tsinghua.edu.cn/simple

如果 requirements.txt 中未包含 PyTorch，请根据您的需求单独安装（以下为 CUDA 11.8 示例）：

pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118

3. 下载统计文件 (可选，用于定量评估)

如果您需要计算 FID 或 Inception Score 等指标，需下载预计算的统计文件。以 CIFAR-10 为例：

mkdir -p assets/stats
# 手动下载 cifar10_stats.npz 并放入 assets/stats/ 目录
# 下载地址: https://drive.google.com/file/d/14UB27-Spi8VjZYKST3ZcT8YVhAluiFWI/view

基本使用

本工具主要通过 main.py 脚本进行模型训练和评估。使用前需准备好配置文件（位于 configs/ 目录下）。

1. 训练模型

选择一个配置文件启动训练。例如，使用默认的 CIFAR-10 NCSN++ 连续时间模型配置：

python main.py --config=configs/default_cifar10_ncsnpp_continuous.py --workdir=./experiments/cifar10_test --mode=train

• --config: 指定配置文件路径。
• --workdir: 指定工作目录，用于保存检查点（checkpoints）、样本和日志。
• --mode=train: 启动训练模式。如果工作目录中已有检查点，将自动恢复训练。

2. 评估与采样

训练完成后（或使用预训练模型），可切换至评估模式生成图像或计算指标：

python main.py --config=configs/default_cifar10_ncsnpp_continuous.py --workdir=./experiments/cifar10_test --mode=eval --eval_folder=eval_results

通过命令行参数控制具体评估任务（利用 ml_collections 特性）：

• 生成样本并计算 FID/IS:

  python main.py ... --mode=eval --config.eval.enable_sampling=True
  

• 计算对数似然 (Log-likelihood):

  python main.py ... --mode=eval --config.eval.enable_bpd=True --config.eval.dataset=test
  

3. 快速体验 (通过 🤗 Diffusers)

如果您仅想快速测试预训练模型而无需配置完整训练环境，可以使用 Hugging Face 的 diffusers 库，仅需几行代码即可生成图像：

安装 diffusers:

pip install diffusers torch accelerate -i https://pypi.tuna.tsinghua.edu.cn/simple

运行示例:

from diffusers import DiffusionPipeline

# 加载预训练模型 (例如 FFHQ 1024px)
model_id = "google/ncsnpp-ffhq-1024"
sde_ve = DiffusionPipeline.from_pretrained(model_id)

# 生成图像
image = sde_ve().images[0]

# 保存图像
image.save("sde_ve_generated_image.png")

更多预训练模型可在 Hugging Face Hub 查找。