Generative Distribution Embeddings: Lifting autoencoders to the space of distributions for multiscale representation learning

Fuente: arXiv
Saved in:
Bibliographic Details
Main Authors: Fishman, Nic, Gowri, Gokul, Yin, Peng, Gootenberg, Jonathan, Abudayyeh, Omar
Format: Preprint
Published: 2025
Subjects:
Online Access:
Tags: Add Tag
No Tags, Be the first to tag this record!
_version_ 1866911457592999936
author Fishman, Nic
Gowri, Gokul
Yin, Peng
Gootenberg, Jonathan
Abudayyeh, Omar
author_facet Fishman, Nic
Gowri, Gokul
Yin, Peng
Gootenberg, Jonathan
Abudayyeh, Omar
contents 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).
format Preprint
id arxiv_https___arxiv_org_abs_2505_18150
institution arXiv
publishDate 2025
record_format arxiv
spellingShingle Generative Distribution Embeddings: Lifting autoencoders to the space of distributions for multiscale representation learning
Fishman, Nic
Gowri, Gokul
Yin, Peng
Gootenberg, Jonathan
Abudayyeh, Omar
Machine Learning
Quantitative Methods
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).
title Generative Distribution Embeddings: Lifting autoencoders to the space of distributions for multiscale representation learning
topic Machine Learning
Quantitative Methods
url https://arxiv.org/abs/2505.18150