🔬因果模型需要因果数据——Xaira 的 X-Cell 药物发现模型(Bo Wang 与 Ci Chu)
Xaira 推出 X-Cell 模型,主张因果建模需要信息丰富的数据;基于 CELLxGENE 数据库,突破基因表达预测瓶颈,并强调数据质量对扩展的决定性作用。
中文处理结果
押注信息
如果在达到 1.5B 参数后测试损失趋于平稳,而训练损失随着规模的扩大持续下降,这告诉你,你的模型受到数据中信息量的限制。
在单个、较小的数据集上训练暴露了一个信息缺口:3.1B 模型偏离了缩放趋势。在这堵墙之后,无论是参数还是计算量都无法提升性能。要预测基因表达的变化,你需要更富含信息的数据。
这正是 Chu 和 Bo 的团队所做的,以下是大致 30 倍信息量带来的效果:
现在我们可以随参数和训练计算量扩展了!我们不知道这项工作花费了多少,但可以猜测数据收集实验和基础设施花费了几千万,而计算 + 人力 + 研究花费了几百万。这一预算看起来像是强化学习(RL)部署的预算,而不是数据丰富的预训练预算。
我们很幸运,在我们的播客中邀请到了这个故事中的两位核心人物。在 Ci Chu 和 Bo Wang 的领导下,Xaira Therapeutics 正押注于信息丰富的数据是 AI 驱动药物开发的关键。Chu 最近被提升为首席发现官,Bo 被提升为首席 AI 科学家1,这突显了 Xaira 对这一赌注的高度重视。
逆向工程人类细胞
如果你必须弄清楚人类细胞是如何工作的,你会怎么做?一个不错的起点是记录在不同类型的细胞、不同情况下哪些基因被表达(例如,存在哪些 RNA)。
这就是 CELLxGENE,一个由 Chan Zuckerberg 研究所构建、包含 1.68 亿个细胞的数据库,它把每个细胞映射到该细胞中检测到的 2 万至 3 万个基因的出现次数,并附有每个细胞的详细元数据。这是一个约 4 万亿条目的矩阵。
如果说蛋白质数据银行(PDB)解锁了结构生物学模型(Boltz 特辑、ESM/BioHub 特辑),CELLxGENE 则为虚拟细胞模型做了同样的事情。与 PDB 一样,CELLxGENE 启发了大量 RNA 表达的 AI 模型,以至于 RNA 表达模型已成为虚拟细胞模型的代名词。Bo Wang 构建了其中最有影响力的模型之一 scGPT,它成为了 Xaira 新模型的起点。
RNA 表达 ≠ 虚拟细胞
在 CELLxGENE 上训练的模型描述了细胞类型与细胞状态之间的关系,但它们并不擅长预测如果我们改变 RNA 表达会发生什么。基因表达的变化高度相关,在大多数情况下,很难(甚至不可能)弄清因果关系。
原始正文
🔬Causal Models Need Causal Data - Xaira’s X-Cell model for Drug Discovery (Bo Wang & Ci Chu, Chief Discovery Officer & Chief AI Scientist)
Bet on information
If test loss flatlines after 1.5B parameters while training loss continues to drop as you scale, that tells you that your model is limited by the amount of information in your data.
Training on a single, smallish data set exposed an information gap: the 3.1B model falls off the scaling trend. Neither parameters nor compute will improve performance past this wall. For predicting changes to gene expression, you need more information rich data.
This is what Chu and Bo’s teams have done, and here is what ~30x the information buys you:
Now we can scale with parameters and training compute! We don’t know how much this effort costed, but we can guess that data collection experiments and infrastructure was a few tens of millions, and compute + headcount + research was a few million. The budget looks like a RL rollout budget, rather than a data rich pre-training one.
We were lucky enough to have the two central figures in this story on our podcast. Taking the lead from Ci Chu and Bo Wang, Xaira Therapeutics is betting that information rich data is the key to AI-driven drug development. Chu was recently promoted to Chief Discovery Officer and Bo to Chief AI Scientist1, underscoring just how strategic Xaira considers this bet.
Reverse engineering the human cell
If you had to figure out how a human cell works, what would you do? A good place to start might be by documenting what genes are expressed (e.g. what RNA is floating around) in different kinds of cells, in different circumstances.
That is CELLxGENE, a database of 168M cells built by Chan Zuckerberg Institute that maps each cell to a count of how many times 20K-30K genes were detected in that cell, plus detailed metadata about every cell. A ~4 trillion-entry matrix.
If the Protein Data Bank (PDB) unlocked structural biology models (Boltz Episode, ESM/BioHub Episode), CELLxGENE has done the same thing for Virtual Cell models. Like PDB, CELLxGENE has inspired a zoo of AI models of RNA expression; so much so that RNA expression models have become synonymous with Virtual Cell models. Bo Wang built one of the most influential, scGPT, that became the starting point for Xaira’s new model.
RNA expression ≠ Virtual Cell
Models trained on CELLxGENE describe the relationship between cell types and cell states, but they are not good at predicting what will happen if we make changes to RNA expression. Changes in gene expression are highly correlated, and its is difficult (impossible) to figure out what causes what in most cases.
If you could “turn the dial down” on one gene at a time, however, then you would be able to observe what is upstream and downstream of a given gene2. You could tell if A → B & C or B → A & C or B → A, C → B → … If you did this for all of the genes, then maybe you could train a model that could predict what would happen to a cell if you change a gene (e.g. with a drug or a gene edit). Or maybe you could figure out the least invasive way to change a particular gene’s expression.
X-Atlas → X-Cell
This is exactly what Chu and Bo’s teams have done. The data set is called X-Atlas and the model is called X-Cell.
In this episode, we discuss:
- Why the team abandoned autoregression for diffusion
- The CRISPR-based experiments that run millions of tests in parallel, and generate the raw data for X-Atlas and X-cell
- Generalization to real lab experiments in real human cells
- Beating the linear baseline that has outperformed previous models
- Justifying a kitchen-sink of priors, and how that stacks up vs. data and architecture
Bo also shared with us some of the (major) advantages he has as an academic vs. industry leader, and how his labs keep up with the breakneck pace of AI innovation.
Check out the full episode on YouTube, or your favorite podcasting platform!
1
These promotions happened after we recorded the episode
2
There can be cycles in the chain reaction, of course, and there can be second, third, etc. order effects (meaning things that only happen when multiple genes change at once), but the first order effects are a great place to start, and might tell us a lot of what we need to know.