跳转至

Querying Counterfactuals on Tissue Graphs with Supervised Disentanglement

会议: NeurIPS 2026
arXiv: 2606.08493
OpenReview: https://openreview.net/forum?id=YYY27OR76u
代码: https://github.com/PMBio/cellina
领域: 因果推理
关键词: 组织图反事实、反事实预测、监督解耦、空间转录组学、图变分自编码器

一句话总结

提出组织图反事实(Tissue Graph Counterfactuals)框架与双编码器图变分自编码器 Cellina,通过形式化边重连与节点通路扰动两类空间干预,并结合细胞类型分类与空间域对抗解耦,实现了精准的单细胞空间微环境反事实表达预测与亚结构发现。

研究背景与动机

单细胞生物学与空间组学长期致力于预测细胞对微环境扰动的响应,以揭示发育、稳态以及疾病恶化中的细胞行为机制。然而,现有的单细胞扰动预测方法大多依赖于两大理想化假设:一是将实验扰动视为在所有细胞上均匀施加的共享全局刺激,二是假设细胞之间条件独立,从而将响应建模退化为独立同分布(i.i.d.)学习问题。在真实的生物组织中,细胞绝非孤立的个体,其转录状态在很大程度上受到局部微环境(组织小生境)的塑造,取决于空间相邻细胞的细胞类型组成及其分泌的分子配体信号。因此,探索诸如“若将某个细胞置于完全不同的组织微环境,或其邻居细胞的信号通路活性被激活/抑制,该细胞的基因表达将发生何种变化”这一反事实问题,构成了组织空间生物学建模的核心挑战。

然而,在组织图(Tissue Graph)上进行此类反事实预测面临严峻的理论与建模障碍。从理论上看,反事实查询要求严格将细胞自身的固有细胞状态(Intrinsic State)与由外部微环境诱导的空间变异(Extrinsic Microenvironment Variation)解耦开来;但在纯无监督设置下,非线性隐变量模型的因子解耦在数学上并不具备可辨识性(Identifiability),变分自编码器的固有隐变量极易吸收空间结构信号而产生因果混淆。从建模范式上看,现存的空间转录组分析工具(如 MISTy、SIMVI)主要关注在已知静态组织中拟合可解释的空间方差分解,无法支持假想环境下的反事实生成;而经典扰动模型(如 scGen、CPA)仅能处理离散的粗粒度标签平移,无法表征拓扑邻居网络与邻域内连续变化的分子信号。

为化解这一困境,本文提出将组织空间干预系统性地形式化为组织图上的因果反事实推断问题,并构建了具备软归纳偏置监督解耦能力的双编码器图生成模型 Cellina。核心 idea:将组织图反事实明确形式化为拓扑边重连与邻居节点特征变换两类空间干预,借助双编码器结构将单细胞表达解耦为固有表征 \(z\) 与空间微环境表征 \(s\),通过细胞类型分类锚定 \(z\) 并对抗剥离空间域信息,以解耦后的 \(s\) 作为反事实条件输入实现高保真度的虚拟细胞空间干预模拟。

方法详解

整体框架

Cellina 采用双编码器图变分自编码器(Graph VAE)架构,将每个细胞在组织中的高维基因表达谱分解为两个相互正交的潜变量:表征细胞固有身份属性(与空间位置无关)的固有隐变量 \(z \in \mathbb{R}^k\),以及表征局部微环境空间影响(组织小生境)的外在空间隐变量 \(s \in \mathbb{R}^k\)。在生成过程中,解码器以 \(z\) 与 \(s\) 的拼接向量 \([z; s]\) 为输入,参数化负二项式(Negative Binomial)计数分布以重构基因表达;在推理过程中,\(z\) 仅由焦点细胞自身的基因表达编码,而 \(s\) 则从局部空间邻域特征中提取。为了解决无监督隐变量无法可信解耦的问题,Cellina 引入生物学软监督:利用已知细胞类型标签锚定 \(z\) 的生物身份,并通过对抗判别器强行剥离 \(z\) 中的空间组织域信号,迫使所有空间微环境诱导的方差流向 \(s\)。在反事实推断阶段,模型固定细胞固有的 \(z\),仅对 \(s\) 所依赖的邻域拓扑或邻域表达施加虚拟干预,进而解码出反事实状态。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    Input["组织图输入<br/>单细胞表达量与空间坐标图"] --> N1["双隐变量图变分推断<br/>分离固有细胞状态与空间微环境"]
    N1 --> N2["双重监督对抗解耦<br/>类型分类锚定与空间域对抗剥离"]
    N2 --> N3["组织图反事实查询<br/>边拓扑重连与节点通路靶向干预"]
    N3 --> Output["负二项式解码输出<br/>反事实表达谱与生物学效应预测"]

关键设计

1. 双隐变量图变分推断:分离固有细胞状态与空间微环境 针对现有方法将细胞固有状态与环境诱导状态混为一谈的问题,Cellina 建立了显式双隐变量生成体系 \(p(x, z, s) = p(x \mid z, s) p(z) p(s)\),其中先验分布均设为标准高斯分布 \(\mathcal{N}(0, I_k)\)。观测变量 \(x\) 服从单细胞计数建模中标准的负二项式分布 \(x \sim \text{NB}(\mu, r)\),解码器 \(\text{Dec}([z; s])\) 以连接后的隐表征和测序深度 \(\ell = \log \sum_g x_g\) 预测基因表达率均值 \(\mu = \exp(\ell)\mu_\theta([z; s])\) 与逆离散度 \(r = r_\theta([z; s])\)。在推断端,固有隐变量通过 MLP 编码器近似后验 \(q(z \mid x) = \mathcal{N}(\mu_z(x), \sigma_z^2(x))\);而针对空间隐变量 \(s\),本文提出了两种互补机制:基础版 Cellina 对空间近邻图(由高斯核权重 \(W_{uv} = \exp(-d(u, v)^2 / (2\sigma^2))\) 赋予)内的归一化对数表达量进行度归一化聚合,得到紧凑小生境特征 \(\varphi(v) = (\sum_{u \in \mathcal{N}(v)} W_{uv} \tilde{x}_u) / (\sum_{u \in \mathcal{N}(v)} W_{uv})\),再经由轻量 MLP 编码出 \(q(s \mid \varphi_v)\),实现与非空间基线媲美的超高训练吞吐;而进阶版 Cellina-GAT 则摒弃固定聚合,采用不包含自环的局部子图多层 GATv2 图注意力网络自适应聚合邻域消息,在稍高的算力成本下捕获高阶图拓扑依赖。

2. 双重监督对抗解耦:类型分类锚定与空间域对抗剥离 单纯优化变分下界(ELBO)无法阻止固有隐变量 \(z\) 擅自记忆空间位置或吸收微环境变异,导致因果反事实推理失效。为在不破坏生成灵活性的前提下引入因果解耦归纳偏置,Cellina 在 \(z\) 上构建了由分类器和判别器组成的双重监督约束机制。首先,设立细胞类型分类器 \(f_{\text{clf}}(z)\),以交叉熵损失 \(\mathcal{L}_{\text{clf}} = \mathbb{E}[-\log f_{\text{clf}}(y \mid z)]\) 强制 \(z\) 捕获细胞的核心类型身份(如上皮细胞、成纤维细胞、T 细胞等);其次,设立空间组织域对抗判别器 \(f_{\text{disc}}(z)\),采用两步交替优化策略:第一步固定 VAE 编码器,在停止梯度的 \(z\) 上训练判别器以最小化域分类交叉熵 \(\mathcal{L}_{\text{disc}} = \mathbb{E}[-\log f_{\text{disc}}(d \mid \text{sg}(z))]\);第二步固定判别器,优化 VAE 编码器最大化判别器的预测熵(通过最小化 \(-\mathcal{L}_{\text{adv}}\) 实现),强行抹去 \(z\) 中包含的空间域信息(如肿瘤区域与正常组织区域的区别)。此外,在 Cellina-GAT 中,还对空间隐变量 \(s\) 施加了图监督对比损失 \(\mathcal{L}_{\text{spatial}}\),将空间近邻作为正样本、跨域非近邻作为负样本,促使 \(s\) 形成平滑且具域区分度的流形。为了避免多任务权重失衡,模型在首轮训练后自动计算自适应归一化系数 \(\alpha\),无需繁琐的超参数调优即可保证稳健训练。

3. 组织图反事实查询:边拓扑重连与节点通路靶向干预 基于组织图 \(\mathcal{G} = (\mathcal{V}, \mathcal{E}, W)\) 的形式化,组织微环境的突变可精准归结为图拓扑结构(边集合 \(\mathcal{E}\) 与权重 \(W\))或邻域节点特征(基因表达矩阵 \(X\))的变动,对应两类因果反事实查询: - 边扰动(Edge Perturbation):形式化为将目标细胞 \(v\) 的邻域 \(\mathcal{N}(v)\) 替换为全新的反事实邻居集合 \(\mathcal{N}'\)。在组织域迁移评测中,针对源域 \(d\)(如正常黏膜)中类型为 \(y\) 的细胞,从目标域 \(d'\)(如结直肠癌肿瘤病灶)中与 \(y\) 相邻的其他类型细胞中采样虚拟邻居构建 \(\mathcal{N}'\),重新计算空间小生境特征后送入编码器得到反事实空间隐变量 \(s^{\text{cf}}\),与该细胞原始固有隐变量 \(z\) 拼接解码,模拟细胞迁入新微环境后的转录应答。 - 节点扰动(Node Perturbation):保持图拓扑结构不变,对邻居节点特征应用基因特异性变换 \(x_{u,g}^{\text{cf}} = T_g(x_{u,g}) = x_{u,g} \cdot \delta_g\)(针对特定基因子集 \(\mathcal{S}\))。这种机制允许研究者利用生物学先验(如 PROGENy 信号通路),针对性地上调或下调局部微环境中成纤维细胞或免疫细胞的 TGF\(\beta\)、NF\(\kappa\)B 等信号通路权重 \(\delta_g\),从而进行高精度的微环境信号通路因果干预仿真。

一个完整示例

以结直肠癌(CRC)临床切片中的成纤维细胞(Fibroblasts)为例,探究其从正常黏膜(REF)迁入肿瘤区域(CRC)后的反事实激活表型: 1. 固有身份编码:提取正常结肠黏膜成纤维细胞 \(v\) 的原始基因表达计数 \(x_v\),经 MLP 编码器映射至 \(z_v \in \mathbb{R}^{64}\)。在双重监督解耦作用下,\(z_v\) 牢固绑定了“成纤维细胞”的谱系特异性标记基因(如 COL1A1、ACTA2),但被剔除了正常黏膜这一空间位置指纹。 2. 反事实邻域采样与空间编码:在病理标注的肿瘤核心区(CRC)中,检索所有与肿瘤成纤维细胞相邻的非成纤维细胞池 \(\mathcal{P}_{\setminus \text{Fib}}\)(富含肿瘤上皮细胞、巨噬细胞等),为细胞 \(v\) 采样并构建虚拟恶性肿瘤微环境 \(\mathcal{N}'\)。经高斯距离加权聚合计算出肿瘤小生境特征 \(\varphi(v)^{\text{cf}}\),进而编码得到反事实微环境潜变量 \(s_v^{\text{cf}} \in \mathbb{R}^{64}\)。 3. 反事实解码与基因表达重建:将 \([z_v; s_v^{\text{cf}}]\) 输入共享负二项式解码器,按测序深度还原各基因表达率。模型成功预测出关键肿瘤相关成纤维细胞(CAF)激活标志物 IGF2、FN1、MMP3 的显著上调,其表达分布与真实观测的肿瘤区域成纤维细胞高度吻合,且保留了单个细胞固有的转录异质性。

损失函数 / 训练策略

模型综合优化目标包含负变分下界与多项加权正则化损失: $\(\mathcal{L} = \mathcal{L}_{\text{VAE}} + \lambda_{\text{clf}} \alpha_{\text{clf}} \mathcal{L}_{\text{clf}} + \lambda_{\text{spatial}} \alpha_{\text{spatial}} \mathcal{L}_{\text{spatial}} - \lambda_{\text{adv}} \alpha_{\text{adv}} \mathcal{L}_{\text{adv}}\)$ 其中 \(\mathcal{L}_{\text{VAE}}\) 包含重构负对数似然与带有线性预热(warmup)的 KL 散度约束;自适应缩放因子 \(\alpha\) 在首个 epoch 结束时根据各项损失相对于 \(\mathcal{L}_{\text{VAE}}\) 的均值比值固定,使得默认超参数 \(\lambda_{\text{clf}} = \lambda_{\text{disc}} = \lambda_{\text{adv}} = 1\) 在各类切片上均具备高度稳健性。训练采用两阶段交替梯度更新,判别器与 VAE 参数完全隔离,有效避免对抗崩塌。

实验关键数据

主实验

主实验在 6 例结直肠癌(CRC)临床切片(共约 240 万个单细胞,18,000 基因全转录组测序)上展开,采用留一细胞类型交叉验证(Leave-One-Cell-Type-Out)协议:在训练阶段完全留出肿瘤区域(CRC)中的某一细胞类型,测试时将正常区域(REF)的同类型细胞重连至肿瘤邻域,评估前 50 个差异表达基因(DEGs)的反事实预测质量。评估指标包含 log 对数差异倍数 Pearson 相关系数(\(r\))、带符号精确度(Signed Precision)、局部能量距离(Local E-distance)以及 log 对数差异倍数均方根误差(\(\text{RMSE}_{\text{LFC}}\))。

方法 扰动类型 Pearson \(r\) ↑ Precision\(_{\text{signed}}\) ↑ E-distance ↓ \(\text{RMSE}_{\text{LFC}}\) ↓
Mean shift 边扰动 0.51 ± 0.26 0.18 ± 0.15 29.89 ± 10.49 5.08 ± 2.51
scGen 边扰动 0.50 ± 0.37 0.19 ± 0.20 4.65 ± 3.49 1.99 ± 0.84
CPA 边扰动 0.68 ± 0.19 0.22 ± 0.19 6.44 ± 2.27 1.50 ± 0.56
MintFlow 边扰动 0.65 ± 0.31 0.23 ± 0.21 13.01 ± 1.76 1.51 ± 0.43
Cellina (无监督消融) 边扰动 0.70 ± 0.26 0.30 ± 0.22 5.07 ± 1.71 1.50 ± 0.91
Cellina 边扰动 0.82 ± 0.18 0.40 ± 0.19 7.55 ± 1.14 1.29 ± 0.64
Cellina-GAT 边扰动 0.85 ± 0.15 0.40 ± 0.18 9.35 ± 1.75 1.14 ± 0.62
SpatialProp 节点扰动 0.40 ± 0.32 0.07 ± 0.16 33.07 ± 10.46 6.44 ± 1.03
Cellina 节点扰动 0.85 ± 0.16 0.41 ± 0.18 7.82 ± 1.23 1.23 ± 0.70
Cellina-GAT 节点扰动 0.73 ± 0.21 0.32 ± 0.20 8.32 ± 1.80 1.47 ± 0.71

跨组织与跨物种泛化(成年小鼠全脑 MERFISH 数据集)

为检验模型在不同物种与组织系统中的跨域泛化能力,论文在包含约 14.6 万细胞的小鼠全脑 MERFISH 数据集上执行跨解剖域反事实预测,将丘脑(Thalamus)作为源域,将纤维束(Fiber-tracts)与等皮质(Isocortex)作为独立留出的目标域。

方法 扰动类型 Pearson \(r\) ↑ Precision\(_{\text{signed}}\) ↑ E-distance ↓ \(\text{RMSE}_{\text{LFC}}\) ↓
Mean shift 边扰动 0.43 ± 0.23 0.12 ± 0.09 25.16 ± 4.69 10.98 ± 2.78
scGen 边扰动 0.77 ± 0.17 0.21 ± 0.13 5.74 ± 4.81 6.53 ± 4.05
CPA 边扰动 0.82 ± 0.16 0.42 ± 0.13 7.02 ± 2.33 6.40 ± 4.76
MintFlow 边扰动 0.81 ± 0.17 0.29 ± 0.16 19.58 ± 1.73 7.24 ± 5.41
Cellina (无监督消融) 边扰动 0.79 ± 0.15 0.37 ± 0.17 8.96 ± 5.31 6.86 ± 4.48
Cellina 边扰动 0.83 ± 0.16 0.47 ± 0.15 8.01 ± 1.36 6.25 ± 4.81
Cellina-GAT 边扰动 0.85 ± 0.15 0.52 ± 0.14 8.69 ± 1.45 5.80 ± 4.50
SpatialProp 节点扰动 0.69 ± 0.14 0.07 ± 0.07 22.40 ± 3.16 6.94 ± 2.56
Cellina 节点扰动 0.82 ± 0.16 0.47 ± 0.15 9.07 ± 1.86 6.31 ± 4.76
Cellina-GAT 节点扰动 0.85 ± 0.15 0.51 ± 0.14 8.60 ± 1.43 5.86 ± 4.43

关键发现

  • 监督解耦为反事实精度带来核心增益:相比于完全无监督的消融版本 Cellina (ablated),引入细胞类型锚定与空间域对抗剥离使结直肠癌反事实预测的 Pearson \(r\) 从 0.70 提升至 0.82,Precision\(_{\text{signed}}\) 从 0.30 跃升至 0.40。这有力证明了对抗剥离空间域混杂能够迫使固有状态 \(z\) 纯净化,使空间微环境因果影响集中由 \(s\) 承载。
  • 连续空间编码显著击败离散标签方法:即使是不带监督的消融模型 Cellina (ablated),其在两项主基准上的基因效应相关度也普遍高于基于离散标签平移的 CPA 与 scGen,表明对每个细胞具体局部邻域分子组成的连续建模远优于粗粒度组织标签。
  • 无监督发现肿瘤亚结构与通路级因果驱动:在解耦后的空间隐空间 \(s\) 上进行聚类,模型无须监督即可发现肿瘤内部具有截然不同信号活性的微环境亚区(如富集 TGF\(\beta\) 的 CRC1 亚区与富集 NF\(\kappa\)B/MAPK 的 CRC2 亚区)。利用 PROGENy 通路先验对邻居细胞施加定向节点扰动,高度还原了这些亚区特异的成纤维细胞活化特征(相关性达到 \(r = 0.76 \sim 0.77\)),且成功预测出 FN1 与 MMP3 等标志性基因的上调。

亮点与洞察

  • 将组织空间干预上升为图因果反事实理论:创造性地将复杂的单细胞组织微环境效应形式化为图结构上的边扰动(拓扑重构)与节点扰动(特征调节),填补了空间转录组学缺乏严格反事实问题定义的理论空白。
  • 极简线性邻域聚合实现极致计算可扩展性:基础版 Cellina 仅采用离线度归一化邻居均值聚合配合轻量 MLP,便在各项反事实指标上逼近甚至超越高开销的 GATv2 模型,不仅单张切片训练仅需 1 小时,且能够无缝扩展至数百万细胞规模的大样本队列。
  • 连接纯数据驱动模型与先验生物机制通路:通过支持针对特定基因子集的节点缩放干预,无缝融合了 PROGENy 等生物学先验信号网络,展现了从“黑盒反事实预测”走向“靶向药物/免疫通路虚拟筛选”的广阔应用前景。

局限与展望

  • 反事实性质尚属操作性模拟而非严格 SCM:作者坦承当前反事实基于输入上下文条件的生成模拟,尚未建立基于 Pearl 严格结构因果模型(SCM)并满足非线性因果解耦可辨识性保证的理论框架。
  • 依赖粗粒度标注与细胞分割误差:监督解耦仍需离散细胞类型与组织域标注作为锚点,难以捕捉连续细胞状态转换;同时空间转录组成像的细胞分割伪影(跨边界分子串扰)可能对邻域因果分析引入系统偏差。
  • 实验验证依赖上下文迁移替代真实物理干预:受限于当前空间 CRISPR 筛选通量与分辨率,现阶段评测仍采用留一细胞类型上下文迁移,未来需结合真实的高通量空间扰动荧光实验进一步实证闭环。

相关工作与启发

  • vs scGen / CPA: scGen 与 CPA 依赖细胞间独立同分布假设,通过离散扰动向量在潜空间平移模拟应答,完全忽略了连续空间邻居结构与局部微环境异质性;Cellina 显式利用图拓扑与连续空间隐变量 \(s\),能够捕获每个细胞特异的微环境变异。
  • vs SIMVI: SIMVI 采用无监督 MMD/互信息惩罚尝试解耦固有与空间状态,但在百万级细胞规模下面临显存爆炸,且无法原生支持外分布反事实边/节点干预查询;Cellina 通过软监督对抗解耦不仅消除了因果混淆,且基础架构具备近乎线性的内存与计算扩展性。
  • vs SpatialProp / MintFlow: SpatialProp 局限于同域掩码图推断且缺乏双隐变量解耦;MintFlow 虽引入连续流匹配但无法外推至未见细胞类型;Cellina 首次联合定义并解决了组织图上的边扰动与节点扰动,实现了真正的外分布反事实迁移。

评分

  • 新颖性: ⭐⭐⭐⭐⭐ 首次系统性提出组织图反事实概念,并成功将因果解耦引入空间转录组微环境推断中。
  • 实验充分度: ⭐⭐⭐⭐⭐ 覆盖超过 250 万单细胞真实数据(CRC 全转录组与小鼠全脑 MERFISH),评估了外分布留一验证与通路靶向干预。
  • 写作质量: ⭐⭐⭐⭐⭐ 数学形式化清晰严谨,生物学机制阐释深入,图文与指标定义极具学术水准。
  • 价值: ⭐⭐⭐⭐⭐ 为空间生物学提供了一个兼具高精度、强解释性与百万级单细胞扩展性的虚拟细胞因果扰动平台。