简单有效的掩码扩散语言模型 (NeurIPS 2024)

作者：Subham Sekhar Sahoo、Marianne Arriola、Yair Schiff、Aaron Gokaslan、Edgar Marroquin、 Justin T Chiu、Alexander Rush、Volodymyr Kuleshov

[更新日期：2025年4月14日：改进后的实现已在此处提供： DUO Github仓库。]

[更新日期：2025年6月3日：带有KV缓存的MDM： Eso-LMs Github仓库。]

我们提出了MDLM，一种基于掩码的离散扩散语言模型，其特点是采用了一种新颖的（SUBS）替换参数化方法，将吸收态扩散损失简化为经典掩码语言建模损失的混合形式。通过这种方式，我们在LM1B和OpenWebText数据集上实现了扩散模型中的SOTA困惑度指标，同时在多个数据集上与SOTA自回归模型相比也达到了具有竞争力的零样本困惑度。我们在此提供了演示：可在该笔记本或中查看，还有一个视频教程如下：

在这个仓库中，我们发布了：

• MDLM框架。
  1. 基于替换的参数化方法
  2. 简化的掩码扩散过程损失计算
• 基线实现 [示例]：
  1. 自回归模型，在LM1B数据集上达到SOTA自回归性能。
  2. 基于分数熵的离散扩散SEDD。
  3. 吸收态D3PM的高效实现，在LM1B数据集上超越了之前的SOTA文本扩散模型SEDD。
• 采样器
  1. D3PM中提出的祖先采样法。
  2. SEDD中提出的解析采样法。
  3. 我们提出的高效采样器，它
     • 使MDLM的速度比现有扩散模型快约3到4倍。[示例]
     • 支持半自回归(SAR)生成。[示例]

代码组织

1. main.py: 训练和评估流程
2. noise_schedule.py: 噪声调度
3. diffusion.py: 正向/反向扩散
4. dataloader.py: 数据加载器
5. utils.py: 学习率调度器、日志记录、fsspec处理
6. models/: 去噪网络架构。支持DiT、AR变换器和Mamba
7. configs/: 数据集/去噪网络/噪声调度/学习率调度的配置文件
8. scripts/: 用于训练/评估的Shell脚本

在此仓库中开始使用

要开始使用，请创建一个包含所需依赖项的conda环境。

conda env create -f requirements.yaml
conda activate mdlm

创建以下目录以存储保存的模型和slurm日志：

mkdir outputs
mkdir watch_folder

然后以批处理作业的形式运行训练：

sbatch scripts/train_owt_mdlm.sh

检查点

我们已将基于OpenWebText数据集训练了100万步的MDLM模型上传至Huggingface Hub 🤗： kuleshov-group/mdlm-owt 此外，我们还在这个Google Drive文件夹中发布了在OpenWebText数据集上训练的AR和SEDD基线的检查点。

实验复现

下面我们将描述复现论文中实验所需的步骤。 在整个过程中，运行实验的主要入口是main.py脚本。我们还在scrips/目录中提供了用于启动预训练和下游微调实验的示例slurm脚本。

生成样本

sampling.predictor 参数指定了采样器，其取值如下：

• ddpm_cache: 我们提出的采样器，比 D3PM 和 SEDD 中提出的采样器快 ~3-4 倍。
• ddpm: D3PM 中提出的祖先采样法。
• analytic: SEDD 中提出的解析采样器。

下表报告了在单块 A5000 GPU 上以 batch_size=1 生成 64 个样本的壁时。$T$ 表示逆向过程的时间离散化步数。

	$T=5k (\downarrow)$	$T=10k (\downarrow)$
SEDD	127.1	229.3
MDLM + ddpm	113.8	206.6
MDLM +ddpm_cache	40.1	60.4

要从预训练模型生成样本，请使用以下命令之一：

Huggingface 模型

python main.py \
  mode=sample_eval \
  eval.checkpoint_path=kuleshov-group/mdlm-owt \
  data=openwebtext-split  \
  model.length=1024  \
  sampling.predictor=ddpm_cache  \
  sampling.steps=1000 \
  loader.eval_batch_size=1 \
  sampling.num_sample_batches=10 \
  backbone=hf_dit

本地检查点

python main.py \
  mode=sample_eval \
  eval.checkpoint_path=/path/to/checkpoint/mdlm.ckpt \
  data=openwebtext-split  \
  model.length=1024  \
  sampling.predictor=ddpm_cache  \
  sampling.steps=10000 \
  loader.eval_batch_size=1 \
  sampling.num_sample_batches=1 \
  backbone=dit

半自回归样本生成

MDLM 还可以以半自回归 (SAR) 方式生成任意长度的样本。我们在单块 3090 GPU 上生成了 200 条长度为 2048 个标记的序列，并在预训练的 GPT-2 模型下评估生成困惑度。下表显示，除了获得更好的生成困惑度外，与 SSD-LM 相比，MDLM 的 SAR 解码速度还快了 25-30 倍。

	生成困惑度 ($\downarrow$)	每序列秒数 ($\downarrow$)
SSD-LM	35.43	2473.9
MDLM +ddpm_cache	27.18	89.3

生成困惑度：Generation Perplexity，每序列秒数：Seconds per Sequence

python main.py \
  mode=sample_eval \
  eval.checkpoint_path=kuleshov-group/mdlm-owt \
  data=openwebtext-split \
  parameterization=subs \
  model.length=1024  \
  sampling.predictor=ddpm_cache  \
  sampling.steps=1000 \
  loader.eval_batch_size=1 \
  sampling.num_sample_batches=2 \
  sampling.semi_ar=True \
  sampling.stride_length=512 \
  sampling.num_strides=2 \
  backbone=hf_dit

训练

要在 OpenWebText 数据集上从头开始训练 MDLM，请使用以下命令：

python main.py \
  model=small \
  data=openwebtext-split \
  wandb.name=mdlm-owt \
  parameterization=subs \
  model.length=1024 \
  eval.compute_generative_perplexity=True \
  sampling.steps=1000

参数 loader.batch_size 和 loader.eval_batch_size 允许您控制全局批量大小和每 GPU 的批量大小。如果 loader.batch_size * num_gpus 小于全局批量大小，PyTorch Lightning 将采用梯度累积策略。您也可以使用 sbatch scripts/train_owt_mdlm.sh 命令在 Slurm 上启动训练任务。用于训练自回归和 SEDD 基线的 Slurm 脚本分别为：scripts/train_lm1b_ar.sh、scripts/train_owt_sedd.sh。

评估

要计算测试困惑度，请使用 mode=ppl_eval。示例脚本见 scripts/ 目录。以下是在 OpenWebText 上进行困惑度评估的示例命令：

python main.py \
  mode=ppl_eval \
  loader.batch_size=16 \
  loader.eval_batch_size=16 \
  data=openwebtext-split \
  model=small \
  parameterization=subs \
  backbone=dit \
  model.length=1024 \
  eval.checkpoint_path=/path/to/checkpoint/mdlm.ckpt \
  +wandb.offline=true

基线评估

我们在此 Google Drive 文件夹 中发布了在 OpenWebText 上训练的 SEDD 和 AR 基线的检查点。请下载检查点：ar.ckpt、sedd.ckpt，并使用以下命令计算测试困惑度：

自回归

python main.py \
  mode=ppl_eval \
  loader.batch_size=16 \
  loader.eval_batch_size=16 \
  data=openwebtext-split \
  model=small-ar \
  parameterization=ar \
  backbone=ar \
  model.length=1024 \
  eval.checkpoint_path=/path/to/checkpoint/ar.ckpt \
  +wandb.offline=true

SEDD

python main.py \
  mode=ppl_eval \
  loader.batch_size=16 \
  loader.eval_batch_size=16 \
  data=openwebtext-split \
  model=small \
  parameterization=sedd \
  backbone=dit \
  model.length=1024 \
  eval.checkpoint_path=/path/to/checkpoint/sedd.ckpt \
  time_conditioning=True \
  sampling.predictor=analytic \
  +wandb.offline=true

致谢

本仓库基于 SEDD 构建。

引用

@inproceedings{
sahoo2024simple,
title={Simple and Effective Masked Diffusion Language Models},
author={Subham Sekhar Sahoo and Marianne Arriola and Aaron Gokaslan and Edgar Mariano Marroquin and Alexander M Rush and Yair Schiff and Justin T Chiu and Volodymyr Kuleshov},
booktitle={The Thirty-eighth Annual Conference on Neural Information Processing Systems},
year={2024},
url={https://openreview.net/forum?id=L4uaAR4ArM}
}

MDLM 快速上手指南

MDLM (Masked Diffusion Language Models) 是一个基于掩码扩散机制的语言模型框架，通过创新的 SUBS 参数化方法，在保持生成质量的同时显著提升了推理速度（比现有扩散模型快 3-4 倍）。

环境准备

• 操作系统: Linux (推荐 Ubuntu 18.04+)
• 硬件要求: NVIDIA GPU (支持 CUDA)，显存建议 16GB 以上以进行训练或长序列生成。
• 软件依赖:
  • Python 3.8+
  • Conda (推荐用于环境管理)
  • PyTorch (需匹配你的 CUDA 版本)
  • Slurm (可选，仅用于集群批量训练)

安装步骤

1. 克隆仓库

   git clone https://github.com/s-sahoo/mdlm.git
   cd mdlm
   

2. 创建并激活 Conda 环境 使用项目提供的 requirements.yaml 文件安装所有依赖。 (注：若下载速度慢，可配置 conda 使用清华或中科大镜像源)

   conda env create -f requirements.yaml
   conda activate mdlm
   

3. 创建必要目录 创建用于存储模型输出和日志的文件夹：

   mkdir outputs
   mkdir watch_folder
   

基本使用

1. 快速生成文本样本 (使用预训练模型)

最简单的使用方式是从 Hugging Face 加载已在 OpenWebText 上预训练的模型进行采样。以下命令使用高效的 ddpm_cache 采样器生成文本：

python main.py \
  mode=sample_eval \
  eval.checkpoint_path=kuleshov-group/mdlm-owt \
  data=openwebtext-split  \
  model.length=1024  \
  sampling.predictor=ddpm_cache  \
  sampling.steps=1000 \
  loader.eval_batch_size=1 \
  sampling.num_sample_batches=10 \
  backbone=hf_dit

参数说明：

• sampling.predictor=ddpm_cache: 推荐使用此采样器，速度比标准 DDPM 快 3-4 倍。
• sampling.steps: 去噪步数，步数越少速度越快，步数越多质量通常越高。
• backbone=hf_dit: 指定从 Hugging Face 加载 DiT 架构模型。

2. 半自回归生成长文本 (Semi-Autoregressive)

MDLM 支持以半自回归方式生成任意长度的序列，且速度极快。以下示例生成长度为 2048 的文本：

python main.py \
  mode=sample_eval \
  eval.checkpoint_path=kuleshov-group/mdlm-owt \
  data=openwebtext-split \
  parameterization=subs \
  model.length=1024  \
  sampling.predictor=ddpm_cache  \
  sampling.steps=1000 \
  loader.eval_batch_size=1 \
  sampling.num_sample_batches=2 \
  sampling.semi_ar=True \
  sampling.stride_length=512 \
  sampling.num_strides=2 \
  backbone=hf_dit

3. 从头开始训练 (可选)

若需在 OpenWebText 数据集上从头训练模型，可运行以下命令：

python main.py \
  model=small \
  data=openwebtext-split \
  wandb.name=mdlm-owt \
  parameterization=subs \
  model.length=1024 \
  eval.compute_generative_perplexity=True \
  sampling.steps=1000

提示：如需在 Slurm 集群上运行训练任务，可使用 sbatch scripts/train_owt_mdlm.sh。