DiffuSeq

DiffuSeq: 基于扩散模型的序列到序列文本生成 和 DiffuSeq-v2: 桥接离散与连续文本空间以加速 Seq2Seq 扩散模型 的官方代码库。

我们条件扩散语言模型 DiffuSeq 的扩散过程。

加速版 DiffuSeq 的扩散过程。

亮点

• 我们提出的 DiffuSeq 作为一种条件语言模型，以无分类器的方式进行端到端训练。
• 我们建立了 AR、NAR 和 DiffuSeq 模型之间的理论联系（参见我们的原始论文）。
• DiffuSeq 是一种强大的文本生成模型，在质量和多样性方面可媲美甚至超越具有竞争力的 AR、迭代式 NAR 以及大型 PLM 模型。

我们的研究通过这种新的序列到序列学习范式取得了令人鼓舞的成果。

更新：我们的增强版本有效地将训练收敛速度提升了 4 倍，并且能够在 800 倍更快的时间内生成质量相近的样本，使其更接近实际应用。

环境搭建：

该代码基于 PyTorch 和 HuggingFace transformers。

pip install -r requirements.txt 

数据集

准备数据集并将其放置在 datasets 文件夹下。以 datasets/CommonsenseConversation/train.jsonl 为例。我们在论文中使用了四个数据集。

任务	数据集	训练样本	来源	是否用于 DiffuSeq
开放域对话	Commonsense Conversation	338.2万	CCM	下载
问题生成	Quasar-T	11.7万	OpenQA	下载
文本简化	Wiki-alignment	67.7万	Wiki-auto	下载
同义句改写	QQP	14.4万	Kaggle	下载

DiffuSeq 训练

cd scripts
bash train.sh

参数说明：

• --dataset: 数据集名称，仅用于标注
• --data_dir: 已保存数据集文件夹的路径，包含 train.jsonl,test.jsonl,valid.jsonl
• --seq_len: 序列 $z$ ($x\oplus y$) 的最大长度
• --resume_checkpoint: 若不为空，则恢复该检查点并继续训练
• --vocab: 使用 bert 初始化分词器，或加载您自己预处理过的词汇表（例如使用 BPE）

在 4 张 NVIDIA A100 80G 显卡上训练 QG 和 QQP 的 DiffuSeq 模型大约需要 2 天时间，训练步数应根据训练集的大小相应增加。为了复现我们论文中表 1 的结果，我们建议在训练每个数据集时采用以下配置。

更新：

新增参数：

• --learned_mean_embed: 设置是否使用学习得到的软吸收态。
• --denoise: 设置是否添加离散噪声
• --use_fp16: 设置是否使用混合精度训练
• --denoise_rate: 设置去噪率，默认为 0.5

在 2 张 NVIDIA A100 80G 显卡上训练 QQP 模型仅需约 11 小时。

python -m torch.distributed.launch --nproc_per_node=4 --master_port=12233 --use_env run_train.py --diff_steps 2000 --lr 0.0001 --learning_steps 50000 --save_interval 10000 --seed 102 --noise_schedule sqrt --hidden_dim 128 --bsz 2048 --dataset qqp --data_dir {datasets/QQP} --vocab bert --seq_len 128 --schedule_sampler lossaware --notes qqp

python -m torch.distributed.launch --nproc_per_node=4 --master_port=12233 --use_env run_train.py --diff_steps 2000 --lr 0.0001 --learning_steps 40000 --save_interval 2000 --seed 102 --noise_schedule sqrt --hidden_dim 128 --bsz 2048 --microbatch 64 --dataset qg --data_dir {datasets/QG} --vocab bert --seq_len 128 --schedule_sampler lossaware --notes qg

python -m torch.distributed.launch --nproc_per_node=7 --master_port=12233 --use_env run_train.py --diff_steps 2000 --lr 0.0001 --learning_steps 140000 --save_interval 20000 --seed 102 --noise_schedule sqrt --hidden_dim 128 --bsz 2048 --microbatch 64 --dataset dialogue --data_dir {datasets/Conversation} --vocab bert --seq_len 128 --schedule_sampler lossaware --notes dialogue

python -m torch.distributed.launch --nproc_per_node=8 --master_port=12233 --use_env run_train.py --diff_steps 2000 --lr 0.0001 --learning_steps 80000 --save_interval 20000 --seed 102 --noise_schedule sqrt --hidden_dim 128 --bsz 2048 --microbatch 64 --dataset dialogue --data_dir {datasets/TS} --vocab bert --seq_len 128 --schedule_sampler lossaware --notes ts

经验表明，较大的批次大小（此处为更大的 microbatch）可以获得更高的 BLEU 分数（无需 MBR）。如果您希望将训练损失同步到 wandb，请在 train.py 中自定义您的 wandb 设置（添加您自己的 API KEY）。

DiffuSeq 解码

您需要修改在训练阶段获得的 model_dir 路径。

cd scripts
bash run_decode.sh

为了复现我们论文中表 1 的结果，我们建议将 MBR 候选集的大小设置为 10（使用不同的随机种子运行 10 次）。经验表明，更大的规模可以获得更高的 BLEU 分数。对于多样性指标，在计算时 MBR 候选集的大小应为 3。

加速解码

我们针对 DiffuSeq 定制了 DPM-Solver++ 的实现，以加快其采样速度。

cd scripts
bash run_decode_solver.sh

评估 & MBR

您需要指定解码文本的文件夹。该文件夹应包含来自同一模型但使用不同随机种子进行采样的解码文件。如果未附加 mbr，我们将从该文件夹中的文件计算多样性分数；否则我们将进行 MBR 解码：

cd scripts
python eval_seq2seq.py --folder ../{your-path-to-outputs} --mbr

注意：如果您想将此评估脚本用于其他模型的输出文件，请确保这些输出文件中的同一条记录对应相同的数据。否则多样性分数可能会不准确。

更新

• 2023年10月10日更新：我们对DiffuSeq-v2进行了更新，重点提升了训练和采样的速度。详细信息请参见新分支diffuseq-v2。
• 2023年5月22日更新：我们已在该链接中准备好了剩余任务的检查点及采样结果。
• 2022年11月28日更新：我们已在该链接中准备好了QQP数据集10个随机种子的检查点及采样结果。
• 2023年2月14日更新：我们更新了评估脚本以及论文的最终定稿版本。

如果您有任何问题，欢迎交流讨论。

引用

若我们的论文或代码对您的研究有所帮助，请添加以下引用：

@inproceedings{gong2022diffuseq,
  author = {Gong, Shansan and Li, Mukai and Feng, Jiangtao and Wu, Zhiyong and Kong, Lingpeng},
  booktitle = {国际表征学习大会, ICLR},
  title = {{DiffuSeq}: 基于扩散模型的序列到序列文本生成},
  year = 2023
}

@article{gong2023diffuseqv2,
  title={DiffuSeq-v2: 桥接离散与连续文本空间以加速Seq2Seq扩散模型},
  author={Gong, Shansan and Li, Mukai and Feng, Jiangtao and Wu, Zhiyong and Kong, Lingpeng},
  journal={arXiv预印本 arXiv:2310.05793},
  year={2023}
}

ICLR 2023的DiffuSeq海报。

DiffuSeq 快速上手指南

DiffuSeq 是一个基于扩散模型（Diffusion Models）的序列到序列（Seq2Seq）文本生成工具。它支持端到端训练，无需分类器引导，并在文本生成的质量和多样性上媲美甚至超越传统的自回归（AR）和非自回归（NAR）模型。最新版本 DiffuSeq-v2 显著提升了训练收敛速度和采样效率。

1. 环境准备

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

• 操作系统: Linux (推荐 Ubuntu)
• Python: 3.8 或更高版本
• 深度学习框架: PyTorch (需支持 CUDA)
• 硬件要求:
  • 训练建议使用多张 NVIDIA GPU (如 A100)。
  • 原文示例中使用了 4-8 张 A100 80G 进行大规模训练，小规模实验可适当减少显存需求。
• 前置依赖库:
  • transformers (HuggingFace)
  • 其他依赖详见项目 requirements.txt

提示：国内用户安装 Python 包时，建议指定清华或阿里镜像源以加速下载： pip install -r requirements.txt -i https://pypi.tuna.tsinghua.edu.cn/simple

2. 安装步骤

2.1 克隆代码与安装依赖

首先克隆仓库并安装必要的 Python 包：

git clone https://github.com/Shark-NLP/DiffuSeq.git
cd DiffuSeq
pip install -r requirements.txt

2.2 准备数据集

将处理好的数据集放置在 datasets 文件夹下。数据文件需包含 train.jsonl, test.jsonl, valid.jsonl。

项目论文中使用的四个主要数据集及下载方式如下（需自行下载并放入对应目录）：

任务	数据集名称	说明	下载链接
开放域对话	Commonsense Conversation	CCM 数据集	Google Drive
问题生成	Quasar-T	OpenQA 数据集	Google Drive
文本简化	Wiki-alignment	Wiki-auto 数据集	Google Drive
复述生成	QQP	Kaggle 数据集	Google Drive

注：由于原链接为 Google Drive，国内用户可能需要借助代理或寻找国内镜像资源。

3. 基本使用

3.1 模型训练 (Training)

进入 scripts 目录并运行训练脚本。以下以 QQP (复述生成) 任务为例，展示如何在 4 张 GPU 上启动训练：

cd scripts
bash train.sh

若需自定义参数（如数据集路径、序列长度、优化器等），可直接使用 torch.distributed.launch 调用 run_train.py。以下是复现论文结果的推荐配置示例：

QQP 数据集训练命令：

python -m torch.distributed.launch --nproc_per_node=4 --master_port=12233 --use_env run_train.py --diff_steps 2000 --lr 0.0001 --learning_steps 50000 --save_interval 10000 --seed 102 --noise_schedule sqrt --hidden_dim 128 --bsz 2048 --dataset qqp --data_dir {datasets/QQP} --vocab bert --seq_len 128 --schedule_sampler lossaware --notes qqp

关键参数说明：

• --dataset: 数据集名称标识。
• --data_dir: 数据集所在路径（需包含 train/test/valid.jsonl）。
• --seq_len: 序列最大长度 ($x\oplus y$)。
• --vocab: 分词器类型，推荐使用 bert 或加载预处理的 BPE 词表。
• --learned_mean_embed: (DiffuSeq-v2) 是否使用学习的软吸收状态。
• --use_fp16: (DiffuSeq-v2) 是否开启混合精度训练以加速。

注意：完整训练耗时较长（QQP/QG 任务在 4 张 A100 上约需 2 天）。DiffuSeq-v2 版本通过优化可将训练时间缩短至约 11 小时（2 张 A100）。

3.2 模型推理 (Decoding)

训练完成后，修改脚本中的 model_dir 指向保存的模型检查点路径，然后运行推理脚本：

cd scripts
bash run_decode.sh

• MBR (Minimum Bayes Risk) 解码：为了获得最佳 BLEU 分数，建议设置 MBR 候选集大小为 10（即使用不同随机种子运行 10 次）。
• 多样性评估：计算多样性指标时，MBR 候选集大小通常设为 3。

3.3 加速推理 (Speed-up Decoding)

DiffuSeq 集成了定制的 DPM-Solver++ 以大幅加速采样过程。使用以下命令启动加速推理：

cd scripts
bash run_decode_solver.sh

3.4 结果评估 (Evaluation)

对生成的文本进行评估。指定包含不同种子采样结果的输出文件夹路径：

cd scripts
python eval_seq2seq.py --folder ../{your-path-to-outputs} --mbr

• 若添加 --mbr 参数，则执行 MBR 解码评估。
• 若不加该参数，则直接计算文件夹内文件的多样性得分。

提示：若评估其他模型的输出，请确保不同文件中的同一行对应相同的数据样本，否则多样性评分可能不准确。