Elucidating the Design Space of Flow Matching for Cellular Microscopy

TL;DR

This paper systematically analyzes flow matching models for cellular microscopy, develops a scalable recipe, and achieves state-of-the-art results by scaling and fine-tuning the model with molecular embeddings.

Contribution

It introduces a simplified, stable, and scalable approach to flow matching for cell microscopy, scaling models significantly and enhancing performance with molecular embedding fine-tuning.

Findings

Many popular techniques are unnecessary or detrimental.

Scaling the model improves FID and KID metrics.

Fine-tuning with molecular embeddings achieves state-of-the-art results.

Abstract

Flow-matching generative models are increasingly used to simulate cell responses to biological perturbations. However, the design space for building such models is large and underexplored. We systematically analyse the design space of flow matching models for cell-microscopy images, finding that many popular techniques are unnecessary and can even hurt performance. We develop a simple, stable, and scalable recipe which we use to train our foundation model. We scale our model to two orders of magnitude larger than prior methods, achieving a two-fold FID and ten-fold KID improvement over prior methods. We then fine-tune our model with pre-trained molecular embeddings to achieve state-of-the-art performance simulating responses to unseen molecules. Code is available at https://github.com/valence-labs/microscopy-flow-matching