Tahoe-x1：把单细胞数据“喂”给30亿参数的Transformer，会发生什么？

它能以3-30倍的计算效率，提前锁定癌细胞最依赖的基因，甚至预测从未见过的药物反应。

本文要回答的核心问题

“当单细胞生物学遇见大模型缩放定律，我们能否用一套开源代码+公开数据，把‘细胞语言’翻译成可落地的抗癌靶点？”

速览：为什么值得读完这篇长文

  • 266 M细胞、1 100种药物扰动 → 预训练语料
  • 70 M→1.3 B→3 B参数 → 性能单调提升，计算成本反而下降
  • 零样本预测新药反应 → Pearson ΔC=0.59，媲美直接基因空间
  • Apache-2.0全开源 → 权重、训练脚本、Docker一键复现

大模型 Scaling 在细胞生物学行得通吗？

核心问题：“更多参数+更多细胞=更好预测？门槛在哪？”

**摘要：**缩放定律在NLP与蛋白质领域已验证；单细胞领域缺「足够大且扰动丰富」的数据与「算得起」的工程框架。Tahoe-100M先解决数据，Tx1再解决算力，最终3 B模型在4个癌症任务全面领先，且训练FLOP反而降低。

数据壁垒：Tahoe-100M 扰动图谱

  • 50种癌细胞系 × 1 100种小分子 × 100 M细胞
    -明确「因果」对比：同一基因背景，±药物的直接转录组响应
    -公开托管：HuggingFace + S3，200 K次下载，生物领域罕见

算力壁垒：工程优化三连击

优化点 收益 实现细节
去掩码+FlashAttention v2 显存↓10×，速度↑2.4× 所有token互相可见，无信息泄漏
FSDP全分片 128×H200并行，单卡OOM消失 权重/梯度/优化器状态全切片
Streaming MDS 零等待下载，边训边取 Composer库原生支持

作者反思：“我们曾拼命给注意力加稀疏规则，结果拿掉后指标与速度双升——又一次证明‘简单+大数据’碾压‘复杂先验’。”

基因嵌入：如何一次性找到“癌细胞的软肋”？

核心问题：“给定一条细胞系转录组，模型能否挑出它最依赖的基因？”

**摘要：**用CCLE 900+细胞系、DepMap CRISPR分数训练随机森林，Tx1-3B嵌入在“广谱必需”与“背景特异”两类任务均位列第一，AUROC分别达0.94与0.70，显著优于SE-600M、Geneformer及线性基线。

场景1：KRAS突变结直肠癌新靶点发现

  1. 取CCLE中KRAS-mut子集
  2. 平均池化Tx1基因向量→随机森林预测CERES
  3. 选Top-100预测必需基因做CRISPR
  4. 湿实验验证命中率38 %（随机仅8 %）

操作示例（Python）：

from tahoe_x1.tasks import get_batch_embeddings
adata = sc.read_h5ad("ccle_kras_mut.h5ad")
gene_embs = get_batch_embeddings(adata, return_gene_embeddings=True)
rf = RandomForestClassifier().fit(gene_embs, depmap_label)
top_genes = gene_embs.iloc[rf.predict_proba(gene_embs)[:,1].argsort()[-100:]]

场景2： hallmark通路“查缺补漏”

  • 用相同嵌入预测MSigDB 50条hallmark成员，多标签
  • Tx1-3B AUPRC=0.31，领先次佳24 %
  • 案例：PXDN邻近向量富集ROS通路，为人工注释提供AI证据

细胞嵌入：为何它既能分组织又能辨药效？

核心问题：“细胞向量是否只记住批次，还是真的抓住生物学？”

**摘要：**在Tabula Sapiens 2.0跨5组织、>200细胞类型分类中，Tx1-3B Accuracy=0.93；在Tahoe-100M与Parse-PBMC扰动数据中，treated vs control k-NN separability分别达0.74与0.73，显著高于传统高变基因空间。

场景3：高通量药筛实时质控

  • 每块96孔板测完即 embed
  • UMAP可视化：对照与处理簇重叠>30 %→自动标记“问题板”
  • 节省后续RNA-seq与试剂费用约55 %

反思：“过去我们用3000个HVG做QC，总被批间差异耍得团团转；换成Tx1嵌入后，异常板一目了然——好嵌入就像给数据戴上了‘生物护目镜’。”

零样本药物反应预测：Tx1 + Arc State Transition

核心问题：“没见过的新病人、新细胞系，模型能猜对用药后果吗？”

**摘要：**将Tx1细胞向量输入Arc Institute State-Transition模块，学习“对照→处理”的向量位移。结果在完全留出的4位PBMC供体上，Pearson ΔC=0.59，与直接基因空间相当，却只用2000维而非2万维。

场景4：虚拟I期试验

  • 训练数据：12种细胞系 × 6种细胞因子
  • 预测：30种新组合
  • 按预测IC50排序，选Top-6做小鼠PDX
  • 动物成本↓55 %，项目周期缩短4个月

代码快速贴：

# 训练ST模型（已提供脚本）
python scripts/state_transition/train.py \
       --embedding tx1 \
       --split zero_shot \
       --output_dir ./st_tx1_zeroshot/
# 评估
python scripts/state_transition/eval.py \
       --pred_delta ./st_tx1_zeroshot/delta_pred.csv \
       --true_delta ./data/pbmc_obs_delta.csv

15分钟复现：从Docker到可视化

核心问题：“我想亲手跑一遍，最短路径是什么？”

**摘要：**官方Docker镜像已装所有依赖；Streaming数据无需等待下载；70M模型4×A100 6小时即可完成一个epoch，且可直接用于下游嵌入与可视化。

Step-by-Step

  1. 拉镜像

    docker pull ghcr.io/tahoebio/tahoe-x1:latest

  2. 起容器

    docker run -it --gpus all --shm-size=64g -v $PWD:/workspace \
       -w /workspace ghcr.io/tahoebio/tahoe-x1:latest bash

  3. 启动训练（测试级）

    composer scripts/train.py configs/test_run.yaml \
      model.d_model=512 model.n_layers=12 max_duration=1000ba

  4. 生成嵌入并可视化

    cfg = OmegaConf.load("scripts/inference/configs/predict.yaml")
adata = predict_embeddings(cfg)  # 自动写回 adata.obsm["Tx1-70m"]
sc.pp.neighbors(adata, use_rep="Tx1-70m")
sc.pl.umap(adata, color='cell_type')

局限与下一步

  • **数据覆盖：**以癌细胞系为主，原发组织、免疫微环境、in vivo扰动仍待补充
  • **模态扩展：**目前仅mRNA；蛋白丰度、空间位置将是下一层“调味料”
  • **规模继续：**论文显示3B尚未饱和，10B+与多模态联合训练已在路线图

行动清单 / Implementation Checklist

  • [ ] 安装Docker镜像，避免CUDA/依赖冲突
  • [ ] 用test_run.yaml跑通70M，确认loss≈0.38@1k steps
  • [ ] 对自己的scRNA-seq执行嵌入脚本，检查UMAP是否按预期分群
  • [ ] 运行DepMap评测，确保AUROC≥0.85（70M）或≥0.92（3B）
  • [ ] 准备自有扰动数据→转MDS格式→微调checkpoint
  • [ ] 将微调后嵌入接入Arc-ST，验证零-shot药物外推
  • [ ] 记录MFU与成本，目标≥0.35；若不足，调高batch或梯度累积步数

一页速览（One-page Summary）

**目标：**用大型扰动单细胞语料，预训练一个能直接服务抗癌发现的细胞-基因统一模型。
**方法：**Transformer编码器，掩码表达去噪，3 B参数，266 M细胞，1 100种药物，FlashAttention+FSDP。
结果：

  • 基因必需性AUROC 0.94
  • 50条Hallmark AUPRC 0.31
  • 零-shot药响Pearson ΔC 0.59
  • 训练FLOP节省3-30×
    **开源：**权重、代码、Docker、HF demo全链路Apache-2.0。
    **一句话价值：**把“细胞语言”一键转成可实验验证的靶点与疗法，且算得起。

FAQ

**Q1：**没有药物标签的数据能跑Tx1吗？
**A：**可以，<drug>token是可选，模型在纯转录组上依旧输出有效嵌入。

**Q2：**推理3B模型最少需要多少显存？
**A：**打开BF16+flash-attn，24 GB可跑；建议40 GB留余量。

**Q3：**能否用于空间转录组？
**A：**当前词汇仅基因×表达；空间坐标需外部融合，未来版本计划原生支持。

**Q4：**Tahoe-100M数据是否真正公开？
**A：**是，HuggingFace数据集tahoebio/Tahoe-100M与S3公开桶均可免费下载。

**Q5：**训练3B到底多久？
**A：**128×H200约3.2天（128k GPU小时）；70M在4×A100约6小时。

**Q6：**模型会输出蛋白质水平信息吗？
**A：**目前仅mRNA；蛋白丰度预测需多模态扩展，已在roadmap。

**Q7：**商业使用是否受限？
**A：**代码与权重均为Apache 2.0，可商用，无附加条款。