Generative Distribution Embeddings

Much of biology is measured one sample at a time but interpreted at the level of populations. A single-cell atlas is read as a statement about donors, not nuclei; a Perturb-seq screen is read as a statement about perturbations, not cells; a lineage-tracing experiment is read as a statement about clones. Standard autoencoders embed individual points, so population-level signal has to be recovered afterward by pooling — which discards everything about a distribution beyond its mean.

Generative Distribution Embeddings, from Harvard University, MIT, the Wyss Institute, and collaborators at Beth Israel Deaconess Medical Center, Harvard Medical School, and Mass General Brigham, lift the autoencoder to the space of distributions. The encoder consumes a set of samples and returns a single latent vector; the decoder is replaced by a conditional generative model tasked with reproducing the input distribution from that vector. The encoder must satisfy a property the authors call distributional invariance, which is what makes the resulting embedding a function of the distribution rather than of the particular sample draw.

GDE is a framework rather than a single checkpoint: each application trains its own encoder-generator pair with architectures suited to the data. What makes the trained instances useful beyond their fitting sets is that they generalize to conditions never seen in training. The work appeared at NeurIPS 2025.

#Key Features

• Distributionally invariant encoders: The encoder acts on sets, so its output depends on the underlying distribution and not the ordering or size of the sample drawn from it, which the authors formalize and prove permits mean-aggregation over minibatches of samples.
• Wasserstein geometry in the latent space: Latent distances approximately recover the W2 distance between distributions (Spearman rho of 0.96 on Gaussians), and latent interpolation approximately traces optimal transport trajectories.
• Architecture-agnostic pairing: Any distributionally invariant encoder can be coupled to any conditional generator — deep sets, ResNet-GNN, and 1D/2D convolutional GNN encoders are paired with CVAE, DDPM, HyenaDNA, and ProGen2 generators across the six biological applications.
• Zero-shot generalization to held-out conditions: Perturbation embeddings can be predicted from a text or sequence embedding of a gene the model never saw, then decoded into a full response distribution.
• Learning directly from raw reads: One experiment recovers tissue-specific DNA methylation signatures from bisulfite sequencing reads with no alignment, reference genome, or assay-structure information supplied.

#Technical Details

On synthetic benchmarks GDE beats both methods it generalizes — kernel mean embeddings with a DDPM decoder, and Wasserstein Wormhole — on reconstruction error for Gaussians (0.02 vs 0.04 and 0.20), Gaussian mixtures (1.82 vs 2.17 and 2.88), and image-set distributions. Across biology it was applied to six problems: donor representations from 6.3M prefrontal-cortex nuclei across 1,494 donors, where a semi-supervised GDE beats a supervised deep-sets baseline on accuracy (0.8887 vs 0.8791), ROC AUC, and F1; clone-level embeddings from 150K lineage-traced mouse hematopoietic cells, exceeding Wasserstein Wormhole by over 2 bits of mutual information between early and late clonal representations; transcriptional perturbation response on roughly 1M Perturb-seq cells, where predicting the GDE embedding from GenePT gene embeddings via ridge regression reaches R2 = 0.458 against 0.421 for scVI and 0.378 for the mean baseline; morphological response across 20M single-cell images from a pooled CRISPR screen of 5,072 essential genes, where holding out the 30% most perturbative perturbations still yields R2 = 0.7055 zero-shot on nuclear signal intensity, including the held-out RACGAP1 knockout; yeast promoter design over 34M random 80-nucleotide sequences with a HyenaDNA decoder, recovering transcription-factor binding site motif distributions with the top five expression quantiles held out; and spatiotemporal modeling of 1M SARS-CoV-2 spike sequences with a 50M-parameter ESM encoder and a conditional ProGen2 generator, predicting sampling month of held-out distributions with 1.83-month mean absolute error against 2.24 for mean-pooled ESM embeddings.

#Applications

GDE fits any setting where the analytical unit is a population: donor- or patient-level phenotyping from single-cell atlases, clone-level fate prediction in lineage tracing, in-silico prediction of transcriptional or morphological responses to unscreened genetic perturbations, regulatory sequence design against a target expression distribution, and surveillance-style modeling of viral sequence populations over time and geography. The zero-shot pathway is the practical draw: prioritizing which perturbations to run next without having run them.

#Impact

GDE gives a single, provably grounded formulation to a set of problems previously handled by unrelated bespoke methods, and demonstrates it at scale on six distinct biological modalities. Its main caveat is architectural: GDE names a method and a family of task-specific trained instances, not one downloadable checkpoint that transfers across all six applications, and the released code provides training and benchmarking pipelines rather than consolidated pretrained weights. The repository carries no license file, which limits reuse in derivative work. All validation is computational, on held-out splits of existing public datasets.