离散分数熵扩散

本仓库包含了论文《通过估计数据分布比率的离散扩散建模》的PyTorch实现，作者为Aaron Lou、Chenlin Meng和Stefano Ermon。

设计选择

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

noise_lib.py：噪声调度
graph_lib：前向扩散过程
sampling.py：采样策略
model/：模型架构

安装

只需运行

conda env create -f environment.yml

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

使用预训练模型

下载模型

我们的预训练模型托管在huggingface上（small，medium）。但是，模型也可以在本地加载（例如训练后）。所有功能都在load_model.py中。

# 加载预训练模型
pretrained_small_model, graph, noise = load_model("louaaron/sedd-small")
pretrained_medium_model, graph, noise = load_model("louaaron/sedd-medium")
# 加载本地实验
local_model, graph, noise = load_model("exp_local/experiment")

这样加载会得到模型以及图和噪声（用于损失/采样设置）。

运行采样

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

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

我们也可以使用以下命令进行条件采样

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）
training.accum            累积步数，small设为1，medium设为2（假设使用8x80GB节点）
noise.type                可选geometric或loglinear 
graph.type                可选uniform或absorb
model                     可选small或medium
model.scale_by_sigma      如果graph.type=uniform则设为False（尚未配置）

一些示例命令包括

# 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上训练，只需运行

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}
}