[Paper Review] 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 Biology3 citations

TL;DR

This paper introduces ViTally Consistent, a 1.9-billion-parameter Vision Transformer (ViT-G/8) masked autoencoder pre-trained on 8 billion curated microscopy image crops from the Phenoprints-16M dataset. It achieves state-of-the-art performance in biological representation learning, showing a 60% improvement in linear separability of genetic perturbations and superior replicate consistency and whole-genome relationship recall, with intermediate transformer blocks yielding more biologically meaningful features than final layers.

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.

Motivation & Objective

• To develop a large-scale foundation model for cell microscopy that learns biologically meaningful and consistent representations across experimental replicates.
• To improve the linear separability of genetic perturbations and whole-genome biological relationship recall in high-throughput screening data.
• To identify optimal intermediate layers in Vision Transformers for downstream biological evaluation, rather than relying on final layers.
• To establish a scalable, biologically motivated training and probing strategy for self-supervised learning in biological imaging.
• To demonstrate that model scale and curated data significantly enhance consistency and biological relevance in representation learning.

Proposed method

• Pretraining a Vision Transformer (ViT-G/8) as a Masked Autoencoder (MAE) on over 8 billion microscopy image crops from the curated Phenoprints-16M dataset.
• Using a biologically motivated linear probing strategy to evaluate representations across all transformer blocks, selecting the best-performing intermediate layer for downstream tasks.
• Applying a curated dataset of statistically significant positive samples to improve model generalization and biological relevance.
• Employing a multi-task linear probing framework on gene group classification tasks to identify representations with strong correlation to whole-genome benchmarks.
• Training on 48,000 H100 GPU hours to achieve billion-parameter scale, enabling robust representation learning at scale.
• Validating performance using replicate consistency, whole-genome relationship recall, and linear separability metrics on large-scale HCS data.

Experimental results

Research questions

• RQ1 Can a large-scale self-supervised Vision Transformer model trained on curated microscopy data achieve superior consistency and biological relevance in representation learning?
• RQ2 Do intermediate transformer blocks in Vision Transformers yield more biologically meaningful representations than final layers for cell microscopy data?
• RQ3 How does model scale and data curation impact linear separability of genetic perturbations and recall of known biological relationships?
• RQ4 Can biologically motivated linear probing tasks effectively predict performance on large-scale whole-genome benchmarks?
• RQ5 What is the impact of training on a curated dataset of positive phenotypic samples on downstream representation quality?

Key findings

• The proposed MAE-G/8 model achieves a 60% improvement in linear separability of genetic perturbations compared to a prior ViT-L/8 MAE.
• The model achieves state-of-the-art performance on both whole-genome biological relationship recall and replicate consistency benchmarks.
• Intermediate transformer blocks consistently outperform final blocks in downstream biological evaluation, indicating they capture more relevant biological concepts.
• Performance on biologically motivated linear probing tasks is strongly correlated with performance on large-scale whole-genome benchmarks, enabling efficient model evaluation.
• Scaling training FLOPs shows a strong, consistent linear trend with improved performance across all benchmarks, confirming the scalability of the approach.
• The curated Phenoprints-16M dataset significantly improves both recall of known gene-gene relationships and embedding consistency for gene knockout perturbations.