Lecture 05:Discrete Diffusion Models and Discrete Flow Matching

本讲主线

前几讲主要讨论连续空间 上的 flow / diffusion:用 ODE、SDE、vector field、score 来从 noise 生成 data。本讲把这些思想搬到离散数据上,例如语言 token、蛋白质序列、DNA 序列等。

关键变化是:离散空间里没有真正的 ODE / SDE,也没有“沿着某个方向移动”的 vector field。我们改用 continuous-time Markov chain, CTMC,并把连续 flow matching 中的 probability path、conditional target、marginalization trick、simulation-free training 翻译成离散版本。


0. 和前几讲的连接

连续空间里的生成模型目标是:

或者使用 SDE:

这里 ,所以可以谈“速度”“方向”“梯度”“扩散”。但语言模型里的对象通常是 token 序列:

其中 是词表, 是序列长度。这个空间是离散的,所以不能直接写:

因为 token 没有连续坐标,也没有可以求导的路径。

本讲核心替换

连续空间:vector field / ODE / SDE 负责让样本在空间中移动。

离散空间:rate matrix / CTMC 负责让样本在不同离散状态之间跳转。


1. 离散 diffusion 的基本设定

1.1 数据空间

设离散状态空间为:

对于长度为 的 token 序列,通常有:

数据分布是离散概率质量函数:

训练时我们不知道完整的 ,只拿到样本:

例如:

  • 自然语言:互联网上的文本序列。
  • 蛋白质:氨基酸序列。
  • DNA:碱基序列。

目标仍然是生成:

1.2 为什么不能直接说 diffusion / flow?

严格数学意义上:

  • 离散空间没有 Brownian motion 驱动的 SDE。
  • 离散空间没有 ODE 中的连续轨迹。
  • 离散空间没有“从当前位置沿某个方向走一点”的 vector field。

所以这里的 “discrete diffusion models” 是机器学习中的叫法;数学上更准确的模型是:

直觉

连续 diffusion:概率质量像流体一样在空间里移动、扩散。

离散 diffusion:概率质量不是连续移动,而是在状态之间“跳跃”或“传送”。slides 中也强调:probability is teleported, not moved through space。


2. Continuous-Time Markov Chain(CTMC)

2.1 CTMC 的直觉

CTMC 是一个连续时间随机过程:

它的轨迹是分段常数的:

  • 大部分时间停在某个离散状态。
  • 随机时刻发生跳跃。
  • 跳到另一个离散状态。

例如状态空间是 时,一条路径可能是:

这和连续 ODE 轨迹非常不同:CTMC 的路径不是平滑曲线,而是 piecewise constant jumps。

2.2 Rate matrix:离散版 vector field

CTMC 的核心对象是 rate matrix:

含义是:在时间 ,当前状态为 时,跳到状态 的瞬时速率。

rate matrix 满足两个条件:

因此对固定起点 :

记号方向很重要

表示从 跳到 的 rate。不要把 和 的方向写反。

如果把 看成矩阵,并用列表示当前状态 ,那么每一列求和为 。

2.3 小时间步近似

对于很小的 ,CTMC 的转移概率可以近似写成:

忽略高阶项时得到 Euler 近似:

这就是 CTMC 采样的基本数值方法:当前位置是 ,构造离散分布 ,然后从这个 categorical distribution 采样下一状态。

2.4 二状态 CTMC 例子

设:

若从 跳到 的 rate 是 ,从 跳到 的 rate 也是 ,则 rate matrix 可以写成:

经过时间 的转移概率为:

当 时,转移概率收敛到:

也就是说,不管起点是 还是 ,长期分布都会变成 uniform distribution。


3. 用 CTMC 做生成模型

3.1 目标:从 noise 到 data

离散生成模型的目标与连续情形完全平行:

这里:

  • 是简单的离散初始分布,例如 uniform distribution 或全 mask 分布。
  • 是神经网络参数化的 rate matrix。
  • 采样时从 初始化,然后模拟 CTMC 到 。

语言模型中的例子

对于 masked diffusion language model,可以令初始分布是所有位置都是 [MASK]:

然后 CTMC 逐渐把 mask token 替换成真实 token,最后得到完整文本。

3.2 神经网络输出什么?

一个 general CTMC model 需要输出:

也就是:给定当前状态 ,输出从 跳到所有可能 的 rate。

但对于序列空间:

这个空间太大,不可能直接输出所有 的 rate。因此需要结构化假设:factorized CTMC。


4. Factorized CTMC:只允许一次改一个 token

4.1 Neighbor 的概念

对 token 序列:

如果 和 只在一个位置不同,就称它们是 neighbors。

例如:

是 neighbors,因为只改了第二个 token。

但:

和 不是 neighbors,因为改了两个位置。

4.2 Factorized rate matrix

Factorized CTMC 只允许单 token 更新:

当 和 不是 neighbors 时。

这让模型不需要输出整个 上的 rate,只需要对每个位置 输出把当前 token 改成每个 vocabulary token 的 rate。

记:

其中 表示:当前序列是 ,把第 个位置改成 token 的 rate。

为什么 factorized 很重要?

对语言序列,完整状态空间大小是 ,不可处理。

Factorized CTMC 把输出规模降到 ,这和普通语言模型输出每个位置的 vocabulary logits 类似。

4.3 Factorized CTMC 的 Euler / tau-leaping 采样

设当前序列为:

步长为:

对每个位置 ,令当前 token 为:

构造 per-position transition distribution:

然后对所有位置并行采样:

和真实 CTMC 的关系

真实 CTMC 在一个极小时间段内几乎只会发生一个 jump。Factorized parallel update 允许多个位置在同一步中同时更新,因此是 tau-leaping 近似。步长足够小时,这个近似在一阶上是合理的。


5. Kolmogorov Forward Equation(KFE)

5.1 KFE 是什么?

在连续 flow matching 中,我们用 continuity equation 描述概率密度如何随 vector field 演化。离散 CTMC 中对应的方程是 Kolmogorov Forward Equation:

含义:状态 的概率变化率等于所有其他状态流入 的净概率流。

其中:

  • :从 跳到 的 rate。
  • :当前处在 的概率。
  • 求和表示所有状态 对 的贡献。

5.2 KFE 的作用

KFE 给出了判断一个 CTMC 是否 follow 某个 probability path 的标准:

当且仅当 和 满足 KFE。

类比关系

连续 flow matching:

离散 flow matching:

所以 KFE 是离散版 continuity equation。


6. Discrete Flow Matching 的整体矩阵

连续 flow matching 的核心对象是:

  1. Conditional probability path。
  2. Marginal probability path。
  3. Conditional vector field。
  4. Marginal vector field。
  5. Conditional / marginal flow matching loss。

离散 flow matching 中对应为:

  1. Conditional probability path。
  2. Marginal probability path。
  3. Conditional rate matrix。
  4. Marginal rate matrix。
  5. Discrete flow matching loss。

一句话版

把连续空间中的 vector field 换成离散空间中的 rate matrix,其他 flow matching 的思想几乎不变。


7. Discrete probability path

7.1 Conditional probability path

给定一个数据点 ,conditional probability path 是一族分布:

它需要满足:

也就是说:

  • 时是 noise distribution。
  • 时集中到目标数据点 。

7.2 Marginal probability path

和连续情形一样,marginal path 通过对数据点 做 marginalization 得到:

因此:

关键直觉

Conditional path 是“从 noise 到某个固定数据点 ”。

Marginal path 是“从 noise 到整个数据分布 ”。


8. Factorized mixture path

8.1 定义

最重要的离散 probability path 例子是 factorized mixture path。

设 schedule:

对序列 和目标数据 ,定义:

解释:每个位置 独立地在 noise token distribution 和真实 token 之间插值。

  • 当 ,,所有 token 来自 noise distribution。
  • 当 ,,所有 token 都等于数据 token。

8.2 采样方式

从 采样可以写成:

然后:

最后:

不要误解

这个写法在 slides / notes 中用于表达选择关系,但 token 本身不是数字。更安全的理解是:如果 ,选真实 token ;如果 ,选 noise token 。

8.3 和 Gaussian probability path 的类比

连续 Gaussian path 是:

它逐渐降低 noise、增加 data signal。

离散 factorized mixture path 是:

它逐渐降低 noise token 的概率、增加真实 token 的概率。

不同点是:

  • Gaussian path 中 在连续空间中移动。
  • Factorized mixture path 中 token 不移动,只是在 noise token 与 clean token 之间随机切换。

9. Conditional rate matrix

9.1 定义目标

给定 conditional path ,我们希望找到 conditional rate matrix:

使得对应 CTMC 满足:

也就是说,conditional rate matrix 是离散版的 conditional vector field。

9.2 Factorized mixture path 的 conditional rate

对于 factorized mixture path,conditional rate matrix 也是 factorized 的,只允许单 token 更新。

对位置 ,候选 token ,有:

等价地,可以写成分段形式:

直觉:

  • 如果当前位置已经是正确 token,即 ,就不需要跳。
  • 如果当前位置不正确,就以 rate 跳到正确 token 。
  • diagonal 项为负数,保证 rate matrix 的列和为 。

9.3 rate explosion

若选择简单 schedule:

则:

当 时 rate 会爆炸。

数值注意

slides 中专门指出 rates explode at 。实际实现中通常需要避免精确取到 ,或者使用截断、稳定化的 schedule / sampler。


10. Discrete marginalization trick

10.1 Marginal rate matrix

连续 flow matching 中,conditional vector field 可以通过 marginalization 得到 marginal vector field。离散情形中完全类似:

定义 posterior:

则:

这就是离散版 marginalization trick

Conditional rate matrix 会把 noise 变成某个固定数据点 。

Marginal rate matrix 对所有可能 做 posterior 加权,最终把 noise 变成整个数据分布 。

10.2 为什么它 follow marginal path?

如果 follow conditional path:

那么由 KFE 可以证明,marginal rate matrix follow marginal path:

因此:

证明思路和连续版本类似:

  1. 写出 marginal path。
  2. 对时间求导。
  3. 对每个 conditional path 使用 KFE。
  4. 交换求和并整理。
  5. 得到 marginal KFE。

11. Factorized mixture path 的 marginal rate

11.1 从 conditional rate 到 posterior classifier

对 factorized mixture path,marginal rate 也保持 factorized,并且有一个非常关键的形式:

这里唯一未知的是:

这是在给定当前 noisy sequence 和时间 后,第 个位置最终 clean token 是 的 posterior probability。

关键结论

学习 marginal rate matrix 等价于学习每个 token 位置的 posterior classifier。

也就是说,离散 flow matching 最终变成一个分类问题:预测 clean token。

11.2 和 denoising 的关系

训练时我们知道原始 clean sequence ,也知道如何从 生成 noisy 。所以可以让网络根据 预测每个位置的 clean token:

这和 denoising diffusion 的思想完全一致:输入被腐蚀的数据,预测干净数据。


12. Discrete Flow Matching loss

12.1 Posterior probability network

令神经网络输出每个位置上的 vocabulary distribution:

经过 softmax 得到:

12.2 Loss:每个 token 一个 cross-entropy

Discrete Flow Matching 的训练目标是:

这就是对每个 token 位置做 cross-entropy classification。

和连续 FM 的差别

连续 FM:回归一个 target vector field。

离散 FM:分类一个 clean token posterior。

12.3 训练算法

一次训练迭代可以写成:

  1. 采样真实序列:
  1. 采样时间:
  1. 计算 schedule:
  1. 从 factorized mixture path 采样 noisy sequence:
  1. 网络输出每个位置的 clean-token logits:
  1. softmax 得到 posterior prediction:
  1. 最小化 token-wise negative log-likelihood:

13. Mask Diffusion Language Models

13.1 引入 [MASK] token

Masked diffusion language model 在 vocabulary 中加入一个特殊 token:

令初始分布为全 mask:

于是 时:

13.2 Factorized mixture path 在 mask 模型中的形式

因为 noise token 固定为 [MASK],所以每个位置只有两种情况:

训练时,模型看到部分 mask、部分真实 token 的序列,然后预测所有位置的原始 token。

13.3 采样过程

采样时从全 mask 开始:

随后模拟 CTMC / discrete sampler,逐渐把 mask 替换成 token。slides 中的例子展示了:

  • :几乎全是 mask。
  • :少量 token 被填入。
  • :句子结构逐渐出现。
  • :大部分内容已经可读。
  • :完整文本生成。

和自回归语言模型的差异

自回归模型通常从左到右生成:

Mask diffusion 可以以任意顺序填 token,甚至多个 token 并行生成。


14. Discrete diffusion vs autoregressive models

14.1 可能的优势

Discrete diffusion / masked diffusion LLM 的潜在优势:

  • 并行生成多个 token:可能比逐 token 自回归更快。
  • 任意顺序生成:不必固定 left-to-right。
  • 适合文本编辑:可以 mask 掉局部内容再重生成。
  • 可以设计新的 probability paths:也许能设计更符合语义的 noising / denoising 过程。

14.2 主要挑战

也有明显困难:

  • 没有 KV caching:自回归 Transformer 的推理加速优势不一定能直接用。
  • 学习更难:模型要学会在任意 mask pattern 下恢复 token。
  • left-to-right 顺序本身有语义合理性:自然语言强依赖前文,AR factorization 很适合这个结构。
  • 需要多步迭代采样:如果步数很多,并行优势可能被抵消。

一个核心研究问题

Discrete diffusion 能否在质量、速度、可控编辑能力之间找到比 autoregressive LLM 更好的 trade-off?这仍然是活跃研究方向。


15. 统一视角:Generator Matching

本讲最后指出:flow matching 的思想不仅适用于:

  • Euclidean ODE / SDE。
  • 离散 CTMC。

它其实可以推广到更一般的 Markov processes。这个统一框架被称为 Generator Matching。

在这个视角下:

  • 连续空间中的 vector field 是一种 generator。
  • 离散空间中的 rate matrix 是一种 generator。
  • 更一般的 Markov process 也有自己的 generator。

因此,flow matching 的核心思想是:

构造一个从 noise 到 data 的 probability path,然后学习一个 Markov process 的 generator,使该过程 follow 这条 path。

这个视角可以扩展到:

  • smooth manifolds 上的数据。
  • 混合状态空间,例如图像 + 文本。
  • jump processes。
  • 几何数据或科学数据。

16. 连续 FM 与离散 DFM 对照表

连续 Flow Matching离散 Flow Matching
状态空间 状态空间 或
ODE / SDECTMC
vector field rate matrix
conditional vector field conditional rate matrix
marginal vector field marginal rate matrix
continuity equationKolmogorov Forward Equation
regression losscross-entropy classification loss
Gaussian probability pathfactorized mixture path
denoise continuous noisepredict clean token posterior

17. 常见易错点

17.1 “Discrete diffusion” 不是真的 SDE diffusion

在离散状态空间里没有 Brownian motion,也没有 SDE。这里的 diffusion 是机器学习术语,数学对象是 CTMC。

17.2 的方向不要写反

表示从 跳到 ,不是从 跳到 。

KFE 中写的是:

这里 是从 流入 的 rate。

17.3 条件记号用 \mid

在 Obsidian 表格和 LaTeX 中,建议写:

不要写成 x|z,避免和 Markdown 表格分隔符冲突。

17.4 Factorized mixture path 不是线性插值 token

公式中的:

只是表达“二选一”。token 本身不是数值向量时,不应该把它理解为线性加法。

17.5 rate explosion

当 时:

时会爆炸,采样和训练实现中需要处理。


18. 本讲速记

最小记忆版

  1. 离散数据不能用 ODE / SDE;用 CTMC。
  2. CTMC 的核心是 rate matrix 。
  3. KFE 是离散版 continuity equation。
  4. Discrete flow matching 依然是:先设 probability path,再找 conditional rate,再 marginalize,再训练。
  5. Factorized mixture path 对每个 token 独立地在 noise 和 clean token 之间混合。
  6. Factorized mixture path 的 marginal rate 只需要 posterior 。
  7. 因此 DFM loss 就是 token-wise cross-entropy。
  8. Mask diffusion LLM 从全 [MASK] 开始,多步并行填 token。

19. 复习问题

  1. 为什么离散空间中不能直接使用 ODE / SDE?
  2. 的数学含义是什么?为什么 diagonal entry 是负数?
  3. KFE 和 continuity equation 的对应关系是什么?
  4. Factorized CTMC 为什么可以把输出规模从 降到 ?
  5. Factorized mixture path 中 起什么作用?
  6. 为什么 conditional rate matrix 在 时可能爆炸?
  7. Discrete marginalization trick 中 posterior 的意义是什么?
  8. 为什么 DFM loss 是 cross-entropy,而不是 MSE?
  9. Mask diffusion LLM 和 autoregressive LLM 的生成顺序有什么不同?
  10. Generator Matching 如何统一连续和离散生成模型?

20. 一页公式总结

CTMC rate matrix 条件:

小时间步转移近似:

Kolmogorov Forward Equation:

Discrete conditional probability path:

Marginal probability path:

Factorized mixture path:

Conditional rate for factorized mixture path:

Discrete marginalization trick:

Posterior definition:

Marginal rate for factorized mixture path:

Discrete Flow Matching loss: