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

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

---

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

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

该代码的功能是什么？

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

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

PyTorch版本

请在此处查看PyTorch实现：[https://github.com/yang-song/score_sde_pytorch]，该实现支持除使用预训练分类器进行类别条件生成之外的所有功能。

JAX与PyTorch对比

一般来说，PyTorch版本的内存消耗较少，但运行速度比JAX慢。以下是在4块Nvidia Tesla V100 GPU（32GB）硬件上训练NCSN++连续模型时的基准测试结果：

框架	每步耗时（秒）	总内存占用（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) 列）。前者对应于训练过程中每 50k 次迭代时达到的最小 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 快速上手指南

Score-SDE 是一个基于随机微分方程（SDE）的分数生成模型官方实现，支持图像生成、修复、着色等任务。本指南帮助开发者快速搭建环境并运行基础示例。

环境准备

系统要求

• 操作系统: Linux (推荐) 或 macOS
• GPU: 推荐使用 NVIDIA GPU (需安装 CUDA)，JAX 版本对显存要求较高，PyTorch 版本相对节省显存。
• Python: 3.7 或更高版本

前置依赖

本项目主要依赖 JAX/Flax (官方首选) 或 PyTorch。

• JAX 版: 性能更优，但显存占用较大。
• PyTorch 版: 显存占用较小，运行速度稍慢，不支持部分基于预训练分类器的类别条件生成功能。

注意：国内用户建议在安装 Python 包时使用清华或阿里镜像源加速下载。

安装步骤

1. 克隆代码库

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

2. 安装依赖包

使用 pip 安装所需依赖（推荐配置国内镜像源）：

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

注：若需使用 PyTorch 版本，请前往 score_sde_pytorch 仓库按照其说明安装。

3. 下载评估统计文件 (可选)

若需进行定量评估（如计算 FID、Inception Score），需下载 CIFAR-10 的统计文件：

# 创建目录
mkdir -p assets/stats

# 下载 cifar10_stats.npz (需手动从 Google Drive 下载或使用替代链接)
# 链接: https://drive.google.com/file/d/1fXgBupLzThTGLLsiYCHRQJixuDsR1bSI/view?usp=sharing
# 下载后将其移动到 assets/stats/ 目录下

4. 获取预训练模型 (可选)

如需直接推理而非重新训练，可从 Google Drive 下载预训练 checkpoints 并放入对应工作目录。

基本使用

所有训练和评估操作均通过 main.py 脚本执行。

1. 训练模型

选择一个配置文件开始训练。配置文件位于 configs/ 目录，命名规则为 {数据集}_{模型}_{是否连续时间}.py。

以下示例使用 CIFAR-10 数据集和 NCSN++ 连续时间模型进行训练：

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

• --config: 指定配置文件路径。
• --workdir: 指定实验输出目录（保存检查点、样本等）。
• --mode=train: 启动训练模式。若目录下存在断点，将自动恢复训练。

2. 评估与采样

训练完成后（或使用预训练模型），可切换至评估模式生成样本并计算指标。

以下示例生成样本并计算 FID 和 Inception Score：

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

若需计算对数似然值 (Log-likelihood / bits/dim)：

python main.py --config=configs/cifar10_ncsnpp_continuous.py --workdir=./experiments/cifar10_test --mode=eval --eval_folder=eval_results --config.eval.enable_bpd=True --config.eval.dataset=test

3. 关键参数说明

通过命令行覆盖配置文件参数（基于 ml_collections）：

• --config.eval.enable_sampling: 启用图像采样。
• --config.eval.enable_bpd: 启用 bits/dim 计算。
• --config.eval.dataset: 指定评估数据集 (train 或 test)。
• --eval_folder: 评估结果保存的子文件夹名称。

快速体验 (Colab)

如果本地环境配置困难，推荐使用官方提供的 Colab 教程直接运行：

• JAX 版本: Open in Colab (JAX)
• PyTorch 版本: Open in Colab (PyTorch)