Flow Matching for Count Data¶

会议: NeurIPS 2026
arXiv: 2605.07746
OpenReview: https://openreview.net/forum?id=5N0iGXXzEJ
代码: https://github.com/weigcdsb/count-flow-matching
领域: 图像生成
关键词: 流匹配, 计数数据, 连续时间马尔可夫链, 生死过程, 最优传输

一句话总结¶

针对高维计数数据离散非负、跨度大且难以连续化的难点，本文提出基于局部单步生死过程的流匹配框架 count-FM，通过条件二项式桥实现无模拟边缘转移率匹配，在单细胞测序与神经脉冲生成中以更少参数实现更高质量的分布传输与生成。

研究背景与动机¶

高维计数数据（如单细胞 RNA 测序的基因表达计数、多脑区神经脉冲发放序列）在生物医药和神经科学中极为常见。这类数据天然具有严格的非负整数离散性、高稀疏性以及复杂的非线性多变量相关结构。在发育生物学或认知行为实验中，研究者不仅需要拟合静态边缘分布，更迫切需要建立不同实验批次、发育时间节点或脑状态之间的高维分布传输映射。然而，现有的深度生成模型主要针对连续欧氏空间设计，难以直接外推至离散几何空间。

将扩散与流匹配拓展至离散数据是当前生成模型的前沿。然而现有方案面临两难困境：一类方法采用分类扩散模型（如 D3PM、tauLDR、SEDD），将每一个计数值离散化为一个独立的类别状态，但这需要预先截断最大计数值，在输出端维护维度高达 $\sum_i (C_i + 1)$ 的概率分布，不仅参数量随计数范围暴增，更彻底破坏了计数天然的序关系与度量距离；另一类方法通过去量化（dequantization）强行将计数映射到连续空间进行连续流拟合，但平滑噪声会严重模糊零膨胀稀疏性与离散物理边界，且中间轨迹脱离了整数语义。

面对离散几何保持与计算可扩展性之间的核心矛盾，本文提出直接在非负整数格点空间 $\mathbb{N}_0^d$ 上构建保持序关系的局部单步跳跃动力学。核心 idea：基于连续时间单步生死马尔可夫链构建条件二项式桥，仅用 $2d$ 维输出参数化局部出生与死亡转移率，实现无模拟的边际率匹配训练与可解释的计数空间平滑传输。

方法详解¶

整体框架¶

count-FM 的目标是学习从源计数分布 $p_0(x_0)$ 到目标计数分布 $p_1(x_1)$ 的连续时间传输过程，其中 $x_0, x_1 \in \mathbb{N}_0^d$。整体流水线分为端点耦合构造、条件二项式桥采样、局部生死转移率回归训练以及一阶跳跃离散化采样四个阶段。模型无需预设最大截断值，输出仅包含每个维度的出生率与死亡率（共 $2d$ 个通道），从而在轨迹全程严格保持整数计数值。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["源计数与目标计数<br/>x0 ~ p0, x1 ~ p1"] --> B["端点耦合机制<br/>独立耦合 / 批次最优传输 Γ*"]
    B --> C["条件二项式桥采样<br/>Xt^(i) = x0^(i) + sgn(Δx)*Bt, Bt ~ Binomial"]
    C --> D["局部生死率匹配目标<br/>回归瞬时出生率 λθ 与死亡率 μθ"]
    D --> E["一阶局部跳跃采样器<br/>τ-leaping 步长内每维度至多单步 ±1 跳跃"]
    E --> F["目标计数生成与跨分布传输<br/>X_τ 落在目标计数流形"]

关键设计¶

1. 局部单步生死过程参数化：以 $2d$ 维极简输出保留计数几何序关系

传统离散分类流模型需要为每个特征预测所有可能计数值的分类 logits，当计数跨度达到数百时参数量不可承受且丢失相邻大小关系。count-FM 将传输动力学受限在局部单步邻域跳跃（$\pm e_i$）。在无穷小时间步长 $h>0$ 内，坐标 $i$ 的计数值仅允许发生单步增加（出生）或单步减少（死亡）： $$\mathbb{P}(X_{t+h} = x + e_i \mid X_t = x) = h \lambda_{t,i}(x) + o(h), \quad \mathbb{P}(X_{t+h} = x - e_i \mid X_t = x) = h \mu_{t,i}(x) + o(h)$$ 为了保证计数非负边界条件（当 $x^{(i)}=0$ 时死亡率严格为零），模型将网络输出显式参数化为非负出生率 $\lambda_\theta(x, t)$ 和非负死亡系数 $\beta_\theta(x, t)$，并定义实际死亡率为 $\mu_\theta(x, t) = x \odot \beta_\theta(x, t)$。这一设计让输出维度恒定为 $2d$，与最大计数值完全解耦，既大幅精简参数，又让传输轨迹具备局部平滑演化的物理可解释性。

2. 闭式条件二项式桥：解析匹配边缘柯尔莫哥洛夫向前方程

直接在边缘离散分布路径上训练连续时间马尔可夫链（CTMC）极其困难，因为边缘转移概率难以解析求得。count-FM 给定端点对 $(x_0, x_1) \sim \pi$ 构造了坐标独立的条件二项式概率桥： $$X_t^{(i)} = x_0^{(i)} + \operatorname{sgn}\bigl(x_1^{(i)} - x_0^{(i)}\bigr) B_t^{(i)}, \quad B_t^{(i)} \sim \mathrm{Binomial}\left(\left|x_1^{(i)} - x_0^{(i)}\right|, t\right)$$ 该桥使得每个维度的条件均值随时间 $t \in [0, 1]$ 从 $x_0^{(i)}$ 严格线性演化至 $x_1^{(i)}$。通过代入一维局部马尔可夫过程的质量守恒方程（柯尔莫哥洛夫向前方程），论文推导出维持该二项式边缘分布的精确条件出生率与死亡率： $$\lambda_{t,i}(x \mid x_0, x_1) = \frac{(x_1^{(i)} - x)_+}{1 - t}, \qquad \mu_{t,i}(x \mid x_0, x_1) = \frac{(x - x_1^{(i)})_+}{1 - t}$$ 当时间接近端点 $t=1$ 时公式分母趋于奇异，因此算法设置截止截断时间 $\tau := 1 - \varepsilon_t$（理论证明截断导致的分布全变差误差仅为 $\mathcal{O}(\varepsilon_t)$）。在边缘化后，由于神经网络以全局状态 $x$ 为输入，模型能够捕捉特征间的强相关性，生成非单调的复杂计数动力学。

3. 最优传输端点耦合与分类器免引导：低曲率直线传输与可控条件生成

在端点配对策略上，模型支持独立耦合 $\pi_{\mathrm{ind}}(x_0, x_1)$ 以及基于批次最优传输（OT）的耦合 $\Gamma^\star$。针对计数数据的离散统计特性，count-FM-OT 引入对称泊松发散成本函数 $c(x, y) = \sum_{i=1}^d [x_i \log \frac{x_i+\varepsilon_c}{y_i+\varepsilon_c} + y_i \log \frac{y_i+\varepsilon_c}{x_i+\varepsilon_c}]$，有效拉直了源-目标点对的传输路径，降低采样曲率与所需函数评估次数（NFE）。对于带有协变量 $y$ 的条件生成任务（如神经脉冲与动物位置轨迹的条件建模），论文将无分类器引导（CFG）拓展至转移率空间： $$\lambda_\theta^{\mathrm{cfg}} = \bigl[\lambda_\theta^{\mathrm{uncond}} + w (\lambda_\theta^{\mathrm{cond}} - \lambda_\theta^{\mathrm{uncond}})\bigr]_+, \quad \mu_\theta^{\mathrm{cfg}} = \bigl[\mu_\theta^{\mathrm{uncond}} + w (\mu_\theta^{\mathrm{cond}} - \mu_\theta^{\mathrm{uncond}})\bigr]_+$$ 通过调节引导尺度 $w \ge 0$，在无需训练额外辅助判别器的情况下，灵活权衡边缘校准精度与条件特征对比度。

损失函数 / 训练策略¶

训练目标源于条件计数桥在路径空间（Path-space）上的相对熵（KL 散度）最小化。在无穷小极限下，路径 KL 散度等价于在时间 $t \sim \mathrm{Unif}[0, \tau]$ 和状态 $x \sim p_t(\cdot \mid x_0, x_1)$ 下的泊松似然损失： $$\mathcal{L}_{\mathrm{train}}(\theta) = \mathbb{E}\left[\sum_{i=1}^d \ell(\lambda_{t,i}(x \mid x_0, x_1), \lambda_{\theta,i}(x, t)) + \sum_{i=1}^d \ell(\mu_{t,i}(x \mid x_0, x_1), \mu_{\theta,i}(x, t))\right]$$ 其中对数似然核函数为 $\ell(u, v) = v - u \log(v + \varepsilon_\ell)$。在推理生成阶段，模型采用类似 $\tau$-leaping 的一阶离散化采样器：步长 $\Delta = \tau / K$，在每一时间步根据总跳跃率 $r_i = \lambda_{\theta,i} + \mu_{\theta,i}$ 采样保持、出生 $(+1)$ 或死亡 $(-1)$ 三种互斥结果，在保持离散状态的同时使全局截断误差收敛于 $\mathcal{O}(K^{-1})$。

实验关键数据¶

主实验¶

论文在二维合成伽马-泊松混合分布（2D Simulation）、单细胞 RNA 测序（Dentate Gyrus scRNA-seq，13,913 基因）以及大鼠海马与内嗅皮层多脑区神经脉冲计数（hc-3 dataset）上进行了系统评测。

在二维模拟与单细胞基因表达生成基准上，count-FM 及其最优传输版本展示出卓越的生成保真度与极致的参数效率：

数据集 / 任务	评测模型	参数量 (Params)	$W_1 \downarrow$ / $W_2 \downarrow$	$\mathrm{MMD}^2_{\mathrm{RBF}} \downarrow$
2D 仿真 (Target: 2-mode Gamma-Poisson)	count-FM-OT (本文)	2,372	1.685 ($W_1$) / 2.341 ($W_2$)	0.0005
2D 仿真	count-FM (本文)	2,372	1.866 ($W_1$) / 2.583 ($W_2$)	0.0009
2D 仿真	SEDD (离散扩散基线)	15,306	1.747 ($W_1$) / 3.051 ($W_2$)	0.0002
2D 仿真	discrete-FM (分类流匹配)	14,250	1.942 ($W_1$) / 2.684 ($W_2$)	0.0007
2D 仿真	D3PM (分类扩散模型)	15,306	1.982 ($W_1$) / 2.812 ($W_2$)	0.0010
2D 仿真	Poisson-JUMP (泊松跳跃模型)	67,330	2.974 ($W_1$) / 3.794 ($W_2$)	0.0035
Dentate Gyrus scRNA 无条件生成	count-FM (本文)	9,857,792	20.456 ($W_2$)	0.0185
Dentate Gyrus scRNA 无条件生成	count-FM-OT (本文)	9,857,792	20.553 ($W_2$)	0.0195
Dentate Gyrus scRNA 无条件生成	scDiffusion (连续潜在扩散)	34,704,897	20.545 ($W_2$)	0.0225
Dentate Gyrus scRNA 无条件生成	scVI (深度隐变量基线)	18,155,426	20.687 ($W_2$)	0.0229
Dentate Gyrus scRNA 无条件生成	scLDM (潜在扩散模型)	10,426,955	20.739 ($W_2$)	0.0202
Dentate Gyrus scRNA 无条件生成	CFGen (单细胞生成流模型)	29,165,380	22.226 ($W_2$)	0.0361

在单细胞 P12 到 P35 的发育轨迹传输实验中，通过谱系受限（lineage-restricted）的 OT 耦合，count-FM 生成的单细胞计数状态沿平滑路径渐进演变，精准还原了神经母细胞向幼年与成熟颗粒细胞发育的连续转录组重塑，并在星形胶质细胞等侧向分枝中保持自洽闭环。

消融实验：神经脉冲计数条件生成与引导尺度分析¶

在大鼠 hc-3 线性轨道位置调制神经脉冲数据集上，论文评估了在不同无分类器引导尺度 $w$ 下模型对位置编码发放率均值 ($\mathrm{RMSE}_\mu$)、方差校准 ($\mathrm{RMSE}_{\mathrm{var}}$)、零发放稀疏性 ($\mathrm{RMSE}_0$) 以及神经元群体间协方差结构误差 ($\mathrm{Cov}_{\mathrm{F}}$) 的表现：

| 配置 / 引导强度 | $\mathrm{RMSE}_\mu \downarrow$ | $\mathrm{RMSE}_{\mathrm{var}} \downarrow$ | $\mathrm{RMSE}_0 \downarrow$ | $\mathrm{Cov}_{\mathrm{F}} \downarrow$ | Contrast (位置场锐度) | 说明 | |---|---|---|---|---|---| | count-FM ($w=1$) [完整条件模型] | 0.026 | 0.048 | 0.016 | 0.344 | 3.209 | 最优平衡：均值、稀疏度与群体协方差全面最优 | | count-FM ($w=0$) [无条件模型] | 0.070 | 0.110 | 0.043 | 0.451 | 0.109 | 丢失空间位置条件信号，仅学到边缘平均发放率 | | count-FM ($w=2$) [过度外推引导] | 0.061 | 0.094 | 0.036 | 0.451 | 4.521 | 位置场高度锐化，但分布标定与多样性受损 | | Poisson MLP [泊松回归基线] | 0.027 | 0.115 | 0.027 | 0.564 | 3.134 | 无法建模超泊松方差，大幅低估跨神经元协同相关 | | MLP mean [确定性回归基线] | 0.026 | – | – | – | 3.087 | 仅预测条件均值，缺乏概率生成与噪声采样能力 |

关键发现¶

参数量显著降低且不受计数上限约束：在 2D 模拟中，count-FM 仅使用 2,372 个参数，相较分类离散基线（约 15k 参数）削减超 84%，在 scRNA-seq 上比 scVI 减少 45.70% 的参数、比 scDiffusion 减少 71.6% 的参数，彻底规避了高维计数类别发散的维度灾难。
最优传输（OT）耦合拉直几何轨迹并加速采样：对称泊松代价的批次 OT 显著减少了中间路径的摆动与曲率。中间状态演化曲线证明，分类扩散方法在轨迹前期突变严重，而 count-FM 呈现连续平稳演化，在低采样步数预算下可以更快收敛到高质量生成指标。
无分类器引导的有效折中：在神经脉冲实验中，标准条件生成 $w=1$ 达成最佳分布校准；适度增大引导（$w=2$）可强化位置野的响应反差比（Contrast 提高至 4.521），但过高的引导会压低分布多样性并引起过饱和。

亮点与洞察¶

将流匹配的连续速度场转化为离散转移率匹配：巧妙地将连续时间马尔可夫链的无穷小生成元与流匹配的条件边缘匹配思想结合，建立了 $\mathbb{N}_0^d$ 格点上的无模拟流匹配理论，数学形式优雅且闭式可解。
单步跳跃先验保证了状态全程停留在物理计数空间：避开了传统连续去量化加噪导致的真实物理边界模糊与零截断失真，在单细胞测序发育分析中赋予中间生成状态真实生物学意义。
可广泛迁移至运筹优化与排队系统建模：局部出生-死亡流匹配方法不仅适用于生物组学，更可直接迁移至交通流量计数、工业库存状态轨迹预测、互联网用户点击等大规模整数计数动力学建模中。

局限与展望¶

二项式桥的动力学路径相对刚性：当前条件二项式桥在给定端点时为坐标独立的单峰演化，中间轨迹随机探索能力有限，面对多模态复杂分化命运时主要依赖终端状态的边际平均，未来可探索 Beta-二项式或潜变量随机插值桥。
一阶跳跃采样的离散化误差：目前采用的一阶 $\tau$-leaping 近似在跳跃率极高时必须增加步长以控制 $\mathcal{O}(K^{-1})$ 的截断误差，采样开销仍受端点间绝对距离影响，未来可结合一致性模型（Flow Map Matching）或自适应高阶积分器进行快速单步采样。

评分¶

新颖性: ⭐⭐⭐⭐⭐ 首次提出非负整数格点上的单步生死流匹配理论，条件二项式桥闭式解设计精妙。
实验充分度: ⭐⭐⭐⭐☆ 涵盖二维合成、单细胞测序无条件生成/跨阶段发育传输、以及多脑区神经脉冲条件建模，指标充实。
写作质量: ⭐⭐⭐⭐⭐ 数学推导严谨自洽，对柯尔莫哥洛夫方程与路径空间 KL 散度的联系阐述透彻。
价值: ⭐⭐⭐⭐⭐ 为高维离散计数数据的深度生成与分布传输提供了全新理论基石，在生物信息与神经计算中极具实用价值。