多模态生物数据分析的可解释性多任务学习.

paper：Explainable multi-task learning for multi-modality biological data analysis

0. TL; DR

UnitedNet是一个可解释的单细胞多组学分析框架，采用交替训练策略同时优化联合细胞群识别（即多模态整合与聚类）和跨模态预测。实验证明，这种多任务学习机制能相互促进，使得两个任务的性能都优于单任务学习。UnitedNet采用编码器-解码器-判别器架构，能够隐式地学习和逼近各种复杂模态（如RNA、ATAC、电生理、空间信息等）的统计特性，并实现高质量的跨模态数据生成。

使用SHAP量化UnitedNet不同模态特征之间的关联，能够发现具有细胞类型特异性的调控关系。在多种真实的多组学数据集上，UnitedNet在多模态整合和跨模态预测任务上均取得了与当前最先进方法相当甚至更优的性能。

1. 背景介绍

随着生物技术的发展，现在可以在同一个细胞中同时测量基因表达（RNA）、DNA可及性（ATAC）、蛋白丰度、电生理活性甚至空间位置信息。这些多模态数据为理解复杂的生命过程提供了新的分析挑战。一个全面的多组学分析流程通常涉及以下几个核心任务：

联合细胞群识别 (Joint Group Identification): 将来自不同模态的数据进行整合，以获得比单一模态更精确、更鲁棒的细胞类型或状态定义。
跨模态预测 (Cross-modal Prediction): 利用一种或多种易于测量的模态，来预测另一种难以测量或成本高昂的模态。例如，用基因表达谱来预测细胞的电生理特性。
跨模态关联挖掘 (Relevance Discovery): 揭示不同模态特征之间的调控关系。例如，找到影响特定基因表达的DNA可及性位点（ATAC peaks）。

将这些任务整合到一个统一的框架中，面临着两大难题：

异质性数据处理：不同模态的数据具有截然不同的统计特征（如稀疏性、噪声水平、数据分布等），需要一个足够灵活的模型来适应它们。
目标函数冲突：细胞聚类和数据预测是两个目标迥异的任务，前者的目标是最大化类间距离、最小化类内距离，而后者的目标是最小化重构误差。

2. UnitedNet 模型

UnitedNet 基于编码器-解码器-判别器 (Encoder-Decoder-Discriminator) 架构，通过交替训练策略来协同完成多项任务。

2.1 编码器-解码器-判别器结构

假设有 $V$ 种模态的数据，对于第 $i$ 个细胞，其在第 $v$ 种模态的数据为 $x_i^{(v)}$。

编码器 (Encoder, $\text{Enc}_v$): 每种模态都有一个专属的编码器，它将高维的原始输入数据 $x_i^{(v)}$ 压缩成一个低维的、信息浓缩的模态特异性潜在编码 (modality-specific latent code) $z_i^{(v)}$。

\[z_i^{(v)} = \text{Enc}_v(x_i^{(v)})\]

解码器 (Decoder, $\text{Dec}_v$): 每种模态也有一个专属的解码器，作用是将一个潜在编码（可以是来自任何模态的）“翻译”并重构回原始的高维数据空间。例如可以用第 $v_1$ 种模态的潜在编码 $z_i^{(v_1)}$ 去预测第 $v_2$ 种模态的数据：

\[\hat{x}_i^{(v_1, v_2)} = \text{Dec}_{v_2}(z_i^{(v_1)})\]

判别器 (Discriminator, $\text{Dis}_v$): 每个模态也配备一个判别器，它分辨输入的数据是“真实数据” ($x_i^{(v)}$) 还是由解码器生成的“虚假数据” ($\hat{x}_i^{(v)}$)。通过与判别器的对抗博弈，编码器和解码器（合称生成器）生成越来越逼真的数据。

2.2 两个任务的交替训练

UnitedNet 并非同时优化所有目标，而是通过交替训练两个核心任务来避免目标冲突，并实现互相增益。

⚪ 任务1: 联合细胞群识别 (Joint Group Identification)

这个任务的目标是整合所有模态的信息，得到最准确的细胞分组。首先将来自同一个细胞 $i$ 的所有模态特异性编码 $z_i^{(1)}, \dots, z_i^{(V)}$ 通过自适应加权方案融合成一个统一的共享潜在编码 $z_i$。

\[z_i = \sum_{v=1}^{V} \eta_v z_i^{(v)} \quad (\sum \eta_v = 1)\]

权重 $\eta_v$ 是可学习的，模型会自动给信息量更大、噪声更小的模态分配更高的权重。

接着，在共享潜在编码 $z_i$ 的空间上进行聚类（无监督）或分类（有监督）。损失函数 ($L_{group}$)主要基于 DDC (Deep Divergence-based Clustering) 思想，包含三个部分：

$L_{c_1}$: 促使不同簇的软分配概率向量相互正交，以增大簇间分离度。
$L_{c_2}$: 促使每个细胞的软分配概率向量更接近one-hot向量，以增强簇内紧凑度。
$L_{c_3}$: 基于熵的正则化，防止所有细胞都挤在少数几个簇里，保证聚类结果的均衡性。

这个任务的目标是精准地实现不同模态间的相互转换。可以利用编码器和解码器进行各种组合的模态内重构和跨模态预测。

损失函数 ($L_{PGC}$)包括：

预测损失 ($L_{Wpredict}$, $L_{Cpredict}$): 标准的重构损失，通常是均方误差(MSE)或交叉熵(BCE)，用于衡量预测值与真实值之间的差距。
生成器/判别器损失 ($L_{Gen}$, $L_{Dis}$): GAN的对抗损失。
对比损失 ($L_{Con}$): 拉近同一个细胞不同模态编码 $(z_i^{(v_1)}, z_i^{(v_2)})$ 的距离，同时推远不同细胞编码的距离，强制模型学习一个跨模态共享的语义空间：

\[l_{i}^{(v_1, v_2)} = -\log \frac{\exp(s_{i,i}^{(v_1, v_2)} / \tau)}{\sum_{s' \in \text{Neg}} \exp(s' / \tau)}\]

其中 $s$ 代表余弦相似度，$\tau$ 是温度超参数，Neg 包含所有负样本对（不同细胞的编码）。

⚪ 交替训练流程

在每个训练周期中，模型先执行若干步“任务1”的更新（优化 $L_{group}$），再执行若干步“任务2”的更新（优化 $L_{PGC}$）。这两个任务通过共享的编码器和潜在空间相互连接，从而实现了性能的相互增强。

2.3 基于SHAP的可解释性

作者引入了 SHAP (SHapley Additive exPlanations) 算法，它能公平地为每个输入特征分配一个“贡献值”（Shapley value），以表示该特征对模型某个特定输出的贡献大小。

在UnitedNet中，可以利用SHAP来回答两类关键的生物学问题：

特征-分组关联 (Feature-to-Group Relevance): 对于某个特定的细胞类型，哪些输入特征（基因、ATAC peak等）的贡献最大？
特征-特征关联 (Feature-to-Feature Relevance): 在某个特定细胞类型内，某个模态的某个特征（如基因A的表达）与另一个模态的哪个特征（如蛋白B的丰度）关联最强？

3. 实验分析

作者在多种模拟和真实的多组学数据集上对UnitedNet进行了全面验证。

3.1 模拟数据集Dyngen

使用Dyngen模拟器生成了一个包含四种模态（DNA, pre-mRNA, mRNA, Protein）的数据集，并进行了基准测试和消融实验。

与只进行联合识别任务或只进行跨模态预测任务的UnitedNet相比，进行多任务交替训练的完整版UnitedNet在两个任务上都取得了更好的性能。

UMAP可视化清晰地显示，多任务训练能够更好地对齐不同模态的潜在编码，并形成更清晰的簇结构。

随着训练的进行，模态间潜在编码的距离越近，联合识别的准确率（ARI）和跨模态预测的精度（R²）就越高，三者呈现显著的正相关关系。这直接证明了两个任务是通过共享潜在空间的对齐而互相促进的。

3.2 真实多传感数据 (Patch-seq)

作者应用UnitedNet分析了一个包含形态学（M）、电生理学（E）和转录组学（T）三种模态的Patch-seq神经元数据集。

UnitedNet在无监督聚类中识别出的MET-types（综合三种模态定义的细胞类型）与文献报道的参考标签高度一致（主要类型ARI=0.82，亚型ARI=0.41），性能优于WNN、MOFA等方法。

UnitedNet能够高保真地在三种模态间进行相互预测。

通过SHAP分析，不仅识别出了对Pvalb神经元类型最重要的基因、电生理和形态学特征，还进一步量化了这些特征间的跨模态关联。例如，模型发现基因Lrrc38的表达与Pvalb神经元的平均放电频率高度相关，这与Lrrc38调控BK钾离子通道（影响神经元放电）的已知生物学功能完全吻合。这充分展示了UnitedNet挖掘生物学机制的强大能力。

3.3 大规模多组学数据 (Multiome ATAC+Gene)

作者挑战了一个包含13个批次、跨越不同供体和组织位点的大规模BMMC数据集，以测试UnitedNet处理复杂、大规模数据的能力。

UMAP可视化显示，原始的单一模态数据存在显著的批次效应，而经过UnitedNet整合后的共享潜在空间中，来自不同批次的细胞被均匀地混合在一起，批次效应被有效去除。在监督学习模式下，UnitedNet可以利用12个批次的带标签数据作为参考图谱，对第13个未标记的批次进行高精度的细胞类型注释，性能优于专用的注释迁移工具scANVI。