扩散语言模型

概述

扩散模型（Diffusion Models）在图像生成领域取得了巨大成功（如 DDPM、Stable Diffusion、DALL-E）。近年来，研究者开始探索将扩散模型应用于语言建模，形成了**离散扩散语言模型（Discrete Diffusion Language Models, dLLM）**这一新兴范式。¹

与传统的自回归模型（Autoregressive Models, ARM）相比，扩散语言模型具有以下潜在优势：

特性	自回归模型	扩散语言模型
生成方式	逐token顺序生成	并行token生成
上下文利用	仅利用前缀上下文	利用完整双向上下文
生成控制	相对受限	更灵活的约束控制
推理速度	受限于序列长度	潜在更快的并行生成

---

离散扩散 vs 连续扩散

连续扩散模型

标准的扩散模型（如 DDPM）在连续空间中操作：

• 前向过程：逐步向数据添加高斯噪声
• 反向过程：学习去噪神经网络 $p_{\theta }(x_{t-1}|x_t)$

$$x_t=\sqrt{1-{\beta }_t}{x}_{t-1}+\sqrt{{\beta }_t}ϵ,ϵ\sim \mathcal{N}(0,I)$$

离散扩散模型

语言数据是离散的（token序列），因此需要特殊处理。离散扩散模型在离散状态空间中操作：

1. 状态空间定义

对于词汇表大小为 $V$ 的语言模型，token $x\in \{1,2,\dots ,V\}$。可以使用one-hot编码：

$$\mathbf{x}\in \{0,1{\}}^V,\text{其中 }{\mathbf{x}}_i=1⟺x=i$$

2. 过渡矩阵

定义从时间步 $t$ 到 $t-1$ 的过渡矩阵 $Q_t$，其中 $(Q_t{)}_{ij}=P(x_t=j|x_{t-1}=i)$。

3. 分类分布的反向过程

给定噪声状态 $x_t$，反向过程需要预测原始分布：

$$p_{\theta }(x_{t-1}|x_t)=\text{Softmax}(f_{\theta }(x_t,t))$$

---

Masked Diffusion Model (MDM)

Masked Diffusion Model（掩码扩散模型）是离散扩散语言模型的主要框架之一。¹²

核心思想

┌─────────────────────────────────────────────────────────────┐
│  前向Masking过程（可并行）：                                    │
│                                                              │
│  原始文本: "The cat sat on the mat"                         │
│      ↓                                                      │
│  Step 1: "The [MASK] sat on the mat"   (10% masked)        │
│      ↓                                                      │
│  Step 2: "The [MASK] [MASK] on the mat"  (20% masked)    │
│      ↓                                                      │
│  ...                                                        │
│      ↓                                                      │
│  Step T: "[MASK] [MASK] [MASK] [MASK] [MASK]" (100%)      │
│                                                              │
│  反向去噪过程（可并行）：                                      │
│                                                              │
│  全 [MASK] → 预测所有token → 并行生成完整序列                  │
└─────────────────────────────────────────────────────────────┘

数学框架

前向过程

前向过程逐步将token替换为 [MASK]：

$$q(x_t|x_{t-1})=\{\begin{array}{ll}{x}_t=\text{MASK}& \text{以概率 }{\alpha }_t\\ x_t=x_{t-1}& \text{以概率 }1-{\alpha }_t\end{array}$$

等价于：每个token在每步有固定的mask概率。

反向过程

模型预测每个位置应该恢复的token：

$$p_{\theta }(x_{t-1}|x_t)={\text{Dist}}_{\theta }(x_t,t)$$

其中 ${\text{Dist}}_{\theta }$ 是词汇表上的分类分布。

训练目标

通过变分推断可以得到训练目标：

$$\mathcal{L}={\mathbb{E}}_{t,x_0,x_t}\left[-\log p_{\theta }(x_{t-1}|x_t)\right]$$

Transformer参数化

LLaDA¹使用标准Transformer作为骨干网络：

# 简化的MDM前向过程伪代码
class MaskedDiffusionLM(nn.Module):
    def __init__(self, vocab_size, d_model, n_heads, n_layers):
        super().__init__()
        self.transformer = TransformerEncoder(
            d_model=d_model,
            n_heads=n_heads, 
            n_layers=n_layers
        )
        self.output_proj = nn.Linear(d_model, vocab_size)
        
    def forward(self, x_masked, timestep_t):
        # x_masked: 已mask的部分token
        # timestep_t: 当前时间步 (0-1归一化)
        
        # 加入时间步嵌入
        t_emb = self.time_embedding(timestep_t)
        h = x_masked + t_emb
        
        # Transformer编码
        h = self.transformer(h)
        
        # 预测原始token
        logits = self.output_proj(h)
        return logits
    
    def training_loss(self, x0, mask_schedule):
        # x0: 原始token序列
        t = sample_t()  # 随机时间步
        x_t = mask_tokens(x0, t, mask_schedule)  # mask过程
        logits = self.forward(x_t, t)
        return F.cross_entropy(logits, x0)

---

生成过程

标准的逐步去噪

def generate_ddm(model, prompt, n_steps, temperature=1.0):
    """
    离散扩散模型的生成过程
    """
    seq_len = len(prompt)
    x = [MASK] * seq_len  # 初始化为全MASK
    
    # 可选：将prompt位置固定
    x[:len(prompt)] = prompt
    
    for step in reversed(range(n_steps)):
        t = step / n_steps
        logits = model(x, t)
        
        if step > 0:  # 非最后一步，只恢复部分token
            probs = F.softmax(logits / temperature, dim=-1)
            # 采样并mask掉一些token
            for i in range(seq_len):
                if random.random() < (step / n_steps):
                    x[i] = torch.multinomial(probs[i], 1).item()
        else:  # 最后一步，贪婪解码
            x = torch.argmax(logits, dim=-1)
    
    return x

并行生成的优势

步骤	自回归模型	扩散模型（MDM）
Token 1	需要	需要
Token 2	需要Token 1	需要全序列
Token 3	需要Token 1,2	需要全序列
…	顺序依赖	可并行

---

与自回归模型的关键区别

1. 建模目标

自回归模型：

$$p(x_1,x_2,\dots ,x_n)=\prod _{i=1}^{n}p(x_i|x_1,\dots ,x_{i-1})$$

扩散模型：

$$p(x_1,\dots ,x_n)=p_{\theta }(x_0)\cdot \prod _{t=1}^T{p}_{\theta }(x_{t-1}|x_t)$$

2. 推理计算量

• AR模型：$O(n)$ 次前向传播（n = 序列长度）
• 扩散模型：$O(T)$ 次前向传播（T = 扩散步数）

然而，MDM的$T$通常远小于$n$（如$T=8$或16），且可高度并行化。

3. 上下文利用

AR模型：x1 → x2 → x3 → ... → xn
        ↓     ↓     ↓
       仅能   仅能   仅能
       看到   看到   看到
       前缀   前缀   前缀

扩散模型：x1, x2, x3, ..., xn (全序列同时建模)
          ↕    ↕    ↕    ↕
         可利用完整上下文信息

---

主要研究方向

1. 架构优化

• Mamba骨干：DiffuMamba 使用状态空间模型替代Transformer
• 混合架构：结合注意力机制与线性复杂度模型

2. 推理加速

• Unmasking策略学习：dUltra 使用强化学习优化解mask顺序
• Block-level扩散：避免全序列扩散的计算开销

3. 多模态扩展

• 视觉-语言模型：LaViDa 将扩散扩展到多模态场景

4. 大规模训练

• LLaDA：首个8B参数规模的扩散语言模型
• LLaDA2.0：100B参数的扩散语言模型

---

参考

Footnotes

1. LLaDA: Large Language Diffusion Models ↩ ↩² ↩³

2. Masked Diffusion Models: A Study on Discrete Token Prediction ↩