scNODE: 时序单细胞转录组数据预测的生成模型.
0. TL; DR
测量不同时间点的单细胞基因表达是研究细胞发育的关键。然而，由于实验资源的限制和技术挑战，研究人员通常只能在离散且稀疏的时间点上对基因表达进行分析。这种时间点信息的缺失阻碍了下游的细胞发育研究。
为了解决这个问题，作者提出了scNODE，一个端到端的深度学习模型，能够in silico（计算机模拟）预测未观测时间点的单细胞基因表达。scNODE整合了变分自编码器（variational autoencoder, VAE）和神经普通微分方程（neural ordinary differential equations, neural ODEs），在一个连续且非线性的隐空间中预测基因表达。重要的是，scNODE引入了一个动态正则化项（dynamic regularization term），以学习一个能够抵抗分布偏移（distribution shifts）的隐空间，从而在预测未观测时间点的单细胞基因表达时保持稳健。
在三个真实的scRNA-seq数据集上的评估表明，scNODE的预测性能优于当前最优的方法。此外，作者还证明了scNODE的预测有助于在缺失时间点的情况下进行细胞轨迹推断，并且其学习到的隐空间对于沿发育路径对相关基因进行in silico扰动分析也很有用。
1. 背景介绍
理解基因表达如何随时间变化是细胞发育研究中的一个基本挑战。高通量单细胞RNA测序（scRNA-seq）技术使我们能够在单细胞分辨率上常规地分析基因表达，揭示了同一组织类型或发育阶段内广泛存在的异质性。通过在多个时间点进行scRNA-seq分析，我们可以通过推断细胞轨迹、识别细胞命运转变和表征差异表达基因来理解生物过程的动态。
然而，时间序列scRNA-seq实验存在显著的局限性。由于时间和资源成本巨大，研究人员通常只能在离散且稀疏的时间点上进行采样。这种非连续的观测导致了在连续测量的时间点之间存在信息丢失。因此，开发一个能够in silico预测任何时间点（无论是插值还是外推）基因表达的模型，对于进行更全面的单细胞时序分析至关重要。
目前，已有一些计算方法被提出来解决这个问题：
- 单时间点生成模型：一些方法使用生成对抗网络（GANs）或统计框架来增强数据集，但它们是为单个时间点的快照数据设计的，没有考虑发育的动态过程。
- 轨迹推断方法：大多数这类方法推断的是“伪时间”（pseudo-times），它基于低维空间中的几何距离来分析细胞分化，而不能准确地模拟细胞在真实物理时间下的发育过程。
- RNA velocity方法：这类方法通过分析剪接和未剪接的转录本数量来揭示发育趋势，但它们通常需要深度测序的数据，并且大多不是生成式模型。
- 插值方法：Waddington-OT和TrajectoryNet等方法可以在两个观测时间点之间插值生成细胞，但无法外推到观测范围之外的时间点。
- 多时间点生成模型：
- PRESCIENT：该模型使用PCA进行降维，然后在低维空间中用neural ODEs建模。然而，PCA是一种线性降维方法，可能无法捕捉scRNA-seq数据中复杂的非线性关系，并可能导致不同细胞类型在低维空间中混合在一起（overcrowding）。
- MIOFlow：该模型使用测地线VAE（geodesic VAE）来学习非线性的低维表示。但它在建模前就固定了低维空间，作者假设，如果未观测时间点的分布与观测时间点有很大差异，这种固定的表示可能无法很好地泛化。
为了克服上述方法的局限性，作者提出了scNODE。这是一个整合了VAE和neural ODE的框架，旨在通过一个动态的低维空间来准确预测任何未测量时间点的基因表达。
2. scNODE 方法
scNODE（single-cell neural ODE）框架的核心思想是结合VAE的非线性降维能力和neural ODE的连续动态建模能力，并通过一个创新的动态正则化项来确保模型的泛化性能。
2.1 scNODE使用VAE学习复杂的低维空间
scNODE首先使用一个变分自编码器（VAE）来学习单细胞基因表达数据的一个低维表示（即隐空间，latent space）。与PCA相比，VAE利用神经网络的非线性能力，能更有效地捕捉复杂的细胞关系和异质性。
具体来说，scNODE首先将所有观测时间点$t \in T$的细胞基因表达数据$X(t)$合并成一个大的数据集$X_{ALL}$，然后用一个VAE对其进行预训练。VAE由两部分组成：
- 编码器（Encoder） $Enc_\phi$：将高维的基因表达数据$X_{ALL}$映射到一个低维的隐空间，并参数化为一个高斯分布$N(\mu, \sigma)$。
- 解码器（Decoder） $Dec_\theta$：将从隐空间高斯分布中采样的点$Z_{ALL}$映射回高维的基因空间，以重构原始输入$\hat{X}_{ALL}$。
其过程可以表示为：
\[\mu_{ALL}, \sigma_{ALL} = Enc_\phi(X_{ALL}) \\ Z_{ALL} \sim N(\mu_{ALL}, \sigma_{ALL}) \\ \hat{X}_{ALL} = Dec_\theta(Z_{ALL})\]预训练的损失函数$L_{pre}$由两部分组成：
\[L_{pre} = \text{MSE}(X_{ALL}, \hat{X}_{ALL}) + \lambda \text{KL}[N(\mu_{ALL}, \sigma_{ALL}), N(0, 1)]\]- 重构损失：输入与重构输出之间的均方误差（MSE）。
- KL散度（KL divergence）：隐空间分布与标准高斯先验分布之间的距离。这一项强制编码器学习一个结构良好、易于分离的隐空间。
通过这个预训练步骤，scNODE获得了一个能够有效表征所有观测细胞变化的初始隐空间。
2.2 scNODE使用Neural ODE建模细胞动态
接下来，scNODE使用neural ODEs在VAE学习到的隐空间中对细胞的发育动态进行建模。一个ODE描述了一个量如何随另一个独立变量变化，因此非常适合用来描述基因表达如何随时间变化。
scNODE将细胞隐变量$Z(t)$随时间$t$的变化用一个neural ODE来参数化：
\[dZ(t) = \text{Drift}_\omega(Z(t), t) \cdot dt\]其中，$\text{Drift}_\omega$是一个由参数$\omega$化的神经网络，它模拟了细胞在隐空间中的发育“速度场”，即细胞状态转变的方向和强度。
为了从离散的时间点数据中学习这个连续的动态过程，scNODE采用如下策略：首先使用预训练好的编码器$Enc_\phi$将起始时间点$t=0$的细胞$X(0)$编码到隐空间，得到初始隐变量$Z(0)$；然后使用一阶欧拉法（或其它ODE求解器）来逐步积分上述ODE，以预测后续任何时间点$t$的细胞状态：
\[Z(t+\Delta t) = Z(t) + \text{Drift}_\omega(Z(t), t)\Delta t\]在每个观测时间点$t \in T$，将通过ODE模拟得到的隐变量$Z(t)$用解码器$Dec_\theta$映射回基因空间，得到预测的基因表达$\hat{X}(t)$。
由于无法在真实细胞和模拟细胞之间建立一一对应关系，scNODE使用Wasserstein距离来衡量在每个观测时间点上，真实数据分布$X(t)$与预测数据分布$\hat{X}(t)$之间的差异。Wasserstein距离$Wass(X, \hat{X})$衡量了将一个分布变换为另一个分布的“最小代价”。
2.3 scNODE使用动态正则化学习稳健的隐空间
传统方法的一个弊端是，它们通常先固定一个低维空间，然后再学习动态。如果未观测时间点的细胞分布与训练时有很大差异（即存在分布偏移，distribution shift），那么这个固定的隐空间可能无法很好地泛化。
为了解决这个问题，scNODE引入了一个关键的创新：动态正则化项（dynamic regularization term）。其核心思想是，让VAE学习到的隐空间不仅要能很好地重构数据，还要与neural ODE学习到的动态流形保持一致。具体来说，在每个观测时间点$t \in T$，scNODE计算由VAE直接编码得到的隐变量$Z(t)_{enc}$的分布与由ODE模拟得到的隐变量$Z(t)$的分布之间的Wasserstein距离，并将这些距离的总和作为正则化项$R(T)$：
\[R(T) = \sum_{t \in T} \text{Wass}(Z(t)_{enc}, Z(t))\]其中，\(Z(t)_{enc}\) 是从真实数据$X(t)$通过编码器 \(Enc_\phi\) 得到的。
最终，scNODE的整体损失函数$L_{dyn}$是联合优化VAE参数（$\phi, \theta$）和neural ODE参数（$\omega$）的：
\[L_{dyn} = \sum_{t \in T} \text{Wass}(X(t), \hat{X}(t)) + \beta R(T)\]第一项是真实数据与重构数据在基因表达空间的Wasserstein距离。第二项是动态正则化项，由超参数$\beta$控制其强度。
通过这个联合优化的过程，neural ODE学习到的动态流形会反过来“指导”VAE的隐空间，使其不仅能捕捉细胞的静态变化，还能适应动态的发育过程。这使得scNODE的隐空间对于预测分布发生变化的未观测时间点更加稳健。
3. 实验分析
作者在三个公开的scRNA-seq数据集上对scNODE进行了全面的评估，这些数据集覆盖了不同物种和组织，并且都包含超过10个时间点。
3.1 scNODE能准确预测未观测时间点的基因表达
作者设计了三种难度的任务来测试模型的插值（interpolation）和外推（extrapolation）能力：简单任务（只插值）、中等任务（只外推）和困难任务（既插值又外推）。
在困难任务中，作者将预测的基因表达和真实数据在2D UMAP空间中进行可视化。结果直观地显示，scNODE的预测结果（蓝色点）与真实数据（红色点）的分布和形状高度吻合，而基线方法PRESCIENT和MIOFlow的预测则存在明显偏差。
作者使用Wasserstein距离来量化预测的准确性（值越小越好）。在所有三个数据集的困难任务中，scNODE在大多数测试时间点上都取得了最优或次优的性能。特别是在需要外推的远期时间点（例如SC数据集的$t=15$），scNODE的优势非常明显，其Wasserstein距离（132.86）远低于PRESCIENT（150.53）和MIOFlow（162.12）。
作者进一步分析了模型的性能与其面对的分布偏移程度的关系。结果显示，scNODE相对于基线模型的性能提升程度与测试时间点和训练时间点之间的分布偏移水平呈正相关（Spearman相关系数最高可达0.93）。这有力地证明了动态正则化项的有效性：当预测任务更具挑战性（即分布偏移更大）时，scNODE的优势就越明显。
3.2 scNODE的预测有助于恢复细胞轨迹
由于scRNA-seq数据在时间上的稀疏采样，直接从观测数据中推断细胞轨迹可能会因为时间点之间的“断层”而失败。作者检验了scNODE的预测是否能“填补”这些断层，从而帮助恢复更连续的细胞轨迹。
作者使用PAGA算法来构建细胞群的拓扑连接图。G_true：使用所有时间点的真实数据构建的图，代表了“黄金标准”的细胞发育拓扑结构。G_removal：移除测试时间点后，仅用训练数据构建的图。可以看到，图的连通性被破坏，出现了多个不相连的组件。G_scNODE, G_MIOFlow, G_PRESCIENT：在训练数据的基础上，加入各自模型预测的测试时间点数据后构建的图。
作者使用Ipsen-Mikhailov (IM)距离来量化不同图之间的相似性（值越小越相似）。结果显示，G_removal与G_true的IM距离为0.200，而G_scNODE与G_true的IM距离仅为0.093，远小于G_MIOFlow（0.133）和G_PRESCIENT（0.170）。这表明，scNODE的预测数据能够最有效地帮助PAGA算法恢复出与完整数据最接近的、连续的细胞发育轨迹。
3.3 scNODE的可解释隐空间有助于扰动分析
最后，作者展示了scNODE学习到的隐空间不仅可以用于预测，还具有一定的生物学可解释性，能够用于in silico的扰动分析。
- 图A (最可能路径识别)：作者在斑马鱼胚胎（ZB）数据集上训练了scNODE。利用学习到的隐空间和动态，作者使用最小作用路径（Least Action Path, LAP）方法，计算了从起始细胞到两个不同终末细胞群——体节前中胚层（PSM）和后脑（Hindbrain）——的最可能发育路径。
- 图B (驱动基因识别)：沿着这两条路径，作者识别出了差异表达最显著的基因。对于PSM路径，排名最高的基因是TBX16；对于Hindbrain路径，排名最高的是SOX3。这两个基因都已被文献报道在相应的细胞命运决定中起关键作用。
- 图C (In silico扰动)：作者在起始细胞中模拟了对这些驱动基因的过表达和敲除。结果显示，过表达TBX16能够显著增加最终分化为PSM细胞的比例（从约40%增加到超过70%）。而扰动一个随机挑选的非差异表达基因则对细胞命运比例没有影响。
这些结果表明，scNODE学习到的隐空间不仅在数学上是稳健的，而且在生物学上是可解释的，能够捕捉到驱动细胞命运的关键基因，并用于模拟基因扰动对发育结果的影响。
