Latent Process Generator Matching¶
会议: NeurIPS 2026
arXiv: 2605.20547
OpenReview: https://openreview.net/forum?id=2ghWjnDjkj
论文: https://arxiv.org/abs/2605.20547
领域: 学习理论
关键词: 生成扩散模型, 流匹配, 生成元匹配, 隐随机过程, 测度前推
一句话总结¶
将生成元匹配(Generator Matching)从静态隐变量拓展到一般时间非齐次 Feller 隐随机过程,证明了前推 Kolmogorov 前向方程与条件损失函数的梯度等价性,统一并拓展了离散与连续状态空间下的辅助动力学消除问题。
研究背景与动机¶
近年来,流匹配(Flow Matching)与连续时间马尔可夫链(CTMC)扩散生成模型广泛引入富含辅助信息的随机动力学来定义条件目标。例如在变长离散序列生成中,Edit Flows、OneFlow 与 Flowception 分别通过增广空标记、交错标记序列与插入速率等离散潜变量来构建联合路径;在连续与流形生成中,Branching Flows 引入离散树状分支指示变量以解决分支归属模糊性,Any-order 掩码扩散同样借助辅助过程追踪 token 插入次序。然而,这些模型在生成采样阶段通常无法、也无需模拟高维或离散的辅助状态,而只需在目标可观测状态空间上生成与联合过程单时刻边缘分布一致的样本。
这一范式面临的核心理论瓶颈在于:现有文献中的状态投影定理与损失函数推导均针对特定模型“一题一议”,缺乏统一的数学基石。更关键的是,已有投影收敛性证明几乎全部严格受限于增广状态空间中的离散分量消除(即要求隐过程取值于有限或离散状态集),无法直接推广到流形约束或高维连续潜动力学;而现有的生成元匹配(Generator Matching)通用框架仅支持对静态随机变量(如起始与终止点对)的条件化,无法囊括具备时变动力学特性的隐随机过程。
本文的切入角度是将可观测生成状态视为可追踪马尔可夫过程在一组确定性映射下的前推测度(Pushforward Measure)。核心 idea:建立前推生成元匹配理论,证明对任意 Polish 空间上的时间非齐次 Feller 过程经过可测映射后的边缘测度满足前推 Kolmogorov 前向方程,且基于样本轨迹的条件生成元匹配损失与边际生成元损失保持严格梯度等价,从而将静态生成元匹配拓展为时变隐随机过程的普适框架。
方法详解¶
整体框架¶
Latent Process Generator Matching 旨在通过训练一个神经网络生成元,在可观测目标空间 \(\mathcal{X}\) 上模拟边缘过程 \(X_t\),使其在任意时间 \(t \in [0, 1]\) 的单时刻边缘分布与富含隐辅助变量的联合过程 \(Y_t\) 经确定性映射 \(\Phi\) 前推后的测度严格一致。
整个方法流程分为三步:首先在增广或隐状态空间 \(\mathcal{Y}\) 上定义一个易于模拟的时间非齐次 Feller 过程 \(Y_t\)(其生成元为 \(W_t\));然后通过可测映射 \(X_t = \Phi(Y_t)\) 将其投射至观测空间 \(\mathcal{X}\),由前推 Kolmogorov 前向方程诱导出观测边际生成元 \(L_t\);最后,借助条件生成元在线性参数化与 Bregman 散度下的期望交换性质,在不显式计算边缘期望的情况下,仅凭 \(Y_t\) 的采样样本优化条件匹配损失。
%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
A["隐空间联合过程 Y_t<br/>生成元 W_t (Feller 过程)"] --> B["前推映射与条件生成元定义<br/>X_t = Φ(Y_t) 映射至观测空间"]
B --> C["条件线性参数化构建<br/>解耦结构算子与可学目标参数"]
C --> D["双重条件化生成元匹配损失<br/>Bregman 散度优化与梯度等价保证"]
D --> E["观测空间边缘生成元 L_t<br/>用于推理采样 (匹配 Φ(Y_1) 目标分布)"]
关键设计¶
1. 前推 Kolmogorov 前向方程与边际生成元诱导:奠定隐过程投影的连续流理论
针对已有工作针对不同任务重复推导测度投影方程的问题,论文建立了测度前推机制。设 \((\mathcal{Y}, d_{\mathcal{Y}})\) 为 Polish 空间,\((Y_t)_{0 \le t \le 1}\) 为具备无穷小生成元 \(W_t\) 的时间非齐次 Feller 过程,对应时间边缘测度为 \(p_t^{\mathcal{Y}}\)。给定可测映射 \(\Phi: \mathcal{Y} \to \mathcal{X}\),定义观测空间的边缘测度 \(p_t^{\mathcal{X}} := \Phi_\# p_t^{\mathcal{Y}}\)。对于任意属于检验函数空间 \(f \in \mathcal{T}(\mathcal{X})\),定义观测边际生成元为作用于复合检验函数的条件期望: $\(L_t f(x_t) := \mathbb{E}[W_t(f \circ \Phi)(Y_t) \mid \Phi(Y_t) = x_t]\)$ 利用积分推导及 Markov 核分解定理(Regular Conditional Distribution),论文证明了边缘分布 \(p_t^{\mathcal{X}}\) 严格满足以 \(L_t\) 为生成元的 Kolmogorov 前向方程(KFE):\(\partial_t \langle p_t^{\mathcal{X}}, f \rangle = \langle p_t^{\mathcal{X}}, L_t f \rangle\)。若 \(\mathcal{Y} = \mathcal{X} \times \mathcal{Z}\) 且 \(\Phi = \pi_{\mathcal{X}}\) 为规范正交投影,即可将现有的离散 CTMC、欧氏或流形投影收敛性均统一为该推论的特例。
2. 条件线性参数化解耦:分离结构算子与可学目标场
为使神经网络能够学习抽象的算子,必须将边际生成元转化为易于参数化的形式。论文构建了条件线性参数化框架:对于每个点 \(x_t = \Phi(y_t)\),将条件生成元 \(L_t^{y_t} f(x_t) := W_t(f \circ \Phi)(y_t)\) 分解为内积空间 \((V_{t, x_t}, \langle \cdot, \cdot \rangle)\) 上的结构映射与目标向量: $\(L_t^{y_t} f(\Phi(y_t)) = \langle \mathcal{K}_{t, \Phi(y_t)} f, F_t^{y_t}(\Phi(y_t)) \rangle_{V_{t, \Phi(y_t)}}\)$ 其中线性映射 \(\mathcal{K}_{t, x_t}\) 承载生成元的先验物理结构(例如扩散过程的微分算子、CTMC 的差分算子),且仅依赖于观测变量 \(x_t\);而目标函数 \(F_t^{y_t}(x_t)\) 则属于闭凸集 \(\Omega_{t, x_t}\) 并包含全部可学习参数。由于内积的线性性质与条件期望可交换,边际生成元的目标参数恰好为条件目标关于隐状态的条件期望: $\(F_t(x_t) = \mathbb{E}[F_t^{Y_t}(x_t) \mid \Phi(Y_t) = x_t]\)$ 这一设计使得对复杂隐过程的消除完全被吸收在目标向量的加权平均中,网络架构无需改变。
3. Bregman 散度下的条件匹配损失与梯度精确等价:消除边缘化的训练闭环
在实际训练中,直接优化边缘目标 \(F_t(x_t)\) 需要对隐状态后验分布积分,计算不可行。论文定义了前推条件生成元匹配损失(Pushforward CGM Loss): $\(L_{\text{cgm}}(\theta) := \mathbb{E}_{t \sim U[0, 1], y_t \sim p_t^{\mathcal{Y}}}[D_{t, \Phi(y_t)}(F_t^{y_t}(\Phi(y_t)), F_t^\theta(\Phi(y_t)))]\)$ 其中 \(D_{t, x_t}\) 为严格凸函数诱导的 Bregman 散度(如均方误差或 KL 散度)。利用 Bregman 散度在凸共轭下的线性期望均值性质,严格证明了梯度等价定理: $\(\nabla_\theta L_{\text{cgm}}(\theta) = \nabla_\theta L_{\text{gm}}(\theta)\)$ 这表明,仅凭从联合过程 \(Y_t\) 采集的单条轨迹样本,直接将神经网络 \(F_t^\theta(x_t)\) 与条件目标 \(F_t^{y_t}\) 计算残差并反向传播,其梯度期望完全等于直接拟合难解边际生成元 \(F_t(x_t)\) 的梯度。
4. 静态端点与时变过程的双重条件化:端到端采样的构造机制
在生成建模中,扩散路径通常依赖于端点先验分布(如噪声与数据样本对 \(z = (x_0, x_1)\))。本文将静态端点条件化与时变隐过程条件化进行融合,推导出具有实用可操作性的双重条件匹配损失: $\(L_{\text{cgm}}(\theta) = \mathbb{E}_{t \sim U[0, 1], z \sim p_{\mathcal{Z}}, y_t \sim p^{\mathcal{Y}}(\cdot \mid z)}[D_{t, \Phi(y_t)}(F_t^{z, y_t}(\Phi(y_t)), F_t^\theta(\Phi(y_t)))]\)$ 该公式明确了高效训练流程:先从静态先验中抽取端点对 \(z\),进而沿条件路径模拟单步隐过程状态 \(y_t\),最后仅以投影值 \(x_t = \Phi(y_t)\) 作为输入计算损失,在生成推理时模型摆脱所有辅助变量,以恢复的纯边际向量场求解常微分方程或 SDE。
损失函数 / 训练策略¶
训练时,时间步 \(t \sim U[0, 1]\),端点对 \((x_0, x_1) \sim q(x_0, x_1)\),在给定端点条件下联合采样状态 \((x_t, z_t) \sim p_t(\cdot \mid x_0, x_1)\)。针对常见的连续流速度预测与端点预测,损失采用均方误差: $\(L_{\text{cgm}}(\theta) = \mathbb{E}_{t, (x_0, x_1), (x_t, z_t)} \left[ \| z_t x_1 - x_1^\theta(x_t) \|^2 \right]\)$ 网络仅接收 \(x_t\) 作为输入预测终点 \(x_1^\theta(x_t)\)。采样时,根据解析关系恢复边际漂移向量 \(u_t^\theta(x) = \frac{x_1^\theta(x) - x}{1 - t}\),配合标准 ODE 求解器(如 Euler 或 RK4)进行去噪生成。
实验关键数据¶
主实验¶
为了在严苛设定下验证理论的正确性,论文设计了一维机制转换扩散过程(Regime-Switching Diffusion):连续状态 \(X_t\) 受到一个离散马尔可夫跳跃过程 \(Z_t \in \{-1, +1\}\) 的调控,条件漂移剧烈反转(在 \(Z_t=+1\) 时漂向 \(x_1\),在 \(Z_t=-1\) 时漂向 \(-x_1\)),该动力学无法由传统静态隐变量生成元匹配覆盖。实验对比了标准 Flow Matching、静态 Generator Matching 以及本文的 Latent Process Generator Matching(LPGM)。
| 模型配置 | 状态空间 | 隐过程支持 | 训练损失梯度等价性 | \(t=1\) 边缘分布匹配度 |
|---|---|---|---|---|
| Flow Matching (Lipman et al., 2023) | 欧氏空间 \(\mathbb{R}^d\) | 无(仅端点插值) | 满足(仅限静态) | 无法拟合切换路径 |
| Generator Matching (Holderrieth et al., 2025) | Polish 空间 | 仅限静态隐变量 \(Z\) | 满足(静态隐变量) | 无法消除动态跳跃 |
| Edit Flows (Havasi et al., 2025) | 离散有限集 \(\mathcal{X} \times \mathcal{Z}\) | 离散 CTMC | 满足(特定离散推导) | 仅限离散序列任务 |
| LPGM (本文) | 任意 Polish 空间 \(\mathcal{Y}\) | 通用 Feller 隐随机过程 | 满足 (Theorem 3.12 严格证明) | 精确恢复平滑边缘目标 |
消融实验¶
论文通过正则性违背反例(Example 3.15)消融分析了 Kolmogorov 前向方程充分性假设(Assumption A.19)对于唯一解的必要性。当光滑映射 \(\Phi(x) = e^{-1/x} \cdot \mathbb{I}_{\{x > 0\}}\) 作用于标准布朗运动时,检验方程在原点的退化现象以及一参数族解构造对收敛性的影响。
| 理论构造 / 验证配置 | 生成元边界条件 \(L_t f(0)\) | 测度在原点测度分布 | KFE 边际解唯一性 | 结论与分析 |
|---|---|---|---|---|
| 满足假设 A.19 (紧致/正则投影) | 满足非退化漂移/扩散 | 连续测度或严格确定性点转移 | 唯一保证 | 隐过程平滑积分,网络稳定收敛至目标分布 |
| 违背假设 A.19 (平坦边界映射) | \(\forall f \in \mathcal{T}, L_t f(0) = 0\) | \(t>0\) 时 \(\mathbb{P}[\Phi(B_t)=0] = \frac{1}{2}\) | 存在单参数族不唯一解 \(q_t^{(c)}\) | 瞬时微观动力学无法唯一确定边缘测度,产生退化 |
关键发现¶
- 在一维动态跳跃实验中,尽管训练样本的条件轨迹呈现强烈的折返与不连续漂移,LPGM 训练出的模型成功将潜变量的跳跃影响边缘化,学到的边际轨迹高度光滑,在 \(t=1\) 处的生成样本直方图与真实目标分布完全重合。
- 理论证明了当隐空间 \(\mathcal{Z}\) 为紧致空间(例如有限离散状态或紧致黎曼流形 \(\mathrm{SO}(3)\))时,检验函数在纤维上的可积性与相容性自动成立,免除了严苛的分析验证。
- 论文指出了投影映射可能诱发生成元在零测集处退化从而导致 KFE 解不唯一的理论边界,为未来流匹配架构的边界设计提供了明确判据。
亮点与洞察¶
- 将生成元匹配从静态变量跨越至时变随机过程:以往的扩散与流匹配理论大多基于确定性插值或静态端点条件化,本文建立了时变马尔可夫隐过程投射的理论完备体系,将静态理论作为其特例收敛。
- Bregman 散度在投影算子下的梯度不变性:无需估计复杂的隐状态后验分布,只需利用单次随机采样的条件残差梯度即可无偏逼近边缘生成元,为非马尔可夫复杂观测系统的建模提供了极低开销的训练范式。
- 拓展非投影式几何前推的前景:论文以多链蛋白质结构生成为例,提出将主链刚体变换群与氨基酸局部柔性分别建模在隐空间 \(\mathcal{Y} = \mathrm{SO}(3) \times \mathbb{R}^3 \times \mathrm{SO}(3)^n \times (\mathbb{R}^3)^n\) 中,通过非平凡群作用映射 \(\Phi\) 推导至真实世界坐标,展示了其在复杂物理与结构生物学中的潜力。
局限与展望¶
- KFE 解的唯一性判据较为抽象:对于高维非线性映射 \(\Phi\),验证 Assumption A.19 是否被满足仍然具有较高的理论难度,目前缺少针对通用神经网络映射的快速经验检验准则。
- 非投影案例尚未进行全面实证验证:论文对蛋白质多链刚体/柔性分离的非投影几何生成(Section 4)仅停留在理论建模与前推映射设计层面,尚未落地为完整的大规模实验与开源代码。
- 未来方向:放宽 Feller 过程的正则性约束,并将该框架正式应用到全原子蛋白质复合物的流匹配生成模型中。
相关工作与启发¶
- vs Generator Matching (Holderrieth et al., 2025): 原始生成元匹配将条件分布严格局限于静态潜变量 \(Z\);本文将条件目标推广为时间非齐次的 Feller 随机过程 \(Y_t\),涵盖了动态过程相互作用。
- vs Edit Flows (Havasi et al., 2025): Edit Flows 针对离散 CTMC 和添加空标记的序列投影推导了生成率投影;本文证明其为一般前推生成元匹配在有限离散状态空间上的直接推论。
- vs Branching Flows (Billera et al., 2025b): Branching Flows 仅解决了连续/流形主干与离散树分支标签的联合投影;本文的前推框架完全放开了隐空间必须为离散的限制,支持高维连续或流形隐变量。
评分¶
- 新颖性: ⭐⭐⭐⭐⭐ 首次建立通用隐随机过程的前推生成元匹配理论,彻底打破了以往仅支持离散潜变量消除的局限。
- 实验充分度: ⭐⭐⭐☆☆ 以理论推导为主,实验部分主要提供了一维机制切换过程的机制验证与反例消融,缺少大规模视觉或文本基准评测。
- 写作质量: ⭐⭐⭐⭐⭐ 测度论阐述严密规范,从 Bogachev 测度论基础到 Markov 核前推定理层层递进,推导清晰自洽。
- 价值: ⭐⭐⭐⭐☆ 为变长离散生成、混合模态流模型及生物大分子分层几何生成提供了强有力的统一理论支撑。