离散扩散模型：用扩散的原理造语言模型

推导：$(\star)$ 等价于一阶展开式

把 $h\mapsto p_{t+h\mid t}(y\mid x)$ 在 $h=0$ 处做一阶泰勒展开。首先注意零时刻的取值：

$$ p_{t\mid t}(y\mid x) \;=\; \mathbb P(X_t = y\mid X_t = x)\;=\;\delt{x}(y), $$

因为「过了零时间之后还在原地」是必然事件。于是由 $(\star)$，

$$ \boxed{\ \mathbb P(X_{t+h}=y\mid X_t=x)\;=\;\delt{x}(y)\;+\;h\,\Rate_t(y\mid x)\;+\;o(h)\ } \tag{2.1} $$

反过来，若 (2.1) 成立，两边减去 $\delt{x}(y)$ 再除以 $h$ 并令 $h\to0$，立刻得回 $(\star)$。所以微分形式 $(\star)$ 与展开形式 (2.1) 完全等价。

(2.1) 是本讲最该背下来的一行。它的读法：

• 取 $y \ne x$：$\mathbb P(\text{在 }h\text{ 时间内跳到 } y) = h\,\Rate_t(y\mid x)+o(h)$。速率乘以时长就是概率——这就是「速率」二字的含义。
• 取 $y=x$：$\mathbb P(\text{还在原地}) = 1 + h\,\Rate_t(x\mid x)+o(h) = 1 - h\,\lambda_t(x)+o(h)$。

把 $\R^{\abs{\Vocab}}$ 中的行向量记法用起来，(2.1) 就是 $p_{t+h\mid t}(\cdot\mid x) = e_x + h\,\mathbf{\Rate}_t(x,\cdot) + o(h)$：转移概率 = one-hot 向量 + $h$ 乘以速率矩阵的第 $x$ 行。和欧拉法 $x + h\,\vf_t(x)$ 的形状一模一样，只不过「$+$」发生在概率单纯形里而不是状态空间里。

推导：条件 (1) 与条件 (2) 的验证

条件 (1)：非对角元非负。固定 $y\ne x$。函数 $h\mapsto p_{t+h\mid t}(y\mid x)$ 是概率，故恒 $\ge 0$；而它在 $h=0$ 的取值是 $\delt{x}(y)=0$（因为 $y\ne x$）。一个从 0 出发、之后始终非负的可导函数，其右导数必然非负：

$$ \Rate_t(y\mid x)=\lim_{h\downarrow0}\frac{p_{t+h\mid t}(y\mid x)-0}{h}\ \ge\ 0 . $$

条件 (2)：行和为零。直接计算：

$$ \begin{aligned} \sum_{y\ne x}\Rate_t(y\mid x) &\overset{(i)}{=} \sum_{y\ne x}\frac{\dd{}}{\dd{h}}\Big|_{h=0} p_{t+h\mid t}(y\mid x)\\ &\overset{(ii)}{=} \frac{\dd{}}{\dd{h}}\Big|_{h=0}\sum_{y\ne x} p_{t+h\mid t}(y\mid x)\\ &\overset{(iii)}{=} \frac{\dd{}}{\dd{h}}\Big|_{h=0}\Big(1-p_{t+h\mid t}(x\mid x)\Big)\\ &\overset{(iv)}{=} -\frac{\dd{}}{\dd{h}}\Big|_{h=0} p_{t+h\mid t}(x\mid x)\;=\;-\,\Rate_t(x\mid x). \end{aligned} $$

其中 (i) 是 $(\star)$；(ii) 交换有限求和与求导（$\Vocab$ 有限，恒可交换）；(iii) 用了「概率之和为 1」，即 $\sum_{y\in\Vocab}p_{t+h\mid t}(y\mid x)=1$，把 $y=x$ 那一项移到右边；(iv) 常数 1 求导为 0，再用 $(\star)$。

所以 $\Rate_t(x\mid x)=-\sum_{y\ne x}\Rate_t(y\mid x)$，正是条件 (2)。结论：每个 CTMC 至少有一个满足 $(\star)$ 的速率矩阵，而且它自动是合法的速率矩阵。

推导：定理 33 的证明（对应讲义附录 C）

第 1 步：把转移概率变成一个线性 ODE 的解。固定出发点 $x$ 与出发时刻 $t$，把 $t'\mapsto \big(p_{t'\mid t}(y\mid x)\big)_{y\in\Vocab}\in\R^{\abs{\Vocab}}$ 看成一个向量值函数。计算它的导数：

$$ \begin{aligned} \frac{\dd{}}{\dd{t'}}p_{t'\mid t}(y\mid x) &\overset{(i)}{=} \frac{\dd{}}{\dd{h}}\Big|_{h=0} p_{t'+h\mid t}(y\mid x)\\ &\overset{(ii)}{=} \frac{\dd{}}{\dd{h}}\Big|_{h=0}\sum_{z\in\Vocab} p_{t'+h\mid t'}(y\mid z)\,p_{t'\mid t}(z\mid x)\\ &\overset{(iii)}{=} \sum_{z\in\Vocab}\Big[\frac{\dd{}}{\dd{h}}\Big|_{h=0}p_{t'+h\mid t'}(y\mid z)\Big]\,p_{t'\mid t}(z\mid x)\\ &\overset{(iv)}{=} \sum_{z\in\Vocab} \Rate_{t'}(y\mid z)\,p_{t'\mid t}(z\mid x). \end{aligned} $$

(i) 只是把对 $t'$ 的导数写成对时间增量 $h$ 的导数；(ii) 是 Chapman–Kolmogorov 方程（对中间时刻 $t'$ 的状态 $z$ 做全概率分解，这一步用到马尔可夫性）；(iii) 交换有限求和与求导，注意 $p_{t'\mid t}(z\mid x)$ 与 $h$ 无关；(iv) 用 $(\star)$。

初值条件是 $p_{t\mid t}(y\mid x) = \delt{x}(y)$。所以：对固定的 $(x,t)$，转移概率向量是一个已知初值的线性常微分方程组的解，系数矩阵就是 $\mathbf{\Rate}^\top_{t'}$。

第 2 步：唯一性。线性 ODE（系数有界连续）在给定初值下解唯一（第一讲定理 3 / Picard–Lindelöf）。所以满足 $(\star)$ 的转移概率至多有一族。

第 3 步：存在性 —— 先解 ODE。反过来，对任意速率矩阵 $\Rate_t$，上述线性 ODE 一定有解 $p_{t'\mid t}(y\mid x)$。剩下要验证它确实是一族合法的转移核，即三条性质：

$$ \text{(a) } \sum_{y}p_{t'\mid t}(y\mid x)=1;\quad \text{(b) } p_{t'\mid t}(y\mid x)\ge 0;\quad \text{(c) Chapman–Kolmogorov 成立.} $$

(a) 归一性。在 $t'=t$ 时 $\sum_y \delt{x}(y)=1$ 成立。再看它的导数：

$$ \frac{\dd{}}{\dd{t'}}\sum_{y}p_{t'\mid t}(y\mid x) =\sum_{y}\sum_{z}\Rate_{t'}(y\mid z)p_{t'\mid t}(z\mid x) =\sum_{z}\Big[\underbrace{\sum_{y}\Rate_{t'}(y\mid z)}_{=\,0\ \text{（行和为零）}}\Big]p_{t'\mid t}(z\mid x)=0 . $$

总和的导数恒为 0 且初值为 1，故恒等于 1。这正是「行和为零」这个条件的用处：它保证概率质量守恒。

(b) 非负性。在 $t'=t$ 时非负。假设某时刻某个分量触到 0，即 $p_{t'\mid t}(y\mid x)=0$，此刻它的导数是

$$ \frac{\dd{}}{\dd{t'}}p_{t'\mid t}(y\mid x) =\sum_{z}\Rate_{t'}(y\mid z)p_{t'\mid t}(z\mid x) =\underbrace{\Rate_{t'}(y\mid y)\cdot 0}_{=0}+\sum_{z\ne y}\underbrace{\Rate_{t'}(y\mid z)}_{\ge 0}\underbrace{p_{t'\mid t}(z\mid x)}_{\ge 0}\ \ge\ 0 . $$

即：一旦某个分量降到 0，它只能不降或上升，无法穿过 0 变负。所以解始终非负。这正是「非对角元非负」这个条件的用处。

(c) Chapman–Kolmogorov。对 $t_0\le t_1\le t_2$，定义

$$ q_{t_2\mid t_0}(y\mid x) := \sum_{z\in\Vocab}p_{t_2\mid t_1}(y\mid z)\,p_{t_1\mid t_0}(z\mid x). $$

要证 $q_{t_2\mid t_0} = p_{t_2\mid t_0}$。先看初值：$t_2 = t_1$ 时 $q_{t_1\mid t_0}(y\mid x)=\sum_z\delt{z}(y)p_{t_1\mid t_0}(z\mid x)=p_{t_1\mid t_0}(y\mid x)$，与 $p_{t_2\mid t_0}$ 在 $t_2=t_1$ 处相同。再看它作为 $t_2$ 的函数满足的 ODE：

$$ \begin{aligned} \frac{\dd{}}{\dd{t_2}}q_{t_2\mid t_0}(y\mid x) &=\sum_{z}\Big[\frac{\dd{}}{\dd{t_2}}p_{t_2\mid t_1}(y\mid z)\Big]p_{t_1\mid t_0}(z\mid x)\\ &=\sum_{z}\sum_{\bar z}\Rate_{t_2}(y\mid \bar z)\,p_{t_2\mid t_1}(\bar z\mid z)\,p_{t_1\mid t_0}(z\mid x)\\ &=\sum_{\bar z}\Rate_{t_2}(y\mid\bar z)\Big[\sum_{z}p_{t_2\mid t_1}(\bar z\mid z)p_{t_1\mid t_0}(z\mid x)\Big] =\sum_{\bar z}\Rate_{t_2}(y\mid\bar z)\,q_{t_2\mid t_0}(\bar z\mid x). \end{aligned} $$

这与 $p_{t_2\mid t_0}$ 满足的 ODE 完全相同，且初值相同。再次由线性 ODE 解的唯一性，$q\equiv p$。$\square$

整个证明的味道：CTMC 的全部理论，被「转移概率满足一个 $\abs{\Vocab}$ 维线性 ODE」这一句话统摄。速率矩阵的两个条件恰好各自负责一条概率公理（行和为零 → 归一；非对角非负 → 非负）。

推导：把转移概率算出闭式

记 $f(h):=\mathbb P(X_{t+h}=a\mid X_t=a)$。由第 1 步得到的线性 ODE（这里 $\Rate$ 不依赖时间）：

$$ \begin{aligned} f'(h) &= \sum_{z\in\set{a,b}}\Rate(a\mid z)\,p_{t+h\mid t}(z\mid a)\\ &= \Rate(a\mid a)\,f(h) + \Rate(a\mid b)\,\big(1-f(h)\big) \quad&&\text{(利用 } p(b\mid a)=1-f)\\ &= -\lambda f(h) + \lambda\big(1-f(h)\big)\\ &= \lambda - 2\lambda f(h), \end{aligned} $$

初值 $f(0)=1$。这是一维线性 ODE。令 $g(h) := f(h)-\tfrac12$，则 $g' = -2\lambda g$，$g(0)=\tfrac12$，故 $g(h)=\tfrac12 e^{-2\lambda h}$，于是

$$ f(h) = \frac{1}{2}\Big(1+e^{-2\lambda h}\Big),\qquad \mathbb P(X_{t+h}=b\mid X_t=a) = 1-f(h)=\frac{1}{2}\Big(1-e^{-2\lambda h}\Big). $$

由 $a\leftrightarrow b$ 的对称性，整张转移概率表是

$$ \begin{pmatrix} 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{pmatrix} =\frac12\begin{pmatrix} 1+e^{-2\lambda h} & 1-e^{-2\lambda h}\\ 1-e^{-2\lambda h} & 1+e^{-2\lambda h}\end{pmatrix}. $$

验证 $(\star)$。对 $h$ 求导并取 $h=0$：$\frac{\dd{}}{\dd{h}}\big|_0 \frac12(1+e^{-2\lambda h}) = \frac12(-2\lambda) = -\lambda = \Rate(a\mid a)$ ✓；$\frac{\dd{}}{\dd{h}}\big|_0\frac12(1-e^{-2\lambda h}) = \lambda = \Rate(b\mid a)$ ✓。

读出两件事。（i）$h=0$ 时矩阵是单位阵——没过时间就没跳转。（ii）$h\to\infty$ 时 $e^{-2\lambda h}\to0$，

$$ P(h)\ \longrightarrow\ \begin{pmatrix}\tfrac12&\tfrac12\\[2pt]\tfrac12&\tfrac12\end{pmatrix}, $$

即链忘记了自己从哪出发，收敛到均匀平稳分布。指数项 $e^{-2\lambda h}$ 正是「对初始状态的记忆」的衰减；$\lambda$ 越大，翻转越频繁，遗忘越快。这是离散版的「OU 过程收敛到高斯平稳分布」。

推导（一）：直接展开法

这是最能看清楚「KFE 从哪来」的写法。设 $X_t\sim p_t$，即 $p_t(x)=\mathbb P(X_t=x)$。对增量 $h>0$ 用全概率公式按 $t$ 时刻的状态分解：

$$ \begin{aligned} p_{t+h}(x) &\overset{(i)}{=}\sum_{y\in\Vocab}\mathbb P\big(X_{t+h}=x,\ X_t=y\big)\\ &\overset{(ii)}{=}\sum_{y\in\Vocab} p_{t+h\mid t}(x\mid y)\,p_t(y)\\ &\overset{(iii)}{=}\sum_{y\in\Vocab}\Big[\delt{y}(x) + h\,\Rate_t(x\mid y) + o(h)\Big]p_t(y)\\ &\overset{(iv)}{=}\underbrace{\sum_{y}\delt{y}(x)p_t(y)}_{=\,p_t(x)}\; +\; h\sum_{y}\Rate_t(x\mid y)p_t(y)\;+\;o(h). \end{aligned} $$

其中 (i) 是把事件 $\set{X_{t+h}=x}$ 按 $X_t$ 的取值分割（这些事件互斥且穷尽）；(ii) 是条件概率的定义；(iii) 代入速率矩阵的一阶展开式 (2.1)；(iv) 分配求和，并用 $\sum_y\delt{y}(x)p_t(y)=p_t(x)$（只有 $y=x$ 那项非零），以及 $\Vocab$ 有限使得 $\sum_y o(h)p_t(y)=o(h)$。

移项、除以 $h$：

$$ \frac{p_{t+h}(x)-p_t(x)}{h} \;=\; \sum_{y\in\Vocab}\Rate_t(x\mid y)\,p_t(y)\;+\;\frac{o(h)}{h}. $$

令 $h\to0$，右端第二项消失，左端按定义收敛到 $\frac{\dd{}}{\dd{t}}p_t(x)$，得

$$ \boxed{\ \frac{\dd{}}{\dd{t}}p_t(x)=\sum_{y\in\Vocab}\Rate_t(x\mid y)\,p_t(y)\ } $$

这就是 KFE。整个推导只用了三样东西：全概率公式、马尔可夫性（藏在转移概率的定义里）、速率矩阵的一阶展开。

推导（二）：讲义的写法（等价，更简洁）

讲义给的是同一件事的紧凑版本，四个等号：

$$ \begin{aligned} \frac{\dd{}}{\dd{t}}p_t(x) &\overset{(i)}{=}\frac{\dd{}}{\dd{h}}\Big|_{h=0}p_{t+h}(x)\\ &\overset{(ii)}{=}\frac{\dd{}}{\dd{h}}\Big|_{h=0}\sum_{y}p_{t+h\mid t}(x\mid y)\,p_t(y)\\ &\overset{(iii)}{=}\sum_{y}\Big[\frac{\dd{}}{\dd{h}}\Big|_{h=0}p_{t+h\mid t}(x\mid y)\Big]p_t(y)\\ &\overset{(iv)}{=}\sum_{y}\Rate_t(x\mid y)\,p_t(y). \end{aligned} $$

(i) 把对 $t$ 的导数改写成对时间偏移 $h$ 的导数；(ii) 用转移概率的定义做全概率分解（注意 $p_t(y)$ 与 $h$ 无关）；(iii) 交换有限求和与求导；(iv) 用速率矩阵的定义 $(\star)$。

推导：KFE 是充分条件

设某族分布 $(q_t)$ 满足 KFE 且 $q_0=p_0$。把 KFE 写成向量形式

$$ \frac{\dd{}}{\dd{t}}q_t = \mathbf{\Rate}_t^\top q_t,\qquad q_t\in\R^{\abs{\Vocab}} , $$

这是有限维空间 $\R^{\abs{\Vocab}}$ 上的线性 ODE，系数矩阵 $\mathbf{\Rate}_t^\top$ 有界且关于 $t$ 连续，初值 $q_0=p_0$ 给定。由第一讲的定理 3（ODE 解的存在唯一性），这样的解唯一。

另一方面，由 3.1 的必要性部分，CTMC 的真实边际 $p_t(x):=\mathbb P(X_t=x)$ 也满足同一个 ODE、同一个初值。两个解相等：$q_t = p_t$ 对所有 $t$。故 $X_t\sim q_t$。$\square$

这正是本课第二个母题的答案在离散情形的形态：为什么两个看起来不同的过程会给出同一个边际分布？因为边际分布由一个线性 ODE 唯一决定，只要两个速率矩阵都满足同一个 KFE，它们的边际就必然逐点相同——速率矩阵中不影响 KFE 右端的那部分自由度，对边际完全没有影响。

推导：净流入形式 $$ \begin{aligned} \sum_{y\in\Vocab}\Rate_t(x\mid y)p_t(y) &\overset{(i)}{=}\sum_{y\ne x}\Rate_t(x\mid y)p_t(y)\;+\;\Rate_t(x\mid x)p_t(x)\\ &\overset{(ii)}{=}\sum_{y\ne x}\Rate_t(x\mid y)p_t(y)\;-\;\Big[\sum_{y\ne x}\Rate_t(y\mid x)\Big]p_t(x)\\ &\overset{(iii)}{=}\sum_{y\ne x}\Big[\underbrace{\Rate_t(x\mid y)\,p_t(y)}_{\text{从 }y\text{ 流入 }x}\;-\;\underbrace{\Rate_t(y\mid x)\,p_t(x)}_{\text{从 }x\text{ 流出到 }y}\Big]. \end{aligned} $$

(i) 拆出对角项；(ii) 用条件 (2)（行和为零）把 $\Rate_t(x\mid x)$ 换成 $-\sum_{y\ne x}\Rate_t(y\mid x)$；(iii) 合并两个对 $y\ne x$ 的求和。

于是 KFE 读作：

$$ \frac{\dd{}}{\dd{t}}p_t(x)=\sum_{y\ne x}\Big[\Rate_t(x\mid y)p_t(y)-\Rate_t(y\mid x)p_t(x)\Big]. $$

状态 $x$ 上概率质量的变化率 = 所有邻居流进来的 − 从 $x$ 流出去的。这就是「守恒律」的离散形态：把 $J_t(y\to x):=\Rate_t(x\mid y)p_t(y)$ 叫做从 $y$ 到 $x$ 的概率流（probability flux），则上式是 $\partial_t p_t(x) = -\big(\text{离散散度}\big)(x)$。在 $\R^d$ 上，流是 $j_t = p_t\vf_t$，守恒律是 $\partial_t p_t = -\divg j_t$；在 $\Vocab$ 上，流是矩阵 $J_t(y\to x)$，「散度」是对所有邻居求净流出。把 $\divg$ 换成「对邻居求和」，这就是从连续到离散的整个翻译字典的第一条。

推导：质量守恒（一个必须做的健全性检查）

KFE 必须保证 $\sum_x p_t(x)$ 恒为 1，否则 $p_t$ 就不是分布了。验证：

$$ \frac{\dd{}}{\dd{t}}\sum_{x\in\Vocab}p_t(x) =\sum_{x}\sum_{y}\Rate_t(x\mid y)p_t(y) =\sum_{y}p_t(y)\underbrace{\sum_{x}\Rate_t(x\mid y)}_{=\,0} =0 . $$

内层求和是「从固定出发点 $y$ 出去的所有速率之和」，正是条件 (2)。所以「行和为零」在 KFE 里的角色，与「向量场的散度形式保证 $\int\partial_t p_t = 0$」完全一样：它就是概率守恒。

推导：$\tilde p$ 什么时候真的是一个分布

要求两件事。

(a) 归一。精确成立，不需要任何条件：

$$ \sum_{y\in\Vocab}\tilde p_{t+h\mid t}(y\mid x)=\underbrace{\sum_y \delt{x}(y)}_{=1}+h\underbrace{\sum_y \Rate_t(y\mid x)}_{=0}=1 . $$

(b) 非负。分两种情况：

• $y\ne x$：$\tilde p = h\,\Rate_t(y\mid x)\ge 0$ 自动成立（条件 (1)）。
• $y=x$：$\tilde p = 1 + h\,\Rate_t(x\mid x) = 1 - h\,\lambda_t(x)$，要求

$$ \boxed{\;h\ \le\ \frac{1}{\lambda_t(x)}=\frac{1}{-\Rate_t(x\mid x)}\;} $$

也就是说：步长必须小于「在当前状态的平均停留时间」。速率越大（越急着跳走），步长必须越小。这是 Euler 格式在离散世界里的稳定性条件，与 $\R^d$ 上 Euler 法的 CFL 型条件同源。第 8 节会看到混合路径的速率在 $t\to1$ 时形如 $\dot\kappa_t/(1-\kappa_t)\to\infty$，正好在终点附近逼近这个约束边界——这是实现里必须小心的地方。

推导：停留时间服从指数分布

设当前 $X_t=x$，速率矩阵时齐（$\Rate_t\equiv\Rate$）。令

$$ g(s) := \mathbb P\big(X_r = x \text{ 对所有 } r\in[t,t+s]\ \big|\ X_t=x\big) $$

为「在 $s$ 时长内一次都没跳」的概率。由马尔可夫性，把 $[t,t+s+h]$ 拆成 $[t,t+s]$ 与 $(t+s,t+s+h]$：

$$ g(s+h) = g(s)\cdot\mathbb P(\text{在 } h \text{ 内不跳}\mid X_{t+s}=x) = g(s)\big(1-h\,\lambda(x)+o(h)\big), $$

这里用了 (2.1) 的 $y=x$ 情形。移项除以 $h$ 取极限：

$$ g'(s) = -\lambda(x)\,g(s),\qquad g(0)=1 \quad\Longrightarrow\quad g(s)=e^{-\lambda(x)\,s}. $$

所以停留时间 $\tau\sim\mathrm{Exp}(\lambda(x))$，期望停留时长 $\E[\tau]=1/\lambda(x)$。这正好解释了 4.2 里步长约束 $h\le 1/\lambda_t(x)$ 的含义：Euler 步长不能超过平均停留时间。

其次，跳到哪？在「已经决定要跳」的条件下，跳到 $y\ne x$ 的概率与速率成正比：

$$ \mathbb P(\text{跳到 } y\mid \text{发生跳转}) = \frac{\Rate(y\mid x)}{\lambda(x)} = \frac{\Rate(y\mid x)}{\sum_{y'\ne x}\Rate(y'\mid x)} . $$

直观理由：在极短的 $h$ 内，跳到 $y$ 的概率是 $h\Rate(y\mid x)$，发生任何跳转的概率是 $h\lambda(x)$，条件概率取比值即得（$h$ 约掉）。

时变情形：把 $\lambda$ 换成时间的函数，同样的推导给出 $\mathbb P(\tau>s)=\exp\big(-\int_t^{t+s}\lambda_r(x)\ud r\big)$，采样需要解这个积分方程或用 thinning（拒绝采样）。这也是为什么在实践中，时变速率的离散扩散模型基本都用 Euler 而不用 Gillespie。

import torch

def euler_ctmc(Q, p0, n_steps=200, T=1.0, bs=200_000):
    """用 Euler 格式模拟时齐 CTMC。
    Q : (S, S)  速率矩阵，约定 Q[x, y] = 从 x 跳到 y 的速率；行和为 0
    p0: (S,)    初始分布
    返回 X_T 的样本 (bs,)
    """
    S = Q.shape[0]
    assert torch.allclose(Q.sum(1), torch.zeros(S), atol=1e-6), "行和必须为 0"
    h = T / n_steps
    P = torch.eye(S) + h * Q                     # (S, S)  单步转移概率 delta_x + h Q
    assert (P >= 0).all(), "步长过大：需要 h <= 1 / max_x |Q[x, x]|"
    x = torch.multinomial(p0, bs, replacement=True)   # (bs,)  X_0 ~ p_init
    for _ in range(n_steps):
        x = torch.multinomial(P[x], 1).squeeze(1)     # (bs,)  按第 x 行采下一个状态
    return x

# --- 验证第 2.5 节的两状态例子 ---
lam = 1.5
Q  = torch.tensor([[-lam,  lam],
                   [ lam, -lam]])                    # (2, 2)
p0 = torch.tensor([1.0, 0.0])                        # 确定性地从状态 a 出发
xT = euler_ctmc(Q, p0, n_steps=200, T=1.0)
emp = (xT == 0).float().mean().item()                # 经验的 P(X_1 = a | X_0 = a)
exact = 0.5 * (1 + torch.exp(torch.tensor(-2 * lam * 1.0))).item()
print(f"empirical={emp:.4f}  exact={exact:.4f}")     # 两者应吻合到蒙特卡洛误差内

推导：因子化后需要多少参数

非零速率只出现在「把第 $j$ 个位置换成某个 $v\in\mathcal V$」这种跳转上。用

$$ \Rate^\theta_t(v, j\mid x) := \Rate^\theta_t\big(\underbrace{(x_1,\dots,x_{j-1},\,v,\,x_{j+1},\dots,x_N)}_{y}\ \big|\ x\big) $$

记这个速率。于是网络的全部输出可以排成一个矩阵

$$ x\ \longmapsto\ \big\{\Rate^\theta_t(y\mid x)\big\}_{y\in\mathcal N(x)} =\begin{pmatrix} \Rate^\theta_t(v_1,1\mid x) & \cdots & \Rate^\theta_t(v_V,1\mid x)\\ & \cdots & \\ \Rate^\theta_t(v_1,N\mid x) & \cdots & \Rate^\theta_t(v_V,N\mid x) \end{pmatrix}\in\R^{N\times V}. $$

输出形状从 $V^N$ 变成 $N\times V$：从指数变成线性。取 $V=32000,N=1024$，这是约 $3.3\times10^7$ 个数——一个普通的 logits 张量而已。

合法性条件也逐位置地写：对每个位置 $j$，第 $j$ 行必须自成一个合法速率矩阵的一行，

$$ \Rate^\theta_t(v,j\mid x)\ \ge\ 0\ \text{ 若 } v\ne x_j, \qquad \Rate^\theta_t(x_j, j\mid x) = -\sum_{v\ne x_j}\Rate^\theta_t(v,j\mid x). $$

校验一下这与原始的行和为零一致：原始条件要求 $\sum_{y\in\Vocab}\Rate^\theta_t(y\mid x)=0$。因为非邻居项全为 0，求和只剩邻居，而邻居按「改哪个位置」分成 $N$ 组（$y=x$ 这一项要小心，它属于每一组的对角元）。若我们约定「对角总速率」为 $\Rate^\theta_t(x\mid x)=\sum_{j=1}^N \Rate^\theta_t(x_j,j\mid x)$，则

$$ \sum_{y\in\Vocab}\Rate_t^\theta(y\mid x) =\sum_{j=1}^{N}\Big[\sum_{v\ne x_j}\Rate^\theta_t(v,j\mid x)+\Rate^\theta_t(x_j,j\mid x)\Big] =\sum_{j=1}^N 0 = 0 .\quad\checkmark $$

换句话说：整条序列的总跳出速率，等于各位置跳出速率之和；$\lambda_t(x)=\sum_j\lambda_t^{(j)}(x)$，其中 $\lambda^{(j)}_t(x)=\sum_{v\ne x_j}\Rate_t^\theta(v,j\mid x)$。

推导：并行 per-token Euler 与完整 Euler 的关系

完整的 Euler 步 (4.1) 要从 $\delt{x}+h\Rate_t(\cdot\mid x)$ 中采一个状态，这个分布把 $1-h\lambda_t(x)$ 的质量留在 $x$，把 $h\Rate_t(v,j\mid x)$ 的质量分给各个邻居——它一次最多改一个位置。

并行版本则对每个位置 $j$ 独立地做一次一维 Euler 步：

$$ \tilde p_{j,t}(v\mid x_j)= \begin{cases} h\,\Rate_t(v,j\mid x), & v\ne x_j,\\[4pt] 1-h\sum_{v'\ne x_j}\Rate_t(v',j\mid x), & v = x_j, \end{cases} \qquad X^{(j)}_{t+h}\sim \tilde p_{j,t}(\cdot\mid X^{(j)}_t)\ \text{（各 } j \text{ 独立）}. $$

两者有什么差别？算一下「这一步里恰好有一个位置 $j$ 被改成 $v$」的概率。并行版本给出

$$ \underbrace{h\,\Rate_t(v,j\mid x)}_{\text{位置 }j\text{ 改了}}\cdot\prod_{j'\ne j}\underbrace{\Big(1-h\lambda^{(j')}_t(x)\Big)}_{\text{其它位置没改}} = h\,\Rate_t(v,j\mid x)\;+\;O(h^2), $$

与完整 Euler 的 $h\Rate_t(v,j\mid x)$ 在 $h$ 的一阶上一致。差别只出现在二阶：并行版本以 $O(h^2)$ 的概率让多个位置在同一步同时改变，而完整 Euler 永远不会。

这就是讲义那句「agrees with the full CTMC Euler step up to first order in $h$, but allows for a $O(h^2)$ probability of simultaneous updates」的具体含义。

算法 7：从因子化 CTMC 模型采样

输入：因子化速率网络 $\Rate^\theta_t$、初始分布 $\simple$、步数 $n$。

1. $t\leftarrow 0$，$h\leftarrow 1/n$。
2. 采 $X_0=(X_0^{(1)},\dots,X_0^{(N)})\sim\simple$。
3. 对 $i=1,\dots,n$：
   1. 一次网络前向，得到全部因子化速率 $\{q_j(v)\}_{j=1..N,\,v\in\mathcal V}\leftarrow \Rate^\theta_t(\cdot\mid X_t)$，形状 $N\times V$。
   2. 对 $j=1,\dots,N$ 并行：令 $x\leftarrow X^{(j)}_t$，构造 $$\tilde p_{j,t}(v\mid x)=\begin{cases} h\,q_j(v), & v\ne x,\\ 1-h\sum_{v'\ne x}q_j(v'), & v=x,\end{cases}$$ 然后采 $X^{(j)}_{t+h}\sim\mathrm{Categorical}\big(\{\tilde p_{j,t}(v\mid x)\}_{v\in\mathcal V}\big)$。
   3. $t\leftarrow t+h$。
4. 返回 $X_1$。

注意每一步只有一次网络前向——$N$ 个位置共享同一次前向的输出。这是与自回归解码（每 token 一次前向）最本质的区别。

推导：边际路径确实连接噪声与数据

两个端点各验证一次：

$$ p_0(x)=\sum_{z}p_0(x\mid z)\data(z)=\sum_z \simple(x)\data(z)=\simple(x)\underbrace{\sum_z\data(z)}_{=1}=\simple(x), $$ $$ p_1(x)=\sum_{z}p_1(x\mid z)\data(z)=\sum_z\delt{z}(x)\data(z)\overset{(\ast)}{=}\data(x). $$

$(\ast)$ 处：$\delt{z}(x)$ 只在 $z=x$ 时为 1，所以整个求和塌缩成 $\data(x)$。（在 $\R^d$ 上这一步需要「$\delta$ 的筛选性质」这种分布论语言；在有限集合上它只是「求和里只剩一项」，离散情形反而更严格、更简单。）

另外要确认 $p_t\in\Delta(\Vocab)$：非负性显然，归一性由 $\sum_x p_t(x)=\sum_z\data(z)\sum_x p_t(x\mid z)=\sum_z\data(z)=1$。$\square$

推导（一）：它确实是一条合法的条件概率路径

第 1 步：每个因子是 $\Delta(\mathcal V)$ 里的点。固定 $j$，记 $q^{(j)}_t := (1-\kappa_t)\,\simple^{(j)} + \kappa_t\,e_{z_j}\in\R^{V}$。 因为 $0\le\kappa_t\le1$，这是两个概率向量的凸组合，故非负且分量和为 $(1-\kappa_t)+\kappa_t=1$。✓

第 2 步：乘积是 $\Delta(\Vocab)$ 里的点。$N$ 个独立分布的乘积仍是分布：

$$ \sum_{x\in\Vocab}p_t(x\mid z)=\sum_{x_1\in\mathcal V}\cdots\sum_{x_N\in\mathcal V}\prod_{j}q^{(j)}_t(x_j)=\prod_{j}\Big[\sum_{x_j}q^{(j)}_t(x_j)\Big]=1 .\ \checkmark $$

第 3 步：$t=0$ 端点。$\kappa_0=0$，(6.1) 变成 $\prod_j\simple^{(j)}(x_j)=\simple(x)$。✓

第 4 步：$t=1$ 端点。$\kappa_1=1$，(6.1) 变成 $\prod_j\delt{z_j}(x_j)$。这个乘积当且仅当每个位置都满足 $x_j=z_j$ 时为 1，即 $\prod_j\delt{z_j}(x_j)=\delt{z}(x)$。✓ $\square$

推导（二）：等价的采样过程（这才是实现时真正用的东西）

幻灯片上写着「Check for yourself that this is a cond. prob. path!」——我们把这个检查做完。断言：从 $p_t(\cdot\mid z)$ 采样等价于

$$ \begin{aligned} &m_j \sim \mathrm{Bernoulli}(\kappa_t),\qquad \xi_j\sim \simple^{(j)},\qquad (j=1,\dots,N \text{ 独立})\\ &x_j = m_j\,z_j + (1-m_j)\,\xi_j,\\ &x=(x_1,\dots,x_N). \end{aligned} $$

证明。固定位置 $j$ 与任意 $v\in\mathcal V$，按 $m_j$ 的取值做全概率分解：

$$ \begin{aligned} \mathbb P(x_j = v) &= \mathbb P(m_j=1)\,\mathbb P(z_j = v) + \mathbb P(m_j=0)\,\mathbb P(\xi_j=v)\\ &= \kappa_t\,\delt{z_j}(v) + (1-\kappa_t)\,\simple^{(j)}(v), \end{aligned} $$

正是 (6.1) 的第 $j$ 个因子。又因为 $(m_j,\xi_j)$ 关于 $j$ 独立，所以 $(x_1,\dots,x_N)$ 相互独立，联合分布等于各边际之积，即 (6.1)。$\square$

读法：$m_j=1$ 表示「第 $j$ 个位置保留真实 token」，概率 $\kappa_t$；$m_j=0$ 表示「第 $j$ 个位置被噪声替换」，概率 $1-\kappa_t$。所以 $\kappa_t$ 是「保留率」，$1-\kappa_t$ 是「破坏率」。$t=0$ 时全被破坏，$t=1$ 时全部保留。这三行就是第 11 节训练代码的第 3–8 行。

推导（第一步）：$\Rate_t$ 确实是合法速率矩阵

讲义直接断言了这一点，我们补上。设 $p_t(x)>0$（否则 (7.2) 无定义，见下面的注记）。

(1) 非对角元非负。对 $y\ne x$：

$$ \Rate_t(y\mid x)=\sum_z \underbrace{\Rate^z_t(y\mid x)}_{\ge 0\ (\text{每个 }\Rate^z_t\text{ 合法})}\underbrace{p_{1\mid t}(z\mid x)}_{\ge0\ (\text{概率})}\ \ge\ 0 .\ \checkmark $$

(2) 行和为零。交换求和次序：

$$ \sum_{y\in\Vocab}\Rate_t(y\mid x)=\sum_{z}p_{1\mid t}(z\mid x)\underbrace{\sum_{y}\Rate^z_t(y\mid x)}_{=\,0}=0 .\ \checkmark $$

注意这两条都只用到「凸组合保持速率矩阵的两个条件」——因为 $p_{1\mid t}(\cdot\mid x)$ 是一个概率分布，(7.2) 就是一族合法速率矩阵的凸组合。合法速率矩阵的集合是凸的，这是边际化技巧能成立的代数根源。$\square$

推导（第二步）：验证 $\Rate_t$ 满足边际路径的 KFE

由命题 2 的充分性，只需验证 $\frac{\dd{}}{\dd{t}}p_t(x)=\sum_y \Rate_t(x\mid y)p_t(y)$。

$$ \begin{aligned} \frac{\dd{}}{\dd{t}}p_t(x) &\overset{(i)}{=}\frac{\dd{}}{\dd{t}}\sum_{z\in\Vocab}p_t(x\mid z)\,\data(z)\\ &\overset{(ii)}{=}\sum_{z\in\Vocab}\frac{\dd{}}{\dd{t}}p_t(x\mid z)\,\data(z)\\ &\overset{(iii)}{=}\sum_{z\in\Vocab}\Big[\sum_{y\in\Vocab}\Rate^z_t(x\mid y)\,p_t(y\mid z)\Big]\data(z)\\ &\overset{(iv)}{=}\sum_{y\in\Vocab}p_t(y)\Big[\sum_{z\in\Vocab}\Rate^z_t(x\mid y)\,\frac{p_t(y\mid z)\,\data(z)}{p_t(y)}\Big]\\ &\overset{(v)}{=}\sum_{y\in\Vocab}p_t(y)\,\Rate_t(x\mid y). \end{aligned} $$

逐步说明：

• (i) 边际概率路径的定义。
• (ii) 交换求和与求导（$\Vocab$ 有限，恒可交换；这一步在 $\R^d$ 上需要控制收敛定理，离散情形免费）。
• (iii) 对条件路径用 KFE (7.1)——这一步正是「条件速率矩阵」定义的全部用处。
• (iv) 交换 $y$ 与 $z$ 的求和次序，然后同乘同除 $p_t(y)$。这是整个证明的技术核心：凭空造出一个 $p_t(y)$ 提到外面，剩下的部分就恰好凑成了后验 $p_{1\mid t}(z\mid y)$。
• (v) 括号里正是 (7.2) 在 $x\leftarrow y$（出发点）、$y\leftarrow x$（到达点）时的定义式。

于是边际路径 $p_t$ 满足 $\Rate_t$ 的 KFE。由命题 2 的充分性，速率矩阵为 $\Rate_t$、初值 $p_0=\simple$ 的 CTMC 满足 $X_t\sim p_t$ 对所有 $t$。取 $t=1$ 得 $X_1\sim p_1=\data$。$\square$

推导：例 37 的证明

第 0 步：约化到 $N=1$。混合路径 (6.1) 关于位置完全因子化，$p_t(x\mid z)=\prod_j q^{(j)}_t(x_j)$，而 (8.1) 给出的速率矩阵也只在单个位置上动作。因此只需对每个位置分别验证 KFE，不妨设 $N=1$，把 $x,z$ 直接当成 $\mathcal V$ 中的元素，路径为

$$ p_t(x\mid z)=(1-\kappa_t)\,\simple(x)+\kappa_t\,\delt{z}(x), $$

待验证的速率矩阵为 $\Rate^z_t(y\mid x)=\frac{\dot\kappa_t}{1-\kappa_t}\big(\delt{z}(y)-\delt{x}(y)\big)$。

第 1 步：算左端。下面这串等号是整个推导的主体，八步：

$$ \begin{aligned} \frac{\dd{}}{\dd{t}}p_t(x\mid z) &\overset{(i)}{=}\frac{\dd{}}{\dd{t}}\Big[(1-\kappa_t)\simple(x)+\kappa_t\delt{z}(x)\Big]\\ &\overset{(ii)}{=}\dot\kappa_t\,\delt{z}(x)-\dot\kappa_t\,\simple(x)\\ &\overset{(iii)}{=}\frac{\dot\kappa_t}{1-\kappa_t}\Big(\delt{z}(x)-\big[(1-\kappa_t)\simple(x)+\kappa_t\delt{z}(x)\big]\Big)\\ &\overset{(iv)}{=}\frac{\dot\kappa_t}{1-\kappa_t}\Big(\delt{z}(x)-p_t(x\mid z)\Big)\\ &\overset{(v)}{=}\frac{\dot\kappa_t}{1-\kappa_t}\delt{z}(x)\big(1-p_t(x\mid z)\big)+\frac{\dot\kappa_t}{1-\kappa_t}\big(\delt{z}(x)-1\big)p_t(x\mid z)\\ &\overset{(vi)}{=}\sum_{y\ne x}\frac{\dot\kappa_t}{1-\kappa_t}\delt{z}(x)\,p_t(y\mid z)+\frac{\dot\kappa_t}{1-\kappa_t}\big(\delt{z}(x)-1\big)p_t(x\mid z)\\ &\overset{(vii)}{=}\sum_{y\ne x}\Rate^z_t(x\mid y)\,p_t(y\mid z)+\Rate^z_t(x\mid x)\,p_t(x\mid z)\\ &\overset{(viii)}{=}\sum_{y\in\Vocab}\Rate^z_t(x\mid y)\,p_t(y\mid z). \end{aligned} $$

这正是 KFE (7.1)。逐步说明：

• (i) 代入 $N=1$ 的混合路径定义。
• (ii) 求导：$\simple(x)$ 与 $\delt{z}(x)$ 都与 $t$ 无关，只有系数含 $t$，$\frac{\dd{}}{\dd{t}}(1-\kappa_t)=-\dot\kappa_t$，$\frac{\dd{}}{\dd{t}}\kappa_t=\dot\kappa_t$。
• (iii) 纯代数配凑，是全场唯一需要「想一下」的一步。验证方法是把右端展开： $$\frac{\dot\kappa_t}{1-\kappa_t}\Big[(1-\kappa_t)\delt{z}(x)-(1-\kappa_t)\simple(x)\Big]=\dot\kappa_t\delt{z}(x)-\dot\kappa_t\simple(x),$$ 其中中括号内用了 $\delt{z}(x)-\kappa_t\delt{z}(x)=(1-\kappa_t)\delt{z}(x)$。与 (ii) 的结果一致 ✓。这一步的动机：我们想让 $p_t(x\mid z)$ 本身出现在表达式里，因为 KFE 的右端是关于 $p_t(\cdot\mid z)$ 的线性式。
• (iv) 中括号内正是 $p_t(x\mid z)$ 的定义。到这里已经得到一个漂亮的中间结论： $$\frac{\dd{}}{\dd{t}}p_t(x\mid z)=\frac{\dot\kappa_t}{1-\kappa_t}\big(\delt{z}(x)-p_t(x\mid z)\big),$$ 即「路径以正比于『离目标还差多少』的速度向 $\delt{z}$ 逼近」——一个标准的一阶弛豫方程，和 $\dot X=\frac{z-X}{1-t}$ 长得一样。
• (v) 把 (iv) 拆成两项。验证：把右端展开， $$\tfrac{\dot\kappa}{1-\kappa}\big[\delt{z}(x)-\delt{z}(x)p_t(x\mid z)+\delt{z}(x)p_t(x\mid z)-p_t(x\mid z)\big]=\tfrac{\dot\kappa}{1-\kappa}\big[\delt{z}(x)-p_t(x\mid z)\big]\ \checkmark$$ 中间两项恰好抵消。动机：第一项要变成「流入」（对 $y\ne x$ 求和），第二项要变成「流出」（对角项）。
• (vi) 用 $1-p_t(x\mid z)=\sum_{y\ne x}p_t(y\mid z)$（因为 $\sum_{y\in\Vocab}p_t(y\mid z)=1$）。
• (vii) 用 (8.1) 的定义把系数认出来：
  • 对 $y\ne x$：$\Rate^z_t(x\mid y)=\frac{\dot\kappa_t}{1-\kappa_t}(\delt{z}(x)-\delt{y}(x))=\frac{\dot\kappa_t}{1-\kappa_t}\delt{z}(x)$，因为 $y\ne x$ 时 $\delt{y}(x)=0$；
  • 对角项：$\Rate^z_t(x\mid x)=\frac{\dot\kappa_t}{1-\kappa_t}(\delt{z}(x)-\delt{x}(x))=\frac{\dot\kappa_t}{1-\kappa_t}(\delt{z}(x)-1)$，因为 $\delt{x}(x)=1$。
• (viii) 把对角项并回求和。

KFE 得证。再由命题 2 的充分性，$\Rate^z_t$ 的 CTMC 从 $\simple$ 出发就服从 $p_t(\cdot\mid z)$。$\square$

推导：(8.1) 是合法速率矩阵吗？

必须检查，否则前面的一切都不成立。

非对角元非负：取 $v\ne x_j$，则 $\delt{x_j}(v)=0$，于是

$$ \Rate^z_t(v,j\mid x)=\frac{\dot\kappa_t}{1-\kappa_t}\,\delt{z_j}(v)\ \ge\ 0 $$

——当且仅当 $\dot\kappa_t\ge0$ 且 $\kappa_t<1$。这解释了例 35 里为什么要求调度器单调不减：$\dot\kappa_t\ge0$ 不是为了好看，而是速率矩阵合法性的充要条件。（$\kappa_t=1$ 只在 $t=1$ 端点发生，那时过程已经结束。）

行和为零：

$$ \sum_{v\in\mathcal V}\Rate^z_t(v,j\mid x)=\frac{\dot\kappa_t}{1-\kappa_t}\Big[\underbrace{\sum_v\delt{z_j}(v)}_{=1}-\underbrace{\sum_v\delt{x_j}(v)}_{=1}\Big]=0 .\ \checkmark $$

本质上：「one-hot 减 one-hot」永远是一个分量和为零的向量，这正好是速率矩阵一行所需的形状。$\square$

推导：速率爆炸是必须的

取 $\kappa_t=t$，速率因子是 $\frac{1}{1-t}$，在 $t\to1$ 时发散。这不是 bug。原因：某个位置在时刻 $t$ 仍然错误的概率是 $1-\kappa_t$，而我们要求它在 $t=1$ 时以概率 1 变正确。设 $u(t):=\mathbb P(x_j\ne z_j)=1-\kappa_t$。由动力学，$u$ 只减不增，且

$$ \frac{\dd{}}{\dd{t}}u(t)=-\frac{\dot\kappa_t}{1-\kappa_t}\,u(t)=-\frac{\dot\kappa_t}{1-\kappa_t}(1-\kappa_t)=-\dot\kappa_t\ \checkmark $$

确实与 $u=1-\kappa_t$ 自洽。若速率因子有界，比如恒为常数 $c$，则 $u(t)=e^{-ct}>0$ 对所有有限 $t$ 成立，永远到不了 0。要在有限时间内把「还错着」的概率精确压到 0，速率必须以 $\frac{1}{1-t}$ 的量级发散。这与第二讲 CondOT 向量场 $\frac{z-x}{1-t}$ 在 $t\to1$ 时爆炸是同一个现象、同一个原因。

实践后果：Euler 步长约束 $h\le 1/\lambda_t(x)$ 在 $t\to1$ 时变得苛刻。第 10 节会看到，对 $\kappa_t=t$ 与均匀步长 $h=1/n$，第 $i$ 步的单位置跳转概率恰为 $\frac{1}{n-i}$，在最后一步等于 1 —— 边界正好被取到，不多不少。这是一个漂亮的巧合，也是均匀步长在这里格外好用的原因。

推导：定理 38 的证明

第 1 步：边际速率矩阵也是因子化的。由边际化技巧 (7.2)，

$$ \Rate_t(y\mid x)=\sum_{z\in\Vocab}\Rate^z_t(y\mid x)\,p_{1\mid t}(z\mid x). \tag{9.2} $$

若 $y$ 与 $x$ 相差多于一个位置，则对每一个 $z$ 都有 $\Rate^z_t(y\mid x)=0$（例 37 的条件速率矩阵是因子化的），所以整个求和为 0，$\Rate_t(y\mid x)=0$。因子化性质在取凸组合下被保持。✓

第 2 步：算非零的那些分量。取 $y$ 为「把 $x$ 的第 $j$ 位换成 $v$」，代入 (8.1)：

$$ \begin{aligned} \Rate_t(v,j\mid x) &=\sum_{z\in\Vocab}\Rate^z_t(v,j\mid x)\,p_{1\mid t}(z\mid x)\\ &\overset{(i)}{=}\sum_{z\in\Vocab}\frac{\dot\kappa_t}{1-\kappa_t}\Big(\delt{z_j}(v)-\delt{x_j}(v)\Big)p_{1\mid t}(z\mid x)\\ &\overset{(ii)}{=}\frac{\dot\kappa_t}{1-\kappa_t}\Big(\sum_{z\in\Vocab}\delt{z_j}(v)\,p_{1\mid t}(z\mid x)\;-\;\delt{x_j}(v)\underbrace{\sum_{z\in\Vocab}p_{1\mid t}(z\mid x)}_{=1}\Big)\\ &\overset{(iii)}{=}\frac{\dot\kappa_t}{1-\kappa_t}\Big(p_{1\mid t}(z_j=v\mid x)-\delt{x_j}(v)\Big). \end{aligned} $$

逐步说明：

• (i) 代入例 37 的条件速率矩阵公式。
• (ii) 把与 $z$ 无关的 $\delt{x_j}(v)$ 提出求和，剩下 $\sum_z p_{1\mid t}(z\mid x)=1$（后验是一个分布）。
• (iii) 关键一步——边缘化。注意 $\delt{z_j}(v)=\mathbb 1[z_j=v]$，所以

$$ \sum_{z\in\Vocab}\delt{z_j}(v)\,p_{1\mid t}(z\mid x)=\sum_{z:\,z_j=v}p_{1\mid t}(z\mid x)=p_{1\mid t}(z_j=v\mid x). $$

也就是说，那个对整个 $\Vocab$（$V^N$ 项！）的求和，塌缩成了后验在单个位置上的边缘分布——一个 $V$ 维向量。$\square$

这一步为什么如此重要：它把「需要知道整条序列的后验分布 $p_{1\mid t}(z\mid x)\in\Delta(\Vocab)$」（不可表示）降格成「只需要知道 $N$ 个位置各自的边缘后验 $p_{1\mid t}(z_j=\cdot\mid x)\in\Delta(\mathcal V)$」（形状 $N\times V$，完全可表示）。因子化的速率矩阵只需要因子化的信息。

推导：(9.1) 的合法性与「已定稿的位置不会再动」

合法性。非对角元（$v\ne x_j$）：$\Rate_t(v,j\mid x)=\frac{\dot\kappa_t}{1-\kappa_t}p_{1\mid t}(z_j=v\mid x)\ge0$ ✓。行和：$\frac{\dot\kappa_t}{1-\kappa_t}\big(\sum_v p_{1\mid t}(z_j=v\mid x)-\sum_v\delt{x_j}(v)\big)=\frac{\dot\kappa_t}{1-\kappa_t}(1-1)=0$ ✓。只要网络输出的是一个合法的概率分布（比如 softmax 的结果），(9.1) 自动是合法速率矩阵——不需要任何额外的投影或裁剪。这是这个参数化最漂亮的地方。

已定稿的位置。假设在时刻 $t$，第 $j$ 位的后验已经塌缩成点质量 $p_{1\mid t}(z_j=\cdot\mid x)=e_{x_j}$（即「我确信原本就是 $x_j$」）。代入 (9.1)：

$$ \Rate_t(\cdot,j\mid x)=\frac{\dot\kappa_t}{1-\kappa_t}\big(e_{x_j}-e_{x_j}\big)=0 . $$

整行为零，这个位置再也不会改变。第 10 节会看到，掩码模型中所有「已经被解掩」的位置都自动落入这一情形——这就是 MDLM 里「一旦生成就不再修改（carry-over）」这一行为的数学来源，它不是实现技巧，而是公式的推论。

推导：条件损失 = 边际损失 + 常数

第 1 步：把期望改写成「先采 $x$ 再采 $z$」。固定 $t$。$(z,x)$ 的联合分布是

$$ \mathbb P(z,x)=\data(z)\,p_t(x\mid z). $$

按贝叶斯把它反过来分解（这就是塔性质 / tower property 的用法）：

$$ \data(z)\,p_t(x\mid z)=p_t(x)\cdot\underbrace{\frac{\data(z)p_t(x\mid z)}{p_t(x)}}_{=\,p_{1\mid t}(z\mid x)} . $$

第 2 步：代入损失，对每个位置分别处理。取 (9.3) 中固定的 $t$ 与固定的位置 $j$：

$$ \begin{aligned} \E_{z,x}\Big[-\log p^\theta_{1\mid t}(z_j\mid x)\Big] &\overset{(i)}{=}\sum_{z\in\Vocab}\sum_{x\in\Vocab}\data(z)p_t(x\mid z)\Big[-\log p^\theta_{1\mid t}(z_j\mid x)\Big]\\ &\overset{(ii)}{=}\sum_{x\in\Vocab}p_t(x)\sum_{z\in\Vocab}p_{1\mid t}(z\mid x)\Big[-\log p^\theta_{1\mid t}(z_j\mid x)\Big]\\ &\overset{(iii)}{=}\sum_{x\in\Vocab}p_t(x)\sum_{v\in\mathcal V}\Big[\sum_{z:\,z_j=v}p_{1\mid t}(z\mid x)\Big]\Big[-\log p^\theta_{1\mid t}(v\mid x)\Big]\\ &\overset{(iv)}{=}\E_{x\sim p_t}\Big[\ \sum_{v\in\mathcal V} -\,p_{1\mid t}(z_j=v\mid x)\,\log p^\theta_{1\mid t}(v\mid x)\Big]\\ &\overset{(v)}{=}\E_{x\sim p_t}\Big[\;\underbrace{\KL\Big(p_{1\mid t}(z_j{=}\cdot\mid x)\ \Big\|\ p^\theta_{1\mid t}(\cdot\mid x)\Big)}_{\text{我们真正想最小化的}}\;+\;\underbrace{H\Big(p_{1\mid t}(z_j{=}\cdot\mid x)\Big)}_{\text{与}\theta\text{ 无关}}\Big]. \end{aligned} $$

逐步说明：

• (i) 期望写成显式求和。
• (ii) 第 1 步的贝叶斯重排，交换求和次序。这就是塔性质：把「先采 $z$ 再采 $x$」换成「先采 $x$ 再采 $z$」，边际化在采样过程里自动完成，我们从头到尾没有计算过 $p_{1\mid t}$。
• (iii) 被求和的项只通过 $z_j$ 依赖 $z$，于是按 $z_j=v$ 把 $z$ 分组。
• (iv) 中括号里正是边缘后验 $p_{1\mid t}(z_j=v\mid x)$。这一行就是交叉熵 $H(p_{1\mid t},p^\theta)$ 的定义。
• (v) 交叉熵的标准分解 $H(p,q)=\KL(p\Vert q)+H(p)$。

第 3 步：结论。对 $j$ 求和、对 $t$ 取期望，得

$$ \mathcal L_{\mathrm{DFM}}(\theta)=\underbrace{\E_{t\sim\mathrm{Unif}[0,1]}\ \E_{x\sim p_t}\sum_{j=1}^{N}\KL\Big(p_{1\mid t}(z_j{=}\cdot\mid x)\ \Big\|\ p^\theta_{1\mid t}(\cdot\mid x)\Big)}_{=:\ \mathcal L_{\mathrm{marginal}}(\theta)\ \text{（不可计算）}}\;+\;\underbrace{C}_{\text{与 }\theta\text{ 无关}} . $$

因此

$$ \nabla_\theta\mathcal L_{\mathrm{DFM}}(\theta)=\nabla_\theta\mathcal L_{\mathrm{marginal}}(\theta), $$

两个损失有完全相同的梯度，因而有相同的优化轨迹与相同的最优解。而 $\mathcal L_{\mathrm{marginal}}$ 显然在且仅在

$$ p^\theta_{1\mid t}(\cdot\mid x)=p_{1\mid t}(z_j=\cdot\mid x)\qquad\text{对所有 }t,\ \text{所有 }x\text{ 满足 }p_t(x)>0 $$

时取到最小值 0（KL 非负，等于 0 当且仅当两分布相同）。$\square$

算法 8：训练因子化 CTMC 模型（离散扩散）

输入：数据集（序列 $z=(z_1,\dots,z_N)\in\mathcal V^N$）、初始 token 边际 $\simple^{(j)}$、调度器 $\kappa_t$、返回逐位置 logits 的网络 $f_\theta$、优化器。

1. 采一个数据点 $z\sim\data$。
2. 采时间 $t\sim\mathrm{Unif}[0,1]$，计算 $\kappa\leftarrow\kappa_t$。
3. 采噪声状态 $x\sim p_t(\cdot\mid z)$，对 $j=1,\dots,N$ 并行：
   • $m_j\sim\mathrm{Bernoulli}(\kappa)$；
   • $\xi_j\sim\simple^{(j)}$；
   • $x_j\leftarrow m_j z_j+(1-m_j)\xi_j$。
4. 前向：$\ell_j(\cdot)\leftarrow f_\theta(x,t)_j$，$\ p^\theta_{1\mid t}(v\mid x)_j=\mathrm{softmax}(\ell_j)(v)$。
5. 损失：$\mathcal L_{\mathrm{DFM}}(\theta)\leftarrow\sum_{j=1}^{N}\big[-\log p^\theta_{1\mid t}(z_j\mid x)_j\big]$。
6. 更新：$\theta\leftarrow\mathrm{Opt.step}(\nabla_\theta\mathcal L_{\mathrm{DFM}}(\theta))$。

注意没有任何模拟：不需要跑 CTMC，不需要反向过程，第 3 步一次采样直接跳到时刻 $t$。这就是「免模拟（simulation-free）」训练，和第 2 讲的 flow matching 完全一致。

推导：把 $\simple=\delt{[\mathrm{MASK}]^N}$ 代进所有公式

(a) 条件概率路径。代入 (6.1)，第 $j$ 个因子变成

$$ q^{(j)}_t=(1-\kappa_t)\,e_{[\mathrm{MASK}]}+\kappa_t\,e_{z_j}. $$

这是一个只在两个 token 上有质量的分布：第 $j$ 位以概率 $\kappa_t$ 是真值 $z_j$，以概率 $1-\kappa_t$ 是 [MASK]，没有第三种可能。写 $\alpha_t:=\kappa_t$（「已揭示比例」），则每个位置被掩的概率是 $1-\alpha_t$，与本讲开头的表述一致。

(b) 支撑集。$p_t(\cdot\mid z)$ 的支撑集是

$$ \mathrm{supp}\,p_t(\cdot\mid z)=\big\{x:\ x_j\in\set{z_j,[\mathrm{MASK}]}\ \forall j\big\}, $$

共 $2^N$ 个状态（而不是 $(V+1)^N$）。模型永远不会遇到「某个位置是一个错误的真实 token」这种输入——这是掩码构造相对于「均匀噪声」构造的巨大简化。

(c) 条件速率矩阵。把 (8.1) 特化。分两种情形：

• $x_j = z_j$（已揭示）：整行为 0，不动。
• $x_j=[\mathrm{MASK}]\ne z_j$（仍被掩）：

$$ \Rate^z_t(v,j\mid x)=\frac{\dot\kappa_t}{1-\kappa_t}\begin{cases} +1,& v=z_j\quad(\text{解掩到真值})\\ -1,& v=[\mathrm{MASK}]\quad(\text{对角，跳出})\\ 0,& \text{其它 } v . \end{cases} $$

结论：非零速率只出现在 $[\mathrm{MASK}]\to z_j$ 这一个方向上。反方向 $z_j\to[\mathrm{MASK}]$ 的速率恒为 0，因为按 (8.1)，$x_j=z_j$ 时整行为零。所以：

$$ \boxed{\text{解掩是不可逆的：一旦某个位置被填上，它永远不会退回 }[\mathrm{MASK}]\text{，也永远不会被改成别的 token。}} $$

用马尔可夫链的语言说，「$x_j=z_j$」对第 $j$ 个位置的子链是吸收态，所以这类构造在文献里也叫「吸收态扩散（absorbing-state diffusion）」。

(d) 边际速率矩阵。把 (9.1) 特化。对已揭示的位置（$x_j\ne[\mathrm{MASK}]$），我们已经论证过后验是点质量：因为在支撑集里 $x_j\ne[\mathrm{MASK}]$ 只可能来自 $x_j=z_j$，故 $p_{1\mid t}(z_j=v\mid x)=\delt{x_j}(v)$，于是

$$ \Rate_t(\cdot,j\mid x)=\frac{\dot\kappa_t}{1-\kappa_t}\big(e_{x_j}-e_{x_j}\big)=0 . $$

对仍被掩的位置（$x_j=[\mathrm{MASK}]$）：

$$ \Rate_t(v,j\mid x)=\frac{\dot\kappa_t}{1-\kappa_t}\Big(p_{1\mid t}(z_j=v\mid x)-\delt{[\mathrm{MASK}]}(v)\Big) =\frac{\dot\kappa_t}{1-\kappa_t}\times\begin{cases} p_{1\mid t}(z_j=v\mid x),& v\ne[\mathrm{MASK}]\\ -1,& v=[\mathrm{MASK}] . \end{cases} $$

（用到 $\sum_{v\in\mathcal V}p_{1\mid t}(z_j=v\mid x)=1$，且真实 token 上不含 [MASK]。）

整个采样过程于是只有一件事可做：以速率 $\frac{\dot\kappa_t}{1-\kappa_t}$ 挑一个还被掩的位置，按网络给出的分布 $p^\theta_{1\mid t}(\cdot\mid x)$ 抽一个 token 填进去。$\square$

推导：均匀步长下每步的解掩概率恰好是 $\tfrac{1}{n-i}$

取 $\kappa_t=t$，则速率因子为 $\frac{\dot\kappa_t}{1-\kappa_t}=\frac{1}{1-t}$。用 $n$ 步均匀 Euler，$h=1/n$，第 $i$ 步（$i=0,1,\dots,n-1$）位于 $t=i/n$。对一个仍被掩的位置，这一步内发生解掩的概率是

$$ h\cdot\frac{1}{1-t}=\frac{1/n}{1-i/n}=\frac{1}{n-i}. $$

读一读这个结果：

• $i=0$：概率 $1/n$。第一步只揭示大约 $N/n$ 个 token。
• $i=n-1$（最后一步）：概率 $\frac{1}{1}=1$。所有仍被掩的位置全部一次性解掩。

所以这个格式恰好在最后一步用尽 Euler 的非负性预算（$h\lambda=1$，对角项正好为 0），既不会溢出（概率超过 1），也不会在 $t=1$ 时留下未解掩的位置。这不是运气：$\frac{\dot\kappa_t}{1-\kappa_t}$ 的分母 $1-\kappa_t$ 正是「还剩多少没揭示」，Euler 步长 $h$ 与它的比值自然给出「这一步该揭示其中的多大比例」。

核对期望：仍被掩的位置数的期望在第 $i$ 步后应是 $N(1-i/n)$。归纳：第 0 步前是 $N$；每步保留比例 $1-\frac{1}{n-i}=\frac{n-i-1}{n-i}$，于是 $i$ 步后为 $N\prod_{k=0}^{i-1}\frac{n-k-1}{n-k}=N\cdot\frac{n-i}{n}$（望远镜求积）$=N(1-t)$ ✓，与 $1-\kappa_t=1-t$ 一致。

推导：为什么「更快」要打问号

把两者的实际计算量摆出来。设一次网络前向的代价为 $C(N)$（Transformer 大致 $\Theta(N^2)$ 或用 FlashAttention 后的等价量级）。

• 自回归 + KV 缓存：第 $j$ 步只需为新的一个 token 做注意力，代价约 $\Theta(j)$；总代价 $\sum_{j=1}^N\Theta(j)=\Theta(N^2)$。
• 离散扩散：每步都要对整条序列重算，代价 $C(N)=\Theta(N^2)$；共 $n$ 步，总代价 $\Theta(nN^2)$。

于是加速比大约是 $N/(nN)=1/n$ ······ 只有当 $n\ll 1$ 时才更快，而 $n\ge1$，所以在这个粗略模型下离散扩散反而更慢。这正是讲义把「No KV caching → Less Speed?!」列为缺点的原因。

那为什么大家还认为它有速度潜力？因为上面的比较忽略了两件事：(a) 自回归的 $N$ 次前向是严格串行的，每次都要等上一次结束，硬件利用率极低（受限于内存带宽而非算力）；离散扩散的 $n$ 次前向虽然每次更重，但每次都是一个饱和的大 batch 矩阵乘，在 GPU 上吞吐率高得多。(b) $n$ 可以远小于 $N$。所以真正的结论是：离散扩散把「延迟受限于串行深度」换成了「吞吐受限于总算力」，这是一个硬件相关、任务相关的权衡，不是一个普适的「更快」。讲义的问号是诚实的。

推导：为什么「更难学」

自回归模型只需要建模 $N$ 个条件分布 $p(z_j\mid z_{<j})$，对应 $N$ 种「已知前缀」的情形。

MDLM 的网络在训练中会遇到任意的掩码模式：给定一个已揭示位置的子集 $S\subseteq\set{1,\dots,N}$，网络必须给出 $p(z_j\mid \set{z_i}_{i\in S})$ 对所有 $j\notin S$。这样的子集有 $2^N$ 个。换句话说，一个 MDLM 在参数量相同的情况下，要同时扮演所有 $N!$ 种生成顺序下的自回归模型（任意顺序自回归）。这是一个严格更难的学习任务，通常表现为在同等算力下困惑度略逊于自回归模型。这就是讲义「Need to learn how to generate Tokens in any order → Harder to learn?!」的含义。

反过来，这也正是它的优势来源：正因为学了所有顺序，才能任意顺序生成、才能做中间填空与编辑。能力和难度是同一枚硬币的两面。