[Paper Review] Masked Autoencoders for Microscopy are Scalable Learners of Cellular Biology

Oren Kraus, Kian Kenyon-Dean+15 more|arXiv (Cornell University)|Apr 16, 2024

Cell Image Analysis Techniques10 citations

TL;DR

The paper shows that self-supervised masked autoencoders (MAEs) with Vision Transformer backbones scale to large microscopy datasets and outperform weakly supervised methods in recalling known biological relationships; it also introduces a channel-agnostic MAE (CA-MAE) that generalizes across different channel configurations.

ABSTRACT

Featurizing microscopy images for use in biological research remains a significant challenge, especially for large-scale experiments spanning millions of images. This work explores the scaling properties of weakly supervised classifiers and self-supervised masked autoencoders (MAEs) when training with increasingly larger model backbones and microscopy datasets. Our results show that ViT-based MAEs outperform weakly supervised classifiers on a variety of tasks, achieving as much as a 11.5% relative improvement when recalling known biological relationships curated from public databases. Additionally, we develop a new channel-agnostic MAE architecture (CA-MAE) that allows for inputting images of different numbers and orders of channels at inference time. We demonstrate that CA-MAEs effectively generalize by inferring and evaluating on a microscopy image dataset (JUMP-CP) generated under different experimental conditions with a different channel structure than our pretraining data (RPI-93M). Our findings motivate continued research into scaling self-supervised learning on microscopy data in order to create powerful foundation models of cellular biology that have the potential to catalyze advancements in drug discovery and beyond.

Motivation & Objective

• Scale self-supervised learning on high-content screening microscopy (HCS) data to learn transferable cellular embeddings.
• Compare MAEs with weakly supervised learning (WSL) and ImageNet-pretrained baselines on biological-relationship recall.
• Develop and evaluate a channel-agnostic MAE (CA-MAE) that handles different input channel configurations at inference.
• Assess transferability of MAE representations to datasets with differing channels and experimental conditions.
• Analyze training dynamics and reconstruction losses to improve MAE stability for large backbones.

Proposed method

• Pretrain MU-Net-based and ViT-based MAEs on progressively larger HCS datasets (RxRx, RPI-52M, RPI-93M).
• Use 256x256 crops with channel-wise self-standardization as input; experiment with 8x8 and 16x16 patches and 75%/25% masks.
• Incorporate a Fourier-domain reconstruction loss to stabilize training and improve texture reconstruction (alpha = 0.01 in L_MAE+ = (1−alpha)L_MAE + alpha L_F).
• Adopt Large-scale ViT encoders (ViT-S/B/L) with MAE training, averaging final-layer patch embeddings for image-level representations.
• Develop CA-MAE: treat channels as separate modalities with a shared tokenizer and per-channel decoders, enabling inference with different channel counts and orders.
• Apply Typical Variation Normalization (TVN) and chromosome-arm bias corrections for batch correction and robust biological-relationship recall.

Experimental results

Research questions

• RQ1 Does scaling MAEs with larger model backbones and bigger pretraining datasets improve recall of known biological relationships in HCS data?
• RQ2 Does a Fourier-domain reconstruction loss stabilize training and improve downstream biological-relationship recall for large ViT-based MAEs?
• RQ3 Can a channel-agnostic MAE (CA-MAE) generalize to microscopy datasets with different channel configurations at inference?
• RQ4 How do MAEs compare to weakly supervised methods and ImageNet-pretrained models in recalling perturbation-based biological relationships?
• RQ5 What is the transferability of MAE-derived embeddings to external datasets like JUMP-CP and their Perturbation/Sibling retrieval performance?

Key findings

• MAEs outperform pretrained ImageNet and most WSL models, with up to 11.5% relative improvement over the best WSL model on recalling biological relationships (ViT-L/8+ trained on RPI-93M vs ViT-L/16 trained on RxRx1-2M).
• Recall of known biological relationships scales with training FLOPS (model size × dataset size); larger MAEs and datasets yield higher recall across databases (CORUM, hu.MAP, Reactome, StringDB).
• Incorporating a Fourier-domain reconstruction loss stabilizes training of large MAE ViTs and enables crossing saddle points in loss landscapes; it yields robust double-descent behavior.
• Channel-agnostic MAEs enable generalization to datasets with different channel configurations; CA-MAE ViT-L/16+ on RPI-93M achieves strong cross-dataset recall (e.g., CORUM 0.614, hu.MAP 0.424, Reactome 0.264, StringDB 0.478).
• Transfer to JUMP-CP shows competitive perturbation retrieval with CA-MAE and MAE models; CA-MAE achieves high perturbation retrieval (e.g., 0.95 average precision on Perturbation retrieval on CPJUMP1 subset).
• MAE embeddings better predict a wide range of CellProfiler features than a strong WSL model, indicating richer morphological representation learning.