ViTally Consistent: Scaling Biological Representation Learning for Cell Microscopy

## Abstract

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. 1 1 1 Correspondence: kian.kd@recursion.com , info@rxrx.ai

## 1 Introduction

Large-scale cell microscopy assays are used to discover previously unknown biological processes (Przybyla & Gilbert, ; Bock et al., ; Rood et al., ) and identify novel drug candidates and targets (Vincent et al., ). Labs are now able to achieve extremely high throughput by leveraging high content screening (HCS) systems that combine automated microscopy with robotic liquid handling (Boutros et al., ). Extracting meaningful features from microscopy images in large-scale screens has become increasingly difficult as this scale has increased. Public datasets like RxRx3 (Fay et al., ) and JUMP-CP (Chandrasekaran et al., ) now include millions of cellular images across 100,000s of unique chemical and genetic perturbations. In addition to limitations in expressiveness of the features that can be derived from them, traditional methods relying on customized pipelines for segmentation, feature extraction, and downstream analysis (Caicedo et al., ) struggle to handle this scale effectively  (Chandrasekaran et al., ; Carpenter et al., ).

The size and complexity of large-scale microscopy data demands image models that can extract rich biological features and do so consistently across experimental replicates, both of which are crucial for downstream biomedical applications. Rich, biologically meaningful representations reveal relationships between genes or compounds to drive the discovery of novel targets and drug candidates, while consistency in features extracted across replicates ensures that findings are reproducible and reliable for therapeutic development.

Foundation models have been developed for representing high-dimensional unstructured biological data such as protein structures (Jumper et al., ) and transcriptomics (Hao et al., ), but the scale and dimensionality of large-scale microscopy data present unique challenges for generating representations that are both biologically informative and consistent across replicates. HCS datasets are often confounded by complex noise known as batch effects (Caicedo et al., ), stemming from differences between experimental batches and biological variability. These batch effects – including natural variation in cell populations – obscure the biological effects of perturbations and make it challenging to isolate the specific effects of the perturbations applied (Yang et al., ). Overcoming these obstacles with a model capable of generating robust, biologically meaningful representations can empower HCS to systematically interrogate gene function and identify novel drug candidates (Rood et al., ).

State-of-the-art (SOTA) deep learning methods for microscopy leverage Vision Transformers (ViT) (Dosovitskiy et al., ) trained with self-supervised learning (SSL) techniques (Balestriero et al., ) to learn unbiased representations from large-scale screens (Doron et al., ; Kim et al., ; Bourriez et al., ).
Recent studies have demonstrated that ViTs trained as Masked Autoencoders (MAEs) (He et al., ) can effectively scale beyond previous approaches and outperform various supervised and smaller SSL models in capturing biologically informative representations of cell images (Kraus et al., ). However, the level of consistency found in these representations across a large number of experimental replicates was not previously reported.
Furthermore, compared to recent multi-billion parameter transformers developed for natural images (Dehghani et al., ) and natural language (Llama3, ), model scale in microscopy lags behind (Kraus et al., ; Chen et al., ) despite the existence of massive datasets.

In this work, we developed the largest foundation model to date for cell microscopy images, achieving SOTA results in both replicate consistency and biological recall of known gene-gene relationships. Specifically our work offers the following contributions:

- •
We demonstrate that training on a curated microscopy dataset of statistically significant positive samples, named Phenoprints-16M, improves both recall of known gene-gene relationships and consistency of embeddings for gene knockout perturbations (Figure 1A). We describe components of this curation strategy that can be generalized to other scientific datasets (§ 3.1).
- •
We present a new foundation model, MAE-G/8, a 1.86 billion parameter ViT-G/8 MAE trained on Phenoprints-16M over 48,000 H100 GPU hours on more than 8 billion samples from the curated dataset (Figure 1A, § 3.2).
- •
We propose new set of biological linear probing tasks to evaluate representations learned by intermediate ViTs blocks for microscopy data (§ 4).
Performance on these linear probing tasks are strongly correlated with performance on important whole-genome scale evaluation metrics while requiring significantly less resources to compute (Figure 4).
- •
We find that using intermediate layers leads to better performance on these downstream whole-genome benchmarks at a lower computational inference cost, across SSL ViTs trained on microscopy or natural images. By taking advantage of our linear probing proxy task, we are able to cheaply find the best performing intermediate block (Eq. 1).

Our results indicate that the biological scaling properties first identified by Kraus et al. () extend to the multi-billion parameter regime (§ A.9). We show that our MAE-G/8 model produces a nearly 60% more phenotypically linearly separable latent space compared to previous approaches using the final block of MAE-L/8 (Figure 4), correlating with significant improvements in both recall and replicate consistency when benchmarking across the whole genome (Figure 1B).

Figure: Figure 1: (A) Overview of performance gain from different foundation model pretraining and inference strategies. (B) Example whole-genome results for replicate consistency and biological relationship recall on StringDB for models trained with different combinations of strategies, by model name and dataset (left to right): MAE-L/8 (RPI-93M, block $b=24$), MAE-L/8 trimmed to block $b^{*}=15$, MAE-L/8 (Phenoprints-16M, block $b^{*}=20$), MAE-G/8 (Phenoprints-16M, $b^{*}=38$), where $b^{*}$ is the optimal block according to linear probes as defined in Equation 1.
Refer to caption: x1.png

## 2 Related work

#### Evaluating Representations for Drug Discovery.

Evaluating the quality of biological representation learning methods for drug discovery remains challenging, as ground truth data is sparse, noisy, biased to well-studied diseases and pathways, and poorly annotated. Metrics have been proposed that use mean average precision (Kalinin et al., ) or AUC ROC (Sivanandan et al., ) to assesses how similar related samples are represented, including replicates of the same perturbation or different perturbations with similar annotated biological activities. Recently, Celik et al. () introduced terminology for describing perturbative “maps of biology”, in which representations of perturbations in HCS data can be placed in unified, relatable embedding spaces allowing for the generation of genome-scale sets of pairwise comparisons. Here we leverage the biological relationship recall benchmark proposed by Celik et al. (), which assess how well known relationships between pairs of perturbations are recalled among the most similar or dissimilar embeddings. Computing reliable versions of these relationship benchmarks with HCS data is particularly expensive as they require genome-wide embeddings to be inferred for hundreds of millions of image crops from the genome-wide RxRx3 microscopy screen (Fay et al., ).

Figure: Figure 2: Samples for subset of groups in Anax 40-class functional gene group classification task.
Refer to caption: x2.png

#### Dataset Curation for Foundation Models.

Dataset curation is crucial for enhancing the efficiency of foundation models, especially in large-scale contexts. Usual approaches to dataset construction are inspired by the image retrieval community (Weinzaepfel et al., ; Radenović et al., ; Berman et al., ). Existing methods often utilize pre-trained models for filtering and pruning, such as vision-language models to discard irrelevant pairs (Schuhmann et al., ), semantic deduplication to remove redundancy (Abbas et al., ), and prototypicality-based approaches to retain representative data (Sorscher et al., ). However, these techniques are less effective for HCS, where redundancy, variability, and subtle morphological differences make conventional filtering challenging. Our work addresses these limitations by building on Celik et al. ()’s perturbation consistency framework to curate a balanced dataset of images across semantic classes, which is vital for effective learning under the masked objectives (Zhang et al., ).

#### Layer-wise Analysis of Deep Neural networks.

Recent work suggests that intermediate layers (or, blocks) in large ViTs may achieve superior performance on certain linear probing tasks compared to the final encoder layer (Evci et al., ; Dehghani et al., ). For example, Alkin et al. () reported that intermediate layers in large MAE-ViTs (ViT-L, ViT-H) have superior ImageNet-1K k-NN accuracy. They attributed this property to the later encoder layers becoming more optimized for the reconstruction task.

## 3 Vision Transformers for Microscopy Images

We train and evaluate various vision transformers (ViTs, Table 1) as encoders to extract feature embeddings from $256\times 256\times 6$ (HxWxC) microscopy image crops (Figure 2).

### 3.1 Training Dataset Curation

Many academic and industry labs have adopted the Cell Painting imaging protocol (Bray et al., ), which multiplexes fluorescent dyes to reveal eight broadly relevant cellular components. The datasets used here contain a six-channel implementation of Cell Painting (Figure 2), as well as brightfield images, spanning 100,000s of chemical and genetic perturbations applied to dozens of cell types (Kraus et al., ).
In these datasets, cells that look like unperturbed cells tend to be very over-represented because many perturbations do no induce a morphological change. Some morphological changes are also far more common (e.g. many perturbations will kill cells, resulting in a relatively high proportion of dead cell morphological phenotype). This results in significant imbalance in the morphological phenotypes that the models learn to reconstruct.

To address this, we constructed an aggressively curated training dataset (§ A.1).
To learn an initial representation, we began by reproducing the MAE-L/8 model of Kraus et al. () on a dataset of similar size consisting of 93 million HCS images.
Using this representation, we first filtered perturbations that did not induce consistent morphological changes to cells. To perform this filtering, we utilized Celik et al. ()’s non-parametric perturbation consistency test (§ A.3) Typical Variation Normalization (Ando et al., ; Kraus et al., ). This test was applied within each experiment for computational efficiency, and we restricted the analysis to wells containing single perturbations.
This consistency was computed for CRISPR guides, siRNAs, and particular concentrations of small molecules across replicates of the same perturbation. P-values were computed for each gene and each (perturbation, concentration) pair. When multiple experiments existed for the same condition, we combined p-values using the Cauchy Combination test (Liu & Xie, ).

We repeated this procedure with a weakly supervised learning (WSL) model trained on RxRx1 (Sypetkowski et al., ) and filtered to perturbations where any condition had a p-value $<0.01$ in either the MAE-L/8 or WSL model.
This process reduced our original dataset of 93M samples to 16M, which we refer to as Phenoprints-16M.
While some redundancy remains when distinct perturbations have the same effect, the proportion of samples with that differ from negative controls increased substantially with little decrease in overall diversity. We believe that iteratively repeating this process with the best models from previous iterations to guide data selection for subsequent models may be a viable strategy.

### 3.2 Models

**Table 1: Overview of vision transformer (ViT) encoders used and evaluated in this work.**
| Model Name | Parameters | Blocks | Model Dim | Pretraining Data |
| --- | --- | --- | --- | --- |
| Baselines |  |  |  |  |
| Untrained ViT-S/16 | 25M | 12 | 384 | N/A |
| Dino-V2 ViT-S/14 | 25M | 12 | 384 | Natural images |
| Dino-V2 ViT-L/14 | 307M | 24 | 1024 | Natural images |
| Dino-V2 ViT-G/14 | 1,100M | 40 | 1536 | Natural images |
| ViT-L/16 WSL | 307M | 24 | 1024 | Imagenet-21k |
| ViT-L/16 MAE | 307M | 24 | 1024 | Imagenet-21k |
| MAEs for microscopy |  |  |  |  |
| CA-MAE-S/16 | 25M | 12 | 384 | RxRx3 |
| MAE-L/8 | 307M | 24 | 1024 | RPI-93M |
| MAE-L/8 | 307M | 24 | 1024 | Phenoprints-16M |
| MAE-G/8 | 1,860M | 48 | 1664 | Phenoprints-16M |

#### Baselines.

We compare to several non-finetuned baseline ViT image encoders: three different Dino-v2 backbones (Oquab et al., ) (with 4 register tokens (Darcet et al., )) trained on a curated non-biological natural image dataset; a weakly supervised (WSL) classifier ViT-L/16 trained on Imagenet-21k (Ridnik et al., ); a MAE ViT-L/16 trained on Imagenet-21k (He et al., ); and an untrained ViT-S/16. Preliminary investigations found that channel-wise self-standardization worked best as the image normalization preprocessing for these baselines, and that the class token was slightly better than the global pool of the patch tokens (except for MAE). Convolutional weights in the patch embedding layer were repeated to embed 6 channel images when using models trained on RGB datasets (Wightman, ).

#### Prior work.

Our primary point of comparison is with respect to the best pretrained foundation model presented by Kraus et al. (), the MAE-ViT-L/8+ trained on RPI-93Mrained for approximately 40 epochs, learning from over 3.5 billion image crops, using the L2 mean squared error loss function plus an additional Fourier domain reconstruction loss term.

#### CA-MAE-S/16 trained on RxRx3.

We trained a new channel-agnostic MAE (Kraus et al., ) ViT-S/16 on the RxRx3 dataset (Fay et al., ) for 100 epochs. Channel-agnostic ViTs tokenize each image channel separately with shared patch embedding weights and leverage the dynamic sequence length of transformers with repeated positional encodings to train ViTs that can process images with varying numbers of channels (Bao et al., ; Bourriez et al., ; Kraus et al., ). Kraus et al. () demonstrate that the large MAEs with 8x8 patch size perform either better or the same as the 16x16 channel-agnostic variants for consistently 6-channel data, so we opted to train standard MAEs for the following two new models since they require fewer tokens at inference time.

#### MAE-L/8 trained on Phenoprints-16M.

Holding the model backbone constant compared to the MAE-ViT-L/8 by Kraus et al. (), we assess the impact of our curated dataset in contrast to the 93M dataset by training a new ViT-L/8 MAE for 500 epochs on Phenoprints-16M.

#### MAE-G/8 trained on Phenoprints-16M.

Holding the dataset constant compared to MAE-L/8 above, we assess the impact of increased model scale in terms of parameters by training a new ViT-Gigantic MAE with nearly 1.9 billion parameters for 500 epochs on Phenoprints-16M. Training this model required 256 H100 GPUs running in parallel for over 1 week. See § A.2 for other hyperparameter settings we used for model training.

Figure: (a) RxRx1 siRNA knockdown classification.
Refer to caption: x3.png

## 4 Linear probing representation learning across ViT blocks

Figure: Figure 4: Correlations between validation set linear probing (Figure 3) on Anax and RxRx1 for best and last blocks (Eq. 1) compared to downstream whole-genome benchmarks (Table 2) for biological relationship recall on StringDB at 0.05-0.95 threshold and replicate consistency KS statistic.
Refer to caption: x5.png

We improve the quality of our learned image representations by leveraging previous findings that suggest intermediate blocks within an encoder can provide better representation compared to the final block (Alkin et al., ). Unfortunately, it is infeasible to search for the best block by simply performing whole-genome evaluation on each block of a large model because the evaluation is extremely time-consuming and resource intensive. For example, evaluating the final block of MAE-G/8 required 4,000 L4 GPU hours just for inference (§ 5). We demonstrate that using block-wise linear probes provides insights into the quality of biological features extracted by these models in their intermediate blocks, allowing us to trim the model to an earlier block to both reduce inference costs and improve representation quality.

Our block-wise search consists of training a logistic regression model (linear probe) on the output features of each transformer block to predict either the gene that was perturbed or the functional group that the gene belongs to, and test performance on held-out experiments (§ A.4). We define the optimal block $b^{*}$ for a probing task as the block whose output features achieve the highest test balanced accuracy when trained on the probing task, across all $N$ blocks of the encoder,

$$ $b^{*}=\operatorname*{arg\,max}_{b\in\{1,2,\dots,N\}}\text{BalancedAccuracy}( \mathbf{z}^{(b)}),$ (1) $$

where $\mathbf{z}^{(b)}$ are output features from block $b$ of a ViT. Performance on our linear probing tasks can be viewed as a measure of linear separability of a feature space across experimental batches.

#### RxRx1 1139-class siRNA genetic perturbation classification.

We expect high quality representations of cell images to generate similar embeddings for cells with the same perturbation, hence a simple linear probe should be able to predict gene perturbation from these representation reasonably well. We train linear probes on the publicly-available RxRx1 dataset in Sypetkowski et al. () which consists of 125,510 high-resolution fluorescence microscopy images of human cells under 1,138 siRNA-induced gene knockdowns (plus unperturbed controls) across four cell types (HEPG2, HUVEC, U2OS, RPE). These gene knockdowns produce strong phenotypes which makes the prediction task more feasible.

We found that, for MAE-G/8 , the best features came from intermediate block $b^{*}=38$ (out of $48$) of the encoder, achieving a balanced accuracy ($0.51$) that is $8.5\%$ greater compared to its final block’s output features (3(a), left). Additionally, these features achieved $60\%$ greater accuracy than the typically used final block of MAE-L/8+ (Kraus et al., ). We observed similar trends for ViT models pretrained on natural images. For example, DINO-G/14 and ViT-L/16 MAE trained on non-biological natural image data have their best features at blocks that are positioned within the first half of the encoder. For ViT-L/16 MAE, the performance of the best block is $27\%$ higher compared to its final block output features that are typically used for downstream tasks. The higher performance observed for intermediate blocks does not appear to be an intrinsic feature of the ViT architecture as an untrained ViT did not exhibit such a parabolic trend (3(a), right).

#### Anax 40-class functional gene group classification.

Biologically meaningful representation of microscopy images of genetically perturbed cells should capture functional relationships between genes, hence a simple linear probe should be able to predict functional gene groups when trained on these representations.
We curated a small subset of 80,000 wells from RxRx3 (Fay et al., ) to evaluate linear probes on functional group prediction. We also evaluated similar whole genome knockout screens with ARPE-19 and an additional population of HUVEC cells with soluble TNF-$\alpha$ added to all wells.
We manually curated Anax, a set of 40 functionally-diverse gene groups containing 348 genes, with details provided in (§ A.8). Examples of groups include major protein complexes (e.g. proteasome, ribosome-small/large), metabolic pathways (e.g. Krebs cycle) and signaling pathways (e.g. calcium signaling) (Figure 2). These groups span broad biological processes that are conserved across cell types – linear separability of these groups would likely indicate that representations are biologically meaningful regardless of cell type.

As shown in 3(b), MAE-G/8 significantly outperforms other models in Anax group linear probe classification. The best representations once again are obtained from an intermediate block, achieving a balanced accuracy ($0.32$) that is $5\%$ greater compared to its final block’s output features. We observed similar trends for ViT models pretrained on natural images and representations computed from microscopy images of other cell types/conditions (§ A.5, Figure 6).

In Figure 4, we observe that performance on this novel linear probing task correlates strongly with downstream whole-genome benchmarks across all models (Table 2), whether they are trained on microscopy data or natural images, achieving an overall rank correlation $\rho=0.97$ with whole-genome StringDB recall and $\rho=0.91$ with whole-genome replicate consistency. This strong correlation is crucial as it allows us to trim our model to the block with the best linear probe performance as a way to improve the quality of our representations for the whole-genome (Table 2).

## 5 Whole-genome benchmarking

**Table 2: Multivariate known biological relationship recall and univariate replicate consistency benchmarks by model, encoding block $b$, benchmark database, and test statistic. The trimmed models used linear probes to select an earlier block as the feature encoder (Fig. 3). Results are computed over all whole-genome CRISPR knockout perturbation images in RxRx3, after applying TVN and chromosome arm bias correction. For relationship recall, we report results over four databases (React stands for Reactome-PPI (Gillespie et al., )). Best overall result is in bold.**
| Model backbone | $b$ | CORUM | hu.MAP | React | StringDB | KS | CVM |
| --- | --- | --- | --- | --- | --- | --- | --- |
| Baseline ViTs |  |  |  |  |  |  |  |
| ViT-S/16, Untrained | $12$ | .45 | .34 | .205 | .36 | .30 | 4.3 |
| ViT-S/14, Dino-V2 | $12$ | .48 | .345 | .20 | .38 | .34 | 5.6 |
| trimmed | $5$ | .51 | .36 | .21 | .40 | .35 | 6.0 |
| ViT-L/16, ImageNet WSL | $24$ | .52 | .35 | .21 | .39 | .34 | 5.5 |
| ViT-L/14, Dino-V2 | $24$ | .49 | .34 | .21 | .38 | .34 | 5.3 |
| trimmed | $12$ | .55 | .37 | .22 | .41 | .36 | 5.9 |
| ViT-L/16, ImageNet MAE | $24$ | .53 | .355 | .215 | .40 | .34 | 5.1 |
| trimmed | $11$ | .53 | .36 | .22 | .40 | .35 | 5.8 |
| ViT-G/14, Dino-V2 | $40$ | .44 | .31 | .20 | .35 | .29 | 3.8 |
| trimmed | $16$ | .53 | .35 | .22 | .40 | .33 | 5.2 |
| MAEs for microscopy |  |  |  |  |  |  |  |
| CA-MAE-S/16 , RxRx3 | $12$ | .55 | .37 | .23 | .43 | .47 | 10.4 |
| MAE-L/8 , RPI-93M | $24$ | .61 | .43 | .25 | .47 | .52 | 12.3 |
| trimmed | $15$ | .60 | .43 | .255 | .475 | .57 | 15.2 |
| MAE-L/8 , PP-16M | $24$ | .60 | .43 | .255 | .48 | .59 | 16.2 |
| trimmed | $20$ | .60 | .435 | .26 | .48 | .59 | 16.2 |
| MAE-G/8 , PP-16M | $48$ | .62 | .44 | .26 | .49 | .60 | 16.4 |
| trimmed | $38$ | .615 | .44 | .26 | .49 | .63 | 18.2 |

Table 2 presents our benchmarks computed across the whole-genome. These evaluate the genomic representations obtained for each model by aggregating millions of embeddings of cell images spanning $>$100,000 of genetic knockout perturbations (17,063 genes $\times$ 6 single guide RNAs each) on HUVEC cells from RxRx3 (Fay et al., ). Computing these benchmarks for HCS screens typically requires inferring 140 million crops from the genome-wide RxRx3 microscopy screen (Kraus et al., ) (64 tiled crops per each of the 2.2 million wells), but, to reduce compute costs, we discard the outer ring of crops, leaving the 36 center non-edge crops for each well. This requires 80 million forward passes to comprehensively evaluate a new encoder.
After inference, we use typical variation normalization (Ando et al., ) and chromosome arm bias correction (Lazar et al., ) to post-process the embeddings and aggregate them to the gene-level.

We present the multivariate biological relationship recall benchmarks proposed by Celik et al. () and originally evaluated for MAEs by Kraus et al. (, ). These metrics evaluate how many annotated pair-wise relationships are recalled from public databases (CORUM, hu.MAP, Reactome-PPI, StringDB) in the extremities of a ranked list of cosine similarities of all pair-wise post-processed embeddings (details in § A.6).
To ensure embeddings represent technical replicates of perturbations consistently, we also evaluate model performance on replicate consistency based on the experimental design used in the RxRx3 dataset. Specifically, we compare the similarity of the embedding for corresponding wells across different experiments via a non-parametric statistical test.
The test statistic measures the difference between the perturbation replicates’ similarity distribution and an empirical null distribution, with larger values indicating greater consistency (details in § A.7).

In order to compare models, we summarize the resulting statistics over all technical replicates in RxRx3 by taking their median, as reported in columns KS and CVM in Table 2, and visualized in Figure 5. MAEs pretrained on microscopy data show improved performance compared the baseline models. Furthermore, training on the Phenoprints-16M dataset improves the performance of the MAEs significantly and trimmed MAE-G/8 achieves the best overall performance.

Our linear probing analysis (Figure 3) allowed us to trim our models to better encoding blocks. Comparing models on their best respective blocks, MAE-G/8 improves on MAE-L/8 with a 16% improvement in Anax functional gene group classification (.27$\rightarrow$.31) and a 24% improvement in RxRx1 perturbation classification (.41$\rightarrow$.51). Compared to the best published result for whole-genome benchmarks (MAE-L/8 trained on RPI-93M (Kraus et al., )), MAE-G/8 obtains a 20% improvement in replicate consistency KS (.52$\rightarrow$.63) and 4.3% improvement in StringDB recall (.472$\rightarrow$.492). When using our linear probes to select outputs from block $b^{*}=15$ (Equation 1) from that MAE-L/8, the gain for MAE-G/8 changes to 9.2% in KS and 3.5% in StringDB recall.

Similarly, linear probing to select optimal ViT blocks led to significant improvements even when applied to frozen Dino-V2 based models pretrained on natural images. Dino-V2 ViT-G obtains a nearly 20% improvement CORUM recall (.44$\rightarrow$.53) by using the embeddings extracted at $b^{*}=16$ (chosen by linear probes) rather than the final embedding from $b=40$ (which performs worse than a random untrained ViT-S). Dino-V2 ViT-S also observes improvements by using $b^{*}=5$ rather than $b=12$ and outperforms Dino-V2 ViT-G in replicate consistency.

Figure: Figure 5: Replicate consistency whole-genome results between perturbations from cosine similarity distributions on RxRx3 post-TVN. Comparing baseline Dino-V2 ViT-G/14 at its typically used final block (left) versus MAE-G/8 at the best block found via linear probing (right).
Refer to caption: x6.png

## 6 Discussion and Conclusions

This work demonstrates that:
(1) within the context of biological imaging, trimming many ViTs to an earlier block leads to stronger biological linearity and improved performance on downstream tasks in addition to cheaper inference costs (Figure 3); (2) linear probing performance on a subset of genetic perturbations correlates strongly with downstream performance on whole-genome benchmarks and can be used to optimize which block is selected for representing the whole-genome (Figure 4); (3) the most scaled model, MAE-G/8 , obtains the overall best performance across all benchmarks and linear probes, providing further evidence for the scaling hypothesis in biological image data (Table 2).
This demonstrates that intentionally scaling training compute and parameters of SSL models for microscopy can benefit downstream biological relationship recall, whole-genome replicate consistency, and biological linear separability on smaller datasets (see § A.9 for scaling plots).

More broadly, this work proposes a reusable recipe for training and extracting optimal representations from fully self-supervised models trained on experimental data. The pattern we use can be applied to other domains that contain data from repeated experiments but without accurate ground truth labels. Specifically, we recommend: (1) curating the training set by identifying diverse sets of samples that are represented consistently, e.g., by using a pre-existing model to select such samples; (2) training a scaled transformer-based model using a self-supervised learning technique, such as masked autoencoding; and, (3) evaluating the performance of the trained transformer at every block to identify the optimal layer for representing the data.

## Appendix A Appendix

### A.1 Training dataset curation details

In order to produce Phenoprint-16M, we curated 93M using the following steps:

- 1.
Filtering out data that did not pass data quality filters related to the focus of the image, quantity of dead cells, assay conditions, and presence of strong anomalous imaging artifacts.
- 2.
Filtering out data with missing information about the perturbations applied, data with more than 3 perturbations applied, and data of unusual size (in the image dimension or number of channels).
- 3.
Filtering out perturbation conditions that had been in less than 3 distinct experiments or 20 distinct wells so as to capture a variety of batch effects and have a broad sample of positives per class.
- 4.
Under-sampling perturbation conditions that were clearly over-represented in the dataset. Our experiment designs contain positive controls, negative controls, and wells without perturbation within each experiment. At this step, we keep 10% of positive controls and wells without any perturbation, 30% of negative controls, and all other perturbation conditions.
- 5.
Filtering out wells where none of the perturbation conditions had a phenoprint (§A.3) (across different map types) in any experiment it had been run in.

### A.2 Training hyperparameters

**Table 3: Training hyperparameters for the new models presented in this work. Each used a one-cycle cosine learning rate decay schedule with 10% warm-up using the Lion optimizer from Chen et al. () with betas (0.9, 0.95) and weight decay of 0.05, with additional ViT settings such as LayerScale as proposed by Dehghani et al. (). ^∗Note that MAE-G/8 had multiple restarts during training due to challenges associated with massive model training on large-scale shared distributed compute clusters.**
| Hyperparameter | CA-MAE-S/16 | MAE-L/8 | MAE-G/8 |
| --- | --- | --- | --- |
| Vision transformer backbone | ViT-S | ViT-L | ViT-G (Zhai et al., ) |
| Pretraining Data | RxRx3 | Phenoprints-16M | Phenoprints-16M |
| Training epochs | 100 | 500 | 500^∗ |
| Learning rate | 1e-4 | 3e-5 | 3e-5 |
| Global batch size | 2048 | 16384 | 8192 |
| Stochastic depth | 0.1 | 0.3 | 0.6 |
| # GPUs | 16 A100s | 128 H100s | 256 H100s |
| # GPU-hours | 400 | 15,360 | 48,000 |

Table 3 provides the hyperparameters used for training the new vision transformers presented in this work. Each model was trained using a 75% mask ratio and the standard decoder architecture for MAEs (He et al., ). Each model was trained with the standard L2 MAE loss and the Fourier-space loss function implemented by Kraus et al. () with a weight of $\alpha=0.01$. We note, however, that the details presented by Kraus et al. () do not precisely correspond with the implementation provided in their Github repository; when reshaping the tokens to a shape compatible with the 2D Fourier transform, the permute operation resulted in adjacent pixels being from different channels of the input, resulting in the high frequency components of the loss being a function of the relationships between input channels. An initial investigation with a ViT-L/8 showed that changing the implementation to the one described in the paper did not dramatically change probing results. As such, we used the implementation as-is and leave additional analysis of loss function design for MAEs to future work.

### A.3 Perturbation Consistency

In order to assess the consistency of the induced morphology on the cells by the perturbations, we used a non-parametric perturbation consistency test similar to the one introduced in Celik et al. (). Let $x_{g,1},x_{g,2},\cdots,x_{g,n}$ be the embeddings for replicates of perturbation $x_{g}$ on experiment (batch) $e$. As the test statistic for perturbation consistency, $\bar{s}_{g}^{e}$ is defined as the mean of the cosine similarities across all pairs of replicates of $x_{g}$.

$$ $\bar{s}_{g}^{e}=\frac{1}{n^{2}}\sum_{i=1}^{n}\sum_{j=1}^{n}\frac{\langle x_{g, i},x_{g,j}\rangle}{||x_{g,i}||||x_{g,j}||}.$ (2) $$

where $\langle.\rangle$ and $||.||$ denote dot product and $L_{2}$ norm.

Statistical significance of $\bar{s}_{g}^{e}$ is assessed using a permutation test comparing it against an empirical null distribution generated using the same statistic for a set of randomly selected perturbations in experiment $e$, $\{\bar{s}^{\prime}_{1},\cdots,\bar{s}^{\prime}_{K}\}$. The p-value for $\bar{s}_{g}^{e}$ is computed as follows

$$ $p_{g}=\frac{\max\biggl{\{}\#\{\bar{s}^{\prime}_{k}\geq\bar{s}_{g}^{e}\},1 \biggl{\}}}{K}.$ (3) $$

When multiple experiments existed for the same perturbation, we combined p-values using the Cauchy Combination test (Liu & Xie, ).

### A.4 Training linear probes

In this section, we provide details about the training process and preprocessing steps used in our logistic regression models. These models were trained on output features derived from various Vision Transformer (ViT) blocks.

The data was split by experiments, ensuring that the test data originated from experiments distinct from those used for training. This approach helps to validate the generalization performance of our models across different experimental conditions.

For both RxRx1 gene prediction and Anax group prediction, we apply StandardScaler from the scikit-learn library as the only preprocessing step to standardize the features prior to training linear probes. StandardScaler transformation was fitted on data from the train split.
We trained the logistic regression models using scikit-learn’s LogisticRegression class. The following parameters and settings were used during model optimization:

- •
Solver: lbfgs
- •
Maximum Iterations: 2000
- •
Class Weight: balanced

For RxRx1 gene prediction, we trained logistic regression models to predict one of $1139$ possible perturbation labels ($1138$ genetic perturbation and non-perturbed control). For Anax group prediction, we trained logistic regression models to predict one of $40$ possible function group labels (§ A.8). We report the balanced test accuracy as the main evaluation metric for all linear probing experiments.

### A.5 Anax classification for other cell lines/treatment conditions: ARPE19 and HUVEC with TNF-alpha background

We performed linear probing on imaging data obtained for a retinal pigment epithelia (RPE) cell line, ARPE19, and HUVEC cells treated with an inflammatory cytokine, TNF$\alpha$. We similarly observed that intermediate blocks often have the most linearly separate features compared to the final block.

Figure: Figure 6: Layerwise validation set linear probe performance on Anax functional gene group classification beyond RxRx3: CRISPR knockouts in the ARPE-19 immortalized epithelial cell-line (left), and in HUVEC cells with a TNF-$\alpha$ background (right).
Refer to caption: x7.png

### A.6 Biological Relationship Recall

A valuable use of large-scale HCS experiments is to perform large-scale inference of biological relationships between genetic perturbations.
We evaluate each model’s ability to recall known relationships by using the biological relationship recall benchmark described in Celik et al. (). First, we correct for batch effects using Typical Variation Normalization (TVN) (Ando et al., ), and also correct for possible chromosome arm biases known to exist in CRISPR-Cas9 HCS data (Lazar et al., ). To infer biological relationships, we compute the aggregate embedding of each perturbation by taking the spherical mean over its replicate embeddings across experiments.
We use the cosine similarity of a pair of perturbation representations as the relationship metric, setting the origin of the space to the mean of negative controls.
We compare these similarities with the relationships found in the following public databases: CORUM (Giurgiu et al., ), hu.MAP (Drew et al., ), Reactome (Gillespie et al., ), and StringDB Szklarczyk et al. () (with $>$95% combined score). Table 2 reports the recall of known relationships amongst the top and bottom 5% of all cosine similarities between CRISPR knockout representations in RxRx3 (Fay et al., ).

### A.7 Replicate Consistency

In order to assess the reproducibility of the perturbations across their technical replicates, we compare the distributions of the similarities for same perturbations across replicates against an empirical null distribution.
Specifically, for technical replicate experiments $e^{i}_{a}$ and $e^{i}_{b}$, we calculate the cosine similarity between the embeddings of perturbation $x_{j}$ in them, denoted as $s^{x_{j}}$.The query distribution $q^{e_{i}}$ is constructed by computing the cosine similarities for all perturbations that have a matching well on experiments $e^{i}_{a}$ and $e^{i}_{b}$. An empirical null distribution of identical cardinality is created by computing cosine similarity, $r^{x_{k},x_{l}}$, between random pairs from $e^{i}_{a}$ and $e^{i}_{b}$ such that no pair corresponds to the same perturbation, $p^{e_{i}}_{0}$. Using non-parametric statistical tests, namely Kolmogorov-Smirnov (KS) and Cramer Von-Mises (CVM), we can evaluate the hypothesis that $q^{e_{i}}$ and $p^{e_{i}}_{0}$ are drawn from the same distribution. Formally, let $Q^{e_{i}}(x)$ and $P^{e_{i}}_{0}(x)$ be the cumulative distribution functions for $q^{e_{i}}$ and $p^{e_{i}}_{0}$ respectively, then the KS statistic for the two-sample case of technical replicate experiments $e^{i}_{a}$ and $e^{i}_{b}$ is defined as:

$$ $\text{KS}^{e_{i}}=\text{sup}_{x}|Q^{e_{i}}(x)-P^{e_{i}}_{0}(x)|.$ (4) $$

The Cramér–von Mises test statistic (CVM) for experiments $e^{i}_{a}$ and $e^{i}_{b}$ is computed as:

$$ $\text{CVM}^{e_{i}}=\frac{1}{2N^{2}}\sum_{m=1}^{N}\biggl{[}(r_{m}-m)^{2}+(s_{m} -m)^{2}\biggl{]}-\frac{4N^{2}-1}{12N}.$ (5) $$

where $N$ is the cardinality of $q^{e_{i}}$ and $p^{e_{i}}_{0}$ and $s_{m}$ and $r_{m}$ are ranks of similarities $s^{x_{j}}$ and $r^{x_{k},x_{l}}$ in the combined distribution of $q^{e_{i}}$ and $p^{e_{i}}_{0}$ when ordered. In order compare models, we use the median of $\text{CVM}^{e_{i}}$ and $\text{KS}^{e_{i}}$ over all technical replicate experiment pairs $e_{i}$.

Since the pairs are randomly selected for $p^{e_{i}}_{0}$, the embeddings would be mostly orthogonal thus the distribution would be centered around 0, similar to what Figure 5 illustrates. Given that not all CRISPR knockouts would induce a morphological change in the cells, it’s plausible for distribution $q^{e_{i}}$ to exhibit a peak around 0. As the model approaches the precision of an oracle, we would anticipate the mass situated around this peak to shift towards higher cosine similarity values.

### A.8 Anax Group Prediction Details

The Anax probing task introduced in this paper is intended to balance capturing a diverse range of biology that is broadly conserved between cell types with a reduced cost of execution. The name “Anax” is a reference to Anaximander, the 6th century B.C. philosopher credited with making the first world map.

In curating these genes, we analyzed the sources listed in § A.6 as well as internal gene expression data to produce “functional” groups corresponding to biological processes, cellular components, and molecular functions. Not all genes within each group are expected to have the same knockout phenotype, but are classified by humans as having related function – linear separability of these genes would indicate that a model has learned similar concepts to those deemed significant by biologists.

The gene groups we use for the 40-class Anax group classification task (§ A.4) are listed in Table 4.

Table: Table 4: Anax groups and their associated genes. This table presents a comprehensive list of gene groups and their corresponding genes.

### A.9 Correlation between model scale and benchmark results

In Figure 7 we show the correlations between training FLOps (floating point operations) and downstream results. Over all benchmarks we observe a very strong consistent linear trend where scaling training FLOps improves overall pwerformance. This work provides the next log step in scale as we enter into the billion-parameter model regime with MAE-G/8. These results therefore provide additional evidence that the trend initially discovered by Kraus et al. () between FLOps and relationship recall actually extends both to billion-parameter models and even moreso for other biologically meaningful benchmarks pertaining to linear probes on small experiments and to replicate consistency on the whole-genome.

Figure: Figure 7: Relationship between FLOPs and benchmark evaluation results for the six whole-genome tasks (Table 2) and the two linear probing tasks (Figure 3).
Refer to caption: x9.png

## Peer Reviews

confidence: 4
contribution: 2
initial_rating: 3
presentation: 3
questions: - Please address the weakness mentioned above, especially regarding the model evaluation and model training algorithm.
- Are the data, model, and training code going to be publicly available?
soundness: 3
strengths: - Cell microscopy analysis is an important problem and has a large impact on the future of computational pathology and bioinformatics research.
- The authors curated a dataset to only include relevant data using biological informed methods, which may be a worthwhile contribution.
- The authors performed careful tuning of model architectural and training details.
summary: The authors presented a new 1.9 billion-parameter ViT-G/8 MAE model, trained on over 8 billion microscopy image crops. The authors demonstrated performance boost over a smaller ViT-L/8 model in various downstream tasks.
weaknesses: - The methodological contribution is limited. The authors applies MAE, and evaluated on various downstream tasks. It is also known that newer algorithms such as iBOT or DINOv2 that are more effective in learning visual representations. The authors should benchmark these methods as well.
- The author’s model performance boost over a much smaller MAE-L/8 trained on RPI-93M model is very marginal, and the benchmarks do not have error bars.

---

confidence: 5
contribution: 2
initial_rating: 5
presentation: 3
questions: * Dataset Availability: Will the 384-gene manually curated Anax dataset and curated Phenoprints-16M datasets be released as part of this submission?
* Model Availability: Will the ViT-G/8 model weights be released as part of this submission?
soundness: 2
strengths: * This paper explores the following important aspects of training and scaling large foundational representation models for microscopy images such as 
  - dataset curation, 
  - proxy linear probing tasks for reasonable evaluation of models during training that generalizes to downstream performance and 
  - selection of optimal intermediate layer representation for specific downstream tasks
* The paper shares details of a manually curated dataset of 40 functionally-diverse gene groups containing 348 genes which provide a useful dataset for linear probing proxy tasks and evaluation of the models
summary: In this work, the authors present a large 1.9B parameter masked auto-encoder model (MAE-G/8) trained on over 8 billion microscopy image crops. In addition to scaling dataset and model size, they also explore data curation strategies and evaluation on proxy tasks for searching for the optimal representation for biologically relevant tasks from the transformer architecture. They introduce a data curation strategy that filters perturbations that do not have a statistically significant consistent morphological change in cells compared to negative controls on smaller pre-trained networks (WSL and MAE-L/8). 

Through their experiments, they demonstrate that intermediate layer representations capture 60% more linearly separable latent space which is also more biologically informative (improvements in both recall and replicate consistency) across all models pre-trained on ImageNet and models trained on microscopy datasets. The paper also presents a manually curated 40-class dataset named Anax and compare performance of their proposed methods against baselines across Anax, RxRx1 and whole genome dataset.
weaknesses: 1. While the paper shows performance improvements with model scale on microscopy images, the overall contribution is incremental and lacks sufficient ablation studies.
2. While the authors hypothesize why their dataset curation might be better compared to other dataset curation strategies, there is no experimental comparison of different data curation strategies on microscopy datasets
3. There is very little novel contribution in the selection of optimal block for specific downstream task based on balanced accuracy on linear classification task
4. The gains in biological recall metric (across all 4 gene network databases) for the larger MAE-G/8 seem very modest compared to the MAE-L/8 baselines (Table 2). 
5. The biological recall metrics for trimmed vs untrimmed versions of the MAE models trained on microscopy images are mostly the same which is counterintuitive to the argument that intermediate block representations are more useful. 
6. The paper describes a data curation and a methodology for trimming blocks of ViT, it specifically explores the two options for training MAE models. While DINO-ViTs trained on ImageNet are a useful reference, additional experiments with DINO-ViT trained on microscopy images would be required to demonstrate the general applicability of the proposed data curation and ViT trimming strategies.
7. The paper does not explicitly state that the dataset and model weights will be available to the public. Metrics shared on private datasets (manually curated datasets) are insufficient to use for evaluation of scientific contribution.
8. While CA MAE was used as a baseline, the results and lower performance of CA-MAE compared to other baselines is not discussed.

---

confidence: 3
contribution: 3
initial_rating: 8
presentation: 4
questions: Will the ViT-G/8 MAE foundation model be made publicly available? I see no mention of this.
The bolding in table 2 is inconsistent.
soundness: 3
strengths: The paper is very well written and provides clear explanations of the experiments and dataset curation. I believe that this will help future researchers better curate their own datasets.
Their foundation model outperforms the previous SOTA and would be valuable to the scientific community if made publicly available.
The required effort to produce the curated dataset and training of the large-scale model was substantial.
Although it has been shown in other domains, the evidence that mid-level representations can be of higher quality than deeper layers is valuable for this field as it reduces the cost of inference and potential fine-tuning.
summary: The authors:

Curate a large dataset (Phenoprint-16M) using a unique combination of previously introduced curation methods.
Introduce a ViT-G/8 MAE foundation model for cell microscopy imaging which outperforms the previous STOA model.
Show, using linear-probing, that mid-level representations from their model can outperform the representations from deeper layers enabling cheaper inference and improved performance.
Show that scaling to a 1.9 billion parameter MAE model continues to improve the performance for cell microscopy imaging / HCS.
weaknesses: It’s already well-established that scaling up models + dataset curation helps performance, these findings are not novel.
Only one method to evaluate the quality of the representations is used (linear-probing). It would be more convincing if other measures were given as well (kNN, fine-tuning). Fine-tuning evaluations on down-stream tasks using a varying number of blocks would have been useful as well. I understand that this may have been too expensive for all layers, but could have been done for a select number of layers.
No comparison is made between DINO-V2 pretrained on microscopy images versus MAE pretrained on microscopy images. Are the MAE results better because of the training method or just the training data? There’s nothing in this paper that convinces me that an MAE is more suitable than any other SSL techniques.
Why is Figure 5 comparing the ViT-G/8 MAE model against the DINO-V2 G/14 pretrained on natural images? The results in Table 2 already show that pretraining on microscopy images is critical. It feels like an unnecessary comparison. I believe a more fair comparison would be against MAE-L/8 , RPI-93M, the previous STOA.

---

confidence: 4
contribution: 1
initial_rating: 1
presentation: 4
questions: I do not have questions about the paper per se. I do not see novel and reproducible enough to be published.
soundness: 1
strengths: This paper provides an analysis of the critical role dataset curation plays in achieving high-quality representations in microscopy imaging, including the effects of aggressive curation strategies. The work demonstrates how curatorial decisions impact the fidelity and accuracy of data representation.

The authors have presented their findings in a well-organized and accessible manner, with each section building logically on the last. The clarity of the language and structure ensures that concepts are conveyed effectively.

One of the strengths of this paper lies in its visual aids. The figures and images are thoughtfully designed to illustrate core concepts and provide visual clarity, enhancing the reader's understanding of the methodology and outcomes. 

The experimental results are convincing to some extent.
summary: The study presents a novel dataset specifically designed to train foundational models for the Cell Painting assay, a high-content imaging technique used for profiling cellular phenotypes. By curating this dataset, the authors aim to address the unique challenges associated with cellular imaging, such as variations in cellular morphology and staining patterns, which require specialized data to train robust and generalizable models.

In addition to introducing the dataset, the study evaluates the effectiveness of several ImageNet-pretrained models in encoding cellular images from the Cell Painting assay. By leveraging pretrained models originally trained on ImageNet, the authors explore how well these architectures can transfer to the domain of cellular imaging.
weaknesses: Firstly, I have some reservations regarding the model's evaluation protocol. I recommend that the authors consider using a phenotypic screening benchmark [1]. Specifically, for chemical perturbations, the validation splits should ideally be based on the scaffold, ensuring that drugs with similar structures do not appear in both the training and testing sets. Furthermore, an independent dataset is used to validate the method, which would help clearly assess the influence of batch effects. As an independent dataset I would propose to similar approach as in [1], namely, I would use the trained model and do the linear probing for mode of action prediction. 

In terms of contributions, the authors emphasize the importance of data quality. This includes aspects like the overall integrity of the data [2] as well as challenges related to imbalance and undersampling [3]. While these are critical considerations in vision or general datasets, they also apply to microscopy, so the contribution may not be particularly impactful for the research community. I do not see that those issues are specific for microscopy only, they are general for computer vision and transferable to microscopy. The microscopy-related issue is batch effects, and there is not much of their influence on the work. That is why I believe this work's novelty is extremely limited. 

Additionally, it appears that the authors have not made the pretrained model or dataset publicly available. While sharing these resources is not a requirement, their absence limits reproducibility, especially given the lack of detailed descriptions in the paper. Without these, reproducing the model’s training process—such as the steps required to train MAE on this dataset—is challenging and limits the model's utility as a foundational resource for the scientific community. I would envision seeing all the details about the quality control step. The descriptions of Appendix A are too general and to make it reproducible, they should be much more detailed. Probably a few times longer description than it is right now. 

There are also inconsistencies in the comparison of models. The models were trained with different hyperparameters, which raises questions about the validity of direct comparisons. Lastly, the results in Table 2 lack confidence intervals, making it difficult to assess whether the reported differences are statistically significant. Please add confidence intervals in the results. Additionally, would be great to see if the model trained with the same hyperparameters performs differently. 

[1] Borowa, Adriana, et al. "Decoding phenotypic screening: A comparative analysis of image representations." Computational and Structural Biotechnology Journal 23 (2024): 1181-1188.

[2] Cole, Elijah, et al. "When does contrastive visual representation learning work?." Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. 2022.

[3] Johnson, Justin M., and Taghi M. Khoshgoftaar. "Survey on deep learning with class imbalance." Journal of big data 6.1 (2019): 1-54.

Meta-review: The paper presents a new unsupervised vision transformer foundation model of cell microscopy with the goal of using the model for downstream tasks including predicting the cellular phenotype of chemical and genetic perturbations. The main contributions are scaling model size, curation of the training dataset, and the discovery that intermediate layers of the model are more informative for downstream tasks. The reviewers cite the lack of methodological novelty and the marginal benefit from model scaling to largely recommend rejection.

Decision: Reject