ReSFlow: Rectifying Categorical Flows on Statistical Manifolds for One-Step Generation¶

会议: NeurIPS 2026
arXiv: 2610.13225
代码: https://github.com/xiaohuanjia/ReSFlow
领域: 其他
关键词: 离散生成模型、统计流形、流匹配、Reflow、单步生成

一句话总结¶

针对离散分类数据流匹配在概率单纯形上由于 Fisher-Rao 测地线相交弯曲而依赖多步高昂 ODE 积分的难题，本文提出 ReSFlow，通过平方根微分同胚映射到正正交超球面，并在超球面上显式执行轨迹拉直（Reflow）与测地线重耦合，首次实现了高质量的单步离散分类数据生成。

研究背景与动机¶

离散数据（如自然语言词元、分子图结构、启动子 DNA 序列以及二值化图像）的连续流生成模型正受到学术界的广泛关注。传统的连续流匹配（Flow Matching）与连续时间扩散模型在图像生成与连续空间控制中取得了突破，但将其迁移到离散状态空间时会遭遇离散不可微与拓扑失配的本质障碍。近期，统计流匹配（Statistical Flow Matching, SFM）等方法通过 Fisher-Rao 信息度量将离散分类分布视作概率单纯形这一统计流形上的连续点，成功架起了连续概率流与离散采样之间的桥梁。

然而，现有统计流模型在实际应用中受制于巨大的推理计算瓶颈。由于先验分布与目标离散分布的独立耦合（Independent Coupling）在单纯形上诱导出了相互交叉的局部测地线，边际化后合成的聚合速度场高度弯曲，严重背离任何单一测地线轨迹。为了精确模拟弯曲的 ODE 动力学轨迹，推理阶段必须调用自适应或多步数值求解器进行数十至上百步积分，推理延迟极高。尽管欧氏空间中的 Rectified Flow（Reflow）技术可以通过拉直轨迹实现单步或极少步极速采样，但由于概率单纯形配备的 Fisher-Rao 度量具有非欧内在曲率，简单在欧氏单纯形上做线性插值会造成严重的几何畸变与越界；更关键的是，非线性平方根变换意味着欧氏空间中互不相交的直线在投影回统计流形后依然会发生测地线交叉。

面对流形曲率与单步极速生成之间的根本冲突，亟需一种能够完全遵循统计流形内蕴几何规律的流拉直理论体系。本文的核心 idea 是：利用平方根微分同胚将概率单纯形等距嵌入到正正交超球面，将流拉直操作直接置于超球面上展开，通过超球面球面线性插值（Slerp）与 ODE 轨迹重耦合将传输路径渐进拉直为不相交的大圆测地线，最终利用黎曼指数映射实现高保真度的一步极速离散生成。

方法详解¶

整体框架¶

ReSFlow 的核心流程涵盖三个紧密相扣的阶段：几何嵌入与初始化流匹配、流形重耦合拉直（Spherical Reflow），以及基于切空间投影的单步指数映射推理。整个生成过程直接作用于原始数据空间，无需依赖外部预训练的分词器或连续词嵌入。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["先验分布 p0 与目标离散数据 q<br/>概率单纯形 Δ"] --> B["超球面等距嵌入与初始流匹配<br/>平方根同胚 π: μ ↦ √μ 映射到正正交超球面 𝕊+"]
    B --> C["轨迹诱导重耦合 Spherical Reflow<br/>求解当前 ODE 获得端点对 (Z0, Z1) 消除交叉"]
    C --> D["切空间投影速度场训练<br/>约束预测向量位于 T_Xt 𝕊 切平面并匹配 Slerp 导数"]
    D --> E["单步指数映射推理<br/>X̂1 = Exp_X0(v(X0, 0)) 经逆变换 π 取 argmax"]

在输入阶段，先验分类分布 \(\mu_0 \sim p_0\) 与目标离散 one-hot 分布 \(\mu_1 \sim q\) 通过平方根映射 \(X_i = \sqrt{\mu_i}\) 嵌入到单位超球面的正正交象限 \(\mathbb{S}_{+}\) 上；在该空间中，单纯形上的 Fisher-Rao 测地线严格等距转换为超球面上的大圆弧（Great-circle Arc）。针对独立耦合诱导的速度场曲折问题，ReSFlow 迭代运行 ODE 求解生成端点对 \((Z_0, Z_1)\) 进行重耦合训练，不断降低传输动能与方差；在测试部署时，模型仅需在 \(t=0\) 计算一步初始速度向量，并沿超球面指数映射即可直接一步输出目标状态。

关键设计¶

1. 超球面等距嵌入与切空间投影：消除边界奇异性并保证流形约束

当直接在概率单纯形 \(\Delta\) 上利用 Fisher-Rao 内积 \(\langle u, v \rangle_\mu = \sum_i \frac{u_i v_i}{\mu_i}\) 计算度量时，分母包含各分量概率 \(\mu_i\)。对于目标离散数据（如分类任务中的 one-hot 极值点），分母趋向于零，导致度量张量在单纯形边界处发散崩溃，发生严重的数值数值奇异。ReSFlow 引入平方根微分同胚 \(\pi(\mu) = (\sqrt{\mu_1}, \dots, \sqrt{\mu_K}) \in \mathbb{S}_{+}\)，使得单纯形边界平滑映射为单位球面的正正交象限边界。在超球面 \(\mathbb{S}_{+}\) 上，切空间诱导的黎曼内积退化为标准欧几里得点积 \(\sum_i \tilde{u}_i \tilde{v}_i\)，彻底规避了坐标分母除零问题。

为了确保参数化神经网络预测的速度场 \(v_\theta(X_t, t)\) 在动力学积分演化过程中不脱离超球面流形，ReSFlow 显式施加切空间正交投影算子，剔除平行于球面法向量的法向分量：

\[v_\theta(X_t, t) \leftarrow v_\theta(X_t, t) - \langle v_\theta(X_t, t), X_t \rangle X_t\]

这一正交投影操作使速度场处处切于流形切空间 \(\mathcal{T}_{X_t}\mathbb{S}\)，从而保证由其驱动的任何连续时间演化严格封闭在超球面上，逆变换回单纯形后自然满足非负性与归一化概率和为 1 的物理约束。

2. 球面重耦合拉直（Spherical Reflow）：消除测地线交叉与降低输运损耗

在初次训练阶段（\(k=0\)），先验噪声 \(X_0\) 与数据 \(X_1\) 为独立随机配对，所构建的球面线性插值（Slerp）测地线族在球面上形成大量交织与相交。在任意相交点处，边缘速度场 \(v_\theta(X_t, t)\) 实际上逼近的是多条相交测地线切向速度的加权条件期望 \(\mathbb{E}[u_t \mid X_t]\)，这一平均向量往往脱离了原测地线切向，导致积分所得的 ODE 实际轨迹严重弯折。

为了消除相交，ReSFlow 构建了超球面上的轨迹诱导重耦合迭代机制。基于当前训练收敛的速度场 \(v_{\hat{\theta}}\)，求解确定性 ODE 将噪声样本 \(Z_0 \sim \pi(p_0)\) 前向推演至 \(t=1\) 得到对齐的端点 \(Z_1\)。由于 ODE 解的存在唯一性与光滑传输映射的单射性质，空间相邻的源点被平滑映射为空间相邻的目标点，诱导测地线几乎呈平行走向、不再彼此穿梭相交。随后模型基于新配对 \((Z_0, Z_1)\) 最小化球面流匹配损失：

\[\mathcal{L}_{\text{ReSFlow}} = \mathbb{E}_{t, Z_0, Z_1} \left[ \| v_\theta(Z_t, t) - u_t(Z_t \mid Z_0, Z_1) \|_2 \right]\]

其中 \(Z_t = \text{slerp}(Z_0, Z_1; t) = \frac{\sin((1-t)\theta)}{\sin\theta}Z_0 + \frac{\sin(t\theta)}{\sin\theta}Z_1\)，目标切速度为 \(u_t = \frac{\theta}{\sin\theta}[-\cos((1-t)\theta)Z_0 + \cos(t\theta)Z_1]\)，\(\theta = \arccos(Z_0^\top Z_1)\)。理论证明，每次 Reflow 迭代后期望球面传输成本 \(\mathbb{E}[d_{\mathbb{S}}(Z_0, Z_1)]\) 单调不增，且多余动能以 \(\mathcal{O}(1/N_{\text{reflow}})\) 的速率快速衰减收敛，使全局速度场彻底蜕变为匀速测地线场。

3. 封闭解指数映射：实现高精度的单步极速离散生成

传统连续流模型依赖自适应高阶积分器（如 Dopri5）逐步追踪复杂的弯曲向量场，耗费大量函数评估次数（NFE）。而在 ReSFlow 经过 Spherical Reflow 收敛后，积分曲线已严格逼近常速测地线，其沿途速度向量的大小与方向由初始切速度完全决定。

因此，在采样推理阶段，模型彻底摒弃了迭代数值积分，直接从先验噪声 \(X_0 \sim \pi(p_0)\) 出发，查询模型在 \(t=0\) 处预测的切向量 \(v_0 = v_\theta(X_0, 0)\)，并通过超球面黎曼指数映射（Riemannian Exponential Map）的闭式解一步跳跃至终点：

\[\hat{X}_1 = \mathrm{Exp}_{X_0}(v_0) = X_0 \cos(\|v_0\|_2) + \frac{v_0}{\|v_0\|_2} \sin(\|v_0\|_2)\]

最后，通过平方逆变换 \(\hat{p} = \pi^{-1}(\hat{X}_1) = (\hat{X}_{1,1}^2, \dots, \hat{X}_{1,K}^2) \in \Delta\) 恢复概率单纯形坐标，并在离散空间中取 argmax 得到最终生成符号。整个推理仅需 1 次网络前向传播（NFE=1），彻底打破了离散生成模型在多步采样的计算枷锁。

损失函数 / 训练策略¶

在每一轮 Reflow 迭代周期中，先由当前速度模型通过 ODE 数值积分构造端点耦合数据集，随后利用 AdamW 优化器进行迷你批次蒙特卡洛优化。每个批次中均匀抽取时间步 \(t \sim \mathcal{U}[0, 1]\)，采样插值点 \(Z_t\) 并计算带有切空间投影的均方误差损失：

\[\mathcal{L} = \frac{1}{B}\sum_{j=1}^B \left\| v_\theta(z_{t_j}, t_j) - \langle v_\theta(z_{t_j}, t_j), z_{t_j} \rangle z_{t_j} - u_{t_j} \right\|_2\]

实验表明，通常仅需 1 次或极少数轮次 Reflow 迭代即可完成对流形测地线的高保真拟合，后续即可直接冻结模型投入单步采样。

实验关键数据¶

主实验¶

论文在低维流形合成测试（Swiss Roll on 2-simplex）、高维离散图像生成（Binarized MNIST）、字符级语言建模（Text8）以及生物学长程启动子 DNA 序列设计（Promoter Sequence Design）四个代表性任务上展开对比。对比基线包括单纯形欧氏线性流（LinearFM）、统计流匹配（SFM）、欧氏拉直流匹配（R-LFM）以及基于预训练分词嵌入的连续拉直离散流模型（ReDi）。

表 1 汇总了 Binarized MNIST（Bernoulli 离散分布图像）与启动子 DNA 序列设计（4 元离散序列，长度 1024）的定量指标对比：

任务 / 数据集	方法	ODE 采样质量 (FID ↓ / SP-MSE ↓)	ODE 耗时 / 1k 样本 (s) ↓	单步生成质量 (NFE=1) (FID ↓ / SP-MSE ↓)	单步耗时 / 1k 样本 (s) ↓
Binarized MNIST	LinearFM	5.187 ± 0.129	87.13 ± 0.44	25.317 ± 0.517	0.406 ± 0.178
	SFM	4.748 ± 0.092	112.81 ± 8.20	35.640 ± 0.487	0.719 ± 0.175
	R-LFM	5.681 ± 0.252	61.19 ± 0.28	6.261 ± 0.379	0.583 ± 0.244
	ReSFlow (本文)	5.541 ± 0.252	36.07 ± 0.60	5.973 ± 0.442	0.452 ± 0.243
Promoter DNA	LinearFM	0.027734 ± 0.000897	119.04 ± 16.57	0.029449 ± 0.001673	0.410 ± 0.004
	SFM	0.028015 ± 0.000257	82.18 ± 4.26	0.035564 ± 0.007768	0.410 ± 0.005
	R-LFM	0.052146 ± 0.000781	67.88 ± 4.17	0.060504 ± 0.000351	0.409 ± 0.001
	ReSFlow (本文)	0.028743 ± 0.000755	67.94 ± 0.35	0.028718 ± 0.000366	0.409 ± 0.003

在 Text8 字符级文本建模基准上（以 ODE 似然比特每字符 BPC 评估，越低越好）： - LinearFM: \(2.15448 \pm 0.00024\) - SFM: \(1.82998 \pm 0.00028\) - R-LFM: \(3.26687 \pm 0.00012\)（欧氏拉直破坏了离散流形几何，导致似然严重劣化） - ReDi: \(2.30562 \pm 0.01010\)（依赖连续预训练 token embedding） - ReSFlow (本文): \(1.26146 \pm 0.00036\)（显著超越所有基线，展现出在长程高熵离散符号建模上的优势）

消融实验¶

通过几何建模（欧氏 vs. 球面流形）与轨迹优化（无拉直 vs. Reflow）的 \(2 \times 2\) 交叉消融设计，量化各组件贡献：

实验配置	几何度量空间	轨迹拉直机制	Binarized MNIST 单步 FID ↓	Promoter 单步 SP-MSE ↓	说明
LinearFM	欧氏单纯形 \(\Delta\)	否 (独立耦合)	25.317	0.029449	缺乏内蕴几何，且未拉直轨迹，单步完全溃散
SFM	超球面 \(\mathbb{S}_{+}\)	否 (独立耦合)	35.640	0.035564	具有内蕴测地线但轨迹严重交叉，无法单步采样
R-LFM	欧氏单纯形 \(\Delta\)	是 (欧氏 Reflow)	6.261	0.060504	欧氏拉直后投影回流形仍相交，复杂序列严重失真
ReSFlow (完整模型)	超球面 \(\mathbb{S}_{+}\)	是 (球面 Reflow)	5.973	0.028718	正确流形 + 球面拉直，单步生成指标全面领先

关键发现¶

几何建模与流拉直缺一不可（互补性）：未拉直的 SFM 虽然在完整 ODE 下表现优异，但在 NFE=1 时 FID 暴跌至 35.64；而欧氏拉直的 R-LFM 在复杂长序列建模中（如 Text8 和 DNA）会导致几何误差累积并发生剧烈灾难性退化（Text8 BPC 退化至 3.266，DNA SP-MSE 增至 0.0605）。ReSFlow 证实了必须在内蕴流形空间内执行拉直操作。
推理步数极致压缩（250x 加速）：在 Binarized MNIST 任务中，SFM 达到相当质量（FID 5.922）需要 34 次函数评估，耗时超过 100 秒；ReSFlow 单步评估（NFE=1）耗时仅 0.452 秒，直接实现两百余倍的端到端生成加速，且质量几无损耗。
端到端无需外部 Tokenizer：相较于 ReDi 必须强依赖外部预训练模型（如 GPT-2 分词器）才能构造连续嵌入，ReSFlow 直接在离散状态原始概率单纯形上建立流场，在缺乏分词器的生物序列（DNA/蛋白质）与位图任务中具备显著通用性与轻量优势。

亮点与洞察¶

突破离散空间单步生成的固有壁垒：将经典的“直而快”（Straight and Fast）原理从欧几里得平坦空间优雅地拓展至非欧统计流形，在保持 Fisher-Rao 几何不变性的同时赋予模型单步黎曼指数映射生成能力。
超球面平方根同胚的理论闭式优势：借助平方根微分同胚 \(\pi(\mu)=\sqrt{\mu}\)，将通常难以计算的大圆积分和测地线向量转化为极简的三角函数 Slerp 插值，既彻底消除了边界处的除零数值发散，又使得指数映射具有直接的初等解析表达式。
广阔的跨模态拓展潜力：这一球面拉直机制天然具备通用性，未来可无缝迁移至具备方向统计特性、单位超球面表示的连续模态任务中，例如单位范数视觉嵌入、蛋白质骨架二面角流形以及三维旋转 \(SO(3)\) 动力学生成。

局限与展望¶

似然评估仍需完整 ODE 积分：连续时间变化率（Change-of-variables formula）在数学上严格依赖连续轨迹对流场散度的沿线积分。尽管单步指数映射可以完成极速数据采样，但若要精确计算测试集负对数似然（NLL）或文本 BPC，依然需要调用自适应高阶积分器。
大词表状态下的计算开销扩展：当应用于具有数万至数十万词表的超大型自然语言模型时，每个 token 位置的概率单纯形维度 \(K\) 极其庞大。在庞大维度下直接维护全维度超球面切向量预测，可能引发显存与计算占用的线性扩张，亟需结合局部切空间降维或低秩分解技术。

评分¶

新颖性: ⭐⭐⭐⭐⭐ 首次提出在正正交超球统计流形上执行 Reflow 流拉直，理论自洽且推导严谨。
实验充分度: ⭐⭐⭐⭐⭐ 涵盖一维几何合成分布、二值视觉图像、长程生物 DNA 序列及字符级文本建模，多角度消融极为扎实。
写作质量: ⭐⭐⭐⭐⭐ 动机阐述自然流畅，数学推导与实验对比条理清晰，图文对照直观。
价值: ⭐⭐⭐⭐⭐ 为离散生成模型摆脱百步扩散积分束缚、走向实时低延迟部署提供了极具前景的通用基础范式。