Flow Matching Guide and Code: Discrete Flow Matching

1. 公式1：条件概率路径

通俗理解

这个公式描述的是：在已知源 token $x_0^i$ 和目标 token $x_1^i$ 的情况下，在时刻 $t$，第 $i$ 个位置上的 token $x^i$ 的概率分布。

• $\kappa_t$ 是一个从 0 到 1 单调递增的函数，比如 $\kappa_t = t$ 或 $\kappa_t = t^2$。它告诉我们“已经变成目标的进度”。
• $\delta(a,b)$ 是克罗内克函数，当 $a=b$ 时为 1，否则为 0。所以 $\delta(x^i, x_1^i)$ 表示“$x^i$ 是否等于目标”，$\delta(x^i, x_0^i)$ 表示“$x^i$ 是否等于源”。

意思就是：在时间 $t$，这个位置上的 token 要么是源 $x_0^i$（概率 $1-\kappa_t$），要么是目标 $x_1^i$（概率 $\kappa_t$），没有其他可能。而且，随着时间增加，$\kappa_t$ 变大，变成目标的概率越来越大，最终在 $t=1$ 时 $\kappa_1=1$，必然变成目标。

举例：假设源 token 是 [MASK]，目标 token 是 “猫”，$\kappa_t = t$。那么在 $t=0.3$ 时，该位置有 70% 的概率还是 [MASK]，30% 的概率已经是 “猫”。这非常直观：我们只是简单地把源和目标按比例混合。

2. 公式2：条件速度（转移速率）

通俗理解

这个公式给出的是：在已知源和目标的情况下，如果当前 token 是 $x^i$，那么它跳转到另一个 token $y^i$ 的“速率”。速率可以理解为“单位时间内发生跳转的概率”。

• $\dot\kappa_t$ 是 $\kappa_t$ 对时间的导数，表示变化快慢。
• $\frac{\dot\kappa_t}{1-\kappa_t}$ 是一个随时间变化的跳转强度。当 $\kappa_t$ 接近 1 时，分母很小，强度变得很大，这样才能保证在最后时刻所有 token 都变成目标。
• 括号里的部分 $\delta(y^i, x_1^i) - \delta(y^i, x^i)$ 决定了跳转的方向：
  • 如果 $y^i$ 等于目标 $x_1^i$，且当前 $x^i$ 不等于目标，则括号 = 1，得到一个正的速率（表示向目标跳转）。
  • 如果 $y^i$ 等于当前 $x^i$，则括号 = -1，表示从当前状态流出的速率（这是负的，代表离开）。
  • 其他情况括号 = 0，表示没有直接跳转。

直观上：只要当前 token 还不是目标，它就会以很高的速率向目标跳转；一旦变成目标，就停在那里不再动。跳转的速率由调度函数决定，确保在 $t=1$ 时所有人都到达目标。

举例：沿用上面的例子，$\kappa_t = t$，则 $\dot\kappa_t = 1$，强度 $= 1/(1-t)$。假设当前 $t=0.3$，当前 token 是 [MASK]（源）。那么：

• 向 “猫” 跳转的速率 = $1/(0.7) \times 1 \approx 1.43$（高）。
• 保持 [MASK] 的速率 = $-1.43$（表示流出）。
• 向其他词跳转的速率 = 0。

如果当前 token 已经是 “猫”，那么所有速率为 0，不再变化。

3. 如何从公式1推导出公式2？

我们想构造一个连续时间马尔可夫链（CTMC），它的概率分布恰好是公式1。由于每个坐标独立，只需考虑一个坐标，状态空间只有两个可能：源 $x_0$ 和目标 $x_1$。我们只允许从源到目标的单向跳转（不可逆），因为一旦变成目标就不会再变回源。

设从源到目标的转移速率为 $\lambda_t$（随时间变化）。那么这个 CTMC 的主方程（Kolmogorov 前向方程）为：

其中 $p_t(x_0)$ 和 $p_t(x_1)$ 是当前时刻处于源和目标的概率。

根据公式1，$p_t(x_0) = 1-\kappa_t$，$p_t(x_1) = \kappa_t$。代入方程：

两个方程等价，解得：

现在写出速率矩阵 $u(y|x)$（即从状态 $x$ 到状态 $y$ 的速率）。对于我们的两状态链：

• 当 $x = x_0$ 时，$u(x_1|x_0) = \lambda_t$，且 $u(x_0|x_0) = -\lambda_t$（因为所有行和为零）。
• 当 $x = x_1$ 时，$u(x_1|x_1) = 0$，$u(x_0|x_1) = 0$（吸收态）。

为了用一个公式统一表达，我们引入克罗内克函数：

验证：

• 若 $x = x_0$，$y = x_1$：$\lambda_t(1-0)=\lambda_t$ ✓
• 若 $x = x_0$，$y = x_0$：$\lambda_t(0-1)=-\lambda_t$ ✓
• 若 $x = x_1$：不论 $y$ 是 $x_1$ 还是 $x_0$，括号都是 $\delta(y,x_1)-\delta(y,x_1)=0$ ✓
• 其他 $y$ 不在状态空间中，但公式仍成立（括号=0）。

将 $\lambda_t$ 代入即得公式2。

总结

• 公式1 给出了一个简单的“混合”概率路径：每个位置要么是源要么是目标，概率随时间线性变化。
• 公式2 是通过 CTMC 的主方程推导出的转移速率，它描述了从源到目标的跳转过程，速率大小由调度函数决定，确保概率路径正好是公式1。
• 这两个公式共同构成了混合路径的核心，是离散流匹配中常用的一种简单而有效的概率路径构造方法。

这个公式是离散流匹配中边际速度的表达式。它告诉我们：在只知道当前状态 $x$（而不直接知道源 $x_0$ 和目标 $x_1$）的情况下，如何计算从当前 token 跳转到其他 token 的速率。

1. 符号含义

• $i$：第 $i$ 个位置（比如句子中的第 $i$ 个词）。
• $x$：当前时刻 $t$ 所有位置的 token 状态（但公式中只关心第 $i$ 个位置，所以 $x^i$ 是该位置的 token）。
• $y^i$：我们想要知道转移到这个 token 的速率。
• $x_1^i$：可能的目标 token（即最终在 $t=1$ 时该位置的 token）。
• $p_{1|t}^i(x_1^i|x)$：给定当前状态 $x$，最终目标 token 是 $x_1^i$ 的概率（后验分布）。这是模型要学习的东西。
• $\kappa_t$：调度函数，单调从 0 到 1。$\dot\kappa_t$ 是它的导数。
• $\frac{\dot\kappa_t}{1-\kappa_t}$：一个时间相关的标量，代表“跳转强度”。
• $\delta(a,b)$：克罗内克函数，当 $a=b$ 时为 1，否则为 0。

2. 这个公式在做什么？

在离散流匹配中，我们有一个“条件速度”公式（依赖于具体的源 $x_0^i$ 和目标 $x_1^i$）：

这个公式告诉我们：如果我知道源 token 和目标 token 是什么，那么从当前 token 出发，只会跳转到目标 token，且速率是 $\frac{\dot\kappa_t}{1-\kappa_t}$（如果当前 token 不是目标）。这个速度依赖于具体的 $x_1^i$。

但在实际中，我们不知道 $x_1^i$ 是什么，因为训练时我们只能观察到当前状态 $x$，而目标 token 是隐藏的。我们需要一个不依赖于 $x_1^i$ 的边际速度，只由当前状态 $x$ 决定。

怎么办呢？我们使用平均的思想：把条件速度对所有可能的目标 token 进行加权平均，权重就是给定当前状态 $x$ 时，该目标 token 出现的概率 $p_{1|t}^i(x_1^i|x)$。这就是上面公式的含义：

3. 通俗解释

想象你在做一个填空题。

• 你知道当前的单词是某个词（比如 [MASK]）。
• 你不知道最终要填什么词，但你可以根据上下文猜测每个候选词的概率（后验）。
• 每个候选词都有一个“理想的变化路径”：如果最终答案是这个词，那么当前的 [MASK] 应该以多大的速率跳变成它。
• 但你不知道真正的答案，所以你用所有可能答案的平均速率作为最终的变化速率。

这个公式就是做这个平均：

• 对每个可能的答案 $x_1^i$，先假设它就是目标，计算一个“如果答案是它”时的速度；
• 再乘以这个答案的可能性 $p_{1|t}^i(x_1^i|x)$；
• 最后把所有可能加起来，得到最终的速度。

4. 公式中的细节

• 括号里的 $\delta(y^i, x_1^i) - \delta(y^i, x^i)$ 决定了：
  如果 $y^i$ 等于当前 token $x^i$，贡献是负的（表示流出）；
  如果 $y^i$ 等于某个候选目标 $x_1^i$，贡献是正的（表示流入）；
  否则为 0。

• 乘上 $\frac{\dot\kappa_t}{1-\kappa_t}$ 后，就得到了速率的大小。

• 求和把所有这些候选目标的贡献按概率加权平均，就得到了最终的速率。

最终的结果是：

即我们熟悉的边际速度形式。

5. 一个简单例子

假设词表只有三个词：A、B、C。当前 token $x^i = A$。模型预测的后验概率为：

• $p_{1|t}(A|x) = 0.1$
• $p_{1|t}(B|x) = 0.6$
• $p_{1|t}(C|x) = 0.3$

调度函数取 $\kappa_t = t$，则 $\frac{\dot\kappa_t}{1-\kappa_t} = \frac{1}{1-t}$。假设 $t=0.5$，则强度 $= 2$。

我们想计算从 A 跳转到 B 的边际速度 $u_t^i(B, x)$：

• 对 $x_1^i=A$：$\delta(B,A)=0$，$\delta(B,B)=0$，括号=0-0=0，贡献0。
• 对 $x_1^i=B$：$\delta(B,B)=1$，$\delta(B,A)=0$，括号=1-0=1，贡献 $2 \times 1 \times 0.6 = 1.2$。
• 对 $x_1^i=C$：$\delta(B,C)=0$，$\delta(B,A)=0$，括号=0-0=0，贡献0。

所以 $u_t^i(B,x) = 1.2$。
这意味着在 $t=0.5$，当前是 A 时，跳转到 B 的速率为 1.2（单位时间内的概率）。

同理，跳转到 C 的速率：只有 $x_1^i=C$ 时 $\delta(B,C)=0$，所以为 0。
保持 A 的速率：$u_t^i(A,x)$ 用类似方法可得为 $-1.2 - 0 = -1.2$（因为总和为 0）。

6. 总结

这个公式是从条件速度到边际速度的桥梁。它把依赖于具体目标 token 的“理想”速度，通过后验概率平均，变成了只依赖于当前状态的“实际”速度。正是这个速度驱动了离散流匹配中的连续时间马尔可夫链，使我们能够在不知道最终目标的情况下，从源分布逐步生成目标分布。

离散流匹配（DFM）中广义 KL 损失的推导，源于将条件流匹配损失与混合路径的具体形式相结合，并通过选择特定的 Bregman 散度，将速度场的拟合问题转化为对后验分布的优化。下面从基础出发，逐步推导该损失。

1. 条件流匹配损失（CDFM）

在离散流匹配中，我们有一族条件概率路径 $p_{t|Z}(x|z)$，每个 $z$ 对应一个条件速度 $u_t(\cdot,\cdot|z)$，它生成该条件路径。我们希望参数化边际速度 $u_t^\theta$，使其逼近真实边际速度 $u_t$。

条件离散流匹配损失定义为：

其中 $D_x$ 是定义在速率向量空间 $\Omega_x$ 上的 Bregman 散度。Bregman 散度的一般形式为：

其中 $f_x$ 是一个凸函数。选择不同的 $f_x$ 会得到不同的损失函数。

2. 为混合路径选择合适的 Bregman 散度

对于混合路径，条件速度具有特殊形式：

为了将速度的学习转化为后验分布的学习，我们选择 广义 KL 散度（也称为“生成散度”）作为 Bregman 散度。具体地，对于速率向量 $u, v \in \Omega_x$，定义：

当 $u$ 和 $v$ 都是概率分布时，这退化为 KL 散度；但这里 $u$ 和 $v$ 是速率（可为负），实际上需要将其分解为正负部分。对于混合路径，由于速率只涉及两个状态，该散度可以简化为可解析计算的形式。

在实际推导中，研究者通常利用 Bregman 散度的性质，直接计算条件速度与参数化速度之间的散度，并证明它可以写成关于后验分布的损失。以下给出一个简化的推导思路。

3. 代入条件速度与参数化速度

参数化的边际速度 $u_t^\theta$ 由后验分布 $p_{1|t}^\theta$ 构造：

而条件速度 $u_t(\cdot|x|Z)$ 依赖于 $x_1^i$：

现在将这两个表达式代入散度中。由于 Bregman 散度是凸的且满足一定的可加性，我们可以计算每个坐标的散度，并利用边际化技巧将条件期望转化为后验的期望。

最终，经过代数运算（涉及指数和对数的性质），可以得到：

4. 得到广义 KL 损失

忽略与 $\theta$ 无关的常数，并对所有坐标求和、对 $t, Z, X_t$ 取期望，我们得到：

这就是代码中实现的损失表达式。注意，通常我们将 $\lambda_t$ 提到期望外，并在实现时采样一个批次来近似该期望。

5. 损失的意义

• 当 $X_t \neq X_1$ 时（即当前 token 还不是目标），损失为 $\lambda_t \bigl( -\log p_{1|t}(X_1|X_t) \bigr)$，即交叉熵项，鼓励模型正确预测目标 token。
• 当 $X_t = X_1$ 时（当前 token 已是目标），损失为 $\lambda_t \bigl( 1 - p_{1|t}(X_t|X_t) \bigr)$，即惩罚模型对当前 token 的自信度过高，防止过早停止（因为理想情况下，一旦到达目标，应该不再改变，但后验概率 $p_{1|t}(X_t|X_t)$ 应接近 1，所以该项很小）。

因此，广义 KL 损失既包含了对目标 token 的预测监督，也包含了对当前状态的“置信度”调节，使得 CTMC 的采样更平稳。

6. 与 ELBO 的关系

将广义 KL 损失展开，可以证明它等于一个 ELBO 的负值：

因此最小化该损失等价于最大化对数似然的下界。这为模型评估提供了理论依据。

7. 总结

广义 KL 损失是从条件流匹配损失出发，针对混合路径选择特定 Bregman 散度（广义 KL 散度）后，经过边际化技巧和代数简化得到的结果。它直接作用于后验分布，避免了显式计算速度场，同时与 ELBO 相联系，是离散流匹配中常用的一种高效、稳定的训练目标。