[论文解读] ViTally Consistent: Scaling Biological Representation Learning for Cell Microscopy

Kian Kenyon-Dean, Zitong Jerry Wang+11 more|arXiv (Cornell University)|Nov 4, 2024

Cell Image Analysis TechniquesBiochemistry, Genetics and Molecular Biology被引用 3

一句话总结

本论文提出 ViTally Consistent，一个基于 19 亿参数视觉 Transformer（ViT-G/8）的掩码自编码器，其在来自 Phenoprints-16M 数据集的 80 亿张精心筛选的显微镜图像裁剪样本上进行预训练。该模型在生物表征学习方面达到最先进性能，线性可分性提升 60%，在实验重复一致性与全基因组关系召回方面表现更优，且中间 Transformer 块提取的特征比最终层更具生物学意义。

ABSTRACT

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.

研究动机与目标

• 开发一个大规模细胞显微成像基础模型，实现跨实验重复的生物意义明确且一致的表征学习。
• 提升高通量筛选数据中基因扰动的线性可分性以及全基因组生物学关系的召回率。
• 识别视觉 Transformer 中最适合下游生物评估的中间层，而非依赖最终层。
• 建立一种可扩展的、基于生物学动机的自监督学习训练与探针策略，适用于生物成像。
• 证明模型规模与精心筛选的数据显著提升表征学习中的一致性与生物学相关性。

提出的方法

• 在来自精心筛选的 Phenoprints-16M 数据集的超过 80 亿张显微镜图像裁剪样本上，对视觉 Transformer（ViT-G/8）进行掩码自编码器（MAE）预训练。
• 采用基于生物学动机的线性探针策略，评估所有 Transformer 块的表征，并选择在下游任务中表现最佳的中间层。
• 利用统计显著的正样本集合进行数据筛选，以提升模型泛化能力与生物学相关性。
• 在基因组分组分类任务上采用多任务线性探针框架，识别与全基因组基准高度相关的表征。
• 通过 48,000 个 H100 GPU 小时的训练实现十亿参数规模，支持大规模下稳健的表征学习。
• 使用重复一致性、全基因组关系召回率与线性可分性指标，在大规模 HCS 数据上验证性能。

实验结果

研究问题

• RQ1 在精心筛选的显微镜数据上预训练的大规模自监督视觉 Transformer 模型，能否在表征学习中实现更优的一致性与生物学相关性？
• RQ2 在细胞显微成像数据中，视觉 Transformer 的中间 Transformer 块是否比最终层产生更具生物学意义的表征？
• RQ3 模型规模与数据筛选如何影响基因扰动的线性可分性以及已知生物学关系的召回率？
• RQ4 基于生物学动机的线性探针任务能否有效预测模型在大规模全基因组基准上的表现？
• RQ5 在精心筛选的表型正样本数据集上进行训练，对下游表征质量有何影响？

主要发现

• 所提出的 MAE-G/8 模型相较于先前的 ViT-L/8 MAE，在基因扰动的线性可分性方面提升 60%。
• 该模型在全基因组生物学关系召回与重复一致性基准上均达到最先进性能。
• 在下游生物评估中，中间 Transformer 块始终优于最终层，表明其捕捉了更相关的生物学概念。
• 基于生物学动机的线性探针任务表现与大规模全基因组基准表现高度相关，支持高效模型评估。
• 训练 FLOPs 的扩展在所有基准上均呈现稳定且显著的线性增长趋势，证实该方法具有良好的可扩展性。
• 精心筛选的 Phenoprints-16M 数据集显著提升了已知基因-基因关系的召回率以及基因敲除扰动的嵌入一致性。