MIT 6.S184 on My Site

MIT_6.S184_5 — Discrete Diffusion Models

Mon, 01 Jun 2026 09:53:07 +0800

引

终于来到了本次课程的最后一讲。在前四讲中，我们介绍了欧氏空间中扩散模型的基本原理、训练方法、神经网络架构以及一些常用的数学工具。核心假设是：数据可以用一个向量在$R^d$空间中表示，且不同数据之间是连续、光滑的。这正是 flow matching 和 score matching 能够借助微分方程建模的基础。但在实际应用中，很多数据并非天生处于连续空间——例如文本生成、蛋白质组成等场景，数据是由一个个离散 token 构成的序列。我们无法得到一个连续的场，让这些 token 在其中平滑流动。因此，需要一套专门处理离散数据的生成框架，这就是 Discrete Diffusion Models。
从这一设定出发，ODE、SDE 等连续微分方程工具将不再适用。取而代之的是连续时间马尔可夫链（Continuous-Time Markov Chains, CTMCs）。接下来，我们将介绍 CTMCs 的基本概念与原理，并展示如何用它构建一个 LLM 式的离散生成模型。
温馨提示：本讲对应的课程视频我认为老师讲的不太好，有部分内容讲得比较模糊？也跳了不少，所以还是推荐大家认真看一遍note。

Continuous-Time Markov chain (CTMC) models

首先，对于离散数据，我们定义一个字母表 $vocabulary = {v_1, v_2, \cdots, v_V}$, 其中每一个$v_i$ 都是一个独立的token，于是我们可以定义一个离散空间 $ S = V^d $, 其中d是序列长度，对于$x \in S$, $x = (x_1, x_2, \cdots, x_d)$, 其中每一个$x_i$ 都是vocabulary中的一个token。此时，你就应该能明白为什么ODE,SDE等连续微分方程在这里不起作用了吧。

补充一下Markov process的概念，马尔可夫过程是一个无记忆的过程，即未来的状态只与当前状态有关，而与过去的状态无关。假设存在一个随机过程$X_t \in S$, $t \in (0,1)$, 满足：

$$ \underbrace{p(X_{t+h} \mid X_t, X_{t_1}, \cdots, X_{t_k})}_{\text{prob. of future given present and past}} = \underbrace{p(X_{t+h} \mid X_t)}_{\text{prob. of future given present}} \qquad (\forall\, 0 < h,\ 0 \leq t_1 < t_2 < \cdots < t_k < t) $$

需要知道的是，ODE和SDE都是Markov process，因为由公式可知，下一个时间点的概率分布可以由当前的概率分布、当前的向量场、当前的的噪声和score函数唯一确定，而与之前的状态无关。对于离散空间的Markov process, 我们称之为Markov chain, 故命名为Continuous-Time Markov Chain (CTMC)。

我们称$p(X_{t+h}|X_t)$为transition probabilities, 由马尔可夫性质知道，一个CTMC过程完全由$X_0 \sim p_0$和transition probabilities唯一确定。因此，后面提到CTMC，我们默认就是transition probabilities了。

此时，有一个很严重的问题：在欧式空间中，对于两个数据点，我们可以通过一条连续的路径将它们连接起来，即二者坐标作差得到的向量就是当前最优的路径方向；但是在上述定义的离散空间中，如何衡量两个点之间的距离呢？又如何定义一条从A到B的路径呢？对此，我们定义一个rate matrix $Q_t(y|x) \in R$, which effectively summarizes the rate of jumping (or switching) from state $x \in S$ to state $y \in S$ at time t. 这个定义和欧式空间中的是有本质区别的，欧式空间中我们是通过一个方向场来从A到B的，而在离散空间，我们不得不一步到位，直接从A突变到B了。

此外，这个rate matrix还必须满足如下两个性质：

(1) $\quad $ Outgoing rates are positive: $Q_t(y|x) \geq 0$ $\quad$ for all $x \neq y$.

(2) $\quad $ Rate staying equals negative outgoing rate: $Q_t(x|x) = -\sum_{y \neq x} Q_t(y|x)$ $\quad$ for any x.

这两个性质说明，从x变化到其他状态的rate必然是非负的，如果不可能switch到某个状态，那么Q就是0；同时，对于x状态，它自己的rate必然和其他状态switch到x的rate之和抵消。故对角线上的元素是非正的，非对角线元素是非负的。

有了rate matrix，我们希望它能发挥和vector field在欧式空间中类似的作用, 有如下奇妙定义： $$ \left. \frac{d}{dh} p_{t+h|t}(X_{t+h} = y \mid X_t = x) \right|_{h=0} = Q_t(y \mid x), \qquad \forall x, y \in S,\ 0 \le t $$ 直觉理解左侧是从x到y的transition probability的瞬时变化率，似乎和rate matrix有点符合哈，但是我们还需要证明一下这个定义是符合上述两个性质的.

证明如下：

对于第一条，当h=0时，显然$X_t$来不及变化，所以如果$y \neq x$, 那么此时概率为0，而当h>0时，概率必然是非负的，所以说明在0处导数非负。

对于第二条，推导如下：

$$ \sum_{y \neq x} \frac{d}{dh} p(X_{t+h} = y \mid X_t = x) \big|_{h=0} = \sum_{y \neq x} Q_t(y \mid x) = \frac{d}{dh} \left(1 - p(X_{t+h} = x \mid X_t = x)\right) \big|_{h=0} = - Q_t(x \mid x) $$

ok，由上面的定义我们可以知道，对于一个CTMC过程，我们一定能找到一个满足那两条性质的rate matrix(直接求导即可)；反过来，如果我们有一个满足那两条性质的rate matrix，是否能找到一个CTMC过程使得那个公式成立呢？如果能找到，是否是唯一的呢？答案是肯定的，证明略😆，在note中的Section C 自行阅读。这种唯一存在性质让我们可以先用一个神经网络预测出$Q_t(y|x)$，然后必然会有一个唯一的Markov chain 与之对应。

note中的Example 34还是不错的，结合图像可以生动理解rate matrix的含义。

Simulation of CTMC

ok，现在让我们尝试用上述定义的rate matrix来模拟一个CTMC过程。令h>0是一个时间步，$p_{init}$是S的一个分布，在之前的ODE和SDE中，我们通常使用高斯分布作为初始分布，但是在离散空间中，不存在什么高斯分布，所以我们使用$p_{init} = Uniform(S)$, 即在S上均匀分布。采样$X_0 \sim p_{init}$, 对于每一个时间步，有$X_{t+h} \sim p_{t+h|t}(\cdot | X_t)$. 于是，只要我们知道$p_{t+h|t}(\cdot | X_t)$, 即transition probabilities, 就可以模拟了。但是对于大多数CTMC来说，transition probabilities是没有closed-form solution的，所以我们能用的只有rate matrix了, 由于rate matrix反应了transition probabilities的瞬时变化率，所以我们可以使用泰勒展开来近似transition probabilities:

$$ p_{t+h|t}(X_{t+h} = y \mid X_t = x) = p_{t|t}(X_t = y \mid X_t = x) + h Q_t(y \mid x) + R_t(h) = \mathbf{1}_{y=x} + h Q_t(y \mid x) + R_t(h) $$

其中$R_t(h)$是高阶误差项，当h很小时可以忽略。

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

于是，我们对S中所有的y进行上式计算得到$p_{t+h|t}(\cdot | x)$, $X_{t+h}$将按照不同x的概率进行采样。

$$ X_{t+h} \sim \tilde{p}_{t+h|t}(\cdot \mid x) = \big( \mathbf{1}_{y=x} + h Q_t(y \mid x) \big)_{y \in S} $$

这里有一个问题哈，就是我们说transition probabilities没有闭式解，但是Q应该也没有闭式解啊，所以上面的Q从哪里来呢？从后面看，这个Q应该是用神经网络预测出来的？这里其实就是告诉你可以通过变化公式来用Q进行sample next state。

CTMC model

现在，让我们用神经网络来参数化一个CTMC过程。A CTMC model (or discrete diffusion model) is given by an initial distribution $p_{init}$ over S and a neural network $Q_t^{\theta}$ with parameters θ such that for every input $x \in S$ the model returns a single column of the rate matrix : $$x \rightarrow \{ Q_t^{\theta}(y|x) \}_{y \in S}$$ 如此一来，我们就可以用神经网络预测出an entire column, 而这个正是上面写的next state的公式中需要的。
千万不要沾沾自喜，我们看似构建了从state to next state的模型，但是需要知道的是S是一个非常非常非常大的空间，所以上面的y的可能取值是$V^d$个，指数爆炸，没有一个计算机的内存能装下这个。于是，我们需要对这个state -> next state的过程进行一些压缩和限制，一个比较好的想法是factorized（绝大多数CTMC模型都采用这种方法），即每次最多改变一个token，这样就大大减少了next state的可能取值，也就不用存一个巨大的column了。

A factorized CTMC model is given by a CTMC model $Q_\theta^t$ such that for all $y = (y_1, \ldots, y_d), x = (x_1, \ldots, x_d) \in S = V^d$ $\qquad$ it holds :

\[ Q_\theta^t(y \mid x) = 0 \quad \text{whenever } y_i \neq x_i \text{ for more than one position } i. \]

于是，我们可以把那些和x最多有一个token不同的y叫做Neighbors N(x) of x,故上面的公式可以改写成：

$$ x \rightarrow \left\{ Q^\theta_t(y \mid x) \right\}_{y \in N(x)} = \begin{pmatrix} Q^\theta_t(v_{1}, 1 \mid x) & \cdots & Q^\theta_t(v_{V}, 1 \mid x) \\ & \cdots & \\ Q^\theta_t(v_{1}, d \mid x) & \cdots & Q^\theta_t(v_{V}, d \mid x) \end{pmatrix} $$

其中，$Q^\theta_t(v_i, j \mid x)$表示对于x的第j个token，我们把它变成vacabulary中的第i个token的rate, 得到$y = (x_1, \cdots, x_{j-1}, v_i, x_{j+1}, \cdots, x_d)$. 虽然是用神经网络预测的Q，但是我们仍然需要满足： (1) $\qquad$ $Q^\theta_t(v_i, j \mid x) \geq 0$ $\quad$ if $v_i \neq x_j$

(2) $\qquad$ $Q^\theta_t(x_j, j \mid x) = -\sum_{v \neq x_j} Q^\theta_t(v, j \mid x)$ $\quad$

很容易发现，使用factorized 后的next state的维度是$d \times V$, 和d，V都是成正比的，远远优于之前的$V^d$!!

需要注意的是，在上面的定义中，我们标注了两条需要满足的性质，在实际网络构架中，我们可以使用transformer作为骨干，在输出d x V的矩阵后，使用一些类似ReLU或softmax的激活函数满足第一条，然后再计算没行中not switch的rate来满足第二条。

Simulating a CTMC model

现在，我们已经有了一个用Q进行simulation的公式了，有了用$\theta$参数化的Q了，并且使用factorized的方式来限制next state的维度，于是我们可以对这个神经网络模型进行simulatioin了，伪代码如下：

不知道你有没有发现一个问题，我们之前不是说，我们只考虑Neighbors of x吗？也就是说每一步，我们最多有一个token的变化，但是在上面的代码中，我们内层对d进行循环，这样就可能出现有多个token变化的情况了。如果不考虑它违反了neighbor原则，从直观上理解这个反而是合理的，毕竟每一个时间步，我们肯定是希望每一个位置的token都有机会发生变化，不可能说固定哪个token变化吧，就算是固定了第一步变化第一个token，第二部变化第二个token，。。。，那么每一个时间步都最多变化一个token，十分低效啊，可能需要很多很多时间步才能从$x_0$到最终目标z。所以，在考虑到所有token都有机会变化(这个也是discrete difussion model 和 gpt类auto regressive model的一个重大区别，后者只能一个一个确定)和提高每一个时间步的计算效率，这个伪代码很合理。

note中提到 This approximation agrees with the full CTMC Euler step up to first order in h, but allows for a O(h2) probability of simultaneous updates to multiple tokens. 也就是说，上面的伪代码确实违反了neighbor原则，但是误差是二阶的，可以接受。为什么是二阶的呢？

对于一个token来说，如果更新的$v \neq x$, 它的概率是$h q_j(v)$. 当有两个token同时发生变化时，概率就是$h^2 q_j(v) q_{i}(v_1)$, 其中q就是Q，默认是一个常数，所以这个概率就是$O(h^2)$的。在一阶精度可以接受，其实就是说误差都是二阶及以上。

Training CTMC models

ok，现在我们已经有了一个可以参数化的CTMC模型了，并且学会了如何进行simulation了，如何训练这个模型呢？下面展示了训练的三个步骤：

The principles are the same as for flow matching:

We construct a probability path interpolating between noise and data.

We derive a conditional rate matrix and marginal rate matrix.

We learn the marginal rate matrix in a simulation-free manner.

Our goal is to train the CTMC model $Q_\theta^t$ such that

\[ X_0 \sim p_{\text{init}}, \quad X_t \text{ CTMC of } Q_\theta^t \quad \Rightarrow \quad X_1 \sim p_{\text{data}} \]

其中$p_{data}$是我们想要生成的离散数据的分布，比如互联网上的文本数据，但是显然我们不可能对$p_{data}$有一个表达式，所以我们在训练的时候逐个采样具体的数据点z。所以，你应该意识到，这个训练思路和之前的flow matching几乎一样，我们只是把中间的迭代过程从ODE变成了CTMC而已。

Conditional and Marginal Probability Path

我们定义$\delta_z(x)$ : 当$x=z$时为1，否则为0.（注意它是一个分布，只不过所有的概率质量都集中到了一个点上）故 (discrete) conditional probability path is given by a set of distributions $p_t(x \mid z)$ for $x, z \in S$ and $0 \le t \le 1$ such that \[ p_0(\cdot \mid z) = p_{\text{init}}, \qquad p_1(\cdot \mid z) = \delta_z. \] 故 (discrete) marginal probability path is then given by \[ p_t(x) = \sum_{z \in S} p_t(x \mid z) \, p_{\text{data}}(z) \] 简单带入0和1可知： $p_0(x) = p_{init}$, $p_1(x) = p_{data}$.

Note中定义：Factorized mixture path (independent noising per token)

Let $S = V^d$ and let $p_{\text{init}}(x) = \prod_{j=1}^d p_{\text{init}}^{(j)}(x_j)$ be a factorized initial distribution. Fix a scheduler $0 \le \kappa_t \le 1$ such that $\kappa_0 = 0$, $\kappa_1 = 1$ with $\frac{d}{dt} \kappa_t \ge 0$. Define the conditional path by

\[ 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). \]

这个公式是将x的d个token完全独立进行变化的结果，与我们前面提到的discrete diffusion model中每一个token的生成不遵循顺序，而是同时进行的思想是一致的。不过这里命名为Factorized mixture path, 这里的factorized应该是指每个token的变化都是独立的，和之前factorized CTMC model中每次最多改变一个token的factorized含义是不同的。

它和我们之前提到的Gaussian path的区别在于，Gaussian path是以一个概率分布为主体，从init状态，不断变形移动，最终变成data distribution; 而Factorized mixture path是有很多不同的state点z，我们从一个init点出发，每一步都根据概率选择跳转到哪个z上去，最终希望能跳到data distribution上去。总而言之，前者是一个分布连续变形过程，后者是一个离散随机跳转过程。 ps: 如果你能意识到这个点，那么说明你对整个课程的理解很不错哈🥰

Conditional and Marginal Rate Matrix

we call it a conditional rate matrix if
$$ X_0 \sim p_{\text{init}}, \quad X_t \text{ CTMC of } Q_t^z \quad \Rightarrow \quad X_t \sim p_t(\cdot \mid z) $$
先直接给出marginal rate matrix的定义：

Theorem 36 (Discrete marginalization trick). The marginal rate matrix defined by

$$ Q_t(y \mid x) = \sum_{z \in S} Q_t^z(y \mid x) \, \frac{p_t(x \mid z) \, p_{\text{data}}(z)}{p_t(x)} = \sum_{z \in S} Q_t^z(y \mid x) \, p_{1|t}(z \mid x) $$

where

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

is a valid rate matrix and fulfills the following condition:

$$ X_0 \sim p_{\text{init}}, \quad X_t \text{ CTMC of } Q_t \quad \Rightarrow \quad X_t \sim p_t. $$

这里的道理和之前flow matching中的marginalization trick是一样的，就是这个公式看起来合理，但是需要证明它正确性，我们需要证明一个CTMC过程的基本等式 —— Kolmogorov Forward equation。

Kolmogorov Forward Equation (KFE).

Let $p_t$ be a set of distributions on $S$ for every $0 \le t \le 1$. Further, let $X_t$ be a CTMC with matrix $Q_t$ and initial distribution $p_0$. Then $X_t \sim p_t$ for all $0 \le t \le 1$ if and only if the Kolmogorov Forward Equation (KFE) holds:

$$ \frac{d}{dt} p_t(x) = \sum_{y \in S} Q_t(x \mid y) \, p_t(y). $$

于是，我们只要证明先前定义的marginal rate matrix满足KFE就可以了。

有：

$$ \begin{aligned} \frac{d}{dt} p_t(x) &= \frac{d}{dt} \sum_{z \in S} p_t(x \mid z) \, p_{\text{data}}(z) \\ &= \sum_{z \in S} \frac{d}{dt} p_t(x \mid z) \, p_{\text{data}}(z) \\ &= \sum_{z \in S} \sum_{y \in S} Q_t^z(x \mid y) \, p_t(y \mid z) \, p_{\text{data}}(z) \\ &= \sum_{y \in S} \left( \sum_{z \in S} Q_t^z(x \mid y) \, \frac{p_t(y \mid z) \, p_{\text{data}}(z)}{p_t(y)} \right) p_t(y) \\ &= \sum_{y \in S} Q_t(x \mid y) \, p_t(y) \end{aligned} $$

上面的证明用到了$\frac{d}{dt} p_t(x \mid z) = \sum_{y \in S} Q_t^z(x \mid y) , p_t(y \mid z)$, 由于一个conditional path($target=z$)必然也满足KFE，对比KFE的定义可知这个就是它的conditional版本，所以这个等式是成立的。这个证明过程基本和flow matching中的一样。

不过，还有一个问题，就是为什么满足KFE就说明这个CTMC过程的分布是$p_t$呢(或者说二者等价)，在flow matching中，我们似乎没有严谨证明这个事，只是直观理解。下面我们来证明一下：

首先证明KFE是必要条件： To show that the KFE is necessary, assume that $p_t(x)$ are the true marginals of the CTMC, i.e. $X_t \sim p_t$ for every $0 \le t \le 1$. Then we can compute:

$$ \begin{aligned} \frac{d}{dt} p_t(x) &\stackrel{\text{(i)}}{=} \left. \frac{d}{dh} \right|_{h=0} p_{t+h}(x) \\ &\stackrel{\text{(ii)}}{=} \left. \frac{d}{dh} \right|_{h=0} \sum_{y} p_{t+h|t}(x \mid y) \, p_t(y) \\ &\stackrel{\text{(iii)}}{=} \sum_{y} \left. \frac{d}{dh} \right|_{h=0} p_{t+h|t}(x \mid y) \, p_t(y) \\ &\stackrel{\text{(iv)}}{=} \sum_{y} Q_t(x \mid y) \, p_t(y) \end{aligned} $$

where in (i) we simply use a time offset, in (ii) we use the definition of the transition probabilities, in (iii) we swap the sum and derivative, and in (iv) we use the definition of the rate matrix.

接下来证明KFE是充分条件： To show that the KFE is sufficient, we can rewrite the KFE in matrix form:

$$\frac{d}{dt} p_t = Q_t \, p_t$$

Note that the above is a linear ODE over vector space $R^S$. Its initial condition is fixed by $p_0$ as stated in the theorem.

$p_t$就是所有S中可能点x的概率分布（可以看成一个vector），而Q就是S中所有y到x的rate matrix的大矩阵，对于t时刻，右边直接就算出来了，所以这个ODE其实就是一个一阶线性微分方程，在初始$p_0$给定的条件下，这个ODE解出来的$p_t$必然是唯一的。

需要注意的是，之前定义中的等价的两个事件是：如果$p_t$是某个CTMC的边际分布$<=>$ $p_t$满足KFE。当前我们只证明了如果$p_t$满足KFE，那么这个$p_t$一定是唯一的，似乎没有说明它是某个CTMC的边际分布啊？产生这个疑惑的原因是，由于公式中掺杂了Q，导致我们无法直接看到一个CTMC的过程, 下面是我的一个理解：

如果$p_t$是一个CTMC过程的分布，则：

$$ p_{t+h}(x) = (p_{t+h|t}(x \mid y) \, p_t(y))_{y \in S} $$

求导得到：

$$ \frac{d}{dh} p_{t+h}(x) \big|_{h=0} = \left( \frac{d}{dh} p_{t+h|t}(x \mid y) \big|_{h=0} \, p_t(y) \right)_{y \in S} = (Q_t(x \mid y) \, p_t(y))_{y \in S} $$

而$\frac{d}{dh} p_{t+h}(x) \big|_{h=0} = \frac{d}{dt} p_t(x)$, 替换得到：

$$\frac{d}{dt} p_t(x) = \sum_{y \in S} Q_t(x \mid y) \, p_t(y)$$

而这个式子我们已经有了，再积分回去即能看出$p_t$是满足CTMC的定义了。

然而note的证明更绝：设 $q_t$ 为 CTMC 的真正边际分布。
由必要性（之前证明过的部分），我们知道 $q_t$ 一定满足 KFE：

$$ \frac{d}{dt} q_t = Q_t q_t $$

我们又构造了一个 $p_t$，并且它也满足：

$$ \frac{d}{dt} p_t = Q_t p_t $$

由 ODE 解的唯一性（给定 $p_0$ 相同），所以：

$$ p_t = q_t $$

因此 $p_t$ 就是那个 CTMC 的边际分布。太美妙了，我的证明是🐷，之所以会捣鼓出这个🐷证明是因为我一开始没看到note中那个绝妙的证明😭

ok，现在我们已经搞定了Conditional and Marginal Rate Matrix了，接下来我们以factorized mixture path为例子计算一下它的conditional rate matrix。

Example 37 (Conditional rate matrix for factorized mixture path). Set $\frac{d}{dt}\kappa_t = \dot{\kappa}_t$. The factorized mixture path has a factorized conditional rate matrix given by

$$ Q_t^z(y \mid x) = \bigl( Q_t^z(v_i, j \mid x_j) \bigr)_{v_i, j} $$

$$ Q_t^z(v_i, j \mid x_j) = \frac{\dot{\kappa}_t}{1-\kappa_t} \bigl( \delta_{z_j}(v_i) - \delta_{x_j}(v_i) \bigr) = \frac{\dot{\kappa}_t}{1-\kappa_t} \begin{cases} 0 & \text{if } x_j = z_j, \\[4pt] 1 & \text{if } v_i = z_j, \; x_j \neq z_j, \\[4pt] 0 & \text{if } v_i \neq z_j, \; x_j \neq z_j, v_i \neq x_i, \\[4pt] -1 & \text{if } v_i = x_j, \; x_j \neq z_j. \end{cases} $$

上面这个公式就非常美妙，我们最终的目标是每一个$x_i$能变成$z_i$, 所以可以看到只有当$v_i = z_i \text{ and }x_j \neq z_j$时，rate是1；其他情况下，rate要么是0，要么是-1.

如何证明这个公式就是factorized mixture path的conditional rate matrix呢？根据前面的结论，我们需要证明它满足KFE等式，由于factorized mixture path每一个token都是独立的，所以我们不妨就考虑一个token的情况：

$$ \begin{aligned} \frac{d}{dt} p_t(x \mid z) &\stackrel{\text{(i)}}{=} \frac{d}{dt} \left[ (1-\kappa_t) p_{\text{init}}(x) + \kappa_t \delta_z(x) \right] \\ &\stackrel{\text{(ii)}}{=} \dot{\kappa}_t \delta_z(x) - \dot{\kappa}_t p_{\text{init}}(x) \\ &\stackrel{\text{(iii)}}{=} \frac{\dot{\kappa}_t}{1-\kappa_t} \left( \delta_z(x) - \left[ (1-\kappa_t) p_{\text{init}}(x) + \kappa_t \delta_z(x) \right] \right) \\ &\stackrel{\text{(iv)}}{=} \frac{\dot{\kappa}_t}{1-\kappa_t} \left( \delta_z(x) - p_t(x \mid z) \right) \\ &\stackrel{\text{(v)}}{=} \frac{\dot{\kappa}_t}{1-\kappa_t} \delta_z(x) \left( 1 - p_t(x \mid z) \right) + \frac{\dot{\kappa}_t}{1-\kappa_t} (\delta_z(x)-1) p_t(x \mid z) \\ &\stackrel{\text{(vi)}}{=} \frac{\dot{\kappa}_t}{1-\kappa_t} \delta_z(x) \sum_{y \neq x} p_t(y \mid z)+ \frac{\dot{\kappa}_t}{1-\kappa_t} (\delta_z(x)-1) p_t(x \mid z)\\ &\stackrel{\text{(vii)}}{=} \sum_{y \neq x} Q_t^z(x \mid y) \, p_t(y \mid z) + Q_t^z(x \mid x) \, p_t(x \mid z) \\ &\stackrel{\text{(viii)}}{=} \sum_{y \in S} Q_t^z(x \mid y) \, p_t(y \mid z) \end{aligned} $$

(i) uses the definition of the factorized mixture path for $d=1$

(ii) is obtained by taking derivatives and setting $\frac{d}{dt}\kappa_t = \dot{\kappa}_t$

(iii) follows by simple algebra

(iv) by the definition of the factorized mixture path

(v) by simple algebra

(vi) follows from the fact that $\sum_{y \in S} p_t(y \mid z) = 1$

(vii) by the definition of the rate matrix, and (viii) by simple algebra.

其中需要关注的步骤是(vii)，这一步我们用到了先前对factorized mixture conditional rate matrix的定义，对于第一项$y \neq x$，$Q_t^z(x \mid y)$, 只有当$x=z$时跳转，rate=1；当$x\neq z$时，不跳转，rate=0. 对于第二项$y = x$, $Q_t^z(x \mid x)$, 当$x=z$时，不跳转，rate=0；当$x \neq z$时，rate=-1. 每一种结果都和$\delta_z(x)$，$\delta_z(x) - 1$是一致的，所以(vii)是ok的。于是，我们证明了这个定义下的conditional rate matrix满足KFE等式，所以它就是该factorized mixture path的conditional rate matrix了。

Learning the Marginal Rate Matrix

ok，现在我们已经有了factorized mixture path的conditional rate matrix了，而这个path正是当前绝大多数discrete diffusion/flow matching model采用的，所以我们接下来都将以这个path为例子来说明如何学习它的marginal rate matrix.
首先，我们需要根据之前的marginalization trick的定义来计算一下它的marginal rate matrix，首先给出结论：

Theorem 38 (Marginalization trick for factorized mixture path). The marginal rate matrix of the factorized mixture path is factorized and has the form

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

where $p_{1|t}(z_j = v_i \mid x)$ is the conditional probability of the $j$-th position ($j$-th token in the sequence) being equal to $v_i$ given the full noisy sequence $x$. 它表示在时刻t，给定当前状态x，token位置j变成v_i的rate。

如何证明这个结论呢？

$$ \begin{aligned} Q_t(v_i, j \mid x) &= \sum_{z \in S} Q_t^z(v_i, j \mid x) \, p_{1|t}(z \mid x)\\ &\stackrel{\text{(i)}}{=} \sum_{z \in S} \frac{\dot{\kappa}_t}{1-\kappa_t} \bigl( \delta_{z_j}(v_i) - \delta_{x_j}(v_i) \bigr) \, p_{1|t}(z \mid x)\\ &\stackrel{\text{(ii)}}{=} \frac{\dot{\kappa}_t}{1-\kappa_t} \left( \sum_{z \in S} \delta_{z_j}(v_i) \, p_{1|t}(z \mid x) - \delta_{x_j}(v_i) \right)\\ &\stackrel{\text{(iii)}}{=} \frac{\dot{\kappa}_t}{1-\kappa_t} \left( p_{1|t}(z_j = v_i \mid x) - \delta_{x_j}(v_i) \right) \end{aligned} $$

至此，我们已经得到了factorized mixture path的marginal rate matrix了，即： $$ Q_t(v_i, j \mid x) = \frac{\dot{\kappa}_t }{1 - {\kappa}_t}\left( p_{1|t}(z_j = v_i \mid x) - \delta_{x_j}(v_i) \right) $$ 对于这个式子，哪个部分是需要我们训练的呢？首先${\kappa_t}$是一个scheduler，是我们人为设定的，括号中的第二项$\delta_{x_j}(v_i)$也是一个常数，唯一一个无法知道的是括号中的第一项$p_{1|t}(z_j = v_i \mid x)$, 这是一个条件概率，由于我们不知道z，也不知道它的分布，所以无法直接计算它，所以它就是我们需要训练的部分了！

在上图中，我们的训练思路是给定一个z，给定一个时刻t，给定一个x，我们希望有一个神经网络能够预测$x_j$变成$z_j$的概率能尽量接近1，然后将d个位置的loss乘起来就是我们的Disccrete Flow Matching loss了。

ok，现在万事具备，可以开始训练了，但是别急，我们先来看看$p_{init}$的另一种选择，之前我们说过，离散空间中不存在高斯分布，所以我们使用了S每一个点的均匀采样作为初始分布，但是这个方法包含了一些随机信息，有没有更简答的初始化呢？有的，有的：

Example 39 (Masked Diffusion Language Model). A specific case of the above method is masked diffusion language models (MDLMs). The idea of MDLMs is that we can extend the vocabulary of tokens $V = {v_1, \ldots, v_V}$ with a new token $[\text{mask}]$ that indicates that this token is missing (or was masked). Specifically, we set $V = {v_1, \ldots, v_V, [\text{mask}]}$ and the initial point is simply $[\text{mask}]^d$, i.e. the sequence that is all-masked. Formally, this means setting $p_{\text{init}} = \delta_{[\text{mask}]^d}$ in the above framework. The sampling procedure is illustrated in Figure 20.

就是说在初始阶段，我们直接把所有token填充成mask，这样就不会包含任何随机信息了，训练会更简单。

终于，终于，我们给出完整的训练伪代码(如下图所示)，需要注意的是对t时刻的x进行采样时，我们万万不能直接用${\kappa}_t$对noise和z进行差值采样，因为我们是离散空间，两个点之间根本无法进行差值，所以我们先用$m_j \sim Bernoulli(\kappa_t)$来决定第j个token是noise还是z的第j个token，进行d次得到x。

Discrete Diffusion Models vs Autoregress Models

Advantages

Generate Multiple Tokens in Parallel -> More Speed ?!
Generate Tokens in any order -> Text editing ?!
New probability paths -> Can we design ones that make semantic sense?

Disadvantages

No KV caching -> Less Speed ?!
Need to learn how to generate Tokens in any order -> Harder to learn?!
Autogressive order (left-to-right) makes semantic sense -> Is it worth it ?

So ？

在不同场景下，它们各有优势和劣势，至于哪个更好，让field和future决定吧！

结束了吗？

至此，本次课程基本结束，后面还有一些补充的证明附录，感兴趣的朋友可以去看一下。本讲内容很多，我们相当于是把前几讲的内容从连续空间应用到了离散空间，很多方法和概念都是一样的，当下SOTA的扩散语言模型使用的就是我们提到的这些内容（使用之前提到过的一些transformer架构）！
note的最后提出有一个Generator Matching framework, 这个框架基于Markovo processes, 希望能统一flow matching, score matching, discrete model等不同空间和数据类型的生成模型😭
本系列笔记到此结束了，如果你觉得有用的话，欢迎分享哦！如果你对笔记内容有任何问题或者想法，欢迎在评论区留言，我们一起讨论！
我将在课程大横评_0中对MIT_6.S184_5进行总结和评价，敬请期待！

MIT_6.S184_4 — Building Large-Scale Image or Video Generators

Thu, 21 May 2026 15:23:00 +0800

引

通过前面几节课的学习，我们已经基本掌握了flow matching和score matching的原理和算法，并且学会了如何根据prompt来训练一个conditional flow matching model, 接下来我们将学习如何如何使用具体的神经网络实现这些模型。
本篇内容主要分为3个部分：首先介绍具体的神经网络架构；其次介绍一下隐空间以及如何使用VAE来学习一个好的latent space; 最后介绍两个当前SOTA的生成模型。

Neural Network Architectures

我们现在需要训练的是一个guided vector field, 也就是$u^{\theta}_t(x|y)$, 这个神经网络接受3个输入，t、x和y，所以我们首先需要将这三个输入转化成向量，至于我们具体使用什么神经网络，在之前的lab2中我们使用了MLP，成功学习到了一个格子图的生成。但是对于更高维度的复杂图像，比如蛋白质或者视频，我们需要用到更加复杂的网络结构，在这里我们将介绍U-Net和difussion transformer两种网络结构。

Embedding the Conditioning Variables

Embedding Time

对于t，之前的toy model中我们可以简单的将t和x进行拼接，然后传入神经网络即可。In practice, the scalar time is often embedded in a higher dimensional space using Fourier features, allowing the model to more faithfully capture high-frequency time dependence（呃呃，看似很有道理）。具体的构造如下图所示：

Embedding Class Labels

对于y，如果它只是一个离散的类别标签，那么我们可以仿照LLM中，为每一个标签设计一个类似词向量的标签向量，note提到我们可以把这个embedding向量作为参数的一部分，然后在训练过程中同步更新，以便得到一个更好反应标签之间关系的embedding向量。

Embedding Textual Input

如果y是一段文本输入，那么情况会比较复杂，我们可以使用一个预训练的text embedding model来将文本转化为一个vector，推荐使用CLIP模型，CLIP模型在训练的时候会同时将图像和文本进行embedding，并且将它们进行对齐(using a training loss designed to encourage image embeddings to be close to their corresponding prompts, while being farther from the embeddings of other images and prompts), 于是有$y = \text{CLIP}(y_{\text{raw}}) \in \mathbb{R}^{d_{\text{CLIP}}}$.
当然如果y是一个非常非常长的文本输入，我们显然不可能用一个单vector去表示它，我们可以考虑使用一个transformer对整个text进行逐个embedding，note中也提到我们可以结合多种模型的embedding结果来得到一个更好的文本embedding，最后得到$\text{PromptEmbed}(y_{\text{raw}}) \in \mathbb{R}^{S \times k}$, 其中S是文本长度，k是embedding维度。

Diffusion Transformer

这个部分没啥好说的，直接看note，我建议是把note中的DiT Block看一遍，能get到它由3个子模块组成即可。

we note that class-conditioned DiT’s, such as the one implemented in the lab, are typically simpler and eschew the cross attention layer in favor of a time and class-based AdaNorm conditioning. 对于比较简单的条件输入y(比如类别标签)，我们没必要整一个单独的cross attention layer来将图片信息和prompt信息进行融合，简单的用t，y做一个MLP得到几个参数和x乘一乘就够了，这样就大大简化了模型的复杂度和计算量。

U-Net

直接看note吧，这种网络架构的设计没啥意思（

Working in Latent Space: (Variational) Autoencoders

现在我们有了一个完整的深度学习结构用于训练，问题是当前的embedding都是d维度的，假如我们有一张1024*1024的RGB图像，那么它的embedding结果就是$NC^{\prime}$维度的，其中$C^{\prime} = 3 * patch^2$, $N = \frac{1024^2}{patch^2}$, 故总维度就是$1024^2 * 3$, 这是一个非常非常大的向量，在image classification中，我们还可以通过卷积、pooling等操作最后得到一个维度更小的结果，但是在difussion中，我们最后的输出就是guided vector field，和输入是一样的，所以直接炸了。这仅仅是image生成啊，如果是视频生成话，我们将增加一条时间维度，那么维度再次翻倍，😇。How can we model high-dimensional images within a reasonable memory and computation budget😭

Standard Autoencoders

A natural answer to this question lies in compression: perhaps the actual space of images, for example, lies near a much lower-dimensional manifold of the high dimensional image space. More concretely, we might consider an encoder $µ_ϕ$ : $R_d \rightarrow R_k$, together with some decoder $µ_θ$ : $R_k \rightarrow R_d$, which together map raw images $x \in R_d$ to and from latents $z \in R_k$, respectively. The dimension k is typically chosen to be much smaller than d. 其中$\mu_{\phi},\mu_{\theta}$被称为autoencoder, 它们的训练目标是让$\mu_{\theta}(\mu_{\phi}(x))$尽可能的接近x, 也就是让重构误差最小化：
$$\mathcal{L}_{\text{Recon}}(\phi, \theta) = \mathbb{E}_{x \sim p_{\text{data}}} \left\|\mu_{\theta}(\mu_{\phi}(x)) - x \right\|^2$$
这个方法吧，理论上是没啥大毛病的，可以通过不断训练来让重构误差变的很小，但是但是但是我们似乎没有添加任何对$\mu_{\phi}(x)$的限制，如下图所示，这种神经网络的映射可能导致$\mu_{\phi}(x)$的分布十分混乱(可能聚成很多小簇，彼此之间不连续、有空洞, 如果我们随机输入一个x，得到的z可能根本没有落到这些小簇上，decoder由于没有遇到过这种z，可能得到奇怪的输出)，

Variational Autoencoders

于是，VAE应运而生，它和standard autoencoder的区别在于不再使用deterministic function, 而是得到一个概率分布： $$q_\phi(z|x) = \mathcal{N}(z; \mu_\phi(x), \text{diag}(\sigma_\phi^2(x))), \quad p_\theta(x|z) = \mathcal{N}(x; \mu_\theta(z), \sigma_\theta^2(z) I_d)$$ 如果我们想要采样： $$z \sim q_\phi(\cdot \mid x) \quad \text{(encode)}, \qquad x \sim p_\theta(\cdot \mid z) \quad \text{(decode)}$$ 当$\sigma_\phi^2(x) \rightarrow 0$ and $\sigma_\theta^2(z) \rightarrow 0$, VAE就退化成了standard autoencoder了。它的损失函数可以定义为： $$\mathcal{L}_{\text{VAE-Recon}}(\phi, \theta) = -\mathbb{E}_{x \sim p_{\text{data}}(x), \, z \sim q_\phi(\cdot|x)} \left[ \log p_\theta(x|z) \right]$$ 这个 loss反应了 how likely would our original data point x be if we encoded and decoded it - and we take all possible decodings/encodings into account as things have become random now. 在高斯分布的情况下，这个公式可以写为： $$\mathcal{L}_{\text{VAE-Recon}}(\phi, \theta) = \mathbb{E}_{x \sim p_{\text{data}}(x), \, z \sim q_\phi(z|x)} \left[ \frac{1}{2\sigma_\theta^2(z)} \| x - \mu_\theta(z) \|^2 + \frac{d}{2} \log \sigma_\theta^2(z) \right] + \text{const}$$

这里需要介绍一下我们在生成模型中常用的一个PDF(probability density function) : d-dimensional isotropic Gaussian. 公式如下：

$$\mathcal{N}(x; \mu, \sigma^2 I) = (2\pi\sigma^2)^{-\frac{d}{2}} \exp\left( -\frac{\|x - \mu\|^2}{2\sigma^2} \right)$$

由于这个分布的协方差矩阵是$\sigma^2 I$, 所以它在各个维度都是独立的，并且方差都一样，所以我们称它为isotropic Gaussian. 它的PDF就是将根据独立性将每个维度的PDF相乘得到的结果，即上面的公式。对它取负对数后将常量分离,再用网络预测的均值和方差带入即为$\mathcal{L}_{\text{VAE-Recon}}(\phi, \theta)$.

下面是note中的一段： The second term depending on the decoder variance controls the tradeoff between reconstruction accuracy and predictive uncertainty. Many implementations, including that in the lab, fix $\sigma_{\phi}(x)$ and $\sigma_{\theta}(z)$ to learned scalar constants (that is, independent of x and z, respectively), thereby avoiding pathological behavior and numerical stability when learning variances. Therefore, the VAE reconstruction loss in this case then becomes basically the standard autoencoder reconstruction loss up to stochasticity in the encoding and constants. 这段话提到了两个方差的影响，我们注意到$\sigma_{\theta}^2(z)$同时出现在了损失函数中第一项的分母和第二项的分子，这是一种tradeoff，如果$\sigma_{\theta}^2(z)$很大，那么第一项重建误差就会很小，精度损失大；如果$\sigma_{\theta}^2(z)$很小，那么第一项重建误差就会很大，精度控制好，但是最终采样的方差小，生成的样本会回到了之前standard autoencoder的问题，碎片化、不连续。因此，我们通常会将$\sigma_{\phi}(x)$和$\sigma_{\theta}(z)$设置为一个常数，这样就避免了学习方差时可能出现的不稳定问题，同时也让VAE的重建损失基本上退化成了standard autoencoder的重建损失。

即：

$$\mathcal{L}_{\text{VAE-Recon}}(\phi, \theta) = \mathbb{E}_{x \sim p_{\text{data}}(x), \, z \sim q_\phi(z|x)} \left[ \frac{1}{2\sigma_\theta^2} \| x - \mu_\theta(z) \|^2 \right] + \text{const}$$

重新回到VAE诞生的初衷，我们希望得到一个连续，容易学习和推广的隐空间分布，我们已经通过引入高斯分布让隐空间变得连续一些了，但是可能仍然会有空白部分，于是我们引入一个prior distribution - $p_{prior}(z)$，这里使用$p_{\text{prior}} = \mathcal{N}(0, I_k)$ (a isotropic Gaussian),然后我们希望所有x在编码后的分布都往这个prior分布靠近，这样就能保证不同x的编码结果更加连续，没有空洞。利用KL散度来定义loss有： $$ L_{\text{VAE-Prior}}(\phi) = \mathbb{E}_{\mathbf{x} \sim p_{\text{data}}(\mathbf{x})} \left[ D_{\text{KL}}(q_{\phi}(\cdot \mid \mathbf{x}) \parallel p_{\text{prior}}) \right]$$ 下图是对KL散度的补充：

让我们证明一下KL散度的非负性以及等号成立的条件：设 $P$ 和 $Q$ 为定义在离散概率空间 $\mathcal{X}$ 上的两个概率分布，且 $P(x) > 0$，$Q(x) > 0$ 对所有 $x \in \mathcal{X}$ 成立。

KL 散度的定义为：

$$ D_{\text{KL}}(P \parallel Q) = \sum_{x \in \mathcal{X}} P(x) \ln \frac{P(x)}{Q(x)} $$

考虑随机变量 $Y = \frac{Q(X)}{P(X)}$，其中 $X \sim P$。由于 $\ln$ 是凹函数，根据 Jensen 不等式有：

$$ \mathbb{E}_{P}[\ln Y] \le \ln \mathbb{E}_{P}[Y] $$

具体写出：

$$ \sum_{x} P(x) \ln \frac{Q(x)}{P(x)} \le \ln \left( \sum_{x} P(x) \cdot \frac{Q(x)}{P(x)} \right) $$

左边计算：

$$ \sum_{x} P(x) \ln \frac{Q(x)}{P(x)} = -\sum_{x} P(x) \ln \frac{P(x)}{Q(x)} = -D_{\text{KL}}(P \parallel Q) $$

右边计算：

$$ \sum_{x} P(x) \cdot \frac{Q(x)}{P(x)} = \sum_{x} Q(x) = 1 $$

因此：

$$ \ln \left( \sum_{x} Q(x) \right) = \ln 1 = 0 $$

代入 Jensen 不等式得：

$$ -D_{\text{KL}}(P \parallel Q) \le 0 $$

两边乘以 $-1$（注意不等号方向反转）：

$$ D_{\text{KL}}(P \parallel Q) \ge 0 $$

等号成立的条件：由于 $\ln$ 是严格凹函数，Jensen 不等式取等当且仅当 $Y$ 为常数，即 $\frac{Q(X)}{P(X)} = c$ 几乎处处成立。结合 $\sum Q(x) = \sum P(x) = 1$ 可得 $c = 1$，从而 $Q(x) = P(x)$ 几乎处处成立。

对于连续情形，只需将求和换为积分，证明过程完全类似。

于是，结合重建损失和先验损失，我们有： \[ L_{\text{VAE}}(\phi, \theta) = L_{\text{VAE-Recon}}(\phi, \theta) + \beta L_{\text{VAE-Prior}}(\phi) \] \[ = -\mathbb{E}_{\mathbf{x} \sim p_{\text{data}}(\mathbf{x}), \mathbf{z} \sim q_{\phi}(\mathbf{z} \mid \mathbf{x})} \left[ \log p_{\theta}(\mathbf{x} \mid \mathbf{z}) \right] + \beta \mathbb{E}_{\mathbf{x} \sim p_{\text{data}}(\mathbf{x})} \left[ D_{\text{KL}}(q_{\phi}(\cdot \mid \mathbf{x}) \parallel p_{\text{prior}}) \right] \]

下面我们计算一下KL Divergence Between Isotropic Gaussians：对于任意d的情况，由之前的isotropic Gaussian的PDF我们可以得到：

$$\mathcal{N}(x; \mu, \sigma^2 I) = (2\pi\sigma^2)^{-\frac{d}{2}} \exp\left( -\frac{\|x - \mu\|^2}{2\sigma^2} \right)$$

有$\log q(x) = -\frac{d}{2} \log(2\pi\sigma_q^2) - \frac{|x - \mu_q|^2}{2\sigma_q^2}$, $\log p(x) = -\frac{d}{2} \log(2\pi\sigma_p^2) - \frac{|x - \mu_p|^2}{2\sigma_p^2}$, 于是：

$$ D_{\text{KL}}(q \parallel p) = \mathbb{E}_{\mathbf{x} \sim q} \left[ \log q(\mathbf{x}) - \log p(\mathbf{x}) \right] = \frac{1}{2} \left[ d \log \frac{\sigma_p^2}{\sigma_q^2} + \mathbb{E}_{\mathbf{x} \sim q} \left[ \frac{\|\mathbf{x} - \mu_p\|^2}{\sigma_p^2} - \frac{\|\mathbf{x} - \mu_q\|^2}{\sigma_q^2} \right] \right]$$

又：

$$ \mathbb{E}_{\mathbf{x} \sim q} \left[ \frac{\|\mathbf{x} - \mu_p\|^2}{\sigma_p^2} \right] = \frac{1}{\sigma_p^2} \mathbb{E}_{\mathbf{x} \sim q} \left[ \|\mathbf{x} - \mu_q + \mu_q - \mu_p\|^2 \right] = \frac{1}{\sigma_p^2} \left( d\sigma_q^2 + \|\mu_q - \mu_p\|^2 \right)$$

上式用到了：

$$ \mathbb{E}_{\mathbf{x} \sim q} \left[\mathbf{x} - \mu_q\right] = 0$$

又：

$$ \mathbb{E}_{\mathbf{x} \sim q} \left[ \frac{\|\mathbf{x} - \mu_q\|^2}{\sigma_q^2} \right] = \frac{1}{\sigma_q^2} \mathbb{E}_{\mathbf{x} \sim q} \left[ \|\mathbf{x} - \mu_q\|^2 \right] = \frac{1}{\sigma_q^2} \cdot d\sigma_q^2 = d $$

综上有：

$$D_{\text{KL}}(q \parallel p) = \frac{1}{2} \left[ d \log \frac{\sigma_p^2}{\sigma_q^2} + \frac{d\sigma_q^2 + \|\mu_q - \mu_p\|^2}{\sigma_p^2} - d \right] = \frac{1}{2} \left[ d \left( \frac{\sigma_q^2}{\sigma_p^2} -\log \frac{\sigma_q^2}{\sigma_p^2} - 1 \right) + \frac{\|\mu_q - \mu_p\|^2}{\sigma_p^2} \right]$$

和note中的公式其实是一样的，在各向同性高斯分布下，各个维度的方差相同，所以没有必要写成note中那个格式。

由于$F(x) = x - \log x - 1$在x=1处取得最小值0，所以当$\frac{\sigma_q^2}{\sigma_p^2} = 1$时取最小；当$\mu_p = \mu_q$时，KL散度的第二部分最小为0。这也符合了之前提到的KL散度非负性以及等号成立的条件。

我们将标准高斯分布带入上式有： $$L_{\text{VAE-Prior}}(\phi) = \mathbb{E}_{\mathbf{x} \sim p_{\text{data}}(\mathbf{x})} \left[ D_{\text{KL}}(q_{\phi}(\cdot \mid \mathbf{x}) \parallel \mathcal{N}(0, I_k)) \right] = \mathbb{E} \left[ \frac{1}{2} \mathcal{K}(\sigma_{\phi}^2(\mathbf{x})) + \frac{1}{2} \| \boldsymbol{\mu}_{\phi}(\mathbf{x}) \|^2 \right]$$ 其中$\mathcal{K}(a) = d(a - \log a - 1)$, 将$\mu_p=0和\sigma_p=1$带入即可得到上式。最终有： $$ \mathcal{L}_{\text{VAE}}(\phi, \theta) = \mathcal{L}_{\text{VAE-Recon}}(\phi, \theta) + \beta \mathcal{L}_{\text{VAE-Prior}}(\phi) = \mathbb{E}_{\mathbf{x} \sim p_{\text{data}}(\mathbf{x}), \, \mathbf{z} \sim q_\phi(\mathbf{z} \mid \mathbf{x})} \left[ \frac{1}{2\sigma_\theta^2(\mathbf{z})} \|\mathbf{x} - \mu_\theta(\mathbf{z})\|^2 + \frac{d}{2} \log \sigma_\theta^2(\mathbf{z}) + \frac{\beta}{2} \mathcal{K}(\sigma_\phi^2(\mathbf{x})) + \frac{\beta}{2} \|\mu_\phi(\mathbf{x})\|^2 \right] $$

The four terms of the above loss function are very intuitive: 第一项是重建误差，第二项是decoder的方差，可以理解为不确定性，第三项是make latent variance 1, 第四项是make latent mean 0.

Training a VAE

ok，我们已经有了一个不错的loss定义，现在需要考虑如何训练它。现在有一个最大的问题是我们需要对z进行采样，而这个待采样的分布又是由神经网络预测的，这就导致我们的backward中间遇到了一个依概率采样的操作，这个操作压根不可以微分，所以我们需要一个reparameterization trick来解决这个问题。对于 $$q_\phi(\mathbf{z} \mid \mathbf{x}) = \mathcal{N}(\mathbf{z}; \mu_\phi(\mathbf{x}), \sigma_\phi^2(\mathbf{x}) \mathbf{I}_k)$$ 我们可以换个方法采样： $$\boldsymbol{\epsilon} \sim \mathcal{N}(0, \mathbf{I}_k), \quad \mathbf{z} = \mu_\phi(\mathbf{x}) + \sigma_\phi(\mathbf{x}) \boldsymbol{\epsilon} \;\Rightarrow\; \mathbf{z} \sim q_\phi(\cdot \mid \mathbf{x})$$ 这样一来我们每次都先采样好一个$\epsilon$, 然后直接计算出z，这样就可以微分了。于是我们改进loss函数为： $$\mathcal{L}_{\text{VAE}}(\phi, \theta) = \mathbb{E}_{\mathbf{x} \sim p_{\text{data}}(\mathbf{x}), \, \boldsymbol{\epsilon} \sim \mathcal{N}(0, \mathbf{I}_k)} \left[ \frac{1}{2\sigma_\theta^2(\mathbf{z})} \|\mathbf{x} - \mu_\theta(\mu_\phi(\mathbf{x}) + \sigma_\phi(\mathbf{x})\boldsymbol{\epsilon})\|^2 + \frac{d}{2} \log \sigma_\theta^2(\mathbf{z}) + \frac{\beta}{2} \mathcal{K}(\sigma_\phi^2(\mathbf{x})) + \frac{\beta}{2} \|\mu_\phi(\mathbf{x})\|^2 \right]$$

再结合之前说过的可以固定方差，于是：

$$\mathcal{L}_{\text{VAE}}(\phi, \theta) = \mathbb{E}_{\mathbf{x} \sim p_{\text{data}}(\mathbf{x}), \, \boldsymbol{\epsilon} \sim \mathcal{N}(0, \mathbf{I}_k)} \left[ \frac{1}{2\sigma^2} \|\mathbf{x} - \mu_\theta(\mu_\phi(\mathbf{x}) + \sigma_\phi(\mathbf{x})\boldsymbol{\epsilon})\|^2 + \frac{\beta}{2}\mathcal{K}(\sigma_\phi^2(\mathbf{x})) + \frac{\beta}{2} \|\mu_\phi(\mathbf{x})\|^2 \right]$$

本篇需要推导的部分到此结束，请自行查看note部分的VAE训练伪代码。note还提到了一些工程上的优化方法，比如加一个GAN的判别器，使用预训练模型提取特征来计算误差，$\beta$参数的调节等等。并且指出当今几乎所有SOTA生成模型都遵循先训练一个良好的VAE自编码器，然后将数据在latent中进行flow or difussion model训练，最后再decode。所以最终生成图片的指令不仅看flow or difussion model的训练效果，还要看autoencoder compression的效果以及最终recover的情况。

Stable Diffusion 3 and Meta Moive Gen Video

自行阅读note🤣

MIT_6.S184_3 — How To Condition on a Prompt

Mon, 18 May 2026 09:31:42 +0800

引

在上一讲中我们已经介绍了difussion model是如何定义，如何进行训练的，现在我们可以用flow matching 或者 score matching来训练一个generative model，并且用ODE或者SDE进行采样。但是在真正使用时，我们肯定是希望能够给定一个prompt来进行条件生成的，而我们上述的模型都是直接在$z \in data$上进行训练的，至于最终的采样结果完全是z中一个随机点，所以我们需要一些方法来让我们的模型能够根据一个prompt来进行条件生成。

Vanilla Guidance

我们首先来介绍一下最简单的vanilla guidance方法，定义如下：
$$ \begin{align*} & \textbf{Initialization: } \qquad X_0 \sim p_{\text{init}} \\ & \textbf{Simulation: } \qquad \quad dX_t = u_t^\theta(X_t \mid y) \, dt + \sigma_t \, dW_t \\ & \textbf{Goal: } \qquad \qquad \qquad X_1 \sim p_{\text{data}}(\cdot \mid y) \end{align*} $$
当这里的$\sigma_t = 0$时，我们称之为guided flow model, 我们接下来谈谈如何在guided flow model中进行训练。一个很简单的想法就是在原先的训练方式上我们就加上一个y的限制，其他都不变：给定一个prompt y, 我们就从$z \in p_{data}(\cdot|y)$中sample一个z, 然后按照之前的方式从$z$出发进行训练，类似这样：
$$\mathbb{E}_{z \sim p_{\text{data}}(\cdot|y), \, x \sim p_t(\cdot|z)} \left\| u_t^\theta(x \mid y) - u_t^{\text{target}}(x \mid z) \right\|^2$$
Note that the label y does not affect the conditional probability path pt(·|z) or the conditional vector field $u_t^{\text{target}}(x|z)$ (although in principle, we could make it dependent). Expanding the expectation over all such choices of y, we thus obtain a guided conditional flow matching objective.
$$\mathcal{L}_{\text{CFM}}^{\text{guided}}(\theta) = \mathbb{E}_{(z,y) \sim p_{\text{data}}(z,y), \, t \sim \text{Unif}[0,1], \, x \sim p_t(\cdot|z)} \left\| u_t^\theta(x|y) - u_t^{\text{target}}(x|z) \right\|^2$$
可以注意到这个公式和我们之前定义的flow matching的loss函数非常相似，唯一的区别就是这里是从$p_{data}$中取一个(z,y)对来训练，同时我们的模型接受的输入是(x,y)。

这个方法理论上是可行的，但是实际训练结果证明效果并不好（相关但是不够准确），可能是模型没有学习到正确的marginal vector field, 或者是网络上的(image, text)数据本身有问题。为了强化生成结果和prompt的相关性，我们介绍一种当前SOTA(state-of-the-art)的训练方法。

Classifer-Free Guidance

Classifier Guidance

我们首先介绍Classifier Guidance，在gaussian probability path的例子中，我们有vector field和score function之间的一个线性关系，用(x|y)替换后有： $$u_t^{\text{target}}(x|y) = a_t \nabla \log p_t(x|y) + b_t x \tag{1}$$

这里的看似显然的替换值得思考，实际上$p_t(x|y)$和$p_t(x)$是一个东西，后者理解为对所有的$z \in data$进行一个加权，前者也可以看成是对所有的$z \in data(y)$进行一个加权，只不过前者是后者数据范围的一个子集罢了。于是乎，利用加权积分的线性性质在$z \in data(y)$上展开一次就得证了。

对于prompt下的概率路径$p_t(x|y)$, 我们有：

$$p_t(x|y) = \frac{p_t(x) p_t(y|x)}{p_t(y)}$$

进一步有：（where we used that the gradient $\nabla$ is taken with respect to the variable x, so that $\nabla \log p_t(y) = 0$）

$$ \nabla \log p_t(x|y) = \nabla \log \frac{p_t(x) p_t(y|x)}{p_t(y)} = \nabla \log p_t(x) + \nabla \log p_t(y|x) $$

将这个式子代入上面(1)式中，有：

$$u_t^{\text{target}}(x|y) = b_t x + a_t \left( \nabla \log p_t(x) + \nabla \log p_t(y|x) \right) = u_t^{\text{target}}(x) + a_t \nabla \log p_t(y|x)$$

这个式子可以看成两部分，第一部分是我们之前的unconditional marginal vector field, 第二部分可以理解为给定x，y分类的score function, 它的作用就是让x更倾向于分类y。之前我们提到使用vanilla guidance方法训练出来的结果不够符合prompt的要求，一个自然的想法就是我们增大第二部分的系数，即：

$$\tilde{u}_t(x|y) = u_t^{\text{target}}(x) + w a_t \nabla \log p_t(y|x), \quad \text{(classifier guidance)}$$

where w > 1 is known as the guidance scale.

How can we learn the term log $p_t(y|x)$? Note that this can be considered as a sort of classifier of noised data (i.e. it gives the log-likelihoods of y given x). So we can simply learn it via supervised learning. 这里的supervised learning的意思应该我们从data中取一对(z,y), 然后在z上加一个噪声，$x = \alpha_t z + \beta_t \epsilon$, 然后让神经网络根据这个x来预测y。

这里有一个细节是，我们用了大量的(z,y)训练一个很好的分类器之后，在第二步训练时我们是对x求解梯度，通过改变x来让它更倾向于分类y！所以我们称这个方法为classifier guidance. 由于这个方法需要额外训练一个分类器模型，需要两倍的工作量，并且当y不是一个类别标签而是一个高维度文本描述时，这个分类器的训练会非常困难，求梯度更是难上加难。

其实，只有当w=1的时候，classifier guidance才是真正的guided vector field, 当w不等于1时，虽然它的采样结果更符合prompt，但是 it is not the “true” guided vector field. （偏离真实vector field,但结果更fit，这就是reinforce）

Classifier-Free Guidance

由于上述提到的两个困难，我们需要一个classifier-free guidance的方法。 Classifier-free guidance results in the theoretically equivalent effect as classifier guidance but without having to train a separate classifier.
我们利用这个等式： $$\nabla \log p_t(x|y) = \nabla \log p_t(x) + \nabla \log p_t(y|x)$$ 代入之前的$\tilde{u}_t(x|y)$的结果中，有： $$\begin{aligned} \tilde{u}_t(x|y) &= u_t^{\text{target}}(x) + w a_t \nabla \log p_t(y|x) \\ &= u_t^{\text{target}}(x) + w a_t \left( \nabla \log p_t(x|y) - \nabla \log p_t(x) \right) \\ &= u_t^{\text{target}}(x) - \left( w b_t x + w a_t \nabla \log p_t(x) \right) + \left( w b_t x + w a_t \nabla \log p_t(x|y) \right) \\ &= (1-w) u_t^{\text{target}}(x) + w u_t^{\text{target}}(x|y). \end{aligned}$$ 你会发现这个公式也是分为两个部分，第一部分依旧是unguided vector field, 第二部分是guided vector field. 也许你会自然的想到我们可以训练两个vector field神经网络来估计它们，然后加权即可。 Wait , wait , 也许我们只用训练一个网络呢？We may augment our label set with a new, additional ∅ label that denotes the absence of conditioning. We can then treat $u^{target}_t (x) = u^{target}_t (x|∅)$. With that, we do not need to train a separate model to reinforce the effect of a hypothetical classifier. This approach of training a conditional and unconditional model in one (and subsequently reinforcing the conditioning) is known as classifier-free guidance (CFG)!

Note that the construction $\tilde{u}_t(x|y) = (1−w)u^{target}_t (x) + wu^{target}_t (x|y)$, is equally valid for any choice probability path‼️, not just a Gaussian one. When w = 1, it is straightforward to verify that $\tilde{u}_t(x|y) = u^{target}_t (x|y)$. Our derivation using Gaussian paths was simply to illustrate the intuition behind the construction, and in particular of amplifying the contribution of a hypothetical “classifier” $\nabla \log p_t(y|x)$.

Training with Classifier-Free Guidance

由上面的推导，我们现在只需要训练一个神经网络来预测$u^{target}_t(x|y)$即可，但是这里有一个问题：我们从data中取得(z,y)时，y不可能是$\emptyset$, 所以需要人工制造空标签。
$$\mathcal{L}_{\text{CFM}}^{\text{CFG}}(\theta) = \mathbb{E}_{(z,y) \sim p_{\text{data}}(z,y), \; t \sim \text{Unif}[0,1], \; x \sim p_t(\cdot|z), \; y \leftarrow \varnothing \text{ w.p. } \eta} \left\| u_t^\theta(x|y) - u_t^{\text{target}}(x|z) \right\|^2$$
其中$y \leftarrow \varnothing \text{ w.p. } \eta$的意思是replace y = $\emptyset$ with probability $\eta$.
可以自行欣赏一下note中的训练伪代码，需要注意的是使用ODE最终前向得到的X1并不 aligned with X1 $\in p_{data}(·|y)$ if we use a weight w>1 (之前提到过 the guided vector field not “true” guided vector field, 故从噪音还原到数据的过程必然偏航), 然而, 从结果上看这个方法更能符合prompt提示。所以说Classifier-free guidance 是一种启发式方法，其有效性主要体现在其优秀的实际结果上。事实上现在的AI生成图片或者视频都高度依赖于$w>=4$的classifier-free guidance!
我们上述的讨论都是基于flow model的，对于difussion model来说，如下图所示，利用(55)式将$u^{\theta}_t(X_t)$换成$u^{\theta}_t(X_t|y)$即可进行SDE。

MIT_6.S184_2 — Score Functions and Score Matching

Sun, 17 May 2026 14:49:20 +0800

引

还记得我们曾经在lab1中用过一个$∇\log𝑝(𝑋_𝑡)$来进行ODE吗，本文将详细介绍与之有关的知识。
在上一讲中我们已经推导出了vector field的公式，并且给出了flow matching的训练算法，现在我们将学习如何处理diffusion model 并且如何使用score matching来训练它。

Conditional and Marginal Score Functions

什么是score function? 对于一个概率分布p(x)，它的score function定义为： $$s(x) = \nabla \log p(x)$$
这个函数的意义在于它告诉我们在x这个点上，对于x的log-likelihood的最快增长方向，也就是它的梯度。

Conditional Score Function

于是，我们定义Conditional Score Function为：
$$s_t(x|z) = \nabla_x \log p_t(x|z)$$
我们可以以之前提到过的Gaussian conditional path为例来求解一下它的conditional score function。已知：
\[p_t(x|z) = \mathcal{N}(\alpha_t z, \beta_t^2 I_d)\]
所以
\[p_t(x|z) = \frac{1}{(2\pi)^{d/2} \beta_t^d} \exp\left(-\frac{1}{2\beta_t^2} \|x - \alpha_t z\|^2\right)\]

\[\log p_t(x|z) = -\frac{d}{2} \log(2\pi) - d \log \beta_t - \frac{1}{2\beta_t^2} \|x - \alpha_t z\|^2\]

\[s_t(x|z) = \nabla_x \log p_t(x|z) = -\frac{1}{\beta_t^2} (x - \alpha_t z)\]
我们会发现：哎嘿，这个score function恰好是x和z的一个线性组合，而我们之前求出的gaussian conditional path的vector field也是x和z的一个线性组合，我们完全可以将它们之间进行一个转化。已知：(推导过程见上一篇)
\[u_t^{\text{target}}(x|z) = \frac{\dot{\beta_t}}{\beta_t} x + \left( \dot{\alpha_t} - \frac{\dot{\beta_t}}{\beta_t} \alpha_t \right) z \]

令：

\[a_t = \beta_t^2 \frac{\dot{\alpha_t}}{\alpha_t} - \dot{\beta_t}\beta_t\]

\[b_t = \frac{\dot{\alpha_t}}{\alpha_t}\]

则有：

\[u_t^{\text{target}}(x|z) = a_t s_t(x|z) + b_t x\]

Marginal Score Function

同理定义一个
$$s_t(x) = \nabla_x \log p_t(x)$$
我们可以先凭感觉定义一个Marginal Score Function:
\[ s_t(x) = \nabla_x \log p_t(x) = \int \nabla_x \log p_t(x|z) \frac{p_t(x|z) \, p_{\text{data}}(z)}{p_t(x)} \, dz \]
这个公式依然看起来很合理，对x求marginal score,我们就先固定x，然后对所有的z进行一个后验概率加权平均。
证明：
$$ \begin{align*} \nabla \log p_t(x) &= \frac{\nabla p_t(x)}{p_t(x)} \\ &= \frac{\nabla \int p_t(x|z) p_{\text{data}}(z) \, dz}{p_t(x)} \\ &= \int \frac{\nabla p_t(x|z) p_{\text{data}}(z)}{p_t(x)} \, dz \\ &= \int \nabla \log p_t(x|z) \, \frac{p_t(x|z) p_{\text{data}}(z)}{p_t(x)} \, dz \end{align*} $$
直接就推导出来了，奇妙，中间运用了两次$log$的微分法则。
然后我们可以将上面conditional score function在gaussian conditional path的例子中求出的结果带入$u_t^{\text{target}}(x)$的公式，得到：
\[u_t^{\text{target}}(x) = a_t s_t(x) + b_t x\]

conditional and marginal denoiser

本段视频中老师没讲，纯纯个人理解，如果有错误欢迎指正。
在gaussian probility path的例子中，我们的$u_t^{\text{target}}(x|z)$和$s_t(x|z)$都是x和z的线性组合，而它们的marginal version则是在它们的基础上进行一个对z的后验概率加权积分。
于是note中提到它们都可以看成是 a linear reparameterization of the posterior mean $E[z|x]$, 也就是说我们给定一个噪声x，然后我们通过后验概率加权平均的方式得到一个对z的估计，然后就能求出score function 或者是 vector field了。这个东西有什么用呢，note中说一是形式简单，二是可能会数值/训练稳定好。
conditional denoiser就是$D_t(x|z) = z$， marginal denoiser就是$D_t(x) = E[z|x] = \int z , \frac{p_t(x|z) p_{\text{data}}(z)}{p_t(x)} , dz \stackrel{(i)}{=} \frac{1}{\dot{\alpha}_t \beta_t - \alpha_t \dot{\beta}_t} \left( \beta_t u_t^{\text{target}}(x_t) - \dot{\beta}_t x_t \right)$，它们都是对z的一个估计。其中(i)是通过带入$u_t^{\text{target}}(x|z)$与z的关系式就得到了。

The denoiser has a very intuitive interpre-tation: it is the expected value of clean data z given noisy data x.People often call such models denoising diffusion models as learning D_t and learning u_target are theoretically equivalent. 这里note中提到了一个很有意思的点，will the denoiser always output a “clean” data point? Why or why not, and what might this depend on? 这个让我想起了EAI课程中wh老师举过一个例子，在无人驾驶中，如果路中间有一个障碍物，那么我们记录左转和右转都是有效，写出一个loss函数是 (x-1)^2 + (x+1)^2, 但问题是这个loss函数最优解是x=0，完全撞上了障碍物，大错特错了。这里的道理是一样的，如果我们的data distribution有多个峰值，那么对于一个噪声数据点x来说，它的后验概率可能在多个峰值上都有比较大的值，那么它的denoiser输出的就是这些峰值的一个加权平均，可能并不是一个“clean”的数据点了。

Sampling with SDEs

我们之前展示过如何利用vector field来进行ODE sampling了，那么对于diffusion model来说，我们就需要用到SDE了。
SDE Extension Trick, 仿照flow matching的ODE版本，我们定义： $$ \begin{aligned} X_0 &\sim p_{\text{init}}, \\ dX_t &= \left( u_t^{\text{target}}(X_t) + \frac{\sigma_t^2}{2} \nabla \log p_t(X_t) \right) dt + \sigma_t \, dW_t, \\ \Rightarrow \quad X_t &\sim p_t \quad (0 \le t \le 1). \end{aligned} $$ 你可能觉得很奇怪，我们原本的SDE想法是一个确定的vector field加上一个noise项，为什么这里要加上一个score function呢？这一点其实很好理解，我们一般是先确定了一个probability path, 然后计算vector field来进行ODE sampling的，但是对于SDE来说，加上了一个noise，会导致中间的$p_t$发生变化，所以我们需要加上这个score function来修正noise带来的影响。

加上这个修正项，可以保证每个时刻的p_t总概率分布是正确的，只不过每一个x采样的路径可能会受到noise影响而歪歪扭扭。这里是对于任何一个σ_t都成立的，但是由于训练中的误差，采样时的误差，我们可以挑选一个最佳的σ_t来得到最好的采样效果。

更进一步的，可以代入marginal score function：（Only for gaussian probability path‼️）
$$X_0 \sim p_{\text{init}}, \quad dX_t = \left( (a_t + \frac{\sigma_t^2}{2}) \nabla \log p_t(X_t) + b_t X_t \right) dt + \sigma_t \, dW_t$$

$$\Rightarrow \quad X_t \sim p_t \quad (0 \le t \le 1)$$
如何证明这个SDE Extension Trick呢？只要证明它满足Fokker-Planck equation就可以了。

至于为什么满足这个定理就可以做到消除噪音的影响，请自行参考note中的附录Section B. 下面我们来计算在满足Fokker-Planck equation的条件下，新的$u_t$是什么。

首先定义 $\Delta$ ：

$$ \Delta w_t(x) = \sum_{i=1}^d \frac{\partial^2}{\partial x_i^2} w_t(x) = \operatorname{div}(\nabla w_t)(x) $$

for scalar field $w_t: \mathbb{R}^d \to \mathbb{R}$ 其实就是对一阶导数算div。

证明：

$$ \begin{align} \partial_t p_t(x) &\stackrel{(i)}{=} -\operatorname{div}\left(p_t u_t^{\text{target}}\right)(x) \\ &\stackrel{(ii)}{=} -\operatorname{div}\left(p_t u_t^{\text{target}}\right)(x) - \frac{\sigma_t^2}{2} \Delta p_t(x) + \frac{\sigma_t^2}{2} \Delta p_t(x) \\ &\stackrel{(iii)}{=} -\operatorname{div}\left(p_t u_t^{\text{target}}\right)(x) - \operatorname{div}\left(\frac{\sigma_t^2}{2} \nabla p_t\right)(x) + \frac{\sigma_t^2}{2} \Delta p_t(x) \\ &\stackrel{(iv)}{=} -\operatorname{div}\left(p_t u_t^{\text{target}}\right)(x) - \operatorname{div}\left(p_t \frac{\sigma_t^2}{2} \nabla \log p_t\right)(x) + \frac{\sigma_t^2}{2} \Delta p_t(x) \\ &\stackrel{(v)}{=} -\operatorname{div}\left(p_t \left( u_t^{\text{target}} + \frac{\sigma_t^2}{2} \nabla \log p_t \right) \right)(x) + \frac{\sigma_t^2}{2} \Delta p_t(x), \end{align} $$

符合Fokker-Planck equation的格式，于是新的$u_t$就是 $u_t^{\text{target}} + \frac{\sigma_t^2}{2} \nabla \log p_t$。

Langevin Dynamics

小表格中提到了Langevin Dynamics, 其实就是一个$p_t = p$的特殊情况，也就是说我们在采样过程中保持$p_t$不变，此时我们令$u_t^{\text{target}} = 0$ (带入Fokker-Planck equation你会发现是成立的)，于是得到了Langevin Dynamics的SDE： $$ dX_t = \frac{\sigma^2}{2} \nabla \log p(X_t) dt + \sigma dW_t $$
这个SDE的意义在于它可以让我们采样到一个固定的分布p，而且它具有收敛性，即便$p_{init}$离p十万八千里，只要有足够的时间，最后一定会趋向于p.
开头提到的我们曾经在lab1中就使用了这个公式从一个随机噪声分布收敛到一个高斯分布！用的就是这个公式了。

Optional: GLASS Flows, Stochastic evolution with ODEs

The remarkable property of SDE sampling (compared to ODEs) is that the evolution becomes stochastic, i.e. the initial point X0 does not fully determine Xt for t > 0. Perhaps surprisingly, it is also possible to get the same stochastic transitions purely via ODEs via a simple sampling trick called GLASS Flows [20]. This allows to exploit the stochastic nature of SDEs (e.g. via search algorithms) while keeping the efficiency of ODEs.
SDE能够让采样更多样性，可以在采样中跳过一些训练误差带来的不良影响，但是效率更低；GLASS Flows则是一个采样技巧，可以让我们在保持ODE效率的同时，也能得到SDE的采样多样性。

Score Matching

首先定义loss函数： $$\begin{align} \mathcal{L}_{\text{SM}}(\theta) &= \mathbb{E}_{t \sim \text{Unif}, \, z \sim p_{\text{data}}, \, x \sim p_t(\cdot|z)} \left\| s_t^\theta(x) - \nabla \log p_t(x) \right\|^2 \\ \mathcal{L}_{\text{CSM}}(\theta) &= \mathbb{E}_{t \sim \text{Unif}, \, z \sim p_{\text{data}}, \, x \sim p_t(\cdot|z)} \left\| s_t^\theta(x) - \nabla \log p_t(x|z) \right\|^2 \end{align}$$ 问题是这个$\nabla \log p_t(x)$是一个关于z的积分，算个🥚，所以第一个loss直接放弃，和之前flow matching类似的我们可以证明：

$$ \mathcal{L}_{\text{SM}}(\theta) = \mathcal{L}_{\text{CSM}}(\theta) + C, $$

where $C$ is independent of parameters $\theta$. Therefore, their gradients coincide:

$$ \nabla_\theta \mathcal{L}_{\text{SM}}(\theta) = \nabla_\theta \mathcal{L}_{\text{CSM}}(\theta). $$

In particular, for the minimizer $\theta^$, it will hold that $s_t^{\theta^} = \nabla \log p_t$.

note中又以gaussian为例子详细描述了Denoising Diffusion Models 以及最后的训练算法，这里最重要的insight就是为什么要叫denoising, 通过带入公式我们可以发现，如果$\beta_t$比较小，那么很容易数值爆炸(也就是要想训练出来的话必须加很多的噪声)。我们可以先提取出$\beta_t$, 舍掉它，然后你会发现优化项里是$\beta_t s_t^\theta(x) + \epsilon$, 于是可以令
$$-\beta_t s_t^\theta(x) = \epsilon_t^\theta(x) \quad \Rightarrow \quad \mathcal{L}_{\text{DDPM}}(\theta) = \mathbb{E}_{t \sim \text{Unif}, \, z \sim p_{\text{data}}, \, \epsilon \sim \mathcal{N}(0, I_d)} \left\| \epsilon_t^\theta(\alpha_t z + \beta_t \epsilon) - \epsilon \right\|^2$$
这实际上就是用$\theta$ 来在固定z，固定t，固定x的情况下去预测原始噪声$\epsilon$，所以叫denoising了。
note的最后也是给出了训练的伪代码，当然loss函数我们既可以用score版本也可以用noise版本。其实，在使用gaussian probability path的情况下，这些量之间都是线性关系，都可以相互转化，我们完全可以只搞一个最简单的训练器，比如(noise或者D(z|x)),然后转化得到vector field和score function, 最后应用SDE采样就好了。

MIT_6.S184_1 — Flow Matching

Mon, 04 May 2026 10:42:29 +0800

Lec2

第二节课主要讲了Flow Matching的原理，具体包括了conditional path 和 marginal path，以及它们对应的loss，最终的Flow Matching训练算法。本节课对应的视频部分地方可能讲的不是很清楚，所以一些数学推导细节可能需要自己结合文档来进一步理解。

Conditional and Marginal Probability Path

Conditional path的定义：

$$p_0(\cdot|z) = p_{\text{init}},\quad p_1(\cdot|z) = \delta_z \quad \text{for all } z\in\mathbb{R}^d$$

这里的$\delta_z$是一个Dirac delta distribution，表示只在z这个点上有概率，是一个确定性的分布。简单理解就是让一个随机初始化的分布经过一系列变化最终收缩到一个确定性的z上。当然这里的z你不要以为真的是一个点了，它是从一个$p_{data}$中sample得到的，如果$p_{data}$是一个点集，那么z就是一个点；如果$p_{data}$是一个图像集合，那么z就是一张图片。
Conditional path 就是从0-1过程中每一个t时，$p_{init}$和$p_{\delta_z}$之间的一个插值分布，一种过渡状态。

One example: Gaussian Conditional Probability Path

我们希望得到一个从标准高斯分布到一个确定性的z的conditional path，即
$$p_t(\cdot|z) = \mathcal{N}(\alpha_t z,\; \beta_t^2 I_d)$$或者写成 $$z \sim p_{\text{data}},\; \epsilon \sim p_{\text{init}} = \mathcal{N}(0, I_d) \;\Rightarrow\; x = \alpha_t z + \beta_t \epsilon \sim p_t$$
其中t=0时，$\alpha=0, \beta=1$, t=1时，$\alpha=1, \beta=0$, 所以 t=0时， $p_0(\cdot|z) = \mathcal{N}(0, I_d)$, t=1时，$p_1(\cdot|z) = \mathcal{N}(z, 0) = \delta_z$。

marginal probability path

在知道了conditional path大概是个什么东西后，我们可以定义marginal path了， $$z \sim p_{\text{data}},\; x \sim p_t(\cdot|z) \;\Rightarrow\; x \sim p_t $$ 对于dataset中的每一个z，我们都能得到一个$p_t(\cdot|z)$, marginal path就是按照z的概率对它们的conditional path进行积分加权。 $$ p_t(x) = \int p_t(x|z) \, p_{\text{data}}(z) \, dz $$

Conditional and Marginal Vector Fields ‼️

对于conditional path, 我们按照定义有：
$$\frac{d X_t}{d t} = u^{\text{target}}_t (X_t | z)$$
以之前提到的 Gaussian conditional path 为例，我们来求解一下它的 vector field。

已知：

\[ x = \alpha_t z + \beta_t \epsilon \]

直接求导得到：

\[ \frac{d x}{d t} = \dot{\alpha_t} z + \dot{\beta_t} \epsilon \]

其中 $\dot{\alpha_t} = \frac{d \alpha_t}{d t}, \dot{\beta_t} = \frac{d \beta_t}{d t}$。

由于

\[ \epsilon = \frac{x - \alpha_t z}{\beta_t} \]

所以

\[ \frac{d x}{d t} = \dot{\alpha_t} z + \dot{\beta_t} \frac{x - \alpha_t z}{\beta_t} = \dot{\alpha_t} z + \frac{\dot{\beta_t}}{\beta_t} x - \frac{\dot{\beta_t}}{\beta_t} \alpha_t z = \frac{\dot{\beta_t}}{\beta_t} x + \left( \dot{\alpha_t} - \frac{\dot{\beta_t}}{\beta_t} \alpha_t \right) z \]

对于marginal path,
$$u^{\text{target}}_t (x) = \int u^{\text{target}}_t (x|z) \, \frac{p_t(x|z) \, p_{\text{data}}(z)}{p_t(x)} \, dz$$
其中
$$\frac{p_t(x|z) \, p_{\text{data}}(z)}{p_t(x)}$$
is “posterior over data points z given noisy data x” where pdata(z) is the prior distribution. The marginal vector field then is simply a average: for every possible data point z it takes the velocity ut(x|z) - i.e. the direction that would bring us to z - and then weighs this velocity by how much we believe that x comes from z. Averaging over all data points, we obtain the marginal vector field.
对于conditional vector field，我们毫无疑问直接对X求导得到u。但是对于marginal vector field,我们这样定义的u，你是否怀疑过它的正确性？我们在上一个模块marginal path中定义了:
$$ p_t(x) = \int p_t(x|z) \, p_{\text{data}}(z) \, dz $$
我认为这个定义是合理的，所以我们需要证明u的定义和它是自洽的，怎么证明呢？唉哎，只要我们能证明下面这个式子，那么u的定义就是合理的。
$$ \partial_t p_t(x) = -\operatorname{div}\left(p_t \, u_t^{\text{target}}\right)(x), \quad \forall x \in \mathbb{R}^d,\; 0 \le t \le 1 $$

The left-hand side $\partial_t p_t(x)$ describes how much the probability $p_t(x)$ at $x$ changes over time. Intuitively, the change should correspond to the net inflow of probability mass. For a flow model, a particle $X_t$ follows along the vector field $u^{\text{target}}_t$. As you might recall from physics, the divergence measures a sort of net outflow from the vector field. Therefore, the negative divergence measures the net inflow. Scaling this by the total probability mass currently residing at $x$, we get that the net $-\operatorname{div}(p_t \cdot u_t)$ measures the total inflow of probability mass. Since probability mass is conserved (always integrates to 1), the left-hand and right-hand side of the equation should be the same!

上面的note原文解释了这个式子的原理，就是概率质量守恒，只要我们的u能让概率质量守恒，那么说明它作为vector field就是正确的。证明如下： $$ \begin{aligned} \partial_t p_t(x) &= \partial_t \int p_t(x|z) \, p_{\text{data}}(z) \, dz \quad &\text{(i)} \\ &= \int \partial_t p_t(x|z) \, p_{\text{data}}(z) \, dz \quad &\text{(ii)} \\ &= -\int \operatorname{div}\left( p_t(\cdot|z) \, u_t^{\text{target}}(\cdot|z) \right)(x) \, p_{\text{data}}(z) \, dz \quad &\text{(iii)} \\ &= -\operatorname{div} \int p_t(x|z) \, u_t^{\text{target}}(x|z) \, p_{\text{data}}(z) \, dz \quad &\text{(iv)} \\ &= -\operatorname{div} \left( p_t(x) \int u_t^{\text{target}}(x|z) \, \frac{p_t(x|z) \, p_{\text{data}}(z)}{p_t(x)} \, dz \right) \quad &\text{(v)} \\ &= -\operatorname{div} \left( p_t \, u_t^{\text{target}} \right)(x) \quad &\text{(vi)} \end{aligned} $$

这里面由于对x，z，t的微分和积分变量都是独立的，所以可以交换顺序，同时前后各引入了$p_t(x)$来凑出我们定义的u的形式，最终得到了我们想要的结果。

Learning the Marginal Vector Field

上一部分的理解我认为是本次课程最硬核的部分，现在我们在定义了vector field之后，接下来我们就可以用神经网络来学习这个vector field了，想要训练神经网络，首先要定义损失函数。

flow matching loss

$$ \mathcal{L}_{\text{FM}}(\theta) = \mathbb{E}_{t \sim \text{Unif}, \, x \sim p_t} \left[ \left\| u^\theta_t(x) - u^{\text{target}}_t(x) \right\|^2 \right] $$

这个看似很合理，但是有一个巨大的问题：我们需要计算$u^{\text{target}}_t(x)$，而这个是一个关于z的积分，而z所在的dataset有可能非常大，甚至是无穷的，无法计算！

conditional flow matching loss

由于上面的loss无法计算，所以我们引入了一个conditional flow matching loss，来代替上面的loss： $$ \mathcal{L}_{\text{CFM}}(\theta) = \mathbb{E}_{t \sim \text{Unif}, \, z \sim p_{\text{data}}, \, x \sim p_t(\cdot|z)} \left[ \left\| u^\theta_t(x) - u^{\text{target}}_t(x|z) \right\|^2 \right] $$
唉哎，不对啊，我们不是要让$u^\theta_t(x)$去逼近那个marginal vector field吗，用这conditional vector field来训练，难道没有问题吗？可以证明： $$ \mathcal{L}_{\text{FM}}(\theta) = \mathcal{L}_{\text{CFM}}(\theta) + C $$其中C是一个和$\theta$无关的常数，即 $$ \nabla_\theta \mathcal{L}_{\text{FM}}(\theta) = \nabla_\theta \mathcal{L}_{\text{CFM}}(\theta) $$ 这里的证明思路主要是先用$|a - b|^2 = |a|^2 - 2a^T b + |b|^2$将loss展开，然后第三项明显是与$\theta$无关的C1，第二项要经过一系列神奇的妙妙变换直接和CFM的第二项相等了🤣，配凑一下得证。具体证明见Note(实在肝不动了，还是看note轻松)。

训练算法

这部分比较简单，不写了。
结合note中gaussian conditional path的例子理解即可。

MIT_6.S184_0

Fri, 01 May 2026 08:29:30 +0800

写在前面的

在上完cs336的5个模块，并且完成了其中的3个半作业后(hw1，hw2，hw4均已基本完成，hw3由于不想苦苦调参数只做了一半的样子，hw5需要训练一个RL，由于社团的h100集群网络问题非常搞宣告放弃，但是就算能连上，考虑到RL训练的困难，想要收敛也不是一个容易的事情，不过也让我意识到了基础设施建设的重要性)。给我最大的感受是当今人们对于LLM的学习和认识都大都浮于表面的，只是知其名，甚至能把原理说的头头是道，但是真的落到实处，如何从小小的数据结构开始构建一个完整的Transformer, data processing, DDP, benchmark, RL, scaling, 不亲自体验一遍是无法认识到其中的困难的，而其中的细节我认为恰恰是最重要的(如果你不想只会调调包，调调参数，试试各种agent，或者只是和网页版玩的不亦乐乎的话) （不过话又说回来了，和llm相关的硬课基本上都需要不少的财力算力支持，也需要运维管理，听可信机器学习的老师说贵校曾经开过类似的课程，不过开了一段时间宣告失败，我觉得应该就是上述原因）。这就和cs336第一节课里的一句话对上了，llm在显著降低程序员的思考和动手能力，非常有必要深入研究LLM的底层细节，当今各种agent、harness、牛鬼蛇神一天一个，我个人觉得这种应用级别的东西在出现一个最强者完成大一统之前，把时间用在尝试这些东西完全是浪费的，它们的实现细节、数学原理才是大浪底下的黄金。为此开始学习MIT Class 6.S184: Generative AI With Stochastic Differential Equations, 2026，深入理解生成模型！

宗门心法

All of these generative models generate objects by iteratively converting noise into data.
至于如何将噪声变成有意义的数据，需要研究一些微分方程(常微分，随机微分)。我们的Diffusion model, Flow Matching, 就是将这些微分方程的思想和神经网络结合起来的结果。
当我看到一份70多页的在线文档，上面是非常硬的课程note，我对mit的教学工作者唯有敬意！太牛逼了，很多人不理解一份完整且优质的教学文档是多么耗时耗力！

A Reminder on Probability Theory

文档中提到有一些前置的数学知识可能会用到(Section A)，下面是一些重点：
均值的定义不仅仅是
$$\int x \, p(x) \, dx$$
可以证明它也是
$$ \mathbb{E}[X] = \arg\min_{z \in \mathbb{R}^d} \int \|x - z\|^2 \, p_X(x) \, dx = \int x \, p_X(x) \, dx $$
二次误差期望的最小参数就是均值。
对于两个随机变量X和Y，它们的联合概率对其中一个变量积分得到的是另一个变量的边缘概率。
文档中花了不少篇幅讲述E(X|Y)的定义。本人之前的理解就是Y是条件，那么要先确定条件，外层对y积分，然后在y确定的条件下对X求期望。文档中通过一个二次误差的最优函数来定义：
$$ g_* := \arg\min_{g: \mathbb{R}^d \to \mathbb{R}^d} \mathbb{E}\left[\|X - g(Y)\|^2\right] = \arg\min_{g: \mathbb{R}^d \to \mathbb{R}^d} \int \|x - g(y)\|^2 \, p_{X,Y}(x,y) \, dx dy = \arg\min_{g: \mathbb{R}^d \to \mathbb{R}^d} \int \left[ \int \|x - g(y)\|^2 \, p_{X \mid Y}(x \mid y) \, dx \right] p_Y(y) \, dy $$
对于这个我们可以理解为先遍历Y，对于每一个y，都可以对应唯一的一个值g(y)，可以由上面一阶期望的结论知道内层对x的积分误差如果想要取最小，那么一定有
$$ \mathbb{E}[X \mid Y = y] := g_*(y) = \int x \, p_{X \mid Y}(x \mid y) \, dx $$
由于Y是一个随机变量，每一个y对应一个g(y)，那么g(Y)也是一个随机变量。
$$ \mathbb{E}[X \mid Y] := g_*(Y) $$
请注意这里的E(X|Y)和E(X|Y=y)不是一个东西，后者是前者的一个特殊位置的值，前者则是关于Y的一个函数, 这个函数的意义是：对于每一个可能的 y，它给出一个数值，这个数值是在已知Y=y 的条件下，用二次距离来衡量最接近X的那个数。(唉唉，有点绕就对了😇)
一个性质：
$$ \mathbb{E}[\mathbb{E}[X \mid Y]] = \mathbb{E}[X] $$
证明如下：
$$ \mathbb{E}[\mathbb{E}[X \mid Y]] = \int \left[ \int x \, p_{X \mid Y}(x \mid y) \, dx \right] p_Y(y) \, dy = \int \int x \, p_{X,Y}(x,y) \, dx \, dy = \int x \, p_X(x) \, dx = \mathbb{E}[X] $$

Lec1

第一节课主要讲了生成模型的目标 ( p_init -> p_data ), 以及如何使用ODE和SDE来实现simulate这个过程，简单的理解就是将ODE中的常数项参数表示为$∇\log𝑝(𝑋_𝑡)$此类的关于当前分布在目标分布上的概率(或者score)，从而将轨迹朝着增大data分布概率的方向前进。至于SED中的噪声项，暂时没看出有什么深远的意义，或许是增大选择路径，丰富最终的收敛结果？

Lab0

友情提醒，lab中的一些函数没有具体实现，只有一个pass的不用写，它们的作用是函数示范。只有标注了raise NotImplementedError的函数才是需要我们自己实现的。
这个lab个人感觉没有难度，基本上都是照着公式填就对了，但是官方的代码提供了非常牛逼的演示效果(包括过程图，视频)，以及几种函数分布供我们自由组合，可以说非常的亲民，做到了用最小的任务量，理解最核心的知识点，展示最生动的效果！
核心是理解$∇\log𝑝(𝑋_𝑡)$作为$U_t(X)$的作用即可。