Joint work with Charlotte Bunne, Gabriele Gut, Jacobo Sarabia del Castillo, Lucas Pelkmans and David Alvarez-Melis.

We were invited to give a spotlight presentation about our work at the NeurIPS workshop Learning Meaningful Representations of Life (2022).

Tracking the development of cells over time is a major challenge in biology, as measuring cells usually requires their destruction. Optimal transport (OT) can help solve this challenge by learning an optimal coupling of samples taken at different points in time, thus enabling the reconstruction of pairwise correspondences between cells of different measurements. However, the classical formulation of OT assumes conservation of mass, which is violated in unbalanced scenarios in which the population size changes, e.g., when cells die or proliferate. In this work, we present NubOT, a neural unbalanced optimal transport model that learns a parameterized optimal transport map between unbalanced distributions. To model variation of mass, we rely on the formalism of semi-couplings and propose a novel parameterization and algorithmic scheme that is based on a cycle-consistent learning procedure. We apply our model to the challenging task of predicting the responses of heterogeneous cells to cancer drugs on the level of single cells. By accurately modeling cell proliferation and death, our method yields notable improvements over previous neural optimal transport methods.

Check out our paper here.