Openphenom

🚀 Model Card for OpenPhenom-S/16

This is a channel-agnostic image encoding model, CA-MAE, with a ViT-S/16 encoder backbone. It's specifically designed for microscopy image featurization. The model utilizes a vision transformer backbone with channelwise cross-attention over patch tokens to generate contextualized representations for each channel separately.

🚀 Quick Start

You should be able to successfully run the following tests, which demonstrate how to use the model at inference time.

import pytest
import torch

from huggingface_mae import MAEModel

# huggingface_openphenom_model_dir = "."
huggingface_modelpath = "recursionpharma/OpenPhenom"


@pytest.fixture
def huggingface_model():
    # This step downloads the model to a local cache, takes a bit to run
    huggingface_model = MAEModel.from_pretrained(huggingface_modelpath)
    huggingface_model.eval()
    return huggingface_model


@pytest.mark.parametrize("C", [1, 4, 6, 11])
@pytest.mark.parametrize("return_channelwise_embeddings", [True, False])
def test_model_predict(huggingface_model, C, return_channelwise_embeddings):
    example_input_array = torch.randint(
        low=0,
        high=255,
        size=(2, C, 256, 256),
        dtype=torch.uint8,
        device=huggingface_model.device,
    )
    huggingface_model.return_channelwise_embeddings = return_channelwise_embeddings
    embeddings = huggingface_model.predict(example_input_array)
    expected_output_dim = 384 * C if return_channelwise_embeddings else 384
    assert embeddings.shape == (2, expected_output_dim)

We also provide a notebook for running inference on RxRx3-core.

✨ Features

• Create biologically useful embeddings of microscopy images.
• Create contextualized embeddings of each channel of a microscopy image (set return_channelwise_embeddings=True).
• Leverage the full MAE encoder + decoder to predict new channels / stains for images without all 6 CellPainting channels.
• A determined ML expert could fine-tune the encoder for downstream tasks such as classification.

📚 Documentation

Model Details

Model Description

This model is a channel-agnostic masked autoencoder trained to reconstruct microscopy images over three datasets:

1. RxRx3
2. JUMP-CP overexpression
3. JUMP-CP gene-knockouts

Model Sources

• Repository: https://github.com/recursionpharma/maes_microscopy
• Paper: Masked Autoencoders for Microscopy are Scalable Learners of Cellular Biology

Uses

⚠️ Important Note

Model embeddings tend to extract features only after using standard batch correction post-processing techniques.

💡 Usage Tip

We recommend, at a minimum, after inferencing the model over your images, to do the standard PCA-CenterScale pattern or better yet Typical Variation Normalization:

1. Fit a PCA kernel on all the control images (or all images if no controls) from across all experimental batches (e.g. the plates of wells from your assay).
2. Transform all the embeddings with that PCA kernel.
3. For each experimental batch, fit a separate StandardScaler on the transformed embeddings of the controls from step 2, then transform the rest of the embeddings from that batch with that StandardScaler.

Direct Use

• Create biologically useful embeddings of microscopy images.
• Create contextualized embeddings of each channel of a microscopy image (set return_channelwise_embeddings=True).
• Leverage the full MAE encoder + decoder to predict new channels / stains for images without all 6 CellPainting channels.

Downstream Use

• A determined ML expert could fine-tune the encoder for downstream tasks such as classification.

Out-of-Scope Use

• Unlikely to be especially performant on brightfield microscopy images.
• Out-of-domain medical images, such as H&E (maybe it would be a decent baseline though).

Bias, Risks, and Limitations

• Primary limitation is that the embeddings tend to be more useful at scale. For example, if you only have 1 plate of microscopy images, the embeddings might underperform compared to a supervised bespoke model.

Training, evaluation and testing details

See paper linked above for details on model training and evaluation. Primary hyperparameters are included in the repo linked above.

Environmental Impact

BibTeX

@inproceedings{kraus2024masked,
  title={Masked Autoencoders for Microscopy are Scalable Learners of Cellular Biology},
  author={Kraus, Oren and Kenyon-Dean, Kian and Saberian, Saber and Fallah, Maryam and McLean, Peter and Leung, Jess and Sharma, Vasudev and Khan, Ayla and Balakrishnan, Jia and Celik, Safiye and others},
  booktitle={Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition},
  pages={11757--11768},
  year={2024}
}

Model Card Contact

• Kian Kenyon-Dean: kian.kd@recursion.com
• Oren Kraus: oren.kraus@recursion.com
• Or, email: info@rxrx.ai