跳转至

scTrilemma: Balancing Identity, Invariance, and Fidelity in Single-Cell Representation Learning

会议: NeurIPS 2026
arXiv: 2609.38840
代码: https://github.com/yunhak0/scTrilemma
领域: 计算生物
关键词: 单细胞转录组学、表征三难困境、变分自编码器、信息路由、零样本泛化

一句话总结

针对单细胞无标签表征学习中生物身份、批次不变性与表达保真度三者不可兼得的「表征三难困境」,scTrilemma 提出潜空间瓶颈 VAE 架构,通过表达门控编码(E-Gate)、细胞表征解码路由(C-Route)与伪群体先验条件化(PB-Cond)实现三路解耦信息路由,在纯重建目标下兼顾了细胞状态表征与基因级表达预测。

研究背景与动机

单细胞 RNA 测序(scRNA-seq)数据的表征学习本质上面临着无监督与角色依赖的复杂挑战。在实际分析中,同一维度的表达变异往往因研究目标的不同而在「生物信号」与「技术噪声」之间灵活切换:例如供体(donor)之间的微弱差异在跨批次、跨机构匹配通用细胞类型时是必须强力消除的批次混杂因素,但在对比患者群体病理反应与耐药状态时却是最核心的生物信号;同时,细胞类型与功能状态必须在聚类与注释后才能确定,预训练阶段模型只能接触到各类技术混杂与生物调控深度纠缠的高维稀疏计数特征。

现有单细胞基础模型与经典集成方法大多受制于统一特征空间的表征假设,试图将所有类型的变化压缩到单一的全局细胞嵌入向量中。这种「保留或消除」的硬二元规则必然导致严重的权衡冲突,即本文形式化提出的表征三难困境(Representation Trilemma): 1. 生物身份(Biological Identity):必须保留精确的细胞类型、发育分化层级以及疾病诱导的细胞状态结构; 2. 环境不变性(Context Invariance):必须在无标签条件下抑制技术批次、采样组织、测序平台与供体带来的技术变异,避免批次主导相似度; 3. 表达保真度(Expression Fidelity):必须保留单基因层面的丰富表达变异与基因共调控关系,以支持准确的下游差异表达分析与生物通路重构。

当模型强制要求单一全局嵌入同时服务于细胞聚类、批次消除与全基因表达重建时,三者不可调和:过度追求不变性会抹去细微的疾病状态表达,过度保留点对点保真度又会导致嵌入被高表达技术噪声支配。本文的切入视角在于认识到这种张力并非单细胞数据内在的数学不可行性,而是源于「所有信号强行走同一条嵌入通道」的架构挤压。本文的核心 idea 是:将单细胞表征三难困境重新形式化为信息路由问题,基于潜瓶颈变分自编码器将表达变异分别解耦路由至编码器特征门控、解码器查询调制与潜空间先验分布,在统一无监督重建目标下实现三难平衡。

方法详解

整体框架

scTrilemma 的整体架构建立在固定容量潜瓶颈变分自编码器(Latent-Bottleneck VAE)之上,模型整体处理流程包含输入基因表达构建、Perceiver 风格跨注意力编码、潜变量变分后验推断、条件先验对齐以及跨注意力重构解码五个主要阶段。

对于单个细胞观测,输入为稀疏原始基因表达计数。模型首先提取标准化与对数转换后的表达量,通过表达门控模块(E-Gate)对可学习的基因身份嵌入进行逐特征增益缩放,生成具备软上限约束的基因-表达 token 序列;随后,编码器使用 \(K=64\) 个可学习潜查询向量,通过多层预归一化交叉注意力块从 \(4096\) 个随机裁剪的基因 token 中自适应提取高阶细胞特征,映射为潜变量的高斯后验分布参数;在潜空间中,后验均值的确定性均值池化被提取为供下游细胞比对评估的显式细胞表征;重构解码阶段,解码路由机制(C-Route)提取后验采样的池化向量来调制解码器中的基因查询向量,同时保留完整的 \(K\) 个潜变量 token 作为交叉注意力键值记忆;与此同时,基于无标签伪群体谱聚类的先验条件化(PB-Cond)将无监督供体/批次上下文直接作用于先验均值与方差,仅通过 KL 散度约束隐式引导后验正则化。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}, 'subGraphTitleMargin': {'top': 8, 'bottom': 16}}}%%
flowchart TD
    A["单细胞基因计数 x_i 与伪群体 p_b"] --> B["E-Gate 表达门控编码<br/>基因特征与表达值逐维增益缩放"]
    A -.->|无标签分组聚合| C["PB-Cond 伪群体先验条件化<br/>K-Means软分配映射对齐高斯先验"]
    B --> D["Perceiver 潜瓶颈跨注意力编码<br/>K=64 潜查询压缩为多Token潜变量"]
    D --> E["变分后验 q(Z_i|x_i)<br/>均值确定性池化得评估细胞表征 r_i"]
    C -->|KL散度正则化约束| E
    E --> F["C-Route 双路径重构解码<br/>池化均值调制查询 + 全潜Token充当记忆"]
    F --> G["ZINB 零膨胀负二项计数重构输出"]

关键设计

1. E-Gate 表达门控编码:约束表达强度的非线性特征增益调制

传统单细胞模型(如 scGPT)在构建基因输入时,通常将可学习的基因身份嵌入与表达值离散分桶/投影向量进行简单的加性相加。这种线性加法赋予了表达值独立的向量自由度,极易大幅改变基因 token 的空间方向与范数大小,使模型在编码初期被高表达管家基因或批次特异的全局表达偏移主导,冲淡了生物标记基因的特异性。

E-Gate 引入基于 Sigmoid 的逐特征非线性乘性门控机制。对于细胞 \(i\) 中的基因 \(g_{ij}\) 及其标准化表达值 \(\tilde{x}_{ij}\),输入 token 定义为: $$ \mathbf{e}{ij}^{\mathrm{gate}} = \mathbf{u})\right) $$ 其中 }} \odot \sigma\left(f_{\mathrm{gate}}(\tilde{x}_{ij\(\mathbf{u}_{g_{ij}} \in \mathbb{R}^d\) 为预训练中端到端学习的固定基底基因身份锚点,\(f_{\mathrm{gate}}: \mathbb{R} \to \mathbb{R}^d\) 为仿射变换层,\(\sigma\) 为逐元素 Sigmoid 激活函数。由于门控取值严格限制在 \((0, 1)\) 区间内,输入向量范数恒满足 \(\|\mathbf{e}_{ij}^{\mathrm{gate}}\| \leq \|\mathbf{u}_{g_{ij}}\|\)。这一硬性归纳偏置确保了基因身份始终扮演核心空间锚点角色,而表达量仅作为特征尺度的增益控制器。实验机制分析表明,训练后的 E-Gate 对广泛高表达但无特异性的基因产生自适应衰减阻尼,精准放行细胞类型特异的标志基因信号,从源头上切断了共享表达幅值对潜变量瓶颈的噪声污染。

2. C-Route 解码器查询调制:双路解耦重构梯度与下游细胞表征

在采用多 token 潜瓶颈的自编码器架构中,如果解码器仅依赖全部潜变量 token 的注意力记忆完成重构,则下游评估所用的均值池化细胞向量 \(\mathbf{r}_i\) 仅属于无梯度的后处理产物,无法保证其承载核心生物调控语义;而若强行将所有基因预测全部压入单一的均值池化向量中,又会瞬间退化为信息瓶颈,抹杀基因级别的异质性变化细节。

C-Route 创新性地设计了互补的双路重构路径。在特征提取侧,评估细胞嵌入严格采用无参数的后验均值简单平均 \(\mathbf{r}_i = K^{-1} \sum_{k=1}^K \boldsymbol{\mu}_{ik}\),避免任何下游投影头引入的经验偏置;而在解码侧,重构计算保留完整的采样后潜 token \(\mathbf{Z}_i \in \mathbb{R}^{K \times d}\) 作为跨注意力键值记忆以保留细粒度表达细节,同时提取后验采样的均值向量 \(\mathbf{s}_i = K^{-1} \sum_{k=1}^K \mathbf{z}_{ik}\)(其数学期望恰好等于评估表征 \(\mathbf{r}_i\)),通过两层带 GELU 激活的轻量 MLP \(h_{\mathrm{query}}\) 动态调制每个待重构基因的基底查询向量: $$ \mathbf{q}{ij} = \mathbf{u}_i)\right) $$ 该设计建立了两条清晰分工的重构通道:全量潜 token 记忆负责提供高容量的逐基因表达细节,而查询调制路径则使评估表征 }} \odot \sigma\left(h_{\mathrm{query}}(\mathbf{s\(\mathbf{r}_i\) 直接承受 ZINB 似然的显式梯度回传,在不引入额外辅助损失的前提下强制细胞表征学习跨细胞类型的整体生物调控指纹。

3. PB-Cond 伪群体先验条件化:在无标签场景下将批次环境引向隐式先验

常规批次校正方法(如 scVI)高度依赖显式的技术批次编号(batch ID),但在大规模零样本迁移到全新发布的数据集时,未见数据集根本没有已知的一致批次标识符;此外,直接将批次信息注入编码器或解码器极易形成批次旁路捷径,诱导模型直接利用批次特征拟合输出而丢失供体本身的真实病理状态。

PB-Cond 完全放弃任何先验生物注释,仅利用数据集与供体的粗粒度数据集合并结构,计算该组细胞的无标签平均表达谱即伪群体(pseudo-bulk)特征 \(\mathbf{p}_b\)。在预训练阶段,模型在预训练伪群体集上拟合 \(M=32\) 个固定的 K-Means 表达原型中心;针对任意输入细胞所属的群组,计算其与 32 个中心的欧氏距离标准化软分配向量 \(\mathbf{t}_b \in \Delta^{M-1}\)。此软分配向量通过两层线性层映射为 VAE 先验高斯分布的均值与对数方差: $$ p_\eta(\mathbf{Z}i \mid \mathbf{t}}) = \prod_{k=1}^K \mathcal{N}\left(\mathbf{z{ik};\; \boldsymbol{\mu}\eta(\mathbf{t}{b_i}),\; \operatorname{diag}\left(\boldsymbol{\sigma}\eta^2(\mathbf{t}_{b_i})\right)\right) $$ 最关键的架构约束在于:该群组上下文仅用于计算 KL 散度目标 \(\mathcal{L}_{\mathrm{KL}, i}^{\mathrm{PB}} = D_{\mathrm{KL}}\left[q_\phi(\mathbf{Z}_i \mid \mathbf{x}_i, S_i) \,\|\, p_\eta(\mathbf{Z}_i \mid \mathbf{t}_{b_i})\right]\),绝不输入给编码器输入端、解码器查询或解码器记忆。在测试阶段,零样本推断完全基于后验路径 \(q_\phi\) 进行,完全不依赖测试集的伪群体先验代码。这种先验引导的正则化迫使后验分布在训练阶段主动适应并剥离供体层级的粗粒度技术偏移,实现了优雅的即插即用式跨供体不变性。

损失函数 / 训练策略

scTrilemma 全程使用统一端到端的半变分重构目标进行训练,无需任何对抗判别器或多任务监督损失。总损失函数由两部分构成: $$ \mathcal{L}{\mathrm{pretrain}, i} = \mathcal{L} $$ 其中 }, i} + \lambda_{\mathrm{KL}} \mathcal{L}_{\mathrm{KL}, i}^{\mathrm{PB}\(\mathcal{L}_{\mathrm{ZINB}, i}\) 采用零膨胀负二项(Zero-Inflated Negative Binomial, ZINB)分布的负对数似然损失,解码器输出包含 Softmax 归一化均值参数、全基因独立的逆离散度参数以及逐基因的零膨胀概率;KL 散度系数 \(\lambda_{\mathrm{KL}}\) 设置为 \(5 \times 10^{-4}\),并在前 2,500 步(占总步数 25,000 步的 10%)采用线性预热调度。模型采用 AdamW 优化器,学习率为 \(2 \times 10^{-4}\),余弦退火衰减,在 4 张 GPU 上采用 bf16 混合精度与 DDP 分布式训练,单卡 batch size 为 256。

实验关键数据

主实验

在跨版本严格零样本评估协议下(在 2025 年 1 月版 CZ CELLxGENE Census 的 6,260 万个细胞上预训练,在 2025 年 11 月新收录的 89 个独立数据集、3,040 万个细胞上测试),scTrilemma 与主流单细胞基线模型在生物身份保持与环境不变性上的全景对比如下:

方法 模型架构与预训练 生物身份: NMI (↑) 生物身份: ARI (↑) 身份: ASW (↑) 不变性: BRAS (↑) 不变性: PCR (↑) 不变性: iLISI (↑)
scVI 计数变分自编码器 (8.0M) 0.640 (0.138) 0.476 (0.155) 0.570 (0.049) 0.730 (0.093) 0.195 (0.135) 0.218 (0.174)
Geneformer 排序基因 Transformer (316.3M) 0.558 (0.166) 0.377 (0.156) 0.545 (0.037) 0.732 (0.106) 0.112 (0.115) 0.167 (0.169)
scGPT 表达值生成式大模型 (51.9M) 0.645 (0.141) 0.470 (0.158) 0.575 (0.052) 0.728 (0.091) 0.215 (0.107) 0.196 (0.173)
CellPLM 组织感知多细胞模型 (65.8M) 0.595 (0.142) 0.404 (0.150) 0.559 (0.077) 0.677 (0.087) 0.284 (0.175) 0.236 (0.178)
scPRINT 元数据全监督模型 (101.1M) 0.405 (0.154) 0.240 (0.139) 0.510 (0.017) 0.748 (0.105) 0.025 (0.063) 0.138 (0.135)
scTrilemma 潜瓶颈路由 VAE (58.3M) 0.651* (0.149) 0.484* (0.152) 0.540 (0.032) 0.840*** (0.073) 0.378*** (0.104) 0.181 (0.165)

在涵盖肾脏、肝脏、大脑、肌腱与结肠五个跨组织、跨供体配对正常 vs 疾病控制队列中,对表征三难困境的统一综合评估如下:

评估维度 核心量化指标 scVI Geneformer scGPT CellPLM scPRINT scTrilemma
生物身份与状态 细胞类型平衡准确率 (BA - chance) (↑) 0.6535 0.6112 0.6667 0.6041 0.4475 0.6697
局部近邻纯度 (Neighbor Purity) (↑) 0.6817 0.6006 0.7041 0.6459 0.4327 0.6795
疾病病理状态准确率 (Disease BA) (↑) 0.5583 0.5785 0.5744 0.5507 0.5432 0.5856
环境批次不变性 残差供体不变性 (Residual Donor Invar) (↑) 0.8889 0.8864 0.8778 0.9161 0.9016 0.9336
条件供体 BRAS (Conditional BRAS) (↑) 0.7113 0.7240 0.7078 0.6194 0.7835 0.8435
局部供体混合度 (Local Donor Mixing) (↑) 0.8216 0.7879 0.7982 0.8653 0.7582 0.7855
表达保真度 差异表达基因重叠 (DEG Jaccard) (↑) 0.1138 — — 0.1670 0.2202 0.2688
差异表达符号一致性 (DEG Sign Conc) (↑) 0.7160 — — 0.7051 0.7423 0.7861
生物通路富集重叠 (Pathway Jaccard) (↑) 0.2525 — — 0.2539 0.2581 0.3451
通路活性相关系数 (Pathway Score \(\rho\)) (↑) -0.1108 — — -0.0378 -0.0824 0.2288

消融实验

通过对 89 个零样本测试集上各路由模块系统移除后的标准化性能变化(相对于完整模型标准差归一化 \(\Delta / \mathrm{SD}\)),量化各核心组件的具体职责:

配置变更 生物身份 (NMI / ARI) 环境不变性 (BRAS / PCR) 表达保真度 (Pearson / DEG Jaccard) 核心机制作用说明
完整模型 (scTrilemma) 基准性能 (0.651 / 0.484) 基准性能 (0.840 / 0.378) 基准性能 (高相关 / 0.2688) 三路解耦实现表征三难的最佳全局平衡
移除 E-Gate (改为直接加性相加) 显著崩塌 (各集均全面暴跌) 严重退化 (批次混杂全面渗透) 严重退化 (高表达噪声淹没特异性) 作为整个路由管线的基础前置条件,缺少门控导致所有指标一致性崩溃
移除 C-Route (取消查询调制) 显著下降 (梯度不再约束池化表征) 轻微上升 (表征仅被动池化) 线性相关度大幅下降,点对点 MSE 降低 证明查询调制牺牲部分纯点对点拟合,换取评估嵌入的生物语义与相对排序保真度
移除 PB-Cond (替换为标准先验) 中度下降 (后验流形组织粗糙) 供体方差不变,BRAS 轻度下降 重构表达及通路指标显著受损 证明伪群体先验通过隐式对齐优化了后验流形,为复杂重构提供先验组织支持

关键发现

  • 三难破局的实证验证:对比基线模型可以发现明显的偏科现象——scGPT 在生物身份分离度(ASW)上领先但在环境不变性与跨供体对比上滞后;CellPLM 激进地混合局部邻域(iLISI 领先)却导致条件 BRAS 垫底,抹杀了真实的生物异质性;scPRINT 虽具备较强的基因排序却牺牲了全局细胞分类;唯有 scTrilemma 在维持极高生物状态识别(Disease BA 达 0.5856)的同时,实现了断层领先的环境不变性(BRAS 0.840,相比次优提升超 11%)与生物通路保真度(Pathway Score \(\rho\) 达到唯真正相关 0.2288)。
  • 潜在干预实验揭示核心权衡本质:在测试期对潜变量进行的单向干预实验表明,人为向环境不变性施加干预(消除供体均值偏置)对生物身份与表达保真度几乎无负面损耗,表明批次环境方差在路由解耦下完全可以被低成本剔除;真正的本质张力仅存在于「生物身份」与「基因表达保真度」之间(强行强化聚类中心会损耗单基因精细排序),而这二者均由 C-Route 共同调控。

亮点与洞察

  • 单细胞基础模型视角的范式转移:跳出了「做更大参数的序列自回归模型」或「设计更复杂对比学习损失」的传统思路,首次深刻指出单细胞表征退化源于将冲突信号强塞入同一嵌入向量的物理挤压,将问题优雅转化为架构级的「信息路由」。
  • 即插即用的零样本批次不变性:PB-Cond 利用无标签伪群体在测试集完全不依赖批次 ID,仅作为训练期 KL 散度约束,既避免了显式批次编码在推理时的泛化崩溃,又防止了解码器形成批次泄露捷径。
  • 端到端纯自编码目标的极简性:无需复杂的细胞本体图谱监督、无需多源任务混合加权,仅依靠 ZINB 重构与标准高斯 KL 正则,即全面压制了包含 1 亿参数及元数据监督的大型模型。

局限与展望

  • 计算内存开销较大:由于采用 Perceiver 风格多潜变量注意力机制,且每步需要重构 4,096 个高维基因位置,单卡峰值显存占用达到 19.79 GB,推理通量(310.9 cells/s)相比轻量 MLP 架构的 scVI 存在一定差距。
  • 空间转录组与多组学扩展尚未覆盖:当前验证完全集中在单细胞 RNA-seq 解离组织数据,未包含空间坐标(Spatial Transcriptomics)以及单细胞多组学(如 scATAC-seq 染色质可及性)的联合路由。
  • 极端孤立稀有细胞类型的聚类边界:数据特征分析表明,scTrilemma 的聚类增益主要伴随细胞类型丰富度与长尾细胞比例的提升而增长,但在完全孤立无近邻的罕见极小群样本上分离度仍有待进一步挖掘。

相关工作与启发

  • vs scVI: scVI 是单细胞生成模型的基石,依赖显式批次注入解码器进行批次校正。在跨发布版本的无标签零样本场景下,scVI 无法获得未知批次参数;而 scTrilemma 通过 PB-Cond 的自适应原型与 E-Gate/C-Route 路由,实现了全无监督下的鲁棒跨批次泛化。
  • vs scGPT / Geneformer: 现有单细胞大模型将细胞视作基因序列,使用类似加性词嵌入或排序嵌入,导致模型极易受高表达噪声与管家基因扰乱;scTrilemma 提出的乘性 E-Gate 与 Perceiver 潜瓶颈能够在同等参数量下更精准地提取标记基因模式。
  • vs CellPLM / scPRINT: CellPLM 引入组织级图建模,但过度平滑导致局部供体混合时抹杀了生物状态;scPRINT 引入复杂的元数据全监督训练。scTrilemma 证明了仅靠合理的解耦路由设计,纯无监督自编码器即可在所有核心指标上全面超越大参数监督模型。

评分

  • 新颖性: ⭐⭐⭐⭐⭐ 首次提出单细胞表征三难困境并将其重新定义为架构级信息路由问题,概念高度清晰深刻。
  • 实验充分度: ⭐⭐⭐⭐⭐ 采用跨版本严格时间分割零样本协议,涵盖 89 个独立测试集与多个真实疾病配对队列,并包含扎实的潜在干预分析。
  • 写作质量: ⭐⭐⭐⭐⭐ 概念论述严密,动机与架构设计前后呼应,数学形式化清晰自洽。
  • 价值: ⭐⭐⭐⭐⭐ 为单细胞基础模型架构演进提供了极具启发性的信息路由范式,开源代码与评估基准实用性极高。