通过训练代码来理解DLLM扩散语言模型

1. 引入Diffusion-LLMDLLM 是扩散语言模型Diffusion Large Language Model它与LLM有什么区别呢1共同点都是以Transformer结构为主的模型2差异点LLM是 “从左到右的自回归Autoregressive生成” 逻辑就是基于已生成的左侧上下文逐一生成下一个令牌直到触发终止符DLLM不是预测下一个token。而是采用迭代去噪的扩散生成生成过程分为前向腐蚀和反向去噪两个阶段核心是 “从全噪声 / 掩码序列中通过多步迭代逐步恢复出有效序列”生成过程是双向、并行、可迭代修正的。听上去有些抽象我们下面从代码层面上来理解DLLM。2. DLLM最简化代码2.1. DLLM简化代码下面是一个DLLM最简单的训练代码包括了简单的数据集、处理、训练、生成过程importtorchimporttorch.nnasnnimporttorch.optimasoptim# 超参数 vocab_size1000embed_dim64seq_len16hidden_dim128timesteps20batch_size4lr1e-3# 简易词汇表与英文句子 # 真实英文句子已截断/补齐到 seq_lensentences[i love deep learning and diffusion models,language models can generate coherent text,diffusion models work by denoising gradually,transformer is the backbone of modern llm,we train a diffusion language model today]# 简易分词按空格 构建词表wordslist({wforsinsentencesforwins.split()})word2idx{w:i1fori,winenumerate(words)}# 0 padidx2word{i:wforw,iinword2idx.items()}vocab_sizelen(word2idx)1# 句子转 token并补齐到 seq_lendeftokenize(s):tokens[word2idx[w]forwins.split()ifwinword2idx]tokenstokens[:seq_len]tokens[0]*(seq_len-len(tokens))returntorch.tensor(tokens)datasettorch.stack([tokenize(s)forsinsentences])print(dataset.shape{0}.format(dataset.shape))# torch.Size([5, 16])# 嵌入层 embeddingnn.Embedding(vocab_size,embed_dim)# Diffusion 去噪模型 classDenoiseTransformer(nn.Module):def__init__(self):super().__init__()self.time_embnn.Embedding(timesteps,embed_dim)self.layernn.TransformerEncoderLayer(d_modelembed_dim,nhead2,dim_feedforwardhidden_dim,batch_firstTrue,activationgelu)self.outnn.Linear(embed_dim,embed_dim)defforward(self,x,t):# x: [B, L, D]# t: [B]t_embself.time_emb(t).unsqueeze(1)# [B,1,D]xxt_emb xself.layer(x)returnself.out(x)modelDenoiseTransformer()optoptim.Adam(model.parameters(),lrlr)# 扩散前向过程 defforward_process(x0,t):noisetorch.randn_like(x0)# 简单线性调度alphatorch.linspace(0.05,0.99,timesteps)[t]alphaalpha.view(-1,1,1)xttorch.sqrt(alpha)*x0torch.sqrt(1-alpha)*noisereturnxt,noise# 训练 print(开始训练 DLLM...\n)forstepinrange(2000):idxtorch.randint(0,len(dataset),(batch_size,))x_idsdataset[idx]x0embedding(x_ids)# 干净的词嵌入 [B,L,D]ttorch.randint(0,timesteps,(batch_size,))xt,noiseforward_process(x0,t)pred_noisemodel(xt,t)loss(pred_noise-noise).pow(2).mean()opt.zero_grad()loss.backward()opt.step()ifstep%2000:print(fstep{step:04d}| loss{loss.item():.4f})# DLLM 反向生成采样print(\n DLLM 生成句子 )xtorch.randn(1,seq_len,embed_dim)# 从纯噪声开始print(x.shape{0}.format(x.shape))# x.shapetorch.Size([1, 16, 64])# 刚开始全是噪声logitsx embedding.weight.T pred_idslogits.argmax(-1).squeeze().tolist()pred_words[idx2word.get(i,pad)foriinpred_idsifi!0]print(初始噪声, .join(pred_words))fort_stepinreversed(range(timesteps)):withtorch.no_grad():pred_nmodel(x,torch.tensor([t_step]))alphatorch.linspace(0.05,0.99,timesteps)[t_step]# 去噪一步x(x-torch.sqrt(1-alpha)*pred_n)/torch.sqrt(alpha)# 映射回单词logitsx embedding.weight.T pred_idslogits.argmax(-1).squeeze().tolist()pred_words[idx2word.get(i,pad)foriinpred_idsifi!0]print(生成句子, .join(pred_words))运行后程序输出dataset.shapetorch.Size([5,16])开始训练 DLLM... step 0000|loss1.3353step 0200|loss0.7625step 0400|loss0.6536step 0600|loss0.5462step 0800|loss0.5811step1000|loss0.4207step1200|loss0.3732step1400|loss0.5663step1600|loss0.3291step1800|loss0.3746DLLM 生成句子x.shapetorch.Size([1,16,64])初始噪声 of is by models we train of love text deep is learning the the today work 生成句子 by language language language language models by by by by可以从中看到DLLM最有特色的地方是针对输入的噪声x一个句子进行去噪后直接生成一个最终的句子。这个过程一开始全是随机噪声x模型一步步去掉噪声最后直接输出一整句话。整个过程是并行生成一步一步去噪不是预测下一个词。而普通LLMGPT、LLaMA输入比如是“I love”输出则是预测下一个词 deep再输入I love deep再输出learning。LLM是逐词生成串行自回归模型。2.2 DLLM特点1前向扩散加噪defforward_process(x0,t):noisetorch.randn_like(x0)xtsqrt(alpha)*x0sqrt(1-alpha)*noisereturnxt,noise2预测噪声pred_noisemodel(xt,t)3噪声损失 MSEloss(pred_noise-noise).pow(2).mean()4反向采样去噪生成xtorch.randn(...)# 从纯噪声开始fortinreversed(...):x(x-sqrt(1-alpha)*pred_n)/sqrt(alpha)3. 区别LLM与DLLM根据上面的解释我们再来理解这几个差异点就容易一些生成内容的方式不同LLM是 “从左到右的自回归Autoregressive生成” 逻辑就是基于已生成的左侧上下文逐一生成下一个令牌直到触发终止符DLLM不是预测下一个token。而是采用迭代去噪的扩散生成生成过程分为前向腐蚀和反向去噪两个阶段核心是 “从全噪声 / 掩码序列中通过多步迭代逐步恢复出有效序列”生成过程是双向、并行、可迭代修正的。训练的目标不同LLM训练目标主要是预测下一个单词损失函数一般为CrossEntropyLossDLLM训练目标是预测加进去的高斯噪声损失函数一般为 MSE预测噪声 - 真实噪声比如上面代码中loss (pred_noise - noise).pow(2).mean()数据处理方式不同LLM直接对 token ID 建模输出单词概率。LLM 不会对连续向量做扩散去噪。DLLM对 词嵌入向量continuous embedding 加噪、去噪。最后才映射回单词。4. 总结DLLM✅ 使用扩散✅ 对句子向量加噪✅ 训练模型预测噪声✅ 从噪声生成整句✅ 不是自回归✅ 不是预测下一个词普通 LLM❌ 无扩散❌ 无加噪❌ 无去噪✅ 逐词生成Diffusion 加噪声 去噪声 的整套流程5. 参考Large Language Diffusion Models. https://arxiv.org/abs/2502.09992https://github.com/Diffusion-LLM/Awesome-DiffusionLLMhttps://zhuanlan.zhihu.com/p/1913691243197752405