分数熵离散扩散

此仓库包含论文 Discrete Diffusion Modeling by Estimating the Ratios of the Data Distribution 的 PyTorch 实现，作者为 Aaron Lou, Chenlin Meng 和 Stefano Ermon。

设计选择

本代码库采用模块化构建以促进未来研究（相较于更适合应用的紧凑框架）。主要文件包括

1. noise_lib.py: 噪声调度 (noise schedule)
2. graph_lib: 前向扩散过程 (forward diffusion process)
3. sampling.py: 采样策略 (sampling strategies)
4. model/: 模型架构 (model architecture)

安装

只需运行

conda env create -f environment.yml

这将创建一个名为 sedd 的环境并安装相关包。注意，此处安装使用 CUDA 11.8，不同 CUDA 版本需手动安装。最关键的是确保 torch 和 flash-attn 包使用相同的 CUDA 版本（更多信息见 here）。

使用预训练模型

下载模型

我们的预训练模型 (pretrained models) 托管在 huggingface (small, medium)。然而，模型也可以本地加载（例如训练后）。所有功能均在 load_model.py 中。

# load in a pretrained model
pretrained_small_model, graph, noise = load_model("louaaron/sedd-small")
pretrained_medium_model, graph, noise = load_model("louaaron/sedd-medium")
# load in a local experiment
local_model, graph, noise = load_model("exp_local/experiment)

此加载方式提供模型，以及图结构和噪声（用于损失/采样设置 (loss/sampling setup)）。

运行采样

我们可以使用以下命令运行采样 (sampling)

python run_sample.py --model_path MODEL_PATH --steps STEPS

我们也可以使用条件采样 (conditional sampling)

python run_sample_cond.py --model_path MODEL_PATH --step STEPS --prefix PREFIX --suffix SUFFIX

训练新模型

运行训练

我们提供训练代码，可通过以下命令运行

python run_train.py

这将创建一个新目录 direc=exp_local/DATE/TIME，结构如下（兼容在本地运行采样实验）

├── direc
│   ├── .hydra
│   │   ├── config.yaml
│   │   ├── ...
│   ├── checkpoints
│   │   ├── checkpoint_*.pth
│   ├── checkpoints-meta
│   │   ├── checkpoint.pth
│   ├── samples
│   │   ├── iter_*
│   │   │   ├── sample_*.txt
│   ├── logs

此处，checkpoints-meta 用于在中断后重新加载运行，samples 包含运行过程中生成的图像，logs 包含运行输出。参数可通过 ARG_NAME=ARG_VALUE 添加，重要参数包括：

ngpus                     训练中使用的 GPU 数量（使用 PyTorch DDP (DistributedDataParallel)）
training.accum            累积步数，小模型设为 1，中模型设为 2（假设节点为 8x80GB）
noise.type                几何 (geometric) 或 loglinear 之一 
graph.type                均匀 (uniform) 或 absorb 之一
model                     small 或 medium 之一
model.scale_by_sigma      如果 graph.type=uniform 则设为 False（尚未配置）

一些示例命令包括

# training hyperparameters for SEDD absorb
python train.py noise_lib=loglinear graph.type=absorb model=medium training.accum=2
# training hyperparameters for SEDD uniform
python train.py noise_lib=geometric graph.type=uniform model=small model.scale_by_sigma=False

其他功能

SLURM 兼容性

要在 SLURM 上训练，只需运行

python train.py -m args

引用

@article{lou2024discrete,
  title={Discrete diffusion modeling by estimating the ratios of the data distribution},
  author={Lou, Aaron and Meng, Chenlin and Ermon, Stefano},
  journal={arXiv preprint arXiv:2310.16834},
  year={2024}
}

致谢

本仓库大量基于 score sde, plaid, 和 DiT。

Score-Entropy-Discrete-Diffusion 快速上手指南

本仓库包含论文 Discrete Diffusion Modeling by Estimating the Ratios of the Data Distribution 的 PyTorch 实现。该工具支持离散扩散模型的训练与采样。

环境准备

• 系统要求: Linux / macOS / Windows (推荐 Linux + GPU)
• 硬件: NVIDIA GPU (建议使用 CUDA 11.8 环境)
• 依赖: Python 3.x, Conda, PyTorch, Flash Attention
• 注意: 确保 torch 和 flash-attn 包使用的 CUDA 版本一致（默认安装为 CUDA 11.8）。

安装步骤

1. 创建并激活 Conda 环境：

conda env create -f environment.yml
conda activate sedd

2. 下载预训练模型 (可选): 模型托管于 HuggingFace。若网络受限，建议配置国内镜像加速下载。

# 示例：加载预训练模型
from load_model import load_model
pretrained_small_model, graph, noise = load_model("louaaron/sedd-small")
pretrained_medium_model, graph, noise = load_model("louaaron/sedd-medium")

基本使用

运行采样 (Sampling)

使用命令行进行无条件采样：

python run_sample.py --model_path MODEL_PATH --steps STEPS

使用命令行进行条件采样 (Conditional Sampling)：

python run_sample_cond.py --model_path MODEL_PATH --step STEPS --prefix PREFIX --suffix SUFFIX

训练新模型 (Training)

启动训练脚本：

python run_train.py

训练将生成目录结构如下，包含检查点、样本和日志：

├── direc
│   ├── .hydra
│   │   ├── config.yaml
│   │   ├── ...
│   ├── checkpoints
│   │   ├── checkpoint_*.pth
│   ├── samples
│   │   ├── iter_*
│   │   │   ├── sample_*.txt
│   ├── logs

常用参数调整示例：

# SEDD absorb 类型训练
python train.py noise_lib=loglinear graph.type=absorb model=medium training.accum=2

# SEDD uniform 类型训练
python train.py noise_lib=geometric graph.type=uniform model=small model.scale_by_sigma=False

SLURM 集群支持

如需在 SLURM 上训练，添加 -m 参数：

python train.py -m args