🧬 biology

Flow Matching for Count Data

本文介绍了 count-FM，这是一种基于连续时间生灭过程的无模拟流匹配框架，能够高效生成和传输高维计数数据（如单细胞 RNA 测序和神经脉冲序列），同时在样本质量和建模效率上优于现有基线方法。

原作者： Ganchao Wei, John Pearson

发布于 2026-05-11

📖 1 分钟阅读☕ 轻松阅读

原作者： Ganchao Wei, John Pearson

原始论文根据 CC0 1.0（http://creativecommons.org/publicdomain/zero/1.0/）发布到公有领域。 ⚕️ 这是一篇未经同行评审的预印本的AI生成解释。这不是医疗建议。请勿根据此内容做出健康决定。阅读完整免责声明

想象一下，你正试图将一大群人从一个房间移动到另一个房间。在这群人中，每个人都持有特定数量的苹果。有些人持有零个，有些人持有一个，有些人持有五十个，而且没有人能持有半个苹果。这就是科学家所称的计数数据：以整数形式出现的信息，例如细胞中活跃基因的数量或神经元放电的次数。

本文介绍了一种名为count-FM的新工具，旨在帮助将这些“持有苹果”的人群从一个状态转移到另一个状态（例如，从年轻细胞转移到老年细胞，或从静息大脑转移到活跃大脑）。

以下是其工作原理，使用简单的类比：

1. 问题：“像素”与“桶”

现有处理此类数据的方法通常试图做以下两件事之一，但两者都很笨拙：

“像素”方法：它们将每一个可能的苹果数量（0、1、2、3……直到 100）视为完全独立、互不相关的类别，就像将“红色”、“蓝色”和“绿色”视为完全不同的颜色一样。这使得数学计算极其繁重且缓慢，尤其是当人们可以持有数千个苹果时。
“桶”方法：它们假装苹果实际上是液体（连续的水），以便简化数学计算，然后试图在之后将它们倒回成完整的苹果。但这模糊了界限；你失去了“不能有 4.5 个苹果”这一事实。

count-FM 说：“让我们停止假装。让我们始终将苹果保持为完整的苹果。”

2. 解决方案：“出生与死亡”电梯

count-FM 没有将苹果变成液体，也没有将数字视为无关的颜色，而是使用连续时间出生 - 死亡过程。

想象一个巨大的电梯系统，人们只能一次移动一步，向上或向下。

出生：一个人获得一个苹果。
死亡：一个人失去一个苹果。

该模型学习这些“出生”和“死亡”何时发生的规则。它不试图通过一次巨大的飞跃来猜测最终目的地。相反，它模拟了一段旅程，随着时间的推移，人们一次一个地慢慢获得或失去苹果，直到达到目标人群的分布。

3. 为何高效：“局部”地图

大多数其他模型试图为每个人的每一个可能目的地绘制一张巨大的地图。如果你有 10,000 个基因（变量），而每个基因可以有 100 个值，那么这张地图大得不可思议。

count-FM 就像一个局部 GPS。它只问：“如果我现在有 5 个苹果，我获得一个苹果的概率是多少？我失去一个苹果的概率是多少？”

它忽略遥远的可能性。
它只关注紧邻的下一步（+1 或 -1）。
结果：与其他方法相比，它使用的计算机内存（参数）微乎其微，但在移动人群方面效果一样好，甚至更好。

4. “桥梁”类比

为了训练模型，作者使用了一种称为条件二项式桥梁的方法。
想象你有一个人从持有 10 个苹果开始，你希望他最终持有 20 个。“桥梁”是一条预先规划好的平滑路径，它规定：“在 10% 的进度时，你应该有 11 个苹果。在 50% 时，你应该有 15 个。”
模型学习模仿这条平滑路径。由于这条路径在数学上很简单（就像图表上的一条直线），模型能够非常快速且高效地学习规则。

5. 现实世界测试

作者在两个真实的生物数据集上测试了这个“移动苹果”的系统：

单细胞 RNA 测序：将细胞从早期发育阶段（P12）转移到较晚阶段（P35）。该模型成功展示了细胞如何随时间逐渐改变其基因计数（即它们的“苹果计数”），生成了一部平滑、可解释的发育过程影像，而非跳跃且令人困惑的影像。
神经脉冲序列：根据大鼠的位置预测大脑神经元放电的次数。该模型能够生成逼真的脑活动模式，匹配不同神经元之间复杂的关联性，而更简单的模型未能做到这一点。

总结

count-FM 是一种生成和移动离散计数数据（如基因计数或神经元脉冲）的新方法。它不将这些数字强行纳入液体模具，也不将它们视为无关的类别，而是将它们建模为一系列小的、局部的步骤（一次增加或减少一个单位）。这使得该系统速度更快、对计算机内存的需求更低，并且在保留数据天然的“整数”特性方面更加准确。

技术摘要：计数数据的流匹配（count-FM）

问题陈述

高维计数数据广泛存在于单细胞 RNA 测序（scRNA-seq）和神经脉冲序列分析等领域，为生成建模带来了独特的挑战。这类数据的特点是离散性、稀疏性以及复杂的关联结构。现有的深度生成方法在此领域往往表现不佳：

类别状态方法：将计数视为类别变量（例如 token）的方法，其输出维度需随最大计数范围缩放。这导致参数效率低下，并将相邻计数视为无关状态，忽略了计数空间固有的几何结构。
连续变换：通过去量化或潜在表示将计数映射到连续空间的方法，用连续密度替代了离散概率质量，可能会模糊底层的离散结构。
现有计数模型的局限性：虽然存在一些特定的计数跳跃模型（例如 Poisson-JUMP），但它们通常并非设计用于任意计数分布之间的传输，或者缺乏现代流方法的高效训练和生成特性。

本文旨在解决对生成框架的需求，该框架需直接在计数空间中运行，尊重数据的离散几何结构，并能在不产生类别状态模型参数开销的情况下，实现任意源分布与目标分布之间的高效传输。

方法论：count-FM

作者提出了count-FM，这是一个专为计数数据设计的流匹配框架，基于具有局部单位跳跃的连续时间出生 - 死亡过程。

1. 连续时间出生 - 死亡过程

该模型将源分布 $p_0$ 与目标分布 $p_1$ 之间的传输定义为计数向量空间 $\mathbb{N}_0^d$ 上的连续时间马尔可夫跳跃过程（CTMC）。该过程通过每个坐标 $i$ 的局部单位出生（ $+1$ ）和死亡（$-1$）进行演化：

出生率： $\lambda_{t,i}(x)$
死亡率： $\mu_{t,i}(x)$
动力学：在微小时间间隔 $h$ 内，从状态 $x$ 跃迁至 $x + e_i$ 的概率为 $h \lambda_{t,i}(x)$ ，跃迁至 $x - e_i$ 的概率为 $h \mu_{t,i}(x)$ 。当计数为零时，禁止死亡跃迁。

这种参数化方式效率极高：模型无需预测所有可能计数水平的完整类别分布（其规模随最大计数增长），而只需预测每个维度的两个速率（ $2d$ 个输出），无论计数范围如何。

2. 条件二项式桥接与训练

为了训练模型，作者定义了固定端点 $(x_0, x_1)$ 之间的条件桥接。对于每个坐标 $i$ ，桥接利用二项分布构建：
$X_t^{(i)} = x_0^{(i)} + \text{sgn}(x_1^{(i)} - x_0^{(i)}) B_t^{(i)}, \quad B_t^{(i)} \sim \text{Binomial}(|x_1^{(i)} - x_0^{(i)}|, t)$
该桥接确保条件均值从 $x_0$ 线性移动到 $x_1$ ，同时保持中间状态为整数值。

通过在该条件桥接上强制质量守恒（柯尔莫哥洛夫前向方程），作者推导出了条件出生率和死亡率的闭式表达式：
$\lambda_{t,i}(x | x_0, x_1) = \frac{(x_1^{(i)} - x)^+}{1-t}, \quad \mu_{t,i}(x | x_0, x_1) = \frac{(x - x_1^{(i)})^+}{1-t}$
其中 $(\cdot)^+$ 表示正部函数（正整流）。

3. 训练目标

模型通过最小化条件桥接过程与模型学习过程之间的路径空间 Kullback-Leibler (KL) 散度进行训练。这分解为局部速率匹配目标：
$\mathcal{L}_{train}(\theta) = \mathbb{E} \left[ \sum_{i=1}^d \ell(\lambda_{t,i}, \lambda_{\theta,i}) + \sum_{i=1}^d \ell(\mu_{t,i}, \mu_{\theta,i}) \right]$
其中 $\ell(u, v) = v - u \log v$ 是广义 KL 项。这使得模型能够通过学习可识别的条件速率，在无模拟的情况下训练边际跃迁速率。

4. 采样与耦合

采样：样本通过从 $t=\epsilon_t$ 到 $t=1-\epsilon_t$ 前向模拟学习到的出生 - 死亡过程生成，使用一阶局部跳跃离散化。
端点耦合：本文探讨了两种用于配对源样本和目标样本的耦合策略：
- 独立耦合：随机配对（ $p_0(x_0)p_1(x_1)$ ）。
- 最优传输（OT）耦合：使用对称泊松成本进行小批量 OT 配对。这倾向于诱导曲率更低的跃迁，从而提高可解释性和采样效率。

5. 条件生成

对于条件任务，模型引入协变量 $y$ 并利用无分类器引导（CFG）。在训练期间，条件变量被随机丢弃。在采样时，条件速率与无条件速率通过引导尺度 $w$ 结合，以控制条件的强度。

主要贡献

count-FM 框架：一个基于局部出生 - 死亡动力学的计数数据流匹配框架，支持直接在计数空间中进行传输，且中间状态保持为计数值。
可处理的训练方案：一种通过条件二项式桥接学习边际跃迁速率的方法，在保留流匹配效率的同时，避免了类别状态建模的复杂性。
参数效率：一种仅预测局部出生率和死亡率（ $2d$ 个输出）而非完整类别分布的公式，使其适用于具有大计数范围的高维数据。

实验结果

1. 模拟（2D Gamma-Poisson 混合）

性能：与包括 D3PM、discrete-FM、tauLDR、SEDD 和 Poisson-JUMP 在内的基线相比，count-FM 和 count-FM-OT 实现了最佳的样本质量（最低的 2-Wasserstein 距离和 MMD $^2$ ）。
效率：与类别状态基线（14,000–67,000+ 参数）相比，count-FM 以显著更少的可训练参数（2,372）实现了这些结果。
几何结构：中间状态在计数空间中平滑演化。OT 耦合版本遵循明显更直的传输路径，并以更少的函数评估次数（NFEs）达到目标质量。

2. 单细胞 RNA-seq（齿状回）

无条件生成：在一个包含 2,930 个细胞和 13,913 个基因的数据集上，count-FM 在样本质量（W2 和 MMD $^2$ ）方面优于潜在变量基线（scVI）和潜在扩散模型（scDiffusion, scLDM），同时使用了更少的参数（竞争对手约为 29M–39M，而 count-FM 约为 9.8M）。
传输（P12 到 P35）：该模型成功模拟了从出生后第 12 天到第 35 天的发育传输。使用谱系限制的 OT 耦合，传输路径尊重了已知的发育结构（例如，神经母细胞到颗粒细胞的进展），生成了保持在计数空间内的可解释中间分布。

3. 神经脉冲序列（海马/内嗅皮层）

条件生成：该模型被应用于根据线性化位置和运行方向生成脉冲计数。
比较：引导尺度 $w=1$ 的 count-FM 优于确定性 MLP 回归器和 Poisson MLP 基线。
关联结构：与低估种群相关性的 Poisson MLP 不同，count-FM 更好地保留了跨神经元的依赖结构。
引导权衡： $w=1$ 在空间调谐和分布校准之间提供了最佳平衡。 $w=2$ 锐化了位置场，但扭曲了幅度校准。

意义与主张

本文主张，count-FM 通过直接在计数空间中实现生成和传输，为高维计数数据的建模提供了一个有效框架。其主要意义在于：

几何保真度：通过局部单位跳跃模拟传输，该方法尊重计数数据的离散几何结构，避免了连续近似带来的模糊效应。
可扩展性：参数效率使其适用于高维应用（例如数千个基因），而在这些应用中，类别状态方法在计算上变得不可行。
可解释性：传输路径中的中间状态是有效的计数样本，允许在发育生物学等应用中检查有意义的过渡分布。
通用性：该框架支持无条件生成、任意分布之间的传输以及带有可控引导的条件生成。

作者承认了局限性，指出条件二项式桥接相对简单，可能无法探索所有合理的中间路径，且当前的一阶离散化可以通过更精确的模拟方案加以改进。然而，跨模拟和生物医学数据集的实证结果表明，与竞争方法相比，count-FM 以更高的建模效率实现了高质量的样本生成。