前言

学习 MIT S.184 课程，本篇文章记录课程第五讲：离散扩散模型，记录下个人学习笔记，和大家一起分享交流😄

Website：https://diffusion.csail.mit.edu/

Course Notes：An Introduction to Flow Matching and Diffusion Models

1. Class Overview

好的，欢迎来到第五讲，这也是本课程的最后一讲。本讲我们将讨论 离散扩散模型 和 离散流匹配。前几讲中，我们主要围绕连续空间中的生成模型展开，包括流模型、扩散模型、流匹配、score matching、guidance，以及如何构建图像和视频生成器。到了这一讲，我们会稍微转换一下方向，将之前学到的思想推广到离散状态空间中，特别是语言、蛋白质序列等离散序列数据。

在前面的课程中，Lecture 1 介绍了流模型和扩散模型的基本框架；Lecture 2 讨论了作为训练算法的 flow matching；Lecture 3 讲解了 score matching 和 guidance，即如何对生成过程施加条件控制；Lecture 4 则进一步介绍了如何构建图像生成器，包括 latent space 和网络结构设计。而今天的 Lecture 5 将关注一个新的方向：discrete diffusion models and discrete flow matching，即如何把连续空间中的扩散和流匹配思想推广到离散序列数据上。

2. Discrete diffusion models and discrete flow matching

最近，离散扩散模型和扩散语言模型在新闻中出现得越来越频繁。

例如，Google 提出了与 Gemini Diffusion 相关的工作，一些初创公司也开始尝试使用扩散式语言模型生成文本，ByteDance Seed 等团队也在探索基于离散扩散的语言模型。相比传统自回归语言模型，这类模型的一个重要特点是：它们不一定按照从左到右的顺序逐 token 生成文本，而是可以以更加灵活的方式对文本进行逐步去噪和填充。

那么，什么是扩散语言模型呢？我们可以先从直觉上理解：传统的 ChatGPT、Gemini 这类语言模型通常是自回归模型，也就是按照从左到右的顺序，一个 token 接一个 token 地生成文本。而扩散语言模型则试图采用另一种方式生成文本：它可以从一个被破坏、被 mask 或者高度噪声化的序列开始，然后通过多步去噪，逐渐恢复出完整文本。因此，它的生成顺序不再一定是严格的从左到右，而可以是任意位置逐步被恢复。

不过需要特别注意的是，虽然这类模型常被称为 “扩散模型”，但它们并不是连续空间中由 SDE 描述的经典扩散过程。在图像生成中，我们通常把图像看作连续变量，并通过 SDE 或 ODE 描述样本在连续空间中的演化。但对于文本、蛋白质序列等离散数据而言，状态空间本身是离散的，不存在通常意义下的连续轨迹，也不存在 “沿着某个方向移动” 的向量场。因此，离散扩散模型更准确地说，是将连续扩散模型中的 “逐步破坏数据、再学习反向恢复” 的思想推广到离散状态空间。

本讲关注的对象是 discrete sequence data，也就是离散序列数据。例如自然语言可以看作由词表中的 token 构成的序列，蛋白质也可以看作由氨基酸字母表构成的序列。只要数据可以被表示成有限词表上的离散序列，就可以尝试用类似的方法建模。

但是，这里必须澄清一个关键点：在离散空间中，没有通常意义下的 diffusion/SDE，也没有 flow/ODE。因为 SDE 和 ODE 都依赖连续状态空间中的微小位移，而离散状态之间不存在无限小的连续移动。我们不能说一个 token 沿着某个向量场连续移动到另一个 token，因为 token 之间只能发生离散跳变。

因此，本讲真正要讨论的不是直接把连续 SDE 或 ODE 搬到离散空间中，而是将 flow matching 和 denoising diffusion 的学习原则推广到离散数据上。为了描述离散状态空间中的随机演化，我们会使用一个新的数学工具：Continuous-Time Markov Chains，即连续时间马尔可夫链，简称 CTMC。

今天的内容主要包括两个部分：首先介绍 CTMC 模型，然后讨论如何利用离散流匹配来学习这类模型。离散流匹配部分又会进一步涉及 discrete probability paths、discrete marginalization trick 和 discrete FM objective。

3. Continuous-Time Markov Chain, CTMC

我们先从连续时间马尔可夫链开始。

在离散序列建模中，首先需要定义一个词表，也可以称为字母表。设词表为：

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

其中 $V$ 表示词表大小。这里的词表不一定是英文字母，也可以是 BPE token、单词、氨基酸，或者其他任意离散符号集合。

接下来定义状态空间。前几讲中，我们讨论的状态空间通常是连续向量空间，例如 ${\mathbb{R}}^d$ 。而在本讲中，状态空间是长度为 $d$ 的离散序列集合：

$$S={\mathcal{V}}^d$$

也就是说， $S$ 中的每个元素都是一个长度为 $d$ 的 token 序列：

$$X\in S,X=(X_1,\cdots ,X_d)$$

其中每个 $X_i$ 都是词表 $\mathcal{V}$ 中的一个 token。

为了描述状态之间如何随时间变化，我们需要引入 rate matrix，即速率矩阵。速率矩阵可以看作一个随时间变化的函数：

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

其中 $Q_t(y\mid x)$ 表示：在时间 $t$ 附近，当前状态为 $x$ 时，链跳转到状态 $y$ 的瞬时速率。这里一定要注意方向， $Q_t(y\mid x)$ 表示的是 从 $x$ 跳到 $y$ 的速率，而不是从 $y$ 跳到 $x$ 的速率。

在连续空间中，我们可以用向量场描述样本应该朝哪个方向移动。例如在 ODE 中，向量场告诉我们 $x_t$ 应该沿着哪个方向连续演化。但在离散空间中，这种 “方向” 并不存在。例如，从一个 token 变成另一个 token 时，我们不能说它沿着某个连续方向移动，而只能说它从一个离散状态跳到了另一个离散状态。因此，在离散空间中，我们不再使用向量场，而是使用速率矩阵来描述状态之间跳变的强度。

速率矩阵需要满足两个基本条件：

$$\begin{array}{rl}(1)Q_t(y\mid x)& \ge 0,\text{whenever }x\ne y\\ (2)Q_t(x\mid x)& =-\sum _{y\ne x}{Q}_t(y\mid x),\text{for all }x\end{array}$$

第一个条件说明，对于非对角项 $y\ne x$ ，从 $x$ 跳到 $y$ 的速率必须是非负的。速率可以为 0，表示不允许从 $x$ 直接跳到 $y$，但不能为负。

第二个条件说明，对角项 $Q_t(x\mid x)$ 并不表示真正意义上的 “从 $x$ 跳到 $x$”，而是被定义为所有离开 $x$ 的速率之和的相反数。这样做的目的是保证概率质量守恒。换句话说，固定当前状态 $x$，对所有可能的目标状态 $y$ 求和，可以得到：

$$\sum _y{Q}_t(y\mid x)=0$$

如果我们把 $Q_t(y\mid x)$ 写成矩阵，并且约定 $x$ 是列索引、$y$ 是行索引，那么这个条件就表示每一列的和为 0。当然，如果采用相反的矩阵约定，也可能写成行和为 0。本质上这只是符号约定的差异，关键是对每个当前状态 $x$，所有目标状态对应的速率之和必须为 0。

---

Q & A

Q：这里到底是行和为 0，还是列和为 0？

A：这取决于矩阵索引的约定。在这里我们采用 $Q_t(y\mid x)$ 的写法，其中 $x$ 是当前状态， $y$ 是目标状态。如果把 $y$ 作为行索引、 $x$ 作为列索引，那么固定 $x$ ，对所有 $y$ 求和，就是列和为 0：

$$\sum _y{Q}_t(y\mid x)=0$$

如果采用另一种约定，把当前状态放在行索引上，也可以写成行和为 0。所以这里不必纠结行还是列，关键是：对于每个当前状态，所有流出和保留项加起来必须保证总概率质量不变。

---

现在我们可以正式理解 CTMC。连续时间马尔可夫链是一个连续时间随机过程。对每个时间 $t$ ， $X_t$ 都是一个取值于状态空间 $S$ 的随机变量：

$$X_t\in S$$

马尔可夫性意味着未来状态的分布只依赖于当前状态，而不依赖于更早的历史。因此，CTMC 可以通过转移概率来刻画：

$$p_{t+h\mid t}(X_{t+h}=y\mid X_t=x)$$

它表示在时间 $t$ 处于状态 $x$ 的条件下，经过时间间隔 $h$ 后，在时间 $t+h$ 处于状态 $y$ 的概率。这类似于离散时间马尔可夫链中的转移概率，只不过这里时间是连续的。

那么，转移概率和速率矩阵之间是什么关系呢？速率矩阵给出了转移概率在无穷小时间尺度上的变化率：

$${\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$$

这说明 $Q_t(y\mid x)$ 描述的是：当时间从 $t$ 增加一个极小量时，从状态 $x$ 跳到状态 $y$ 的概率增长速度。

这与连续空间中的 ODE 有一个类比关系。在 ODE 中，向量场指定样本轨迹如何连续演化；而在 CTMC 中，速率矩阵指定转移概率如何随时间演化。不同的是，ODE 描述的是连续轨迹，而 CTMC 描述的是离散状态之间的随机跳变。

上图展示了一条 CTMC 轨迹。状态空间中只有三个状态 $S_1,S_2,S_3$ 。随着时间推移，随机过程会在这些离散状态之间发生跳变。例如，一开始可能处于 $S_3$ ，在时间 $t_1$ 跳到 $S_1$ ，之后又在时间 $t_2$ 跳到 $S_2$ ，再在时间 $t_3$ 跳回 $S_1$ ，最后在时间 $t_4$ 跳到 $S_3$ 。可以看到，CTMC 的路径是分段常数的：在两个跳变时刻之间，状态保持不变；只有在某些随机时刻，状态才会突然切换。

接下来我们看一个最简单的二状态 CTMC 例子。

设状态空间为：

S = { a , b } S = \{a,b\} S={a,b}

也就是说，系统只有两个可能状态：$a$ 和 $b$ 。假设从 $a$ 跳到 $b$ 的速率为 $\lambda$ ，从 $b$ 跳到 $a$ 的速率也为 $\lambda$ 。那么对应的速率矩阵可以写成：

$$Q=\left(\begin{array}{cc}-\lambda & \lambda \\ \lambda & -\lambda \end{array}\right)$$

这里的非对角项 $\lambda$ 表示两个状态之间相互跳转的速率，而对角项 $-\lambda$ 则由概率质量守恒条件决定。对于当前状态 $a$ ，离开 $a$ 的速率是 $\lambda$ ，所以 $Q(a\mid a)=-\lambda$ ；对于当前状态 $b$ ，离开 $b$ 的速率也是 $\lambda$ ，所以 $Q(b\mid b)=-\lambda$ 。

通过求解对应的演化方程，可以得到时间间隔为 $h$ 时的转移概率矩阵：

$$P(h)=\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)$$

从这个表达式可以看出，当 $h$ 很小时，系统大概率仍然留在原状态；随着 $h$ 增大，指数项 $e^{-2\lambda h}$ 会逐渐衰减，状态之间的转移概率会逐渐趋于平衡。

当 $h\to \infty$ 时：

$$e^{-2\lambda h}\to 0$$

因此转移概率矩阵收敛为：

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

这意味着无论初始状态是 $a$ 还是 $b$ ，经过足够长的时间后，系统都会收敛到在两个状态上均匀分布的平稳分布。换句话说，随着时间推移，初始状态的信息会逐渐消失，最后系统在两个状态之间达到完全混合。

这个例子说明了 CTMC 的基本机制：速率矩阵决定状态之间的瞬时跳转行为，而转移概率描述经过有限时间后状态变化的整体结果。

4. CTMC models

接下来，我们讨论如何将 CTMC 用于机器学习中的生成建模。

所谓 CTMC model，就是用一个带参数的速率矩阵来描述离散状态空间中的随机演化。这个速率矩阵通常由神经网络参数化，记为：

$$Q_t^{\theta }(y\mid x)$$

其中 $\theta$ 是神经网络参数， $x$ 是当前状态， $y$ 是目标状态， $t$ 是时间。模型的目标是学习一个合适的速率矩阵，使得从简单分布出发，经过 CTMC 的随机演化后，可以生成目标数据分布中的样本。

不过，在离散序列建模中，完整状态空间非常大。如果词表大小为 $V$ ，序列长度为 $d$ ，那么状态空间大小为：

$$|S|=V^d$$

这意味着所有可能序列的数量会随着序列长度指数级增长。对于语言模型而言， $V$ 可能是几万甚至几十万， $d$ 也可能非常大，因此完整的速率矩阵根本无法显式存储或直接学习。

为了解决这个问题，我们引入一个重要约束：factorized condition，也就是因子分解条件。它要求只有当两个序列只在一个位置上不同时，它们之间的直接跳转速率才可能非零。如果两个序列在超过一个位置上不同，则直接跳转速率被设为 0：

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

也就是说，如果 $x$ 和 $y$ 的汉明距离大于 1，那么模型不允许从 $x$ 一步直接跳到 $y$ 。

上图给出了一个邻居关系的例子。这里的序列长度为 6，每个位置上的 token 可以理解为二值变量。序列 $x$ 和 $y$ 只在一个位置上不同，因此它们是邻居；序列 $y$ 和 $z$ 也只在一个位置上不同，因此它们也是邻居；但是序列 $z$ 和 $x$ 在两个位置上不同，因此它们不是邻居。

这里所谓的邻居，指的是与当前序列 $x$ 只在一个位置上不同的状态。换句话说，如果 $x$ 和 $y$ 的汉明距离为 1，那么 $y$ 是 $x$ 的一个邻居；如果它们在两个或更多位置上不同，那么在 factorized CTMC 中，从 $x$ 直接跳到 $y$ 的速率被设为 0。

这个约束非常重要。它并不是说模型永远无法从 $x$ 到达那些相差多个 token 的状态，而是说模型不能一步直接跳过去。经过多次跳变，模型仍然可以从一个序列逐步变化到另一个完全不同的序列。例如，先改变第一个 token，再改变第二个 token，再改变第三个 token，最终仍然可以到达任意目标序列。factorized condition 只是限制每一次瞬时跳变只能改变一个位置。

在这个约束下，神经网络不需要为所有 $V^d$ 个状态之间的跳转都预测速率，而只需要预测 “单位置替换” 的速率。对于当前序列：

$$x=(x_1,\cdots ,x_d)$$

模型只需要考虑：选择某个位置 $j$ ，然后把这个位置上的 token 替换为词表中的某个 token $v_i$ 。因此，网络一次前向传播可以输出一个 $d\times V$ 的矩阵：

$$x\mapsto {Q_t^{\theta }(y\mid x)}_{y\in N(x)}=\left(\begin{array}{ccc}{Q}_t^{\theta }(v_1,1\mid x)& \cdots & Q_t^{\theta }(v_V,1\mid x)\\ ⋮& & ⋮\\ Q_t^{\theta }(v_1,d\mid x)& \cdots & Q_t^{\theta }(v_V,d\mid x)\end{array}\right)$$

也可以写成：

$$Q_t^{\theta }(v_i,j\mid x),v_i\in \mathcal{V},j=1,\cdots ,d$$

其中 $Q_t^{\theta }(v_i,j\mid x)$ 表示：在时间 $t$ ，将当前序列 $x$ 的第 $j$ 个位置替换为 token $v_i$ 的跳转速率。

这样一来，模型输出规模从完整状态空间的指数级规模降到了线性规模。原本完整状态空间大小是 $V^d$ ，如果允许任意状态之间直接跳转，就需要处理极其庞大的速率矩阵；而在 factorized CTMC 中，模型只需要输出 $d\times V$ 个速率值，这对于实际的语言建模或蛋白质序列建模才是可行的。

这也是 factorized CTMC 的核心意义：它把高维离散序列空间中的复杂跳转结构，拆解成每次只改变一个位置的局部跳转结构。虽然单步跳转被限制了，但通过连续时间中的多次跳转，模型仍然能够在整个状态空间中移动。

上图展示了 general CTMC 和 factorized CTMC 的区别。左边是一般的 CTMC。在一般 CTMC 中，从当前状态 $x$ 可以直接跳到许多不同的状态 $y$ ，因此速率矩阵需要描述所有可能状态之间的跳转关系。这种方式在低维状态空间中可以接受，但在高维离散序列空间中几乎不可行，因为状态数量是指数级的。

右边是 factorized CTMC。在这种模型中，从当前状态 $x$ 出发，只允许沿着某一个坐标方向发生跳变，也就是每次只改变序列中的一个位置。这样，所有可能跳转可以被分解为 “选择位置 $j$” 和 “选择替换 token $v_i$” 两个部分，因此模型只需要预测 $Q_t^{\theta }(v_i,j\mid x)$ 。这种结构大幅降低了建模复杂度，使得 CTMC 可以用于实际的离散序列生成任务。

---

Q & A

Q：这里为什么说是字母表？难道只能是英文字母吗？

A：这里的字母表是一个抽象概念，指的是离散 token 的集合。它不一定是英文字母，也可以是 BPE token、单词、氨基酸、DNA 碱基，或者任意离散符号。在数学上，我们只需要一个有限词表 $\mathcal{V}$ ；在实际应用中，选择什么样的词表会对模型效果产生重要影响。

---

5. Sampling from a CTMC model

接下来，我们继续讨论如何从已经学到的 CTMC 模型中进行采样。

在连续扩散模型或流模型中，采样通常从一个简单分布开始，例如高斯分布。然后通过反向 SDE、ODE 或 flow，把简单噪声逐步转换为数据样本。但在离散状态空间中，高斯分布不再适用，因为状态空间不是 ${\mathbb{R}}^d$ ，而是由有限词表构成的离散集合：

$$S={\mathcal{V}}^d$$

因此，在 CTMC 模型中，我们需要选择一个定义在离散状态空间 $S$ 上的初始分布，记为：

$$p_{\text{init}}$$

一个最简单、最通用的选择是状态空间上的均匀分布：

$$p_{\text{init}}={\text{Unif}}_S$$

也就是说，状态空间 $S$ 中的每一个离散序列都有相同的初始概率。如果 $S$ 是有限集合，那么均匀分布可以写成：

$$p_{\text{init}}(z)=\frac{1}{|S|}$$

当然，在实际语言建模中，初始分布不一定必须是均匀分布，也可以是全 mask 序列、mask token 分布，或者其他容易采样的噪声分布。关键是：它应该是一个简单、已知、容易采样的离散分布。

在离散空间中，我们还会频繁使用一个 $\delta$ 函数，也就是离散版本的 Dirac delta。对于 $x,y\in S$ ，定义：

$${\delta }_y(x)=\{\begin{array}{ll}1,& \text{if }y=x,\\ 0,& \text{otherwise}.\end{array}$$

它表示一个集中在状态 $y$ 上的点质量分布。如果 $x=y$ ，概率为 1；如果 $x\ne y$ ，概率为 0。换句话说， ${\delta }_y$ 就是“确定性地取状态 $y$”的概率分布。

现在考虑如何模拟一个 CTMC。假设我们已经有一个由神经网络参数化的速率矩阵：

$$Q_t^{\theta }(y\mid x)$$

采样过程从初始分布开始：

$$X_0\sim p_{\text{init}}$$

然后希望按照 CTMC 的转移概率从 $X_t$ 更新到 $X_{t+h}$ ：

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

这里 $h>0$ 是一个小的时间步长。这个式子表示：给定当前时刻 $t$ 的状态 $X_t$ ，经过一个小时间间隔 $h$ 后，根据转移概率采样下一个状态。

但是问题在于，我们通常并不知道精确的转移概率 $p_{t+h\mid t}(y\mid X_t)$ 。CTMC 模型直接给出的不是有限时间转移概率，而是瞬时速率矩阵 $Q_t^{\theta }(y\mid x)$ 。因此，我们需要用速率矩阵来近似转移概率。

根据 CTMC 的定义，速率矩阵是转移概率在 $h=0$ 处的导数：

$${\frac{d}{dh}{p}_{t+h\mid t}(y\mid x)|}_{h=0}=Q_t^{\theta }(y\mid x)$$

于是，我们可以对转移概率在 $h=0$ 附近做一阶泰勒展开：

$$p_{t+h\mid t}(y\mid x)\approx p_{t\mid t}(y\mid x)+h{\frac{d}{dh}{p}_{t+h\mid t}(y\mid x)|}_{h=0}$$

当没有时间流逝时，如果当前状态是 $x$ ，那么下一时刻仍然确定地处于 $x$ 。因此：

$$p_{t\mid t}(y\mid x)={\delta }_x(y)$$

或者等价地写成 ${\delta }_y(x)$ 。再代入速率矩阵的定义，就得到：

$$p_{t+h\mid t}(y\mid x)\approx {\delta }_x(y)+h{Q}_t^{\theta }(y\mid x)$$

这就是 CTMC 采样中的 Euler 近似。它的含义很直观：如果 $y\ne x$ ，那么从 $x$ 跳到 $y$ 的概率近似为 $h{Q}_t^{\theta }(y\mid x)$ ；如果 $y=x$ ，那么保持在原状态的概率就是 1 减去所有离开 $x$ 的跳转概率。

具体来说，对于 $y\ne x$ ：

$$p_{t+h\mid t}(y\mid x)\approx h{Q}_t^{\theta }(y\mid x)$$

而对于 $y=x$ ：

$$p_{t+h\mid t}(x\mid x)\approx 1+h{Q}_t^{\theta }(x\mid x)=1-h\sum _{y\ne x}{Q}_t^{\theta }(y\mid x)$$

由于速率矩阵满足：

$$\sum _y{Q}_t^{\theta }(y\mid x)=0$$

所以这个近似分布对所有 $y$ 求和仍然等于 1。也就是说， ${\delta }_x(y)+h{Q}_t^{\theta }(y\mid x)$ 在 $h$ 足够小时可以看作一个合法的离散概率分布。

这里要注意，步长 $h$ 不能太大。如果 $h$ 太大，保持在原状态的概率：

$$1-h\sum _{y\ne x}{Q}_t^{\theta }(y\mid x)$$

可能变成负数。因此，实际采样时需要选择足够小的步长，或者使用更稳定的采样方法。课程这里使用的是 Euler / $\tau$-leaping 形式的近似采样。

OK，那么让我们看看如何采样：

上图给出了从 factorized CTMC model 中采样的算法。输入包括三个部分：

首先是一个 factorized rate network $Q_t^{\theta }$ ，它根据当前序列 $X_t$ 和时间 $t$，输出每个位置、每个候选 token 的跳转速率。

其次是初始分布 $p_{\text{init}}$ ，我们从这个分布中采样初始序列 $X_0$ 。

最后是采样步数 $n$，对应的步长为 $h=\frac{1}{n}$ ，采样从 $t=0$ 开始，经过 $n$ 步后到达 $t=1$ 。

具体过程如下。首先从初始分布采样：

$$X_0\sim p_{\text{init}}$$

其中：

$$X_0=(X_0^{(1)},\cdots ,X_0^{(d)})\in {\mathcal{V}}^d$$

也就是说， $X_0$ 是一个长度为 $d$ 的离散 token 序列。

在每一个时间步中，模型根据当前状态 $X_t$ 计算 factorized jump rates：

$${q_j(v)}_{j=1,\cdots ,d,v\in \mathcal{V}}\leftarrow Q_t^{\theta }(\cdot \mid X_t)$$

这里 $q_j(v)$ 表示：在当前位置序列为 $X_t$ 时，把第 $j$ 个 token 跳转为词表中 token $v$ 的速率。

由于 CTMC 采用 factorized 结构，每个位置可以单独构造一个局部的 Euler 转移分布。对于第 $j$ 个位置，设当前 token 为：

$$x=X_t^{(j)}$$

则定义该位置的转移概率：

$${\tilde{p}}_{j,t}(v\mid x)=\{\begin{array}{ll}h{q}_j(v),& v\ne x,\\ 1-h\sum _{v^{\prime }\in \mathcal{V}\setminus x}{q}_j(v^{\prime }),& v=x.\end{array}$$

这个式子和前面的 CTMC Euler 近似完全一致。对于不同于当前 token 的候选 token $v$ ，跳转概率是步长 $h$ 乘以对应速率；对于保持当前 token 不变的情况，概率则由归一化条件决定。

然后对每个位置 $j$ ，从这个 categorical distribution 中采样新的 token：

$$X_{t+h}^{(j)}\sim \text{Categorical}\left({{\tilde{p}}_{j,t}(v\mid x)}_{v\in \mathcal{V}}\right)$$

所有位置更新完成后，时间推进一步：

$$t\leftarrow t+h$$

重复这个过程 $n$ 次，最终得到：

$$X_1$$

这就是模型生成的离散序列样本。

需要特别注意的是，算法中每个位置的更新可以并行执行。这是 factorized CTMC 的重要优势之一。因为模型把整体状态跳转分解成了每个位置上的 token 替换，所以在一个 Euler / $\tau$-leaping 步中，可以同时为所有位置采样新的 token。

当然，这种并行更新本质上是一种近似。严格的 CTMC 在无穷小时间内通常只发生一次跳变，而并行更新允许一个时间步内多个位置同时变化。这对应于 $\tau$-leaping 近似：在小时间窗口内，用局部速率近似多个可能跳变的联合结果。只要步长足够小，这种近似通常是合理的。

上图展示了离散扩散语言模型的采样过程。初始时，序列可能由 mask token 或噪声 token 组成。随着采样步骤推进，模型会根据当前序列和速率矩阵，逐步更新不同位置上的 token。最终，原本被破坏或被 mask 的序列被逐渐恢复成一段完整文本。

这与自回归语言模型的生成方式不同。自回归模型通常从左到右逐 token 生成，而这里的模型可以在不同位置上并行更新 token。也就是说，文本不是按照固定顺序生成的，而是在多轮迭代中逐渐被填充和修正。

---

Q & A

Q：如果不使用 factorized 结构，还能这样采样吗？

A：原则上可以，但在高维离散序列空间中几乎不可行。完整状态空间大小是 $V^d$ ，如果不做 factorization，就需要处理从当前序列到所有其他序列的跳转速率，计算和存储成本都是指数级的。factorized 条件的作用是把一步跳转限制为只改变一个 token，从而让模型只需要输出 $d\times V$ 个速率。这样，采样时也只需要对每个位置分别构造 categorical distribution。

从一阶近似角度看，速率矩阵描述的是无穷小时间内的瞬时跳转。在 factorized CTMC 中，一阶跳转只改变一个 token。更高阶的效应，例如在一个时间间隔内连续跳两次、改变多个 token，会被 Euler / $\tau$-leaping 近似忽略或近似处理。这正是数值采样误差的来源。因此，步长越小，近似越接近真实 CTMC。

Q：条件生成在这里怎么做？例如给定 prompt 生成文本时，和连续模型中的 guidance 有什么关系？

A：条件生成的思想仍然类似：模型需要在生成过程中利用额外条件信息。在连续扩散模型中，我们可以通过 classifier-free guidance 对向量场或 score 做线性组合；而在离散 CTMC 中，直接相加 “向量场” 不再适用，因为这里没有连续向量场，只有速率矩阵。

不过，条件信息仍然可以作为神经网络的输入，例如把 prompt、类别标签或其他上下文输入给速率网络：

$$Q_t^{\theta }(y\mid x,c)$$

其中 $c$ 表示条件。对于文本生成而言，prompt 本身就是一段已经给定的 token 序列。实践中可以固定 prompt 部分，只对需要生成的部分进行 mask 和更新；也可以让速率网络在预测跳转速率时显式依赖 prompt 上下文。这样，模型就能在条件信息的约束下进行离散去噪和生成。

---

6. Training CTMC models

现在我们进一步讨论如何用 CTMC 做生成建模。

我们的目标仍然是从数据分布中采样。设数据分布为：

$$p_{\text{data}}(z),z\in S$$

在语言建模中， $p_{\text{data}}$ 可以理解为互联网上自然文本的分布；在蛋白质生成中，它可以理解为真实蛋白质序列的分布。总之，它是我们希望模型最终学会生成的目标分布。

同时，我们需要一个简单的初始分布：

$$p_{\text{init}}(z),z\in S$$

例如状态空间上的均匀分布：

$$p_{\text{init}}(z)=\frac{1}{|S|}$$

或者其他容易采样的噪声分布。

生成建模的目标可以概括为：

$$X_0\sim p_{\text{init}}\stackrel{\text{CTMC}}{\to }{X}_1\sim p_{\text{data}}$$

也就是说，我们希望学习一个 CTMC，使得从简单噪声分布出发，经过时间区间 $[0,1]$ 上的随机演化后，最终得到服从数据分布的样本。

这和前面连续流模型中的目标非常类似。区别只在于：连续流模型中，我们使用 ODE 或 SDE 把高斯噪声转换为数据；而这里，我们使用 CTMC 把离散噪声分布转换为离散数据分布。

因此，可以把 CTMC 看成离散空间中的生成动力系统。它不像 ODE 那样沿连续路径移动样本，而是通过一系列随机跳变，把初始离散序列逐步转化为目标数据序列。

为了训练这样的 CTMC 模型，我们可以借鉴前面学习过的 continuous flow matching。

在连续流匹配中，训练框架大致包括以下几个部分：

首先定义一个 conditional probability path，它描述在给定数据样本 $z$ 的条件下，如何从噪声分布逐渐过渡到这个数据点。

然后通过对数据分布边缘化，得到 marginal probability path，它描述整体分布如何从初始分布演化到数据分布。

接着，根据概率路径推导对应的 conditional vector field 和 marginal vector field。

最后，训练一个神经网络向量场，使其匹配这些目标向量场。这就得到 flow matching loss。

在离散空间中，我们希望构造一个类似的框架。但由于没有连续向量场，vector field 需要被替换成 rate matrix。也就是说，连续空间中的 “速度场” 对应离散空间中的 “跳转速率矩阵”。

因此，离散流匹配的整体结构可以类比为：

首先定义 conditional probability path 和 marginal probability path。然后根据这些概率路径构造 conditional rate matrix 和 marginal rate matrix。

最后训练一个参数化速率矩阵：

$$Q_t^{\theta }(y\mid x)$$

使其匹配目标速率矩阵，从而得到 discrete flow matching loss。

所以，continuous flow matching 和 discrete flow matching 的核心区别是：

在连续空间中，我们学习的是向量场 $u_t^{\theta }(x)$ ，它告诉样本应该朝哪个方向连续移动。而在离散空间中，我们学习的是速率矩阵 $Q_t^{\theta }(y\mid x)$ ，它告诉样本应该以多大速率从当前状态 $x$ 跳到另一个状态 $y$。

6.1 Conditional and Marginal Probability Path

现在进入训练 CTMC 模型的核心步骤。第一步是定义离散空间中的条件概率路径和边缘概率路径。

在连续 flow matching 中，我们定义过条件概率路径。离散空间中也是类似的，只不过概率密度变成了概率质量函数。

给定数据样本 $z\in S$，定义条件概率路径：

$$p_t(x\mid z),0\le t\le 1,x,z\in S$$

这里 $p_t(x\mid z)$ 表示：在给定终点数据样本 $z$ 的条件下，时间 $t$ 时处于状态 $x$ 的概率。

由于 $S$ 是离散集合，所以 $p_t(x\mid z)$ 是定义在 $S$ 上的概率质量函数，需要满足：

$$\sum _{x\in S}{p}_t(x\mid z)=1,p_t(x\mid z)\ge 0$$

也就是说，对任意固定的 $t$ 和 $z$ ，$p_t(\cdot \mid z)$ 都是一个合法的离散概率分布。

这个条件概率路径需要满足两个端点条件：

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

第一个条件表示，在时间 $t=0$ ，分布应该是初始噪声分布，并且与数据样本 $z$ 无关。也就是说，一开始我们还没有任何关于目标数据点的信息。

第二个条件表示，在时间 $t=1$ ，分布应该集中到给定的数据样本 $z$ 上。换句话说，给定 $z$ 后，终点状态应该确定为 $z$ 。

有了条件概率路径之后，可以通过对数据分布 $p_{\text{data}}(z)$ 边缘化，得到边缘概率路径：

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

这和连续情形中的积分完全对应。只不过在连续空间中，我们对 $z$ 做积分；而在离散空间中，我们对所有可能的 $z\in S$ 求和。

这个边缘路径连接了初始分布和数据分布。首先，当 $t=0$ 时：

$$p_0(x)=\sum _{z\in S}{p}_0(x\mid z)p_{\text{data}}(z)=\sum _{z\in S}{p}_{\text{init}}(x)p_{\text{data}}(z)=p_{\text{init}}(x)$$

因为：

$$\sum _{z\in S}{p}_{\text{data}}(z)=1$$

所以：

$$p_0=p_{\text{init}}$$

当 $t=1$ 时：

$$p_1(x)=\sum _{z\in S}{p}_1(x\mid z)p_{\text{data}}(z)=\sum _{z\in S}{\delta }_z(x)p_{\text{data}}(z)=p_{\text{data}}(x)$$

因此：

$$p_1=p_{\text{data}}$$

这说明，通过构造合适的条件概率路径，我们就得到了一个从初始分布到数据分布的边缘概率路径：

$$p_0=p_{\text{init}},p_1=p_{\text{data}}$$

这一步和连续 flow matching 中的思想完全一致。我们先构造一条从噪声到数据的概率路径，再根据这条路径推导目标动力学，最后训练模型去匹配这个动力学。

接下来，我们看一个具体的条件概率路径例子：factorized mixture path。

首先定义一个时间调度函数：

$$0\le {\kappa }_t\le 1,{\kappa }_0=0,{\kappa }_1=1$$

可以把 ${\kappa }_t$ 理解为数据权重，随着时间从 0 增加到 1， ${\kappa }_t$ 从 0 逐渐增加到 1。对应地，噪声权重为：

$$1-{\kappa }_t$$

因此，在生成方向上，时间越靠近 0，样本越接近初始噪声；时间越靠近 1，样本越接近数据。

factorized mixture path 的核心思想是：对序列中的每个 token 位置独立地进行插值。对于长度为 $d$ 的序列 $x=(x_1,\cdots ,x_d)$ 和目标数据样本 $z=(z_1,\cdots ,z_d)$ ，定义：

$$p_t(x\mid z)=\prod _{j=1}^d\left[(1-{\kappa }_t)p_{\text{init}}^{(j)}(x_j)+{\kappa }_t{\delta }_{z_j}(x_j)\right]$$

其中 $p_{\text{init}}^{(j)}$ 表示第 $j$ 个位置上的初始噪声分布， ${\delta }_{z_j}(x_j)$ 表示集中在目标 token $z_j$ 上的点质量分布。

这个式子可以这样理解：对于每个位置 $j$，时间 $t$ 时的 token $x_j$ 来自两个来源的混合。一部分来自噪声分布 $p_{\text{init}}^{(j)}(x_j)$ ，其权重为 $1-{\kappa }_t$ ；另一部分来自真实数据 token $z_j$ ，其权重为 ${\kappa }_t$ 。

因此，随着 $t$ 增大，噪声项的权重逐渐减小，数据项的权重逐渐增大。到了 $t=0$ ， ${\kappa }_0=0$ ，路径完全等于初始噪声分布；到了 $t=1$ ， ${\kappa }_1=1$ ，路径完全退化为目标数据样本 $z$ 。

从这个条件路径中采样也很简单。对于每个位置 $j$ ，先采样一个伯努利随机变量：

$$m_j\sim \text{Bernoulli}({\kappa }_t)$$

然后从初始分布中采样一个噪声 token：

$${\xi }_j\sim p_{\text{init}}^{(j)}$$

最后根据 $m_j$ 决定当前位置的 token：

$$x_j=\{\begin{array}{ll}{z}_j,& m_j=1,\\ {\xi }_j,& m_j=0.\end{array}$$

也就是说，以概率 ${\kappa }_t$ 选择真实数据 token $z_j$ ，以概率 $1-{\kappa }_t$ 选择噪声 token ${\xi }_j$ 。对所有位置独立执行这个过程，就得到完整序列：

$$x=(x_1,\cdots ,x_d)$$

这和公式中的乘积结构完全对应：每个位置独立地在 “噪声 token” 和 “数据 token” 之间进行混合。

在时间 $t=0$ 时， ${\kappa }_0=0$ ，所以所有 $m_j$ 都为 0，序列完全来自初始噪声分布。

在时间 $t=1$ 时 ，${\kappa }_1=1$ ，所以所有 $m_j$ 都为 1，序列完全等于数据样本 $z$ 。

中间时刻则对应部分 token 来自数据，部分 token 来自噪声。这就像一个逐步去噪的过程：随着时间推进，越来越多的位置被替换为真实数据 token。

OK，让我们看一个例子：

我们之前展示过这个例子，左边是初始分布 $p_{\text{init}}$ ，可以理解为噪声分布；右边是目标数据分布 $p_{\text{data}}$ 。在连续流模型中，我们通常会把初始分布中的样本通过连续轨迹搬运到数据分布中，也就是样本点在空间中沿着向量场移动。

但是，在离散空间中，情况有所不同。离散状态之间没有连续路径，样本不能沿着一条光滑轨迹从一个点移动到另一个点。它们只能在离散状态之间发生跳变。因此，离散概率路径不再表示连续空间中的质量搬运，而更多表示不同离散状态上的概率权重如何随时间变化。

上图展示了离散概率路径的直观效果。我们可以把状态空间想象成一组离散网格点。随着时间从 0 增加到 1，每个离散点上的概率质量发生变化：初始分布附近的权重逐渐降低，数据分布附近的权重逐渐升高。

这和连续流模型非常不同。在连续流中，概率质量像流体一样被向量场推动，在空间中连续移动；而在离散流匹配中，我们没有这种连续移动的概念。这里更像是在逐步调整不同离散状态的概率权重：一些状态的概率被降低，另一些状态的概率被提高。

因此，离散概率路径可以理解为一种 “概率权重插值” 或 “离散状态上的混合过程”。它告诉我们在每个时间 $t$ ，哪些离散状态应该更可能出现，哪些状态应该逐渐被淘汰。

目前为止，我们只是构造了一条从噪声到数据的概率路径。这个路径本身还不是最终模型，但它是训练模型的基础。接下来要做的是：根据这条离散概率路径，推导出对应的速率矩阵，也就是告诉 CTMC 应该以什么样的跳转速率在离散状态之间演化，才能实现这条概率路径。

换句话说，接下来我们要从 $p_t(x\mid z)$ 推导出 $Q_t(y\mid x)$ ，然后训练神经网络 $Q_t^{\theta }(y\mid x)$ 去匹配这个目标速率矩阵。这就是离散流匹配的核心。

6.2 Conditional and Marginal Rate Matrix

接下来，我们讨论如何从前面构造的离散概率路径中进一步得到实际可训练的模型。前面我们已经定义了离散空间中的 conditional probability path 和 marginal probability path，现在要做的是：为这些概率路径构造对应的 rate matrix，也就是 CTMC 的速率矩阵。

在连续 flow matching 中，我们通过概率路径推导出向量场；而在离散 flow matching 中，向量场被替换成速率矩阵。因此，这里的核心问题是：给定一条离散概率路径，如何找到一个 CTMC，使得它的边缘分布正好沿着这条路径演化？

首先定义条件速率矩阵。对于每一个数据点 $z\in S$ ，令 $Q_t^z(y\mid x)$ 表示依赖于数据点 $z$ 的 rate matrix。如果从初始分布出发，并按照这个速率矩阵运行 CTMC：

$$X_0\sim p_{\mathrm{init}},X_t\text{is a CTMC with rate matrix }{Q}_t^z$$

能够保证：

$$X_t\sim p_t(\cdot \mid z)$$

那么我们就称 $Q_t^z$ 是对应于条件概率路径 $p_t(\cdot \mid z)$ 的 conditional rate matrix。

这和连续 flow matching 中的 conditional vector field 完全对应。在连续空间中，conditional vector field 的作用是：如果样本沿着这个向量场演化，那么它的分布会跟随条件概率路径。在离散空间中，conditional rate matrix 的作用是：如果样本按照这个 CTMC 发生随机跳变，那么它的分布会跟随条件概率路径。

也就是说，可以把连续 flow matching 中的对应关系：

$$\text{probability path}\Rightarrow \text{vector field}$$

替换成离散空间中的对应关系：

$$\text{probability path}\Rightarrow \text{rate matrix}$$

接下来是关键的一步：discrete marginalization trick，也就是离散版本的边缘化技巧。

在第二讲中，我们曾经通过对 conditional vector field 做边缘化，得到 marginal vector field。这里也可以做完全类似的事情，只不过对象从 vector field 换成了 rate matrix。

定义边缘速率矩阵：

$$Q_t(y\mid x)=\sum _{z\in S}{Q}_t^z(y\mid x)\frac{p_t(x\mid z)p_{\mathrm{data}}(z)}{p_t(x)}$$

其中：

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

同时定义后验分布：

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

于是边缘速率矩阵可以简洁地写成：

$$Q_t(y\mid x)=\sum _{z\in S}{Q}_t^z(y\mid x)p_{1\mid t}(z\mid x)$$

这里的 $p_{1\mid t}(z\mid x)$ 可以理解为：在时间 $t$ 观察到当前 noisy state 为 $x$ 时，最终数据样本是 $z$ 的后验概率。也就是说，我们不再固定终点 $z$ ，而是对所有可能的终点 $z$ 进行加权平均，权重就是它们在当前状态 $x$ 下的后验概率。

这个公式的意义非常重要。条件速率矩阵 $Q_t^z$ 是依赖于具体数据点 $z$ 的，因此它本身不能直接作为生成模型，因为生成时我们并不知道目标数据点 $z$ 是什么。而边缘速率矩阵 $Q_t$ 不再显式依赖 $z$ ，它只依赖当前状态 $x$ 和时间 $t$，因此可以作为真正的生成模型使用。

如果这个边缘速率矩阵定义正确，那么从初始分布出发并按照 $Q_t$ 运行 CTMC：

$$X_0\sim p_{\mathrm{init}},X_t\text{is a CTMC with rate matrix }{Q}_t$$

就会得到：

$$X_t\sim p_t$$

特别地，当 $t=1$ 时：

$$X_1\sim p_{\mathrm{data}}$$

这说明，边缘速率矩阵 $Q_t$ 正是我们想要学习的对象。只要学到了它，我们就可以从简单初始分布出发，通过 CTMC 采样得到数据分布中的样本。

---

Q & A

Q：为什么要引入 $Q_t^z$ ？它和 $Q_t$ 有什么区别？

A：$Q_t^z$ 是给定具体数据点 $z$ 时的条件速率矩阵。它描述的是：如果我们已经知道最终要到达的样本是 $z$ ，那么应该如何从噪声状态逐步跳转到这个特定样本。问题是，生成时我们并不知道目标样本 $z$ ，所以不能直接使用 $Q_t^z$ 。

因此，我们对所有可能的 $z$ 做边缘化，得到不依赖具体终点的 $Q_t$ 。这个 $Q_t$ 只依赖当前状态 $x$ 和时间 $t$ ，因此可以作为实际生成模型的速率矩阵。

---

为了证明边缘速率矩阵确实能生成边缘概率路径，我们需要用到一个基本工具：Kolmogorov Forward Equation，简称 KFE。

Kolmogorov 前向方程可以看作离散状态空间中的连续性方程。前面在连续模型中，我们见过 continuity equation 或 Fokker-Planck equation，它们描述概率密度如何随时间变化。而在 CTMC 中，对应的概率演化方程就是 Kolmogorov Forward Equation。

具体来说，如果一个 CTMC 的速率矩阵为 $Q_t$ ，并且它在时间 $t$ 的分布为：

$$X_t\sim p_t$$

那么当且仅当下面的方程成立时，CTMC 会沿着概率路径 $p_t$ 演化：

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

这个式子的直观含义是：状态 $x$ 上的概率变化率等于所有其他状态流入 $x$ 的净概率流。

右侧每一项：

$$Q_t(x\mid y)p_t(y)$$

表示从状态 $y$ 流入状态 $x$ 的概率流。因为当前处于 $y$ 的概率是 $p_t(y)$ ，而从 $y$ 跳到 $x$ 的瞬时速率是 $Q_t(x\mid y)$ ，所以二者相乘就是从 $y$ 到 $x$ 的概率流强度。对所有 $y\in S$ 求和，就得到状态 $x$ 的净流入。

这里也包括 $y=x$ 的项。由于对角项 $Q_t(x\mid x)$ 是负的，它实际上表示从 $x$ 流出的总速率。因此，整个求和项同时包含了流入和流出，最终得到状态 $x$ 上概率质量的净变化。

现在我们用 KFE 来证明前面定义的边缘速率矩阵确实会生成边缘概率路径。

边缘概率路径定义为：

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

对时间求导：

$$\begin{array}{rl}\frac{d}{dt}{p}_t(x)& =\frac{d}{dt}\sum _{z\in S}{p}_t(x\mid z)p_{\mathrm{data}}(z)\\ & =\sum _{z\in S}\frac{d}{dt}{p}_t(x\mid z)p_{\mathrm{data}}(z)\end{array}$$

由于 $Q_t^z$ 是条件概率路径 $p_t(\cdot \mid z)$ 的 conditional rate matrix，因此对条件路径应用 KFE，有：

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

代入上式得到：

$$\begin{array}{rl}\frac{d}{dt}{p}_t(x)& =\sum _{z\in S}\left[\sum _{y\in S}{Q}_t^z(x\mid y)p_t(y\mid z)\right]p_{\mathrm{data}}(z)\\ & =\sum _{y\in S}\sum _{z\in S}{Q}_t^z(x\mid y)p_t(y\mid z)p_{\mathrm{data}}(z)\end{array}$$

接下来，我们在每一项中乘以并除以 $p_t(y)$ ：

$$\begin{array}{rl}\frac{d}{dt}{p}_t(x)& =\sum _{y\in S}{p}_t(y)\left[\sum _{z\in S}{Q}_t^z(x\mid y)\frac{p_t(y\mid z)p_{\mathrm{data}}(z)}{p_t(y)}\right]\end{array}$$

注意括号中的部分正是边缘速率矩阵的定义：

$$Q_t(x\mid y)=\sum _{z\in S}{Q}_t^z(x\mid y)\frac{p_t(y\mid z)p_{\mathrm{data}}(z)}{p_t(y)}$$

因此：

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

这正是 Kolmogorov 前向方程。于是我们证明了：通过离散边缘化技巧得到的 $Q_t$ 确实会使 CTMC 沿着边缘概率路径 $p_t$ 演化。

这个证明和连续 flow matching 中的 marginal vector field 证明非常相似。核心思想都是：先构造依赖于数据点 $z$ 的条件动力学，再用后验权重对它做边缘化，从而得到真正可用于生成的边缘动力学。

接下来，我们把上面的抽象公式具体化到前面介绍过的 factorized mixture path 上。

回顾一下，factorized mixture path 的条件概率路径为：

$$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]$$

其中 ${\kappa }_t$ 是调度函数，满足：

$${\kappa }_0=0,{\kappa }_1=1$$

随着时间 $t$ 从 0 增加到 1， ${\kappa }_t$ 逐渐增加，表示数据项的权重越来越大，而噪声项的权重 $1-{\kappa }_t$ 逐渐减小。

对于这条路径，可以构造一个非常简单的 conditional rate matrix。它是 factorized 的，也就是说它只允许一次跳变改变一个 token。对于序列中的第 $j$ 个位置，设当前 token 为 $x_j$ ，目标 token 为 $z_j$ 。对于候选 token $v_i\in \mathcal{V}$ ，条件速率矩阵可以写成：

$$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)$$

这里 $\frac{{\dot{\kappa }}_t}{1-{\kappa }_t}$ 是一个随时间变化的因子，控制跳转速率的大小。

这个公式看起来有点抽象，但它的含义非常简单：

如果当前 token 已经是正确的，也就是：

$$x_j=z_j$$

那么：

$${\delta }_{z_j}(v_i)-{\delta }_{x_j}(v_i)=0$$

因此所有跳转速率都为 0。直观上，如果当前位置已经等于目标 token，就不需要再修改它。

如果当前 token 是错误的，也就是：

$$x_j\ne z_j$$

那么对于目标 token $v_i=z_j$，有：

$${\delta }_{z_j}(v_i)=1,{\delta }_{x_j}(v_i)=0$$

所以跳到正确 token 的速率为：

$$\frac{{\dot{\kappa }}_t}{1-{\kappa }_t}$$

对于其他既不是当前 token、也不是目标 token 的候选 token，速率为 0，因为模型不应该跳到错误的位置。

对于当前 token 本身 $v_i=x_j$，会得到负的对角项：

$$-\frac{{\dot{\kappa }}_t}{1-{\kappa }_t}$$

这表示从当前错误 token 流出的速率。它不是实际跳转到自己的非对角速率，而是由速率矩阵的对角项定义得到的，用来保证概率质量守恒。

因此，这个 conditional rate matrix 可以理解为：如果当前位置已经正确，就不动；如果当前位置错误，就以一定速率跳到正确 token。

上图展示了 ${\kappa }_t$ 和对应速率因子 $\frac{{\dot{\kappa }}_t}{1-{\kappa }_t}$ 的变化。以最简单的线性调度 ${\kappa }_t=t$ 为例，有：

$${\dot{\kappa }}_t=1$$

因此：

$$\frac{{\dot{\kappa }}_t}{1-{\kappa }_t}=\frac{1}{1-t}$$

当 $t$ 趋近于 1 时，这个因子会趋于无穷大。直观上，这是因为在时间 $t=1$ ，条件概率路径必须完全收敛到目标数据点 $z$ 。如果某个 token 到最后还没有变成正确 token，那么它必须以越来越大的速率跳到正确位置，才能保证终点分布是 ${\delta }_z$ 。

因此，速率因子在 $t\to 1$ 时爆炸，并不是异常现象，而是为了强制路径在终点精确到达数据点。

总结一下，对于 factorized mixture path，我们有：

首先是条件概率路径：

$$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]$$

它在初始分布和数据点 $z$ 之间插值。

然后是条件速率矩阵：

$$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)$$

它会让 CTMC 沿着这条条件概率路径演化。这个速率矩阵的结构非常简单：对每个位置独立判断当前 token 是否正确，如果错误，就推动它跳向目标 token；如果已经正确，就保持不变。

接下来，我们真正需要学习的是边缘速率矩阵。条件速率矩阵依赖于已知终点 $z$ ，但生成时我们不知道 $z$ ，因此必须通过边缘化得到不依赖 $z$ 的 rate matrix。

6.3 Learning the Marginal Rate Matrix

现在我们讨论如何学习边缘速率矩阵。对于 factorized mixture path，边缘速率矩阵仍然是 factorized 的，因为每个条件速率矩阵 $Q_t^z$ 都只允许单 token 更新，而对这些条件速率矩阵做后验加权平均后，非邻居状态之间的跳转速率仍然为 0。

因此，边缘速率矩阵可以写成：

$$Q_t(y\mid x)={\left(Q_t(v_i,j\mid x)\right)}_{v_i,j}$$

其中 $Q_t(v_i,j\mid x)$ 表示：在当前完整 noisy sequence 为 $x$ 时，将第 $j$ 个位置替换为 token $v_i$ 的速率。

根据离散边缘化技巧，可以得到：

$$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)$$

这个公式非常关键。它说明边缘速率矩阵和条件速率矩阵具有相同形式，只是原来条件速率矩阵中的已知目标 token $z_j$ ，现在被替换成了一个后验概率：

$$p_{1\mid t}(z_j=v_i\mid x)$$

这个后验概率表示：给定当前带噪序列 $x$ ，最终干净数据序列的第 $j$ 个 token 等于 $v_i$ 的概率。

也就是说，生成时我们不知道真实终点 $z$ ，所以不能直接说 “跳到 $z_j$” 。但我们可以估计 “$z_j$ 最可能是什么”。于是，模型会根据这个后验概率决定当前 token 应该跳向哪些候选 token。

简单证明如下。根据边缘速率矩阵定义：

$$Q_t(y\mid x)=\sum _{z\in S}{Q}_t^z(y\mid x)p_{1\mid t}(z\mid x)$$

如果 $y$ 和 $x$ 不是邻居，也就是二者在超过一个 token 位置上不同，那么对任意 $z$ 都有：

$$Q_t^z(y\mid x)=0$$

因此：

$$Q_t(y\mid x)=0$$

所以边缘速率矩阵仍然是 factorized 的。

对于单个位置 $j$ 和候选 token $v_i$ ，有：

$$\begin{array}{rl}{Q}_t(v_i,j\mid x)& =\sum _{z\in S}{Q}_t^z(v_i,j\mid x)p_{1\mid t}(z\mid x)\\ & =\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\mid t}(z\mid x)\\ & =\frac{{\dot{\kappa }}_t}{1-{\kappa }_t}\left[\sum _{z\in S}{\delta }_{z_j}(v_i)p_{1\mid t}(z\mid x)-{\delta }_{x_j}(v_i)\sum _{z\in S}{p}_{1\mid t}(z\mid x)\right]\\ & =\frac{{\dot{\kappa }}_t}{1-{\kappa }_t}\left[p_{1\mid t}(z_j=v_i\mid x)-{\delta }_{x_j}(v_i)\right]\end{array}$$

其中使用了：

$$\sum _{z\in S}{p}_{1\mid t}(z\mid x)=1$$

以及：

$$\sum _{z\in S}{\delta }_{z_j}(v_i)p_{1\mid t}(z\mid x)=p_{1\mid t}(z_j=v_i\mid x)$$

证明完成。

这个推导告诉我们，离散 flow matching 的学习目标最终可以转化成一个分类问题：给定 noisy sequence $x$ 和时间 $t$ ，预测每个位置的原始干净 token 是什么。

现在我们就可以定义离散流匹配的训练目标。

由于边缘速率矩阵中唯一未知的部分是后验概率 $p_{1\mid t}(z_j\mid x)$ ，因此我们用一个神经网络来参数化它：

$$p_{1\mid t}^{\theta }(z_j\mid x)$$

这个网络的输入是当前带噪序列 $x$ 和时间 $t$，输出是每个位置上的 token 分类分布。对于每个位置 $j$ ，网络输出一个定义在词表 $\mathcal{V}$ 上的概率分布，用来预测原始数据 token $z_j$ 。

训练这个后验网络非常直接：因为训练时我们知道真实数据样本 $z$ ，所以可以把每个位置的 $z_j$ 当作分类标签，然后使用交叉熵损失。

离散流匹配损失定义为：

$${\mathcal{L}}_{\mathrm{DFM}}(\theta )={\mathbb{E}}_{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]$$

这个损失的含义是：

首先从数据分布中采样一个真实序列：

$$z\sim p_{\mathrm{data}}$$

然后随机采样一个时间：

$$t\sim \mathrm{Unif}[0,1]$$

接着从条件概率路径中采样一个带噪序列：

$$x\sim p_t(\cdot \mid z)$$

最后让网络根据带噪序列 $x$ 和时间 $t$ ，预测每个位置的原始 token $z_j$ ，并对所有位置计算交叉熵损失。

从形式上看，这和普通的序列分类任务非常相似。区别在于，这里的输入 $x$ 是按照时间 $t$ 随机破坏后的序列，而模型需要在不同噪声水平下恢复原始 token。

因此，和连续 flow matching “最终变成回归问题” 类似，离散 flow matching 最终变成了一个分类问题：连续模型学习向量场，通常用均方误差回归；离散模型学习 token 后验分布，通常用交叉熵分类。

一旦训练好了后验网络：

$$p_{1\mid t}^{\theta }(z_j=v_i\mid x)$$

就可以通过下面的公式构造速率矩阵：

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

然后就可以使用前面介绍过的 factorized CTMC sampling 算法进行生成。

OK，我们来看训练算法：

上图给出了训练 factorized CTMC model 的完整算法。这个算法本质上就是离散扩散模型或离散流匹配模型的训练过程。

输入包括：

• 一个离散序列数据集：$z=(z_1,\cdots ,z_d)\in {\mathcal{V}}^d,z\sim p_{\mathrm{data}}$
• 每个位置上的初始噪声分布：$p_{\mathrm{init}}^{(j)}$
• 调度函数：${\kappa }_t\in [0,1]$
• 后验网络 $f_{\theta }$ ，它对每个位置输出词表上的 logits；
• 优化器。

每次训练迭代的过程如下：

首先采样一个真实数据序列：

$$z\sim p_{\mathrm{data}}$$

然后采样时间：

$$t\sim \mathrm{Unif}[0,1]$$

并计算：

$$\kappa ={\kappa }_t$$

接下来，根据 factorized mixture path 构造带噪序列 $x$ 。对于每个位置 $j$ ，采样一个 mask 变量：

$$m_j\sim \mathrm{Bernoulli}(\kappa )$$

再从初始噪声分布中采样一个噪声 token：

$${\xi }_j\sim p_{\mathrm{init}}^{(j)}$$

然后设置：

$$x_j=m_j{z}_j+(1-m_j){\xi }_j$$

这个公式表示：如果 $m_j=1$ ，当前位置保留真实 token $z_j$ ；如果 $m_j=0$ ，当前位置被替换成噪声 token ${\xi }_j$ 。对所有位置执行这个过程，就得到带噪序列：

$$x=(x_1,\cdots ,x_d)$$

然后将带噪序列 $x$ 和时间 $t$ 输入后验网络：

$${\ell }_j(\cdot )=f_{\theta }(x,t{)}_j$$

其中 ${\ell }_j$ 是第 $j$ 个位置上对整个词表的 logits。经过 softmax 后得到：

$$p_{1\mid t}^{\theta }(v\mid x{)}_j=\mathrm{Softmax}({\ell }_j)(v)$$

最后计算 token-wise negative log-likelihood：

$${\mathcal{L}}_{\mathrm{DFM}}(\theta )=\sum _{j=1}^d-\log p_{1\mid t}^{\theta }(z_j\mid x{)}_j$$

并用梯度下降更新参数：

$$\theta \leftarrow Opt.step\left({\nabla }_{\theta }{\mathcal{L}}_{\mathrm{DFM}}(\theta )\right)$$

这个训练过程和普通 masked language modeling 很像：随机破坏输入序列，然后训练模型恢复原始 token。不同之处在于，这里的破坏比例由连续时间 $t$ 和调度函数 ${\kappa }_t$ 控制，因此模型学习的是整个时间路径上的去噪能力，而不是单一噪声水平下的恢复能力。

一种非常常见的离散扩散语言模型是 mask diffusion language model。

它的做法是在词表中引入一个新的特殊 token $[\mathrm{MASK}]$ ，这个 token 表示当前位置的原始 token 被遮盖了。此时，初始分布可以定义为集中在 mask token 上的点质量分布：

$$p_{\mathrm{init}}={\delta }_{[\mathrm{MASK}]}$$

也就是说，在生成开始时，所有位置都是 mask：

$$X_0=([\mathrm{MASK}],\cdots ,[\mathrm{MASK}])$$

这和使用均匀分布作为初始噪声不同。均匀分布会随机给每个位置采样一个 token，而 mask 初始分布则让所有位置从同一个特殊 token 开始。这样做的好处是非常直观：模型一开始看到的是一段完全被遮盖的文本，然后逐步将 mask token 替换成真实文本 token。

在训练时，factorized mixture path 也可以理解为随机 mask 数据序列。对于每个位置，以概率 ${\kappa }_t$ 保留真实 token，以概率 $1-{\kappa }_t$ 替换为 $[\mathrm{MASK}]$。模型的任务就是根据当前部分被 mask 的序列，预测每个位置的原始 token。

这种方式和 BERT 式 masked language modeling 有相似之处，但区别在于：mask diffusion language model 把 mask 比例组织成一条连续时间概率路径，并在采样时通过多步 CTMC 逐渐从全 mask 序列生成完整文本。

上图展示了一个实际的扩散语言模型 Demo，例如 LLaDA 这类 large language diffusion model。给定一个 prompt 后，模型并不是像自回归模型那样从左到右一个 token 一个 token 地生成，而是从一段包含大量 mask 的序列开始，然后在多步迭代中逐渐解除 mask，生成最终文本。

这类模型的一个重要特点是：它可以在多个位置上并行更新 token。因此，文本生成顺序不再必须是从左到右，而可以是任意顺序。某些 token 可能较早被确定，某些 token 可能较晚被确定，模型可以根据当前上下文和后验概率决定哪些位置应该更新。

OK，我们再展示一下完整流程：

生成开始时，序列完全被 mask。此时所有位置的信息都还没有被恢复，模型只知道 prompt 或其他条件信息，以及当前所有待生成位置都是 $[\mathrm{MASK}]$ 。

在中间时刻，例如 $t=0.6$，一部分 token 已经被恢复出来，而另一部分位置仍然保持 mask。此时模型已经形成了部分文本结构，并会利用这些已恢复 token 继续推断剩余位置。

当 $t=1.0$ 时，所有位置都被恢复为真实 token，最终得到完整文本。

这个过程很好地体现了扩散语言模型和自回归语言模型的区别。自回归模型的生成路径是固定的，从左到右逐 token 生成；而 mask diffusion language model 的生成路径更加灵活，可以逐步填充、修改和完善整段文本。

---

Q & A

Q：模型如何决定哪些 token 保留、哪些 token 更新？

A：模型根据当前状态下的速率矩阵来决定。对于每个位置，速率越高，表示该位置越有可能发生变化。在 mask diffusion language model 中，这个速率来自后验概率网络：

$$p_{1\mid t}^{\theta }(z_j=v_i\mid x)$$

如果模型认为某个 mask 位置非常确定应该变成某个 token，那么对应的跳转速率就会较高，这个位置就更可能被更新。如果某个位置已经比较确定，或者保持当前 token 的概率较高，那么它发生变化的概率就较低。

因此，采样过程本质上是在每一步根据 categorical distribution 随机决定每个位置的下一个 token。这个分布由当前序列、时间和模型预测的后验概率共同决定。

---

最后，我们讨论一下离散扩散模型和自回归模型之间的优缺点对比。

首先是优势：

离散扩散模型的第一个潜在优势是可以并行生成多个 token。自回归模型通常必须按照从左到右的顺序逐 token 生成，每一步都依赖前一步的输出，因此生成过程天然是串行的。而离散扩散模型可以在一个采样步骤中同时更新多个位置，因此理论上可能带来更高的生成速度。

第二个优势是生成顺序更加灵活。自回归模型通常固定从左到右生成，而离散扩散模型可以以任意顺序生成 token。这对于文本编辑、代码补全、局部重写等任务非常有吸引力。例如，如果只想修改一段文本或代码中间的一小部分，扩散式模型可以固定不需要修改的位置，只对被 mask 的区域进行生成或修复。

第三个优势是可以设计更加灵活的概率路径。当前我们看到的是比较简单的 factorized mixture path，也就是逐 token 从噪声变成数据。但原则上，我们可以设计更复杂、更有语义意义的 probability path。例如，是否可以设计一种路径，让模型先确定句子结构，再填充细节？或者先生成摘要级语义，再逐步细化为完整文本？这些都是离散扩散模型提供的新可能性。

不过，它也有明显的挑战和缺点：

第一个问题是 KV cache 不如自回归模型直接。现代大语言模型的高效推理很大程度上依赖 KV cache，因为自回归生成时每一步只新增一个 token，前面 token 的 key/value 可以缓存并复用。而扩散语言模型在每一步中可能同时修改多个位置，整个上下文都会变化，因此如何高效缓存和复用中间计算更加困难。这可能抵消并行生成带来的速度优势。

第二个问题是学习难度可能更高。自回归模型的生成顺序非常自然：从左到右，这和人类书写文本的顺序一致，也和语言的因果结构较为匹配。而离散扩散模型需要学习如何以任意顺序生成 token，这可能让训练问题变得更复杂。

第三个问题是，自回归顺序本身具有语义意义。对于自然语言来说，从左到右并不是一个完全任意的顺序。句子通常按照语义和语法结构展开，前文为后文提供上下文。因此，放弃自回归顺序是否值得，还需要根据具体任务和模型能力来判断。

所以，离散扩散模型并不是一定优于自回归模型，而是提供了一种不同的建模方式。它在并行生成、文本编辑、局部修复、任意顺序生成等任务上可能更有优势；而自回归模型在工程效率、训练稳定性和语言建模自然性方面仍然非常强大。

此外，对于某些数据类型，自回归顺序可能并不自然。例如蛋白质序列、某些结构化离散数据，或者图结构数据，并不一定存在类似自然语言那样明确的从左到右语义顺序。在这些任务中，离散扩散或离散流匹配可能会更加合适。

最后，课程提到一个更一般的问题：为什么 flow matching 的配方可以如此自然地从连续数据推广到离散数据和 CTMC？

原因是，flow matching 背后的原则其实比 ODE 或 CTMC 更一般。它并不局限于某一种具体的动力系统，而是可以推广到更一般的 Markov process。

在连续 flow matching 中，动力学过程是 ODE，其生成器对应于向量场；在扩散模型中，动力学过程可以是 SDE，其生成器包含 drift 和 diffusion；在离散 flow matching 中，动力学过程是 CTMC，其生成器就是速率矩阵。

这些看起来不同的对象，本质上都可以放在 Markov process 的统一框架下理解。只要一个随机过程可以通过某种 generator 描述其分布演化，就有可能构造类似的 matching objective 来学习这个 generator。

课程中提到的 Generator Matching 就是试图从更一般的角度统一这些方法。它说明，我们之前学习的 ODE、SDE、CTMC 并不是彼此孤立的模型，而是同一个更大框架下的不同实例。

从这个角度看，本课程的主线可以理解为：从连续流模型和扩散模型出发，我们学习如何用微分方程描述生成过程；然后通过 flow matching，把生成建模转化为学习动力学场；最后在离散空间中，用 CTMC 和 rate matrix 替代连续向量场，将同样的学习原则推广到语言、蛋白质序列等离散数据上。

7. Class Recap

在这一讲中，我们首先定义了 CTMC 和 factorized CTMC，然后讨论了如何从简单初始分布采样，并通过 Euler / $\tau$-leaping 近似模拟 CTMC。接着，我们构造了离散概率路径，特别是 factorized mixture path，用来连接初始噪声分布和数据分布。然后通过 Kolmogorov forward equation 和 discrete marginalization trick，推导出边缘速率矩阵。最后，我们看到离散流匹配的训练目标最终可以写成一个 token-wise classification loss，也就是给定带噪序列，预测每个位置的原始 token。

总结来说，第五讲的核心结论是：离散扩散模型并不是简单地把连续 SDE 搬到离散空间中，而是把扩散模型和流匹配中的学习思想推广到了离散状态空间。CTMC 提供了离散空间中的动力学描述，rate matrix 扮演了向量场的角色，而 posterior token prediction 则成为训练这种模型的关键。这样一来，我们就可以用类似 flow matching 的思想来建模语言、蛋白质序列等离散数据。

至此，第五讲也接近尾声。回顾整门课程：

Lecture 1 介绍了 flow models 和 diffusion models，帮助我们理解生成模型如何从简单噪声分布出发，通过连续动力学过程生成复杂数据。

Lecture 2 重点讲解了 flow matching，把生成建模问题转化为学习一个向量场的问题。通过构造概率路径和对应的向量场，我们可以用简单的回归目标训练生成模型。

Lecture 3 讨论了 score matching 和 guidance，解释了扩散模型中 score 的作用，以及如何通过条件信息控制生成过程，例如 classifier-free guidance。

Lecture 4 进一步进入实际图像和视频生成器的构建，包括 latent space、VAE、DiT、Transformer 架构以及高效生成模型的设计。

Lecture 5 则将前面这些连续空间中的思想推广到离散状态空间。我们看到，文本和蛋白质序列等离散数据不能直接用 SDE 或 ODE 描述，因为离散空间中没有连续轨迹和向量场。但是，我们可以使用 CTMC 来描述离散状态之间的随机跳变，并用 rate matrix 替代 continuous vector field。

整门课程最终形成了一条非常清晰的主线：生成建模的核心不是某一个具体模型，而是如何构造从简单分布到数据分布的演化过程，并学习驱动这个过程的动力学对象。在连续空间中，这个对象可以是向量场或 score；在离散空间中，它可以是 CTMC 的速率矩阵；在更一般的框架下，它们都可以被理解为 Markov process 的 generator。

结语

本节课将前面学习的 Flow Matching 与 Diffusion Models 从连续空间推广到了离散状态空间，介绍了 Discrete Diffusion Models 与 Discrete Flow Matching 的核心思想。由于文本、蛋白质序列等离散数据不存在连续轨迹和向量场，因此我们引入 Continuous-Time Markov Chain（CTMC），使用 rate matrix 描述离散状态之间的随机跳转过程。

通过构造离散概率路径、conditional rate matrix 以及 marginal rate matrix，我们证明了 Flow Matching 的思想并不依赖于连续空间，而可以推广到更一般的随机过程。在连续空间中，我们学习的是 vector field；而在离散空间中，我们学习的是 transition rate，本质上都是在寻找驱动数据分布演化的动力学生成器。

进一步地，我们看到离散 Flow Matching 最终可以转化为一个 token prediction 问题：模型通过预测带噪序列中每个位置的原始 token 后验概率，实现从噪声序列逐步恢复数据序列。这也构成了现代扩散语言模型的核心思想，使文本生成不再局限于传统自回归的逐 token 生成方式。

回顾整门 MIT S.184 课程，从 ODE/SDE、Flow Matching、Score Matching，到 Latent Diffusion、Transformer 架构，再到离散 CTMC 生成模型，我们建立了一个统一的生成建模视角：

生成模型的本质，是学习一个能够推动概率分布从简单分布演化到复杂数据分布的动力学过程。

无论是连续空间中的向量场、score function，还是离散空间中的 rate matrix，它们都是这一核心思想在不同数据空间中的具体表现。

至此，本课程完整展示了现代生成式 AI 背后的数学基础与工程实现，为进一步理解和研究图像、文本、视频以及多模态生成模型奠定了坚实基础。

完结撒花！🤗

参考

• https://diffusion.csail.mit.edu/