扩散模型在图像生成领域早就不是新鲜词，但从图像迁移到文本生成时，很多人会卡在同一个地方：文本是离散符号，没有连续的“加噪”空间，那扩散过程要怎么定义？康奈尔大学相关课程在讲深度生成模型时，专门把扩散语言模型（Diffusion Language Model）从数学直觉讲到代码实现，这套思路对想深入生成式 AI 底层原理的开发者非常友好。本文按“理论 → 设计 → 代码 → 排错”的顺序，带大家从零实现一个可运行的 Mask 扩散语言模型，重点讲清楚离散扩散怎么做、训练目标和采样过程怎么写，并附完整 PyTorch 代码。

如果你已经了解自回归语言模型（比如 GPT 系列）的基本结构，但还没搞懂扩散模型如何生成文本；或者你学过图像扩散模型，但不确定离散 token 该怎么处理噪声；再或者你只是想亲手跑通一个最小可实现的“扩散版 BERT”来加深理解，那么这篇文章都很适合你。读完本文，你会掌握扩散语言模型的核心概念、正向加噪与反向去噪的实现方式、训练 loss 的设计思路，以及迭代式采样解码的代码写法。

1. 扩散语言模型是什么：从连续噪声到离散 Mask

1.1 图像扩散模型的启发

要理解扩散语言模型，先要回顾图像扩散模型的基本思想。以 DDPM（Denoising Diffusion Probabilistic Models）为例，它的思路其实很直观：前向过程逐步向一张干净图片加入高斯噪声，经过足够多步之后，图片几乎变成纯噪声；反向过程则训练一个神经网络，每次预测“上一步的噪声”或“上一步去噪后的图像”，从而一步步把噪声还原成清晰图片。

这个框架有一个关键点：图像本身是连续数据，噪声可以叠加在像素值上。但文本是由离散 token 组成的，不能给“love”这个词加上一个高斯噪声向量再让它变成另一个词。因此，扩散模型迁移到文本领域时，最主流的做法是换一种噪声形式：不使用高斯噪声，而是使用“掩码噪声”，也就是把一部分 token 替换成特殊的[MASK]符号。

1.2 Mask 扩散：把加噪变成“盖住词”

Mask 扩散（Mask Diffusion，有时也叫 Absorbing Diffusion）的思想和 BERT 的 Masked Language Model 非常相似：给定一句文本，随机盖住其中一部分 token，然后训练模型根据上下文预测被盖住的内容。

但它和 BERT 有两个本质区别：

• BERT 训练时 mask 比例是固定的（比如 15%），没有“时间步”的概念；
• 扩散语言模型将 mask 比例建模为一个从 0 到 1 连续变化的时间步，mask 比例越高，噪声越强，越接近“纯噪声”状态。反向过程通过在时间步上逐步降低 mask 比例，迭代地修复被盖住的 token，最终生成完整句子。

这个“由全 mask 到部分 mask，再到全文本”的迭代过程，就是非自回归生成。它和 GPT 那种从左到右逐个生成 token 的自回归方式完全不同。

1.3 为什么需要扩散语言模型

自回归模型固然强大，但它有一个天然限制：生成时必须严格从左到右逐个预测。这导致生成速度受限于序列长度，也难以进行显式的“全局修改”。

扩散语言模型的价值在于：

• 支持并行生成：每个位置可以在同一轮迭代中被同时预测，适合长文本并行解码；
• 支持可控编辑：可以从“已知部分 token + 未知 mask”的状态出发，只修复未知部分，天然适合文本填充、改写、纠错等任务；
• 建模全局依赖：模型在预测每个位置时可以看到整句上下文，而不是只能看到左侧信息。

因此，近年来越来越多的研究开始把扩散模型用于语言建模，代表工作包括 Diffusion-LM、DiffusionBERT、MDLM 等。本文实现的正是其中最简单、最适合入门的 Mask 扩散方案。

2. 扩散语言模型的核心原理

2.1 前向过程：如何对 token 加噪

假设原始文本序列为 x0，长度为 L。前向过程定义一个噪声调度：对于一个随机采样的时间步 t ∈ (0,1]，我们用 mask 概率 η_t 控制每个 token 被替换为[MASK]的概率。

具体来说，对每个位置 i：

• 以概率 η_t 将 x0_i 替换为[MASK]；
• 以概率 1 - η_t 保持 x0_i 不变。

这里 η_t 随 t 单调递增。t 越接近 1，mask 比例越高，序列信息丢失越多；t 越接近 0，mask 比例越低，序列越接近原始文本。当 t=1 时，理论上所有 token 都被 mask 覆盖，序列变成“纯噪声”。

这个过程的含义很直观：扩散语言模型中的“噪声”就是“信息缺失”，模型的任务就是学会根据剩余可见信息，把缺失的部分补回来。

2.2 反向过程：如何修复 token

反向过程是扩散模型的核心。我们需要训练一个神经网络 p_θ，输入是加噪后的序列 x_t 和时间步 t，输出是原始序列 x0 的预测分布。

在 Mask 扩散中，对每个被 mask 的位置，模型需要预测该位置原来的 token；对未被 mask 的位置，模型不需要做任何生成。因此训练目标非常简单：只对被 mask 的位置计算交叉熵损失。

这个损失函数本质上和 BERT 的 MLM 损失一模一样。区别在于：扩散模型在每个 batch 中会随机采样不同的 t，让模型看到不同 mask 比例下的序列，从而学会在不同噪声水平下进行修复。采样时则从全 mask 序列出发，按照时间步从大到小逐步修复部分 token。

为了方便理解，你可以把训练过程类比成：同一个模型既要会做“只遮盖 10% 的完形填空”，也要会做“只保留 20% 信息的高度残缺复原”。模型通过在不同难度之间切换，逐渐学会更全局的语言理解能力。

2.3 训练目标的形式化表示

给定原始序列 x0，随机采样时间步 t，得到 mask 后的序列 x_t。设被 mask 的位置集合为 M，那么训练损失可以写为：

L = E[ - ∑(i∈M) log p_θ(x0_i | x_t, t) ]

也就是，模型只在被 mask 的位置上输出一个关于原始 token 的分布，然后和真实 token 计算交叉熵。未 mask 位置不参与 loss 计算，因为它们的信息已经可见，模型即使预测正确也不能说明什么。

在代码实现中，我们可以通过一个布尔 mask 矩阵选出需要计算 loss 的位置，直接调用F.cross_entropy完成计算。

2.4 采样：迭代式去噪生成

训练完成后，采样过程如下：

1. 初始化一个长度为 L 的全[MASK]序列，即最“嘈杂”的状态；
2. 从高噪声时间步到低噪声时间步迭代。每轮输入当前序列和时间步 t，让模型预测每个 mask 位置的 token 分布；
3. 根据预测置信度，选出本轮要修复的位置，将置信度最高的若干个位置替换为预测 token；
4. 其余位置继续保留[MASK]，进入下一轮；
5. 直到所有[MASK]都被替换为具体 token，生成结束。

这种“先修复确定的部分，把不确定部分留到后面结合更多信息再决策”的策略，就是迭代式解码的核心思想。它和并行解码、自回归解码都不一样，属于典型的非自回归生成方式。

3. 环境准备与实验设计

3.1 环境版本

本文代码基于 PyTorch 编写，不需要额外的深度学习框架。推荐环境如下：

• 操作系统：Windows / Linux / macOS 均可；
• Python：3.9 或更高版本；
• PyTorch：2.x 版本（1.13+ 应该也能运行，但建议直接使用 2.x）；
• 其他依赖：NumPy（PyTorch 自带依赖中通常已包含）。

如果你还没有安装 PyTorch，可以根据自己的 CUDA 版本到 PyTorch 官网选择对应命令安装。没有 GPU 也不用担心，本文模型很小，CPU 完全可以跑通训练和采样流程。

pip install torch

3.2 实验设计思路

为了让代码足够简单、可复现，本文不依赖 HuggingFace 数据集，也不使用 GPT-2 等大型预训练模型。我们直接构建一个字符级 tokenizer，使用一个内置的小型英文语料库进行训练。

这样做有几个好处：

• 词表很小（几十个字符），模型参数量少，CPU 也能快速迭代；
• 没有复杂的数据下载和预处理流程，读者可以一次性跑通；
• 后续替换为自己的数据非常方便，只需把CORPUS换成自己的文本列表。

项目结构大致如下：

diffusion_lm/ ├── train.py # 完整训练 + 采样代码 └── README.md # 可选

代码集中在一个文件中方便演示，实际项目中建议拆分为数据、模型、扩散流程、训练器等模块。

4. 从零实现 Mask 扩散语言模型

下面进入正题。我会按照“数据准备 → 加噪过程 → 模型定义 → 训练循环 → 采样生成”的顺序，逐步实现并解释每一段代码。所有代码可以按顺序保存为train.py。

4.1 导入依赖与定义超参数

import math import torch import torch.nn as nn import torch.nn.functional as F from torch.utils.data import Dataset, DataLoader # ---------- 超参数 ---------- MAX_LEN = 32 # 序列最大长度 BATCH_SIZE = 16 # 批大小 D_MODEL = 128 # 模型隐藏维度 NHEAD = 4 # 多头注意力头数 NUM_LAYERS = 4 # Transformer Encoder 层数 HIDDEN_DIM = 256 # FFN 中间维度 DROPOUT = 0.1 # Dropout 比例 LR = 1e-3 # 学习率 NUM_STEPS = 3000 # 训练步数 T = 1000 # 离散时间步总数，用于采样 t

超参数这里需要解释一下。MAX_LEN=32意味着模型每次最多处理 32 个字符的序列；T=1000不是真实训练 1000 步，而是把时间步离散化成 1000 档，训练时从 0~999 中随机采样一个 step，再映射到(0,1]区间内的噪声水平。这样可以让模型看到连续的噪声水平变化。

4.2 准备语料与字符级 Tokenizer

这里使用 8 句经典英文名言作为训练语料。数据量确实很小，但足够验证整个 pipeline 是通的。你把这里的CORPUS换成任何 txt 文本列表，代码都能直接跑。

CORPUS = [ "the quick brown fox jumps over the lazy dog", "a journey of a thousand miles begins with a single step", "do not go where the path may lead go instead where there is no path and leave a trail", "we are what we repeatedly do excellence then is not an act but a habit", "it always seems impossible until it is done", "the best way to predict the future is to invent it", "life is what happens when you are busy making other plans", "knowing yourself is the beginning of all wisdom", ]

接下来定义CharTokenizer。这是一个极简的字符级 tokenizer，它会统计语料中出现的所有字符，并为每个字符分配一个 ID。

为了配合扩散模型，我们预留两个特殊 ID：

• 0：<pad>，用于填充短句；
• 1：<mask>，用于替换被噪声覆盖的 token。

class CharTokenizer: def __init__(self, texts): chars = set("".join(texts)) self.chars = sorted(chars) # 0 保留给 pad，1 保留给 mask self.pad_id = 0 self.mask_id = 1 self.stoi = {c: i + 2 for i, c in enumerate(self.chars)} self.itos = {i: c for c, i in self.stoi.items()} self.itos[self.pad_id] = "<pad>" self.itos[self.mask_id] = "<mask>" self.vocab_size = len(self.stoi) + 2 def encode(self, text: str, max_len: int): ids = [self.stoi[c] for c in text.lower() if c in self.stoi] if len(ids) > max_len: ids = ids[:max_len] return ids def decode(self, ids): return "".join(self.itos.get(i, "") for i in ids if i > 1)

decode会跳过pad_id和mask_id，这样即使采样过程有些位置没有被成功替换，输出中也不会出现奇怪的占位符。当然，理想情况下采样结束不会留下 mask。

class TextDataset(Dataset): def __init__(self, texts, tokenizer, max_len): self.tokenizer = tokenizer self.max_len = max_len self.samples = [tokenizer.encode(t, max_len) for t in texts] def __len__(self): return len(self.samples) def __getitem__(self, idx): ids = self.samples[idx] if len(ids) > self.max_len: ids = ids[:self.max_len] x = torch.full((self.max_len,), self.tokenizer.pad_id, dtype=torch.long) x[:len(ids)] = torch.tensor(ids, dtype=torch.long) return x

TextDataset的作用是把每句文本转成固定长度的 ID 序列。短句后面用<pad>填充，超出长度则截断。这里要强调：我们在加噪时会刻意避免 maskpad位置，因为 pad 是“无意义位置”，不应该让模型学习预测它。

4.3 实现正弦时间步嵌入

时间步 t 是一个连续值，我们需要把它编码成向量，然后注入 Transformer。这里采用图像扩散模型中经典的 sinusoidal 位置编码思路，把 t 编码到不同频率的正弦和余弦函数上。

def sinusoidal_timestep_embedding(t, d_model): half = d_model // 2 freqs = torch.exp( -math.log(10000) * torch.arange(half, device=t.device).float() / half ) args = t.unsqueeze(-1) * freqs return torch.cat([torch.sin(args), torch.cos(args)], dim=-1)

为什么需要时间步嵌入？因为模型必须知道当前噪声水平。同样一句“the quick brown fox”，如果只 mask 了一个词，模型应该大胆地补全；如果 mask 了一半，模型应该更加保守地结合全局信息。时间步嵌入让模型具备这种“根据噪声强度调整策略”的能力。

随后用一个 MLP 把正弦嵌入映射到模型维度：

class TimestepEmbedding(nn.Module): def __init__(self, d_model): super().__init__() self.net = nn.Sequential( nn.Linear(d_model, d_model * 2), nn.SiLU(), nn.Linear(d_model * 2, d_model), ) def forward(self, t): # t: (batch_size,) emb = sinusoidal_timestep_embedding(t, self.net[0].in_features) return self.net(emb)

4.4 实现加噪过程

加噪过程的核心函数是make_timesteps和add_noise。

make_timesteps负责随机采样时间步：

def make_timesteps(batch_size, T, device): steps = torch.randint(0, T, (batch_size,), device=device) t = (steps.float() + 1) / T # 映射到 (0, 1] return t, steps

这里把 0~999 的整数步映射到(0,1]区间。之所以加 1 再除以 T，是为了避免 t=0 的情况，保证每个 batch 都有真实噪声。

add_noise负责 mask token：

def add_noise(x, mask_id, t): # x: (batch_size, seq_len)，t: (batch_size,) B, L = x.shape rand = torch.rand(B, L, device=x.device) noise_level = t.unsqueeze(-1) # (B, 1) valid = (x != 0) # 不对 pad 位置加噪 is_mask = (rand < noise_level) & valid x_noised = x.clone() x_noised[is_mask] = mask_id return x_noised, is_mask

这里的核心设计是：每个 batch 中的每个样本有一个不同的时间步，因此每个样本的 mask 比例不同。模型必须学会根据时间步信息调整输出，而不是简单记住一种 mask 比例。

4.5 定义 Mask 扩散模型主体

模型主体是一个标准的 Transformer Encoder。输入是加噪后的 token 序列和时间步向量，输出是每个位置的词汇分布 logits。

class MaskDiffusionModel(nn.Module): def __init__(self, vocab_size, mask_id, d_model=128, nhead=4, num_layers=4, hidden_dim=256, dropout=0.1, max_len=64): super().__init__() self.mask_id = mask_id self.pad_id = 0 self.d_model = d_model self.token_embedding = nn.Embedding(vocab_size, d_model, padding_idx=0) self.time_embedding = TimestepEmbedding(d_model) self.pos_embedding = nn.Embedding(max_len, d_model) encoder_layer = nn.TransformerEncoderLayer( d_model=d_model, nhead=nhead, dim_feedforward=hidden_dim, dropout=dropout, batch_first=True, ) self.transformer = nn.TransformerEncoder(encoder_layer, num_layers=num_layers) self.ln_f = nn.LayerNorm(d_model) self.head = nn.Linear(d_model, vocab_size) def forward(self, x, t): # x: (B, L)，t: (B,) B, L = x.shape token_emb = self.token_embedding(x) # (B, L, D) time_emb = self.time_embedding(t).unsqueeze(1) # (B, 1, D) positions = torch.arange(L, device=x.device) pos_emb = self.pos_embedding(positions).unsqueeze(0) # (1, L, D) h = token_emb + pos_emb + time_emb padding_mask = (x == self.pad_id) # (B, L) h = self.transformer(h, src_key_padding_mask=padding_mask) h = self.ln_f(h) logits = self.head(h) # (B, L, vocab_size) return logits

结构非常简单：token embedding 加上位置编码，再加上时间步编码，然后送入 Transformer Encoder。最后通过线性层输出词表大小的 logits。

这里要说明两点：

• 使用双向注意力（Encoder），而不是像 GPT 那样的因果注意力。因为扩散模型在修复某个位置时，需要看到左右两侧的上下文，这和有损补全任务是匹配的。
• padding mask 只屏蔽pad位置，mask位置不应该被屏蔽。因为这些位置是模型要预测的位置，注意力需要集中在它们和其他可见 token 之间的交互上。

4.6 训练循环

训练逻辑可以总结为：随机采样 t → 对输入加噪 → 模型预测 → 只对 mask 位置计算交叉熵。

def train(): tokenizer = CharTokenizer(CORPUS) dataset = TextDataset(CORPUS, tokenizer, MAX_LEN) dataloader = DataLoader(dataset, batch_size=BATCH_SIZE, shuffle=True) device = torch.device("cuda" if torch.cuda.is_available() else "cpu") model = MaskDiffusionModel( vocab_size=tokenizer.vocab_size, mask_id=tokenizer.mask_id, d_model=D_MODEL, nhead=NHEAD, num_layers=NUM_LAYERS, hidden_dim=HIDDEN_DIM, dropout=DROPOUT, max_len=MAX_LEN, ).to(device) optimizer = torch.optim.AdamW(model.parameters(), lr=LR) train_iter = iter(dataloader) model.train() for step in range(NUM_STEPS): try: x0 = next(train_iter) except StopIteration: train_iter = iter(dataloader) x0 = next(train_iter) x0 = x0.to(device) t, _ = make_timesteps(x0.size(0), T, device) x_noised, is_mask = add_noise(x0, tokenizer.mask_id, t) logits = model(x_noised, t) # (B, L, V) loss = F.cross_entropy( logits[is_mask], x0[is_mask] ) optimizer.zero_grad() loss.backward() optimizer.step() if step % 200 == 0: print(f"step {step:5d} | loss {loss.item():.4f}") return model, tokenizer

这里有一个容易踩坑的细节：logits[is_mask]返回的是一个二维张量，第一维是所有被 mask 位置的展平数量。由于is_mask是布尔张量，PyTorch 会自动选取 mask 位置的 logits。同理x0[is_mask]选取被 mask 位置的原始 token。这样能保证模型只在缺失位置学习预测，而不会在可见位置浪费表达能力。

如果你发现 loss 收敛很慢，可以检查是否把pad位置也计入了学习目标。本文代码已经通过valid条件排除了 pad 位置，因此不会出现预测 pad 的无意义学习。

4.7 采样生成

采样函数是扩散模型最有趣的部分。我们采用“置信度优先修复”的策略：每一轮只修复一部分置信度最高的 mask 位置，其余位置留到后续轮次处理。

@torch.no_grad() def sample(model, tokenizer, length=MAX_LEN, num_steps=8, temperature=0.8): model.eval() device = next(model.parameters()).device # 初始化：前 length 个位置为 mask，其余为 pad x = torch.full((1, MAX_LEN), tokenizer.pad_id, dtype=torch.long, device=device) x[:, :length] = tokenizer.mask_id for step in range(num_steps): # 当前 mask 比例作为时间步 mask_prop = x.eq(tokenizer.mask_id).float().mean() t = torch.tensor([mask_prop], device=device) logits = model(x, t) probs = F.softmax(logits / temperature, dim=-1) confidence, preds = probs.max(dim=-1) # (1, L) mask_positions = (x == tokenizer.mask_id) remaining = mask_positions.sum().item() if remaining == 0: break remaining_steps = num_steps - step k = (remaining + remaining_steps - 1) // remaining_steps k = max(1, k) cand_conf = confidence.clone() cand_conf[~mask_positions] = -1e9 topk = torch.topk(cand_conf.flatten(), k).indices x[0, topk] = preds.flatten()[topk] return tokenizer.decode(x[0].tolist())

解释几个关键点：

• temperature控制预测分布的平滑程度。温度越低，模型越倾向于选择概率最高的 token；温度越高，生成越多样。文本生成任务中，一般取值在 0.7 到 1.0 之间比较合适。
• k的计算保证每轮至少修复 1 个位置，并且把剩余 mask 大致均匀分配到后续轮次。
• 时间步 t 直接用当前序列中 mask 的占比，而不是训练时随机采样的那个 t。这样做的好处是采样过程与观测状态保持一致，模型看到的噪声水平就是真实噪声水平。

5. 运行脚本与效果观察

在主函数中调用训练和采样：

if __name__ == "__main__": model, tokenizer = train() print("\n生成结果：") for i in range(3): text = sample(model, tokenizer, length=24, num_steps=8, temperature=0.8) print(f"sample {i+1}: {text}")

运行命令：

python train.py

预期你会看到类似下面的训练日志：

step 0 | loss 3.8912 step 200 | loss 1.4127 step 400 | loss 0.8946 step 600 | loss 0.6823 ... step 2800 | loss 0.2104

由于语料很小且是字符级模型，生成结果不会像 GPT 那样流畅，但你应该能看到高频单词的拼写逐渐合理，比如类似the、is、of这样的字符组合会频繁出现。如果你把CORPUS换成几万条英文句子，并适当增加NUM_STEPS和模型规模，生成质量会有明显提升。

这里提醒一句：不同随机种子、不同超参数下的结果差异很大。不要因为第一次生成效果一般就认为代码错了，可以先观察 loss 是否持续下降。如果 loss 能稳定下降，说明扩散模型的训练流程基本正确。

6. 常见问题与排查思路

在实际运行代码或改编代码时，下面几个问题最容易出现。我把现象、原因和解决方案整理成表格，方便快速定位。

如果你遇到类似“loss 下降但生成效果差”的情况，优先检查训练与采样时的 mask 语义是否一致。例如训练时 t=0.8 意味着每个 token 有 80% 概率被 mask，采样时如果当前序列实际只有 30% 的 mask 位置，却把 t 设成 0.8，模型就会产生错误的预测。本文的采样函数中把 t 设置为当前 mask 占比，正是为了避免这种不一致。

另一个高发问题是：Transformer Encoder 的 padding mask 传错。由于我们同时存在pad_id=0和mask_id=1，很容易在构造 mask 时把两者混在一起。记住一个原则：pad 是模型需要忽略的位置，mask 是模型需要预测的位置，二者必须区分。

7. 最佳实践与工程建议

当你把教学代码迁移到真实项目时，以下几个方向值得重点关注。

7.1 让时间步定义更规范

建议把时间步调度单独抽象成一个函数，甚至在配置文件中定义。比如采用 linear schedule 或 cosine schedule，不同的调度会影响训练难度。

def noise_level_from_step(step, T): return (step + 1) / T

真实项目中，最好把噪声调度定义在独立的模块中，这样训练和采样共用同一个调度函数，避免语义不一致。

7.2 使用 EMA 平滑模型参数

扩散模型训练中一个常见的工程技巧是 Exponential Moving Average（EMA）。简单来说，就是额外维护一份“参数平均值”，在训练后期用平均值做采样，通常能明显提升生成稳定性。实现如下：

class EMA: def __init__(self, model, decay=0.999): self.model = model self.decay = decay self.shadow = {k: v.clone() for k, v in model.state_dict().items()} def update(self): for k, v in self.model.state_dict().items(): self.shadow[k] = self.decay * self.shadow[k] + (1 - self.decay) * v def apply(self): self.model.load_state_dict(self.shadow)

在每一步 optimizer.step() 之后调用ema.update()，采样前调用ema.apply()。注意采样后要不要恢复模型原始参数，可以根据你的代码结构决定。

7.3 使用更好的 Tokenizer

字符级 tokenizer 适合教学，但生产级扩散语言模型不应该使用它。真实项目建议使用 HuggingFace 的 GPT-2 BPE tokenizer 或 WordPiece tokenizer。词表变大之后，head层的输出维度也会变大，计算量会显著增加，需要考虑梯度 checkpoint 或混合精度训练。

7.4 采样策略可以进一步优化

本文使用的是最简单的置信度优先修复。实际上，更好的策略还包括“部分噪声采样”：每轮修复一部分位置后，再随机 mask 回一部分位置，从而增加生成多样性。进阶读者可以阅读 MDLM 论文中的采样细节，并尝试在本文代码基础上扩展。

7.5 控制训练与生成的计算资源

这个模型虽然小，但 Transformer Encoder 的注意力复杂度是 O(L²)。当序列长度从 32 提升到 512 时，内存和耗时都会快速上升。生产环境中建议：

• 对长文本做分块处理；
• 使用混合精度训练，显存和速度都能得到优化；
• 用torch.compile加速模型；
• 采样时限制迭代次数，一般 8 到 16 轮就足够，不需要太多。

7.6 安全与版本管理提醒

如果你把这个模型用于业务场景，需要注意几个数据安全点：

• 训练语料必须经过合法授权，不要直接爬取未授权文本；
• 生成内容具备随机性，发布到公网前需要经过内容审核；
• 模型权重和训练数据要纳入版本管理，方便回滚。

8. 下一步学习建议

这篇文章实现的 Mask 扩散语言模型是一个最小闭环，但它已经包含了扩散模型最核心的组件：前向加噪、反向预测、时间步嵌入、迭代采样。如果你想继续深入，可以从以下几个方向入手：

• 阅读 Diffusion-LM 原文，理解连续扩散在语言模型上的实现；
• 阅读 D3PM 论文，了解更一般的离散扩散框架；
• 阅读 MDLM 论文，学习 Mask 扩散的完整数学推导和更高质量的采样策略；
• 把代码中的 Transformer Encoder 换成预训练 BERT 或 GPT 结构，观察扩散模型与预训练模型的结合效果；
• 尝试把扩散模型应用到可控文本生成、文本填充或语义编辑任务上，看看它相比自回归模型有哪些优势。

动手实验时，建议不要一上来就追求大模型和大语料。先在小数据集上观察 loss 曲线和生成结果的变化，再逐步扩大规模。生成式模型的问题往往不是“模型不够大”，而是“噪声调度、采样策略和训练目标没有对齐”。把这三者想清楚，你对扩散语言模型的理解会比单纯刷几篇文章要深刻得多。希望这篇从实现角度出发的教程，能帮你把扩散模型在文本领域的玩法真正跑通。