Expanding Flow Maps¶
会议: NeurIPS 2026
arXiv: 2607.21585
代码: https://github.com/sophtang/ExpandingFlowMaps
领域: 学习理论
关键词: 流映射 (Flow Maps), 变维度生成 (Variable-Dimensional Generation), 扩展插值过程 (Expanding Interpolant), 分段确定性马尔可夫过程 (PDMP), 少步采样 (Few-Step Generation)
一句话总结¶
本文提出了扩展生成流(EFlows)与扩展流映射(EFMs),将流映射分解为条件增噪扩维的扩展算子与前向去噪的传输映射,统一了连续与离散空间下的自适应维度少步生成。
研究背景与动机¶
流模型与基于随机插值(Stochastic Interpolants)的生成方法在连续和离散数据生成中取得了显著成功,近期的流映射(Flow Maps)与一致性模型更进一步通过学习轨迹间的跳跃映射,将多步数值积分蒸馏至仅需 1 到 2 步的少步函数评估。然而,现有的流模型与流映射在理论和参数化上均严格建立在“固定画布”(Fixed Canvas)的假设之上,即连续空间必须预设固定维度 \(d\),离散序列必须预设固定长度 \(L\)。
这种刚性维度约束与许多真实世界的生成任务存在根本冲突:在分子构象生成中重原子与氢原子的多尺度层级、离散分子图生成中未知的节点数量、以及自然语言生成中动态变化的句子长度,均要求模型能够在推理过程中自适应地调整输出维度。先前非自回归扩散模型中虽有掩码插入或编辑操作等探索,但均脱离了流映射这一高效的少步轨迹蒸馏范式。如何在数学上形式化跨越不同维度概率测度之间的生成流,并构建出能够单步实现“扩维+去噪”的流映射,成为了该领域的核心瓶颈。
本文的切入点是打破不同维度流流形之间的微分同胚障碍,构建维度单调递增的扩展插值过程,将任意两时间步之间的流映射显式解耦为扩维与传输两个子操作。核心 idea:将流映射分解为引入条件噪声提升维度的扩展算子(Expand Operator)与在扩充空间中沿插值轨迹推进的传输映射(Transport Map),结合局部时间坐标与分段确定性马尔可夫过程(PDMP),实现连续与离散数据上维度可自适应学习的少步流式生成。
方法详解¶
整体框架¶
扩展生成流(EFlow)与扩展流映射(EFM)的核心流水线分为扩维阶段与传输去噪阶段。给定全局时间区间 \(s < t\),源状态 \(\bm{x}_s\) 处于较低维度空间 \(\mathbb{R}^{d(s)}\),目标状态 \(\bm{x}_t\) 处于较高维度空间 \(\mathbb{R}^{d(t)}\)。由于两空间维度不同,不存在直接的微分同胚映射,方法首先利用扩展算子 \(\mathcal{E}_{s,t}\) 注入条件噪声 \(\bm{\epsilon} \sim p_{\epsilon|s,t}\),将状态提升至目标维度空间 \(\mathbb{R}^{d(t)}\) 得到增广状态 \(\bm{x}_s^\epsilon\);随后,传输映射 \(X_{s,t}\)(在连续状态下由平均速度 \(v_{s,t}\) 参数化,在离散状态下由单纯形上的平均去噪器 \(\psi_{s,t}\) 参数化)在增广空间中沿着各坐标独立的局部时间时钟驱动状态向前推进。复合映射 \(\Phi_{s,t} = X_{s,t} \circ \mathcal{E}_{s,t}\) 便实现了从时间 \(s\) 到时间 \(t\) 的单步联合扩维与去噪。
%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
A["低维输入状态 x_s ∈ R^d(s)"] --> B["扩展算子:条件增噪与几何布局<br/>注入条件噪声 ϵ,提升至目标维度 R^d(t)"]
B --> C["局部时间坐标系统<br/>重置插入维度局部时钟,映射至 [0, 1]"]
C --> D["扩展流映射与一致性约束<br/>平均速度/去噪器驱动增广空间传输"]
D --> E["离散单纯形二项式插入机制<br/>Poisson/Binomial 采样与交叉熵蒸馏"]
E --> F["自适应维度的目标生成样本 x_t ∈ R^d(t)"]
关键设计¶
1. 扩展算子:条件增噪与几何布局 针对跨维度状态间缺乏可逆微分同胚的理论断裂,扩展算子 \(\mathcal{E}_{s,t}: \mathbb{R}^{d(s)} \times \mathbb{R}^{d(t)-d(s)} \to \mathbb{R}^{d(t)}\) 将源状态 \(\bm{x}_s\) 与来自可计算条件分布 \(p_{\epsilon|s,t}\) 的增广噪声变量 \(\bm{\epsilon}\) 相结合,把状态提升至高维空间 \(\bm{x}_s^\epsilon := \mathcal{E}_{s,t}(\bm{x}_s, \bm{\epsilon})\)。具体实现涵盖三种布局:拼接(Concatenation,适合自回归或点云)、位置插入(Positional Insertion,适合序列填充)以及子代扩展(Child Expansion,适合分子构象粗到细生成,将新坐标绑定在父节点上并添加扰动 \(\bm{x}_s^\epsilon[i] = \bm{x}_s[\text{pa}(i)] + \sigma \bm{\epsilon}_i\))。数学上证明了增广状态的条件测度即为噪声测度在该算子下的前推测度(Pushforward Measure): $\(p_{t|s}(\cdot \mid \bm{x}_s) = (\mathcal{E}_{s,t}(\bm{x}_s, \cdot))_{\#} p_{\epsilon|s,t}(\cdot \mid \bm{x}_s)\)$ 同时,这一结构满足固定噪声可微性与两步插入的组合律,当扩展算子退化为恒等映射时,完全恢复传统固定画布流。
2. 局部时间坐标系统:各维度独立的去噪时钟 在动态扩维过程中,各维度坐标是在不同全局时间点 \(t_i^{\text{ins}}\) 被引入系统的,若统一使用全局时间推进去噪,后插入的纯噪声维度将被迫在极短的剩余时间内完成去噪,导致分布失真。为此,本文为每个插入的维度 \(i\) 定义了局部时间坐标 \(t_i \in [0, 1]\): $\(t_i = \frac{t - t_i^{\text{ins}}}{1 - t_i^{\text{ins}}}\)$ 未插入的维度速度设为零(\(b_t(\bm{x}_t^\epsilon)[i] = \bm{0}\)),已激活的维度则严格在其归一化的局部时间区间内沿各自速度场去噪。在此基础之上,论文进一步在不相交维度的并集流形 \(\mathcal{M} = \bigsqcup_t \mathbb{R}^{d(t)}\) 上,将扩维与传输严格统一为分段确定性马尔可夫过程(PDMP),其无穷小生成元显式分解为确定性传输项与以扩展算子为跃迁核的跳跃项,从而推导出了扩展流的瞬时变量替换(Change-of-Variables)公式。
3. 扩展流映射与一致性约束:跨维度的少步蒸馏 为了从连续时间扩展流中蒸馏出单步或少步跳跃的确定性流映射 \(\Phi_{s,t} = X_{s,t} \circ \mathcal{E}_{s,t}\),模型通过平均速度 \(v_{s,t}\)(在对角线 \(s=t\) 处对应瞬时速度 \(b_t\))进行参数化。论文严谨导出了扩展流映射在跨维度条件下必须满足的理论一致性恒等式: $\(\partial_t \Phi_{s,t}(\bm{x}) = v_{t,t}(\Phi_{s,t}(\bm{x})) \quad \text{(Lagrangian)}\)$ $\(\partial_s \Phi_{s,t}(\bm{x}) + \nabla_{\bm{x}} \Phi_{s,t}(\bm{x}) \tilde{v}_{s,s}(\bm{x}) = 0 \quad \text{(Eulerian)}\)$ $\(\Phi_{u,t}(\Phi_{s,u}(\bm{x})) = \Phi_{s,t}(\bm{x}) \quad \text{(Semigroup)}\)$ 其中欧拉条件中的雅可比矩阵为矩形矩阵 \(\nabla_{\bm{x}} \Phi_{s,t} \in \mathbb{R}^{d(t) \times d(s)}\),通过与低维未提升速度 \(\tilde{v}_{s,s}\) 相乘完成维度对齐。基于这些恒等式构建的自一致性损失函数(LSD/ESD/PSD),使得模型无需逐步求解微分方程,即可一步跨越时间与维度。
4. 离散单纯形二项式插入机制:针对图与序列的离散扩展 在离散类别序列(如语言或分子图)中,状态被表示为单纯形向量 \(\Delta^{V-1}\)。离散扩展算子需要预测在现有词元之间的槽位(Gap)中插入多少个新词元。模型引入插入预测头 \(\hat{\mathcal{I}}_{s,t}(\bm{x}_s)[i]\) 拟合区间内的插入期望,并采用二项分布对插入长度进行采样: $\(\ell_i \sim \text{Binomial}\left(L - n_s, \frac{\hat{\mathcal{I}}_{s,t}(\bm{x}_s)[i]}{L - n_s}\right)\)$ 当步数较多时,该分布自然逼近泊松过程;而在单步或极限少步采样中(\(\rho_{s,t} \to 1\)),二项分布的方差趋向于零,能够严格保证长度预算不溢出并有效抑制方差爆炸。配合单纯形上的离散半群一致性损失,离散平均去噪器 \(\psi_{s,t}\) 与插入头联合训练,成功将流映射推广至可变长离散序列。
损失函数 / 训练策略¶
离散扩展流映射(DEFM)的总训练目标为类别交叉熵蒸馏损失与区间插入损失的加权和: $\(\mathcal{L}_{\text{total}} = \mathcal{L}_{\text{DEFM}}(\hat{\mathcal{E}}, \hat{\psi}) + \lambda_{\text{insert}} \mathcal{L}_{\text{insert}}(\hat{\mathcal{I}})\)$ 其中插入损失 \(\mathcal{L}_{\text{insert}}\) 采用泊松负对数似然(\(\phi(a, b) = b - a + a \log(a/b)\)),同时监督对角线瞬时缺失计数与非对角跨步插入计数;\(\mathcal{L}_{\text{DEFM}}\) 则由离散半群组合目标与对角线基准去噪目标构成,使用 stop-gradient(\(\text{sg}\))操作符稳定两步一致性梯度的反向传播。
实验关键数据¶
主实验¶
论文在三项代表性任务上进行了全面评测:连续分子构象生成(GEOM-QM9 / GEOM-Drugs)、离散分子图生成(QM9)以及可变长语言建模(LM1B)。
表 1 展示了在 GEOM 连续分子构象生成任务中的表现(重原子骨架先生成后冻结,再扩展生成氢原子坐标,+R 表示使用轻量坐标优化网络额外迭代 10 步):
| 模型 | 采样步数 | QM9 COV-R(%) ↑ | QM9 AMR-R(Å) ↓ | Drugs COV-R(%) ↑ | Drugs AMR-R(Å) ↓ |
|---|---|---|---|---|---|
| GeoDiff | 500 | 87.80 / 93.66 | 0.3179 / 0.3216 | 64.12 / 75.56 | 1.1444 / 1.1296 |
| SubgDiff | 500 | 89.40 / 94.39 | 0.2543 / 0.2601 | 74.30 / 77.87 | 1.0003 / 0.9905 |
| MSGEN | 1000 | 88.36 / 95.10 | 0.2606 / 0.2332 | 85.06 / 97.16 | 0.9538 / 0.9478 |
| EFlow (Ours) | 20 | 88.51 / 93.01 | 0.3849 / 0.3585 | 81.78 / 90.91 | 0.9971 / 0.9826 |
| EFlow+R (Ours) | 30 | 88.87 / 94.05 | 0.3212 / 0.2924 | 85.80 / 93.85 | 0.9069 / 0.8855 |
| EFM (Ours) | 4 | 83.62 / 86.91 | 0.4165 / 0.3954 | 81.26 / 91.13 | 1.0201 / 1.0029 |
| EFM (Ours) | 2 | 72.83 / 74.25 | 0.4703 / 0.4436 | 77.21 / 88.67 | 1.0740 / 1.0499 |
| EFM (Ours) | 1 | 69.21 / 69.23 | 0.4884 / 0.4635 | 63.28 / 72.47 | 1.1921 / 1.1721 |
| EFM+R (Ours) | 14 | 87.79 / 92.01 | 0.3271 / 0.3049 | 87.97 / 95.45 | 0.9029 / 0.8836 |
消融实验¶
表 2 评估了离散分子图生成任务(QM9 数据集,测试 10,000 个分子),对比了未经蒸馏的多步 EFlow 与基准 DeFoG,以及少步蒸馏下的 EFM 与离散流映射基线 CFM:
| 机制类型 | 模型配置 | 步数 | 有效性 Valid(%) ↑ | 独特性 Unique(%) ↑ | FCD 距离 ↓ | 说明 |
|---|---|---|---|---|---|---|
| 多步对比 | DeFoG | 100 | 99.2 | 96.2 | 0.134 | 经典多步图离散流基线 |
| 多步对比 | EFlow (Ours) | 100 | 98.9 | 95.9 | 0.116 | 动态节点扩展,FCD 更优 |
| 多步极低步 | DeFoG | 4 | 53.6 | 63.1 | 3.008 | 步数骤减时基线性能严重崩溃 |
| 多步极低步 | EFlow (Ours) | 4 | 91.7 | 92.8 | 1.780 | 未蒸馏下依然保持良好有效性 |
| 少步流映射 | CFM | 2 | 91.1 | 97.6 | 0.49 | 固定维度范式流映射 |
| 少步流映射 | EFM (Ours) | 2 | 90.0 | 97.9 | 0.40 | 动态扩维流映射,分布保真度更佳 |
| 单步生成 | CFM | 1 | 90.0 | 91.8 | 2.14 | 单步单次跳跃 FCD 急剧退化 |
| 单步生成 | EFM (Ours) | 1 | 90.2 | 97.7 | 0.44 | 单步 FCD 降幅达 4.8 倍,极强鲁棒性 |
关键发现¶
- 大分子图与复杂几何上的扩维优势:在包含多达 181 个原子的大分子 GEOM-Drugs 上,EFlow 和 EFM 凭借重原子到氢原子的结构化扩展,其性能超越了以往依赖 500-1000 步的扩散模型,EFM+R 仅用 14 步即在全部四项指标上刷新 SOTA。
- 单步采样的二项式截断效应:在离散图与序列生成中,二项分布相较于无界泊松分布能严格限制插入数量,保证了单步生成(1-step EFM)在 FCD 评价上维持在 0.44 的极优水平,彻底规避了极端少步下由于维度爆炸导致的模式崩塌。
- 语言生成下的少步困境:在 LM1B 文本生成中,EFM 在 2 步和 4 步下能生成高质量文本(Perplexity 分别为 125.95 和 99.85,显著优于 CFM 的 267+),但在单步(1-step)语言生成中表现出一定的模式退化,说明在大词表(\(V=30,522\))下一次性完成序列全长扩维与去噪仍面临理论与网络容量的严苛挑战。
亮点与洞察¶
- 将流映射拓展至变维度流形:打破了过去流模型、扩散模型和流映射依赖静态全画布维度的硬性先验,将“画布尺寸”本身转化为了一个可学习的概率自由度。
- 解耦操作算子与 PDMP 理论统一:创造性地将跨维映射解耦为“扩展算子 + 增广空间传输映射”,并通过分段确定性马尔可夫过程完成了无穷小生成元层面的严密闭环,为变维度扩散理论奠定了数学基石。
- 可复用的局部时钟设计:对不同阶段插入的特征赋予归一化的局部时间坐标 \(t_i \in [0, 1]\),保证了各模块即使在不同物理时间点进入系统,也能享受一致的去噪动力学支持。
局限与展望¶
- 单步大词表离散生成的模式塌缩:单步端到端生成长文本时,一步跨越完整长度与语义词表难度极大,样本熵较低,仍需进一步优化时间调度机制。
- 先验噪声假设的拓展空间:当前扩维引入的条件噪声主要假定为高斯分布,未来可将噪声均值与协方差网络化,实现数据自适应的噪声结构。
- 收缩与剪裁操作(Shrinking Flows)的对称性缺失:目前理论仅建立了维度递增的单向扩展流,未来需将剪枝/删除维度的逆过程纳入 PDMP,实现真正意义上的任意双向编辑流。
相关工作与启发¶
- vs Flow Maps (Boffi et al., 2025; Roos et al., 2026): 传统连续与离散流映射均局限于固定输入维度或序列长度;本文将流映射范式扩展为变维度流,原固定维度方法即为其扩展算子为恒等映射的特例。
- vs Edit Flows / FlexMDMs (Havasi et al., 2025; Kim et al., 2026): 先前可变长生成依赖多步马尔可夫离散跳跃或掩码扩散,无法直接应用轨迹蒸馏;本文不仅支持细粒度逐步积分,更首次将变长离散生成推向 1-4 步的高效流映射蒸馏。
评分¶
- 新颖性: ⭐⭐⭐⭐⭐ 首次建立了跨维度流生成与流映射的一般理论,PDMP 建模极其优雅严谨。
- 实验充分度: ⭐⭐⭐⭐⭐ 涵盖连续 3D 分子构象、离散分子图、离散自然语言三大模态,全面验证了 multi-step 到 1-step 的泛化能力。
- 写作质量: ⭐⭐⭐⭐⭐ 理论推导清晰完备,从基础定义到 PDMP 证明一气呵成,结构层次分明。
- 价值: ⭐⭐⭐⭐⭐ 解决了生成模型中长期存在的“固定画布”痛点,对自适应分辨率视觉生成、任意长序列生成和具身图结构生成具有重要借鉴意义。