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

arXiv: https://arxiv.org/abs/2505.18150 · HTML: https://mlsanity.com/paper/2505.18150/

## Key facts

- Cited by 4 papers in the map (data through 14 Oct 2026).
- Builds on 27 of its 71 references that are in the map.
- First posted on arXiv 23 May 2025.

## Builds on (most cited first, 27 of 27)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [Denoising Diffusion Probabilistic Models](https://mlsanity.com/paper/2006.11239/) (2020)
- [Flow Matching for Generative Modeling](https://mlsanity.com/paper/2210.02747/) (2022)
- [Generative Adversarial Networks](https://mlsanity.com/paper/2203.00667/) (2022)
- [Fashion-MNIST: a Novel Image Dataset for Benchmarking Machine Learning Algorithms](https://mlsanity.com/paper/1708.07747/) (2017)
- [HuggingFace's Transformers: State-of-the-art Natural Language Processing](https://mlsanity.com/paper/1910.03771/) (2019)
- [Deep Sets](https://mlsanity.com/paper/1703.06114/) (2017)
- [An Introduction to Variational Autoencoders](https://mlsanity.com/paper/1906.02691/) (2019)
- [Kernel Mean Embedding of Distributions: A Review and Beyond](https://mlsanity.com/paper/1605.09522/) (2016)
- [Learning Generative Models with Sinkhorn Divergences](https://mlsanity.com/paper/1706.00292/) (2017)
- [Generalized Sliced Wasserstein Distances](https://mlsanity.com/paper/1902.00434/) (2019)
- [HyenaDNA: Long-Range Genomic Sequence Modeling at Single Nucleotide Resolution](https://mlsanity.com/paper/2306.15794/) (2023)
- [Towards a Neural Statistician](https://mlsanity.com/paper/1606.02185/) (2016)
- [Dirichlet Flow Matching with Applications to DNA Sequence Design](https://mlsanity.com/paper/2402.05841/) (2024)
- [SODA: Bottleneck Diffusion Models for Representation Learning](https://mlsanity.com/paper/2311.17901/) (2023)
- [A Wasserstein-type distance in the space of Gaussian Mixture Models](https://mlsanity.com/paper/1907.05254/) (2019)
- [Neural Approximate Sufficient Statistics for Implicit Models](https://mlsanity.com/paper/2010.10079/) (2020)
- [Fisher Flow Matching for Generative Modeling over Discrete Data](https://mlsanity.com/paper/2405.14664/) (2024)
- [Universal Approximation of Functions on Sets](https://mlsanity.com/paper/2107.01959/) (2021)
- [How to Build the Virtual Cell with Artificial Intelligence: Priorities and Opportunities](https://mlsanity.com/paper/2409.11654/) (2024)
- [Meta Flow Matching: Integrating Vector Fields on the Wasserstein Manifold](https://mlsanity.com/paper/2408.14608/) (2024)
- [Interpretable Diffusion via Information Decomposition](https://mlsanity.com/paper/2310.07972/) (2023)
- [Wasserstein Flow Matching: Generative modeling over families of distributions](https://mlsanity.com/paper/2411.00698/) (2024)
- [Approximating mutual information of high-dimensional variables using learned representations](https://mlsanity.com/paper/2409.02732/) (2024)
- [Set Norm and Equivariant Skip Connections: Putting the Deep in Deep Sets](https://mlsanity.com/paper/2206.11925/) (2022)
- [Wasserstein Wormhole: Scalable Optimal Transport Distance with Transformers](https://mlsanity.com/paper/2404.09411/) (2024)

## Built on by (most cited first, 4 of 4)

- [Neural Bayesian Filtering](https://mlsanity.com/paper/2510.03614/) (2025)
- [Distribution-Conditioned Transport](https://mlsanity.com/paper/2603.04736/) (2026)
- [Measure-to-measure Regression with Transformers](https://mlsanity.com/paper/2605.28075/) (2026)
- [Belief Embedding Tree Search](https://mlsanity.com/paper/2610.15046/) (2026)

## Abstract (arXiv)

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 W2 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).