MIOFlow 2.0: 从单细胞和空间转录组数据中推断细胞随机动态的统一框架.

0. TL; DR

MIOFlow（Manifold Interpolating Optimal-Transport Flow） 2.0 是一个统一的框架，通过整合流形学习、最优传输和神经微分方程来学习符合生物学信息的细胞动态。MIOFlow 2.0通过明确地建模三个生物学过程，代表了对先前生成模型的重大进步：

  1. 随机性与分叉：通过神经随机微分方程（Neural Stochastic Differential Equations）来建模。
  2. 非守恒的群体动态：通过一个学习到的生长率模型，该模型由非平衡最优传输（unbalanced optimal transport）初始化。
  3. 环境影响：通过构建一个联合的隐空间，将基因表达与空间转录组学数据中的空间特征（如邻近细胞类型组成和配体-受体信号）统一起来。

通过在一个对PHATE距离进行匹配的自编码器的信息丰富的隐空间中操作，MIOFlow 2.0确保了轨迹尊重数据流形的内在几何结构。作者在合成数据集和生物学应用（包括胚状体分化和空间分辨的美西螈脑再生）上验证了MIOFlow 2.0。结果表明，MIOFlow 2.0不仅提高了轨迹的准确性，还揭示了先前隐藏的细胞转变驱动因素，例如调控再生发育的特定信号微环境。

1. 背景介绍

单细胞RNA测序（scRNA-seq）技术彻底改变了我们以空前分辨率分析基因表达的能力。当在多个时间点收集这些数据时，它们捕捉了关键的发育程序、再生反应和疾病进展的“快照”。然而，scRNA-seq技术的一个根本限制是其破坏性：单个细胞无法被连续追踪。因此，研究人员得到的是在离散时间间隔下的群体水平测量，其中每个时间点都包含了处于不同发育阶段的不同细胞的转录谱。

这个实验限制带来了一个重大的计算挑战：我们如何从这些静态的群体快照中推断出连续的、细胞水平的轨迹，而在时间点之间并没有已知的细胞对应关系？这个问题因几个基本的生物学现实而变得更加复杂：

  1. 随机性：细胞的转变过程具有内在的随机性，这源于基因表达的噪声和细胞命运决定的概率性。
  2. 非守恒动态：细胞群体的大小会因增殖和死亡而不断变化，这在癌症治疗或胚胎发育等场景中尤为显著。
  3. 空间环境的影响：细胞的行为受到其空间背景的强烈影响，包括邻近的细胞、信号分子和组织结构。

空间转录组学（spatial transcriptomics）的出现为研究细胞轨迹如何依赖于组织环境提供了新的机会。然而，现有的轨迹推断方法并没有被设计用来利用这些空间信息。它们通常仅基于基因表达来建模动态，忽略了驱动细胞转变的关键环境因素。

现有方法在应对这些综合挑战方面存在不足。许多生成模型专注于从噪声生成数据，而非时间插值。而那些尝试进行轨迹推断的方法，如TrajectoryNet和Waddington-OT，通常做出如高斯分布等限制性假设。最近流行的流匹配（Flow Matching）虽然高效，但标准形式通常假设在环境空间中进行直线插值，忽略了数据所在的弯曲流形几何。

为了解决这些限制，作者提出了MIOFlow 2.0，一个全面的框架，它统一了流形学习、最优传输和神经微分方程，以学习受空间信息调节的细胞动态。

MIOFlow 2.0的核心创新在于它明确地将三个基本的生物学过程整合到一个统一的模型中：

通过在几何感知的自编码器的隐空间中操作，MIOFlow 2.0确保了推断出的轨迹能够尊重细胞状态空间的内在结构。

2. MIOFlow 2.0

作者提出了MIOFlow 2.0，一个用于学习生物学上合理的细胞轨迹的统一框架。该方法在一个流形正则化的自编码器的隐空间中操作，并学习一个能够将样本从一个初始分布$\mu_t$传输到目标分布$\mu_{t+1}$的连续时间演化$X_{u,t}$。该框架由四个核心部分组成。

2.1 流形正则化的自编码器

单细胞数据虽然维度很高，但其内在的生物学变化通常位于一个低维流形M上。为了尊重数据的这种内在几何结构，作者首先使用一个几何感知的自编码器将原始数据（基因表达$x_g$和空间特征$x_s$）映射到一个低维的隐空间$Z$。

作者采用了GAGA框架，它通过一个几何损失函数来训练编码器$E_\phi$，强制隐空间中的欧几里得距离与流形上的测地线距离保持一致。这个测地线距离是通过PHATE（一种基于信息论的距离度量）来近似的。

具体来说，对于一对细胞$x_i, x_j$，模型最小化它们在隐空间中的欧几里得距离 \(\|\|E_\phi(x_i) - E_\phi(x_j)\|\|_2\) 与它们的PHATE距离\(D_{\text{PHATE}}(x_i, x_j)\)之间的差异。通过在这种信息丰富的隐空间$Z$中进行轨迹推断，MIOFlow 2.0确保了推断出的运输成本对应于真实的发育距离，而不是高维噪声的假象。

2.2 用动态最优传输推断轨迹

动态最优传输（Dynamic Optimal Transport）旨在寻找一条连续路径，以最小的总动能将一个分布随时间演化到另一个分布。作者在隐空间$Z$中对细胞状态的这种连续演化进行建模。

目标是学习一个由神经网络参数化的时变向量场$f_\theta(z, t)$，它定义了细胞在隐空间中的瞬时变化率：

\[\frac{dZ_{u,t}}{dt} = f_\theta(Z_{u,t}, t)\]

通过一个神经ODE，可以从初始状态$Z_{u,0}$积分得到在任意时间$t$的状态：

\[Z_{u,t} = Z_{u,0} + \int_0^t f_\theta(Z_{u,\tau}, \tau)d\tau\]

这个过程需要满足在观测时间点$i$的分布约束，即$Z_{u,i} \sim \mu_i$。这个问题可以通过优化一个松弛的损失函数来高效求解，该损失函数包含两项：

\[W_2(\mu, \nu)^2 = \inf_{Z_{u,t}} \mathbb{E}\left[\int_0^1 \|f(Z_{u,t}, t)\|_2^2 dt\right] + \lambda D(\rho_1, \nu), \quad \text{s.t. } Z_{u,0} \sim \mu\]

第一项是路径的动能，鼓励模型学习“最短”的路径。第二项是一个软惩罚项，用于确保在终点$t=1$时，预测的分布$\rho_1$与目标分布$\nu$相似。$D$可以是任何距离度量，如Wasserstein距离。

MIOFlow 2.0的训练目标$L= \lambda_m L_m + \lambda_e L_e + \lambda_d L_d$由三个损失项组成：

  1. 边际匹配损失 ($L_m$)：在每个观测时间点，计算预测分布与真实分布之间的Wasserstein-2距离。
  2. 运输能量正则化 ($L_e$)：惩罚速度场$f_\theta$的大小，以鼓励平滑、高效的路径。
  3. 流形密度损失 ($L_d$)：惩罚那些偏离高密度数据区域的轨迹点，确保轨迹停留在数据流形上。

2.3 用神经SDE建模随机动态

为了捕捉细胞命运决定的随机性和基因表达的噪声，作者将确定性的ODE框架扩展为神经随机微分方程（Neural Stochastic Differential Equations, Neural SDEs）。一个伊藤SDE（Ito SDE）将状态演化描述为：

\[dZ_t = b(Z_t, t; \theta)dt + \sigma(Z_t, t; \phi)dW_t\]

漂移项 (drift term) $b$ 由一个神经网络参数化，代表确定性的发育程序。扩散项 (diffusion term) $\sigma$ 由另一个神经网络参数化，代表随机噪声和在分叉点的不确定性。$W_t$是维纳过程（Wiener process）。

为了使轨迹更稳定并尊重全局动态，作者还在漂移项中引入了一个“动量项”（momentum term），它是过去速度的加权平均，这有助于平滑梯度，使轨迹更加连贯。

2.4 建模细胞增殖与死亡

为了处理非守恒的群体动态，作者引入了一个增殖神经网络$h_\psi(z,t)$，用于估计在状态$z$和时间$t$的细胞的局部生长或死亡率。这个生长率被用作权重，来调整预测的目标分布$\hat{\mu}$：

\[\hat{\mu}(\cdot) = \frac{\sum_j h_\psi(z_j, t)\delta_{z_j}(\cdot)}{\sum_j h_\psi(z_j, t)}\]

这个加权的分布随后被用于计算边际匹配损失$L_m$。为了给增殖网络一个合理的初始值，作者使用静态的非平衡最优传输（unbalanced optimal transport, UOT）在相邻时间点之间进行计算，并将UOT解的源边际作为$h_\psi$的初始估计。

2.5 空间感知的动态：通过联合嵌入

MIOFlow 2.0能够整合空间信息。假设细胞的动态不仅取决于其内在的转录状态，还取决于其局部的微环境。为此，作者构建了一个联合的特征空间。

对于每个细胞，作者从其空间邻域中提取丰富的特征：

  1. 邻域组成：邻域内每种细胞类型的频率。
  2. 相互作用潜能：基于已知的配体-受体对，计算邻域细胞的配体表达与中心细胞的受体表达的乘积之和，以量化信号强度。
  3. 局部表达生态位：邻域细胞基因表达PCA嵌入的平均值。

这些空间特征被拼接并降维后，得到一个空间特征向量$x_s^{(i)}$。

使用两个独立的几何感知自编码器分别对基因表达$x_g^{(i)}$和空间特征$x_s^{(i)}$进行编码，并将它们的隐向量拼接起来，形成一个联合的隐状态$z^{(i)}$。Neural SDE直接在这个联合的流形上学习动态：

\[dZ_t = f_\theta(Z_t, t)dt + \sigma(Z_t, t)dW_t\]

通过这种方式，学习到的速度场$f_\theta$同时依赖于细胞的内在转录状态和外在的空间环境，从而能够模拟依赖于空间位置的不同发育轨迹。

3. 实验分析

作者在一系列合成和真实的生物学数据集上验证了MIOFlow 2.0的性能。

3.1 合成单细胞数据

作者使用SERGIO模拟器生成了两种拓扑结构的合成数据：三向分叉（trifurcation）和S形（S-shaped）轨迹。

在三向分叉数据上，基线方法（SF2M和OT-CFM）的轨迹倾向于在分支之间走“捷径”，穿过了没有数据点的区域。在更复杂的S形数据上，这种偏差更加明显，基线方法的轨迹严重偏离了弯曲的数据流形。相比之下，MIOFlow 2.0在两种情况下都能够精确地贴合数据流形，生成了与真实动态高度一致的轨迹。

在对隐藏时间点进行插值的任务中，MIOFlow 2.0在两种拓扑结构上都取得了最低的轨迹重构误差，显著优于其他方法。

3.2 真实的单细胞胚状体数据

作者在一个人类胚状体（Embryoid Body, EB）分化数据上进行了留一法（leave-one-out）交叉验证。

在隐藏不同的中间时间点（$t=1,2,3$）时，MIOFlow 2.0在W1距离和MMD距离等多个指标上都取得了与OT-CFM和SF2M相当或更优的性能，证明了其在真实数据上的插值能力。

3.3 消融研究

作者通过消融实验来评估MIOFlow 2.0中不同生物学先验的贡献。这些结果表明，明确地建模随机性、增殖和死亡对于学习真实的生物学动态至关重要。

3.4 生物学案例研究：美西螈脑再生中的空间转录组学

为了验证空间调节模块的威力，作者将MIOFlow 2.0应用于一个空间分辨的美西螈（axolotl）脑再生数据集。