Compressing Biology: Evaluating the Stable Diffusion VAE for Phenotypic Drug Discovery

TL;DR

High-dimensional Cell Painting data challenge generative modeling, motivating the use of Stable Diffusion's VAE (SD-VAE) for image reconstruction. The authors present a systematic benchmarking framework that evaluates pixel-level, embedding-based, latent-space, and retrieval-based metrics, using CPJUMP1 and LSUN to assess reconstruction fidelity and biological signal retention. They find that SD-VAE reconstructions preserve phenotypic information with minimal loss and that general-purpose embeddings (e.g., InceptionV3) can match or exceed domain-specific OpenPhenom in retrieval tasks, despite higher latent-space irregularity in microscopy data. The work provides practical guidelines for validating microscopy generative models and supports deploying off-the-shelf components in phenotypic drug discovery workflows.

Abstract

High-throughput phenotypic screens generate vast microscopy image datasets that push the limits of generative models due to their large dimensionality. Despite the growing popularity of general-purpose models trained on natural images for microscopy data analysis, their suitability in this domain has not been quantitatively demonstrated. We present the first systematic evaluation of Stable Diffusion's variational autoencoder (SD-VAE) for reconstructing Cell Painting images, assessing performance across a large dataset with diverse molecular perturbations and cell types. We find that SD-VAE reconstructions preserve phenotypic signals with minimal loss, supporting its use in microscopy workflows. To benchmark reconstruction quality, we compare pixel-level, embedding-based, latent-space, and retrieval-based metrics for a biologically informed evaluation. We show that general-purpose feature extractors like InceptionV3 match or surpass publicly available bespoke models in retrieval tasks, simplifying future pipelines. Our findings offer practical guidelines for evaluating generative models on microscopy data and support the use of off-the-shelf models in phenotypic drug discovery.

Figures (4)

• Figure 1: Overview of our evaluation pipeline. Input images, including Cell Painting and natural images, are encoded into the latent space using the Stable Diffusion VAE (SD-VAE). These latent representations are decoded to reconstruct the images. Original and reconstructed images are compared using channel-wise metrics: mean absolute error (MAE), structural similarity index measure (SSIM), and Earth mover's distance (EMD). Both sets of real and reconstructed images are passed through InceptionV3 and OpenPhenom to extract feature embeddings, which are used to compute the Fréchet Inception distance (FID) and fraction retrieved (FR) of perturbations against negative controls. Latent vectors are further used to compute the Kullback-Leibler divergence (KLD).
• Figure 2: Box plots showing MAE (left) and EMD (right) values across the LSUN and CPJUMP1 datasets and their channels, computed after a single run of SD-VAE applied to all images. The central line within each box shows the median, while the box boundaries represent the 1st and 3rd quartiles.
• Figure 3: FID (left) and KLD (right) scores across the various data subsets, computed after a single run of SD-VAE applied to all images. FID scores are computed using real and reconstructed images. All images were featurized using InceptionV3. KLD scores are presented as the mean values computed across all samples within each dataset. Standard deviations are approximately 19k, 8k, and 10k for Cell Painting, Classroom, and Church images, respectively.
• Figure 4: Box plot showing SSIM values across the LSUN and CPJUMP1 datasets and their respective channels, computed after a single run of SD-VAE applied to all images. The central line within each box shows the median, while the box boundaries represent the 1st and 3rd quartiles.