离散状态马尔可夫链文本扩散模型D3PM详解与实现

引言：D3PM的历史定位与核心贡献

D3PM（Discrete Denoising Diffusion Probabilistic Models）是Austin等人在2021年提出的工作，它首次将扩散模型的框架从连续状态空间系统性地推广到离散状态空间，为扩散模型在文本等离散数据上的应用奠定了理论基础。D3PM不仅在理论上建立了离散扩散的完整概率框架，还在实践中展示了其在文本生成和图像生成上的有效性。

D3PM的核心思想是用离散的马尔可夫转移矩阵替代连续扩散中的高斯噪声添加。在连续扩散中，前向过程通过添加高斯噪声将数据分布转化为标准正态分布；在D3PM中，前向过程通过离散状态转移将数据分布转化为均匀分布或吸收状态分布。这种推广保持了扩散模型的核心优势——可变分推断训练、可迭代去噪采样——同时适配了离散数据的特性。

本文将深入解析D3PM的数学框架，包括离散转移矩阵的设计、前向过程的边际分布推导、反向过程的参数化、训练目标的变分推导，以及多种转移矩阵类型的比较。

离散马尔可夫链基础

离散状态空间与转移矩阵

在D3PM中，数据被建模为离散类别变量，取值范围为{0, 1, …, K-1}，其中K是状态数（在文本生成中即词表大小）。前向过程是一个离散时间马尔可夫链，每一步的转移由一个K乘以K的转移矩阵Q_t定义。

转移矩阵Q_t的每一行是一个概率分布，满足Q_t[i, j]大于等于0且行和为1。Q_t[i, j]表示从状态i转移到状态j的概率。给定当前状态x_{t-1} = i，下一步状态x_t的分布为q(x_t = j | x_{t-1} = i) = Q_t[i, j]。如果将状态表示为one-hot向量，则转移可以写为矩阵乘法x_t ~ Cat(x_{t-1} Q_t)。

累积转移矩阵

与连续扩散类似，离散扩散的前向过程具有良好的边际性质。给定初始状态x_0，经过t步转移后的边际分布为q(x_t | x_0) = Cat(x_0 Q_bar_t)，其中Q_bar_t = Q_1 Q_2 … Q_t是累积转移矩阵。这个性质使得我们可以直接从x_0采样任意时间步的x_t，而不需要逐步执行转移。

累积转移矩阵的计算需要注意数值稳定性。当K很大且T很大时，矩阵乘法的计算和存储成本很高。对于特定类型的转移矩阵（如吸收状态矩阵），Q_bar_t有解析形式，可以避免显式矩阵乘法。

马尔可夫链的平稳分布

离散马尔可夫链的平稳分布满足经过转移后分布不变。在D3PM中，前向过程的平稳分布对应扩散过程的"先验分布"——即t趋近无穷时x_t的分布。不同类型的转移矩阵有不同的平稳分布。均匀转移矩阵的平稳分布是均匀分布。吸收状态转移矩阵的平稳分布是所有质量集中在吸收状态的分布。

平稳分布的选择影响生成过程：采样从平稳分布开始，逐步去噪到数据分布。因此平稳分布应该是容易采样的，这也是均匀分布和吸收状态分布被常用的原因。

D3PM的数学框架

前向过程

D3PM的前向过程定义为q(x_t | x_{t-1}) = Cat(x_{t-1} Q_t)，边际分布为q(x_t | x_0) = Cat(x_0 Q_bar_t)。给定x_0，采样x_t的过程为计算概率向量p = x_0 Q_bar_t，然后从Cat(p)中采样。

反向过程

反向过程的目标是从平稳分布逐步恢复数据分布。反向转移概率参数化为p_theta(x_{t-1} | x_t) = sum_{x_0} q(x_{t-1} | x_t, x_0) p_theta(x_0 | x_t)。这里使用了类似DDPM的分解技巧：先预测x_0的分布，再通过解析公式计算反向转移。

q(x_{t-1} | x_t, x_0)的后验分布可以通过贝叶斯公式计算。对于特定类型的转移矩阵，这个后验有简洁的解析形式。

训练目标的变分推导

D3PM的训练目标通过变分下界推导。对数似然的变分下界包含重建项、先验匹配项和一致性项。通过适当的参数化，这个目标可以简化为L = E[-log p_theta(x_0 | x_t)]，即让模型在给定带噪声的x_t时正确预测x_0。这是一个简单的交叉熵损失，与BERT的MLM目标类似。

转移矩阵的类型与设计

D3PM论文中探讨了多种转移矩阵类型，每种都有不同的特性和适用场景。

均匀转移矩阵

均匀转移矩阵以概率(1-beta_t)保持原状态，以概率beta_t/(K-1)均匀转移到其他任意状态。平稳分布为均匀分布。这种转移矩阵的物理直觉是"随机替换"——以小概率将当前token随机替换为词表中的任意其他token。

均匀转移矩阵的累积转移矩阵有解析形式。设alpha_t = prod(1-beta_s * K/(K-1))，则Q_bar_t[i,j] = alpha_t if i=j, (1-alpha_t)/K if i!=j，即以概率alpha_t保持原状态，以概率(1-alpha_t)均匀分布到所有状态。

吸收状态转移矩阵

吸收状态转移矩阵以概率(1-beta_t)保持原状态，以概率beta_t转移到一个特殊的吸收状态（通常用掩码标记）。平稳分布为所有质量集中在掩码状态。这种转移矩阵的物理直觉是"逐步掩码"——以小概率将当前token替换为掩码。

吸收状态矩阵的累积转移矩阵有简洁的解析形式。设alpha_t = prod(1-beta_s)，则q(x_t | x_0)为以概率alpha_t保持x_0，以概率(1-alpha_t)为掩码。这意味着前向采样只需要一个伯努利试验，非常高效。

离散化高斯转移矩阵

离散化高斯转移矩阵将连续高斯噪声离散化到整数状态上。对于状态有自然顺序关系的场景（如灰度像素值0-255），这种转移矩阵保留了状态的顺序信息。对于文本生成，离散化高斯转移矩阵不太适用，因为token之间没有自然的顺序关系。

后验分布的解析计算

吸收状态转移的后验

对于吸收状态转移矩阵，后验有简洁的解析形式。如果x_t不是掩码（即x_t是原始token），则x_{t-1}必然等于x_t，这是因为吸收状态转移不会改变非吸收状态的值。

如果x_t是掩码，则q(x_{t-1} | x_t, x_0)以概率alpha_{t-1}/alpha_t取x_0，以概率(1-alpha_{t-1}/alpha_t)取掩码。这是因为在前向过程中，x_{t-1}要么是x_0（以概率alpha_{t-1}），要么是掩码（以概率1-alpha_{t-1}）。

这个简洁的后验使得吸收状态扩散的反向过程非常高效：只需要在掩码位置做决策，非掩码位置保持不变。

反向过程的参数化

x_0预测参数化

D3PM最常用的参数化方式是x_0预测：训练一个神经网络预测原始token的分布。给定x_0的预测分布，反向转移通过后验公式计算。对于吸收状态转移，这简化为非掩码位置保持不变，掩码位置以一定概率从预测分布采样token。

网络架构

预测网络通常由Transformer参数化。输入是带噪声的token序列和时间步，输出是对x_0的预测分布（词表上的softmax）。网络架构与BERT类似，使用双向自注意力。时间步信息通过AdaLN注入。输出层是词表大小的线性投影加softmax。

与BERT MLM的联系

D3PM的x_0预测目标与BERT的MLM目标非常相似——都是从被扰动的输入中预测原始token。关键区别在于扰动的方式和程度：BERT使用固定的15%掩码比例，而D3PM使用从0%到接近100%的多层次掩码比例。这种多尺度训练使得D3PM不仅能做填充任务，还能做无条件生成。

采样算法

标准采样

D3PM的标准采样从平稳分布开始，逐步执行反向转移。对于吸收状态扩散，采样过程可以大幅简化。由于非掩码位置在反向过程中保持不变，每步只需要处理掩码位置。这种采样方式允许某些位置在早期步骤就被填充，而其他位置在后期步骤才被填充，实现了"由粗到细"的生成过程。

跳跃采样

与连续扩散类似，D3PM也可以使用跳跃采样加速。选择部分关键时间步进行去噪，跳过中间步骤。但离散扩散的跳跃采样需要更谨慎地选择时间步，因为离散转移的"信息损失"比连续噪声更突然。

PyTorch完整实现

下面实现一个完整的D3PM文本生成模型，支持多种转移矩阵类型。

import torch
import torch.nn as nn
import torch.nn.functional as F
import math
from typing import Optional
from enum import Enum

class TransitionType(Enum):
    """转移矩阵类型"""
    UNIFORM = "uniform"
    ABSORBING = "absorbing"
    DISCRETE_GAUSSIAN = "discrete_gaussian"

class D3PMConfig:
    """D3PM配置"""
    def __init__(self, vocab_size=32000, dim=512, seq_len=128,
                 num_layers=6, num_heads=8, num_timesteps=1000,
                 transition_type="absorbing", mask_id=0,
                 beta_schedule="cosine", dropout=0.1):
        self.vocab_size = vocab_size
        self.dim = dim
        self.seq_len = seq_len
        self.num_layers = num_layers
        self.num_heads = num_heads
        self.num_timesteps = num_timesteps
        self.transition_type = TransitionType(transition_type)
        self.mask_id = mask_id
        self.beta_schedule = beta_schedule
        self.dropout = dropout

def get_beta_schedule(schedule, T):
    """获取beta调度"""
    if schedule == "linear":
        return torch.linspace(0.01, 0.5, T)
    elif schedule == "cosine":
        steps = torch.arange(T + 1, dtype=torch.float32)
        f = torch.cos(((steps / T + 0.008) / 1.008) * math.pi * 0.5) ** 2
        acp = f / f[0]
        betas = 1 - (acp[1:] / acp[:-1])
        return torch.clamp(betas, 1e-4, 0.999)
    elif schedule == "sqrt":
        x = torch.linspace(0, T, T + 1, dtype=torch.float32)
        acp = torch.clamp(1 - torch.sqrt(x / T + 1e-8), 1e-4, 1.0)
        return torch.clamp(1 - (acp[1:] / acp[:-1]), 1e-4, 0.999)
    raise ValueError(f"Unknown schedule: {schedule}")

class TransitionMatrix:
    """转移矩阵管理器，支持均匀、吸收状态和离散化高斯转移"""
    def __init__(self, config):
        self.config = config
        self.K = config.vocab_size
        self.T = config.num_timesteps
        self.mask_id = config.mask_id
        self.transition_type = config.transition_type
        self.betas = get_beta_schedule(config.beta_schedule, self.T)
        alphas = 1.0 - self.betas
        self.alpha_bars = torch.cumprod(alphas, dim=0)
        if self.transition_type == TransitionType.UNIFORM:
            scale = self.K / (self.K - 1)
            alphas_u = torch.clamp(1.0 - self.betas * scale, 1e-6, 1.0)
            self.alpha_bars = torch.cumprod(alphas_u, dim=0)

    def q_sample(self, x_0, t):
        """前向采样：给定x_0和时间步t，采样x_t"""
        if self.transition_type == TransitionType.ABSORBING:
            ab = self.alpha_bars.to(x_0.device)[t].unsqueeze(-1)
            mask = torch.rand_like(x_0, dtype=torch.float) < (1 - ab)
            x_t = x_0.clone()
            x_t[mask] = self.mask_id
            return x_t
        elif self.transition_type == TransitionType.UNIFORM:
            ab = self.alpha_bars.to(x_0.device)[t].unsqueeze(-1)
            replace = torch.rand_like(x_0, dtype=torch.float) < (1 - ab)
            x_t = x_0.clone()
            rand_tokens = torch.randint_like(x_0, 0, self.K)
            rand_tokens = torch.where(rand_tokens == x_0,
                                      (rand_tokens + 1) % self.K, rand_tokens)
            x_t[replace] = rand_tokens[replace]
            return x_t
        else:
            ab = self.alpha_bars.to(x_0.device)[t].unsqueeze(-1)
            mask = torch.rand_like(x_0, dtype=torch.float) < (1 - ab)
            x_t = x_0.clone()
            x_t[mask] = self.mask_id
            return x_t

class D3PMTransformer(nn.Module):
    """D3PM去噪Transformer网络"""
    def __init__(self, config):
        super().__init__()
        self.config = config
        self.token_embed = nn.Embedding(config.vocab_size, config.dim)
        self.pos_embed = nn.Parameter(torch.randn(1, config.seq_len, config.dim) * 0.02)
        self.time_embed = nn.Sequential(
            nn.Linear(config.dim, config.dim * 4), nn.SiLU(),
            nn.Linear(config.dim * 4, config.dim))
        self.layers = nn.ModuleList([self._make_layer() for _ in range(config.num_layers)])
        self.out_norm = nn.LayerNorm(config.dim)
        self.out_proj = nn.Linear(config.dim, config.vocab_size)
        self._init_weights()

    def _make_layer(self):
        dim = self.config.dim
        d = nn.ModuleDict({
            "norm1": nn.LayerNorm(dim),
            "attn": nn.MultiheadAttention(dim, self.config.num_heads,
                                          dropout=self.config.dropout, batch_first=True),
            "norm2": nn.LayerNorm(dim),
            "ffn": nn.Sequential(
                nn.Linear(dim, dim * 4), nn.GELU(), nn.Dropout(self.config.dropout),
                nn.Linear(dim * 4, dim), nn.Dropout(self.config.dropout)),
            "adaLN": nn.Sequential(nn.SiLU(), nn.Linear(dim, 6 * dim)),
        })
        nn.init.zeros_(d["adaLN"][-1].weight)
        nn.init.zeros_(d["adaLN"][-1].bias)
        return d

    def _init_weights(self):
        for m in self.modules():
            if isinstance(m, nn.Linear):
                nn.init.xavier_uniform_(m.weight)
                if m.bias is not None:
                    nn.init.zeros_(m.bias)
            elif isinstance(m, nn.Embedding):
                nn.init.normal_(m.weight, mean=0, std=0.02)

    def _time_emb(self, t):
        dim = self.config.dim
        half = dim // 2
        freq = torch.exp(torch.arange(half, device=t.device, dtype=torch.float) *
                         (-math.log(10000) / max(half - 1, 1)))
        args = t[:, None].float() * freq[None, :]
        emb = torch.cat([torch.sin(args), torch.cos(args)], dim=-1)
        if dim % 2 == 1:
            emb = F.pad(emb, (0, 1))
        return self.time_embed(emb)

    def forward(self, x_t, t):
        """预测x_0分布: [batch, seq], [batch] -> [batch, seq, vocab]"""
        h = self.token_embed(x_t) + self.pos_embed[:, :x_t.size(1), :]
        t_emb = self._time_emb(t)
        for layer in self.layers:
            s1, sc1, g1, s2, sc2, g2 = layer["adaLN"](t_emb).chunk(6, dim=-1)
            nh = layer["norm1"](h)
            nh = nh * (1 + sc1.unsqueeze(1)) + s1.unsqueeze(1)
            attn_out, _ = layer["attn"](nh, nh, nh)
            h = h + g1.unsqueeze(1) * attn_out
            nh = layer["norm2"](h)
            nh = nh * (1 + sc2.unsqueeze(1)) + s2.unsqueeze(1)
            h = h + g2.unsqueeze(1) * layer["ffn"](nh)
        return self.out_proj(self.out_norm(h))

class D3PM(nn.Module):
    """完整的D3PM模型"""
    def __init__(self, config):
        super().__init__()
        self.config = config
        self.transition = TransitionMatrix(config)
        self.net = D3PMTransformer(config)

    def compute_loss(self, x_0):
        """计算训练损失"""
        batch_size = x_0.shape[0]
        device = x_0.device
        t = torch.randint(0, self.config.num_timesteps, (batch_size,), device=device)
        x_t = self.transition.q_sample(x_0, t)
        logits = self.net(x_t, t)
        if self.config.transition_type == TransitionType.ABSORBING:
            mask = (x_t == self.config.mask_id)
            if mask.any():
                return F.cross_entropy(logits[mask], x_0[mask])
            return F.cross_entropy(logits.reshape(-1, self.config.vocab_size), x_0.reshape(-1))
        return F.cross_entropy(logits.reshape(-1, self.config.vocab_size), x_0.reshape(-1))

    @torch.no_grad()
    def sample(self, batch_size=1, seq_len=None, device=None):
        """采样生成"""
        if device is None:
            device = next(self.parameters()).device
        if seq_len is None:
            seq_len = self.config.seq_len
        mask_id = self.config.mask_id
        T = self.config.num_timesteps
        x = torch.full((batch_size, seq_len), mask_id, dtype=torch.long, device=device)
        for t in reversed(range(T)):
            t_batch = torch.full((batch_size,), t, device=device, dtype=torch.long)
            logits = self.net(x, t_batch)
            probs = F.softmax(logits, dim=-1)
            ab_t = self.transition.alpha_bars[t].item()
            ab_prev = self.transition.alpha_bars[t-1].item() if t > 0 else 1.0
            keep_prob = 0.0 if t == 0 else min(ab_prev / max(ab_t, 1e-8), 1.0)
            mask_pos = (x == mask_id)
            if not mask_pos.any():
                continue
            fill_pos = mask_pos & (torch.rand(batch_size, seq_len, device=device) > keep_prob)
            if fill_pos.any():
                sampled = torch.multinomial(probs[fill_pos], 1).squeeze(-1)
                x[fill_pos] = sampled
        return x

    @torch.no_grad()
    def sample_fast(self, batch_size=1, seq_len=None, device=None, num_steps=50):
        """快速跳跃采样"""
        if device is None:
            device = next(self.parameters()).device
        if seq_len is None:
            seq_len = self.config.seq_len
        mask_id = self.config.mask_id
        T = self.config.num_timesteps
        step_indices = torch.linspace(0, T - 1, num_steps, dtype=torch.long)
        x = torch.full((batch_size, seq_len), mask_id, dtype=torch.long, device=device)
        for i in reversed(range(len(step_indices))):
            t = step_indices[i].item()
            t_batch = torch.full((batch_size,), t, device=device, dtype=torch.long)
            logits = self.net(x, t_batch)
            probs = F.softmax(logits, dim=-1)
            ab_t = self.transition.alpha_bars[t].item()
            ab_prev = self.transition.alpha_bars[step_indices[i-1].item()].item() if i > 0 else 1.0
            keep_prob = 0.0 if i == 0 else min(ab_prev / max(ab_t, 1e-8), 1.0)
            mask_pos = (x == mask_id)
            if not mask_pos.any():
                continue
            fill_pos = mask_pos & (torch.rand(batch_size, seq_len, device=device) > keep_prob)
            if fill_pos.any():
                sampled = torch.multinomial(probs[fill_pos], 1).squeeze(-1)
                x[fill_pos] = sampled
        return x

def train_d3pm(model, dataloader, num_epochs=10, lr=3e-4,
               warmup_steps=1000, device="cuda", log_interval=50):
    """训练D3PM模型"""
    model = model.to(device)
    optimizer = torch.optim.AdamW(model.parameters(), lr=lr, weight_decay=0.01, betas=(0.9, 0.99))
    total_steps = num_epochs * len(dataloader)
    def lr_lambda(step):
        if step < warmup_steps:
            return step / max(warmup_steps, 1)
        progress = (step - warmup_steps) / max(total_steps - warmup_steps, 1)
        return 0.5 * (1 + math.cos(math.pi * progress))
    scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda)
    model.train()
    global_step = 0
    for epoch in range(num_epochs):
        epoch_loss = 0.0
        num_batches = 0
        for batch in dataloader:
            x_0 = batch.to(device)
            loss = model.compute_loss(x_0)
            optimizer.zero_grad()
            loss.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
            optimizer.step()
            scheduler.step()
            epoch_loss += loss.item()
            num_batches += 1
            global_step += 1
            if global_step % log_interval == 0:
                print(f"Epoch {epoch+1}, Step {global_step}, Loss: {epoch_loss/num_batches:.4f}")
        print(f"=== Epoch {epoch+1} 完成, 平均损失: {epoch_loss/num_batches:.4f} ===")
    return model

def demo():
    """完整使用示例"""
    config = D3PMConfig(vocab_size=32000, dim=512, seq_len=128, num_layers=6,
                        num_heads=8, num_timesteps=1000, transition_type="absorbing",
                        mask_id=0, beta_schedule="cosine", dropout=0.1)
    model = D3PM(config)
    dummy_data = torch.randint(1, config.vocab_size, (64, config.seq_len))
    dummy_dataloader = [dummy_data[i*8:(i+1)*8] for i in range(8)]
    device = "cuda" if torch.cuda.is_available() else "cpu"
    print("训练D3PM（吸收状态扩散）...")
    trained = train_d3pm(model, dummy_dataloader, num_epochs=5, device=device)
    trained.eval()
    print("\n标准采样（1000步）...")
    tokens = trained.sample(batch_size=4, seq_len=32, device=device)
    print(f"生成token IDs: {tokens[0][:20].tolist()}")
    print("\n快速采样（50步）...")
    tokens_fast = trained.sample_fast(batch_size=4, seq_len=32, device=device, num_steps=50)
    print(f"快速生成token IDs: {tokens_fast[0][:20].tolist()}")

if __name__ == "__main__":
    demo()

D3PM的实验结果与发现

文本生成实验

D3PM论文在文本生成任务上进行了系统实验。在语言建模基准上，吸收状态扩散模型取得了最好的效果，显著优于均匀转移和离散化高斯转移。这表明掩码机制与文本数据的特性最为匹配——文本中的信息是"有或无"的，而非连续渐变的。

D3PM的生成质量虽然仍不及同规模的自回归模型，但显著优于之前的非自回归文本生成方法。特别是在短文本生成任务上，D3PM可以生成语法正确、语义连贯的句子。

转移矩阵类型的影响

实验对比了三种转移矩阵类型的效果。吸收状态转移在文本生成上效果最好，因为掩码操作保留了未掩码位置的精确信息，模型只需要预测掩码位置的内容。均匀转移效果较差，因为随机替换引入了大量噪声，模型需要同时处理"哪些位置被改了"和"改成什么了"两个问题。离散化高斯转移在文本上效果最差，因为token ID没有顺序语义。

采样步数的影响

D3PM的生成质量随采样步数增加而提升，但存在边际递减效应。1000步采样的质量显著优于100步，但1000步和2000步之间的差异很小。通过跳跃采样，50步可以达到接近1000步的质量，这使得D3PM在实际应用中具备可行性。

D3PM与连续扩散的理论联系

统一的变分框架

D3PM和DDPM都可以从变分自编码器的框架推导。两者的核心区别在于前向过程的定义：DDPM使用高斯转移，D3PM使用离散转移。但两者的反向过程参数化和训练目标推导是类似的，都通过变分下界得到简化目标。

极限关系

在特定条件下，离散扩散可以收敛到连续扩散。当状态数K趋近于无穷且转移矩阵选择适当时，离散转移矩阵的极限行为等价于高斯转移。这种极限关系为理解两种扩散模型的联系提供了理论桥梁。

采样策略的对应

D3PM的采样策略与DDPM有直接对应。D3PM的标准采样对应DDPM的 ancestral sampling，D3PM的跳跃采样对应DDIM。这种对应关系使得连续扩散的采样加速技术可以迁移到离散扩散。

D3PM的改进与扩展

DiffusionBERT

DiffusionBERT将预训练的BERT模型作为D3PM的去噪网络，利用BERT的MLM预训练知识加速训练。关键创新在于设计了适合扩散过程的噪声调度，使得BERT的预训练知识能够有效迁移。实验表明，DiffusionBERT在多个文本生成任务上显著优于从头训练的D3PM。

SEDD

SEDD（Score Entropy Discrete Diffusion）提出了一种基于分数匹配的训练目标，避免了D3PM中需要计算后验分布的复杂推导。SEDD直接学习离散分数函数，训练更简单且效果更好。SEDD还提出了更高效的采样算法，将采样步数减少到数十步。

MDLM

MDLM（Masked Diffusion Language Models）对吸收状态扩散进行了简化和优化。MDLM发现，吸收状态扩散可以简化为多尺度掩码语言模型，训练目标可以进一步简化。MDLM在大规模语言建模任务上取得了与自回归模型可比的效果。

实践建议

转移矩阵选择

对于文本生成任务，吸收状态转移矩阵是首选。它简单高效，与掩码语言模型兼容，且实验效果最好。均匀转移矩阵适合需要更随机扰动的场景，但效果通常不如吸收状态。离散化高斯转移不适合文本，仅在有自然顺序的离散数据上使用。

噪声调度调优

余弦调度通常是较好的默认选择。对于短序列，可以考虑更激进的调度（如平方根调度）。对于长序列，应使用更平缓的调度。噪声调度的选择应使得在中间时间步（掩码比例30%-70%）有足够的训练信号。

模型架构选择

Transformer是去噪网络的主流选择。使用Pre-LN架构比Post-LN更稳定。对于大规模模型，考虑使用Flash Attention降低计算成本。可以直接使用BERT/RoBERTa的预训练权重初始化，显著加速训练。

采样策略

标准采样（遍历所有T步）质量最高但速度最慢。跳跃采样（50步）是质量与速度的良好平衡。对于实时应用，可以进一步减少到20-30步，但质量会有明显下降。采样时可以使用温度参数控制多样性。

总结

D3PM是离散扩散模型的奠基性工作，它将扩散模型从连续空间系统性地推广到离散空间，为文本生成提供了新的理论框架。D3PM的核心贡献在于建立了离散马尔可夫链扩散的完整概率框架，包括转移矩阵的设计、前向过程的边际分布、反向过程的后验计算和训练目标的变分推导。

吸收状态转移矩阵是D3PM在文本生成上最成功的方案，它将前向过程建模为逐步掩码，反向过程建模为迭代去掩码，与掩码语言模型天然兼容。D3PM的训练目标是简单的交叉熵损失，训练过程稳定，与预训练语言模型兼容性好。

虽然D3PM的生成质量仍不及自回归模型，但其非自回归的并行生成能力、统一的填充与生成框架、以及灵活的可控性，使其成为文本生成领域的重要研究方向。后续的DiffusionBERT、SEDD、MDLM等工作在D3PM的基础上进一步改进，不断缩小与自回归模型的差距。

本文从数学原理到代码实现，完整地解析了D3PM的各个组件。提供的PyTorch实现支持多种转移矩阵类型，读者可以基于此快速搭建自己的离散扩散文本生成系统，并根据具体任务进行调整和优化。