Generative Distribution Embeddings

Nic Fishman (Harvard University) · Gokul Gowri (Massachusetts Institute of Technology) · Peng Yin (Harvard Medical School, Harvard University) · Jonathan Gootenberg (Beth Israel Deaconess Medical Center, Harvard Universirt) · Omar Abudayyeh (Brigham and Women's Hospital, Harvard University)

$w_2$ distanceautoencodersclonal dynamicscomputational biologyconditional generative modelsdistributional invariancegenerative distribution embeddingslatent gde distancesoptimal transport trajectoriesperturbation effectspredictive sufficient statisticssingle-nuclei rna sequencingspatiotemporal modelingsynthetic datasets wasserstein space

Many real-world problems require reasoning across multiple scales, demanding models which operate not on single data points, but on entire distributions. We introduce generative distribution embeddings (GDE), a framework that lifts autoencoders to the space of distributions. In GDEs, an encoder acts on sets of samples, and the decoder is replaced by a generator which aims to match the input distribution. This framework enables learning representations of distributions by coupling conditional generative models with encoder networks which satisfy a criterion we call distributional invariance. We show that GDEs learn predictive sufficient statistics embedded in the Wasserstein space, such that latent GDE distances approximately recover the $W_2$ distance, and latent interpolation approximately recovers optimal transport trajectories for Gaussian and Gaussian mixture distributions. We systematically benchmark GDEs against existing approaches on synthetic datasets, demonstrating consistently stronger performance. We then apply GDEs to six key problems in computational biology: learning donor-level representations from single-nuclei RNA sequencing data (6M cells), capturing clonal dynamics in lineage-traced RNA sequencing data (150K cells), predicting perturbation effects on transcriptomes (1M cells), predicting perturbation effects on cellular phenotypes (20M single-cell images), designing synthetic yeast promoters (34M sequences), and spatiotemporal modeling of viral protein sequences (1M sequences).