流匹配与扩散模型-07：离散扩散模型

离散动力学

前面的流模型和扩散模型都在欧几里得空间 $\mathbb R^d$ 中演化，向量可以沿任意方向移动，速度或分数描述一次无穷小位移。语言、DNA 等数据由离散符号组成。设词表为

$$\mathcal V=\{v_1,\ldots,v_V\},$$

长度为 $d$ 的序列属于状态空间

$$\mathcal S=\mathcal V^d, \qquad |\mathcal S|=V^d.$$

两个 token 之间没有天然的距离和中间值，把“猫”向“狗”移动 $0.01$ 没有定义，因而不能在 $\mathcal S$ 上直接照搬普通 ODE、布朗运动或 SDE。不过，单条序列虽然只能突然换 token，整个群体落在各序列上的概率却可以随时间平滑变化。离散扩散因此改用连续时间马尔可夫链 Continuous-Time Markov Chain：时间 $t$ 连续变化，状态 $X_t$ 在一段时间内保持不变，只在随机时刻跳到另一个离散状态。连续模型追踪样本在空间中的速度，CTMC 则追踪概率质量从一个状态流向另一个状态的速度。

CTMC 满足马尔可夫性。对任意过去时刻 $t_1<\cdots<t_k<t$ 和 $h>0$，

$$p(X_{t+h}\mid X_t,X_{t_1},\ldots,X_{t_k}) =p(X_{t+h}\mid X_t).$$

马尔可夫性并不否认文本中的长程依赖；它只要求在决定下一次跳变时，当前完整序列 $X_t$ 已经概括所需历史。将“过程记忆”限制在当前状态后，模型无需记录一条序列此前经历过哪些掩码和修改，采样器也只需把当前序列交给网络。Transformer 仍可读取 $X_t$ 的全部位置，利用双向和长距离上下文决定每个位置下一步如何变化。

速率矩阵

CTMC 需要一个与连续向量场相对应的局部规则。由于状态本身不可求导，这个规则改为描述短时间内各种跳转概率增长得多快，即速率矩阵 $Q_t$。采用讲义中的记号，$Q_t(y\mid x)$ 表示从源状态 $x$ 跳向目标状态 $y$ 的瞬时速率，并满足

$$Q_t(y\mid x)\ge0,\qquad y\ne x,$$

$$Q_t(x\mid x)=-\sum_{y\ne x}Q_t(y\mid x).$$

非对角元是离开 $x$ 后跳向各目标的速率，必须非负；对角元是所有离开速率的负和，使固定源状态 $x$ 对应的一列总和为零。若没有这项，其他状态增加的概率没有从 $x$ 中扣除，总概率便会凭空增长。它只是概率守恒的记账项，并非“留在原状态的负概率”。速率也可以大于 $1$，因为它表示单位时间的跳变强度；只有乘上足够短的时间后，才成为 $[0,1]$ 内的概率。

更精确地说，速率是短时转移概率在 $h=0$ 处的导数：

$$\left.\frac{\mathrm d}{\mathrm dh} p_{t+h\mid t}(y\mid x)\right|_{h=0} =Q_t(y\mid x).$$

这一定义把不可微的状态变化转成可微的转移概率变化。保留关于 $h$ 的一阶项，一个很短的 Euler 时间步可以近似为

$$\widetilde p_{t+h\mid t}(y\mid x) =\mathbf 1\{y=x\}+hQ_t(y\mid x).$$

当 $y\ne x$ 时，跳转概率约为 $hQ_t(y\mid x)$；留在 $x$ 的概率为 $1-h\sum_{y\ne x}Q_t(y\mid x)$，所有选项恰好归一化。步长必须足够小，保证留下概率非负；速率越大，允许的 Euler 步越短。采样就是从 $X_0\sim p_{\mathrm{init}}$ 出发，反复把速率转成分类概率并抽取 $X_{t+h}$。较小的 $h$ 更接近连续时间过程，却会增加网络前向次数，这与连续扩散中数值精度和采样成本的取舍相同。

完整速率矩阵要为 $V^d$ 个源状态分别列出 $V^d$ 个目标，甚至无法显式存储。若把一次跳转限制为只修改一个位置，从任意序列出发便只需考虑大约 $dV$ 个相邻状态。实际模型据此使用因子化速率矩阵：

$$Q_t(y\mid x)=0, \qquad \text{若 }x\text{ 与 }y\text{ 在两个或更多位置上不同}.$$

此时网络只需为每个位置 $j$ 输出换成词表中各 token 的速率 $q_j(v)$，输出规模从指数级降为 $d\times V$，与普通语言模型的 logits 形状一致。这里的“因子化”只约束一次无穷小跳变的邻接结构，不代表各 token 在语义上独立；每个 $q_j(v)$ 仍由读取完整序列的 Transformer 计算，因此一个位置的候选词可以取决于所有其他位置。代价是相差许多 token 的两个序列不能一步直接互换，需要经由若干单位置跳变连接起来。

若将 $[0,1]$ 划分为 $n$ 步，$h=1/n$，每一步对位置 $j$ 使用

$$\widetilde p_{j,t}(v\mid x_j)= \begin{cases} hq_j(v), & v\ne x_j,\\[4pt] 1-h\displaystyle\sum_{v'\ne x_j}q_j(v'), & v=x_j, \end{cases}$$

并行抽取所有位置的新 token。这样做利用了加速器擅长的整段并行计算，却似乎违反“一次只改一个位置”的定义。两个指定位置在同一步同时变化，需要两个各为 $O(h)$ 的事件共同发生，概率只有 $O(h^2)$；Euler 法本来就只保留一阶项，因此忽略这类事件仍与原 CTMC 一阶一致。减小步长会让多位置同时改变更少，但会增加 Transformer 前向次数。

离散路径

仅有 CTMC 形式还不知道哪些速率能把噪声变成语言。离散扩散沿用流匹配的办法，先人为设计一条从简单噪声分布到数据分布的概率路径，再反推出实现这条路径的速率。路径不仅决定训练样本如何被破坏，也决定生成时信息以何种节奏恢复。对训练样本 $z\sim p_{\mathrm{data}}$，条件路径满足

$$p_0(\cdot\mid z)=p_{\mathrm{init}}, \qquad p_1(\cdot\mid z)=\delta_z.$$

已知 $z$ 时，构造从噪声到这个终点的路径很容易；真正生成时并不知道终点，因此还要对所有可能的 $z$ 边缘化：

$$p_t(x)=\sum_{z\in\mathcal S}p_t(x\mid z)p_{\mathrm{data}}(z),$$

所以 $p_0=p_{\mathrm{init}}$、$p_1=p_{\mathrm{data}}$。

讲义采用因子化混合路径，因为它能在任意时刻直接从干净序列构造带噪样本，不必先模拟从 $0$ 到 $t$ 的所有跳变。设初始分布按位置分解为

$$p_{\mathrm{init}}(x)=\prod_{j=1}^d p_{\mathrm{init}}^{(j)}(x_j),$$

调度函数满足 $\kappa_0=0$、$\kappa_1=1$、$\dot\kappa_t\ge0$。条件路径定义为

$$p_t(x\mid z)= \prod_{j=1}^d \left[ (1-\kappa_t)p_{\mathrm{init}}^{(j)}(x_j) +\kappa_t\delta_{z_j}(x_j) \right].$$

从中采样无需枚举状态。对每个位置独立抽取

$$m_j\sim\mathrm{Bernoulli}(\kappa_t), \qquad \xi_j\sim p_{\mathrm{init}}^{(j)},$$

再令

$$x_j=m_jz_j+(1-m_j)\xi_j.$$

最后一式表示在 $z_j$ 与噪声 token $\xi_j$ 之间做选择，并非对 token 编号做数值运算。$\kappa_t$ 可以直接读作时刻 $t$ 一个位置已经采用干净 token 的概率：$t=0$ 时所有信息来自噪声，随着 $\kappa_t$ 增大，干净位置的期望比例随之提高，$t=1$ 时序列完全等于 $z$。训练时只要随机选择 $t$ 并抛 $d$ 次伯努利硬币，就能得到正确分布的 $x$；这项可直接采样的性质是后面免模拟训练的基础。

高斯路径会把概率质量沿空间方向连续搬运，离散混合路径没有可用的移动方向，只能逐渐降低噪声状态的权重并提高终点状态的权重。调度 $\kappa_t$ 决定信息何时出现：增长太集中会让速率在一小段时间内很大、数值采样困难，增长较平缓则把去噪任务均匀分散到更多时间步。图中上排是一条已知终点的条件路径，下排是所有终点混合后的边缘路径；后者才是模型最终需要复现的数据分布演化。

速率匹配

下一步要把静态的概率路径变成可执行的跳转规则。先构造条件速率 $Q_t^z$，使知道终点 $z$ 时，CTMC 的时刻 $t$ 分布恰好为 $p_t(\cdot\mid z)$。生成时 $z$ 未知，不能选定其中任何一套条件速率；可行的规则是根据当前 $x$ 判断各终点还有多可信，再对相应速率加权，即离散边缘化：

$$Q_t(y\mid x) =\sum_{z\in\mathcal S}Q_t^z(y\mid x) p_{1\mid t}(z\mid x),$$

其中

$$p_{1\mid t}(z\mid x) =\frac{p_t(x\mid z)p_{\mathrm{data}}(z)}{p_t(x)}$$

是观察到当前噪声状态 $x$ 后，干净终点为 $z$ 的后验概率。若某个终点几乎不可能产生当前 $x$，它的条件速率就几乎不参与决策；多个终点都合理时，边缘速率会保留这种不确定性，而不是过早选定一句文本。这个公式与连续流匹配中的后验平均完全平行：条件问题有现成监督，后验平均把它转换成无需知道真实终点的生成动力学。

CTMC 的概率质量由 Kolmogorov 前向方程控制：

$$\frac{\mathrm d}{\mathrm dt}p_t(x) =\sum_{y\in\mathcal S}Q_t(x\mid y)p_t(y).$$

式中 $Q_t(x\mid y)p_t(y)$ 是从 $y$ 流向 $x$ 的概率流；当 $y=x$ 时，对角元扣除从 $x$ 流向其他状态的部分，因此右侧正好是流入减流出。它在离散空间中承担与连续性方程相同的职责：检查一组局部速度是否真的实现给定的全局概率路径。把条件路径的前向方程对 $z$ 求和，再代入后验加权公式，就得到边缘路径的前向方程，说明后验平均不是经验混合，而会严格保持目标边缘路径。

对于因子化混合路径，位置 $j$ 的条件速率有闭式解：

$$Q_t^z(v_i,j\mid x_j) =\frac{\dot\kappa_t}{1-\kappa_t} \left[ \delta_{z_j}(v_i)-\delta_{x_j}(v_i) \right].$$

若 $x_j\ne z_j$，它只允许当前位置跳到目标 token $z_j$；若 $x_j=z_j$，两项完全抵消，该位置不再跳变。系数 $\dot\kappa_t/(1-\kappa_t)$ 可以从“仍未恢复的位置”理解：时刻 $t$ 尚有比例 $1-\kappa_t$ 的位置保持噪声，为了让下一瞬间恢复比例增加 $\dot\kappa_t\,dt$，每个未恢复位置需要以条件概率约 $\dot\kappa_t\,dt/(1-\kappa_t)$ 跳转。它正是保证路径按指定 $\kappa_t$ 前进的瞬时风险率。对未知终点取后验平均，得到边缘速率

$$Q_t(v_i,j\mid x) =\frac{\dot\kappa_t}{1-\kappa_t} \left[ p_{1\mid t}(z_j=v_i\mid x)-\delta_{x_j}(v_i) \right].$$

这个式子避开了对指数规模速率矩阵的直接监督。网络只需回答熟悉的去噪问题：“看到当前整段 $x$ 后，第 $j$ 个干净 token 最可能是什么？”速率公式再把答案乘上路径规定的时间尺度。对 $v_i\ne x_j$，后验越大，下一小步跳到 $v_i$ 的概率越高；减去 $\delta_{x_j}(v_i)$ 后，当前 token 对应项自动成为所有离开率的负和。随着 $\kappa_t\to1$，待恢复的概率 $1-\kappa_t$ 趋于零，剩余错误必须在越来越短的时间内解决，风险率因而可能发散。实现中不会在 $t=1$ 直接计算它，而会使用终点极限、截断时间或稳定的离散调度，并保证每个 Euler 步仍给出合法概率。

训练与生成

用 Transformer 参数化后验网络

$$p_{1\mid t}^{\theta}(z_j=v_i\mid x), \qquad j=1,\ldots,d,\quad i=1,\ldots,V.$$

网络读取完整噪声序列 $x$ 和时间 $t$，为每个位置输出 $V$ 个 logits，经 Softmax 得到干净 token 的分类分布。真实后验本身需要对整个数据分布求和，无法显式计算，但可以从训练数据产生监督对 $(x,z)$。对这些样本使用逐 token 交叉熵：

$$\mathcal L_{\mathrm{DFM}}(\theta) =\mathbb E_{\substack{z\sim p_{\mathrm{data}},\ t\sim\mathrm{Unif}[0,1]\\x\sim p_t(\cdot\mid z)}} \left[ -\sum_{j=1}^d \log p_{1\mid t}^{\theta}(z_j\mid x) \right].$$

对固定的 $(x,t)$，交叉熵的期望在模型分布等于真实 $p_{1\mid t}(z_j\mid x)$ 时最小，因此它学到的并非某个任意分类器，恰好就是边缘速率公式缺少的后验。一次训练迭代包含四步：抽取干净序列 $z$ 和时间 $t$；按 $\kappa_t$ 独立保留或破坏各位置，构造 $x$；让 Transformer 根据 $(x,t)$ 预测每个原始 token；计算交叉熵并更新参数。由于条件路径允许直接跳到任意时刻，训练无需运行从噪声到数据的 CTMC，随机时刻的监督也能覆盖整条路径。

推理时没有干净的 $z$ 可供直接破坏，才需要把后验预测代入边缘速率公式：从 $X_0\sim p_{\mathrm{init}}$ 开始，在每个时间步计算 $p_{1\mid t}^{\theta}$ 与 $Q_t^\theta$，再按逐位置 Euler 概率并行抽取新 token，直到 $t=1$。训练解决“给定一份由真实文本破坏出的序列，原文可能是什么”，速率公式把这个静态答案转换成“为遵守既定概率路径，下一小步应以多大概率改成什么”。正因为训练有终点而生成没有终点，前者可以单步直接监督，后者必须多步更新并在新上下文下反复修正预测。

掩码扩散语言模型 Masked Diffusion Language Model 是一个常用特例。扩展词表并加入 $[\mathrm{MASK}]$，令

$$p_{\mathrm{init}}=\delta_{[\mathrm{MASK}]^d}.$$

生成从全掩码序列开始，多个位置在连续时间网格上逐渐恢复。掩码使“噪声 token”与真实词明确区分，未恢复位置不会给模型伪造的词义；网络可同时观察已经恢复的左右文和仍待填充的位置，并在同一步为全部掩码位置预测分布。简单的吸收式掩码路径一旦揭示 token 就不再修改它，换来特别简单的破坏过程和稳定目标；使用一般替换噪声与更复杂速率时，则可允许后续修订。

自回归语言模型使用 $p(z)=\prod_jp(z_j\mid z_{<j})$，把联合分布化成有明确顺序的一系列条件概率，训练目标和似然计算直接，推理时却必须等待前一个 token，且通常只能利用左侧已生成内容。离散扩散不固定词序，适合双向补全、长度内并行决策和全局约束，但每轮都要重新处理整段序列，并需要选择时间调度和步数。一次扩散前向能更新许多位置，并不等于采样必然更快；自回归模型可使用 KV cache，而扩散需要多轮全局前向，实际速度取决于序列长度、扩散步数、缓存策略和硬件。因子化 CTMC 也不限于吸收式掩码：若 $p_{\mathrm{init}}$ 是一般替换噪声，已生成 token 可以在后续步骤中被修正，但训练和采样规则也会更复杂。

从更统一的角度看，连续流中的向量场与 CTMC 中的速率矩阵都是马尔可夫过程的无穷小生成元：它们都说明当前分布在下一瞬间怎样变化。两类模型因而遵循同一条训练主线：选择能在任意时刻直接采样的条件概率路径，推导已知终点时容易计算的条件生成元，再用终点后验将其边缘化。真正难算的边缘动力学由一个容易采样的监督任务间接学得，训练时不必模拟完整生成过程。连续空间通常回归速度或分数，离散空间预测 token 后验并换算成跳转速率；状态空间和数值采样器不同，路径设计、后验平均与免模拟训练的逻辑保持一致。