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 的核心对象是:
- Conditional probability path。
- Marginal probability path。
- Conditional vector field。
- Marginal vector field。
- Conditional / marginal flow matching loss。
离散 flow matching 中对应为:
- Conditional probability path。
- Marginal probability path。
- Conditional rate matrix。
- Marginal rate matrix。
- 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:
因此:
证明思路和连续版本类似:
- 写出 marginal path。
- 对时间求导。
- 对每个 conditional path 使用 KFE。
- 交换求和并整理。
- 得到 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 训练算法
一次训练迭代可以写成:
- 采样真实序列:
- 采样时间:
- 计算 schedule:
- 从 factorized mixture path 采样 noisy sequence:
- 网络输出每个位置的 clean-token logits:
- softmax 得到 posterior prediction:
- 最小化 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 / SDE | CTMC |
| vector field | rate matrix |
| conditional vector field | conditional rate matrix |
| marginal vector field | marginal rate matrix |
| continuity equation | Kolmogorov Forward Equation |
| regression loss | cross-entropy classification loss |
| Gaussian probability path | factorized mixture path |
| denoise continuous noise | predict 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. 本讲速记
最小记忆版
- 离散数据不能用 ODE / SDE;用 CTMC。
- CTMC 的核心是 rate matrix 。
- KFE 是离散版 continuity equation。
- Discrete flow matching 依然是:先设 probability path,再找 conditional rate,再 marginalize,再训练。
- Factorized mixture path 对每个 token 独立地在 noise 和 clean token 之间混合。
- Factorized mixture path 的 marginal rate 只需要 posterior 。
- 因此 DFM loss 就是 token-wise cross-entropy。
- Mask diffusion LLM 从全
[MASK]开始,多步并行填 token。
19. 复习问题
- 为什么离散空间中不能直接使用 ODE / SDE?
- 的数学含义是什么?为什么 diagonal entry 是负数?
- KFE 和 continuity equation 的对应关系是什么?
- Factorized CTMC 为什么可以把输出规模从 降到 ?
- Factorized mixture path 中 起什么作用?
- 为什么 conditional rate matrix 在 时可能爆炸?
- Discrete marginalization trick 中 posterior 的意义是什么?
- 为什么 DFM loss 是 cross-entropy,而不是 MSE?
- Mask diffusion LLM 和 autoregressive LLM 的生成顺序有什么不同?
- 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: