Out-of-distribution evaluations of channel agnostic masked autoencoders in fluorescence microscopy

TL;DR

Distribution-shift in fluorescence microscopy hampers generalisation, as typical transfer-learning evaluations mix multiple sources of variation. The authors propose a JUMP-CP-based evaluation scheme to isolate shifts and introduce Campfire, a channel-agnostic masked autoencoder that shares a decoder across channels to scale to diverse fluorescent markers. Across held-out plates, perturbagens, fluorescent channels, and cell types, Campfire learns representations that generalise and transfer with limited fine-tuning, outperforming ImageNet-pretrained baselines. The work provides a framework for evaluating and building foundation models for high-content imaging and suggests future directions to include more channels and cell types for robust, universal representations.

Abstract

Developing computer vision for high-content screening is challenging due to various sources of distribution-shift caused by changes in experimental conditions, perturbagens, and fluorescent markers. The impact of different sources of distribution-shift are confounded in typical evaluations of models based on transfer learning, which limits interpretations of how changes to model design and training affect generalisation. We propose an evaluation scheme that isolates sources of distribution-shift using the JUMP-CP dataset, allowing researchers to evaluate generalisation with respect to specific sources of distribution-shift. We then present a channel-agnostic masked autoencoder $\mathbf{Campfire}$ which, via a shared decoder for all channels, scales effectively to datasets containing many different fluorescent markers, and show that it generalises to out-of-distribution experimental batches, perturbagens, and fluorescent markers, and also demonstrates successful transfer learning from one cell type to another.

Figures (5)

• Figure 2.1: A sketch of how wells are assigned to training/validation/test sets for a given compound in TARGET2 plates. For each compound, we select 14 plates from which to sample a single well treated with that compound for the training set, 2 plates to sample a well for the validation set, and 4 plates for the test set. The exception are 60 randomly selected compounds which are all held-out of training. We also hold 5 TARGET2 plates out of training, with all wells added to the test set.
• Figure 2.2: A sketch of the Campfire architecture.
• Figure 3.3: An example of the reconstruction of a cell-centred tile. From left to right, we show the masked input, followed by six reconstructed images at different epochs, with epoch increasing from left to right, and lastly the original cell-centred tile to be reconstructed.
• Figure 3.4: Accuracy of a linear classifier predicting 1-of-9 control compounds (left panel) or 1-of-60 held-out compounds (right) from single cell embeddings. Each column represents a linear classifier trained on single cell embeddings derived from images comprised of different sets of fluorescent channels. Channels shown: Nucleus (N), Actin+Golgi apparatus+Plasma Membrane (Ac), Mitochondria (M) and their combinations.
• Figure 3.5: $Z'$-score measuring statistical difference between model embeddings from a reference and target compound of stimulation. Model embeddings are derived from Campfire (left) pretrained on JUMP-CP, and DinoViT-S8 (right) pretrained on ImageNet1k. Both models have been finetuned on a macrophage dataset. $Z'$-score is shown for negative (neg) and positive (pos) controls for plates with macrophages in either M1 or M2 polarisation.