离散扩散模型：基于连续时间马尔可夫链的文本与序列生成新范式

1. 项目概述：当离散数据遇上连续时间

在生成式AI的浪潮里，扩散模型无疑是当前最耀眼的明星之一。从生成逼真图像的DALL-E、Stable Diffusion，到合成音频、视频的Sora，其核心都离不开扩散过程。然而，绝大多数人熟悉的扩散模型，无论是DDPM还是其后续变体，都建立在一个默认的假设之上：数据是连续的，比如图像中的像素值可以平滑地从噪声过渡到清晰。但现实世界充满了离散数据——文本、分子结构、代码、用户点击序列，这些数据点之间没有平滑的中间态。你无法想象一个介于单词“猫”和“狗”之间的“半猫半狗”的单词是什么样子。

这就引出了一个核心挑战：如何将强大的扩散生成能力，迁移到离散数据领域？传统的做法，比如引入嵌入层并添加高斯噪声，往往效果不佳或训练不稳定。而“基于连续时间马尔可夫链的离散扩散模型”正是为解决这一难题而生。它巧妙地绕开了“连续中间态”的障碍，将离散数据的生成过程，建模为一个在离散状态空间上、但时间维度连续演化的随机过程。简单来说，它不再试图“模糊”一个单词，而是定义了这个单词以何种速率“跳转”到词典中的其他单词（包括一个特殊的[MASK]状态），整个过程的时间是连续的。这种方法不仅理论优美，在实践中也展现出了生成高质量、多样性离散序列的潜力，为文本生成、分子设计、代码补全等领域提供了新的强大工具。

如果你正在处理自然语言、图结构或任何符号序列的生成任务，并且对传统自回归模型（如GPT系列）的逐 token 生成方式感到局限，或者对VAE在离散数据上的表现不满意，那么理解这套框架将为你打开一扇新的大门。它提供了一种全新的、并行化的生成范式。

2. 核心原理拆解：连续时间马尔可夫链如何驱动离散扩散

要理解这个模型，我们需要拆解两个核心概念：“离散扩散”和“连续时间马尔可夫链”，并看它们是如何结合在一起的。

2.1 离散扩散的本质：状态空间的随机游走

首先，忘掉图像扩散中那幅从模糊到清晰的渐变图。对于离散数据，比如一个来自大小为V的词汇表的单词，它的状态是离散的：x可以是V个可能值中的任何一个。离散扩散过程描述的是这个状态如何随时间t（从0到T）变化。

这个过程通常定义为一个前向过程和一个反向过程：

• 前向过程：从真实数据x_0开始，随着时间t增加，逐步地、随机地破坏它。在连续时间设定下，破坏不是按固定步长发生的，而是在任意时刻，数据点都有一定的概率“跃迁”到其他状态。最终，当t=T时，数据被完全破坏成一个简单的先验分布，例如一个均匀分布或一个吸收态（如全部变成[MASK]）。
• 反向过程：这是生成的核心。我们学习一个模型，使其能够从简单的先验分布（t=T时的状态）开始，沿着时间t反向演化，逐步“去噪”或“修复”数据，最终得到真实数据分布x_0。

关键点在于，这个“破坏”和“修复”的“动作”，是离散状态之间的瞬时切换，而不是连续值的微小扰动。

2.2 连续时间马尔可夫链：为扩散提供数学引擎

连续时间马尔可夫链是描述上述过程的完美数学工具。一个CTMC由两部分定义：

1. 状态空间：对我们来说，就是所有可能的数据状态集合（如所有单词）。
2. 转移速率矩阵 Q(t)：这是一个V x V的矩阵。矩阵中的每个元素Q_{ij}(t)代表了在时间t，从状态i瞬时跃迁到状态j的速率。速率不是概率，它的单位是“每单位时间”。在极短的时间Δt内，从i跳到j的概率近似为Q_{ij}(t) * Δt。

在离散扩散的前向过程中，我们人为设计一个速率矩阵Q_fwd(t)。一个常见且简单的设计是：让每个状态都以一个固定的速率β(t)跃迁到一个共同的“吸收态”或“掩码态”[M]。同时，也可以定义状态之间以某种规则相互跃迁。这个设计保证了随着时间推移，数据最终会趋于一个我们已知的简单分布。

注意：这里β(t)是一个关于时间的函数，类似于连续扩散中的噪声调度。它控制了破坏的“速度”，通常在开始时较小，结束时较大。

2.3 核心桥梁：时间反转与得分匹配

那么，我们如何学习反向生成过程呢？这里需要用到CTMC的一个深刻性质：在一定的正则条件下，一个CTMC的时间反转过程，本身也是一个CTMC。

假设前向过程由Q_fwd(t)定义。理论上存在一个反向过程的速率矩阵Q_rev(t)，它使得从终点T开始、按照Q_rev(t)演化，最终得到的路径分布，正好是前向路径分布的时间反转。我们的生成模型的目标，就是去参数化并学习这个反向速率矩阵Q_rev(t)。

但是，直接学习Q_rev(t)是困难的。一个关键的推导（基于随机过程的反向时间生成器理论）给出了一个可操作的公式：Q_rev_{ij}(t) = (p_t(j) / p_t(i)) * Q_fwd_{ji}(t)其中p_t(i)表示在时间t时状态为i的概率（边缘概率）。这个公式告诉我们，反向速率不仅依赖于前向速率的转置Q_fwd_{ji}，还依赖于两个状态的相对概率密度比p_t(j)/p_t(i)。

这个概率密度比是未知的，但正是我们的神经网络需要学习的东西！我们可以定义一个模型s_θ(i, t)，去估计这个“得分”或者说概率对数的梯度（在离散空间的一种推广）。最终，学习目标可以转化为一种加权的交叉熵损失：让模型预测在给定当前状态和时间下，下一个瞬时可能发生的跃迁类型。

实操心得：理解这一点至关重要。模型的输出层通常是一个大小为V的向量，但它不是直接预测下一个token，而是预测每个token作为“跃迁目标”的“倾向性得分”。训练时，我们从前向过程中采样一个真实跃迁事件（在时间t从状态x_s跳到了x_t），然后让模型去拟合这个事件发生的“速率”。这避免了在离散空间做难以处理的似然最大化。

3. 模型训练全流程详解

理论之后，我们进入实战。训练一个离散扩散模型，步骤比连续扩散模型要更精细地处理时间。

3.1 前向过程设计与速率矩阵构造

第一步是设计前向破坏过程。这里有几个常见的选择，直接影响模型性能和训练稳定性。

1. 均匀吸收：每个token以速率β(t)独立地变成[MASK] token。这是最简单的设计，Q_fwd矩阵非常稀疏：对角线元素为-β(t)（表示离开状态i的总速率），i-> [MASK] 的速率为β(t)，其他为0。

   • 优点：简单，易于实现，反向过程的学习目标清晰。
   • 缺点：破坏过程过于简单，可能无法让模型学到足够复杂的依赖关系来生成高质量样本。

2. 均匀转移：每个token以速率β(t)跃迁到词汇表中的任何其他token，且概率均匀。此时，Q_fwd的非对角线元素（i≠j）为β(t)/(V-1)，对角线元素相应调整。

   • 优点：破坏更“剧烈”，迫使模型在反向过程中学习更强的生成能力。
   • 缺点：计算开销稍大，因为需要考虑所有可能的转移。

3. 基于嵌入相似度的转移：跃迁到另一个tokenj的速率与当前tokeni的嵌入向量相似度成正比（例如，使用点积或余弦相似度）。这更接近“语义模糊”的直觉。

   • 优点：破坏过程更自然，可能提升生成样本的语义连贯性。
   • 缺点：需要预训练的嵌入矩阵，且速率矩阵不再是静态的，计算更复杂。

我的建议：对于初次实现，强烈推荐从均匀吸收开始。它足够简单，能让你快速搭建起整个训练-采样流程，并验证代码的正确性。β(t)函数可以选择线性调度，例如β(t) = 0.01 + (1.0 - 0.01) * (t/T)，其中t从0到1。

3.2 损失函数推导与实现

损失函数是训练的核心。基于之前提到的理论，对于一条从x_s在时间s跃迁到x_t在时间t的路径（其中t = s + Δt，Δt很小），我们可以推导出以下训练目标：

L(θ) = E[ -log π_θ(y=x_t | x_s, s) ]

这里π_θ是模型参数化的反向跃迁概率分布。经过一系列推导（涉及Girsanov定理和重要性采样），这个期望可以转化为一个更易于计算的形式。在实践中，一个稳定且常用的实现方式是加权交叉熵损失：

1. 从训练数据中采样一个真实样本x_0。
2. 从[0, T]中均匀采样一个连续时间t。
3. 根据前向过程Q_fwd(t)，模拟从x_0在时间t内的演化，得到当前被破坏的状态x_t。由于CTMC的模拟需要事件时间，我们可以利用一个技巧：对于均匀吸收，x_t中每个token保持原样的概率是exp(-∫_0^t β(s) ds)，被掩码的概率是1 - exp(...)。我们可以直接按这个概率采样，而无需真正模拟连续时间路径。
4. 将(x_t, t)输入神经网络模型s_θ。模型输出一个形状为(batch_size, seq_len, V)的张量，可以理解为每个位置、每个可能目标token的未归一化“得分”。
5. 计算损失。这里的关键是，损失关注的是模型如何预测“数据原本的样子”。对于被掩码的位置，我们希望模型预测出原始的x_0；对于未被掩码的位置，我们希望模型“保持不动”或预测自身。因此，损失可以构造为：Loss = mean( mask * CE(s_θ(x_t, t), x_0) )其中mask是一个指示哪些位置被破坏的掩码，CE是交叉熵损失。mask的权重可以根据时间t进行调整，通常给中间时间点更高的权重，因为它们最难预测。

代码片段示意（PyTorch风格）：

def compute_loss(model, x0, t, Q_fwd_integral): """ x0: 原始数据，形状 (B, L) t: 连续时间，形状 (B,) Q_fwd_integral: 计算出的累积速率积分，用于计算掩码概率 """ # 1. 计算掩码概率 mask_prob = 1 - torch.exp(-Q_fwd_integral(t)) # 形状 (B,) # 2. 采样掩码 mask = torch.bernoulli(mask_prob.unsqueeze(-1).expand_as(x0)) # 形状 (B, L) # 3. 生成被破坏的状态 x_t x_t = x0.clone() x_t[mask.bool()] = MASK_TOKEN_ID # 将被掩码位置替换为[MASK] # 4. 模型前向传播 logits = model(x_t, t) # 形状 (B, L, V) # 5. 计算加权交叉熵损失 loss = F.cross_entropy(logits.transpose(1, 2), x0, reduction='none') # 形状 (B, L) weighted_loss = (mask * loss).sum() / (mask.sum() + 1e-8) # 只对被破坏的位置计算损失 return weighted_loss

3.3 网络架构选择与输入处理

模型架构没有严格限制，任何能处理序列数据的网络都可以，但需要处理连续时间输入t。

• 主干网络：Transformer Encoder 是最自然的选择，因为它能很好地建模序列内部的依赖关系。对于中等长度的序列，CNN（如Temporal Convolutional Networks）也是一个轻量高效的选项。
• 时间嵌入：连续时间t需要被编码成向量后注入模型。通常采用高斯随机特征映射（Random Fourier Features），即：γ(t) = [sin(ω_1 t), cos(ω_1 t), ..., sin(ω_m t), cos(ω_m t)]其中ω是从某个分布（如高斯分布）中采样的固定频率。然后将γ(t)通过一个线性层，加到每个位置的token嵌入上，或者作为AdaGN（自适应组归一化）的参数。
• 输入表示：模型的输入是已被部分破坏的序列x_t。我们需要一个标准的嵌入层将token ID映射为向量。对于[MASK] token，使用一个可学习的专属嵌入向量。

注意事项：训练初期，由于大部分位置未被掩码，损失可能很小且下降很快，但这不代表模型学得好。要密切关注在验证集上，模型对被掩码部分的预测准确率。一个有效的监控指标是“掩码位置token预测准确率”。

4. 采样算法：从噪声到数据的连续时间逆过程

训练完成后，我们拥有了一个学习了反向速率Q_rev_θ的模型。采样过程，就是从先验分布出发，沿着反向CTMC模拟一条轨迹，最终生成样本。

4.1 采样器类型：模拟与ODE求解

与连续扩散类似，离散扩散的采样也有两种主要思路：

1. 随机采样器：直接模拟反向CTMC。这需要生成随机跃迁事件和时间。可以使用“拒绝采样”或“ thinning”算法来模拟非齐次（时间依赖）的CTMC。这种方法生成的样本多样性好，但速度较慢，且实现复杂。
2. 确定性采样器：通过求解一个概率流常微分方程进行采样。对于离散扩散，也存在一个对应的ODE，其“速度场”由模型预测的得分和速率矩阵决定。使用ODE求解器（如欧拉法、Heun法）可以对其进行数值积分。这种方法速度快、确定性强，适合需要快速生成或插值的场景。

对于大多数应用，确定性采样器是更实用、更推荐的选择。它的步骤可以概括为：

• 初始化：在时间t = T，从先验分布采样初始状态x_T。对于均匀吸收前向过程，先验就是全部为[MASK]的序列。
• 迭代去噪：从t=T到t=0，将时间离散化为多个步长{τ_0=T, τ_1, ..., τ_N=0}。对于每一步k： a. 在时间τ_k，计算当前状态x_{τ_k}对应的模型输出（得分）。 b. 根据ODE公式，计算状态在时间上的导数dx/dt。这个导数本质上描述了每个位置从当前token变为其他token的“倾向性流量”。 c. 使用ODE求解器更新状态：x_{τ_{k+1}} = ODESolverStep(x_{τ_k}, dx/dt, τ_k, τ_{k+1})。
  • 对于离散状态，这个“更新”不是简单的加法。一种常见的方法是计算一个转移概率矩阵P = expm(Δt * Q_rev_θ)，然后根据这个矩阵对每个位置的状态进行采样或取期望（贪心选择概率最大的token）。expm是矩阵指数，对于均匀吸收等简单设计，它有闭式解。

4.2 实用采样步骤与调参

在实际实现中，为了平衡速度和质量，我们通常采用以下步骤：

1. 设定采样步数：例如N=50或100。步数越多，采样越精确，但耗时越长。
2. 时间离散化：将连续时间区间[0, T]划分为N步。可以使用线性间隔，或者更优的、根据噪声调度β(t)调整的非线性间隔（在β(t)大的区域步长更小）。
3. 循环采样：

def deterministic_sample(model, seq_len, steps=50): # 初始化：全部为[MASK] x = torch.full((1, seq_len), MASK_TOKEN_ID).to(device) # 创建时间步 timesteps = torch.linspace(T, 0, steps+1) # 从T到0 for i in range(steps): t = timesteps[i] dt = timesteps[i] - timesteps[i+1] # 负的时间步长 # 1. 获取模型预测的得分/logits with torch.no_grad(): logits = model(x, t.unsqueeze(0)) # (1, L, V) # 2. 根据前向Q_fwd和模型logits构造反向转移矩阵P # 对于均匀吸收，P有一个简洁形式： # P(keep) = exp(-beta_int) , P(switch to j) = (1-exp(-beta_int)) * softmax(logits)[j] # 其中 beta_int 是速率β从t到t+dt的积分 beta_int = integral_of_beta(t, dt) p_keep = torch.exp(-beta_int) p_change_dist = F.softmax(logits, dim=-1) # 模型预测的分布 # 3. 计算每个位置的新状态分布 # 首先，构建一个 (L, V) 的分布矩阵 # 对于当前位置是token k，新分布 = p_keep * one_hot(k) + (1-p_keep) * p_change_dist # 这里需要一点张量操作 eye = torch.eye(V).to(x.device) # (V, V) current_one_hot = eye[x] # (1, L, V) transition_dist = p_keep.unsqueeze(-1) * current_one_hot + (1-p_keep.unsqueeze(-1)) * p_change_dist # 4. 贪心采样：选择每个位置概率最大的token x = transition_dist.argmax(dim=-1) return x

4. 温度调节：在从transition_dist采样时，可以对logits除以一个温度参数τ。τ < 1会使分布更尖锐（确定性更强，多样性降低），τ > 1会使分布更平滑（多样性增加，但可能影响质量）。

实操心得：采样步数N是一个关键的超参数。开始时可以用较少的步数（如20）快速测试采样效果。如果发现生成结果不连贯或重复，尝试增加到50或100。另外，在最后几步（t接近0时），可以尝试减小步长，以获得更精细的调整。

5. 实战挑战与性能优化策略

将理论转化为实践的路上，你会遇到几个典型的挑战。以下是我在复现和实验过程中总结出的问题和解决方案。

5.1 训练不稳定的常见原因与对策

5.2 效率瓶颈分析与优化

离散扩散模型在序列长度L和词汇表大小V较大时，计算和内存开销可能成为瓶颈。

• 计算复杂度：核心操作是模型前向传播和Softmax over Vocabulary。Transformer的复杂度是O(L^2 * d_model)，而最后的Logits矩阵是O(L * V * d_model)。当V很大（如数万）时，最后的线性输出层是主要瓶颈。

  • 优化策略：
    1. 词汇表裁剪：对于特定领域任务，可以考虑使用子词词元（如BPE）或领域相关的精简词汇表。
    2. 自适应Softmax或采样Softmax：在训练时，对于每个批次，只计算目标token所在类别的梯度，可以大幅减少计算量。PyTorch的F.cross_entropy已经非常优化，但对于极大的V，可能需要更高级的技巧。
    3. 模型蒸馏：训练一个大模型后，将其知识蒸馏到一个更小、更高效的模型中进行部署采样。

• 内存占用：存储(B, L, V)的Logits张量会消耗大量显存。

  • 优化策略：使用混合精度训练（AMP），可以显著减少显存占用并加速计算。在PyTorch中，只需几行代码即可启用。

5.3 高级技巧：条件生成与引导

基础的模型是无条件生成。如何实现像“生成一段关于夏天的文字”这样的条件生成？

1. Classifier-Free Guidance：这是目前最流行且有效的方法。在训练时，随机以一定概率（如10%）将条件信息（如类别标签、文本描述）置空。这样模型同时学会了无条件生成p(x)和条件生成p(x|c)。在采样时，使用以下公式进行引导：logits_guided = logits_uncond + guidance_scale * (logits_cond - logits_uncond)其中guidance_scale > 1。这个技巧能显著提升生成样本与条件的对齐质量，但过大的guidance_scale会损害多样性。

2. 后处理编辑：对于文本，可以先使用模型生成一个草案，然后利用一个判别式模型（如情感分类器、语法检查器）对生成结果进行打分，并通过MCMC采样等方法，迭代地优化草案，使其满足特定条件。这种方法更灵活，但计算成本更高。

一个重要的提醒：离散扩散模型的训练和采样，在概念上比自回归模型更复杂，调试周期可能更长。建议从一个极小的数据集（如一个简单的文本数据集）和一个小模型开始，确保前向破坏、损失计算、反向采样整个流程的每个环节都符合预期，再逐步扩展到更大规模的任务上。理解每个张量的形状和含义，是成功复现的关键。