近几年做文本生成相关实践时我越来越常听到一个方向扩散式语言模型。刚开始接触这个概念很多人会下意识觉得它和扩散模型画图是同一套东西直接搬到文本上。实际动手后发现文本是离散符号图像是连续像素想把扩散过程用在语言上关键就在“怎么把离散文本变成可加噪、可去噪的连续表示”。这篇文章就围绕这个核心问题展开先讲清楚扩散式语言模型的基本原理和主流实现思路再给出一个可以跑通小规模实验的代码框架最后整理训练过程中常见的问题和工程建议。1. 扩散式语言模型是什么1.1 从自回归生成说起主流文本生成模型大多采用自回归架构也就是一个 token 一个 token 地预测下一个词。GPT 系列、LLaMA、ChatGLM 等模型的核心思路都是给定前文 x1, x2, ..., x_{t-1}预测 P(xt | x1, x2, ..., x_{t-1})这种方式的优点是生成质量高、训练稳定推理时也容易控制。但它的缺点同样明显生成速度受限于串行解码无法并行输出整段文本。生成过程是单向的很难直接在全局层面进行修改和重构。对于填空、文本修复、受控改写这类任务自回归模型需要专门设计训练目标天然不占优势。扩散式语言模型Diffusion Language ModelDLM则尝试换一条路线先让模型学习一个“从纯噪声逐步还原成完整文本”的去噪过程。训练完成后推理时从一个随机噪声向量出发通过多步迭代逐渐生成整段文本。这种生成方式天然支持并行解码也为全局编辑、属性控制等任务提供了更灵活的空间。1.2 扩散式语言模型的核心思想扩散模型在图像生成中的基本流程是前向过程对真实图像逐步添加高斯噪声经过足够多步后变成纯噪声。反向过程训练一个神经网络学习从噪声逐步还原出原始图像。语言模型不能直接套用这套流程原因是文本是离散的 token 序列比如我喜欢编程对应[我, 喜欢, 编程]这些符号没有天然的连续数值含义不能直接加噪声。扩散式语言模型要解决的核心问题就是如何把离散文本转换成适合扩散过程的连续表示再把连续表示映射回离散词表。主流方案是引入一个可学习的嵌入层把每个 token 映射成连续向量。训练时对连续向量做加噪和去噪推理时通过最近邻查找或 softmax 映射把连续向量还原成词表中的 token。也有研究直接对离散状态定义转移概率比如掩码扩散模型把“加噪”理解为逐步把某些 token 替换成[MASK]。1.3 适用场景与优势扩散式语言模型适合以下几类任务无条件文本生成从随机噪声生成自然语句。文本修复与填空给定部分上下文生成缺失内容。受控文本生成通过引导向量或条件输入控制生成文本的情感、主题、长度。风格迁移与改写在连续空间内对文本向量进行插值或编辑改变属性同时保留语义。它的主要优势可以总结为三点非自回归生成理论上可以并行解码生成速度有提升空间。在连续空间内做编辑方便引入属性控制和语义插值。训练目标可以结合条件信息适合多模态和条件生成任务。当然它也有明显的短板训练开销大、推理步数多、离散映射会损失信息。下面我们会逐步展开说明。2. 扩散式语言模型的核心原理2.1 前向加噪过程在前向过程中假设原始文本为长度为 L 的 token 序列经过嵌入层得到向量序列x0 [e(x1), e(x2), ..., e(xL)]。我们按照扩散模型的公式逐步添加噪声q(xt | x_{t-1}) N(xt; sqrt(1 - beta_t) * x_{t-1}, beta_t * I)其中beta_t是第t步的噪声调度系数通常从小到大递增表示噪声逐渐加强。经过足够多步后xT接近标准高斯分布。这里有一个文本领域特有的问题直接在嵌入向量上加高斯噪声得到的中间表示并不对应任何真实 token。为了让训练更稳定通常会对嵌入向量做归一化或者采用absorbing state吸收态方式把加噪定义为逐步把 token 替换为[MASK]。两种方式各有优势高斯噪声方式连续性好方便做插值、编辑但离散映射存在误差。吸收态方式更贴近文本的离散本质训练更容易但连续编辑能力弱一些。2.2 反向去噪过程反向过程的目标是训练一个神经网络theta预测每一步添加的噪声或者直接预测原始文本。具体来说我们希望模型学习p_theta(x_{t-1} | xt) 或 p_theta(x0 | xt)在实现时通常会用一个基于 Transformer 的 denoiser去噪网络来完成。输入是加噪后的向量序列xt、时间步t、以及可选的条件信息c输出是预测的噪声或者原始 token 分布。训练目标通常是最小化真实噪声与预测噪声之间的均方误差MSE也可以等价地最小化预测原始 token 的交叉熵损失。后者在实践中有时候更容易收敛因为文本的最终目标仍然是离散 token 的分类问题。2.3 采样与解码训练完成后推理过程从标准高斯噪声xT开始用训练好的去噪网络逐步去噪x_{t-1} denoise(x_t, t)经过 T 步迭代后得到接近原始数据分布的连续向量x0。最后一步我们需要把连续向量映射到词表token_i argmax_{v in vocab} sim(x0_i, embed(v))这里的sim可以是余弦相似度或内积。如果训练时使用了额外的映射头projection head也可以通过一个线性层直接输出词表分布再取 argmax 得到 token。3. 环境准备与项目结构3.1 运行环境本文演示的代码以 PyTorch 为基础实现需要安装以下依赖pip install torch einops numpy transformers datasets tensorboard版本说明PyTorch 2.x 均可Transformers 库主要用于加载 BERT 分词器或预训练模型如果网络环境不方便下载模型也可以直接使用简单的字符级分词器完整的示例代码不依赖预训练模型。建议在一张显存不低于 12GB 的 GPU 上运行。如果只有 CPU可以先把序列长度调短、词汇表调小做功能验证。3.2 项目结构为了方便演示我们把项目按下面的目录组织diffusion-lm/ ├── config.py # 配置文件 ├── data_utils.py # 数据加载与分词 ├── diffusion.py # 扩散过程的前向加噪与采样 ├── model.py # 去噪 Transformer 模型 ├── train.py # 训练脚本 ├── generate.py # 推理采样脚本 └── checkpoints/ # 模型保存目录4. 核心代码实现下面我们逐步实现一个简化版扩散式语言模型。这个版本不追求完整复现论文效果重点是让读者看清扩散语言模型的整个数据流离散 token → 嵌入向量 → 加噪 → 去噪 → 映射回 token。4.1 配置模块文件路径config.pyclass Config: # 数据 vocab_size 256 # 字符级词表大小 max_length 32 # 序列最大长度 # 模型 hidden_size 256 num_layers 4 num_heads 4 intermediate_size 512 # 扩散过程 num_timesteps 100 # 扩散步数 beta_start 0.0001 beta_end 0.02 # 训练 batch_size 32 learning_rate 2e-4 num_epochs 50 log_interval 10 # 路径 checkpoint_dir checkpoints这里使用字符级词表vocab_size 256对应 UTF-8 字节或简单字符映射。实际项目中如果要处理中文或英文单词级建模需要换成成熟的 tokenizer。4.2 数据加载与分词文件路径data_utils.pyimport torch from torch.utils.data import Dataset, DataLoader class CharTokenizer: 简单字符级 tokenizer用于演示。 def __init__(self, texts): self.chars sorted(list(set(.join(texts)))) self.stoi {ch: i for i, ch in enumerate(self.chars)} self.itos {i: ch for i, ch in enumerate(self.chars)} self.vocab_size len(self.chars) def encode(self, text): return [self.stoi[ch] for ch in text] def decode(self, ids): return .join([self.itos[i] for i in ids]) class TextDataset(Dataset): def __init__(self, texts, tokenizer, max_length): self.tokenizer tokenizer self.max_length max_length self.data [] for text in texts: ids self.tokenizer.encode(text) if len(ids) max_length: ids ids[:max_length] else: ids ids [0] * (max_length - len(ids)) # 0 作为 padding self.data.append(torch.tensor(ids, dtypetorch.long)) def __len__(self): return len(self.data) def __getitem__(self, idx): return self.data[idx]如果你有现成的文本文件可以按行读取后构造数据集。为了演示这里使用字符级分词器不需要额外下载模型文件。4.3 扩散过程定义文件路径diffusion.py这里定义前向加噪和反向采样所需的基础函数。注意力放在两个核心方法上q_sample用于在训练时快速采样任意时间步的加噪结果p_sample用于推理时逐步去噪。import torch import torch.nn.functional as F def linear_beta_schedule(timesteps, beta_start0.0001, beta_end0.02): return torch.linspace(beta_start, beta_end, timesteps) class DiffusionProcess: def __init__(self, num_timesteps, beta_start, beta_end): self.num_timesteps num_timesteps self.betas linear_beta_schedule(num_timesteps, beta_start, beta_end) self.alphas 1.0 - self.betas self.alpha_cumprod torch.cumprod(self.alphas, dim0) def q_sample(self, x0, t, noise): 前向加噪从原始嵌入向量 x0 出发采样第 t 步的加噪结果。 x0: [batch, length, hidden] t: [batch] noise: 与 x0 同形状的高斯噪声 sqrt_alpha_cumprod torch.sqrt(self.alpha_cumprod[t]).view(-1, 1, 1) sqrt_one_minus_alpha_cumprod torch.sqrt(1 - self.alpha_cumprod[t]).view(-1, 1, 1) return sqrt_alpha_cumprod * x0 sqrt_one_minus_alpha_cumprod * noise def p_sample(self, model, x_t, t, condNone): 反向去噪给定第 t 步向量预测第 t-1 步向量。 这里使用简化采样实际项目中可以用 DDIM 等方式加速。 with torch.no_grad(): pred_noise model(x_t, t, cond) beta_t self.betas[t].view(-1, 1, 1) sqrt_alpha_t torch.sqrt(self.alphas[t]).view(-1, 1, 1) sqrt_one_minus_alpha_cumprod torch.sqrt(1 - self.alpha_cumprod[t]).view(-1, 1, 1) x_prev (x_t - beta_t / sqrt_one_minus_alpha_cumprod * pred_noise) / sqrt_alpha_t if t.min() 0: noise torch.randn_like(x_t) x_prev x_prev torch.sqrt(beta_t) * noise return x_prevq_sample在训练中非常高效它不需要逐步迭代就能直接算出任意时间步的加噪结果。这也是扩散模型训练的一个核心优势。4.4 去噪网络文件路径model.py去噪网络是整个模型的骨干。我们用 Transformer Encoder 作为主干输入是加噪后的向量序列同时把时间步t编码成向量并融入输入。import torch import torch.nn as nn import math class TimeEmbedding(nn.Module): 时间步编码类似 Transformer 中的位置编码。 def __init__(self, hidden_size): super().__init__() self.hidden_size hidden_size def forward(self, t): device t.device half_dim self.hidden_size // 2 emb math.log(10000) / (half_dim - 1) emb torch.exp(torch.arange(half_dim, devicedevice) * -emb) emb t[:, None] * emb[None, :] emb torch.cat([torch.sin(emb), torch.cos(emb)], dim-1) return emb class TransformerDenoiser(nn.Module): def __init__(self, config): super().__init__() self.config config self.time_embedding TimeEmbedding(config.hidden_size) self.input_proj nn.Linear(config.hidden_size, config.hidden_size) self.pos_embedding nn.Parameter(torch.randn(1, config.max_length, config.hidden_size)) encoder_layer nn.TransformerEncoderLayer( d_modelconfig.hidden_size, nheadconfig.num_heads, dim_feedforwardconfig.intermediate_size, batch_firstTrue, ) self.encoder nn.TransformerEncoder(encoder_layer, num_layersconfig.num_layers) self.output_proj nn.Linear(config.hidden_size, config.hidden_size) def forward(self, x_t, t, condNone): # x_t: [batch, length, hidden] time_emb self.time_embedding(t) # [batch, hidden] time_emb time_emb.unsqueeze(1) # [batch, 1, hidden] h self.input_proj(x_t) h h self.pos_embedding h h time_emb if cond is not None: cond cond.unsqueeze(1) # [batch, 1, cond_dim] h h cond h self.encoder(h) return self.output_proj(h)这里的cond是可选的条件向量。如果你需要做情感控制或主题控制可以把条件编码后加进来。4.5 嵌入层与离散映射除了去噪网络我们还需要两个关键的映射组件embedding把离散 token 映射成连续向量。projection把去噪后的连续向量映射回词表分布。文件路径model.pyclass DiffusionLM(nn.Module): def __init__(self, config): super().__init__() self.config config self.token_embedding nn.Embedding(config.vocab_size, config.hidden_size) self.denoiser TransformerDenoiser(config) self.decode_proj nn.Linear(config.hidden_size, config.vocab_size) def forward(self, x0, t, condNone): x0: [batch, length] 离散 token 序列 t: [batch] 时间步 x0_emb self.token_embedding(x0) # 离散 token → 连续向量 noise torch.randn_like(x0_emb) diffusion DiffusionProcess( self.config.num_timesteps, self.config.beta_start, self.config.beta_end, ) xt diffusion.q_sample(x0_emb, t, noise) pred_noise self.denoiser(xt, t, cond) logits self.decode_proj(pred_noise) # 连续向量 → 词表 logits return logits, noise, xt, x0_emb训练时我们有两种等价的目标可以选择预测噪声使用 MSE 损失。预测原始 token使用交叉熵损失。下面要介绍的训练脚本采用第二种方式因为它在实践中的收敛速度通常更快也更容易观察到 loss 下降。4.6 训练脚本文件路径train.pyimport torch import torch.nn.functional as F from torch.utils.data import DataLoader from config import Config from data_utils import CharTokenizer, TextDataset from model import DiffusionLM # 1. 准备示例数据 sample_texts [ hello world, diffusion language model, text generation with diffusion, deep learning is powerful, we love natural language processing, transformers are everywhere, this is a simple example, training ai models requires data, ] config Config() tokenizer CharTokenizer(sample_texts) config.vocab_size tokenizer.vocab_size dataset TextDataset(sample_texts, tokenizer, config.max_length) dataloader DataLoader(dataset, batch_sizeconfig.batch_size, shuffleTrue) # 2. 初始化模型 device torch.device(cuda if torch.cuda.is_available() else cpu) model DiffusionLM(config).to(device) optimizer torch.optim.AdamW(model.parameters(), lrconfig.learning_rate) # 3. 扩散过程实例 diffusion DiffusionProcess( config.num_timesteps, config.beta_start, config.beta_end, ) # 4. 训练循环 model.train() steps 0 for epoch in range(config.num_epochs): total_loss 0.0 for batch in dataloader: batch batch.to(device) batch_size batch.size(0) # 随机采样时间步 t torch.randint(0, config.num_timesteps, (batch_size,), devicedevice) logits, noise, xt, x0_emb model(batch, t) # 交叉熵损失预测原始 token loss F.cross_entropy( logits.view(-1, config.vocab_size), batch.view(-1), ignore_index0, ) optimizer.zero_grad() loss.backward() torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) optimizer.step() total_loss loss.item() steps 1 if steps % config.log_interval 0: print(fEpoch {epoch}, Step {steps}, Loss {loss.item():.4f}) avg_loss total_loss / len(dataloader) print(fEpoch {epoch} finished, Average Loss: {avg_loss:.4f}) if epoch % 10 0: torch.save(model.state_dict(), f{config.checkpoint_dir}/diffusion_lm_{epoch}.pt)训练过程中需要关注 loss 是否持续下降。如果 loss 不降可以先减小学习率或者增加训练数据量。这个简化版本在小数据集上几分钟就能完成训练目的只是验证流程正确性。4.7 推理采样脚本文件路径generate.pyimport torch from config import Config from data_utils import CharTokenizer from model import DiffusionLM from diffusion import DiffusionProcess def generate(model, tokenizer, config, condNone, num_stepsNone): model.eval() device next(model.parameters()).device batch_size 1 length config.max_length diffusion DiffusionProcess( config.num_timesteps, config.beta_start, config.beta_end, ) if num_steps is None: num_steps config.num_timesteps # 从纯高斯噪声开始 x_t torch.randn(batch_size, length, config.hidden_size, devicedevice) # 逐步去噪 for i in reversed(range(num_steps)): t torch.full((batch_size,), i, devicedevice, dtypetorch.long) x_t diffusion.p_sample(model.denoiser, x_t, t, cond) # 映射到词表 logits model.decode_proj(x_t) # [batch, length, vocab] pred_ids torch.argmax(logits, dim-1) # [batch, length] texts [] for ids in pred_ids: texts.append(tokenizer.decode(ids.tolist())) return texts if __name__ __main__: # 需要保证配置与训练时一致 config Config() tokenizer None # 这里需要你在重新加载时恢复 tokenizer # 实际运行时建议将 tokenizer 保存到本地文件避免重新训练 model DiffusionLM(config) model.load_state_dict(torch.load(checkpoints/diffusion_lm_40.pt)) model.eval() texts generate(model, tokenizer, config) for text in texts: print(text)注意示例代码中tokenizer需要单独保存和恢复。建议在训练后把tokenizer.stoi和tokenizer.itos用json或pickle存盘再在推理脚本中加载。5. 常见问题与排查思路在实际训练和调试扩散式语言模型的过程中下面几个问题出现的频率最高。问题现象常见原因解决思路Loss 不下降或震荡学习率过高数据量太小时间步采样不均匀调低学习率、增加数据量、检查时间步分布生成文本全是重复字符模型未充分收敛字符级建模表达能力弱增大训练轮数、换用 BPE 或 WordPiece 分词采样结果与训练分布差异大训练和推理时的去噪步数不一致映射头训练不充分让推理步数与训练步数一致或微调 decode_projGPU 内存不足序列过长batch 过大Transformer 层数过多减小 max_length、降低 batch_size、使用梯度累积生成结果没有语义训练数据太少模型容量不够使用真实语料扩大 hidden_size 和层数从高斯噪声采样时输出全为填充符模型只学会了 padding token 的分布增加有效数据长度、避免过长 Padding、调整损失权重其中训练和推理步数不一致是比较隐蔽的问题。如果训练时使用了 1000 步推理时为了加速只做 50 步模型在很少步数内无法完成从纯噪声到文本的转换。建议先用相同步数验证效果再考虑用 DDIM 等加速采样方法。6. 最佳实践与工程建议6.1 从玩具数据集开始验证扩散式语言模型的训练成本远高于自回归模型。第一次上手时不要直接用大规模语料而是先准备几十条短文本用小模型、小词表、短序列跑通整个流程确认加噪、去噪、采样、离散映射都没有问题后再逐步放大。6.2 选择合适的离散化方案扩散语言模型的一个核心难点是“离散到连续”和“连续到离散”的转换。常用的方案有三种字符级 tokenizer实现简单适合小规模验证但建模能力弱。BPE / WordPiece适合真实语言建模是目前主流选择。预训练嵌入初始化可以使用 BERT 或 GPT 的 embedding 初始化嵌入层加快收敛。使用预训练嵌入初始化时需要保证词表与预训练模型一致并且做适当的归一化。6.3 采样加速标准 DDPM 采样需要几百甚至上千步推理成本很高。实际项目中可以尝试DDIM 采样通过调整采样步数把 1000 步压缩到 50 步左右。一致性模型训练目标改为一步或多步自一致性进一步加速。先自回归初稿、再扩散精修混合架构在生产中更实用。6.4 条件信息融合在做受控生成时条件信息可以加在三个位置拼接进输入向量。作为 cross-attention 的 key/value。通过 AdaIN 或 FiLM 方式调制特征。简单场景下直接把条件向量加到输入上最容易实现但复杂属性控制建议使用 cross-attention。6.5 算力与资源规划扩散语言模型的训练通常比同规模自回归模型更吃显存因为需要保存噪声向量、中间步骤的梯度并且多步去噪也会增加计算量。建议使用梯度累积增大有效 batch size。使用混合精度训练AMP降低显存占用。尽量固定序列长度避免动态长度导致显存波动。先在小规模数据上完成超参实验再申请大规模训练资源。6.6 数据合规与生成安全训练使用的语料必须有合法授权不爬取未经授权的数据。生成式模型存在输出不可控的问题在面向生产环境时需要增加内容过滤、敏感词检测和人工审核机制。涉及用户数据时还要遵守数据最小化原则不采集与任务无关的隐私信息。7. 总结与下一步学习方向这篇文章从离散文本与连续扩散过程之间的矛盾出发介绍了扩散式语言模型的核心概念、训练目标和采样方式并给出了一个完整的 PyTorch 实现框架。通过这个框架你可以掌握以下关键点扩散模型如何通过前向加噪与反向去噪完成文本生成。离散 token 如何通过嵌入层映射为连续向量再通过 decode projection 映射回词表。训练时如何高效采样任意时间步的加噪结果。推理时如何从纯噪声逐步还原文本。常见训练问题的排查思路和工程优化方向。下一步如果继续深入可以优先关注这几个方向阅读 Diffusion-LM、DiffuSeq、SSD-LM 等经典论文理解吸收态扩散和高斯扩散的差异。在真实中文或英文语料上跑通 BPE 级别的扩散语言模型。尝试把扩散模型用于文本属性控制或语义插值任务。研究 DDIM 采样、一致性模型等加速推理的方法。扩散式语言模型目前还处于快速演进阶段离大规模工业应用还有一段距离但它在非自回归生成、全局编辑、可控生成等场景中展现出的潜力值得持续跟进。如果这篇文章对你有帮助可以收藏备用后续我也会继续分享更多关于文本生成与扩散模型的实战细节。 扩散式语言模型：从离散文本到连续去噪的生成新范式
近几年做文本生成相关实践时我越来越常听到一个方向扩散式语言模型。刚开始接触这个概念很多人会下意识觉得它和扩散模型画图是同一套东西直接搬到文本上。实际动手后发现文本是离散符号图像是连续像素想把扩散过程用在语言上关键就在“怎么把离散文本变成可加噪、可去噪的连续表示”。这篇文章就围绕这个核心问题展开先讲清楚扩散式语言模型的基本原理和主流实现思路再给出一个可以跑通小规模实验的代码框架最后整理训练过程中常见的问题和工程建议。1. 扩散式语言模型是什么1.1 从自回归生成说起主流文本生成模型大多采用自回归架构也就是一个 token 一个 token 地预测下一个词。GPT 系列、LLaMA、ChatGLM 等模型的核心思路都是给定前文 x1, x2, ..., x_{t-1}预测 P(xt | x1, x2, ..., x_{t-1})这种方式的优点是生成质量高、训练稳定推理时也容易控制。但它的缺点同样明显生成速度受限于串行解码无法并行输出整段文本。生成过程是单向的很难直接在全局层面进行修改和重构。对于填空、文本修复、受控改写这类任务自回归模型需要专门设计训练目标天然不占优势。扩散式语言模型Diffusion Language ModelDLM则尝试换一条路线先让模型学习一个“从纯噪声逐步还原成完整文本”的去噪过程。训练完成后推理时从一个随机噪声向量出发通过多步迭代逐渐生成整段文本。这种生成方式天然支持并行解码也为全局编辑、属性控制等任务提供了更灵活的空间。1.2 扩散式语言模型的核心思想扩散模型在图像生成中的基本流程是前向过程对真实图像逐步添加高斯噪声经过足够多步后变成纯噪声。反向过程训练一个神经网络学习从噪声逐步还原出原始图像。语言模型不能直接套用这套流程原因是文本是离散的 token 序列比如我喜欢编程对应[我, 喜欢, 编程]这些符号没有天然的连续数值含义不能直接加噪声。扩散式语言模型要解决的核心问题就是如何把离散文本转换成适合扩散过程的连续表示再把连续表示映射回离散词表。主流方案是引入一个可学习的嵌入层把每个 token 映射成连续向量。训练时对连续向量做加噪和去噪推理时通过最近邻查找或 softmax 映射把连续向量还原成词表中的 token。也有研究直接对离散状态定义转移概率比如掩码扩散模型把“加噪”理解为逐步把某些 token 替换成[MASK]。1.3 适用场景与优势扩散式语言模型适合以下几类任务无条件文本生成从随机噪声生成自然语句。文本修复与填空给定部分上下文生成缺失内容。受控文本生成通过引导向量或条件输入控制生成文本的情感、主题、长度。风格迁移与改写在连续空间内对文本向量进行插值或编辑改变属性同时保留语义。它的主要优势可以总结为三点非自回归生成理论上可以并行解码生成速度有提升空间。在连续空间内做编辑方便引入属性控制和语义插值。训练目标可以结合条件信息适合多模态和条件生成任务。当然它也有明显的短板训练开销大、推理步数多、离散映射会损失信息。下面我们会逐步展开说明。2. 扩散式语言模型的核心原理2.1 前向加噪过程在前向过程中假设原始文本为长度为 L 的 token 序列经过嵌入层得到向量序列x0 [e(x1), e(x2), ..., e(xL)]。我们按照扩散模型的公式逐步添加噪声q(xt | x_{t-1}) N(xt; sqrt(1 - beta_t) * x_{t-1}, beta_t * I)其中beta_t是第t步的噪声调度系数通常从小到大递增表示噪声逐渐加强。经过足够多步后xT接近标准高斯分布。这里有一个文本领域特有的问题直接在嵌入向量上加高斯噪声得到的中间表示并不对应任何真实 token。为了让训练更稳定通常会对嵌入向量做归一化或者采用absorbing state吸收态方式把加噪定义为逐步把 token 替换为[MASK]。两种方式各有优势高斯噪声方式连续性好方便做插值、编辑但离散映射存在误差。吸收态方式更贴近文本的离散本质训练更容易但连续编辑能力弱一些。2.2 反向去噪过程反向过程的目标是训练一个神经网络theta预测每一步添加的噪声或者直接预测原始文本。具体来说我们希望模型学习p_theta(x_{t-1} | xt) 或 p_theta(x0 | xt)在实现时通常会用一个基于 Transformer 的 denoiser去噪网络来完成。输入是加噪后的向量序列xt、时间步t、以及可选的条件信息c输出是预测的噪声或者原始 token 分布。训练目标通常是最小化真实噪声与预测噪声之间的均方误差MSE也可以等价地最小化预测原始 token 的交叉熵损失。后者在实践中有时候更容易收敛因为文本的最终目标仍然是离散 token 的分类问题。2.3 采样与解码训练完成后推理过程从标准高斯噪声xT开始用训练好的去噪网络逐步去噪x_{t-1} denoise(x_t, t)经过 T 步迭代后得到接近原始数据分布的连续向量x0。最后一步我们需要把连续向量映射到词表token_i argmax_{v in vocab} sim(x0_i, embed(v))这里的sim可以是余弦相似度或内积。如果训练时使用了额外的映射头projection head也可以通过一个线性层直接输出词表分布再取 argmax 得到 token。3. 环境准备与项目结构3.1 运行环境本文演示的代码以 PyTorch 为基础实现需要安装以下依赖pip install torch einops numpy transformers datasets tensorboard版本说明PyTorch 2.x 均可Transformers 库主要用于加载 BERT 分词器或预训练模型如果网络环境不方便下载模型也可以直接使用简单的字符级分词器完整的示例代码不依赖预训练模型。建议在一张显存不低于 12GB 的 GPU 上运行。如果只有 CPU可以先把序列长度调短、词汇表调小做功能验证。3.2 项目结构为了方便演示我们把项目按下面的目录组织diffusion-lm/ ├── config.py # 配置文件 ├── data_utils.py # 数据加载与分词 ├── diffusion.py # 扩散过程的前向加噪与采样 ├── model.py # 去噪 Transformer 模型 ├── train.py # 训练脚本 ├── generate.py # 推理采样脚本 └── checkpoints/ # 模型保存目录4. 核心代码实现下面我们逐步实现一个简化版扩散式语言模型。这个版本不追求完整复现论文效果重点是让读者看清扩散语言模型的整个数据流离散 token → 嵌入向量 → 加噪 → 去噪 → 映射回 token。4.1 配置模块文件路径config.pyclass Config: # 数据 vocab_size 256 # 字符级词表大小 max_length 32 # 序列最大长度 # 模型 hidden_size 256 num_layers 4 num_heads 4 intermediate_size 512 # 扩散过程 num_timesteps 100 # 扩散步数 beta_start 0.0001 beta_end 0.02 # 训练 batch_size 32 learning_rate 2e-4 num_epochs 50 log_interval 10 # 路径 checkpoint_dir checkpoints这里使用字符级词表vocab_size 256对应 UTF-8 字节或简单字符映射。实际项目中如果要处理中文或英文单词级建模需要换成成熟的 tokenizer。4.2 数据加载与分词文件路径data_utils.pyimport torch from torch.utils.data import Dataset, DataLoader class CharTokenizer: 简单字符级 tokenizer用于演示。 def __init__(self, texts): self.chars sorted(list(set(.join(texts)))) self.stoi {ch: i for i, ch in enumerate(self.chars)} self.itos {i: ch for i, ch in enumerate(self.chars)} self.vocab_size len(self.chars) def encode(self, text): return [self.stoi[ch] for ch in text] def decode(self, ids): return .join([self.itos[i] for i in ids]) class TextDataset(Dataset): def __init__(self, texts, tokenizer, max_length): self.tokenizer tokenizer self.max_length max_length self.data [] for text in texts: ids self.tokenizer.encode(text) if len(ids) max_length: ids ids[:max_length] else: ids ids [0] * (max_length - len(ids)) # 0 作为 padding self.data.append(torch.tensor(ids, dtypetorch.long)) def __len__(self): return len(self.data) def __getitem__(self, idx): return self.data[idx]如果你有现成的文本文件可以按行读取后构造数据集。为了演示这里使用字符级分词器不需要额外下载模型文件。4.3 扩散过程定义文件路径diffusion.py这里定义前向加噪和反向采样所需的基础函数。注意力放在两个核心方法上q_sample用于在训练时快速采样任意时间步的加噪结果p_sample用于推理时逐步去噪。import torch import torch.nn.functional as F def linear_beta_schedule(timesteps, beta_start0.0001, beta_end0.02): return torch.linspace(beta_start, beta_end, timesteps) class DiffusionProcess: def __init__(self, num_timesteps, beta_start, beta_end): self.num_timesteps num_timesteps self.betas linear_beta_schedule(num_timesteps, beta_start, beta_end) self.alphas 1.0 - self.betas self.alpha_cumprod torch.cumprod(self.alphas, dim0) def q_sample(self, x0, t, noise): 前向加噪从原始嵌入向量 x0 出发采样第 t 步的加噪结果。 x0: [batch, length, hidden] t: [batch] noise: 与 x0 同形状的高斯噪声 sqrt_alpha_cumprod torch.sqrt(self.alpha_cumprod[t]).view(-1, 1, 1) sqrt_one_minus_alpha_cumprod torch.sqrt(1 - self.alpha_cumprod[t]).view(-1, 1, 1) return sqrt_alpha_cumprod * x0 sqrt_one_minus_alpha_cumprod * noise def p_sample(self, model, x_t, t, condNone): 反向去噪给定第 t 步向量预测第 t-1 步向量。 这里使用简化采样实际项目中可以用 DDIM 等方式加速。 with torch.no_grad(): pred_noise model(x_t, t, cond) beta_t self.betas[t].view(-1, 1, 1) sqrt_alpha_t torch.sqrt(self.alphas[t]).view(-1, 1, 1) sqrt_one_minus_alpha_cumprod torch.sqrt(1 - self.alpha_cumprod[t]).view(-1, 1, 1) x_prev (x_t - beta_t / sqrt_one_minus_alpha_cumprod * pred_noise) / sqrt_alpha_t if t.min() 0: noise torch.randn_like(x_t) x_prev x_prev torch.sqrt(beta_t) * noise return x_prevq_sample在训练中非常高效它不需要逐步迭代就能直接算出任意时间步的加噪结果。这也是扩散模型训练的一个核心优势。4.4 去噪网络文件路径model.py去噪网络是整个模型的骨干。我们用 Transformer Encoder 作为主干输入是加噪后的向量序列同时把时间步t编码成向量并融入输入。import torch import torch.nn as nn import math class TimeEmbedding(nn.Module): 时间步编码类似 Transformer 中的位置编码。 def __init__(self, hidden_size): super().__init__() self.hidden_size hidden_size def forward(self, t): device t.device half_dim self.hidden_size // 2 emb math.log(10000) / (half_dim - 1) emb torch.exp(torch.arange(half_dim, devicedevice) * -emb) emb t[:, None] * emb[None, :] emb torch.cat([torch.sin(emb), torch.cos(emb)], dim-1) return emb class TransformerDenoiser(nn.Module): def __init__(self, config): super().__init__() self.config config self.time_embedding TimeEmbedding(config.hidden_size) self.input_proj nn.Linear(config.hidden_size, config.hidden_size) self.pos_embedding nn.Parameter(torch.randn(1, config.max_length, config.hidden_size)) encoder_layer nn.TransformerEncoderLayer( d_modelconfig.hidden_size, nheadconfig.num_heads, dim_feedforwardconfig.intermediate_size, batch_firstTrue, ) self.encoder nn.TransformerEncoder(encoder_layer, num_layersconfig.num_layers) self.output_proj nn.Linear(config.hidden_size, config.hidden_size) def forward(self, x_t, t, condNone): # x_t: [batch, length, hidden] time_emb self.time_embedding(t) # [batch, hidden] time_emb time_emb.unsqueeze(1) # [batch, 1, hidden] h self.input_proj(x_t) h h self.pos_embedding h h time_emb if cond is not None: cond cond.unsqueeze(1) # [batch, 1, cond_dim] h h cond h self.encoder(h) return self.output_proj(h)这里的cond是可选的条件向量。如果你需要做情感控制或主题控制可以把条件编码后加进来。4.5 嵌入层与离散映射除了去噪网络我们还需要两个关键的映射组件embedding把离散 token 映射成连续向量。projection把去噪后的连续向量映射回词表分布。文件路径model.pyclass DiffusionLM(nn.Module): def __init__(self, config): super().__init__() self.config config self.token_embedding nn.Embedding(config.vocab_size, config.hidden_size) self.denoiser TransformerDenoiser(config) self.decode_proj nn.Linear(config.hidden_size, config.vocab_size) def forward(self, x0, t, condNone): x0: [batch, length] 离散 token 序列 t: [batch] 时间步 x0_emb self.token_embedding(x0) # 离散 token → 连续向量 noise torch.randn_like(x0_emb) diffusion DiffusionProcess( self.config.num_timesteps, self.config.beta_start, self.config.beta_end, ) xt diffusion.q_sample(x0_emb, t, noise) pred_noise self.denoiser(xt, t, cond) logits self.decode_proj(pred_noise) # 连续向量 → 词表 logits return logits, noise, xt, x0_emb训练时我们有两种等价的目标可以选择预测噪声使用 MSE 损失。预测原始 token使用交叉熵损失。下面要介绍的训练脚本采用第二种方式因为它在实践中的收敛速度通常更快也更容易观察到 loss 下降。4.6 训练脚本文件路径train.pyimport torch import torch.nn.functional as F from torch.utils.data import DataLoader from config import Config from data_utils import CharTokenizer, TextDataset from model import DiffusionLM # 1. 准备示例数据 sample_texts [ hello world, diffusion language model, text generation with diffusion, deep learning is powerful, we love natural language processing, transformers are everywhere, this is a simple example, training ai models requires data, ] config Config() tokenizer CharTokenizer(sample_texts) config.vocab_size tokenizer.vocab_size dataset TextDataset(sample_texts, tokenizer, config.max_length) dataloader DataLoader(dataset, batch_sizeconfig.batch_size, shuffleTrue) # 2. 初始化模型 device torch.device(cuda if torch.cuda.is_available() else cpu) model DiffusionLM(config).to(device) optimizer torch.optim.AdamW(model.parameters(), lrconfig.learning_rate) # 3. 扩散过程实例 diffusion DiffusionProcess( config.num_timesteps, config.beta_start, config.beta_end, ) # 4. 训练循环 model.train() steps 0 for epoch in range(config.num_epochs): total_loss 0.0 for batch in dataloader: batch batch.to(device) batch_size batch.size(0) # 随机采样时间步 t torch.randint(0, config.num_timesteps, (batch_size,), devicedevice) logits, noise, xt, x0_emb model(batch, t) # 交叉熵损失预测原始 token loss F.cross_entropy( logits.view(-1, config.vocab_size), batch.view(-1), ignore_index0, ) optimizer.zero_grad() loss.backward() torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) optimizer.step() total_loss loss.item() steps 1 if steps % config.log_interval 0: print(fEpoch {epoch}, Step {steps}, Loss {loss.item():.4f}) avg_loss total_loss / len(dataloader) print(fEpoch {epoch} finished, Average Loss: {avg_loss:.4f}) if epoch % 10 0: torch.save(model.state_dict(), f{config.checkpoint_dir}/diffusion_lm_{epoch}.pt)训练过程中需要关注 loss 是否持续下降。如果 loss 不降可以先减小学习率或者增加训练数据量。这个简化版本在小数据集上几分钟就能完成训练目的只是验证流程正确性。4.7 推理采样脚本文件路径generate.pyimport torch from config import Config from data_utils import CharTokenizer from model import DiffusionLM from diffusion import DiffusionProcess def generate(model, tokenizer, config, condNone, num_stepsNone): model.eval() device next(model.parameters()).device batch_size 1 length config.max_length diffusion DiffusionProcess( config.num_timesteps, config.beta_start, config.beta_end, ) if num_steps is None: num_steps config.num_timesteps # 从纯高斯噪声开始 x_t torch.randn(batch_size, length, config.hidden_size, devicedevice) # 逐步去噪 for i in reversed(range(num_steps)): t torch.full((batch_size,), i, devicedevice, dtypetorch.long) x_t diffusion.p_sample(model.denoiser, x_t, t, cond) # 映射到词表 logits model.decode_proj(x_t) # [batch, length, vocab] pred_ids torch.argmax(logits, dim-1) # [batch, length] texts [] for ids in pred_ids: texts.append(tokenizer.decode(ids.tolist())) return texts if __name__ __main__: # 需要保证配置与训练时一致 config Config() tokenizer None # 这里需要你在重新加载时恢复 tokenizer # 实际运行时建议将 tokenizer 保存到本地文件避免重新训练 model DiffusionLM(config) model.load_state_dict(torch.load(checkpoints/diffusion_lm_40.pt)) model.eval() texts generate(model, tokenizer, config) for text in texts: print(text)注意示例代码中tokenizer需要单独保存和恢复。建议在训练后把tokenizer.stoi和tokenizer.itos用json或pickle存盘再在推理脚本中加载。5. 常见问题与排查思路在实际训练和调试扩散式语言模型的过程中下面几个问题出现的频率最高。问题现象常见原因解决思路Loss 不下降或震荡学习率过高数据量太小时间步采样不均匀调低学习率、增加数据量、检查时间步分布生成文本全是重复字符模型未充分收敛字符级建模表达能力弱增大训练轮数、换用 BPE 或 WordPiece 分词采样结果与训练分布差异大训练和推理时的去噪步数不一致映射头训练不充分让推理步数与训练步数一致或微调 decode_projGPU 内存不足序列过长batch 过大Transformer 层数过多减小 max_length、降低 batch_size、使用梯度累积生成结果没有语义训练数据太少模型容量不够使用真实语料扩大 hidden_size 和层数从高斯噪声采样时输出全为填充符模型只学会了 padding token 的分布增加有效数据长度、避免过长 Padding、调整损失权重其中训练和推理步数不一致是比较隐蔽的问题。如果训练时使用了 1000 步推理时为了加速只做 50 步模型在很少步数内无法完成从纯噪声到文本的转换。建议先用相同步数验证效果再考虑用 DDIM 等加速采样方法。6. 最佳实践与工程建议6.1 从玩具数据集开始验证扩散式语言模型的训练成本远高于自回归模型。第一次上手时不要直接用大规模语料而是先准备几十条短文本用小模型、小词表、短序列跑通整个流程确认加噪、去噪、采样、离散映射都没有问题后再逐步放大。6.2 选择合适的离散化方案扩散语言模型的一个核心难点是“离散到连续”和“连续到离散”的转换。常用的方案有三种字符级 tokenizer实现简单适合小规模验证但建模能力弱。BPE / WordPiece适合真实语言建模是目前主流选择。预训练嵌入初始化可以使用 BERT 或 GPT 的 embedding 初始化嵌入层加快收敛。使用预训练嵌入初始化时需要保证词表与预训练模型一致并且做适当的归一化。6.3 采样加速标准 DDPM 采样需要几百甚至上千步推理成本很高。实际项目中可以尝试DDIM 采样通过调整采样步数把 1000 步压缩到 50 步左右。一致性模型训练目标改为一步或多步自一致性进一步加速。先自回归初稿、再扩散精修混合架构在生产中更实用。6.4 条件信息融合在做受控生成时条件信息可以加在三个位置拼接进输入向量。作为 cross-attention 的 key/value。通过 AdaIN 或 FiLM 方式调制特征。简单场景下直接把条件向量加到输入上最容易实现但复杂属性控制建议使用 cross-attention。6.5 算力与资源规划扩散语言模型的训练通常比同规模自回归模型更吃显存因为需要保存噪声向量、中间步骤的梯度并且多步去噪也会增加计算量。建议使用梯度累积增大有效 batch size。使用混合精度训练AMP降低显存占用。尽量固定序列长度避免动态长度导致显存波动。先在小规模数据上完成超参实验再申请大规模训练资源。6.6 数据合规与生成安全训练使用的语料必须有合法授权不爬取未经授权的数据。生成式模型存在输出不可控的问题在面向生产环境时需要增加内容过滤、敏感词检测和人工审核机制。涉及用户数据时还要遵守数据最小化原则不采集与任务无关的隐私信息。7. 总结与下一步学习方向这篇文章从离散文本与连续扩散过程之间的矛盾出发介绍了扩散式语言模型的核心概念、训练目标和采样方式并给出了一个完整的 PyTorch 实现框架。通过这个框架你可以掌握以下关键点扩散模型如何通过前向加噪与反向去噪完成文本生成。离散 token 如何通过嵌入层映射为连续向量再通过 decode projection 映射回词表。训练时如何高效采样任意时间步的加噪结果。推理时如何从纯噪声逐步还原文本。常见训练问题的排查思路和工程优化方向。下一步如果继续深入可以优先关注这几个方向阅读 Diffusion-LM、DiffuSeq、SSD-LM 等经典论文理解吸收态扩散和高斯扩散的差异。在真实中文或英文语料上跑通 BPE 级别的扩散语言模型。尝试把扩散模型用于文本属性控制或语义插值任务。研究 DDIM 采样、一致性模型等加速推理的方法。扩散式语言模型目前还处于快速演进阶段离大规模工业应用还有一段距离但它在非自回归生成、全局编辑、可控生成等场景中展现出的潜力值得持续跟进。如果这篇文章对你有帮助可以收藏备用后续我也会继续分享更多关于文本生成与扩散模型的实战细节。