Effective Biological Representation Learning by Masking Gene Expression¶

会议: NeurIPS 2026
arXiv: 2605.31562
代码: https://github.com/recursionpharma/opentxfm
领域: 自监督学习
关键词: 转录组学、自监督学习、掩码自编码器、泊松重构损失、基因表征

一句话总结¶

针对转录组基础模型普遍难以超越线性降维基线的痛点，TxFM 提出基于瓶颈 CLS 与紧凑 MLP 解码器的转录组掩码自编码器，结合有界整流 tanh 激活与泊松重构损失，仅需 140 万精选数据即可在归纳扰动表征与基因功能召回上超越百倍规模的图谱级模型。

研究背景与动机¶

单细胞与批量 RNA 测序（RNA-seq）产生了涵盖海量基因表达的高维度生物数据，揭示了复杂的细胞状态与分子调控网络，在疾病机制解析与药物靶点发现中具有核心价值。然而，高通量表达谱普遍面临严重的文库深度差异、低表达基因的高技术噪声、实验批次效应以及高稀疏度零值。令人困惑的是，近期一系列系统性基准评测（如 Szałata 等、Bendidi 等、Ahlmann-Eltze 等）一致表明，现有动辄训练于数千万乃至上亿单细胞表达谱的大规模转录组基础模型（Transcriptomic Foundation Models, FMs），在未见细胞系或扰动预测等归纳迁移任务上，表现常常落后于简单的 PCA 线性降维甚至标准对数变换后的原始表达计数。

这种失败的根源引发了学界的双重猜想：要么原始基因计数已经承载了几乎所有可恢复的生物学信号，深层表征学习难以带来额外增益；要么现有的自监督预训练范式在模型架构设计、重构优化目标与训练数据构建策略上存在系统性缺陷。传统方法盲目迁移自然语言的顺序自回归或双向掩码语言模型（如 scBERT、scGPT、Geneformer），忽视了基因表达天然属于无序分子的多重丰度测量本质；同时，损失函数通常采用均方误差（MSE）或高斯假设，亦或引入负二项分布（NB/ZINB）导致高离散度基因过估计时梯度严重弥散，使得优化器倾向于用方差参数吸收误差而非改进均值表征。

本文的切入角度是重构转录组自监督学习的基础假设：摒弃复杂的逐 token 序列解码与外部生物先验（如 ESM 蛋白质预训练特征），将基因表达建模为无序集合的高比例掩码重构问题，并通过数学推导构建适配测序计数支持域的数值稳定性机制。本文的核心 idea 是：构建以非对称非序列掩码自编码器为骨干的 TxFM，通过瓶颈 CLS 提取高阶细胞表征并由极轻量 MLP 解码，结合有界整流 tanh 激活函数与泊松对数似然重构损失，并在包含单细胞与批量转录组的 140 万精选语料 DiverseRNA-1.4M 上预训练，以显著更小的数据与算力实现超越图谱级模型的归纳泛化能力。

方法详解¶

整体框架¶

TxFM 将单细胞或批量 RNA-seq 样本建模为无序的高维计数向量 \(\mathbf{x} \in \mathbb{N}^G\)。模型采用非对称掩码自编码架构（MAE），整体流程由数据预处理、高比例随机掩码、Transformer 编码器表征提取以及基于整流 tanh 激活的 MLP 泊松解码器四个阶段串联构成。预处理后，仅有少量未被掩码的基因 token 连同可学习的类别标记（CLS token）输入编码器，全量多头自注意力仅在可见 token 间交互；编码器输出的紧凑 CLS 向量作为整样本的全局生物表征，随后输入轻量级 MLP 解码器，重构样本所测基因的全量对数标准化表达分布。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["输入表达向量<br/>文库标准化 + log1p"] --> B["均匀随机掩码采样<br/>保留 K=2048 个可见基因"]
    B --> C["Transformer 编码器<br/>拼接基因编码与表达量 + CLS"]
    C --> D["整流 Tanh 激活函数<br/>有界非负支撑域约束"]
    D --> E["泊松重构损失优化<br/>相对误差加权与梯度饱和门控"]

关键设计¶

1. 均匀随机掩码采样：消除序列偏置的高压缩比无序子集抽取

针对基因表达数据属于无序分子集合的本质，模型明确避免将基因排序为语言句子的做法。设样本可测基因全集为 \(\{1, \dots, G\}\)，TxFM 采用无偏的均匀随机采样抽取子集 \(\mathbf{k} \subseteq \{1, \dots, G\}\)，子集大小固定为 \(K=2048\)（在不同测序面板下对应 75% 至 94% 的超高掩码率）。未掩码基因 \(i \in \mathbf{k}\) 与其连续表达标量 \(\mathbf{x}_i\) 经由特征拼接构成输入向量：

\[\text{token}_i = [\mathbf{E}_i; \mathbf{x}_i] \in \mathbb{R}^d\]

其中 \(\mathbf{E}_i \in \mathbb{R}^{d-1}\) 为可学习的基因嵌入向量。所有可见 token 与可学习的 CLS token 共同送入 12 层 Transformer 编码器。由于未掩码 token 数量被严格约束在 \(K\)，自注意力计算复杂度从 \(O(G^2)\) 骤降至 \(O(K^2)\)，不仅大幅削减算力消耗，而且迫使全局 CLS token 学习利用稀疏的共表达协同先验，从高度残缺的转录快照中捕获样本级的深层调控态。

2. 整流 Tanh 激活函数：适配文库上界与抑制梯度发散的有界非线性映射

传统回归模型采用 ReLU 或线性层预测表达量，极易在异常高表达或批次偏差下产生数值溢出与梯度爆炸。TxFM 为 MLP 解码器的最终输出层设计了专有的整流双曲正切函数（rectified tanh）：

\[\phi(z) = \log(L + 1) \, \text{ReLU}\left(\tanh\left(\frac{z}{4e}\right)\right)\]

该设计完美契合文库对数标度上限 \(\alpha = \log(L + 1)\)（文库基准大小 \(L = 10^5\)）。当预测 logits \(z \le 0\) 时，ReLU 截断将输出严格钳位为 0（对应原始对数表达计数的下界 0 与泊松速率 \(\lambda = 1\)），天然模拟了单细胞数据的大量零值稀疏分布，无需构建复杂的零膨胀显式混合头；当 \(z > 0\) 时，输出渐进逼近理论上限 \(\log(L + 1)\)。在反向传播中，其导数具有 \(\text{sech}^2(z / 4e)\) 形式，在极端高表达区间平滑衰减，构成了天然的自稳定梯度门控。

3. 泊松重构损失优化：基于相对误差的稳健梯度回传与抗弥散机制

面对表达计数的离散与过度离散特性，常规负二项分布（NB）和零膨胀负二项分布（ZINB）在预测值高估高离散基因时，其梯度公式因包含额外离散度参数而发生快速衰减（\(\lim_{\mu \to \infty} \frac{\partial \mathcal{L}_{\text{NB}}}{\partial \mu} = 0\)），导致优化器通过增大方差参数逃避均值拟合。TxFM 直接采用连续对数目标下的泊松负对数似然损失：

\[\mathcal{L}_{\text{Poisson}}(\mathbf{x}, \hat{\mathbf{x}}) = \frac{1}{G} \sum_{i=1}^G \left( e^{\hat{\mathbf{x}}_i} - \hat{\mathbf{x}}_i e^{\mathbf{x}_i} \right)\]

在泊松损失下，当模型高估均值时，梯度极限严格恒等于 1（\(\lim_{\lambda \to \infty} (1 - x/\lambda) = 1\)），永不发生梯度消失；低估时提供强负梯度促使模型迅速校正。结合激活函数 \(\phi(z)\)，整体损失对 logit \(z\) 的梯度为：

\[\frac{\partial \mathcal{L}}{\partial z} = \begin{cases} \frac{\alpha}{\beta} \text{sech}^2\left(\frac{z}{\beta}\right) \left( e^{\alpha \tanh(z / \beta)} - e^{\mathbf{x}_i} \right), & z > 0 \\ 0, & z \le 0 \end{cases}\]

双曲正割门控因子衰减了过高表达基因的主导权，将模型的拟合容量精确导向构成单细胞表达绝大多数的中低表达基因，迫使编码器将离散度与生物变异隐式压缩至潜在表征空间中。

4. 协同数据精选策略：融入表型判别扰动与混合模态的 DiverseRNA-1.4M

模型表现的飞跃离不开训练语料的高效构建。TxFM 构建了包含约 143 万样本、覆盖 44,349 个独特基因的 DiverseRNA-1.4M 语料库。针对 Replogle 等的大规模 K562 CRISPRi 扰动数据，作者并未盲目全量灌入，而是引入表型指纹（phenoprint）过滤机制，剔除约 75% 缺乏显著转录组效应的无效敲除细胞，极大提升了训练样本的信息密度。更关键的是，语料不仅包含 7 个单细胞数据集，还联合纳入了来自 TCGA 与 GTEx 的 3.4 万例批量 RNA-seq（bulk RNA-seq）样本。在掩码训练中，针对各数据集测序基因集合不重叠的难题，TxFM 摒弃粗暴的零填充，仅对各样本实际测序的基因子集执行掩码并反向传播损失，使模型成功学会跨单细胞与批量转录组的联合表征空间。

损失函数 / 训练策略¶

TxFM 默认采用 Base Transformer 骨干（12 层，12 头，隐藏层维度 768，参数量 159M），随机深度衰减率为 0.1。优化采用 AdamW，最大学习率设为 \(1 \times 10^{-3}\)，动量超参为 \((0.9, 0.999)\)，采用单周期余弦退火（One-cycle Cosine）调度并配合 10% 的预热步数。全局批大小为 1536，在 H100 GPU 集群上训练 200 个 epoch（消耗约 864 GPU 时）。训练过程支持跨数据集不同测序面板动态计算损失，仅对真实存在测序读数的基因计算梯度。

实验关键数据¶

主实验¶

在三个完全未见的测试集（RPE1、HEPG2、Jurkat，包含数十万细胞与 2,392 个 CRISPRi 敲除扰动）上评测零样本归纳扰动表征质量。评测综合 6 项指标平均得分（含批次校正 iLISI、线性分类、KNN、扰动一致性、已知关系召回及表达可重构性）：

模型	预训练数据集与规模	RPE1	HEPG2	Jurkat	综合均分 (Overall)
原始表达计数 (Raw data)	无 (输入基线)	41.93	30.81	33.06	35.27
标准文库+对数归一化 ((Lib+Log)Norm)	无 (输入基线)	41.37	34.12	33.37	36.29
Geneformer-v2	Genecorpus, 104M	20.15	19.89	19.63	19.89
scGPT	Atlases, 33M	30.67	29.35	26.86	28.96
Tahoe-x1-70M	Tahoe+atlases, 266M	33.15	31.36	28.47	30.99
AIDO.Cell-100M	Atlases, 50M	37.45	34.52	31.53	34.50
Cell2Sentence	网页文本+图谱, 50M	39.01	34.49	32.82	35.44
TranscriptFormer-Sapiens	TF-Sapiens, 57M	34.88	31.55	29.92	32.12
STATE-SE (引入 ESM2 先验)	Tahoe+atlases, 170M	40.38	35.10	33.81	36.43
PCA	DiverseRNA, 1.4M	35.72	34.00	30.91	33.54
Ridge PCA	DiverseRNA, 1.4M	39.61	36.76	33.85	36.74
TxFM-B (TF-Sapiens 训练)	TF-Sapiens, 57M	38.57	35.30	32.98	35.62
TxFM-B (精选数据训练, 本文)	DiverseRNA, 1.4M	42.17	38.63	36.52	39.11

消融实验¶

作者在跨数据集扰动一致性（perts）与模型参数内嵌全基因组功能关系召回率（enc: 编码器 codebook; dec: 解码器重建权重）上进行了详尽的逐项消融：

消融维度	具体配置	扰动一致性 (perts)	编码器召回 (enc)	解码器召回 (dec)	结论与机制分析
重构损失	MSE	27.2	31.1	41.8	均方误差受异常值干扰，梯度对极值敏感
	SmoothL1	23.4	23.2	38.1	线性过渡区缺乏相对误差惩罚力度
	NB (负二项)	32.1	29.1	41.7	高离散度基因过估计时梯度严重消失
	ZINB (零膨胀负二项)	31.1	34.2	32.1	零膨胀截断梯度，解码器性能大幅受损
	Poisson (本文默认)	37.3	32.6	43.9	恒定梯度防止弥散，全面取得最优综合性能
激活函数	ReLU (发散)	34.4	26.7	43.9	缺乏上界截断导致训练不稳定甚至发散
	Clamped ReLU	34.6	32.4	44.3	硬截断导致边界梯度不连续，表达表征受阻
	Softmax	27.6	16.2	15.3	破坏独立基因丰度绝对表达量，表征彻底坍塌
	Rectified Tanh (本文)	37.3	32.6	43.9	有界且平滑，sech² 因子对高表达稳健衰减
数据精选	仅 Bulk RNA (3.3万)	21.2	11.6	33.5	缺乏单细胞分辨率导致扰动预测能力受限
	去除 Bulk RNA (1.4M)	34.6	33.1	42.4	丧失批量样本的全转录组宏观共表达约束
	去除 K562 细胞 (93.2万)	31.4	31.5	42.0	证明即使是跨细胞系，扰动共变模式依然通用
	仅保留 K562 对照细胞 (1.0M)	34.3	32.0	42.2	未扰动对照亦能有效辅助学习生理流形
	精选 K562 (DiverseRNA, 本文)	37.3	32.6	43.9	剔除 75% 弱响应细胞，信噪比最高
	全量 K562 细胞 (2.8M)	35.6	31.2	44.9	噪声样本稀释了关键调控梯度的学习效率

关键发现¶

打破图谱规模神话：在相同测试集上，训练于 140 万精选样本的 TxFM-B（综合分 39.11）大幅超越训练于 5700 万样本的 TranscriptFormer（32.12）以及引入蛋白质结构先验、训练于 1.7 亿样本的 STATE-SE（36.43）。数据质量与表型响应显著度远比单纯样本堆积关键。
泊松似然与整流 tanh 的协同效应：将损失从泊松换为 MSE 或 NB 导致扰动一致性暴跌 5–10 个百分点；而整流 tanh 中的 \(\text{sech}^2\) 阻尼因子防止了由于极少数超高表达基因主导梯度回传而淹没弱表达调节基因的现象。
低维流形中的生物规律涌现：参数内嵌特征评估显示，TxFM 学习到的解码器权重天然蕴含蛋白质复合物（CORUM、hu.MAP）与信号通路（Signor、Reactome）等先验关系。经过 PCA 降维至 256 维后，TxFM 解码器关系召回率高达 43.9%，超越各类图谱模型。
单细胞与批量测序的强力协同：在 DiverseRNA 中仅加入 3.3 万个 TCGA 与 GTEx 批量样本，带来的提升幅度与加入 45 万个单细胞扰动样本相当，证实批量测序对矫正单细胞 Dropout 稀疏性与捕捉全局基因调控网络的独特增益。

亮点与洞察¶

去繁就简的架构设计：完全摒弃了在基因序列上堆叠长上下文自回归的沉重设计，将表达重构收敛为“稀疏可见子集输入 \(\to\) 瓶颈 CLS 提取 \(\to\) 极轻量 MLP 输出”，以极其克制的 159M 参数量颠覆了复杂架构的主导地位。
严谨正视并超越线性基线：正面回应了转录组大模型长期落后于 PCA 的行业危机，系统性梳理出梯度弥散与目标函数不匹配的根源，成为业内首个在纯归纳自监督设置下系统性压制 PCA 的深度自编码方案。
迁移与复用价值：提出的整流 tanh 激活函数与泊松相对误差损失形式高度通用，可直接无缝迁移到其他单细胞多组学任务（如 scATAC-seq 染色质可及性重构、空间转录组丰度估计与宏基因组计数建模）。

局限与展望¶

作者指出的局限：DiverseRNA-1.4M 语料目前主要聚焦于肿瘤学与扰动生物学背景，在免疫系统稀有亚群、非肿瘤发育胚胎等组织上的泛化能力可能受到数据分布偏置的限制；此外，尚未完全确立针对转录组自监督学习的清晰 Scaling Laws。
潜在不足：模型推理严重依赖于文库大小归一化与 log1p 这一经典预处理流程，对于极端建库深度不均或特定单细胞技术引入的特异性分子标签（UMI）饱和度效应缺乏端到端的物理建模。
改进方向：可进一步探索将基因调控网络先验以拓扑稀疏图的形式融入注意力机制，或将细胞微环境与组织空间坐标信息与表达自编码联合优化。

评分¶

新颖性: ⭐⭐⭐⭐ [针对基因表达无序离散本质量身定制 MAE 范式，推导了泊松梯度与整流 tanh 激活协同机制]
实验充分度: ⭐⭐⭐⭐⭐ [横跨 16 个前沿基础模型、3 大未见细胞系、6 项下游生物基准及详尽参数消融]
写作质量: ⭐⭐⭐⭐⭐ [动机层层递进，数学推导完备，严密回答了转录组大模型是否优于线性基线这一学界核心痛点]
价值: ⭐⭐⭐⭐⭐ [为生物表征学习确立了兼顾效率与性能的新范式，公开发布的精选数据集与模型权重实用价值巨大]