CellFlow：Flow matching建模cell状态变化

生物系统的复杂性以及潜在扰动的组合空间使得全面的实验表征难以实现。计算模型可以通过从现有扰动数据中学习并外推到未表征的实验条件，帮助系统地探索表型空间。当前单细胞基因组学的扰动预测方法大多集中于预测有限范围内的干预措施在培养系统中有限的分化细胞类型。然而，多细胞组织由异质细胞群体组成，且表型筛选通常需要复杂的实验设置。因此，对单细胞表型的全面预测需要模型能够捕捉异质细胞分布，并整合多种组合干预措施的信息，例如受扰动的基因、施用的药物及其剂量和时间，以及细胞状态的描述符。

CellFlow是一个灵活的框架，用于建模由各种内部或外部信号诱导的单细胞表型。CellFlow借助流匹配（flow matching）这一生成建模技术，该技术已被用于计算机视觉、视频以及分子设计领域生成高质量样本，部分原因在于它能对复杂分布进行建模。CellFlow利用最优传输对未受扰动和受扰动的细胞进行配对，从而能够区分细胞固有的异质性和扰动引起的分布变化。为了表征任意类型的扰动并实现对分布外条件的预测，CellFlow整合了强大的生物实体预训练嵌入。为了以置换不变的方式对任意数量的扰动进行建模，CellFlow采用了包括多头注意力在内的聚合策略，这也是LLM取得成功的关键因素。

在包含近1000万个外周血单核细胞（PBMC）的大型扰动数据集上，CellFlow能准确模拟donor特异的细胞因子处理应答。为了展示CellFlow对高度复杂细胞分布的建模能力，CellFlow预测了在不同发育阶段经不同基因敲除扰动的整个斑马鱼胚胎的单细胞表达谱。在已确立的扰动预测任务中，包括基因敲除和药物处理效应预测，模型表现出最先进的性能。此外，CellFlow展示了预测神经元命运工程和类器官发育过程中产生的异质细胞群的能力。作为原理验证，作者进行了虚拟类器官方案筛选，识别出了此前未经过测试但对类器官发育有显著影响的处理方案。

图1. CellFlow探索细胞表型空间。（A）细胞可响应各种刺激改变其表型，例如药物处理、基因敲除或发育信号。（B）表型筛选有助于探索扰动与表型之间的关系。CellFlow能够从这类筛选中学习，并生成未观测条件下的表型，为各种下游任务提供便利。

来自：CellFlow enables generative single-cell phenotype modelingwithflowmatching
https://github.com/theislab/CellFlow

CellFlow建模扰动细胞表型

CellFlow是一个用于建模受扰动单细胞表型的通用框架。该框架旨在通过将源分布（如对照细胞）有条件地映射到受扰动细胞群体，来预测不同扰动下的单细胞表型。模型首先对实验变量进行编码，并将组合处理方式聚合为一个通用的条件嵌入，随后将该嵌入输入到Flow matching模块中，以指导从源分布到受扰动分布的流动（图1C）。

为了适应多种表型筛选场景，作者定义了一个包含三种变量的通用实验设置：扰动、扰动协变量和样本协变量。"扰动"指可观测的实验干预，例如药物处理或基因修饰。为了能够对训练过程中未见过的干预措施进行预测，作者采用了编码方式，如小分子的分子指纹，以及基因敲除产物和蛋白质的ESM2嵌入。"扰动协变量"指治疗的其他调节因素，如剂量或时间；而样本协变量则指与处理无关的细胞描述符，如细胞系或组织。作者将扰动、扰动协变量和样本协变量的组合称为单一条件（single condition）。

为了同时对多种干预措施进行建模，需要以置换不变的方式对处理组合进行编码。因此，CellFlow实现了两种集合聚合方案，即多头注意力（multi-head attention）和深度集合（deep sets），之后是一个由前馈网络组成的编码器。这使得CellFlow能够将任意的扰动变量组合表示为单一的表征形式。

---

“置换不变”（permutation-invariant）是指在处理一组元素（如多个扰动、特征或数据点）时，模型的输出不受这些元素排列顺序（即 “置换”）的影响。简单来说，无论元素的顺序如何变化，经过处理后得到的结果保持一致。

---

条件嵌入向量随后作为输入传入Flow matching模块，该模块将源分布转换为受扰动的细胞群体。为了在群体和单细胞层面学习有意义的映射，作者利用最优传输对批次样本进行配对，对模型进行端到端训练，从而学习到一个条件表征空间，该空间指导从源群体生成受扰动的细胞状态。

当信噪比高时，流的训练会更顺利。由于单细胞数据可能存在特征检测稀疏和技术变异的问题，CellFlow通过主成分分析（PCA）或变分自编码器（VAEs）获得的低维细胞表征。

• 图1（C）CellFlow会考虑通用实验设置中的各种变量，并采用注意力机制和DeepSets等聚合策略获取聚合的条件嵌入。使用流匹配（将预测流回归到预定义的概率路径上）来学习条件向量场，从而将源分布转换为扰动分布。

CellFlow

给定一组对照细胞和一项实验干预，CellFlow旨在生成一组受扰动的细胞。作者将对照（源）细胞和受扰动（目标）细胞的表型均表示为多维嵌入空间中的向量，该嵌入空间通过主成分分析（PCA）或变分自编码器（VAEs）构建。为了生成预测结果，CellFlow会考虑定义实验干预的所有变量，并将这些变量嵌入到一个单一的表征向量中，该向量随后会指导对照细胞向受扰动细胞的转化。

这种转换由神经常微分方程（neural ODE）控制，该方程由一个依赖时间的神经向量场参数化。神经向量场是一种前馈神经网络，其输入包括编码后的实验条件、需要评估向量场的时间点，以及这些时间点上经过插值的细胞表型向量。在这里，时间是一种抽象的表示方式，用于描述从初始状态（t=0）到受扰动状态（t=1）的细胞状态转变。该网络输出定义这些状态之间流动的速度向量，而受扰动细胞的预测表型是通过使用数值常微分方程求解器，将向量场从t=0（此时细胞表型为给定源细胞的表型）积分到t=1得到的。

作者使用Flow matching和最优传输来训练神经向量场。流匹配是一种无需模拟即可训练神经向量场的方法，能够实现快速且稳定的优化。最优传输通过最小化总位移成本来对齐对照细胞和受扰动细胞，并且已被证明能够加速和改进流匹配的训练。

在每次训练迭代中，会随机采样一批源细胞、一个实验条件（如药物或基因敲除）以及与该条件对应的一批目标细胞。在从[0,1]区间均匀采样的时间点上对神经向量场进行评估。随后，根据用于拉直连接对照细胞和受扰动细胞的神经向量场路径的最优传输解决方案，对采样批次中的源细胞和目标细胞进行配对。之后，对实验条件编码器和神经向量场进行端到端优化，从而隐式学习该条件的函数表征。

一种统一扰动实验设置的表述方式

为涵盖广泛的扰动实验，作者制定了一种通用的表型筛选设置（图1c）。作者将定义扰动实验的变量集划分为三类，即扰动{pi}i∈Ip\left\{p_{i} \right\}_{i \in I_{p}}{pi​}i∈Ip​​、扰动协变量{qi}i∈Iq\left\{q_{i} \right\}_{i \in I_{q}}{qi​}i∈Iq​​和样本协变量{si}i∈Is\left\{s_{i} \right\}_{i \in I_{s}}{si​}i∈Is​​，其中$I_p$、$I_q$和$I_s$表示索引集。扰动包括所有干预措施，如药物处理、基因敲除和细胞因子处理，而扰动协变量则补充了关于扰动的额外信息，例如药物剂量或干预时间。扰动（以及相应的扰动协变量）可以组合使用，例如可以联合使用多种药物。相比之下，作者将样本协变量定义为捕捉与扰动无关的细胞状态，例如细胞系、组织、供体或批次。样本协变量不可组合，例如一个细胞不能同时属于多个批次。

因此，条件$c$是一个定义干预和细胞状态的三元组，即$c=({\left\{p_{c_j}\right\}}_{j\in n_c},{\left\{q_{c_j}\right\}}_{j\in n_c},s_c)$。其中，nc={1,…,n}n_{c}=\{1, \dots, n\}nc​={1,…,n}，$n\in \mathbb{N}$表示条件$c$的扰动（以及相应的扰动协变量）数量。此外，条件$c\in C$唯一对应一对源分布和目标分布$({\mu }_c,{\nu }_c)$。

从所有扰动因素中学习单一表征

为了学习上述定义扰动实验的因素的表征，所有因素都必须转换为机器可读取的形式，随后以置换不变的方式进行聚合。作者采用ESM2嵌入来表征细胞因子和遗传扰动，同时利用分子指纹来表征药物。如果所应用的处理没有现成的表征，可以对其进行one-hot编码，这种方式只能预测未见过的处理组合，而无法预测未见过的单个扰动。

由于扰动协变量通常是标量值（即处理时间、处理剂量），作者将这些值转换为近线性尺度和神经网络常用范围（例如，对剂量进行对数转换）。不过，用户也可以提供任何形式的表征来编码扰动协变量。

与扰动类似，对于所有可能的样本协变量，也没有统一的表征方式。通常利用癌症细胞系百科全书的嵌入来表征癌症细胞系；对于患者的表征，利用供体特异性对照群体的统计数据；在跨数据集学习时，则采用独热编码对其进行编码。

作者将由扰动、扰动协变量和样本协变量定义的实例称为一个“条件”，因此，目标是学习条件的表征，以指导从对照群体到受扰动群体的流转换。

CellFlow使用由前馈神经网络或自注意力层组成的模块来处理扰动、扰动协变量和样本协变量的表征。为了获得所有细胞协变量的单一表征，构成单一条件的所有处理会通过置换不变池化操作进行聚合。为此，CellFlow通过逐特征均值的方式，采用置换不变注意力机制或深度集合。随后，得到的条件向量会通过前馈神经网络或自注意力模块进一步处理，以获得该条件的最终向量表征。

用于细胞配对的最优传输

利用最优传输（OT）对源分布和目标分布中的样本进行配对，已被证明能改善流匹配模型的生成性能，因为它能拉直源群体与目标群体之间的路径。鉴于细胞分布的复杂性和异质性，将对照细胞映射到分布相近的受扰动细胞，是提升CellFlow生成能力的一个有前景的方向。

就目前而言，目标是从给定样本中找到源细胞与目标细胞之间的对应关系，而不涉及对未见过的细胞的泛化，也不引入不同的条件。

Monge Maps

蒙日映射（Monge Maps）是一种确定性函数，能够以最具成本效益的方式将源分布映射到目标分布。形式上，设$\mu \in M_1^{+}({\mathbb{R}}^d)$为对照细胞的源分布，$\nu \in M_1^{+}({\mathbb{R}}^d)$为受扰动细胞的目标分布，其中$M_1^{+}({\mathbb{R}}^d)$表示${\mathbb{R}}^d$上的概率分布集合。我们的目标是学习这两个分布之间的传输映射，即找到一个映射$T:{\mathbb{R}}^d\to {\mathbb{R}}^d$，使得$T♯\mu =\nu$。在这种情况下，$T$模拟了扰动的效应：对于一个服从$\mu$分布的对照细胞$x$，其受扰动后的版本由$\hat{y}=T(x)$给出，且$\hat{y}$服从$\nu$分布。为此，采用最优传输（Optimal Transport，OT）形式化框架。

假设我们有一个连续成本函数$c:{\mathbb{R}}^d\times {\mathbb{R}}^d\to \mathbb{R}$，用于衡量将一个服从$\mu$分布的对照细胞$x$传输到一个服从$\nu$分布的受扰动细胞$y$的成本。在许多应用中，该成本采用数据样本之间的欧氏距离平方，即$c(x,y)=\Vert x-y{\Vert }^2$。本研究也采用这一方式。最优传输的蒙日公式旨在从所有能将$\mu$推前到$\nu$的映射$T$中，找到使期望传输成本最小的映射：
$$\underset{T\#\mu =\nu }{\inf }{\int }_{{\mathbb{R}}^d}\Vert x-T(x){\Vert }^{2}d\mu (x).(公式MP)$$

公式（MP）的任何解都被称为$\mu$和$\nu$之间的蒙日映射。然而，这一问题在实际中往往难以求解，因为其约束集可能具有高度的非凸性，甚至在离散情况下可能为空集。因此，通常会考虑采用松弛形式的表述。

Kantorovich Formulation

康托罗维奇公式（Kantorovich Formulation）与聚焦于传输映射$T$不同，康托罗维奇提出的公式考虑了耦合$\pi \in \Pi (\mu ,\nu )$。这些耦合是${\mathbb{R}}^d\times {\mathbb{R}}^d$上的概率测度，其边缘分布分别为$\mu$和$\nu$。形式上，
$$\Pi (\mu ,\nu ):=\left\{\pi \in {\mathcal{M}}_1^{+}\left({\mathbb{R}}^d\times {\mathbb{R}}^d\right)|p_1\#\pi =\mu ,p_2\#\pi =\nu \right\},$$
其中$p_1:(x,y)\mapsto x$和$p_2:(x,y)\mapsto y$是标准投影映射。从传输映射$T$转向耦合$\pi$，是用一种随机映射替代了确定性映射——在确定性映射中，每个对照细胞$x$被映射到唯一的受扰动细胞$y=T(x)$；而在随机映射中，$x$被映射到一个条件分布$\pi (\cdot |x)$。

根据公式（MP），我们寻求一个耦合$\pi$，使得在$\pi$下的平均位移成本最小：
$$\underset{\pi \in \Pi (\mu ,\nu )}{\min }{\iint }_{{\mathbb{R}}^d\times {\mathbb{R}}^d}\Vert x-y{\Vert }_2^{2}d\pi (x,y).(公式KP)$$

解决公式（KP）的最优耦合${\pi }_{OT}$始终存在。当蒙日映射存在时，蒙日公式与康托罗维奇公式是一致的，此时有${\pi }_{OT}=(Id,T^{\ast })♯\mu$。换句话说，${\pi }_{OT}$是一种确定性耦合，其支撑集为$T$的graph。

因此，若$(x,y)\sim {\pi }_{OT}$，则$y=T(x)$。

非平衡扩展

康托罗维奇公式要求质量守恒。例如，这意味着每个未受扰动的细胞都必须被映射。因此，像细胞凋亡或非同步增殖这样的因素无法被建模。非平衡最优传输（Unbalanced Optimal Transport, UOT）通过$\varphi$散度$D_{\varphi }$惩罚$p_1\#\pi$与$\mu$的偏差以及$p_2\#\pi$与$\nu$的偏差，从而解除了这一约束。作者采用最常用的选择$D_{\varphi }=\text{KL}$（即KL散度），其针对正测度$\rho ,\gamma \in {\mathcal{M}}^{+}({\mathbb{R}}^d)$的定义为：
$$\text{KL}(\rho |\gamma )={\int }_{{\mathbb{R}}^d}\log \left(\frac{d\rho }{d\gamma }\right)d\rho -{\int }_{{\mathbb{R}}^d}d\gamma +{\int }_{{\mathbb{R}}^d}d\rho ,$$

其中$\frac{d\rho }{d\gamma }$是$\rho$相对于$\gamma$的相对密度。

引入${\lambda }_1$、${\lambda }_2>0$以平衡对质量变化的惩罚与传输成本之间的关系，康托罗维奇问题的非平衡版本旨在寻找一个正测度$\pi \in {\mathcal{M}}^{+}({\mathbb{R}}^d\times {\mathbb{R}}^d)$，使得：
$$\underset{\pi \in {\mathcal{M}}^{+}\left({\mathbb{R}}^d\times {\mathbb{R}}^d\right)}{\min }{\int }_{{\mathbb{R}}^d\times {\mathbb{R}}^d}\Vert x-y{\Vert }^{2}d\pi (x,y)+{\lambda }_1\text{KL}\left(p_1\#\pi |\mu \right)+{\lambda }_2\text{KL}\left(p_2\#\pi |\nu \right).\text{（公式UKP）}$$

熵正则化最优传输

康托罗维奇问题（公式KP）的目标函数可通过添加一个熵正则化项进行平滑处理。对于$\epsilon >0$，熵正则化的康托罗维奇问题表述为：
$$\underset{\pi \in \Pi (\mu ,\nu )}{\min }{\iint }_{{\mathbb{R}}^d\times {\mathbb{R}}^d}\Vert x-y{\Vert }^{2}d\pi (x,y)+\epsilon \text{KL}(\pi |\mu \otimes \nu )\text{（EKP）}$$

（注：$\mu \otimes \nu$表示分布$\mu$与$\nu$的张量积，即二者的联合分布）

熵正则化通过避免所谓的最优传输维度灾难，改善了统计特性。此外，当$\mu$和$\nu$为离散分布时，式（EKP）可通过Sinkhorn算法高效求解。这一点尤为重要，因为在实际中，我们只能获取有限的离散细胞样本，即${\hat{\mu }}_n=\frac{1}{n}{\sum }_{i=1}^n{\delta }_{x_i}$和${\hat{\nu }}_n=\frac{1}{n}{\sum }_{i=1}^n{\delta }_{y_i}$，其中$x_1,\dots ,x_n$和$y_1,\dots ,y_n$分别是来自$\mu$和$\nu$的独立同分布样本。

在该离散场景下，熵正则化的康托罗维奇问题（式（EKP））可转化为一个熵正则化的线性规划问题。通过定义成本矩阵$C$，并令
$$U_n=\left\{P\in {\mathbb{R}}_{+}^{n\times n}|P{1}_n=\frac{1}{n}{1}_n,P^{\top }{1}_n=\frac{1}{n}{1}_n\right\},$$
我们得到
$$mi{n}_{P\in U_n}<P,C>+\epsilon KL\left(P|\frac{1}{n^2}{1}_{n\times n}\right).(1)$$

该问题存在一个对偶形式，表现为无约束的强凹最大化问题：
$$\underset{f\in {\mathbb{R}}^n,g\in {\mathbb{R}}^n}{\max }⟨f+g,\frac{1}{n}{1}_n⟩-\epsilon ⟨e^{f/\epsilon },K{e}^{g/\epsilon }⟩,$$

其中，$K$定义为 $K_{i,j}=e^{-C_{i,j}/\epsilon }$（$C$ 为成本矩阵），$e^{f/\epsilon }$ 和 $e^{g/\epsilon }$ 表示逐元素指数运算。由于目标函数是 $\epsilon$-强凹的，因此存在唯一的最优解 $f^{\ast }$、$g^{\ast }$。

算法1（Sinkhorn算法）通过迭代逼近$(f^{\ast },g^{\ast })$。对于矩阵$A=[A_{i,j}{]}_{1\le i,j\le n}$，我们定义按行的$\epsilon$-软极小算子：
$$\underset{\epsilon }{\min }(A):={\left[-\epsilon \log \left(1^{\top }{e}^{-A_{i,\cdot }/\epsilon }\right)\right]}_{1\le i\le n},$$

并令$\oplus$表示张量和，即$f\oplus g=[f_i+g_j{]}_{1\le i,j\le n}$。一旦得到$f^{\ast }$、$g^{\ast }$，通过 primal-dual 关系
$$P_{OT,\epsilon }=\exp \left(\left(C-f^{\ast }\oplus g^{\ast }\right)/\epsilon \right)$$

可恢复相应的最优耦合。由于每次迭代主要涉及矩阵-向量乘法，其时间复杂度为$O(n^2)$；在最一般的情况下，存储矩阵$C$需要$O(n^2)$的内存。然而，当成本具有可分离性时（例如欧氏距离平方成本），内存复杂度可降低至线性水平。

类似的方法也适用于非平衡问题（公式（UKP）），只需添加一个熵惩罚项。形式上，有：
$$\underset{\pi \in {\mathcal{M}}^{+}\left({\mathbb{R}}^d\times {\mathbb{R}}^d\right)}{\min }{\int }_{{\mathbb{R}}^d\times {\mathbb{R}}^d}\Vert x-y{\Vert }^{2}d\pi (x,y)+{\lambda }_1\text{KL}\left(p_1\#\pi |\mu \right)+{\lambda }_2\text{KL}\left(p_2\#\pi |\nu \right)+\epsilon \text{KL}(\pi |\mu \otimes \nu ).$$

在离散情况下，使用与前文类似的符号，我们得到：
$$\underset{P\ge 0}{\min }⟨P,C⟩+\epsilon \text{KL}\left(P|\frac{1}{n^2}{1}_{n\times n}\right)+{\lambda }_1\text{KL}\left(P{1}_n|\frac{1}{n}{1}_n\right)+{\lambda }_2\text{KL}\left(P^{\top }{1}_n|\frac{1}{n}{1}_n\right).(3)$$

（注：$⟨P,C⟩$表示矩阵$P$与$C$的内积；$1_{n\times n}$为$n\times n$维全1矩阵，$1_n$为$n$维全1向量）

与平衡情形下的式（3）类似，该问题存在一个对偶形式，表现为无约束的强凹最大化问题。

$$\underset{f\in {\mathbb{R}}^n,g\in {\mathbb{R}}^n}{\max }-{\lambda }_1⟨\frac{1}{n}{1}_n,e^{-f/{\lambda }_1}⟩-{\lambda }_2⟨\frac{1}{n}{1}_n,e^{-g/{\lambda }_2}⟩-\epsilon ⟨e^{f/\epsilon },K{e}^{g/\epsilon }⟩(4)$$
并且它具有唯一解$(f^{\star },g^{\star })$。

Sinkhorn算法的一个变体只需稍作修改就能解决这个问题。具体来说，我们定义
$${\tau }_1=\frac{{\lambda }_1}{{\lambda }_1+\epsilon }\text{和}{\tau }_2=\frac{{\lambda }_2}{{\lambda }_2+\epsilon }$$

并分别用${\tau }_1$和${\tau }_2$对$f$和$g$进行缩放。直观而言，${\tau }_1$和${\tau }_2$控制着矩阵$P$的行和与列和需要多么接近$\frac{1}{n}{1}_n$。当${\lambda }_1,{\lambda }_2\to \infty$时，我们得到${\tau }_1={\tau }_2=1$，这与平衡情形下的Sinkhorn算法相匹配。算法1着重展示了这一方案，其主要通过矩阵-向量乘法交替进行，时间复杂度为$O(n^2)$。收敛后，原始-对偶关系恢复最优传输计划。与平衡情形不同，行和与列和无需精确匹配$\frac{1}{n}{1}_n$，而是在由${\lambda }_1$、${\lambda }_2$驱动的权衡中保持接近。在实际应用中，直接调整${\tau }_1$和${\tau }_2$往往更为便捷，因为${\tau }_1$、${\tau }_2\in [0,1]$。因此，我们用$\tau =({\tau }_1,{\tau }_2)$重新参数化符号，并相应地记为${\pi }_{\text{UOT},\epsilon ,\tau }$和$P_{\text{UOT},\epsilon ,\tau }$。

作者利用ott-jax计算最优传输耦合，因此，对于特别大的批量大小，通过低秩近似等改进方法，以及为降低内存复杂度而采用的在线成本矩阵评估方法，都能很容易地运行CellFlow。

---

在CellFlow中，$P$是传输计划矩阵，$C$是成本矩阵。

• $P_{i,j}$表示从源分布中第$i$个细胞传输到目标分布中第$j$个细胞的概率，$P$描述了对照组到扰动处理组之间的匹配关系。
• $C_{i,j}$表示从源分布中第$i$个细胞传输到目标分布中第$j$个细胞的成本，$C$描述了对照组到扰动处理组之间的差异度，通常$C_{i,j}=||x_i-y_j|{|}^2$。

CellFlow最小化总传输成本 “$⟨P,C⟩+正则项$” 来学习最优传输计划$P$，另外，最优传输的 $P$ 和 $C$ 为 CellFlow 的流匹配（Flow Matching）模块提供了 “细胞状态转换” 的约束，帮助模型学习从对照组到扰动组的合理表型变化路径，最终生成符合生物学规律的扰动后细胞分布。

---

用于生成受扰动细胞群体的流匹配

在本节中，我们介绍流匹配，并展示如何将其与离散最优传输（OT）技术相结合，以逼近细胞分布之间的最优传输映射。从学习源细胞与目标细胞之间的对应关系（OT），转而学习描述从一个细胞群体到另一个细胞群体转变的向量场（Flow matching）。

连续归一化流（Continuous Normalizing Flows, CNFs）

连续归一化流将源分布$\mu \in M_1^{+}({\mathbb{R}}^d)$转换为目标分布$\nu \in M_1^{+}({\mathbb{R}}^d)$。转换到我们的场景中，源分布对应未受扰动的细胞群体，而目标分布对应受扰动的细胞群体。连续归一化流并非直接将映射$T$参数化为神经网络，而是利用由速度场$v_t:{\mathbb{R}}^d\to {\mathbb{R}}^d$诱导的流${\varphi }_1:{\mathbb{R}}^d\to {\mathbb{R}}^d$这一概念，即${\varphi }_t:{\mathbb{R}}^d\to {\mathbb{R}}^d$满足：
$$\frac{d}{dt}{\varphi }_t(z)=v_t\left({\varphi }_t(z)\right),{\varphi }_0(z)=z$$

连续归一化流（CNFs）利用神经网络$v_{t,\theta }$对向量场进行建模，从而得到一个深度参数化的流模型，该模型通过训练以满足终端条件${\varphi }_{1,\theta }♯\mu =\nu$。

流匹配

虽然训练连续归一化流（CNFs）的标准方法通常需要求解神经常微分方程（ODE），而这种方法计算成本较高，但流匹配缓解了这一限制，从而形成了一种所谓的无模拟技术来训练连续归一化流。它通过对为数据样本（来自$\mu$的样本和来自$\nu$的样本）单独构建的条件路径$p_t(\cdot |x,y)$进行边缘化，构建了一条从$p_0\approx \mu$过渡到$p_1\approx \nu$的概率路径$p_t$。这些条件路径被用作基准真值，由学习到的向量场生成的路径需要逼近这些基准真值。形式上，我们选取任何具有边缘分布$\mu$和$\nu$的耦合$\pi \in \Pi (\mu ,\nu )$（例如独立耦合$\pi =\mu \otimes \nu$），并对$z\in {\mathbb{R}}^d$定义：
$$p_t(z)={\int }_{{\mathbb{R}}^d\times {\mathbb{R}}^d}{p}_t(z|x,y)d\pi (x,y)$$

我们假设这些条件路径可由条件速度场$u_t(\cdot |x,y)$生成。此时，边缘化的速度场可表示为：
$$u_t(z)={\int }_{{\mathbb{R}}^d\times {\mathbb{R}}^d}\frac{u_t(z|x,y)p_t(z|x,y)}{p_t(z)}d\mu (x)d\nu (y)$$

可以证明，该$u_t(z)$能够还原出边缘化路径$p_t$。原则上，通过最小化以下损失函数，我们可以学习到逼近$u_t(z)$的$v_{t,\theta }(z)$，进而生成逼近$p_t(z)$的路径：
$${\mathcal{L}}_{FM}(\theta )={\mathbb{E}}_{(X,Y)\sim \pi ,Z_t\sim p_t}\left[{\Vert v_{t,\theta }\left(Z_t\right)-u_t\left(Z_t\right)\Vert }^2\right](8)$$

然而，由于$u_t$依赖于$p_t$，这个目标函数难以求解。相反，研究表明，仅需对条件速度场进行回归并最小化以下损失函数即可：
$${\mathcal{L}}_{CFM}(\theta )={\mathbb{E}}_{(X,Y)\sim \pi ,Z_t\sim p_t(\cdot |X,Y)}\left[{\Vert v_{t,\theta }\left(Z_t\right)-u_t\left(Z_t|X,Y\right)\Vert }^2\right]$$

因为${\nabla }_{\theta }{L}_{CFM}(\theta )={\nabla }_{\theta }{L}_{FM}(\theta )$。因此，优化易于处理的损失函数$L_{CFM}(\theta )$与优化$L_{FM}(\theta )$是等价的。通过选择形式简单的条件概率路径$p_t(\cdot |X,Y)$和条件向量场$u_t(Z_t|X,Y)$，无需模拟轨迹就能学习到$v_{t,\theta }$。如果该损失为0，那么这个流就是一个推送映射，即$T={\varphi }_1$满足$T♯\mu \approx \nu$。换句话说，这意味着由神经速度场$v_{t,\theta }$诱导的流会将对照组细胞的分布传输到受扰动细胞的分布。

高斯概率路径

在本研究中，我们采用的基准真值概率路径的一种自然构造方式是使用如下形式的高斯条件概率路径：
$$p_t(z|x,y)=\mathcal{N}\left(m_t(z|x,y),{\sigma }_t(z|x,y{)}^2\right)$$

其中，${\mu }_0(x,y)=x$，${\mu }_1(x,y)=y$，且${\sigma }_t(x,y)\ge 0$。假设$m_t$和${\sigma }_t$关于$t$可微，则相应的条件速度场
$$u_t(z|x,y)=\frac{{\sigma }_t^{\prime }(x,y)}{{\sigma }_t(x,y)}\left(z-m_t(x,y)\right)+m_t^{\prime }(z|x,y)$$

可通过仿射条件流
$${\varphi }_t(z|x,y)=\frac{{\sigma }_t(x,y)}{{\sigma }_0(x,y)}\left(z-{\mu }_0(x,y)\right)+m_t(z|x,y)$$

诱导出上述高斯路径$p_t(z|x,y)$。

在本研究中，我们始终选择数据样本之间的线性插值作为条件路径上的均值，即
$$m_t(x,y)=(1-t)x+ty$$

然后我们考虑两种不同的噪声调度：

1. 恒定噪声：${\sigma }_t(x,y)=\sigma \ge 0$
2. 布朗桥：${\sigma }_t(x,y)=\sigma \sqrt{t(1-t)},\sigma >0$

这些调度对应以下条件速度场：

1. 恒定噪声：$u_t(z|x,y)=(x-y)$
2. 布朗桥：$u_t(z|x,y)=\frac{1-2t}{t(1-t)}(z-((1-t)x+ty))+(x-y)$

综上，我们将（含噪声的）直线作为条件概率路径。另一方面，尽管我们在源样本与目标样本对之间得到了各自的直线路径，但由边缘化速度场$u_t$诱导的映射${\varphi }_1$未必是最优传输（OT）映射。

CellFlow

在介绍了针对单一源分布$\mu$和单一目标分布$\nu$的最优传输耦合与流匹配之后，现在将这一设定扩展到多对源分布和目标分布{(μc,νc)}c∈C\left\{ (\mu_{c}, \nu_{c}) \right\}_{c \in C}{(μc​,νc​)}c∈C​。需要说明的是，这其中包含仅存在一个源分布但存在多个目标分布的情况，例如${\mu }_{c_1}={\mu }_{c_2}=\mu$，但${\nu }_{c_1}\ne {\nu }_{c_2}$。如前所述，作者将每一对源分布和目标分布$({\mu }_c,{\nu }_c)$与定义实验干预的变量$({p_{c_j}}_{j\in n_c},{q_{c_j}}_{j\in n_c},s_c)$相关联，这些变量定义了从源分布到目标分布的转换。作者将实验干预的相关信息作为输入提供给学习到的向量场，这样在训练后，当给定新的实验条件时，就可以对该向量场进行评估。

综合以上所有内容，我们现在可以推导出CellFlow的损失函数。下面我们详细阐述这些目标函数的构建过程，并说明模型选择对它们的影响。

1. 首先，我们为条件集合$c\in C$设定一个分布$\rho$，我们选择该分布为均匀分布，即$\rho =\frac{1}{C}{\sum }_{i=1}^C{\delta }_{c_i}$。
2. 然后，我们通过一个聚合器方案$A_{\theta }$，将扰动变量$c=({p_{c_j}}_{j\in n_c},{q_{c_j}}_{j\in n_c},s_c)$嵌入到一个条件向量$A_{\theta }(c)$中。
3. 给定条件$c\sim \rho$，我们从${\mu }_c$和${\nu }_c$中采样细胞批次。对于这些批次，我们通过设定${\pi }_c={\pi }_{OT,\epsilon ,c}$或${\pi }_c={\pi }_{UOT,\epsilon ,\tau ,c}$，并使用算法1对其进行逼近，从而得到一个耦合${\pi }_c\in \Pi ({\mu }_c,{\nu }_c)$。
4. 接下来，我们定义一个（在各条件间共享的）高斯条件概率路径族$p_t(z|x,y)=\mathcal{N}(m_t(x,y),{\sigma }_t(x,y{)}^2)$，用于连接来自条件源分布${\mu }_c$和条件目标分布${\nu }_c$的耦合数据样本$x\sim {\mu }_c$与$y\sim {\nu }_c$。这些路径由通过$m_t(x,y)$和${\sigma }_t(x,y)$得到的条件速度场$u_t(z|x,y)$诱导产生。
5. 最后，我们利用条件概率路径对中间细胞状态$z_t\sim p_t(z|x,y)$进行采样（在图1中记为$x_t$）。给定这些样本和条件向量$A_{\theta }(c)$，我们使用神经向量场$v_{t,\theta }(z_t|A_{\theta }(c))$来逼近$u_t(z)$。

由此得到CellFlow损失：
$${\mathcal{L}}_{\text{CellFlow}}(\theta )={\mathbb{E}}_{t,c\sim \rho ,(x,y)\sim {\pi }_c,Z_t\sim p_t(\cdot |x,y,c)}\left[{\Vert v_{t,\theta }\left(z_t|{\mathcal{A}}_{\theta }(c)\right)-u_t\left(z_t|x,y\right)\Vert }_2^2\right]$$

我们在算法2中重点介绍了CellFlow的训练流程。为了在每个批次中纳入多个条件，我们会在更新参数前，通过$T_{acc}$步进行梯度累积。

CellFlow的参数

CellFlow基于OT-CFM。

---

Improving and generalizing flow-based generative models with minibatch optimal transport（ICML 2023）
Multisample flow matching: Straightening flows with minibatch couplings（2023）
Unbalancedness in neural monge maps improves unpaired domain translation（ICLR 2024）

---

小批量最优传输耦合

离散最优传输耦合的超参数是（线性）最优传输中常用的那些：

• $\epsilon$：熵正则化参数，用于控制耦合的随机性。控制分布和受扰分布越“复杂”，$\epsilon$的值就应选得越小。为避免受数据尺度的影响，我们始终将成本矩阵按其均值进行缩放。默认值：$\epsilon =1.0$
• ${\tau }_a$、${\tau }_b$：左侧和右侧的非平衡参数，取值范围在0（完全非平衡）到1（完全平衡）之间。源分布中的异常值或非期望的分布偏差越多，${\tau }_a$的值就应越小；目标分布中此类情况越多，${\tau }_b$的值就应越小。默认值：${\tau }_a={\tau }_b=1.0$。

编码器模块

在进行拼接之前，每个扰动、协变量和细胞系嵌入都可以通过前馈神经网络（mlp）或自注意力模块（self-attention）进一步处理，作者将这整个过程称为layers_before_pool（拼接前的层）。随后，这些拼接后的向量会通过标准注意力机制（attention_token）、设定注意力机制（attention_seed）或均值池化/深度集合（mean）进行池化操作。当需要合并的向量集大小为1时，注意力机制会退化为自注意力机制，而池化操作则等同于恒等映射。得到的向量会再经过另一个前馈神经网络（layers_after_pool，拼接后的层）处理，之后被投影到大小为condition_embedding_dim（条件嵌入维度）的向量上，这便是编码器模块的输出。总之，相关参数及括号内的默认值如下：

• Pooling (attention_token)
• cond_output_dropout (0.9)
• condition_embedding_dim (256)
• layers_before_pool (specified for perturbation, perturbation covariate and sample covariate):– layer_type (mlp)– dims ([1024, 1024])– dropout_rate (0.0)
• layers_after_pool– layer_type mlp– dims ([1024, 1024])– dropout_rate (0.0)

流匹配模块

流匹配模块在给定条件下将源（对照）分布转换为目标（扰动）分布。time_encoder_dims是一个前馈MLP，用于对时间（非实验时间，而是神经微分方程的时间）的正弦嵌入进行编码，该嵌入的维度为time_freqs。hidden_dims指定了用于处理源（对照）分布中样本（细胞）的MLP架构。decoder_dims定义了MLP的各层，该MLP用于处理以下向量的拼接结果（若设置linear_projection_before_concatenation，则先进行线性投影；若设置layer_norm_before_concatenation，则先进行层归一化）：条件嵌入（来自条件编码器）、时间嵌入（output_dims的输出）以及经处理的细胞嵌入（hidden_dims的输出）。噪声调度可以任意选择，我们建议添加具有标准差flow_noise的恒定高斯噪声（constant_noise），或遵循薛定谔桥（Schroedinger bridge）噪声调度的噪声。

• time_freqs (1024)
• time_encoder_dims ([2048, 2048, 2048])
• time_encoder_dropout (0.0)
• hidden_dims ([4096, 4096, 4096])
• hidden_dropout (0.0)
• linear_projection_before_concatenation (False)
• layer_norm_before_concatenation (False)
• decoder_dims ([4096, 4096, 4096])
• decoder_dropout (0.0)
• flow_type (constant_noise)
• flow_noise (0.1)

训练

与大多数其他深度学习方法不同，由于存在根据最优传输计划进行的重采样步骤，一个批次（大小为batch_size）中的样本并非相互独立。迭代次数（num_iterations）是一个需要优化的超参数；由于推理成本相对较高，早停法的应用不像在单步模型中那样便捷。

multi_steps是最相关的参数之一。它表示梯度更新之间所涉及的条件数量（而非细胞数量）。其名称源于具体实现：由于数据结构的原因，我们首先对源分布（样本协变量）进行采样，然后对相应的扰动及扰动协变量进行采样。在这种条件配置下，我们从对照群体和扰动群体中采样细胞，这构成了训练过程的一个“步骤”。由于这仅相当于条件编码器模块接触到一个数据点，为避免条件编码器模块过拟合，我们将步骤数量增加到multi_steps。