Discrete MeanFlow: 通过条件转移核实现一步生成

# 基于条件转移核的单步生成 来源：https://arxiv.org/html/2605.12805 ###### 摘要 MeanFlow通过学习时间区间上的平均速度而非流匹配的瞬时速度场，实现了连续空间中的单步生成。然而，离散状态空间没有平滑轨迹或空间导数，因此连续公式无法直接应用。我们引入了离散MeanFlow，它用有限状态上概率质量的传输代替了点的运动。我们的关键对象是连续时间马尔可夫链（CTMC）的条件转移核，由此我们定义了一个平均离散速率，用于衡量时间区间内转移概率的平均变化。我们证明了一个离散MeanFlow恒等式，将该有限区间速率与端点处的瞬时CTMC生成元联系起来，其中Kolmogorov正向方程取代了连续MeanFlow的空间链式法则。基于这一恒等式，我们直接通过一种“边界内建”设计来参数化转移核，该设计保证了有效的概率输出和精确的边界条件，无需辅助损失。由于学习到的核本身就是一个概率分布，生成过程简化为一次前向传播加上一次类别采样——这意味着不需要迭代去噪、ODE积分或多步细化。我们在精确的有限状态马尔可夫链上验证了该框架，学习到的核能够高精度地恢复解析真实值；同时，在具有不同字母表大小和序列长度的因式分解合成序列生成任务上也进行了验证。 ## 1 引言 基于流匹配(Lipman et al., 2023 (https://arxiv.org/html/2605.12805#bib.bib1); Liu, 2022 (https://arxiv.org/html/2605.12805#bib.bib2))和扩散(Ho et al., 2020 (https://arxiv.org/html/2605.12805#bib.bib3); Song et al., 2021 (https://arxiv.org/html/2605.12805#bib.bib4))的生成模型在图像、音频和分子设计方面取得了显著成果，但在推理时需要进行大量的ODE或SDE积分步骤。MeanFlow(Geng et al., 2025a (https://arxiv.org/html/2605.12805#bib.bib5))通过学习时间区间上的平均速度并推导出一个将其与端点瞬时速度联系起来的恒等式，降低了连续空间中的计算成本，实现了精确的单步生成。然而，MeanFlow无法直接应用于离散状态空间，从而限制了其在语言生成等任务中的使用。与连续数据不同，词元无法从一个值平滑地移动到另一个值，因为没有连续的轨迹，也没有关于状态的空间导数。现有的离散生成模型(Austin et al., 2021 (https://arxiv.org/html/2605.12805#bib.bib6); Sahoo et al., 2024 (https://arxiv.org/html/2605.12805#bib.bib7); Lou et al., 2024 (https://arxiv.org/html/2605.12805#bib.bib8); Campbell et al., 2024 (https://arxiv.org/html/2605.12805#bib.bib9); Gat et al., 2024 (https://arxiv.org/html/2605.12805#bib.bib10))转而学习瞬时跳跃率函数，并通过模拟一个连续时间马尔可夫链（CTMC）经过多个步骤生成样本。因此，离散空间中的单步生成仍然是一个开放问题。

这项工作的核心思想是，要推广的正确对象不是状态的运动，而是状态上**概率质量**的运动。在离散空间中，点不会移动，但概率会移动。一个CTMC由其转移核完全刻画，该转移核给出了经过给定时间间隔后到达每个状态的概率。我们从这个核中定义了一个**平均离散速率**，用于衡量在有限时间间隔内概率质量在各状态之间重新分布的平均速率。我们推导了一个**离散MeanFlow恒等式**，利用Kolmogorov正向方程将核的时间导数与底层跳跃动力学联系起来，从而将该量与端点处的瞬时CTMC生成元关联起来。

我们参数化中的一个关键设计选择是我们称之为**边界内建**的方法。转移核必须满足一个条件：如果没有时间流逝，系统保持在其当前状态。我们不是通过在损失函数中加入惩罚项来强制这一点，而是将其构建到模型架构中，使得学习到的核是当前状态处的δ函数和神经网络输出的混合，其混合系数在零时间间隔时消失。这保证了输出始终是一个有效的概率分布，边界条件精确成立，并且不需要辅助损失或相关的超参数。

使用我们的模型进行生成是立即可行的，因为学习到的核已经是目标状态上的一个概率分布，因此采样只需要一次前向传播加上一次类别采样。它不需要迭代去噪、ODE求解器或多步细化。

我们在真实值已知的受控设置中验证该框架。图1 (https://arxiv.org/html/2605.12805#S1.F1.7)展示了学习到的3状态环CTMC核与解析真实值。两个矩阵在视觉上无法区分，最大绝对误差低于 $3 \times 10^{-3}$。在字母表大小高达16、长度高达32的因式分解合成序列任务中，单步生成实现了每个位置总变差距离低于0.03。这些结果证实了离散MeanFlow恒等式在数学上是合理的，参数化具有表达能力，并且单步离散生成在实践中是可行的。

图1：在3状态环CTMC上的精确核恢复。学习到的核 $K_\theta$ 与解析核 $K$ 非常接近，最大绝对误差低于 $3 \times 10^{-3}$，并准确恢复了环的非对称转移结构。

我们的主要贡献是：
1. 1. 我们引入了**离散MeanFlow**，这是一种MeanFlow的有限状态公式，用单纯形上的概率传输替代了连续状态轨迹，并推导了相应的恒等式，将有限区间平均转移率与瞬时CTMC生成元联系起来。
2. 2. 我们引入了一种**边界内建**的核参数化方法，通过设计保证了有效的概率输出和精确的边界条件，无需任何辅助损失。
3. 3. 我们演示了**单步离散生成**，并在精确CTMC和合成序列任务上进行了验证，以高精度恢复了真实值核。

## 2 相关工作

#### 扩散与流匹配。
扩散模型(Sohl-Dickstein et al., 2015 (https://arxiv.org/html/2605.12805#bib.bib24); Ho et al., 2020 (https://arxiv.org/html/2605.12805#bib.bib3); Song and Ermon, 2019 (https://arxiv.org/html/2605.12805#bib.bib25); Song et al., 2021 (https://arxiv.org/html/2605.12805#bib.bib4))和流匹配(Lipman et al., 2023 (https://arxiv.org/html/2605.12805#bib.bib1); Liu, 2022 (https://arxiv.org/html/2605.12805#bib.bib2))主要用于连续空间生成建模，能产生高质量样本，但在生成时需要大量顺序函数评估。

#### 单步生成。
一致性模型(Song et al., 2023 (https://arxiv.org/html/2605.12805#bib.bib17); Song and Dhariwal, 2024 (https://arxiv.org/html/2605.12805#bib.bib18))、捷径模型(Frans et al., 2025 (https://arxiv.org/html/2605.12805#bib.bib19))和归纳矩匹配(Zhou et al., 2025 (https://arxiv.org/html/2605.12805#bib.bib20))通过对网络输出施加自一致性约束，将采样步骤减少到一步或少数几步，但需要精心设计的课程调度或额外的损失函数。MeanFlow(Geng et al., 2025a (https://arxiv.org/html/2605.12805#bib.bib5))推导了平均速度与瞬时速度之间的闭式恒等式，提供了原则性的单步目标函数，无需蒸馏或课程调度。它已扩展到机器人学(Sheng et al., 2026 (https://arxiv.org/html/2605.12805#bib.bib11))、强化学习(Wang et al., 2026 (https://arxiv.org/html/2605.12805#bib.bib14); Chen et al., 2025 (https://arxiv.org/html/2605.12805#bib.bib16))和音频(Li et al., 2025 (https://arxiv.org/html/2605.12805#bib.bib15))领域，并有了进一步的理论分析(Geng et al., 2025b (https://arxiv.org/html/2605.12805#bib.bib13); Zhang et al., 2025 (https://arxiv.org/html/2605.12805#bib.bib12))。所有这些方法都在连续欧几里得空间中运行，不适用于速度场未定义的离散空间。

#### 离散空间中的生成模型
D3PM (Austin et al., 2021 (https://arxiv.org/html/2605.12805#bib.bib6))、SEDD (Lou et al., 2024 (https://arxiv.org/html/2605.12805#bib.bib8))和MDLM (Sahoo et al., 2024 (https://arxiv.org/html/2605.12805#bib.bib7))通过学习有限状态空间上的反向转移，将扩散扩展到离散空间；而离散流匹配(Gat et al., 2024 (https://arxiv.org/html/2605.12805#bib.bib10); Campbell et al., 2024 (https://arxiv.org/html/2605.12805#bib.bib9))则用通过Kolmogorov正向方程训练的CTMC速率函数替代了基于ODE的速率场，使得流匹配能够适应具有离散动作的强化学习(Khan et al., 2026 (https://arxiv.org/html/2605.12805#bib.bib31))。这些方法已应用于语言建模(Sahoo et al., 2024 (https://arxiv.org/html/2605.12805#bib.bib7); Lou et al., 2024 (https://arxiv.org/html/2605.12805#bib.bib8))、蛋白质设计(Campbell et al., 2024 (https://arxiv.org/html/2605.12805#bib.bib9))和图生成(Yiming et al., 2025 (https://arxiv.org/html/2605.12805#bib.bib29))。所有这些方法在生成时都需要许多步骤，不支持单步采样。我们的工作通过直接学习转移核来填补这一空白，使得只需一次前向传播即可生成。

## 3 背景
本节回顾我们工作所推广的连续基础——流匹配和MeanFlow，然后引入连续时间马尔可夫链作为离散对应物。图2 (https://arxiv.org/html/2605.12805#S3.F2.3)突出了关键对比：连续模型沿着平滑的、由速度驱动的轨迹移动点，而离散模型描述概率质量在状态之间的流动。

图2：连续与离散生成动力学。连续流匹配遵循平滑轨迹，而离散动力学通过CTMC转移核 $K_{r,t}(y,x)$ 在有限状态上重新分配概率质量。

### 3.1 流匹配
流匹配(Lipman et al., 2023 (https://arxiv.org/html/2605.12805#bib.bib1); Liu, 2022 (https://arxiv.org/html/2605.12805#bib.bib2))学习一个速度场，将先验分布传输到数据分布。给定数据 $x \sim p_{\mathrm{data}}$ 和噪声 $\epsilon \sim p_{\mathrm{prior}}$，构建流路径为 $z_t = (1-t)x + t\epsilon$，条件速度为 $v_t = \epsilon - x$。由于给定的 $z_t$ 可能来自许多不同的 $(x, \epsilon)$ 对，流匹配训练一个网络 $v_\theta$ 来拟合**边缘**速度场 $v(z_t, t) = \mathbb{E}[v_t \mid z_t]$。通过求解ODE $\frac{d}{dt} z_t = v(z_t, t)$ 从 $z_1 = \epsilon$ 到 $z_0 = x$ 来生成样本。即使条件流是直线的，边际速度场也可能导致弯曲轨迹(Lipman et al., 2023 (https://arxiv.org/html/2605.12805#bib.bib1))。粗糙的离散化会导致样本不准确，因此标准流匹配在生成时需要许多ODE求解步骤。

### 3.2 MeanFlow
MeanFlow (Geng et al., 2025a (https://arxiv.org/html/2605.12805#bib.bib5))用时间区间 $[r, t]$ 上的**平均速度**替代了瞬时速度：
$$u(z_t, r, t) \triangleq \frac{1}{t-r} \int_r^t v(z_\tau, \tau) \, d\tau.$$
(1)
该场描述了两个时间步之间流动的平均方向和大小。与瞬时速度不同，平均速度直接编码位移：$z_r = z_t - (t-r) u(z_t, r, t)$。单次评估 $u(\epsilon, 0, 1)$ 即可将噪声一步映射到数据，无需积分ODE。MeanFlow的关键理论结果是一个将平均速度和瞬时速度联系起来的恒等式。对 $u$ 关于 $t$ 的定义进行微分并应用乘积法则，得到：
$$u(z_t, r, t) = v(z_t, t) - (t-r) \frac{d}{dt} u(z_t, r, t).$$
(2)
这就是MeanFlow恒等式，它表明平均速度等于端点处的瞬时速度减去一个与平均速度自身随时间变化率成正比的修正项。总导数 $\frac{d}{dt} u$ 通过链式法则展开为 $v(z_t, t) \partial_z u + \partial_t u$，涉及一个空间雅可比-向量乘积。训练神经网络 $u_\theta$ 以满足该恒等式，其中瞬时速度是唯一的真实值信号，训练过程中无需任何积分计算。

### 3.3 连续时间马尔可夫链
有限状态空间 $\mathcal{S}$ 上的连续时间马尔可夫链（CTMC）(Norris, 1998 (https://arxiv.org/html/2605.12805#bib.bib30))由时间相关的生成元 $u_t(y,x)$ 指定，该生成元给出了过程从状态 $x$ 跳转到状态 $y$ 的瞬时速率。生成元满足 $u_t(y,x) \ge 0$ 对于 $y \neq x$，且 $\sum_{y \in \mathcal{S}} u_t(y,x) = 0$，因此每一列之和为零，概率守恒。

#### 转移核
CTMC的有限时间行为由**转移核**捕获：
$$K_{r,t}(y,x) := \mathbb{P}(X_t = y \mid X_r = x),$$
(3)
它给出了给定过程在时间 $r$ 处于状态 $x$ 的条件下，在时间 $t$ 处于状态 $y$ 的概率。对于每个起始状态 $x$，列 $K_{r,t}(\cdot, x)$ 是一个有效的概率分布。当没有时间流逝时，过程保持不动：$K_{r,r}(y,x) = \delta_{y,x}$。生成元和核通过Kolmogorov正向方程(Lipman et al., 2024 (https://arxiv.org/html/2605.12805#bib.bib23))联系起来：
$$\partial_t K_{r,t}(y,x) = \sum_{z \in \mathcal{S}} u_t(y,z) K_{r,t}(z,x), \qquad K_{r,r}(y,x) = \delta_{y,x}.$$
(4)
该方程描述了条件概率如何演化：处于状态 $y$ 的概率的变化率是所有中间状态 $z$ 的当前处于 $z$ 的概率乘以从 $z$ 到 $y$ 的瞬时速率之和。因此，生成元 $u_t$ 扮演了离散动力学中速度场的角色，决定了每个时刻概率质量在各状态间的重新分布方式。最近的离散生成模型(Austin et al., 2021 (https://arxiv.org/html/2605.12805#bib.bib6); Sahoo et al., 2024 (https://arxiv.org/html/2605.12805#bib.bib7); Lou et al., 2024 (https://arxiv.org/html/2605.12805#bib.bib8); Gat et al., 2024 (https://arxiv.org/html/2605.12805#bib.bib10); Campbell et al., 2024 (https://arxiv.org/html/2605.12805#bib.bib9))学习这样的生成元，并通过在多个时间步上模拟CTMC来采样。我们的目标是完全避免这种迭代模拟，直接学习转移核 $K_{r,t}$，正如MeanFlow通过直接学习平均速度 $u$ 来避免ODE积分一样。

## 4 离散MeanFlow
将单步生成扩展到离散空间并非