ViTally Consistent: Scaling Biological Representation Learning for Cell Microscopy

📄 arXiv: 2411.02572v2 📥 PDF

作者: Kian Kenyon-Dean, Zitong Jerry Wang, John Urbanik, Konstantin Donhauser, Jason Hartford, Saber Saberian, Nil Sahin, Ihab Bendidi, Safiye Celik, Marta Fay, Juan Sebastian Rodriguez Vera, Imran S Haque, Oren Kraus

分类: cs.LG, cs.AI, cs.CV

发布日期: 2024-11-04 (更新: 2025-07-16)

备注: ICML 2025 main-track paper (42nd International Conference on Machine Learning). Formerly appeared as best paper runner-up at NeurIPS 2024 Foundation Models for Science Workshop (38th Conference on Neural Information Processing Systems). 18 pages, 7 figures

💡 一句话要点

ViTally Consistent：扩展细胞显微镜生物表征学习，构建大规模细胞表征基础模型

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture) 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 细胞显微镜 生物表征学习 Vision Transformer 自监督学习 药物发现

📋 核心要点

  1. 现有细胞显微图像分析模型难以一致地表征生物表型，相似生物学效应的扰动可能对应差异大的特征。
  2. 论文核心在于扩展模型规模，并结合生物学知识，优化训练数据和特征提取方式，提升表征一致性。
  3. 实验表明，提出的ViT-G/8 MAE模型在遗传扰动线性可分性上提升60%，并在多项生物学基准上取得最佳性能。

📝 摘要（中文）

大规模细胞显微镜筛选被广泛应用于药物发现和分子生物学研究，以研究数百万种化学和遗传扰动对细胞的影响。为了在下游分析中使用这些图像，我们需要能够将每个图像映射到特征空间的模型，该特征空间能够一致地表示不同的生物表型，即具有相似生物学效应的扰动具有相似的表示。本文提出了迄今为止最大的细胞显微镜数据基础模型，一个拥有19亿参数的ViT-G/8 MAE，它在超过80亿张显微镜图像裁剪上进行了训练。与之前发布的ViT-L/8 MAE相比，我们的新模型在遗传扰动的线性可分性方面提高了60%，并在全基因组生物关系召回和重复一致性基准测试中获得了最佳的整体性能。除了扩展规模之外，我们还开发了两种关键方法来提高性能：（1）在经过精心策划和多样化的数据集上进行训练；（2）使用生物学驱动的线性探测任务来搜索每个Transformer块，以找到全基因组筛选的最佳候选表示。我们发现，许多在自然图像或显微镜图像上预训练的自监督视觉Transformer，在其中间块中比在通常使用的最终块中产生更有生物学意义的显微镜图像表示。更广泛地说，我们的方法和结果为成功构建大规模生物数据基础模型提供了一个通用策略。

🔬 方法详解

问题定义：论文旨在解决细胞显微图像分析中，现有模型无法有效且一致地表征细胞生物学表型的问题。具体来说，即使是具有相似生物学效应的扰动，经过现有模型提取的特征也可能差异很大，这限制了下游分析的准确性和可靠性。现有方法的痛点在于模型规模不足，训练数据缺乏多样性，以及特征提取方式没有充分利用生物学知识。

核心思路：论文的核心思路是通过扩大模型规模，并结合生物学知识来提升细胞表型表征的一致性。具体而言，首先，通过大规模的预训练来学习通用的图像特征；其次，通过精心策划和多样化的数据集来提升模型的泛化能力；最后，通过生物学驱动的线性探测任务来选择Transformer中间层的特征，从而获得更具生物学意义的表征。

技术框架：整体框架基于Vision Transformer (ViT) 架构，采用Masked Autoencoder (MAE) 进行自监督预训练。主要流程包括：1) 数据收集与清洗：收集大规模的细胞显微图像数据，并进行质量控制和预处理；2) 模型预训练：使用MAE在大量图像数据上预训练ViT模型；3) 特征选择：使用生物学驱动的线性探测任务，在预训练模型的不同Transformer块中搜索最佳的特征表示；4) 模型评估：在多个生物学基准测试上评估模型的性能。

关键创新：论文最重要的技术创新点在于结合了大规模预训练和生物学知识来提升细胞表型表征的一致性。与现有方法相比，该方法不仅扩大了模型规模，还通过精心策划的数据集和生物学驱动的特征选择，使得模型能够更好地捕捉细胞的生物学信息。此外，论文还发现Transformer中间层的特征比最终层的特征更具生物学意义，这为后续研究提供了新的思路。

关键设计：模型采用ViT-G/8架构，拥有19亿参数。训练数据集包含超过80亿张显微镜图像裁剪。使用MAE进行自监督预训练，损失函数为均方误差（MSE）。生物学驱动的线性探测任务使用线性分类器来预测基因扰动，并选择在验证集上表现最佳的Transformer块的特征作为最终的细胞表型表征。

📊 实验亮点

实验结果表明，提出的ViT-G/8 MAE模型在多个生物学基准测试上取得了显著的性能提升。与之前的ViT-L/8 MAE模型相比，在遗传扰动的线性可分性方面提高了60%。此外，该模型在全基因组生物关系召回和重复一致性基准测试中获得了最佳的整体性能，证明了其在细胞表型表征方面的优越性。

🎯 应用场景

该研究成果可广泛应用于药物发现、分子生物学研究等领域。通过构建大规模细胞表征基础模型，可以更准确地识别和区分不同的细胞表型，从而加速药物筛选过程，并深入理解细胞的生物学机制。未来，该模型可以进一步扩展到其他生物图像领域，例如病理图像分析，为疾病诊断和治疗提供更强大的工具。

📄 摘要（原文）

Large-scale cell microscopy screens are used in drug discovery and molecular biology research to study the effects of millions of chemical and genetic perturbations on cells. To use these images in downstream analysis, we need models that can map each image into a feature space that represents diverse biological phenotypes consistently, in the sense that perturbations with similar biological effects have similar representations. In this work, we present the largest foundation model for cell microscopy data to date, a new 1.9 billion-parameter ViT-G/8 MAE trained on over 8 billion microscopy image crops. Compared to a previous published ViT-L/8 MAE, our new model achieves a 60% improvement in linear separability of genetic perturbations and obtains the best overall performance on whole-genome biological relationship recall and replicate consistency benchmarks. Beyond scaling, we developed two key methods that improve performance: (1) training on a curated and diverse dataset; and, (2) using biologically motivated linear probing tasks to search across each transformer block for the best candidate representation of whole-genome screens. We find that many self-supervised vision transformers, pretrained on either natural or microscopy images, yield significantly more biologically meaningful representations of microscopy images in their intermediate blocks than in their typically used final blocks. More broadly, our approach and results provide insights toward a general strategy for successfully building foundation models for large-scale biological data.