NeurIPS 2026｜ReSFlow 在统计流形上整流，实现离散数据一步生成

本文作者分别来自内蒙古大学、上海财经大学和卡内基梅隆大学 (CMU)。第一作者贾啸寰是内蒙古大学博士生，通讯作者为内蒙古大学范少华研究员。上海财经大学的徐韧喆助理教授、王晓副教授以及 CMU 的博士生 Jiayun Wu 为论文的共同作者。

用流匹配（Flow Matching）或扩散模型（Diffusion Models）生成一张手写数字图像、一段 DNA 序列，通常需要从噪声出发，经过多步采样才能得到结果。每一步都要调用模型，生成质量逐渐改善，计算开销也随之增加。

能不能只调用一次模型，就完成整个生成过程？

对于字符、二值像素和 DNA 碱基等离散数据，难点还在于：类别概率所在的空间具有自己的几何结构，生成路径该如何在这个空间里 “走直”？

研究团队提出了 Rectified Statistical Flow（ReSFlow）：

将类别概率映射到球面，在球面上重新配对并整流生成轨迹，让模型能够沿接近测地线的路径，通过一次更新完成生成。在二值化 MNIST 上，

ReSFlow 一次模型调用就获得了接近 Statistical Flow Matching（SFM）34 次调用的生成质量。

在条件启动子 DNA 设计任务中，一步采样与完整 ODE 采样的预测活性谱误差接近，生成 1,000 条序列的时间则从

67.94 秒缩短至 0.409 秒。

目前，该工作已经被 NeurIPS 2026 接收。

论文标题： ReSFlow: Rectifying Categorical Flows on Statistical Manifolds for One-Step Generation

论文地址： https://arxiv.org/abs/2610.13225

代码地址： https://github.com/xiaohuanjia/ReSFlow

研究背景

离散数据生成通常从类别概率入手：例如，DNA 中的 A、C、G、T 可以表示为四个非负、和为 1 的概率。模型学习这些概率如何随时间变化，再通过 ODE 推进，最终采样得到离散结果。SFM 引入 Fisher–Rao 几何，用概率空间中的测地线构造训练路径，但要实现一步生成，仍有两个难点。

第一，训练路径是测地线，模型推理时实际走出的路径却未必是。

独立随机配对的条件路径可能在同一时刻交叉，给出不同方向。模型需要综合这些方向，学到的轨迹便可能弯曲，只靠起点处的一次速度预测，很难准确到达终点。

第二，欧氏空间中的「直」，不等于概率空间中的「直」。

Rectified Flow 通过重新配对拉直生成轨迹，但欧氏直线不同于 Fisher–Rao 测地线。两条欧氏插值路径即使互不相交，换成球面测地线后仍可能交叉。

ReSFlow 要解决的正是这两个问题：将类别概率映射到球面，在符合 Fisher–Rao 几何的空间中重新配对和训练，让实际生成轨迹更接近测地线，最终实现一步采样。

图 1：欧氏整流与球面整流的区别。欧氏空间中的线性路径即使不相交，对应的球面测地线仍可能交叉；ReSFlow 则直接在球面上整流。

核心方法

要让生成路径真正「走直」，得先选对空间。ReSFlow 先通过平方根映射，把样本点带到球面上，再让初始流模型从噪声出发，找到各自的终点。用这些新配对重新训练，模型逐渐学会沿测地线前进，最后借助球面指数映射，实现从噪声到数据的一步生成。

球面表示与测地线插值

ReSFlow 对类别概率向量

逐项取平方根，将其从概率单纯形映射到单位超球面的非负区域：

平方根映射使 Fisher–Rao 测地线对应到球面上的大圆弧，两种度量之间仅相差固定常数。之后 ReSFlow 使用球面线性插值（slerp）连接给定球面上的起点分布

和终点分布

：

对这条球面线性插值求导，即得到在球面上训练流模型所需的条件速度

。模型的预测速度被映射到当前位置的切空间中，使模型预测速度时刻遵循黎曼球面几何。

球面整流：从随机配对到模型配对

Stage 1： 使模型学会从噪声走向数据。

这一阶段和传统 flow 模型的训练一样分别从先验分布和数据分布中独立随机采样，映射到球面后构造测地线，让模型预测的切向速度

匹配条件速度：

这一阶段建立了从噪声到数据的生成流，但独立配对对应的条件路径存在大量交叉，模型学到的整体轨迹仍存在弯曲。

Stage 2： 使模型在交叉减少的轨迹上继续学习。

完成初始训练后，ReSFlow 从先验分布中采样起点，再利用当前模型求解球面 ODE，得到该起点实际到达的生成终点：

随后，将起点与生成终点作为新的配对，沿两端之间的球面测地线重新训练。由于这些配对来自模型自身的 ODE 轨迹。在解存在且唯一的条件下，轨迹不能在同一时刻、同一位置沿不同方向继续运动。所以模型会在原有条件路径的交汇处重新组织运动方向。相比独立随机配对，这种配对方式显著减少条件路径之间的交叉，使模型学到的整体轨迹更接近球面的测地线。

一步生成

经过前面的配对调整与训练，模型学到的轨迹逐渐接近测地线。在进行 ODE 时无需沿途反复更新方向，只需在初始噪声处的切空间中预测一次速度，再通过球面指数映射，沿对应的大圆弧直接估计生成终点：

得到球面终点后，将坐标逐项平方，恢复类别概率：

再根据恢复出的概率分布选择或采样离散类别。

实验结果

作者先用 Swiss roll 合成分布，看看模型能否一步还原目标分布，验证模型有效性。再将 ReSFlow 用于图像、文本和 DNA 序列任务等真实任务。除了比较生成结果，论文增加了分析实验，验证整流是否让模型生成的轨迹更加贴合球面测地线。对照基线包括 LinearFM、SFM 和 R-LFM，分别比较欧氏与球面几何，以及整流前后的差别。

验证实验

在 Swiss roll 合成分布上，ReSFlow 生成 1,000 个样本的 ODE 采样耗时为 0.0361 秒，是四种方法中最短的；SFM 则需要 0.0679 秒。ReSFlow 的 NLL 为 −2.25，优于 LinearFM 和 R-LFM，但仍略高于 SFM 的 −2.38，NLL 越低越好。这组结果说明，球面整流缩短了采样耗时，同时保留了较好的概率建模表现。

表 1：Swiss roll 实验结果。比较完整 ODE 采样下的 NLL，以及生成 1,000 个样本的耗时。

一步采样生成的样本更能直观的反映这个结论，真实分布在单纯形上呈现出卷曲的带状结构。未经整流的 LinearFM 和 SFM 无法还原这一形状，生成点集中在较小的区域；而 ReSFlow 保留了完整的卷曲结构，分布也更接近真实样本。

图 2：Swiss roll 真实分布与一步生成样本。依次为真实分布、LinearFM、SFM、R-LFM 和 ReSFlow。

主实验

在二值化 MNIST 数据集上，未经整流的 LinearFM 和 SFM 的一步 FID 分别为 25.317 和 35.640；经过整流后，R-LFM 降至 6.261，ReSFlow 进一步达到 5.973。

完整 ODE 采样时，SFM 的 FID 最低；但压缩到一步后，它的质量损失明显更大。ReSFlow 则将 FID 从完整 ODE 下的 5.541 保持在一步采样的 5.973。

论文还比较了不同模型调用次数：

ReSFlow 一次调用的生成质量，接近 SFM 调用 34 次的结果。

此时，SFM 的 FID 为 5.922，ReSFlow 为 5.973。

表 2：二值化 MNIST 实验结果。比较完整 ODE 与一步采样的 FID 和耗时，时间对应生成 1,000 个样本。

样本图把这组差别展示得更直观。完整 ODE 采样时，四种方法都能生成较清楚的手写数字；减少到一步后，LinearFM 和 SFM 的笔画出现明显缺失，部分数字难以辨认。ReSFlow 则保留了更完整的笔画和轮廓，一步生成的样本与自身完整 ODE 的结果也更接近。

图 3：二值化 MNIST 生成样本。上排为完整 ODE 采样，下排为一步采样；各列依次为 LinearFM、SFM、R-LFM 和 ReSFlow。

在文本生成任务上，ReSFlow 的 BPC 为 1.26146，低于 SFM 的 1.82998。欧氏整流在这里却没能帮上忙：R-LFM 的 BPC 升至 3.26687，比未经整流的 LinearFM 更差。这说明了重新配对并不总能带来改善，路径采用什么几何同样关键。

表 3：Text8 字符级建模结果。BPC 通过完整 ODE 的似然评估计算，报告三个随机种子的均值与标准差。

条件启动子设计任务将生成对象换成了长度为 1,024 的 DNA 序列。ReSFlow 一步采样的 SP-MSE 为 0.028718，完整 ODE 采样为 0.028743，两者的预测活性谱误差接近。

生成 1,000 条序列的平均耗时则

从 67.94 秒缩短至 0.409 秒，约加速 166 倍。

作为对照，SFM 压缩到一步后，SP-MSE 从 0.028015 上升至 0.035564，质量损失更明显。

表 4：条件启动子设计结果。时间对应生成 1,000 条长度为 1,024 的序列，报告三个随机种子的均值与标准差。

图 4：条件启动子设计任务的生成样本。从上到下依次为 LinearFM、SFM、R-LFM 和 ReSFlow。每组上行为完整 ODE 采样，下行为一步采样（Euler），每行展示一条长度为 1,024 个碱基的启动子序列的前 120 个碱基。

分析实验

生成结果变好了，模型生成的轨迹是否也确实更「直」？作者在二值化 MNIST 上记录了整流过程中的传输成本和测地线偏差。

传输成本衡量了噪声起点与生成终点之间的距离代价。ReSFlow 的传输成本随整流明显下降，说明重新配对后，模型倾向于把噪声样本送到几何距离更近的生成终点，减少不必要的远距离搬运。相比之下，R-LFM 的下降幅度较小，也更早进入平台。这样的结果表明，在欧氏空间中改善配对，并不一定能同样有效地降低概率空间中的传输成本。

测地线偏差则检查模型沿途有没有 “偏离方向、改变步速”：在轨迹上的每个位置，将模型实际预测的速度，与沿测地线、在剩余时间内到达同一生成终点所需的速度进行比较。两者越接近，说明模型实际走出的轨迹越接近恒速测地线。随整流次数增多，偏差逐步下降，从实验上支持了 ReSFlow 的核心思路：

球面整流能够让模型实际生成的路径更接近测地线。

路径少了弯折，速度变化也更符合恒速测地线的规律，起点处的一次预测便更有可能描述整段运动。这也解释了为什么 ReSFlow 可以通过一次指数映射，较准确地估计原本需要多步积分才能到达的终点。

图 5：不同整流轮次下的 Fisher–Rao 经验传输成本和测地线偏差。数值越低，表示相应成本或偏差越小。

结语

ReSFlow 为离散数据的一步生成提供了一条路径：从类别概率的 Fisher–Rao 几何出发，在球面上重新配对和整流，让生成轨迹逐渐接近测地线，再通过指数映射完成一步采样。图像与 DNA 序列实验表明，这种设计能够大幅减少采样计算，同时保留较好的生成表现。

一步推理仍需要额外的整流训练，其在更大规模数据和复杂条件任务上的效果也有待探索。但这项工作提供了一个明确的启发：

要让生成路径「走直」，首先需要按照数据所在空间的几何，定义什么是「直」。

© THE END

转载请联系本公众号获得授权

投稿或寻求报道：liyazhou@jiqizhixin.com