使用可扩展插值 Transformer（SiT）探索流模型与扩散模型
_{官方 PyTorch 实现}

论文 | 项目页面 |

此仓库包含我们用于探索基于可扩展 Transformer 的插值模型（SiT）的论文的 PyTorch 模型定义、预训练权重以及训练和采样代码。

使用可扩展插值 Transformer 探索流模型与扩散模型
Nanye Ma, Mark Goldstein, Michael Albergo, Nicholas Boffi, Eric Vanden-Eijnden, Saining Xie
纽约大学

我们提出了可扩展插值 Transformer（SiT），这是一类基于扩散 Transformer（DiT）骨干构建的生成模型。插值框架允许以比标准扩散模型更灵活的方式连接两个分布，从而可以模块化地研究影响基于动力传输的生成模型的各种设计选择：采用离散或连续时间学习、决定要学习的模型、选择连接两个分布的插值路径，以及部署确定性或随机采样器。通过精心引入上述要素，SiT 在条件 ImageNet 256×256 基准上，使用完全相同的骨干网络、参数量和 GFLOPs，在不同模型规模下均优于 DiT。通过探索可独立于学习进行调优的扩散系数，SiT 达到了 FID-50K 得分为 2.06 的优异表现。

本仓库包含：

• 🪐 SiT 的简单 PyTorch 实现
• ⚡️ 在 ImageNet 256×256 上训练的类条件 SiT 预训练模型
• 🛸 使用 PyTorch DDP 的 SiT 训练脚本

设置

首先，克隆并设置仓库：

git clone https://github.com/willisma/SiT.git
cd SiT

我们提供了一个 environment.yml 文件，可用于创建 Conda 环境。如果您只想在本地 CPU 上运行预训练模型，可以从该文件中移除 cudatoolkit 和 pytorch-cuda 依赖项。

conda env create -f environment.yml
conda activate SiT

采样

预训练的 SiT 检查点。 您可以使用 sample.py 从我们的预训练 SiT 模型中采样。根据您使用的模型，我们的预训练 SiT 模型权重将自动下载。该脚本提供了多种参数来调整采样器配置（ODE & SDE）、采样步数、更改无分类器指导尺度等。例如，要使用默认 ODE 设置从我们的 256×256 SiT-XL 模型中采样，可以运行：

python sample.py ODE --image-size 256 --seed 1

为方便起见，我们的预训练 SiT 模型也可在此直接下载：

SiT 模型	图像分辨率	FID-50K	Inception 分数	Gflops
XL/2	256×256	2.06	270.27	119

自定义 SiT 检查点。 如果您已使用 train.py 训练了新的 SiT 模型（见下方 训练 SiT），则可以添加 --ckpt 参数来使用您自己的检查点。例如，要使用 EMA 权重从一个自定义的 256×256 SiT-L/4 模型中进行 ODE 采样，可以运行：

python sample.py ODE --model SiT-L/4 --image-size 256 --ckpt /path/to/model.pt

高级采样器设置

ODE	--atol	float	绝对误差容忍度
	--rtol	float	相对误差容忍度
	--sampling-method	str	采样方法（参考 torchdiffeq）

SDE	--diffusion-form	str	SDE 扩散系数的形式（参考论文中的表 2）
	--diffusion-norm	float	SDE 扩散系数的大小
	--last-step	str	SDE 最后一步的形式
			None - 单步 SDE 积分
			"Mean" - 不含扩散系数的 SDE 积分步骤
			"Tweedie" - Tweedie 去噪步骤
			"Euler" - 单步 ODE 积分
	--sampling-method	str	采样方法
			"Euler" - 一阶积分
			"Heun" - 二阶积分

还有更多选项；详情请参阅 train_utils.py。

训练 SiT

我们在 train.py 中提供了 SiT 的训练脚本。要在单节点上使用 N 张 GPU 来启动 SiT-XL/2 (256x256) 的训练：

torchrun --nnodes=1 --nproc_per_node=N train.py --model SiT-XL/2 --data-path /path/to/imagenet/train

日志记录。 要启用 wandb，首先需要将 WANDB_KEY、ENTITY 和 PROJECT 设置为环境变量：

export WANDB_KEY="key"
export ENTITY="entity name"
export PROJECT="project name"

然后在训练命令中添加 --wandb 标志：

torchrun --nnodes=1 --nproc_per_node=N train.py --model SiT-XL/2 --data-path /path/to/imagenet/train --wandb

插值器设置。 我们还支持不同的插值器和模型预测方式选择。例如，要以 Linear 插值器和 noise 预测方式启动 SiT-XL/2 (256x256) 的训练：

torchrun --nnodes=1 --nproc_per_node=N train.py --model SiT-XL/2 --data-path /path/to/imagenet/train --path-type Linear --prediction noise

恢复训练。 要从自定义检查点恢复训练：

torchrun --nnodes=1 --nproc_per_node=N train.py --model SiT-L/2 --data-path /path/to/imagenet/train --ckpt /path/to/model.pt

注意。 恢复训练会自动恢复模型、EMA 和优化器的状态以及训练配置，使其与检查点中的状态完全一致。

评估（FID、Inception Score 等）

我们包含了一个 sample_ddp.py 脚本，该脚本可以并行地从 SiT 模型中采样大量图像。此脚本会生成一个包含样本的文件夹以及一个 .npz 文件，可以直接与 ADM 的 TensorFlow 评估套件 一起使用，以计算 FID、Inception Score 等指标。例如，在默认的 ODE 采样器设置下，要在 N 张 GPU 上从我们预训练的 SiT-XL/2 模型中采样 50,000 张图像，可以运行以下命令：

torchrun --nnodes=1 --nproc_per_node=N sample_ddp.py ODE --model SiT-XL/2 --num-fid-samples 50000

似然性。 支持似然性评估。要计算似然性，可以在 ODE 采样器中添加 --likelihood 标志：

torchrun --nnodes=1 --nproc_per_node=N sample_ddp.py ODE --model SiT-XL/2 --likelihood

请注意，只有在 ODE 采样器下才能计算似然性；更多细节和设置请参阅 sample_ddp.py。

改进方向

通过以下方式，训练（和采样）的速度有望显著提升：

• 在 SiT 模型中使用 Flash Attention
• 在 PyTorch 2.0 中使用 torch.compile

还有一些不错的基础功能待添加：

• 监控 FID 和其他指标
• AMP/bfloat16 支持

似然性计算的精度可以通过以下方式进一步提高：

• 均匀分布 / 高斯去量化

与 JAX 的差异

我们的模型最初是在 JAX 中使用 TPU 训练的。本仓库中的权重直接从 JAX 模型移植而来。由于在不同平台（TPU vs. GPU）上进行采样，结果可能会存在细微差异。我们观察到，在 TPU 上采样的效果略逊于 GPU（论文中分别为 2.15 和 2.06 的 FID 分数）。

许可证

本项目采用 MIT 许可证。详情请参阅 LICENSE。

SiT 快速上手指南

SiT (Scalable Interpolant Transformers) 是一种基于扩散 Transformer (DiT) 架构的生成模型家族。它通过灵活的插值框架连接分布，支持离散/连续时间学习及确定性/随机采样，在 ImageNet 256x256 基准测试中表现优异。

环境准备

• 系统要求: Linux 或 macOS (Windows 需自行配置 PyTorch CUDA 环境)。
• 硬件建议: 推荐使用 NVIDIA GPU 进行训练和高效采样；仅运行预训练模型可使用 CPU。
• 前置依赖:
  • Python 3.8+
  • Conda (推荐用于环境管理)
  • PyTorch (含 CUDA 支持，如需 GPU 加速)
  • torchdiffeq (用于 ODE/SDE 采样)

注意：国内用户若下载 Conda 包或 PyPI 包较慢，可配置清华源或阿里源加速。

安装步骤

1. 克隆仓库

   git clone https://github.com/willisma/SiT.git
   cd SiT
   

2. 创建 Conda 环境 项目提供了 environment.yml 文件。

   • GPU 版本（默认）：

     conda env create -f environment.yml
     

   • 仅 CPU 版本：请先编辑 environment.yml，移除 cudatoolkit 和 pytorch-cuda 相关行，然后执行上述命令。

3. 激活环境

   conda activate SiT
   

基本使用

1. 快速采样 (使用预训练模型)

无需手动下载权重，脚本会自动根据模型配置下载预训练参数。以下命令使用默认的 ODE 设置从 SiT-XL/2 模型生成 256x256 图像：

python sample.py ODE --image-size 256 --seed 1

• 生成的图片将保存在当前目录。
• 可通过添加 --cfg-scale 调整无分类器引导比例，或修改 --steps 调整采样步数。

2. 自定义模型采样

如果您训练了自己的模型或使用特定检查点，可通过 --ckpt 指定路径：

python sample.py ODE --model SiT-L/4 --image-size 256 --ckpt /path/to/your_model.pt

3. 开始训练 (可选)

在单节点多卡环境下启动训练（以 ImageNet 数据集为例）：

# N 替换为实际 GPU 数量，/path/to/imagenet/train 替换为数据路径
torchrun --nnodes=1 --nproc_per_node=N train.py --model SiT-XL/2 --data-path /path/to/imagenet/train

• 启用 WandB 日志：先设置环境变量 export WANDB_KEY="..." 等，然后在命令后添加 --wandb 标志。
• 断点续训：添加 --ckpt /path/to/checkpoint.pt 即可自动恢复模型、EMA 及优化器状态。