前言

MIT6.S184: Generative AI with Stochastic Differential Equations 最新 2026 年课程笔记 An Introduction to Flow Matching and Diffusion Models 翻译，本篇文章翻译第七章节 Discrete Diffusion Models: Building Language Models with Diffusion 相关内容🤗。

Course Notes：https://diffusion.csail.mit.edu/2026/docs/lecture_notes.pdf

Course Website：https://diffusion.csail.mit.edu/2026/index.html

7. Discrete Diffusion Models: Building Language Models with Diffusion

在前面的章节中，我们研究了 flow 与 diffusion models，它们作为定义在欧几里得空间 ${\mathbb{R}}^d$ 上的生成模型，能够生成表示为向量 $z\in {\mathbb{R}}^d$ 的数据点。然而，并不是所有数据都天然适合建模为欧几里得空间 ${\mathbb{R}}^d$ 中的点。许多数据类型，例如 text 或 DNA，更自然地应被视为离散状态空间 $S$ 中的元素。最重要的是，language 由离散 token 序列构成，而这正是我们希望建模的对象。

那么，我们该如何将 flow 与 diffusion models 应用于这种数据类型呢？事实证明，我们在前面章节中学习的原理，同样能够扩展到这些数据类型。由此得到的模型，在机器学习文献中被称为 discrete diffusion models（离散扩散模型）[5, 16]。

不过，需要牢记的一点是：这里并不存在数学意义上的 diffusion process（因为在离散状态空间中不存在 SDEs）。因此，我们不再使用 ODEs / SDEs，而是使用：连续时间马尔可夫链（continuous-time Markov chain, CTMCs）。

在接下来的内容中，我们将介绍 CTMC models（见 Section 7.1）以及如何学习这些模型（见 Section 7.2），并展示如何利用 flow 与 diffusion models 的原理构建 large language models（LLMs）。

7.1 Continuous-Time Markov chain (CTMC) models

在本节中，我们将介绍 continuous-time Markov chains（CTMCs）。你可以将 CTMCs 看作是 SDEs 的离散版本，我们可以利用它来构建生成离散空间的神经网络模型。此外，我们还将介绍 CTMC models，即利用 CTMCs 来生成离散序列（例如文本）的神经网络模型。

Figure 17：一个 CTMC 轨迹的示意图，其状态空间为 S = { S 1 , S 2 , S 3 } S=\{S_1,S_2,S_3\} S={S1​,S2​,S3​} ，并且序列长度 $d=1$ 。图片改编自文献 [5]。

现在，我们首先来刻画状态空间 $S$ 。设 V = { v 1 , ⋯   , v V } \mathcal{V} = \{v_1, \cdots, v_V\} V={v1​,⋯,vV​} 为我们的 vocabulary（词汇表），则状态空间定义为 $S={\mathcal{V}}^d$ ，其中 $d\in \mathbb{N}$ 表示序列长度（sequence length）， $V\in \mathbb{N}$ 表示词汇表大小（vocabulary size）。

对于 language， { v 1 , ⋯   , v V } \{v_1, \cdots, v_V\} {v1​,⋯,vV​} 可以表示 alphabet 或一组离散 tokens。而 $S$ 则表示所有长度为 $d$ 的 sequences（或 sentences）的集合。对于 DNA， { v 1 , ⋯   , v V } \{v_1, \cdots, v_V\} {v1​,⋯,vV​} 则可以表示全部 4 种 DNA bases，而 $S$ 表示所有长度为 $d$ 的 DNA sequences。

接下来，令 $X_t$ 为定义在 $S$ 上的随机过程，即一个随机轨迹： $X:[0,1]\to S,t\mapsto X_t\in S$ 。我们要求 $X_t$ 是一个马尔可夫过程，即一个 “无记忆（no memory）” 的过程。具体而言，这意味着下面的条件成立：

$$\underset{\text{prob. of future given present and past}}{\underbrace{p(X_{t+h}\mid X_t,X_{t_1},\cdots ,X_{t_k})}}=\underset{\text{prob. of future given present}}{\underbrace{p(X_{t+h}\mid X_t)}}(\text{for all }0<h,0\le t_1<t_2<\cdots <t_k<t)$$

换句话说，未来事件的概率只依赖于当前状态，过去对于未来已经不再重要。注意，虽然 ODEs/SDEs 并不是定义在离散状态空间上，但它们同样也是马尔可夫过程。

这里，由于 $X_t$ 定义在离散空间上，因此它被称为 Markov chain（马尔可夫链），更具体地说，是 Continuous-time Markov chain（CTMC）。其中 quantity $p_{t+h\mid t}(X_{t+h}\mid X_t)$ 被称为 transition probabilities（转移概率）。它们与马尔可夫链的初始分布 $X_0\sim p_0$ 一起完全决定了整个 CTMC。因此，当我们说 CTMC 时，你也可以简单理解为转移概率 $p_{t+h\mid t}(X_{t+h}\mid X_t)$ 。

接下来，我们来推导离散场景下 vector field（向量场）的对应形式。由于现在处于离散场景中，我们只能在 states（状态）之间进行 jump（或 switch），而无法像 ODE 中那样沿着某个连续方向移动。因此，我们定义一个 rate matrix（速率矩阵）$Q_t(y\mid x)$ ，它用于刻画从状态 $x\in S$ 跳转到状态 $y\in S$ 的速率。

形式化地，rate matrix $Q_t$ 是如下有界函数（关于时间连续）：

$$Q:S\times S\times [0,1]\to \mathbb{R},(x,y,t)\mapsto Q_t(y\mid x)\text{\hspace{1em}(84)}$$

其中， $Q_t(y\mid x)$ 描述从 $x$ 切换到 $y$ 的速率，并满足：

$$\begin{array}{rl}(1)\text{Outgoing rates are positives:}{Q}_t(y\mid x)& \ge 0\text{whenever }x\ne y\\ (2)\text{Rate staying equals negative outgoing rate:}{Q}_t(x\mid x)& =-\sum _{y\ne x}{Q}_t(y\mid x)\text{for all }x\end{array}\text{\hspace{1em}(85)(86)}$$

这两个条件是直观的：第一个条件表示从 $x$ 切换到不同状态 $y\ne x$ 的速率只能是非负的（不发生切换对应于 0，因此速率小于 0 是没有意义的）。第二个条件表示停留在 $x$ 的速率 $Q_t(x\mid x)$ 必须与离开 $x$ 的总速率相抵消。它本质上是一个一致性条件，表示你要么停留在 $x$ ，要么离开它，不存在第三种情况。

注意，这些条件特别意味着 $Q_t(x\mid x)\le 0$ 。因此， $Q_t(y\mid x)$ 可以看作一个矩阵，其中对角线元素全部非正，非对角线元素全部非负。

现在，我们可以定义微分方程在离散场景中的对应形式，即要求一个 CTMC “follow” 某个 rate matrix 的条件。其核心思想是 $X$ 的分布或过程应当遵循 rate matrix $Q_t$ 。换句话说，我们要求转移概率满足：

$${\frac{d}{dh}{p}_{t+h\mid t}(X_{t+h}=y\mid X_t=x)|}_{h=0}=Q_t(y\mid x)\text{for all }x,y\in S,0\le t\text{\hspace{1em}(87)}$$

左边表示从 $x$ 切换到 $y$ 的概率的无穷小的变化率。我们要求这些概率按照 rate matrix 所指定的方式变化。

现在，我们简单检查一下这样的条件是否合理。也就是说，如果我们像 公式 (87) 中那样定义 $Q_t(y\mid x)$ ，它是否真的是一个合法的 rate matrix？

当 $h=0$ 时，由于时间尚未流逝，从 $x$ 切换到 $y\ne x$ 的概率为 0，即 $p_{t\mid t}(y\mid x)=0\text{for all }y\ne x$ 。因此，我们知道其导数必须是非负的，从而 $Q_t(y\mid x)\ge 0\text{whenever }y\ne x$ ，这验证了 公式 (85) 中的第一个条件。

进一步地，我们有：

$$\begin{array}{rl}\sum _{y\ne x}{Q}_t(y\mid x)=\sum _{y\ne x}{\frac{d}{dh}p(X_{t+h}=y\mid X_t=x)|}_{h=0}={\frac{d}{dh}\sum _{y\ne x}p(X_{t+h}=y\mid X_t=x)|}_{h=0}& =\frac{d}{dh}(1-p(X_{t+h}=x\mid X_t=x))\\ & =-Q_t(x\mid x)\end{array}$$

这里我们使用了概率之和等于 1。这说明 公式 (86) 成立。因此，我们验证了每一个 CTMC 至少都存在一个满足 公式 (87) 的 rate matrix。但如果反过来呢？也就是说如果我们给定一个 $Q_t$ ，是否一定存在对应的 CTMC？如果存在，它是否唯一？答案是：确实如此。

---

Theorem 33 (CTMC existence and uniqueness)

对于任意的 rate matrix $Q_t$（关于时间 $t$ 有界且连续），都存在唯一的马尔可夫链 $X_t$（即唯一的一组转移概率 $p_{t+h\mid t}(y\mid x)$）使得 公式 (87) 成立。

---

对于感兴趣的读者，我们在 Section C 中提供了一个自洽的证明。这个定理最重要的结论是，对于机器学习的目的，我们可以直接构造一个 rate matrix $Q_t$（例如通过神经网络），并假设存在唯一一个与 $Q_t$ 对应的马尔可夫链。

---

Example 34 (Two-state CTMC with equal jump rates)

设 S = { a , b } S=\{a,b\} S={a,b} ，并考虑一个时间齐次 CTMC $(X_t{)}_{t\ge 0}$ ，它以常数速率 $\lambda >0$ 在两个状态之间切换：

$$Q=\begin{array}{ccc}& a& b\\ a& -\lambda & \lambda \\ b& \lambda & -\lambda \end{array}.$$

则，在时间增量 $h\ge 0$ 上的转移概率同样与时间 $t$ 无关，并由下式给出：

$$\left(\begin{array}{cc}p(X_{t+h}=a\mid X_t=a)& p(X_{t+h}=a\mid X_t=b)\\ p(X_{t+h}=b\mid X_t=a)& p(X_{t+h}=b\mid X_t=b)\end{array}\right)=\frac{1}{2}\left(\begin{array}{cc}1+e^{-2\lambda h}& 1-e^{-2\lambda h}\\ 1-e^{-2\lambda h}& 1+e^{-2\lambda h}\end{array}\right).$$

可以手动验证 公式 (87) 成立，即这些转移概率的确对应于该 rate matrix。实际上，这些速率是非常直观的：该 chain 会以瞬时速率 $\lambda$ 不断在两个状态之间翻转。指数项 $e^{-2\lambda h}$ 描述了初始状态记忆的衰减过程。随着时间趋于无穷即 $h\to \infty$ ，有：

$$P(h)\to \left(\begin{array}{cc}\frac{1}{2}& \frac{1}{2}\\ \frac{1}{2}& \frac{1}{2}\end{array}\right),$$

因此，该 chain 最终会忘记它最初从哪里开始，并且以概率 $\frac{1}{2}$ 处于 $a$ 或 $b$ 。并且，切换速率 $\lambda >0$ 越大，这种收敛发生得越快。

---

Simulation of CTMC.

接下来，我们考虑如何模拟一个 CTMC 的轨迹。设 $h>0$ 为步长， $p_{\mathrm{init}}$ 为定义在 $S$ 上的初始分布。例如 $p_{\mathrm{init}}={\mathrm{Unif}}_S$ 表示 $S$ 上的均匀分布。

随后，我们可以通过如下方式迭代地进行模拟：首先采样 $X_0\sim p_{\mathrm{init}}$ ，然后令：

$$X_{t+h}\sim p_{t+h\mid t}(\cdot \mid X_t).$$

现在，如果我们知道 $p_{t+h\mid t}(\cdot \mid X_t)$ ，那么上述过程当然是可行的。然而，除了最简单的 CTMC 之外，我们通常并不知道封闭式的转移核，而只能访问 rate matrix $Q_t$ 。不过，根据 公式 (87)：

$$p_{t+h\mid t}(X_{t+h}=y\mid X_t=x)=p_{t\mid t}(X_t=y\mid X_t=x)+h{Q}_t(y\mid x)+R_t(h)=1_{y=x}+h{Q}_t(y\mid x)+R_t(h)$$

其中 $R_t(h)$ 是一个误差项，当 $h$ 足够小时可以忽略。因此，对于较小的 $h$ ，我们可以令：

$$p_{t+h\mid t}(X_{t+h}=y\mid X_t=x)\approx 1_{y=x}+h{Q}_t(y\mid x)=:{\tilde{p}}_{t+h\mid t}(y\mid x)$$

可以验证：由于我们对 rate matrix 施加的条件，当 $h$ 足够小时， ${\tilde{p}}_{t+h\mid t}(y\mid x)$ 确实是一个合法的概率分布。因此，我们可以近似地通过如下方式采样下一个点：

$$X_{t+h}\sim {\tilde{p}}_{t+h\mid t}(\cdot \mid x)=(1_{y=x}+h{Q}_t(y\mid x){)}_{y\in S}\text{\hspace{1em}(88)}$$

由于上述只是一个离散分布，因此我们可以使用标准方法轻松进行采样。这提供了一种简单的 CTMC 模拟方法。

CTMC model.

接下来，我们定义如何用神经网络来参数化一个 CTMC。一个 CTMC model（或 discrete diffusion model）由以下部分给出：$S$ 上的初始分布 $p_{\mathrm{init}}$ ，一个带参数 $\theta$ 的神经网络 $Q_t^{\theta }$ ，使得对于每个输入 $x\in S$ 模型返回 rate matrix 的单独一列：

x ↦ { Q t θ ( y ∣ x ) } y ∈ S x \mapsto \{Q_t^\theta(y|x)\}_{y\in S} x↦{Qtθ​(y∣x)}y∈S​

我们希望模型返回完整的一列，因为在 CTMC 的模拟中需要它（公式 (88)），即采样下一个状态。

上述模型的一个复杂之处在于空间 $S$ 可能非常大。特别地 $|S|=V^d$ ，其中 $V$ 是词汇表大小， $d$ 是序列长度。这种指数增长使得在内存中存储 rate matrix 的完整一列基本不可能，也就是说 { Q t θ ( y ∣ x ) } y ∈ S \{Q_t^\theta(y|x)\}_{y\in S} {Qtθ​(y∣x)}y∈S​ 不可能在计算机中表示。

因此，我们必须对模型施加约束。具体来说，几乎所有 CTMC models 都是 factorized（因式化）的（见 Figure 18），这实际上是一种稀疏性约束。具体而言，一个 factorized CTMC model 由一个 CTMC model $Q_t^{\theta }$ 给出，使得对于所有 $y=(y_1,\cdots ,y_d),x=(x_1,\cdots ,x_d)\in S={\mathcal{V}}^d$ 都有：

$$Q_t^{\theta }(y\mid x)=0\text{whenever }{y}_i\ne x_i\text{ for more than one position }i$$

我们将所有与 $x$ 最多只相差一个 token 的 $y$ 称为 $x$ 的邻居 $N(x)$ 。我们可以将这样的 factorized CTMC model 写为：

x ↦ { Q t θ ( y ∣ x ) } y ∈ N ( x ) = ( Q t θ ( v 1 , 1 ∣ x ) ⋯ Q t θ ( v V , 1 ∣ x ) ⋮ ⋮ Q t θ ( v 1 , d ∣ x ) ⋯ Q t θ ( v V , d ∣ x ) ) x \mapsto \{Q_t^\theta(y|x)\}_{y\in N(x)} = \begin{pmatrix} Q_t^\theta(v_1,1|x) & \cdots & Q_t^\theta(v_V,1|x)\\ \vdots & & \vdots\\ Q_t^\theta(v_1,d|x) & \cdots & Q_t^\theta(v_V,d|x) \end{pmatrix} x↦{Qtθ​(y∣x)}y∈N(x)​= ​Qtθ​(v1​,1∣x)⋮Qtθ​(v1​,d∣x)​⋯⋯​Qtθ​(vV​,1∣x)⋮Qtθ​(vV​,d∣x)​ ​

其中 $Q_t(y|x)=Q_t^{\theta }(v_i,j|x)$ 现在表示从 $x=(x_1,\cdots ,x_d)$ 转移到 $x$ 的某个 neighbor 的速率。该 neighbor 是通过将第 $j$ 个元素替换为 $v_i$ 得到的 $y=(x_1,\cdots ,x_{j-1},v_i,x_{j+1},\cdots ,x_d)$ 。每一行对应于每个位置 $i=1,\cdots ,d$ 上的一个 rate matrix，即我们要求：

$$Q_t^{\theta }(v,i|x)\ge 0\text{if }v\ne x_i,Q_t(x_i,i|x)=-\sum _{v\ne x_i}{Q}_t^{\theta }(v,i|x)$$

我们可以很容易地对神经网络的输出施加这些条件。例如，可以使用一个作用在序列长度 $d$ 上的 transformer model，其输出维度为 $V$ 。还需要注意 factorized rate matrix 使输出形状变为 $d\times V$ ，该大小随着维度线性增长，而不是指数增长。

Figure 18：factorized CTMC model 的示意图。Factorized CTMC 只有在起点与终点仅相差一个维度时，其 rate 才是非零的即 $Q_t(y|x)\ne 0$ ，这里的示例中 $d=2$ 。图片取自文献 [26]。

Simulating a CTMC model.

为了从一个 CTMC model 中采样，我们先采样 $X_0\sim p_{\mathrm{init}}$ ，然后执行迭代，在每一步根据 公式 (88) 对下一个状态进行采样。我们在 Algorithm 7 中给出了该算法。

如上所示，对于 factorized CTMC models，可以使用一种并行的 per-token 欧拉近似，其中每个 token 都会在一个较小步长 $h>0$ 下被独立更新。这种近似在关于 $h$ 的一阶项上与完整的 CTMC 欧拉步一致，但它允许出现一个 $O(h^2)$ 概率的事件，即多个 token 被同时更新。

7.2 Training CTMC models

接下来，我们讨论如何学习 CTMC models。其核心原则与 flow matching 相同：

• (1). 我们构造一条在 noise 与 data 之间进行插值的概率路径。
• (2). 我们推导条件 rate matrix 与边缘 rate matrix。
• (3). 我们以 simulation-free 的方式学习边缘 rate matrix。

下面我们将逐步解释这一训练流程。

在本节中，数据分布 $p_{\mathrm{data}}$ 是在状态空间 $S$ 上的一个分布，并由概率质量函数表示。即：$p_{\mathrm{data}}:S\to {\mathbb{R}}_{\ge 0},z\mapsto p_{\mathrm{data}}(z)$ 满足 ${\sum }_{z\in S}{p}_{\mathrm{data}}(z)=1$ 。我们并不知道 $p_{\mathrm{data}}$ 的具体形式，但在训练过程中可以访问其样本 $z\sim p_{\mathrm{data}}$ ，这些样本以数据集的形式给出。例如，互联网上的所有文本数据。我们的目标是学习生成样本 $z\sim p_{\mathrm{data}}$ 。我们的目标是训练 CTMC model $Q_t^{\theta }$ 使得：

$$X_0\sim p_{\mathrm{init}},X_t\text{CTMC of }{Q}_t^{\theta }\Rightarrow X_1\sim p_{\mathrm{data}}$$

因此你会发现，这与欧氏空间 ${\mathbb{R}}^d$ 中的情形（见 Sections 2 与 3）并没有本质区别，只不过这里我们使用的是 CTMC model，而不是 flow/diffusion model。

7.2.1 Conditional and Marginal Probability Path

我们定义 ${\delta }_z(x)$ 为如下函数 ${\delta }_z(x)=0\text{if }x\ne z$ 以及 ${\delta }_z(x)=1\text{if }x=z$ 。一个（离散的）条件概率路径 由一组分布 $p_t(x|z)$ 给出，其中 $x,z\in S,0\le t\le 1$ 并满足：

$$p_0(\cdot |z)=p_{\mathrm{init}},p_1(\cdot |z)={\delta }_z$$

因此，与欧氏空间中的情形类似，一个离散条件概率路径会在一个与 $z$ 无关的分布和一个所有概率质量都集中在 $z$ 上的分布之间进行插值。随后，（离散的）边缘概率路径定义为：

$$p_t(x)=\sum _{z\in S}{p}_t(x|z)p_{\mathrm{data}}(z)$$

可以很容易验证，边缘概率路径在 “noise” 与 data 之间进行插值：

$$p_0=p_{\mathrm{init}},p_1=p_{\mathrm{data}}\text{\hspace{1em}(89)}$$

---

Example 35 (Factorized mixture path (independent noising per token))

设 $S={\mathcal{V}}^d$ ，并令 $p_{\mathrm{init}}(x)={\prod }_{j=1}^d{p}_{\mathrm{init}}^{(j)}(x_j)$ 为一个因子化初始分布。固定一个 scheduler $0\le {\kappa }_t\le 1$ 满足 ${\kappa }_0=0,{\kappa }_1=1,\frac{d}{dt}{\kappa }_t\ge 0$ 。定义条件路径为：

$$p_t(x|z)=\prod _{j=1}^d[(1-{\kappa }_t)p_{\mathrm{init}}^{(j)}(x_j)+{\kappa }_t{\delta }_{z_j}(x_j)].$$

等价地，我们可以通过如下方式采样 $x\sim p_t(\cdot |z)$ ：首先采样 i.i.d. masks m j ∈ { 0 , 1 } m_j\in\{0,1\} mj​∈{0,1} 以及 noise ${\xi }_j\sim p_{\mathrm{init}}^{(j)}$ ，然后设定：

$$\begin{array}{rl}{m}_j& \sim \mathrm{Bernoulli}({\kappa }_t),{\xi }_j\sim p_{\mathrm{init}}^{(j)}\\ x_j& =m_j{z}_j+(1-m_j){\xi }_j,j=1,\dots ,d\\ x& =(x_1,\dots ,x_d).\end{array}$$

我们将上述路径称为 factorized mixture path（因子化混合路径）。上述过程实际上会以概率 $1-{\kappa }_t$ 独立地 “破坏” 序列中每一个位置的第 $j$ 个 token。即当 $t=0$ 时， $1-{\kappa }_t=1$ ，所有信息都被破坏；当 $t=1$ 时， $1-{\kappa }_t=0$ ，没有任何信息被破坏。

注意，这与高斯概率路径（Example 8）是相似的，因为信息都会按照 scheduler ${\kappa }_t$ 所决定的速度逐渐被破坏。然而，它与高斯概率路径也存在不同：因子化混合路径并不会移动 / 运输概率质量（因为在离散空间中不存在方向），它只是从一个分布逐渐淡出，再逐渐淡入另一个分布。

---

Figure 19：当 $d=2$ 时，一个离散概率路径的示意图。第一行：条件概率路径在初始分布与狄拉克分布之间进行插值。第二行：初始分布与数据分布（这里为棋盘格模式）之间的插值。注意它与 Figure 5 的相似性与不同点。这里，概率路径是被 “teleported” 的（我们减小初始分布的权重，并增大 terminal distribution 的权重）。

7.2.2 Conditional and Marginal Rate Matrix

作为下一步，我们现在构造离散 flow matching 的训练目标。首先，我们构造一个条件 rate matrix — 它对应于 flow matching 中的条件向量场。对于每一个数据点 $z\in S$ ，令 $Q_t^z(y|x)$ 为一个 rate matrix。如果满足：

$$X_0\sim p_{\mathrm{init}},X_t\text{CTMC of }{Q}_t^z\Rightarrow X_t\sim p_t(\cdot |z)$$

则我们称其为 conditional rate matrix。

换句话说，条件 rate matrix 所对应的 CTMC 会 “follow” 条件概率路径。条件 rate matrix 作为一个基础构件，用于构造遵循边缘概率路径的边缘 rate matrix：

---

Theorem 36 (Discrete marginalization trick)

定义 marginal rate matrix 为：

$$Q_t(y|x)=\sum _{z\in S}{Q}_t^z(y|x)\frac{p_t(x|z)p_{\mathrm{data}}(z)}{p_t(x)}=\sum _{z\in S}{Q}_t^z(y|x)p_{1|t}(z|x)\text{where }{p}_{1|t}(z|x):=\frac{p_t(x|z)p_{\mathrm{data}}(z)}{p_t(x)}\text{\hspace{1em}(90)}$$

则它是一个合法的 rate matrix，并满足如下条件：

$$X_0\sim p_{\mathrm{init}},X_t\text{CTMC of }{Q}_t\Rightarrow X_t\sim p_t$$

特别地，由 公式 (89) 可知 $X_1\sim p_{\mathrm{data}}$ ，即边缘 rate matrix 对应的 CTMC 会将 noise 转换为 data。

---

为了证明这一结论，我们需要一个关于 CTMC 的基础方程，即所谓的 Kolmogorov Forward equation（科尔莫戈罗夫正方程）：

---

Proposition 2 (Kolmogorov Forward Equation)

设 $p_t$ 是在 $0\le t\le 1$ 上定义于状态空间 $S$ 的一组分布。进一步地，设 $X_t$ 是一个具有矩阵 $Q_t$ 以及初始分布 $p_0$ 的 CTMC。那么， $X_t\sim p_t\text{for all }0\le t\le 1$ 当且仅当 Kolmogorov Forward Equation（KFE） 成立：

$$\frac{d}{dt}{p}_t(x)=\sum _{y\in S}{Q}_t(x|y)p_t(y)$$

---

Proof of KFE. 为了证明 KFE 是必要条件，假设 $p_t(x)$ 是 CTMC 的真实 marginals，即 $X_t\sim p_t\text{for every }0\le t\le 1$ 。那么我们可以计算：

$$\begin{array}{rl}\frac{d}{dt}{p}_t(x)& \stackrel{(i)}{=}\frac{d}{dh}{|}_{h=0}{p}_{t+h}(x)\\ & \stackrel{(ii)}{=}\frac{d}{dh}{|}_{h=0}\sum _y{p}_{t+h|t}(x|y)p_t(y)\\ & \stackrel{(iii)}{=}\sum _y\frac{d}{dh}{|}_{h=0}{p}_{t+h|t}(x|y)p_t(y)\\ & \stackrel{(iv)}{=}\sum _y{Q}_t(x|y)p_t(y)\end{array}$$

其中：

• 在 $(i)$ 中，我们只是使用了一个时间偏移；
• 在 $(ii)$ 中，我们使用了转移概率的定义；
• 在 $(iii)$ 中，我们交换了求和与求导；
• 在 $(iv)$ 中，我们使用了 rate matrix 的定义（见 公式 (87)）。

接下来，为了证明 KFE 是充分条件，我们可以将 KFE 改写为矩阵形式：

$$\frac{d}{dt}{p}_t=Q_t{p}_t$$

其中，在这个方程中，我们将 $p_t=(p_t(x){)}_{x\in S}$ 视为一个向量，并将 $Q_t=(Q_t(y|x){)}_{x,y\in S}$ 视为一个矩阵。注意，上式是在向量空间 ${\mathbb{R}}^S$ 上的一个线性 ODE。其初始条件由定理中的 $p_0$ 给定。

因此，如果任何其他一组 marginals $q_t$ 也满足该方程，那么根据 ODE 的唯一性（见 Theorem 3），我们可以得出 $q_t=p_t$ 。这表明 KFE 同样也是充分条件。

Proof of Theorem 36. 使用 KFE 后，剩下只需要证明定理中定义的 marginal rate matrix（见 公式 (90)）满足 KFE：

$$\begin{array}{rl}\frac{d}{dt}{p}_t(x)& \stackrel{(i)}{=}\frac{d}{dt}\sum _{z\in S}{p}_t(x|z)p_{\mathrm{data}}(z)\\ & \stackrel{(ii)}{=}\sum _{z\in S}\frac{d}{dt}{p}_t(x|z)p_{\mathrm{data}}(z)\\ & \stackrel{(iii)}{=}\sum _{z\in S}\left[\sum _{y\in S}{Q}_t^z(x|y)p_t(y|z)\right]p_{\mathrm{data}}(z)\\ & \stackrel{(iv)}{=}\sum _{y\in S}{p}_t(y)\left[\sum _{z\in S}{Q}_t^z(x|y)\frac{p_t(y|z)p_{\mathrm{data}}(z)}{p_t(y)}\right]\\ & \stackrel{(v)}{=}\sum _{y\in S}{p}_t(y)Q_t(x|y)\end{array}$$

其中：

• 在 $(i)$ 中，使用了边缘概率路径的定义；
• 在 $(ii)$ 中，交换了求和与求导；
• 在 $(iii)$ 中，对条件 rate matrix 使用了 KFE；
• 在 $(iv)$ 中，同时乘以并除以 $p_t(y)$ ；
• 在 $(v)$ 中，使用了边缘 rate matrix $Q_t(y|x)$ 的定义。

这说明 KFE 成立。该结论由 Proposition 2 得出。

现在，我们来推导因子化混合路径的条件 rate matrix 的一个具体例子。

---

Example 37 (Conditional rate matrix for factorized mixture path)

设 $\frac{d}{dt}{\kappa }_t={\dot{\kappa }}_t$ ，因子化混合路径具有如下因子化条件 rate matrix：

$$\begin{array}{rl}{Q}_t^z(y|x)& ={\left(Q_t^z(v_i,j|x_j)\right)}_{v_i,j}\\ Q_t^z(v_i,j|x_j)& =\frac{{\dot{\kappa }}_t}{1-{\kappa }_t}\left({\delta }_{z_j}(v_i)-{\delta }_{x_j}(v_i)\right)\\ & =\frac{{\dot{\kappa }}_t}{1-{\kappa }_t}\{\begin{array}{ll}0& \text{if }{x}_j=z_j\\ 1& \text{if }{v}_i=z_j,x_j\ne z_j\\ 0& \text{if }{v}_i\ne z_j,x_j\ne z_j\\ -1& \text{if }{v}_i=x_j,x_j\ne z_j\end{array}\end{array}$$

注意，这是一个非常简单的 rate matrix：它只允许跳转到 $z^j$ ，也就是说，如果任意 token $j$ 被更新，它必须跳转到终点数据点 $z=(z_1,\cdots ,z_d)$ 的 token value；并且只有在当前还没有到达该 token 时，它才会跳转到 $z^j$ 。

Proof. 注意，因子化混合路径完全分解为独立分量，所提出的条件 rate matrix 也是如此。因此，不失一般性，我们可以假设 $d=1$ 。也就是说，我们只需要逐维进行计算。于是可以推导：

$$\begin{array}{rl}\frac{d}{dt}{p}_t(x|z)& \stackrel{(i)}{=}\frac{d}{dt}\left[(1-{\kappa }_t)p_{\mathrm{init}}(x)+{\kappa }_t{\delta }_z(x)\right]\\ & \stackrel{(ii)}{=}{\dot{\kappa }}_t{\delta }_z(x)-{\dot{\kappa }}_t{p}_{\mathrm{init}}(x)\\ & \stackrel{(iii)}{=}\frac{{\dot{\kappa }}_t}{1-{\kappa }_t}\left({\delta }_z(x)-\left[(1-{\kappa }_t)p_{\mathrm{init}}(x)+{\kappa }_t{\delta }_z(x)\right]\right)\\ & \stackrel{(iv)}{=}\frac{{\dot{\kappa }}_t}{1-{\kappa }_t}\left({\delta }_z(x)-p_t(x|z)\right)\\ & \stackrel{(v)}{=}\frac{{\dot{\kappa }}_t}{1-{\kappa }_t}{\delta }_z(x)\left(1-p_t(x|z)\right)+\frac{{\dot{\kappa }}_t}{1-{\kappa }_t}\left({\delta }_z(x)-1\right)p_t(x|z)\\ & \stackrel{(vi)}{=}\sum _{y\ne x}\frac{{\dot{\kappa }}_t}{1-{\kappa }_t}{\delta }_z(x)p_t(y|z)+\frac{{\dot{\kappa }}_t}{1-{\kappa }_t}\left({\delta }_z(x)-1\right)p_t(x|z)\\ & \stackrel{(vii)}{=}\sum _{y\ne x}{Q}_t^z(x|y)p_t(y|z)+Q_t^z(x|x)p_t(x|z)\\ & \stackrel{(viii)}{=}\sum _{y\in S}{Q}_t^z(x|y)p_t(y|z)\end{array}$$

其中：

• $(i)$ 使用了 $d=1$ 时因子化混合路径的定义；
• $(ii)$ 通过求导并设 $\frac{d}{dt}{\kappa }_t={\dot{\kappa }}_t$ 得到；
• $(iii)$ 由简单代数变形得到；
• $(iv)$ 使用了因子化混合路径的定义；
• $(v)$ 由简单代数变形得到；
• $(vi)$ 使用了 ${\sum }_{y\in S}{p}_t(y|z)=1$ 这一事实；
• $(vii)$ 使用了 rate matrix 的定义；
• $(viii)$ 由简单代数变形得到。

上述推导说明 KFE 成立，因此命题得证。

---

7.2.3 Learning the Marginal Rate Matrix

在本节中，我们推导用于训练 CTMC models 的基本算法。根据 Theorem 36，训练一个 CTMC model $Q_t^{\theta }(y|x)$ 可以通过学习 marginal rate matrix 来实现。

在本节中，我们现在只考虑因子化混合路径（见 Example 35），因为这是目前大多数 discrete diffusion / flow matching models 所使用的路径。在这种情况下，marginal rate matrix 具有非常直观的形式：

---

Theorem 38 (Marginalization trick for factorized mixture path)

因子化混合路径的 marginal rate matrix 是 factorized 的，并且具有如下形式：

$$Q_t(v_i,j|x)=\frac{{\dot{\kappa }}_t}{1-{\kappa }_t}\left(p_{1|t}(z_j=v_i|x)-{\delta }_{x_j}(v_i)\right)$$

其中， $p_{1|t}(z_j=v_i|x)$ 是在给定完整 noisy sequence $x$ 的情况下，第 $j$ 个位置（序列中的第 $j$ 个 token）等于 $v_i$ 的条件概率。

---

Proof. marginal rate matrix 定义为：

$$Q_t(y|x)=\sum _{z\in S}{Q}_t^z(y|x)p_{1|t}(z|x)\text{\hspace{1em}(91)}$$

现在，当 $y$ 与 $x$ 不是 neighbors 时（即超过一个 token 不同），对于任意 $z$ 都有 $Q_t^z(y|x)=0$ 。因此，在这种情况下也有 $Q_t(y|x)=0$ 。这说明 marginal rate matrix 也是因子化的。于是有：

$$\begin{array}{rl}{Q}_t(v_i,j|x)& =\sum _{z\in S}{Q}_t^z(v_i,j|x)p_{1|t}(z|x)\\ & \stackrel{(i)}{=}\sum _{z\in S}\frac{{\dot{\kappa }}_t}{1-{\kappa }_t}\left({\delta }_{z_j}(v_i)-{\delta }_{x_j}(v_i)\right)p_{1|t}(z|x)\\ & \stackrel{(ii)}{=}\frac{{\dot{\kappa }}_t}{1-{\kappa }_t}\left(\sum _{z\in S}{\delta }_{z_j}(v_i)p_{1|t}(z|x)-{\delta }_{x_j}(v_i)\right)\\ & \stackrel{(iii)}{=}\frac{{\dot{\kappa }}_t}{1-{\kappa }_t}\left(p_{1|t}(z_j=v_i|x)-{\delta }_{x_j}(v_i)\right)\end{array}\text{\hspace{1em}(92)(93)(94)(95)}$$

其中：

• $(i)$ 来自条件 rate matrix 的公式（见 Example 37）；
• $(ii)$ 来自 ${\sum }_{z\in S}{p}_{1|t}(z|x)=1$ 这一事实；
• $(iii)$ 来自边缘化。

证明结束。

前面的定理非常值得注意：marginal rate matrix 本质上是对概率 $p_{1|t}(z_j=v_i|x)$ 的一种重新参数化。这实际上无非就是对每一个 token 位置 $j=1,\dots ,d$ 学习一个分类器。换句话说，我们可以简单地定义一个 denoising probabilities network（去噪概率网络）：

$$p_{1\mid t}^{\theta }:\underset{\text{network input}}{\underbrace{x}}\mapsto \underset{\text{network output}}{\underbrace{{\left(p_{1\mid t}^{\theta }(z_j=v_i\mid x)\right)}_{j=1,\dots ,d,v_i\in \mathcal{V}}}}$$

注意，网络输出的 shape 为 $d\times V$ 。我们可以通过简单的 softmax layer 得到每个 token 位置上的概率。网络本身可以是一个标准 sequence-to-sequence network，例如 transformer 就可以工作（见 Section 6.1.2）。

由于这本质上只是对每个位置 $j$ 进行分类，因此我们可以通过每个位置 $j=1,\dots ,d$ 上的交叉熵损失来训练该网络，这便得到如下的 Discrete Flow Matching loss：

$${\mathcal{L}}_{\mathrm{DFM}}(\theta )={\mathbb{E}}_{z\sim p_{\mathrm{data}},t\sim {\mathrm{Unif}}_{[0,1]},x\sim p_t(\cdot |z)}\left[\sum _{j=1}^d-\log p_{1|t}^{\theta }(z_j|x)\right]$$

这一点非常值得注意：为了训练一个生成模型，我们真正需要做的，只是对每个位置 $j$ 训练一个分类模型。就像连续 flow matching 被简化成简单回归问题（见 Section 3）一样，discrete flow matching 以及 discrete diffusion models 也被简化成了简单的分类训练。

在 Algorithm 8 中，我们总结了训练算法。训练完成后，我们可以通过 Algorithm 7 进行采样。

---

Example 39 (Masked Diffusion Language Model)

上述方法的一个特殊情况是 masked diffusion language models（MDLMs）。MDLM 的核心思想是：我们可以将 token 词汇表 V = { v 1 , … , v V } \mathcal{V}=\{v_1,\ldots,v_V\} V={v1​,…,vV​} 扩展一个新的 token $[\mathrm{mask}]$ ，它表示该 token 缺失（或者被 mask 掉）。

具体而言，我们设置 V = { v 1 , … , v V , [ m a s k ] } \mathcal{V}=\{v_1,\ldots,v_V,[\mathrm{mask}]\} V={v1​,…,vV​,[mask]} 并将初始点简单设置为 $[\mathrm{mask}{]}^d$ ，即整个序列全部由 mask token 构成。形式化地，这意味着在上述框架中设置 $p_{\mathrm{init}}={\delta }_{[\mathrm{mask}{]}^d}$ ，采样过程如 Figure 20 所示。

---

Figure 20：MDLM 轨迹示意图

至此，我们已经完成了一个完整的 CTMC 模型训练与采样 pipeline，它能够用于生成诸如文本这样的离散序列。当前最先进的 discrete diffusion models [4] 采用的正是本文中描述的方法：使用神经网络（通常是 transformers）并在 web-scale 数据上进行训练。

---

Remark 40 (Generator Matching)

你可能会好奇为什么 flow/diffusion models 的原理能够如此自然地推广到离散状态空间？事实证明，flow matching 的原理并不局限于 flow，甚至也不局限于 CTMC。更准确地说，这些实际上是利用马尔可夫过程构建生成模型的一般性学习原理。

这便引出了 Generator Matching framework [19]，这是一个能够统一并扩展离散与连续 flow/diffusion models 的框架。generator 可以被看作是向量场 $u_t$ 以及 rate matrix $Q_t$ 的一种泛化。

马尔可夫过程与 generators 可以针对任意数据模态与状态空间进行构建。例如，你可以为 smooth manifolds 构建模型 [8, 10]（例如几何数据）；也可以针对混合状态空间构建模型（例如联合文本与图像生成）[6]；还可以构建其它类型的马尔可夫过程，例如 jump processes [19, 7]。

---