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 轨迹。在解存在且唯一的条件下，轨迹不能在同一时刻、同一位置沿不同方向继续运动。所以模型会在原有条件路径的交汇处重新组织运动方

来源只提供摘要，完整内容请打开原文。