# Flow Matching Guide and Code: Discrete Flow Matching


我们来详细整理并解释这段关于连续时间马尔可夫链（CTMC）的内容，使其更易于理解。

---

## 6. 连续时间马尔可夫链模型

### 核心思想：CTMC 是什么？
CTMC 是一种用于生成**离散数据**（比如文本、类别数据）的模型。你可以把它想象成一个在有限个离散状态之间随时间跳转的“粒子”，它按照一定的“速率”从一个状态跳到另一个状态。这与之前讨论的“流模型”（用于连续数据，如图像）形成对比，CTMC 是后续“离散流匹配”模型的基础。

### 6.1 离散状态空间和随机变量
- **状态空间**：模型的数据点由 $d$ 个坐标（或“标记”）组成，每个坐标可以从 $K$ 个可能的值中选取。例如，在文本中，每个“标记”可能是一个词或一个字符。整个状态空间 $S$ 是所有可能的 $d$ 维组合的集合。
- **随机变量**：我们用 $X$ 表示一个随机状态。它的概率分布由概率质量函数（PMF）$p_X(x)$ 给出，表示 $X$ 恰好等于某个特定状态 $x$ 的概率。
- **狄拉克δ函数**：这是一个指示函数，当两个状态相等时值为1，否则为0。它用于表示确定性事件。

---

### 6.2 CTMC 生成模型

#### 模型如何定义？
CTMC 定义了一个随时间变化的随机变量族 $\{X_t\}_{0 \le t \le 1}$，$t$ 从0到1。它由两个关键要素完全确定：
1.  **初始分布**：在 $t=0$ 时，$X_0$ 服从某个初始概率分布 $p(x)$（通常是我们能轻松采样的分布，如均匀分布或单点分布）。
2.  **转移核**：描述了在极短时间 $h$ 内，从当前状态 $x$ 转移到另一个状态 $y$ 的概率。其数学形式为：
    $$
    p_{t+h|t}(y|x) = \delta(y, x) + h \cdot u_t(y, x) + o(h)
    $$
    其中：
    - $\delta(y, x)$：是“留在原处”的概率主体部分。
    - $u_t(y, x)$：是**速率**或**速度**，它告诉我们从 $x$ 到 $y$ 的概率变化有多快。如果 $u_t(y, x)$ 很大，那么从 $x$ 到 $y$ 的跳转就更可能发生。
    - $o(h)$：是比 $h$ 更高阶的无穷小量，确保 $h$ 很小时，它几乎不影响概率。
    - 速率 $u_t$ 必须满足一个“速率条件”：
        *   **非负性**：对于任何 $y \neq x$，有 $u_t(y, x) \ge 0$（跳转速率非负）。
        *   **守恒性**：$\sum_y u_t(y, x) = 0$。这意味着留在原处的速率 $u_t(x, x) = -\sum_{y \neq x} u_t(y, x)$ 是负的，保证所有概率之和为1。

**形象理解**：如果把概率想象成“质量”，速率就是质量从状态 $x$ 流向状态 $y$ 的“速度”。正速率代表流入，负速率（对角线上）代表流出，总和为零保证了总质量守恒。

#### 如何从 CTMC 采样？
要生成一个样本路径 $\{X_t\}$：
1.  从初始分布 $p(x)$ 中采样 $X_0$。
2.  然后，用近似方法（如欧拉方法）逐步更新。一个常用且保证概率有效的更新规则是：
    $$
    \mathbb{P}(X_{t+h}=y|X_t=x) =
    \begin{cases}
    \exp\left[ h \cdot u_t(x, x) \right] & \text{如果 } y = x \\
    \frac{u_t(y, x)}{|u_t(x, x)|} \left( 1 - \exp\left[ h \cdot u_t(x, x) \right] \right) & \text{如果 } y \neq x
    \end{cases}
    $$
    这个公式保证了无论 $h$ 多大，更新后的概率分布都是有效的。

---

### 6.3 概率路径和科尔莫戈罗夫方程

#### 边际概率如何变化？
CTMC 的核心是描述概率质量函数 $p_t(x)$（即 $X_t$ 的分布）如何随时间演变。这个演变由**科尔莫戈罗夫方程**控制，它类似于连续情况下的“连续性方程”：
$$
\frac{d}{dt} p_t(y) = \sum_x u_t(y, x) p_t(x)
$$
这个方程直观地解释了 $p_t(y)$ 的变化率：
- **流入**：从其他状态 $x$ 转移到 $y$ 的概率（$u_t(y, x) p_t(x)$）。
- **流出**：从 $y$ 转移到其他状态的概率（$u_t(x, y) p_t(y)$）。

因此，方程右边可以重写为：
$$
\frac{d}{dt} p_t(y) = \underbrace{\sum_{x \neq y} u_t(y, x) p_t(x)}_{\text{流入}} - \underbrace{\sum_{x \neq y} u_t(x, y) p_t(y)}_{\text{流出}}
$$
这完美体现了“概率质量守恒”：一个状态的概率增加等于净流入。

#### 关键定理
**定理13（离散质量守恒）** 是一个重要结论，它告诉我们：
> 如果一组速率 $u_t$ 和一条概率路径 $p_t$ 满足科尔莫戈罗夫方程，且 $u_t$ 满足速率条件，那么 $u_t$ 就是生成这条路径的“速度”。反之亦然。

#### 6.3.1 保概率速度
这是一个非常有用的性质。如果我们已经有一组速率 $u_t$ 生成了一条概率路径 $p_t$，那么我们可以**添加任何“无散度”的速度项 $v_t$**，只要它满足：
1.  **速率条件**：$v_t(y, x) \ge 0$ 对于 $y \neq x$，且 $\sum_y v_t(y, x) = 0$。
2.  **无散度条件**：$\sum_x v_t(y, x) p_t(x) = 0$（即净流量为零）。

那么新的速度 $\tilde{u}_t = u_t + v_t$ **生成的边际概率路径 $p_t$ 是完全一样的**。这意味着我们可以在不改变最终生成的数据分布的前提下，自由地调整模型内部的“运动轨迹”。这对于设计更高效或更稳定的生成模型（如离散流匹配）非常有用。

---

### 总结

- **CTMC 的定位**：它是生成离散数据的模型框架，是连接“连续流模型”和“离散流匹配”的桥梁。
- **核心组成**：
    - **状态**：离散的、由多个坐标组成的数据点。
    - **速率**：定义状态间转移快慢的核心参数。
    - **初始分布**：过程的起点。
- **核心方程**：**科尔莫戈罗夫方程**，描述了概率分布如何随时间按“流入-流出”的规律演变。
- **关键性质**：**保概率速度**，允许我们在不改变最终数据分布的情况下，灵活调整模型内部的转移机制。


## 7 离散流匹配（Discrete Flow Matching, DFM）

离散流匹配（DFM）将第4节介绍的**流匹配（Flow Matching）**思想从连续空间（如图像）推广到了**离散状态空间**（如文本、类别数据）。它通过构造一个连续时间马尔可夫链（CTMC），使得在时间 $t=0$ 时样本服从源分布 $p$，在时间 $t=1$ 时样本逼近目标分布 $q$，并利用可学习的速度场（速率）来驱动这个演化过程。训练目标是最小化一个离散版本的流匹配损失。

与连续情况类似，DFM 的流程如下：
1. **定义概率路径** $p_t$：它是从源分布 $p$ 到目标分布 $q$ 的平滑插值。
2. **构建 CTMC 模型**：用参数化的速度 $u_t^\theta$ 描述状态间的转移速率。
3. **训练**：通过最小化离散流匹配损失，使模型生成的概率路径逼近真实路径。

下面我们逐步拆解这些内容。

---

### 7.1 数据与耦合

- **源分布 $p$** 和 **目标分布 $q$**：我们想把从 $p$ 中采样的数据 $X_0$ 逐步转变为从 $q$ 中采样的数据 $X_1$。
- **耦合 $\pi_{0,1}$**：它描述了 $X_0$ 和 $X_1$ 之间的联合分布。两种常见选择：
  - **独立耦合**：$\pi_{0,1}(x_0,x_1)=p(x_0)q(x_1)$，即二者独立。
  - **特定耦合**：例如在翻译任务中，$x_0$ 和 $x_1$ 是同一个句子在不同语言中的表示，因此不是独立的。另外，有时我们会为源分布添加一个特殊标记（如“mask”），使其成为一个固定的“起点”。

---

### 7.2 离散概率路径

概率路径 $p_t(x)$ 是时间 $t$ 时随机变量 $X_t$ 的概率质量函数（PMF）。我们希望 $p_0=p$，$p_1=q$。

为了构建 $p_t$，我们引入一个条件随机变量 $Z$（例如，在后面的混合路径中，$Z=(X_0,X_1)$）。首先定义**条件概率路径** $p_{t|Z}(x|z)$，然后通过边际化得到：
$$
p_t(x) = \sum_{z} p_{t|Z}(x|z) \, p_Z(z).
$$
这样，只要条件路径满足边界条件 $p_{0|Z}=p$，$p_{1|Z}=q$，边际路径也会自动满足边界条件。

---

### 7.3 边际化技巧（Marginalization Trick）

这是从条件速度构造边际速度的核心工具。假设对于每个 $z$，我们有一个条件速度 $u_t(\cdot,\cdot|z)$，它生成了条件概率路径 $p_{t|Z}(\cdot|z)$。那么，边际速度定义为：
$$
u_t(y,x) = \mathbb{E}\bigl[ u_t(y,X_t|Z) \mid X_t = x \bigr] = \sum_z u_t(y,x|z) \, p_{Z|t}(z|x),
$$
其中后验 $p_{Z|t}(z|x) = \frac{p_{t|Z}(x|z)p_Z(z)}{p_t(x)}$。

**定理14（离散边际化技巧）** 说明：如果条件速度生成了条件路径，那么按上述公式计算的边际速度就会生成边际路径 $p_t$。这个定理的证明与连续情况类似，核心是交换微分与求和，并利用贝叶斯公式。

---

### 7.4 离散流匹配损失

我们希望训练一个参数化的速度 $u_t^\theta$，使其尽可能接近真实（或条件）速度。损失函数定义为：
$$
\mathcal{L}_{\text{DFM}}(\theta) = \mathbb{E}_{t, X_t \sim p_t} D_{X_t}\!\left( u_t(\cdot, X_t),\, u_t^\theta(\cdot, X_t) \right).
$$
其中：
- $t$ 均匀采样自 $[0,1]$，
- $X_t$ 从边际概率路径 $p_t$ 中采样，
- $D_x$ 是定义在凸集 $\Omega_x$ 上的 **Bregman 散度**（一种度量两个向量差异的函数）。

$\Omega_x$ 是所有满足速率条件（非负和为零）的速度向量的集合：
$$
\Omega_x = \left\{ v \in \mathbb{R}^S \mid v(y)\ge 0\ (y\neq x),\ v(x)=-\sum_{y\neq x}v(y) \right\}.
$$
实际中我们通常使用条件形式：
$$
\mathcal{L}_{\text{CDFM}}(\theta) = \mathbb{E}_{t, Z, X_t \sim p_{t|Z}} D_{X_t}\!\left( u_t(\cdot, X_t|Z),\, u_t^\theta(\cdot, X_t) \right).
$$
**定理15** 保证了这两个损失有相同的梯度，因此训练时可以使用更易计算的条件版本。

---

### 7.5 因子化路径与速度

在离散空间中，状态空间 $S = [K]^d$ 大小是 $K^d$，当 $d$ 和 $K$ 较大时（如句子长度为 512，词汇量为 10000），直接建模 $u_t(y,x)$ 的维度是不可行的。因此我们采用**因子化**技巧：假设状态变化只能**每次改变一个坐标**，这样速度可以分解为每个坐标独立的速度。

#### 因子化速度的公式
$$
u_t(y,x) = \sum_{i=1}^d \delta(y^{\bar{i}}, x^{\bar{i}})\; u_t^i(y^i, x),
$$
其中 $\delta(y^{\bar{i}}, x^{\bar{i}})$ 指示除了第 $i$ 个坐标外，其余坐标是否相同。这意味着：
- 只有当 $y$ 与 $x$ 仅在第 $i$ 个坐标上不同时，该速度才非零；
- 每个坐标 $i$ 的局部速度 $u_t^i(y^i, x)$ 只关注第 $i$ 个坐标的变化。

此时，模型的输出维度从 $K^d$ 降低为 $d \times K$，变得可行。

#### 因子化 CTMC 的采样
在 $o(h)$ 精度下，转移概率可以因子化为各坐标独立的更新：
$$
\mathbb{P}(X_{t+h}^i = y^i \mid X_t = x) = \delta(y^i, x^i) + h\, u_t^i(y^i, x) + o(h).
$$
这意味着我们可以**按坐标独立采样**：对每个位置，根据其局部速度独立决定是否改变以及改变成什么值。

#### 因子化概率路径
我们还可以构造因子化的条件路径：
$$
p_{t|Z}(x|z) = \prod_i p_{t|Z}^i(x^i|z).
$$
**命题2** 表明，如果每个一维条件路径 $p_{t|Z}^i$ 由速度 $u_t^i(y^i, x^i|z)$ 生成，那么整体路径的生成速度恰好是因子化速度 $u_t(y,x|z) = \sum_i \delta(y^{\bar{i}}, x^{\bar{i}}) u_t^i(y^i, x^i|z)$。

**定理16（离散因子化边际化技巧）** 进一步给出了边际速度的因子化形式，使得我们可以通过条件速度的平均来得到可学习的边际速度。

---

### 7.5.3 因子化速度的条件离散流匹配损失

利用因子化，损失也可以分解为各坐标的 Bregman 散度之和：
$$
\mathcal{L}_{\text{CDFM}}(\theta) = \mathbb{E}_{t, Z, X_t} \sum_i D_{X_t}^i\!\left( u_t^i(\cdot, X_t|Z),\, u_t^{\theta,i}(\cdot, X_t) \right),
$$
其中每个坐标的散度定义在 $\Omega_{x^i}$（一维速率条件集合）上。这样训练时只需学习每个坐标的局部速度模型。

---

### 7.5.4 混合路径（Mixture Paths）

混合路径是一种简单而实用的概率路径构造方法。我们取 $Z = (X_0, X_1)$，即同时包含源和目标样本。定义条件路径为：
$$
p_{t|0,1}^i(x^i|x_0,x_1) = \kappa_t \delta(x^i, x_1^i) + (1-\kappa_t) \delta(x^i, x_0^i),
$$
其中 $\kappa_t$ 是一个从 0 到 1 的单调函数（调度器），满足 $\kappa_0=0,\ \kappa_1=1$。这意味着在每个时刻 $t$，第 $i$ 个坐标要么保持源值 $x_0^i$（概率 $1-\kappa_t$），要么已经变成目标值 $x_1^i$（概率 $\kappa_t$）。整个过程独立于其他坐标。

#### 条件速度的显式形式
对上述路径求导可得生成速度：
$$
u_t^i(y^i, x^i|x_0,x_1) = \frac{\dot\kappa_t}{1-\kappa_t}\bigl[ \delta(y^i, x_1^i) - \delta(y^i, x^i) \bigr].
$$
这非常直观：它只驱动从当前值 $x^i$ 向目标值 $x_1^i$ 的转移。

{{< admonition tip "推导过程" false >}}

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

$$
p_{t|0,1}^i(x^i|x_0,x_1) = \kappa_t \, \delta(x^i, x_1^i) + (1-\kappa_t) \, \delta(x^i, x_0^i)
$$

### 通俗理解
这个公式描述的是：**在已知源 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：条件速度（转移速率）

$$
u_t^i(y^i, x^i|x_0,x_1) = \frac{\dot\kappa_t}{1-\kappa_t}\bigl[ \delta(y^i, x_1^i) - \delta(y^i, x^i) \bigr]
$$

### 通俗理解
这个公式给出的是：**在已知源和目标的情况下，如果当前 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 前向方程）为：

$$
\frac{d}{dt} p_t(x_0) = -\lambda_t \, p_t(x_0), \quad
\frac{d}{dt} p_t(x_1) = \lambda_t \, p_t(x_0)
$$

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

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

$$
\frac{d}{dt}(1-\kappa_t) = -\dot\kappa_t = -\lambda_t (1-\kappa_t)
$$
$$
\frac{d}{dt}\kappa_t = \dot\kappa_t = \lambda_t (1-\kappa_t)
$$

两个方程等价，解得：

$$
\lambda_t = \frac{\dot\kappa_t}{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$（吸收态）。

为了用一个公式统一表达，我们引入克罗内克函数：
$$
u(y|x) = \lambda_t \bigl[ \delta(y, x_1) - \delta(y, x) \bigr]
$$
验证：
- 若 $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。
- 这两个公式共同构成了混合路径的核心，是离散流匹配中常用的一种简单而有效的概率路径构造方法。

{{< /admonition >}}

#### 后验参数化
通过边际化技巧，边际速度可以写成：
$$
u_t^i(y^i, x) = \sum_{x_1^i} \frac{\dot\kappa_t}{1-\kappa_t}\bigl[ \delta(y^i, x_1^i) - \delta(y^i, x^i) \bigr] \, p_{1|t}^i(x_1^i|x),
$$
其中 $p_{1|t}^i(x_1^i|x)$ 是给定 $X_t=x$ 时，目标值 $X_1^i$ 的后验概率。因此，我们只需要学习这个后验分布（即“去噪”模型）即可得到速度。

{{< admonition tip "推导过程" false >}}
这个公式是离散流匹配中**边际速度**的表达式。它告诉我们：在只知道当前状态 $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$）：

$$
u_t^i(y^i, x^i \mid x_0, x_1) = \frac{\dot\kappa_t}{1-\kappa_t}\bigl[ \delta(y^i, x_1^i) - \delta(y^i, x^i) \bigr].
$$

这个公式告诉我们：**如果我知道源 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)$。这就是上面公式的含义：

$$
u_t^i(y^i, x) = \sum_{x_1^i} \underbrace{\frac{\dot\kappa_t}{1-\kappa_t}\bigl[ \delta(y^i, x_1^i) - \delta(y^i, x^i) \bigr]}_{\text{条件速度（假设目标为 $x_1^i$）}} \times \underbrace{p_{1|t}^i(x_1^i|x)}_{\text{目标 $x_1^i$ 的后验概率}}.
$$

---

## 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}$ 后，就得到了速率的大小。

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

最终的结果是：  
$$
u_t^i(y^i, x) = \frac{\dot\kappa_t}{1-\kappa_t}\Bigl( p_{1|t}^i(y^i|x) - \mathbf{1}_{x^i=y^i} \Bigr),
$$
即我们熟悉的边际速度形式。

---

## 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)$：

$$
u_t^i(B,x) = 2 \sum_{x_1^i} \bigl[ \delta(B, x_1^i) - \delta(B,A) \bigr] \, p_{1|t}(x_1^i|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 的“理想”速度，通过后验概率平均，变成了只依赖于当前状态的“实际”速度。正是这个速度驱动了离散流匹配中的连续时间马尔可夫链，使我们能够在不知道最终目标的情况下，从源分布逐步生成目标分布。
{{< /admonition >}}
#### 两种损失形式
1. **直接学习后验**：使用 KL 散度作为 Bregman 散度，损失为：
   $$
   \mathcal{L}_{\text{CM}}(\theta) = -\mathbb{E}_{t,X_0,X_1,X_t} \log p_{1|t}^{\theta,i}(X_1^i|X_t) + \text{const}.
   $$
   这等价于最大化似然。
2. **广义 KL 损失**：使用一种特殊的 Bregman 散度，得到的损失与连续流匹配中的“噪声预测”损失类似，并且可以推导出证据下界（ELBO），用于模型评估。

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

---

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

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

**条件离散流匹配损失**定义为：
$$
\mathcal{L}_{\text{CDFM}}(\theta) = \mathbb{E}_{t, Z, X_t \sim p_{t|Z}} D_{X_t}\!\left( u_t(\cdot, X_t|Z),\; u_t^\theta(\cdot, X_t) \right),
$$
其中 $D_x$ 是定义在速率向量空间 $\Omega_x$ 上的 Bregman 散度。Bregman 散度的一般形式为：
$$
D_x(u, v) = f_x(u) - f_x(v) - \langle \nabla f_x(v), u - v \rangle,
$$
其中 $f_x$ 是一个凸函数。选择不同的 $f_x$ 会得到不同的损失函数。

---

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

对于混合路径，条件速度具有特殊形式：
$$
u_t^i(y^i, x^i | x_0, x_1) = \lambda_t \bigl[ \delta(y^i, x_1^i) - \delta(y^i, x^i) \bigr], \quad \lambda_t = \frac{\dot\kappa_t}{1-\kappa_t}.
$$
为了将速度的学习转化为后验分布的学习，我们选择 **广义 KL 散度**（也称为“生成散度”）作为 Bregman 散度。具体地，对于速率向量 $u, v \in \Omega_x$，定义：
$$
D_x(u, v) = \sum_{y} \left[ u(y) \log \frac{u(y)}{v(y)} - u(y) + v(y) \right],
$$
当 $u$ 和 $v$ 都是概率分布时，这退化为 KL 散度；但这里 $u$ 和 $v$ 是速率（可为负），实际上需要将其分解为正负部分。对于混合路径，由于速率只涉及两个状态，该散度可以简化为可解析计算的形式。

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

---

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

参数化的边际速度 $u_t^\theta$ 由后验分布 $p_{1|t}^\theta$ 构造：
$$
u_t^{\theta,i}(y^i, x) = \lambda_t \bigl( p_{1|t}^{\theta,i}(y^i|x) - \mathbf{1}_{x^i=y^i} \bigr).
$$
而条件速度 $u_t(\cdot|x|Z)$ 依赖于 $x_1^i$：
$$
u_t^i(y^i, x^i | x_0, x_1) = \lambda_t \bigl( \delta(y^i, x_1^i) - \mathbf{1}_{x^i=y^i} \bigr).
$$
现在将这两个表达式代入散度中。由于 Bregman 散度是凸的且满足一定的可加性，我们可以计算每个坐标的散度，并利用边际化技巧将条件期望转化为后验的期望。

最终，经过代数运算（涉及指数和对数的性质），可以得到：
$$
D_{x^i}\!\left( u_t^i(\cdot|x|Z),\; u_t^{\theta,i}(\cdot|x) \right) = \lambda_t \left[ \mathbf{1}_{x^i=x_1^i} - p_{1|t}^{\theta,i}(x^i|x) + (1-\mathbf{1}_{x^i=x_1^i}) \left( -\log p_{1|t}^{\theta,i}(x_1^i|x) \right) \right] + \text{常数（与 } \theta \text{ 无关）}.
$$

---

## 4. 得到广义 KL 损失

忽略与 $\theta$ 无关的常数，并对所有坐标求和、对 $t, Z, X_t$ 取期望，我们得到：
$$
\mathcal{L}_{\text{GKL}}(\theta) = \mathbb{E}_{t, Z, X_t} \sum_i \lambda_t \left[ (1 - \mathbf{1}_{X_t^i = X_1^i}) \bigl( -\log p_{1|t}^{\theta,i}(X_1^i|X_t) \bigr) + \bigl( \mathbf{1}_{X_t^i = X_1^i} - p_{1|t}^{\theta,i}(X_t^i|X_t) \bigr) \right].
$$
这就是代码中实现的损失表达式。注意，通常我们将 $\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 的负值：
$$
\mathcal{L}_{\text{GKL}} = -\mathbb{E} \left[ \log \frac{p_{1|t}^\theta(X_1|X_t)}{p_{\text{true}}(X_t|X_0,X_1)} \right] + \text{常数},
$$
因此最小化该损失等价于最大化对数似然的下界。这为模型评估提供了理论依据。

---

## 7. 总结

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

{{< /admonition>}}


#### 采样算法
给定训练好的后验 $p_{1|t}^{\theta,i}$，采样过程如下（对每个坐标独立进行）：
- 从 $p_{1|t}^{\theta,i}(\cdot|x)$ 采样 $x_1^i$（即预测目标值）。
- 根据混合速度执行欧拉步：
  $$
  \mathbb{P}(X_{t+h}^i=y^i|X_t=x) = \delta(y^i, x^i) + h \frac{\dot\kappa_t}{1-\kappa_t}\bigl[\delta(y^i, x_1^i) - \delta(y^i, x^i)\bigr] + o(h).
  $$
  由于 $h$ 很小，实际常用指数欧拉公式来保持概率非负且和为 1。

#### 概率保持速度（可选）
在独立耦合假设下，我们还可以添加一个**散度无关**的速度分量，它不改变边际分布但可能改善采样质量。具体形式为：
$$
v_t^i(y^i,x^i|x_1) = \frac{\dot\kappa_t}{\kappa_t}\bigl[\delta(y^i, x^i) - p(x^i)\bigr],
$$
将其乘以系数 $c_t$ 加到原速度上，不会影响 $p_t$，但可以调整采样路径。

---

### 总结

离散流匹配（DFM）通过以下关键技术将流匹配思想迁移到离散数据上：
- **CTMC** 作为概率演化的载体。
- **因子化速度** 将高维状态空间的建模分解为独立的坐标级建模，大大降低了计算复杂度。
- **混合路径** 提供了一种简单有效的概率路径构造方法，其速度可以解析写出，且可通过学习后验分布来参数化。
- **边际化技巧** 允许使用条件速度训练，从而避免了对整个状态空间的枚举。

DFM 在文本生成、离散时间序列建模等任务中表现出色，是连接连续流匹配与离散生成模型的重要桥梁。