Generative Distribution Embeddings

TL;DR

The paper introduces generative distribution embeddings (GDEs), a framework that maps sets of samples to latent distributions using distributionally invariant encoders paired with conditional generators. By design, the latent geometry in GDEs approximates Wasserstein-$W_2$ distances and OT geodesics, effectively functioning as predictive sufficient statistics for distribution-level inference. The authors demonstrate strong empirical performance on both synthetic benchmarks and six large-scale computational biology tasks, including lineage-traced scRNA-seq, Perturb-seq, cellular morphology, BS-seq methylation, yeast promoter activity, and spatiotemporal viral sequence distributions, highlighting the method's versatility and scalability. The work also develops a formal statistical-geometric foundation, linking encoder invariance to asymptotic sufficiency and viewing the learned manifold as a constrained Wasserstein submanifold, with prior-weighting to tailor geometry for downstream objectives. Overall, GDEs offer a flexible, scalable approach to population-level inference across domains where distributions—not individual samples—are the primary unit of analysis.

Abstract

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 representations of cell populations from lineage-tracing data (150K cells), predicting perturbation effects on single-cell transcriptomes (1M cells), predicting perturbation effects on cellular phenotypes (20M single-cell images), modeling tissue-specific DNA methylation patterns (253M sequences), designing synthetic yeast promoters (34M sequences), and spatiotemporal modeling of viral protein sequences (1M sequences).

Figures (11)

• Figure 1: GDEs leverage distribution-invariant encoders ($\mathcal{E}$) and conditional generative models ($\mathcal{G}$) to lift autoencoders to statistical manifolds where points correspond to distributions ($\mathcal{M}$).
• Figure 2: Right: Histogram of loss for fixed $P_i$ different set sizes. Left: First two PCs of latent representation of empirical distributions of MNIST data from Fig. \ref{['fig:multinomial']} for fixed $P_i$ for different set sizes.
• Figure 3: $L_2$ in GDE latent space compared to $W_2$ distance. Normalized distances from the center, $p = (1/3, 1/3, 1/3)$. The plots to the left show GDE $L_2$ learned from empirical distributions. MNIST and DNA distributions are constructed by sampling conditional on class label according to a multinomial, for MNIST subsetted to images of (0, 1, 2) and a synthetic DNA dataset with 3 patterns respectively. Rightmost plot shows the Gaussian approximation for the $W_2$ between multinomials.
• Figure 4: Top row: Trajectories between pairs of Gaussians under optimal transport (left) and GDE (right). Bottom row: Similar comparison for Gaussian mixture models, we compute the "OT" by finding the optimal pairing between Gaussians and computing the OT. Inset ternary plots show mixture weights during interpolation.
• Figure 5: Similar to Fig. \ref{['fig:multinomial']} we show the GDE distances of multinomials from $p = (\frac{1}{3},\frac{1}{3}, \frac{1}{3})$. We shift the prior asymmetrically by changing $\alpha_1$ while fixing $\alpha_2=\alpha_3=1$. This shifts the focus of the model, leading to a different learned geometry.

Theorems & Definitions (30)

• Proposition 1
• Proposition 2
• Proposition 3
• Example 1: Gaussian Mean
• Example 2: Gaussian Mixture Model
• Example 3: Uniform$(0,\theta)$
• Definition 1: Distributional Invariance
• Definition 2: Asymptotic Distributional Invariance
• Lemma 1: Strong Law of Large Numbers for Empirical Measures
• Lemma 2: Wainwright’s Rademacher–tail bound wainwright2019hds