Generative Distribution Embeddings: Lifting autoencoders to the space of distributions for multiscale representation learning
Authors: Nic Fishman
Gokul Gowri
Peng Yin
Jonathan S. Gootenberg
Omar Abudayyeh
Venue: NeurIPS 2025 poster ()
Keywords: generative modeling
multiscale
distribution embeddings
representation learning
diffusion models
large language models
variational autoencoders
optimal transport
biology
genomics
Primary area: machine_learning_for_sciences
TL;DR: We introduce generative distribution embeddings (GDEs), a framework that lifts autoencoders to the space of distributions.

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

Generative Distribution Embeddings: Lifting autoencoders to the space of distributions for multiscale representation learning
Nic Fishman" Gokul Gowri"!, Peng Yin', Jonathan Gootenberg 156." and Omar Abudayyeh" 45.6."
Department of Statistics, Harvard University; LLaboratory for Information and Decision Systems, MIT; Wyss Institute for Biologically Inspired Engincering and Department of Systems Biology, Harvard University;
Center for Virology and Vaceine Research, Beth Isracl Deaconess Medical Center, Harvard Medical School; 5 Dept. of Medicine, Div. of Engineering in Medicine, Brigham and Women's Hospital, Harvard Medical School; "Gene and Cell Therapy Institute, Mass General Brigham
*Equal contribution: njwfish@gmail. com, gokulgdmit www.gtttieeeeeeeo edu.
"Senior authors jointly supervised this work: jgootenb@bidmc. hawvard. edu, omardabudayych science.
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 au- toencoders to the space of distributions. In GDEs, an encoder acts 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 call distributional invariance. We show that GDEs learn predictive suf- ficient 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 distribu- tions. We systematically benchmark GDEs against existing approaches synthetic datasets, demonstrating consistently stronger performance. We then apply GDEs 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 (IM 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).
of
Introduction
g
Advancements in science and engineering increasingly depend on our ability reason across multiple scales: modeling not just individual data points, but entire populations those datapoints are drawn from. In ap- plications ranging from single-cell genomics to DNA sequence design, the relevant unit of analysis not an individual sample (e.g., single cell), but the dis- tribution from which it is drawn (e.g. the cell state or the patient they were sampled from). These settings are fundamentally hierarchical: we observe sets samples from latent distributions. which themselves are drawn from a mnta-distribution Without directly modeling these distributions, population-level
M
GDEs leverage distribution- invariant encoders (E) and conditional gener- ative models (G) to lift autoencoders to sta- tistical manifolds where points correspond distributions (M).
Figure 1:
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

signals can be lost underneath unit-level noise. The fundamental challenge we consider is how to learn representations at the level of distributions, not just individual data points.
We introduce generative distribution embeddings (GDEs) (Fig. D. a framework that lifts autoen- coders to the distribution space. In GIDEs, the encoder maps a finite set of samples empirical distribution - to a latent space, while the decoder is replaced by a generative model that reconstructs the distribution by sampling conditional on this latent representation. Our central observation is that strong distributional representations can be learned by coupling conditional generative models with encoders that satisfy a minimal distributional invariance property.
Our framework synthesizes modern generative modeling, classical statistics, and information geom- etry. We show empirically that GDEs behave approximate predictive sufficient statistics [I 2]. capturing distribution-level structure while marginalizing over sampling noise. Moreover, the learned latent spaces exhibit geometrie regularity: latent 122 distances correlate with Wasserstein distances (W2) between underlying distributions, and linear interpolation in latent space approximates optimal transport geodesics [3], such that one can generate synthetic data which smoothly interpolate between observed distributions.
We benchmark GDEs on synthetic datasets with known parametric structure, demonstrating improved generative fidelity and structure preservation relative to baselines. We then scale our approach to multiple domains in computational biology, showcasing GDEs' versatility in modeling distributions defined across diverse organizing principles such as distinct populations, varying experimental conditions, spatial arrangements, and temporal dynamics. We demonstrate six applications: learning donor-level representations from single-nuclei RNA sequencing data (6M cells), capturing clonal dynamics in theege-tracce RNA sequencing data (150K cells), predicting perturbation effects transcriptomes (IM cells), predicting perturbation effects on cellular phenotypes (20M single-cell images), designing synthetic yeast promoters (34M sequences), and spatiotemporal modeling of viral protein sequences (IM sequences). Across these domains, GDEs offer a flexible and scalable framework for distribution-level inference. Code for all experiments is available here
Setting and methods
400
Count
2.1 A motivating example
200
Many modern datasets comprise large groups of exchangeable samples: for example, single cells or DNA sequences collected from individ- ual patients 4 (]. We observe n such groups, each written as St,, { and collect them 2 [S.m 1n with Tzy X. Each group's 'ss samples are i..dd draws from a latent, group-specific distribution P. P(X)) and the P. themselves are i.i.d. draws from a meta- distribution Q: that is, we first sample P, Q and then, conditional on P,, draw Zij ~ P,. In this framework, "patient" (or more generally, any group) is implicitly represented by its probability measure P,.
0.006
0.008
0.010
0.012
DDPM loss
PC1
Set
10
100
1000
10000
Figure Concentration of distribution embed- dings and plug-in in loss. (Left) Distribution of plug- in GDE loss (diffusion generator) for sets of dif- ferent sizes sampled from the same distribution P. over MNIST digits. (Right) First two principal components of embeddings of sample sets of dif- ferent sizes generated by the same P,.
This is a classical hierarchical data generating process, which gives rise to a multiscale problem. The subject of interest here is not only individual cells or DNA sequences Tıj, but the probability measures P.. As we explore in Sec. 3 and concretely demonstrate in Sec. 6, this setting is broadly applicable beyond the particular case of representing patients as a collection of samples.
In practice, there are often two major challenges in modeling this kind of data. First, unit-level data is often inherently noisy. For example, single-cell data suffers from noise due to molecular undersampling [6]. Our goal is to learn patient embeddings which capture distribution-level signal, rather than the sample-level noise. The second challenge that groups can contain millions of samples (in the case of DNA sequencing reads per patient, m can be ~ 108). It is computationally infeasible to train a model on all samples simultaneously, but given the inherent noise at the unit level would benefit from embedding all available samples at inference time.

---

In the remainder of Sec. 2, we will show that both of these practical challenges can be overeome by learning distribution embeddings rather than simply encoding sets of samples. The distributional perspective enables models which distill distribution-level signal, and are able to massively scale at inference time to make use of all available data for precise embeddings.
2.2 Learning generative distribution embeddings
We address this problem with GDEs, which consist of an encoder & that maps finite set of samples St.n to latent representation, and a conditional generator 9 that (given this representation) induces distribution the sample space. Formally, we aim to learn ε, G such that
Algorithm 1 Training GDEs
1:
for each set St.r. do Subsample Si. ~ St.,.
2:
3: 24 E(S
4: & g(;))
G(E(S..m) m=-000 P..
(I)
5: Backprop E
6:
end for
The loss l is the standard training objective for the condi- tional generator (for example, an evidence lower bound for a VAE or a denoising score-man matehig objective for a diffusion model); we do not need backpropagate through the sampling process of
We show that to guarantee Eq. (1) the encoder must satisfy the following two constraints:
1. Permutation invariance: reordering the samples in St.m does not change the embedding.
2. Proportional invariance: duplicating every sample K times does not change the embedding.
We refer to an encoder with these properties as distributionally invariant: the encoder must depend only on the empirical distribution. So for some function & we can write:
1 m7
m
Prmm
E(S..mm (Pp.,.).
We show formally that distributionally invariant encoders can capture any distributional property and furthermore that any non-distributionally invariant architecture can spuriously encode noise features irrelevant to the distribution (formal proofs in App. D.2
Proposition 1. (Informal statement of Corollary 7y To rule out order and set-size artifacts, the encoder should depend on the sample only through its empirical distribution (equivalently: be permutation invariant and invariant to proportional duplication).
Beyond separating signal and noise, distributional invariance, coupled with Hadamard differentiability of the pooling operator, has a second consequence: it enables a central limit theorem for embeddings. As the set size grows, {m...m) concentrates around its population value with Gaussian fluctuations:
Vm (E(S.m)) -((.)) ب Es,))
This result, illustrated empirically in Fig. 2, is what makes encoding massive sets possible. The CLT composes through the plug-in loss that {(Si.m: (.S..mm(E(()) ((E(S..m)) is consistent and asymptotically normal estimator of the population loss. This provides theoretical justification for training GDEs on subsets of larger sample sets: the gradient of the loss computed small sets matches (in expectation) the gradient computed using all samples per set (see App. D.2h:
Proposition 2. (Informal statement of Theorem 2 Fixing P. under mild regularity condi- tions: (i) a distributionally invariant encoder will have asymptotically normal distribution embeddings; (ii) for a suitable divergence, the plug-in loss, E, = E( E(Sm. ((E is consistent and asymptotically normal around the population loss; (iii) global minimizer will recover the true data distribution as m - oo: ((E(S.))) = P. See Fig. 2
Violating distributional invariance (for example, by using sum pooling) causes the embedding to depend on set size and breaks this limit theory causing Eq. (1) to fail. contrast, mean pooling and M/Z-estimmaors satisfy these properties (see App. D.2
3

---

3 From labels to distributions
In the previous section, we focused on a simple motivating example where we have patients and their associated single-cell data. In that setting, the multiscale nature is clear and it is straightforward to define a metadistribution (in other words, how to group samples into sets). In many datasets, a natural hierarchy is not clear: for example, DNA sequences and expression labels are not multiscale in the same sense as our motivating example.
Here we illustrate how to set up the distribution learning problem in a more general framework based ee
on a dataset of unit-level outcomes associated with labels = {(Tx y1:) )+=)]== rather than sets drawn from Q. We will show how to group data points into sets {Iaz [= whose empirical distributions P. approximate draws from Q. The grouping reflects the structure of the label space y and enables us to shape the GDE latent space for downstream applications (see Sec. [
When y is discrete, we can form sets by grouping datapoints with the same label (e.g., our motivating patient example further explored in Sec. 6.2 cells by clone identity in Sec. [6.3] cell transcriptomes by perturbation in Sec. [6.4] or epigenetic samples by tissue in Sec. B.I]. If there is some semantic simi- larity between discrete labels we can define sets proportional to those similarity metrics, paralleling contrastive learning, where labels define semantic neighborhoods.
When Y' is continuous or structured (e.g. spatial coordinates for the Tiy as in Sec. 6.5 or temporal in the viral protein sequences in Sec. 6).). we can use similarity kernel to sample points near target yf: Wikk exp(-d(w) y, "2"²²(2²²") defining a probabilistic neighborhood in the label space, enforcing the consideration of the local structure.
When labels are noisy measurements (e.g. expression associated with DNA sequences as in Sec. 6.6 we can invert the noise model y/ y + c₁ by sampling latent target y and computing likelihood weights wick = p(ys vī), yielding samples that reflect the uncertainty of the data.
All these constructions can be unified as instances of a general framework: let Q09) be a prior over label distributions. Fix a reference measure v on ) (counting for discrete labels, Lebesgue for continuous) and assume P(()
v. For each set i, draw p(Y) Q(3) and compute weights
dP(Y)
Wisk
WiA
dv
(8k),
Wik
Σn Way
and sample zig from D accordingly. This framework subsumes the above examples and gives general set of tools for shaping the GDE latent space, as we will illustrate in Sec.
4 Related work
Several lines of literature have tried to Iearn distribution embeddings summary statistics. Kernel methods, such as kernel mean embedding (KME) and set kernels, provide nonparametric approaches to represent probability measures as points in reproducing kernel Hilbert space, enabling tasks like distributional regression and classification [7] g 9 10 11. GDEs naturally nest these methods as particular choices of distributionally invariant encoders. GDEs also generalize the approach in [12). where they develop a particular encoder and
generator.
Distribution embeddings have also been studied from geometric perspective. Building on theoretical foundations from Amari 13, several works model distributions as points on a manifold imbued with the Fisher-Rao metric [14
15 16 17.. These methods are either not generative or restricted to categorical distributions. Building on the work of Otto [3, others have considered learning flows over Wasserstein spaces [18 19 (see Appendix C.2 for background on Wasserstein spaces), primarily focused on leveraging distribution encodings for transport problems as opposed to GDEs which aim to auto-encode distributions. GDEs are complementary to these works, and can be plugged in to many of these frameworks. One recent method closely related to GDEs, Wasserstein Wormhole 200 aims represent distributions points in space where Euclidean distances match Sinkhorn divergences in the sample space. Wasserstein Wormhole is a particular instantiation of a GDE, using an attention-based encoder and generator that only samples fixed number of points.
25 These methods typically consider a supervised setting with a particular inferential target. For example, in
related body of work aims to learn informative summary statistics 21 22

---

the context of likelihood-free inference, one aims to learn summary statistics which are maximally informative about the parameters of a generative model 23 24
GDEs are distinct from these approaches along several dimensions: first, we generalize these methods under common framework with central objective of re-sampling the encoded distribution (D: second, we develop theory to guide the design and analysis of GDEs toward this objective; third, we show that distribution embedding is deeply related to generative modeling, enabling domain-specific generative models to be bootstrapped into high-quality GDEs to tackle multiscale problems.
On the architectural side, the encoder in the GDE framework requires distributionally invariant model. While distributional invariance is a concept introduced in this work, it requires permutation invariance, which has been well-studied [26] 27 28) Some permutation invariant approaches, such
as deep sets 26]. are not distributionally invariant due to proportional sensitivity, while others, such as mean-pooled attention layers, are also distributionally invariant (as shown in Appendix [D.2].
key contribution of our work the observation that any conditional generative model can be repurposed to learn distributional representations, Recent work in the vision domain has found that conditional diffusion models can induce strong image representations [29] Our work formalizes and generalizes this finding. We demonstrate in practice that number of modern techniques, including variational autoencoders [30], Sinkhorn based generative models [3I. sliced Wasserstein models
[2]. denoising diffusion models [33]. and autoregressive sequence models [34] [35), can be leveraged to learn GDEs. This is by no means exhaustive: any other conditional generative modelling approach
36 37 including those which will emerge in the future, can be used in the GDE framework.
5
Statistical and geometric properties of GDEs
GDEs aim to learn representations that separate the structure of the data- data-generating distribution from snite-sample noise, and to synthesize new data consistent with that structure. We formalize this dual role through two complementary perspectives. First, predictive sufficient statistics, GDEs act like learned Rao-Blackwellizttions that denoise sampling variability. Second, as statistical manifold embeddings, they interpolate between distributions along smooth geometric paths.
5.1 Learning approximate predictive sufficient statistic
The core objective of GDEs is to recover the true Hata-generating distribution P from finite samples by learning an aggregate representation that distills the structure of /' from sampling noise. This objective is captured by the notion of asymptotic predictive sufficiency sufficiency in terms of conditional independence:
2], which reformulates
PP A T(Sm)) P(Z A.A. (.)
m-+00
for all measurable A s X. In our setting, the encoder (S,m) serves as such a statistic, asymptotically determining P and marginalizing over sampling variability in Sm.
Predictive sufficiency implies Blacklla improve-
ment principle: conditioning any predictor on predic-
n
Naive
RB
GDE
tively sufficient statistic cannot increase predictive risk
10
4.72c-03
3.79e-03
3.12c-03
under convex loss (2). To illustrate this empirically, we
100
3.41e-04
2.76e-04
2.71e-04
consider X; ~ Pois(A) and predict IP(X,+1 0). The
1000
3,03e-05
2.81e-05
2.67e-05
10000
3.23e-06
2.64e-06
3.32e-06
baseline uses the observed frequency of X; 0, the GDE estimator conditions on the embedding and draws 10 synthetic samples upon which the bascline estimator is applied, and the Rao-Blackwellizz tions on the sufficient statistic T Σ, X,. In Tab. we can see that the GDE estimator outperforms the RB bound on the data manifold at low sample numbers. This suggests GDEs can act as data-driven analogues of Bto--Bacckeelizzinn
Table 1: MSE of Naive, RB, and GDE estimators for P(X = 0). X~ Poi(A).
(RB) estimator condi-
using synthetic sampling to magnify signal.
predictive sufficient statistic distills the structural properties of the meta-distri while marginal- izing over sampling variability in the observed data. Generative distribution embeddings achieve this practice: they recover consistent representations of underlying distributions, even across diverse domains and observational sample spaces.
5

---

Multinom.
MNIST
DNA
W3
1.00
1 nomdlized distt
0.50
0.00
Figure 3: Lg in GDE latent space compared to W2 distance. Normalized distances from the center, p (1/3, 1/3, 1/3).. The plots to the left show GDE L2 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 synthetic DNA dataset with 3 patterss respectively. Rightmost plot shows the Gaussian approximation for the W2 between multinomials.
We demonstrate this using the multinomial distribution. We learn GDEs of 3-dimensional multinomial distributions using mean saan-p deep sets encoder and a diffusion generator. The model's latent space is able to recover the structure of the multinomial simplex (Fig. B) Next, we use two real- world datasets with discrete class labels and conditionally sample observations according to label identities, which are drawn from the same family of 3-dimensional multinomial distributions. For both three-digit subset of MNIST and set of three synthetic DNA sequence patterns, GDBs (using 2D and ID convolutional encoders and diffusion and HyenaDNA generators, respectively) recover the same structure of the underlying multinomial simplex in the latent space. Despite coming from three different domains and using three vastly different architectures, the latent geometry learned between these experiments is nearly identical demonstrating GDEs capacity to learn signal from noise.
In fact, the learned geometry is rather particular: the L2 distance in GDE latent spaces in all three cases closely resemble W2 distances between multinomials (computed under a Gaussian approximation). This points to geometric interpretation of GDEs, bringing us to our second theoretical perspective.
trajectory
GDE trajectory
5.2 Learning manifold of distributions
From the geometric perspective, can try to understand GDEs by examining the struc- ture of their latent spaces. In Fig. 4 we take a "source" and a "target" distribution and com- pute the corresponding distribution embeddings Asre = = and Zigt- = {(.g). We then com- pute the linear interpolants 2t = (1-t) +tagt and push those back out into distribution space by sampling G(z)). The results are rather re- markable: the paths traced in distribution space by the linear latent interpolants closely resemble the optimal transport paths. This motivates second view of GDEs: they can generate syn- thetic distributions that interpolate smoothly be- tween observed populations. Empirically, la- tent trajectories in GDE space correspond to families of synthetic data that move coherently through probability space, providing control- lable mechanism for exploring or augmenting realistic generative scenarios.
5-
GMM OT trajectory
GDE trajectory
2
0
0
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.
Formally, let X' denote the sample space and P2(X)) the set of probability measures with finite second moment. The 2- W-Wssersstenn distance W2 makes (2()') geodesic metric space: in Euclidean settings also admits Riemannian interpretation [3. For meta-distribution Q over Pa (2)), let M C P2('') be set supporting Q: when M is a smooth submanifold, it inherits the metric induced by W2. Intrinsic geodesics on M are shortest paths constrained lie in M under this induced metric (and need not coincide with ambient W2 geodesics in P2()))

---

The encoder & can be viewed as an (approximate) smooth embedding Ø M + R", while the gener- ator parameterizes an approximate inverse that decodes latent trajectories into synthetic distributions along MM. Ideally, to would preserve the manifold's intrinsic geometry, as approximate isometry that maps Wasserstein distances between distributions to Fuclidean distances in latent space. To further assess this, beyond Fig. 4 we can compute the correlation between the latent and Wasserstein distances: for Gaussian distributions, latent-space Le distances correlate with true W2 distances at p 0.96; for 3-component Gaussian mixtures and W2 distances restricted to the mixture family (38), the correlation remains high (p = 0.76). This highlights a connection between our two perspectives:
Assume Q is supported on a d- dimensional statistical manifold M and the encoder induces a C¹ map é P(X) + RR (with the regularity conditions in App. D.3. Then Φ) M is asymptotically predictively sufficient when M is smooth embedding.
Proposition 3. (Informal statement of Theorem 3
key question this exploration prompts is precisely when GDE latent spaces are endowed with Wasserstein geometry. It is worth noting that classically, sufficient statisties and statistical manifolds are fixed once the model family is specified, and the geometry is independent observing particular distribution. In contrast, our hierarchical model involves meta-distributional prior, endowing the setting with a Bayesian flavor. GDEs' predictive sufficiency is therefore evaluated with respect Q: favoring statistics that preserve predictive information for distributions that are more common under Q. As a result, the learned representation and the synthetic data generated from become Q-weighted, allocating resolution to regions of M according to their probability. This adaptive weighting explains why GDEs can outperform the RB estimator in Tab.
the likelihood of
In earlier examples Q was approximately uni- form over a region of M, yielding a learned ge- ometry close P2 When is non-uniform, the geometry warps: high-density regions ex- pand to preserve finer distinctions, while low- density regions contract. We demonstrate this in our synthetic multinomial setting by train- ing GDEs on empirical distributions sampled from skewed Dirichlet task distributions, (2-5, 1, 1) and = (25, 1. 1). compared to the uniform
-2-5
25
1.00
0.75
0.0
0.2
0.00
Figure 5: Similar Fig. we show the GDE
distances of multinomials from We shift the prior asymmetrically by changing a while fixing a2 = a3 = 1. This shifts the focus of the model, leading to a different learned geometry.
( 3: 4)-
(1, (Fig. 5p. As expected, dis- tances stretch precisely where Q concentrates.
The takeaway is operational: the embedding is not only geometrically faithful but also prior-weighted. By choosing Q strategically (e.g., via the task-informed sampling in Sec. (3). we can bias the model toward capturing distributional properties most relevant to downstream objectives.
6 Applications
We first benchmark our approach and then demonstrate the generality of GDEs on tasks across the biological sciences, spanning several data domains: DNA sequences, protein sequences, gene expression data, and microscopy data. Throughout, we explore different combinations of encoder- generator pairs, see App. for a detailed discussion of architectures and training dynamics.
6.1 Benchmarking GDEs on synthetic distribution datasets
Table 2: Wasserstein reconstruction error across synthetic
The design space of GDEs is large: any distributionally invariant encoder can be coupled to any conditional gen- erative model. To guide our imple- mentation choices, we systematically benchmark architectures using syn- thetic datasets. Included in the bench- marked models are two existing meth-
distributional datasets. Computed as W2 for normal and GMM, and as Sinkhorn divergence for MNIST and FMNIST.
Model
Normal
GMM
MNIST
FMNIST
KME DDPM
0.04
2.17
80.46
111.01
W2 Wormhole
0.20
2.88
263.29
320.18
GDE
0.02
1.82
63.79
102.21
ods that GDEs generalize, kernel mean embeddings and Wasserstein Wormhole ||| [20)
7

---

We benchmark 30 combinations of encoders and generators on multivariate normal distributions in 5 dimensions. For evaluation we compute the Wasserstein reconstruction error from ground truth distribution by estimating means and covariance matrices from generated samples and using the closed-form for W₂ between Gaussians. We find that mean-pooled deep sets with skip-connections coupled with DDPM generators provide the highest quality generations, outperforming existing techniques. For synthetic distributions we present results for this architecture (see App. B.2)
In Table 2 we additionally benchmark this GDE architecture on three more sophisticated datasets: (I) 3-component Gaussian mixtures in 5 dimensions, (2) mixtures of MNIST [39 images according to categorical distributions of 3 classes, and (3) an analogous dataset using Mushion-MNIST [40]. For image datasets, where W2 distances are not tractable, we instead compute the Sinkhorn divergences between pretrained Resnet18 [4]] representations of generated and ground truth samples. In all cases, our chosen GDE architecture outperforms existing approaches.
6.2 GDEs enable semi-supervised distribution-leve representation learning
We next explore GDEs in motivating example: for learning patient representa- tions from single-nucleus RNA-seq atlas of the human prefrontal cortex [4. which profiled over 6.3 million nuclei from 1,494 donors across neurological and psychiatric conditions. We consider each donor's nu- clei as samples from an empirical distribu- tion, and each condition as a label we wish predict. As baseline, we first train supervised model to predict patient labels from nuclei sets using mean-pooled deep sets architecture using 10% of the available labelled data. We compare this with a semisupervised model implemented using GDEs. To construct a GDE model, we combine mean-pooled deep sets encoder with a CVAE generator, and train using the same 10% labelled data available the supervised model, along with the remaining 90% with labels withheld. semi-supervised GDEs outperform supervised baselines across all evaluation metries (Tab. 3)
Table 3: Patient label prediction from single-cell data. Semi-supervised GDEs improve performance.
Metric Supervised Semi-supervised
Accuracy ROC AUC
0.8791
0.8887
0.4872
0.5131
Fl Score
0.1293
0.1479
6.3 Modeling clonal populations in lineage-traced scRNA-seg data
While many methods have been developed for learning representations of single cells from scRNA-so seq data [42), methods for learning representations of cell populations remain relatively underexplored. This task is relevant to the analysis of lineage tracing data, where the unit of interest is clone, population of cells that arise from the same progenitor.
Using lineage-traced scRNA-seq data from mouse hematopoietic stem cells 43), we apply GDEs to learn clone-level representations by treating the set of cells within a clone as samples from an empirical distribu- tion. Following prior frameworks [44], evaluate the ability of representations to predict future clonal gene expression based the mutual information (MI) between a clone's 's representation at an early timepoint during differentiation and its representation a late timepoint. We find that GDEs with a CVAE genera- tor outperform Wasserstein Wormhole embeddings by over bits (Fig. 6). We next ask if this increase in predictive power due improved representations within certain cell types (e.g., neutrophils or monocytes). Decomposing MI estimates into their pointwise contributions [45), reveals contributions across the entire cell state space rather than any particular cell subtype (Fig. 6).
Wormhole
GDE
1.16 bits
3.31 bits
SPRINGZ
Figure 6: 2D embeddings of lineage- tracedd seRNA-seq data, hued by pointwise mutual information between clonal representation at early timepoint and clonal fate.
6.4 Predicting transcriptional responses to genetic perturbations
A central goal in genomics is to predict the transcriptional effects of genetic perturbations 46 We evaluate GDE for genetic perturbation prediction, using the Perturb-seq data of Replogle et al. 49 that profiled gene expression responses to CRISPRi knockdown of thousands of genes.
48

---

We consider the following task: given the identity of perturbation, predict the full distribution of transcriptional responses. We compare two approaches. In the first case, we train a linear model to predict the mean expression profile directly. In the second case, we predict the GDE embedding (trained on sets of cells subject to the same perturbation, via a Resnet Deep Sets encoder and CVAE generator as in Sec. 6.3 of the perturbation- -nnducc expression distribution and then recover the mean via a learned linear projection from the embedding space. both cases, we use a ridge regression on top of GenePT embeddings (50) to enable zero-shot generalization across perturbation conditions,
demonstrating that GDE improves both R2 and MSE in Tab, 4 See Appendix B.5 for full details.
6.5 Learning morphological cellular responses to genetic perturbations
We apply GDEs to pooled image-based CRISPR screening data from Funk et al. [51], which profiles the phenotypic effects of perturbing 5,072 essential human genes in HeLa cells. The dataset includes over 20 million single-cell microscopy images with four stains, capturing diverse phenotypic variation.
Table 4: GenePT predicting held- out perturbations mean expres- sion space and GDE latent space.
R2
MSE
Mean 0.378290
1.854997
Each perturbation induces a distribution over cell morphologies based on perturbation groupings. We treat these empirical distributions and train a GDE model to reconstruct them. To explore the role of inductive biases, we instantiate GDEs with two different priors: a spatial prior that models positional image structure (see App. B.6), and perturbation prior that captures latent variation across perturbation conditions. These approaches capture spatial and perturbation sets, respectively.
se GDE
0.421491
1,551414
0.457941
1.500731
Qualitatively, the model learns reproduce phenotypic features, including nuclear shape, cytoplasmic texture, and boundary sharpness across perturbations (Fig. 7 Quan- titatively, similar to Sec. [6.4]we hold out 30% of the most perturbative perturbations and use ridge regression with GenePT to enable zero-shot generalization across perturba- tions by predicting the GDE embedding. We then sample conditional on the predicted embedding and compute the nuclear signal intensity. The predictions on these held-out
Real
Demerabed
Figure 7: Real/generated DAPI images for the heldout RACGAPI knockout. perturbations achieved an R² 0.7055 and an MSE of 0.00068, indicating a strong zero-shot generalization of phenotypic outcomes.
6.6 Decoding yeast promoter sequence activity with GDEs
Real
Reconstructed
We next consider a large-scale dataset from a massively paral- lel reporter assay measuring tran- scriptional activity across 34 mil- lion randomly generated yeast promoter sequences [52). Each promoter consists of a random 80 nucleotide DNA sequence em- bedded in a fixed DNA scaf- fold and assayed for expression in yeast cells. Because the sequences are randomly sam- pled, there is no shared structure across examples signal lies entirely in how distributions over sequences give rise to distributions over expression levels, due to the presence of transcription factor binding sites (TFBS): short, position-specific DNA motifs that interact with transcription factors and control gene expression [52]
Expression Quantilee
142
PCA2
Density
PCA1
PWM Scores
Figure 8: The PCA (left) of the GDE latent space of quantile embeddings with underlying 34 million promoter sequences and the recovered distribution of TFBS (right) measured by motif counts in both the real and reconstructed data.
unconditional generative models cannot learn anything meaningful. Instead, the
We construct distributional learning task where each training example is a set of sequences sampled from narrow expression quantile; we hold out the top quantiles. We train GDE with convolutional network over the one-hot encoded sequences the encoder and HyenaDNA [35 as

---

the decoder. As shown in Fig. 8 the learned GDE embeddings reflect a smooth gradient across expression quantiles. Using the set of all known yeast TFBS [53] we can identify the motifs present in each of the real and generated sequences. Reconstructed motif distributions closely match those of the input, indicating that the model learns to represent biologically meaningful variation across promoter sets. Further details are available in App. B.8
6.7 Modeling spatiotemporal distributions of viral lineages
Powerful modeling approaches have been developed to represent individual protein sequences [54 55 56 57 34). Here, we show that the GDE framework can naturally lift these modeling approaches to learn representations of distributions of sequences. In particular, of SARS-CoV2 spike protein sequences over time and location. Using a dataset from the Global Initiative on Sharing All Influenza Data (GISAID) [58], we group sequences by sampling month and site location and treat each group as an empirical distribution over protein sequences. We embed these distributions using GDE which couples the ESM architecture [56] to a mean- meea--poled deep sets as the generator.
model distributions
as the encoder and a conditional ProGen2 architecture
As shown in Fig. the learned latent space organizes samples chronologically. suggesting that GDEs capture time-varying signal about sequence distributions. And indeed, this is ob- served quantitatively: ridge regression on GDE representations predicts the month of held out sequence distributions with mean absolute error (MAE) of 1.83 + 0.01 months, an improvement over the baseline of mean-poole ESM embeddings with MAE of 2.24 + 0.01 months (errors reported as mean = s.c.m. over 10 random train/test splits). See App [.9]B.9 for further details.
monbi since 2019-12
UMAP2
20
Similarly, we also observe spatial signal, albeit much weaker. An SVM trained to classify distributions by country achieves 0.28 + 0.001 accuracy from GDE representations, compared to 0.25 + 0.003 from mean-pooled ESM embeddings. Both approaches slightly outperform the baseline of predicting the most common dataset label ('USA' with accuracy 0.21).
Figure 9: GDE representations of protein sequence distributions. Each point corresponds to set of SARS-CoV2 spike sequences ob- tained from one lab in one month.
7 Discussion
We introduce generative distribution embeddings, framework that couples distribution iinirbttioon-inaaiiatt encoders with conditional generators to learn structured representations of distributions. Finite sample sets mapped by smooth embeddings that asymptotically identify the underlying distribution, enabling consistent reconstruction in the large-sample limit. We formalized these properties via connections to predictive sufficiency and statistical manifold embeddings, and proved that broad class of encoder architectures is asymptotically normal and unbiased when trained via plug-in loss.
We demonstrated GDEs across a diverse set of large-scale biological problems. These applications highlight the generality of GDEs and their ability to operate directly on measurement data while modeling population-level structure. Crucially, GDEs support flexible distributional constructions (e.g. spatial neighborhoods, time windows, expression quantiles), showing that a wide range of problems can be cast population-level modeling tasks. Code for model training and dataset preprocessing is available at this Github repository
Limitations GDEs rely on sensible choices of meta-distributional priors (i.e. construction of sets, Sec. 3. often requiring careful, comain-speciffc design. GDEs also pose practical engineering challenges (propagate gradients to the encoder through the generator, scaling large set sizes) discussed in App. A On the theoretical side, the current formalism assumes exchangeable samples, and does not admit non-i.i.d. samples within distribution. Regarding geometry, we provide empirical but not mechanistic evidence that GDEs learn isometries across domains.
Extensions GDEs can serve as tool for generalization (akin to meta flow matching [19), can be expanded settings where the i.i.d. assumption within sets of samples does not hold, and extended to semi-supervised settings. More broadly, GDEs point toward questions at the intersection of empirical process theory, information geometry, and generative modeling; we hope this connection can be explored more deeply in future work.
10

---

Acknowledgements Funding: J.S.G. and O.O.A. are supported by NIH grants R01 RO1-EB031957, R01-AG074932, and 201-GM148745; Harold & Leila Y. Mathers Charitable Foundation; Rett Syndrome Research Trust; The Gordon and Betty Moore Foundation; Impetus Grants; Cystic Fibrosis Foundation Pioneer Grant; Google Deepmind; Sanofi; Yosemite; Michelson Foundation; Hevolution Foundation; American Federation for Aging Research; Pivotal Life Sciences; and the MGB Gene and Cell Therapy Institute. G.G. is supported by a Tayebati Fellowship. N.F. is supported by the NSF GRFP.
References
[1] Kei Takeuchi and Masafumi Akahira. Characterizations of prediction sufficiency (adequacy) in terms of risk functions. The Annals of Statistics, 3(4):1018-1024, 1975.
[2] Erik N Torgersen. Prediction sufficiency when the loss function does not depend on the unknown parameter. The Annals of Statistics, 5(1):155-163, 1977.
[3] Felix Otto. The geometry of dissipative evolution equations: the porous medium equation. 2001.
[4] John F Fullard, Prashant Nm, Donghoon Lee, Deepika Mathur, Karen Therrien, Aram Hong, Clara Casey, Zhiping Shao, Marcela Alvia, Stathis Argyriou, et al. Population-scale cross- disorder atlas of the human prefrontal cortex at single-cell resolution. Scientific Data, 12(1): 954, 2025.
[5] Netanel Loyfer, Judith Magenheim. Ayelet Peretz, Gordon Cann, Joerg Bredno, Agnes Klochendler, Ilana Fox-Fisher, Sapir Shabi-Porat, Merav Hecht, Tsuria Pelet, Joshua Moss, Zeina Drawshy, Hamed Amini, Patriss Moradi, Sudharani Nagaraju, Dvora Bauman, David Shveiky, Shay Porat, Uri Dior, Gurion Rivkin, Omer Or, Nir Hirshoren, Einat Carmon, Alon Pikarsky, Abed Khalaileh, Gideon Zamir, Ronit Grinbaum, Machmud Abu Gazala, Ido Mizrahi, Noam Shussman, Amit Korach, Ori Wald, Uzi Izhar, Eldad Erez, Vladimir Yutkin, Yaacov Samet, Devorah Rotnemer Golinkin, Kirsty L Spalding, Henrik Druid, Peter Arner, A M James Shapiro, Markus Grompe, Alex Aravanis, Oliver Venn, Arash Jamshidi, Ruth Shemer, Yu- val Dor, Benjamin Glaser, and Tommy Kaplan. A DNA methylation atlas of normal human cell types. Nature, 613(7943):355-366) 4 January 2023. ISSN 0028-0836, 0ew.000831.1666610 4776-4677 doi: 10. (0338//415860022-05555533
[6] Gokul Gowri, Peng Yin, and Allon M. Klein. Measurement noise scaling laws for cellular representation learning, 2025. https::s:///arx wrw.://////a//eeee 02726
[7] Alex Smola, Arthur Gretton, Le Song, and Bernhard Schölkopf. A Hilbert space em- bedding for distributions. In Lecture Notes in Computer Science, Lecture notes in com- puter science, pages 13-31. Springer Berlin Heidelberg, Berlin, Heidelberg, 2007. ISBN 9788540052224998834477222
doi: 10. 00007/788-3-54--5225---555
[8] Krikamol Muandet, K Fukumizu, Francesco Dinuzzo, and B Scholkopf. Learning from dis- tributions via support measure machines. Neural Information Processing Systems, 25:10-18, 29 February 2012.
[9] Junier B Oliva, B Póczos, and J Schneider. Distribution to Distribution Regression. International Conference on Machine Learning, (063)):049-1055 16 June 2013.
[10] Zoltan Szabo, Arthur Gretton, Barnabas Poczos, and Bharath Sriperumbudur. Two-stage sampled learning theory on distributions. In Artificial Intelligence and Statistics, pages 948-957. PMLR, 21 February 2015.
[11] Krikamol Muandet, Kenji Fukumizu, Bharath Sriperumbudur, and Bernhard Schölkopf. Kernel mean embedding of distributions: A review and beyond. Foundations and Trendy@ in Machine Learning, (0(1-2):1-141 2017. ISSN 1935-8237,1935-1 1935--2555. doi: 10.1561/2200000000.
[12] Harrison Edwards and Amos Storkey. Towards a Neural Statistician. International Conference on Learning Representations, 6 February 2017.
[13] S Amari and H Nagaoka. Methods of information geometry. 191, 2000.

---

[14] K Carter, R Raich, W Finn, and A Hero. FINE: Fisher Information Nonparametric Embedding. IEEE transactions on pattern analysis and machine intelligence, 31(11):0933-009 14 February 2008. ISSN 0162-8828,1923 1939-3539. doi: 10. 10.1109/TPAMI.2009.66
[15] Yonghyeon Lee, Seungyeon Kim, Jinwon Choi, and F Park. A statistical manifold framework for point cloud data. International Conference Machine Learning, pages 2022.
[16] Hannes Stärk, Bowen Jing, Chenyu Wang, Gabriele Corso, Bonnie Berger, R Barzilay, and T Jaakkola. Dirichlet flow matching with applications to DNA sequence design. International Conference on Machine Learning, 235:46495 46513, 8 February 2024. doi: 10.48550/arXiv. 2402.05841.
[17] Oscar Davis, Samuel Kessler, Mircea Petrache, I Ceylan, Michael M Bronstein, and Bose. Fisher flow matching for generative modeling over discrete data. Neural Information Processing Systems,
139084, 23 May 2024. doi: 10.485500arXiv.2405.. 14664.
[18) Doron Haviv, Aram-Alexandre Pooladian, Dana Pe' and Brandon Amos. Wasserstein flow matching: Generative modeling over families of distributions, 2024. URL https://arxiv org/abs//411.00699
[19] Lazar Atanackovic, Xi Zhang, Brandon Amos, Mathieu Blanchette, Leo J Lee, Yoshua Bengio, Alexander Tong, and Kirill Neklyudov. Meta Flow Matching: Integrating Vector Fields on the Wasserstein Manifold. In The Thirteenth International Conference on Learning Representations, 4 October 2024.
[20] D Haviv, Russell Z Kunes, Thomas Dougherty, Cassandra Burdziak, T Nawy, Anna Gilbert, and D Pr'er. er. Wasserstein Wormhole: Scalable optimal transport distance with transformers. International Conference on Machine Learning, 139777777,, 15 April 2024. doi: 10. 48550/arXiv.2404.009 14.
[21] John Fisher, III, A Ihler, and Paul Viola. Learning informative statistics: nonparametnic approach. Neural Information Processing Systems, pages 900-906, 29 November 1999, doi: 10.5555/3009557.300999
[22] Paul Joyce and Paul Marjoram. Approximately sufficient statistics and bayesian computation. Statistical applications in genetics and molecular biology, 7(1):Article26, 30 August 2008. ISSN 1544-6115,2194-63333 doi: 10.2202/1544.-1155..53
[23] Justin Alsing, Benjamin Wandelt, and Stephen Feeney. Massive optimal data compression and density estimation for scalable, likelihood-free inference in cosmology. Monthly notices of the Royal Astronomical Society, (7743))2874-2888) 1 July 2018. ISSN 0033--871, 115--877,,13652966 1365-2996. doi: 10. 10/33/mmaas//tt811...
[24] Yanzhi Chen, Dinghuai Zhang. Michael U Gutmann, Aaron Courville, and Zhanxing Zhu. Neural Approximate Sufficient Statistics for Implicit Models. International Conference on Learning Representations, 2 October 2020.
[25] Maxime Peyrard and Kyunghyun Cho. Meta-Statistical Learning: Supervised Learning of Statistical Inference, 2025.
[26] M Zaheer, Satwik Kottur, Siamak Ravanbakhsh, B Póczos, R Salakhutdinov, and Alex Smola. Deep Sets. Advances in neural information processing systems, 30, 10 March 2017.
[27] Wagstaff, F Fuchs, Martin Engelcke, Michael A Osbome, and Posner. Universal approxima- tion of functions on sets. Journal of machine learning research: JMLR, 23((51): 5 July 2021. ISSN 1532-4435.1533-7955
[28] Lily Н Zhang, Veronica Tozzo, J Higgins, and R Ranganath. Set norm and equivariant skip connections: Putting the deep Deep Sets. International Conference on Machine Learning, (62:26559-26574 23 June 2022. doi: 10.48550/arxiv.22066 11925.
12

---

[29] Drew A Hudson, Daniel Zoran, Mateusz Malinowski, Andrew K Lampinen, Andrew Jaegle, James L McClelland, Loic Matthey, Felix Hill, and Alexander Lerchner. SODA: Bottleneck diffusion models for representation learning. In 2024 IEEE/CVE Conference Computer Vision and Pattern Recognition (CVPR), pages 231 2-12777. IEEE, 16 June 2024. doi: 10.
11 1109/cvp:52733..22 22.207222773222.2111111
[30] Diederik P Kingma and M Welling. An Introduction to Variational Autoencoders. Found. Trends Mach. Learn., 377))307-332, June 2019. doi: 10. 10.1561/2200000556
[31] Aude Genevay, Gabriel Peyre, and Marco Cuturi. Learning Generative Models with Sinkhorn Divergences. In International Conference Artificial Intelligence and Statistics, pages 1608- 1617. PMLR, 31 March 2018.
[32] Soheil Kolouri, Kimia Nadjahi, Umut Simsekli, R Badcau, and G Rohde, Generalized Sliced Wasserstein Distances. Neural Information Processing Systems, 32:261-272, 1 February 2019.
[33] Jonathan Ho, Ajay Jain, and P Abbeel. Denoising Diffusion Probabilistic Models. Neural Information Processing Systems, abs/2006. 123396684 1231112403660655111 19 June 2020.
[34] Erik Nijkamp, Jeffrey A Ruffolo, Eli N Weinstein, Nikhil Naik, and Ali Madani. ProGen2: Ex- ploring the boundaries of protein language models. Cell systems, (4(11))968-97)) 15 Novem- ber 2023. ISSN 2405-47 440--4722,2405
doi: http6//..cll.22221
[35] Eric D Nguyen, Michael Poli, Marjan Faizi, A Thomas, Callum Birch-Sykes, Michael Wornow, Aman Patel, Clayton M Rabideau, Stefano Massaroli, Y Bengio, Stefano Ermon, S Baccus, and Christopher Ré. HyenaDNA: Long-range genomic sequence modeling at single nucleotide resolution. Neural Information Processing Systems, 08443177-43201 27 June 2023. doi: 00448550/rXxiv.336..555 15794.
[36] Ian Goodfellow, Jean Pouget-Abadie, Mehdi Mirza, Bing Xu, David Warde-Farley. Sherjil Ozair, Aaron Courville, and Yoshua Bengio. Generative adversarial networks. Communications the ACM, 13:111413,,,4, 22 October 2020. ISSN 1001-7772,5557-711, doi: 10. 145/3422222..
[37] Lipman, Ricky T Q Chen, Heli Ben-Hamu, Maximilian Nickel, and Matt Le. Flow Matching for generative modeling. International Conference on Learning Representations, abs/2210.02747, 6 October 2022.
[38] Julie Delon and Agnes Desolneux. A wasserstein-type distance in the space of gaussian mixture models. SIAM Journal on Imaging Sciences, 13(2):9369970, 2020.
[39] Yann LeCun, Corinna Cortes, and CJ Burges. Mnist handwritten digit database. ATT Labs [Online]. Available: http://ynnn. http://yanm.eecum.c////// www..//cm..o///eeb//miissso 2, 2010.
[40] Han Xiao, Kashif Rasul, and Roland Vollgraf. Fashion-mnist: a novel image dataset for benchmarking machine learning algorithms. CoRR, abs/1708.07747,, 2017. URL http://
arxiv. abs/ 1atbas/17700 1708.1700 07747
[41] Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In 2016 IEEE Conference on Computer Vision and Pattern Recognition (CVPR), pages 770-778. IEEE, June 2016. ISBN 2016.90.
doi: 10. 10.11 109/cvpr.
[42] Romain Lopez, Jeffrey Regier, Michael Cole, Michael I Jordan, and Nir Yosef. Deep genera- tive modeling for single-cell transcriptomics. Nature methods, 15(12):1053-100 10588, December doi: 10.
2018. ISSN 1548-7091, 709., 5155-
[43] Caleb Weinreb, Alejo Frrttiuuz-Fraaiielll, Fernando D Camargo, and Allon M Klein. Lineage tracing on transcriptional landscapes links state to fate during differentiation. Science, 367 (6479), 14 February 2020. ISSN 1036-8075,1095-9900 doi: 10. 126/sciience.aw3338))
[44] Gokul Gowri, Xiao-Kang Lun, Allon Klein, and Peng Yin. Approximating mutual informa- tion of high-dimensional variables using learned representations. In A Globerson, I. Mackey, Belgrave, A Fan, Paquet, J Tomezak, and Zhang, editors, Advances in Neural Information Processing Systems, volume 37, pages 132843-13288 Curran Associates, Inc., 2024.
13

---

[45] Xianghao Kong, Ollie Liu, Han Li. Dani Yogatama, and Greg Ver Steeg. Interpretable Diffu- sion via Information Decomposition. In The Twelfth International Conference on Learning Representations, 13 October 2023.
[46] Omar O Abudayyeh and Jonathan S Gootenberg. Programmable biology through artificial intelligence: from nucleic acids to proteins to cells. Nature Methods, 21(8):1384-11 1386, 2024.
[47] Charlotte Bunne, Yusuf Roohani, Yanay Rosen, Ankit Gupta, Xikun Zhang, Marcel Roed, Theo Alexandrov, Mohammed AlQuraishi, Patricia Brennan, Daniel B Burkhardt, et al. How to build the virtual cell with artificial intelligence: Priorities and opportunities. Cell. 187(25): 7045-7063, 2024.
[48] Julia Joung, Sai Ma, Tristan Tay, Kathryn R Geiger-Schuller, Paul Kirchgatterer, Vanessa K Verdine, Baolin Guo, Mario A Arias- Garcia, William E Allen, Ankita Singh, et al. A transcrip- tion factor atlas of directed differentiation. Cell, 186(1):209-22) 229, 2023.
[49] Joseph M Replogle, Reuben A Saunders, Angela N Pogson, Jeffrey A Hussmann, Alexander Lenail, Alina Guna, Lauren Mascibroda, Eric J Wagner, Karen Adelman, Gila Lithwick-Yann -anai, et al. Mapping information-rich genotype-phenotppe landscapes with genome-scale perturb-seq. Cell, (85(1)))2559 2575, 2022.
[50] Yiqun Chen and James Zou. Genept: a simple but effective foundation model for genes and cells built from chatgpt. bioRxiv, pages 2023-10, 2024.
[51] Luke Funk, Kuan-Chung Su, Jimmy Ly, David Feldman, Avtar Singh, Brittania Moodie, Paul C Blainey, and lain M Cheeseman. The phenotypic landscape of essential human genes. Cell, 185 (24):4634 4653, 2022.
[52] Carl G de Boer, Eeshit Dhaval Vaishnav, Ronen Sadeh, Esteban Luis Abeyta, Nir Friedman, and Aviv Regev. Deciphering eukaryotic gene-regulatory logic with 100 million random promoters. Nature biotechnology, (8(1):56-66 65,,666666666 2020.
[53] Carl G De Boer and Timothy R Hughes. Yetfasco: a database of evaluated yeast transcription factor sequence specificities. Nucleic acids research, 40(D1):D169 D179,, 2012.
[54] Adam Riesselman, John B Ingraham, and Debora S Marks. Deep generative models of genetic variation capture the effects of mutations. Nature methods, 15(10):816-82) October 2018. doi: (0.10388/441592-018--0333..
ISSN 1548-7091,1548-777 7105.
[55] Tristan Bepler and Bonnie Berger. Learning the protein language: Evolution, structure, and function. Cell systems, (5666665466). 16 June 2021. ISSN 2405-4712,2405-4222 doi: 10. 00.1016f..ccs.222..0.0111
[56] Alexander Rives, Joshua Meier, Tom Sercu, Siddharth Goyal, Zeming Lin, Jason Liu, Demi Guo, Myle Out, C Lawrence Zitnick, Jerry Ma, and Rob Fergus. Biological structure and function emerge from scaling unsupervised learning to 250 million protein sequences. Proceedings of the National Academy of Sciences of the United States of America, 1 118(15), 13 April 2021. 18.
ISSN (027-8424,1091-6666
doi: 10.
[57] Zeming Lin, Halil Akin, Roshan Rao, Brian Hie, Zhongkai Zhu, Wenting Lu, Nikita Smetanin, Robert Verkuil, Ori Kabeli, Yaniv Shmueli, Allan Dos Santos Costa, Maryam Fazel-Zarandi, Tom Sercu, Salvatore Candido, and Alexander Rives. Evolutionary-scall prediction of atomic- level protein structure with a language model. Science, 379(6637): ISSN 0036-8075, 1095-9920 doi: 10. 10.1126/ciiencc.aee25777
17 March 2023.
[58] Yuelong Shu and John McCauley. GISAID: Global initiative on sharing all influenza data from vision to reality. Euro surveillance bulletin Europeen sur les maladies transmissibles [Euro surveillance European communicable disease bulletin), J. 22(13), 30 March 2017. ISSN 1155-499X11560-9911 doi: 022307715660-917..5...771777777.
[59] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, tukasz Kaiser, and Illia Polosukhin. Attention all you need. Advances in neural information processing systems, 30, 2017.
14

---

[60] Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 770-778, 2016.
[61] Jacob Devlin. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint arXiv: 180.:0880.04805, 2018.
[62] Zeming Lin, Halil Akin, Roshan Rao, Brian Hie, Zhongkai Zhu, Wenting Lu, Nikita Smetanin, Robert Verkuil, Ori Kabeli, Yaniv Shmueli, al. Evolutionary-scale prediction of atomic-level protein structure with language model. Science, 37966637):1123-11111 2023.
[63] Thomas Wolf, Lysandre Debut, Victor Sanh, Julien Chaumond, Clement Delangue, Anthony Moi, Pierric Cistac, Tim Rault, Rémi Louf, Morgan Funtowicz, et al. Huggingface's 's transform- ers: State-of-the-ar natural language processing. arXiv preprint 1xXi::1110.05571,2000 2019,
[64] Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever, et al. Language models are unsupervised multitask learners. OpenAl blog. 18)::9 2019.
[65] Alexander Kraskov, Harald Stögbauer, and Peter Grassberger. Estimating mutual information. Physical review. E, Statistical, nonlinear, and soft matter physics, 69(6 Pt 2):066138, June 2004. ISSN 1539-3755. doi: 10.1103/Phy.RcvE.69.6666.
[66] Netanel Loyfer, Jonathan Rosenski, and Tommy Kaplan, wgbstools: A computational suite for DNA methylation sequencing data representation, visualization, and analysis. bioRxiv, page 15.4.005.8.593332 10 May 2024. doi: 10.1 10.1101/2024005..85591111
[67] Aad W Van Der Vaart and Jon A Wellner, Weak convergence. Springer, 1996.
15

---

A Architectures and Training Dynamies
In this section we outline some general details about the architectures and training dynamics for GDEs. In the following section we will give more detailed explanations about each specific experiment, in addition to full details available the codebase. All of these findings are somewhat provisional, and there is significant scope for future work to further explore these design choices, but we hope this is useful complement to our codebase for researchers trying to train their own GDEs.
A.1 Encoder Architectures
Our framework utilizes permutation iiniutttttiioniniiannntn encoders map input sets S... {I1, ...y Lmh where each æ, R4, fixed-dimensional latent representation R We primarily employ several types of set encoders, including variants based on self-attention, Graph Neural Network (GyNl-style pooling, and residual connections. All encoders typically conclude by applying final pooling operation (c.g., mean pooling) across the element representations, followed by a linear projection and non-linearity (e.g., SELU) produce the final latent vector
A.1.1 Simple Self-Attention Encoder
This encoder provides baseline transformer- approach. It first applies a linear layer followed by a SELU activation project input elements x; into hidden dimension H. It then processes these representations through series of multi- head Self-Attention blocks [59]. This architecture directly models pairwise interactions within the set.
A.1.2 Simple GNN Encoder
The simple GNN-style encoder offers alternative based iterative pooling and non-linear transformations, distinct from the standard DeepSets [26] sum-decooposiition It starts with an MLP projection into the hidden dimension Hl. Subsequently, it applies a sequence of layers, each performing pooling operation across the set followed by an MLP. This structure iteratively refines element representations based aggregated information.
Our theoretical framework (see Appendix D.2) justifies the use of pooling operations that correspond to M/Z-estimators. We focus mean pooling but additionally implement median pooling as an illustrative example. Notably, max pooling is generally not suitable in this context its non-differnntiabiiiyy breaks the convergence guarantees we interested in for Eq. (T) see the remarks in App D.2 for details. Future work might thoroughly explore which pooling operations lead to the greatest flexibility and stability for distribution embedding.
Pooling Operations:
A.1.3 ResNet-GNN Encoder
To improve gradient flow and enable deeper architectures, we enhance the GNN-style encoder with residual connections. This encoder first projects each input element x; into using an MLP. then processes the set through series of blocks where each block & computes representation (x( for each element The core operation within a block uses mean pooling (or median pooling). Inspired by ResNet [60 28 we incorporate skip connections. The input to block k includes the output from the previous block (k-11,, a linear projection of the original input æ, and the output of the initial MLP projection. Formally:
intermediate
h (M) LayerNorm (PooledFC(I) + 1 (k-1) + Linears(x))
where (0) is the output of the initial input projection combined with projection of x, followed by Layer Normalization. This structure ensures the original input signal is preserved.
A.1.4 CosNet-Transformrr Encoder
This variant follows the same residual structure as the ResNet-MLPP encoder but replaces the layers with standard multi-head Self- blocks [59]. This potentially allows the model to learn more complex interactions while benefiting from the improved training dynamics of residual connections. The skip connection mechanism remains identical to the ResNet-MLP version.
16

---

A.1.5 Encoder Comparison
Transformer-based encoders (Simple Self-Attention and ResNet- Transformer) often leverage pre- trained weights effectively and can converge in fewer epochs compared to GNN-style approaches. However, this typically comes at higher computational cost per epoch and during inference due to the quadratic complexity of self-attention with respect to set size m. With sufficient training, we find that the GNN-based architectures, particularly the ResNet-GNN, achieve strong performance, often rivaling the transformer variants while being more computationally efficient for large sets.
Alternative Generative Strategies and Sampling The Wasserstein Wormhole (0 uses a self-attention decoder with fixed positional embeddings that can map the latent back to samples. One potential method replaces fixed positional embeddings with samples drawn from a simple distribution (e.g.. Gaussian) transforming this into true generator. But this incurs substantial computational costs (e.g.. quadratic cost the number generated samples for attention-based sampling decoders), and is not clear this would lead to significant improvements in performance.
also becomes less obvious how to adapt existing generator architectures using this approach. One option is to use self-attention to construct sample-specific condtional signals from the latent and the noise vector, and then condition the generator on this signal. This significantly more complex, and is not clear that this would lead to significant improvements in performance.
A.2 Adapting Pre-trained Models
Our framework is designed to flexibly incorporate pre-trained models, leveraging their learned representations and generative capabilities. We adapt pre-trained models for both the encoder and the generator components.
A.2.1 Encoder Adaptation
For tasks involving complex input modalities like natural language or protein sequences, we can utilize pre-trained transformer-basee encoders such as BERT [6]] ESM 62 as powerful feature extractirs. These pre-trained models can serve as the initial feature extraction layer, whose outputs (h1...., hN) are then fed into the subsequent aggregation layers of our encoders (e.g., ResNet- GNN or ResNet-Trane Transforrrr,,,,, see subsection A.I.
The adaptation process typically involves:
1. Loading Pre-trained Weights: We load the desired pre-trained encoder model using standard libraries like Hugging Faces transf ormers 633
2. Feature Extraction: For each element x; in the input set X = (x1,..,, we pass through the pre-trained transformer obtain a contextualized representation h4. Often, the output embedding corresponding to a special token (like [CLS] in BERT) or the mean/max- pooled output of the final hidden states is used.
3. Set Aggregation: These element-wise feature vectors (fi,..., AN } are then fed into the subsequent layers of our chosen set encoder (e.g., ResNet-MLP or ResNet-Traasforme layers) which perform the permutation-invaaian aggregation to produce the final latent representation 2.
4. Fine-tuning (Optional): Depending on the task and dataset size, the pre-trained encoder's weights might be kept frozen initially or fine-tuned jointly with the rest of the model during end-to-end training.
A.2.2 Generator Adaptation and Conditioning
A core strength of our approach is the ability to use large pre-trained causal language models (LMs), such GPT-2 [64] ProGen2 34 specialized models like HyenaDNA (35), the conditional
generator pe (1(z).
The adaptation involves:
Loading Pre-trained Weights: We load the chosen pre-trained causal LM and its associated tokenizer using "transformers' (63
17

---

2. Prefix Conditioning: The primary challenge is to effectively condition the generators output on the latent set representation produced by the encoder. In practice, we find prefix tuning to be an effective and widely applicable method. The latent vector E R. is projected, typically via a small MIP W., into one or more vectors p = W2(+) that have the same hidden dimension as the LM. These projected vectors p are then treated as continuous "prefix" embeddings prepended to the actual input sequence embeddings E(x<r) before they are processed by the transformer layers. The model learns to interpret this prefix as the conditioning signal specifying the target distribution. Mathematically, the imput embedding sequence the transformer becomes [p: Ex adjusted accordingly to allow all sequence tokens ø<r to attend to the prefix p.
)]. The attention mask is
3. Fine-tuning: The pre-trained generator weights can be either frozen or fine-tuned. Fine- tuiigg the entire model allows the LM adapt its generation process based the condi- tioning prefix Freezing the LM backbone and only training the conditioning projection Wp (and potentially adapter layers) can be more parameter-efficient.
A.3 Training Details and Considerations
A.3.1 Learning Rate Schedule
For simpler models we a fixed learning rate, but for more complex models we typically employ cosine annealing learning rate schedule during training. This involves starting with an initial learning rate and gradually decreasing it towards zero following a cosine curve over the course of training epochs. This schedule is often effective in achieving stable convergence and good final performance. In general we have found that whatever the current state of the art for training the (unconditional) generator is, that will generally give good results when learning the encoder-generane jointly.
A.3.2 Performance and Convergence
Our experiments generally indicate that this training setup, combined with the described architectures and adaptation strategies, leads strong performance across various tasks and datasets presented in the main paper. As noted in subsection A.5. the choice of encoder can impact convergence speed and computational cost.
A.3.3 Set Size and Batching Trade-off's
We observe that achieving optimal performance sometimes necessitates using large input set sizes (N). However, processing large sets can significantly increase the computational and memory requirements per batch, particularly for the attention mechanisms in transformer-basee This often forces reduction in the overall batch size to fit within hardware constraints. Smaller batch sizes can, turn, lead to increased variance the loss gradients, potentially slowing down or destabilizing training. Careful tuning of the set size N, batch size, and learning rate parameters is often required to balance performance and training efficiency for a given task and hardware setup.
encoders or generators.
A.3.4 Gradient Propogation Challenges
potential challenge arises, particularly with deeper encoder and generator architectures. The encoder only receives learning signal indirectly through the generator via the shared latent variable z. If the generator itself struggles to utilize the latent information effectively, or if the dimensionality L of creates information bottleneck, the gradients flowing back to the encoder can become weak or noisy. This can make training deep encoders difficult. Addressing this might require more sophisticated generator architectures capable of integrating the latent information more effectively or alternative training schemes with auxiliary losses directly on the encoder. We found these issues in the simple encoder architectures, but they seemed to be alleviated in the ResNet-based architectures.
A.4 Implementation recipe for GDEs
The GDE framework is instantiated by pairing a distributionally invariant encoder with conditional generative model. The following steps outline a general recipe for building GDEs across diverse data domains:
18

---

1. Sample from the metadistribution (construct sets) Group raw data into sets S {z where each reflects draw from an unknown latent distribution P. Groupings can be based on discrete metadata (e-g., text by author, reviews by rating, images by label, cell clones, gene perturbations) or continuous metadata (e-g-, time, location, expression quantiles). Sets need not be mutually exclusive, meaning single data point can belong to multiple sets.
2. Choose distributionally invariant encoder Select or construct distributionally invariant encoder &. This selection generally involves (1) using an architecture for element- wise embeddings and (2) pooling across element-wit embeddings with sample mean (or other M-estimate). We found that architectures with multiple pooling layers, where each layer's pooled output is concatenated with the element-wise embeddings, were particularly effective. This contrasts with pure DeepSets-style architectures that only pool once at the final layer. For deeper architectures, we have found that including skip-connections improves performance, especially if the generator is also relatively deep network.
3. Build a conditional generator The generator G "decodes" from latent space back to the sample space. It should be conditionable on z E(S)
4. Train via plug-in loss Optimize the generator to minimize the generator loss function E( ( ((E(Sm)) This loss should be the standard training objective for the conditional generator. This plug-in loss encourages reconstruction of the true distribution.
A.5
Encoder and Generator Architectures by Experiment
The encoder-generator pairs used for each application the paper are shown the table below.
Sec.
Task
Set Construc-
Encoder Arch.
Generator Arch.
Notes
tion
6.1
MNIST,
FM-
Same class
image
see Table
Table 1
Synthetic data benchmark
NIST
6.2
Lineage-traced
Same
cell
ResNet-GNN
CVAE
scRNA-seq
clones
6.3
Genetic pertur- bation (scRNA-
Same perturba- tion
ResNet-GNN
CVAE
seq)
6.4
Morphological responses (cell images)
Same perturba- tion
2D Conv-GNN
DDPM (U-Net)
6.5
Tpssue- methylation
Same patient; Same tissue
ID Conv-GNN
HyenaDNA
Uses prefix con- ditioning
type Expression quantile (con- tinuous)
6.6
Yeast promoter quantile decod- ing Viral protein spatiotemporal
ID Conv-GNN
HyenaDNA
Uses prefix con- ditioning
6.7
Same sampling month and loca-
ESM mean pooling
ProGen2
Uses prefix con- ditioning
modeling
tion
19

---

B Experiments
B.1
Determining tissue-specific methylation signatures from bisulfite sequencing reads
Analyzing sequencing data typically extensive preprocessing, including alignment to reference genome. GDEs present alternative, where sequencing reads can be modeled directly- without alignment other preprocessing steps. demonstrate capability, we show that GDEs can detect sasuu--pecific DNA methylation patterns directly from bisulfite sequencing (BS-seq) reads. BS-seq measures methylation indirectly through substitution errors: methylated cytosines remain unchanged. while unmethylated cytosines are substituted as thymines. Using publicly available methylation data from diverse tissues (5), simulate sample-specific BS-seq read distributions imposing corresponding base substitutions the reference genome (see Appendix |B.7 for details).
Critically, we do not provide the GDE model with any explicit information about methylation signals, the structure of the experimental assay, or a reference genome. The model has access only to sets of sequencing reads grouped by both patient and tissue type. For the GDE model architecture, we choose a ID convolutional network encoder, and the decoder is HyenaDNA model [3S]. To support large-scale inference over tens of millions of reads per patient, process 200,000 reads time through the encoder and aggregate the resulting embeddings using simple mean, justified by Theorem 2 This design allows the model to scale efficiently while preserving distributional fidelity.
Our approach enables end-te and-to-ene to-end learning of methylation signatures from tissue-specific read distributions. There are two levels of tissue classification, coarse level with 37 categories and a fine-grained level classification with 83 tissues. Training linear classifier top of the GDE latent space, achieve test accuracy of 60% coarse task 35% on fine-grained classification.
B.2 Additional semi-synthetic experimental results
W2 Distance GDE Distance
GMM W; Distance vs GDE Distance
Spearman Correlation: 0.96
Spearman Correlation: 0.76
W2 Distance
Distance
4
3
2
0
2
GDE Distance
GDE Distance
Figure 10: Left: Distance correlation showing high alignment between latent GDE distances and analytical Wa distances (Spearman P 0.96). Left: Distance correlation showing high alignment between latent GDE distances and the OT-GMM distance 38 which a W2 metric restricted to the subspace of GMMs (Spearman p = 0.76).
20

---

0.20
0.15
1- corr
0.10
0.05
0.00
10
100
10
81
Figure 11: Expanding on Fig. 5 we show that the Pearson correlation between the W2 (computed via normal approximation) and the latent GDE distances decreases as a1 deviates from while keeping fixed ag az = 1.
Table 5: W2 reconstruction error of 30 possible GDE implementations (including two existing methods generalized by GDE, Wasserstein Wormhole and kernel mean embeddings) on 5-dimensional multivariate Gaussians. Covariance matrices sampled from Wishart distribution with scale of 1, and means sampled uniformly from [0, 5]. Further results included Table
Gen. W Enc.
Mean
Kernel mean
GNN
Med.-GNN
ResNet-GNN
SelfAttn.
Sinkhorn
0.05
0.14
0.09
0.10
0.05
0.06
Sliced W2
0.03
0.04
0.07
0.07
0.03
0.04
CVAE
0.16
0.16
0.19
0.20
0.15
0.17
DDPM
0.03
0.04
0.06
0.05
0.02
0.07
Wormhole
0. 1.14
0.15
0.72
0.49
0.14
0.20
Table W22 reconstruction error (mean = s.e.m. over trials) for possible GDE implementations (including two existing methods generalized by GDE, Wasserstein Wormhole and kernel mean embeddings) on 5-dimensional multivariate Gaussians. Covariance matrices sampled from Wishart distribution with scale of 0.1, and means sampled uniformly from [0, 5].
Gen. J W Enc. +
Kernel mean
GNN
ResNet-GNN
Self-Attn.
CVAE
0.15 0.011
0.12 0.006
0.12 0.009
0.11 0.007
DDPM
0,15 # 0.008
0.13 + 0.020
0.09 + 0.003
0,10 + 0.005
Direct SW
0.15 0.008
0.007
0.13 0.009
0.15 =
Direct Sinkhorn
0.29 : 0.008
0.22
0.17 + 0,005
0.19
Wormhole
0.23 0.021
0.72 0.090
0.24 0.011
0.34 # 0.021
B.3 Donor-level representation learning experiments
B.3.1 Data preprocessing
We use singlc-nucleus RNA-seq data from the Population-seaa sccatiion--clcc cross-disorder atlas of the human prefrontal cortex [4, which profiles over 6.3 million nuclei from 1,494 donors across 33 ncurological and psychiatric conditions. The dataset consists of multiple sub-datasets, so to avoid integration issues we subset to the largest sub-dataset which contains 4 million cells. For each donor, raw count matrices were normalized to 10 counts per nucleus, log-transformed, and restricted to the top 000 highly variable genes. We treat each donor's collection of nuclei an empirical distribution over transcriptional states. Donor-level diagnostic metadata were obtained from the accompanying PsychAD clinical annotations, and we restrict prediction targets to the six major disease categories which have at least one positive and negative example in the dataset (Alzheimer's ss Parkinson's diffuse Lewy body, bipolar, schizophrenia. and vascular dementia).
21

---

B.3.2 Model architecture and training
Both the supervised and semi-supervised models use the same ResNet deep sets encoder to aggregate single-cell features into donor-level embeddings. For the supervised variant, the encoder is trained end-to-end with classification head. the supervised GDE, same encoder coupled a conditional variational autoencoder (CVAB) generator, trained jointly to reconstruct cell distributions while predicting disease labels for the 10% of donors with laheled diagnoses. In both cases, we use 64-dimmnsiionll latent space and two hidden layers of size 128. After training, logistic regression classifier is fit on the donor embeddings prediet across the six categories.
multi-label disease
B.3.3 Evaluation
We report donor-level predictive performance using accuracy, balanced accuracy, ROC-AUC, and score on 10% of completely heldout data. Semi-supe superviie GDEs outperform purely supervised models across all metrics (Ta demonstrating that unlabeled donor distributions improve representation quality and generalization in large heterogeneous cohorts.
B.4 Lineage-traced scRNA-see -eq experiments
B.4.1 Data preprocessing details
We use lineage tracing data from Weinreb al. 43 The single-cell RNA sequencing (scRNA-seq)) count matrices were preprocessed following standard procedures. Specifically, counts each cell were normalized by rescaling to 10* counts per cell, followed by log transformation. Finally, the top 10* highly variable genes (IVGss) were selected. Cell-type annotations and two-dimensional SPRING embeddings were obtained directly from the annotations provided in Weinreb al.
B.4.2 Mutual information estimation
We compute mutual information sample mean of pointwise mutual information estimates. To estimate pointwise mutual information the representation space, we use the nonparametric nearest-neighhor estimator introduced by Kraskov et al. [65] with h = 3. This estimator has been shown to be effective in this setting:
model latent spaces with tens of dlimesions 44
B.4.3 GDE modelling architecture
We use Resnet- G architecture the encoder and a CVAE as the generator. We use 64 latent dimensions, with hidden layers of size 128.
B.5
Perturbation Prediction
B.5.1 Data preprocessing details
We use the pre-processed
file from 49 including 10 genes. We compute the 10% most perturbative perturbations by examining the differentially expressed genes and then randomly select 20 of those perturbations hold out. We hold these across all cell types.
B.5.2 GDE modelling architecture
We Resnet-GNN architecture the encoder and CVAE as the generator, similar to the architecture in the lineage-tracti experiment, except we use larger hidden state (1024) and larger latent space (256). include a perturbation prediction loss during training which trains a linear model with pairwise interactions between the control cell distribution embedding and the gene embedding to predict the difference in mean expression through linear head. This structures the latent space for our downstream perturbation prediction task.
B.5.3 Perturbation Prediction
We fit a ridge regression predict (1) the difference in expression and the difference between the perturbed embedding and control for each perturbation using GenePT gene embeddings (50) with cross- validation to perform grid search over A. We then compute the predictions on the held-out perturbations and use a linear head to predict the mean expression from the latent difference. Finally we compute the R² score and the MSE.
22

---

B.6 Optical pooled screening dataset
B.6.1 Data preprocessing details
We use phenotyping images with assigned perturbation barcodes from Funk et al. [ЛT] We analyze only of measured channels: DAPI and GFP. Each image 64x64 bounding surrounding single cell (center-padded or center-cropped from the original bounding box as necessary). Image intensities are normalized to minimum of -1 and maximum of 1. Using the set of perturbative perturbations computed in [51] we randomly select 30% to holdout during training for evaluation.
B.6.2 GDE modelling architecture
the encoder architecture, we extend our GNN approach to convolutional layers, standard image processing. For the generator we use U-net architecture standard diffusion for images, but upscaled in expressivity relative to our MNIST and Fashion-MN MNNIST examples.
B.6.3 Perturbation Prediction
We find that empirically, our diffusion approach struggles model the padded border of the cells. So. inference time we condition on the border to generate our predictions. Using GenePT, we train a ridge regression with grid search (similar to App. B.5) to predict the perturbation distribution embeddings. We also construct a nearest neighbor model using the GenePT embeddings to sample the padding. We then condition on the padding and the predicted latent to sample set of 1,000 cells from each heldout perturbation. We then compute the DAPI intensity and compare with the ground truth, computing the R² and the MSE
B.7
Methylation atlas of human tissues
B.7.1 Simulating raw bisulite--sequeccing reads from methylation patterns
While sample-specific methylation patterns are published [] the raw sequencing reads are not public due to patient privacy considerations. Here, we instead use the published methylation patterns (in the form of pat files) to simulate bisulfite sequencing reads. For each methylation site entry of the .pat file, we use wgbstools (66 to find the 100 preceding bases of the HG38 genome reference, and append to the CpG sequence. We omit all CpG sites with unknown methylation status. We subsample 10" sequencing reads per sample.
B.7.2 GDE modelling architecture
We use a ID convolutional neural network as our encoder, with mean pooling at each layer (analogous to the fully connected GNN with an MLP, but using convolutional layers). For the generator, we use HyenaDNA We additionally include linear classification head on top of the distribution embedding, co-trained with cross-entropy loss.
35
B.8
GPRA
B.8.1
Data processing details
collect sequences in the Gal and Gly conditions from (52) and process them into 100 quantiles by measured expression. totaling 34 million sequences. We one-hot encode these sequences for ACTGN, and tokenize them using the HyenaDNA tokenizer. We break these sequences into 100 quantiles and hold out the top 5 quantiles during training. During training, construct sets selecting "center" quantile and then randomly sampling from that quantile and the adjacent quantiles.
B.8.2 GDE modelling architecture
We use the same architecture as in the methy experiment (App.
B.7
B.8.3 Details for Fig. S
We encode random subsample of 130K sequences from each quantile in the Gal condition to construct the set embeddings (the larger dots). We then compute the PCA of these embeddings. We embed all the DNA sequences sets of size one and project them to the PCA. For the histograms of TFBS motifs we leverage the PWMs from [53]. We wrote a simple unidirectional motif scanning procedure in Torch to facilitate efficient scanning. and used a threshold of to determine hits. We then sum over the motifs to derive the motif' count per sequence, and then compute the histogram by plotting the distribution of these counts quantile.
23

---

B.9
Spatiotemporal distribution of viral lineages
B.9.1
Data preprocessing details
We obtain all SARS-CoV2 spike sequences deposited up to April 2025 in GISAID (58). We group sequences by submission month and lab of collection. We discard sequences with improperly formatted date fields. During tokenization, we truneate sequences to 1000 amino acids.
B.9.2 GDE modelling architecture
The encoder couples the ESM-50M [56] architecture coupled mean-pooled GNN, while the generator uses the Progen2-150M architecture R4 with prefix conditioning. We initialize (but not freeze) protein language models with their pretrained weights. We use 128 dimensional latent space.
24

---

C Background
C.1 Frequentist, Bayesian, and Predictive Sufficiency
Sufficiency classical notion statisties that formalizes when statistic retains all information about parameter or distribution. In this appendix, we distinguish three forms sufficiency relevant to modcrn generative modeling and provide canonical examples.
C.1.1 Frequentist Sufficiency
Let (Po : E } be a parametric family of probability distributions on a sample space X. A statistic T(XI..... ) is frequentist sufficient for if the conditional distribution of the data given T does not depend on 0:
Pa X ... ..)) = (independent of 8).
Intuitively, the likelihood depends on the data only through
C.1.2 Bayesian Sufficiency
Given prior x (0) over the parameter space, statistic T' is Bayesian sufficient for 0 if the posterior depends on the data only through T:
T( , X..) = x(0 T(X)...... X,)).
Bayesian sufficiency holds if and only if T a sufficient statistic in the sense that the posterior conditionally independent of the data given? T.
C.1.3 Predictive Sufficiency
A weaker notion, often relevant in nonparametric and distributional settings, is predictive sufficiency. Assume a joint model (0, X Xnew) (e.g. and X; i...d. Po). A statistic predictive sufficient if the distribution new sample Nnew given the same as given the data:
P(X E B T(X) x.)) = P(Xnew E XI...., X,), VB E B(x)..
This requires only that T contains enough information to match the predictive distribution of future data.
C.1.4 Implications and Comparisons
There is strict hicrarchy among these definitions:
Frequentist sufficiency Bayesian sufficiency + Predictive sufficiency.
first implication follows from the factorization the likelihood, and the second follows because the postcrior predictive a marginal of the posterior. However, the reverse implications do not hold in general, especially in or nonparametric models. In particular, predictive sufficiency may hold in settings where no fiite-dimensional parameter exists.
C.1.5 Examples
Example (Gaussian Mean). Let X1, XI,..., X. ((,, with known o² Then the sample X., sufficient for μ in three senses: frequentist, Bayesian, and predictive. The likelihood, posterior, and predictive distributions all depend on the data only through X..
Example 2 (Uniform(O, 0)). Let X1, .1.......... Xn, Unif(0, 0). Then the sample maximum
T. max .....,X...,....
the minimal sufficient statistic for 0 in both the frequentist and Bayesian senses. also suffices for prediction of future samples, since the predictive distribution under 0 is uniform on [0, 0). and T,, provides all information about 0.
C.1.6 Nonparametric Extensions
In the nonparametric regime where is not indexed by mimte-dimensionall parameter, predictive sufficiency remains well-defined, For instance, under de Finetti (exchangcable) model with latent random measure P II and X, P k.d. P. the empirical measure P., Bayesian and hence predictive sufficient for P. In this setting, stronger finite smite-dimensional forms of sufficiency may not exist, but predictive sufficiency still supports meaningful generative modeling.
>i Óx, (equivalently, the multiset of observations)
25

---

C.2 Otto's Geometry and Statistical Submanifolds
This appendix recalls Otto's Riemannian calculus on the 2-Wasserstein space P2(X² and explains how fmite-dimmnniooal parametric family of measures inherits induced geometry [3 Throughout, statements intended the standard "Otto calculus" sense; rigorous treatments interpret P2 as geodesic metric space and identify tangent objects for absolutely continuous measures.
C.2.1 Wasserstein Space and the Benamou-Brenier Formulation
Let X' C Rd be convex (e.g. == 1e"), and let (2)) be the Borel probability measures on X' with finite second moment. The 2- 2-Wasserstein distance is
inf J xxx rE)(Ho.P1)
( (Mo, pi)
|| ull dy(2,
Benamou-Brenier gives the dynamic formulation
L. Jo [| dp((z) de,
(2 (po, p1)=
inf (Me))0)) Depe+ V- (=(e)))==
where the continuity equation holds in the distributional sense, and E L (px).
C.2.2 Otto's Riemannian Structure
For μ absolutely continuous with density tangent vector can be represented as
(pVø),
for a potential & (defined up to an additive constant). Equivalently, one represents the tangent direction by its minimal kinetic energy velocity field e = VΦ. The Otto (Wasserstein) inner product is
gr (رs p2) = Vel (z) (e)(s) du(x)) J vi (1) (2) du(x).
with Vd, the minimal- nerimall--ommmaee norm representatives. With this metric, constant-speed W2 Wa-geodesics precisely curves of minimal kinetic energy.
If po is ahsolutely continuous, the (Brenier) optimal map from po μr induces the displacement interpolation f4 ((1 f)id + (T)) ) 4o
which is constant-speed W2-geodesic (on convex 2))
C.2.3 Statistical Submanifolds and Induced Wasserstein Geometry
Let Q be distribution over (x). speak a submanifold, we assume is supported on finite-dimensional smooth embedded family
M {ue 0 E ReCk C P2(X),
where 0 + As is smooth and are absolutely continuous.
The induced (pullback) Wasserstein metric on parameters is defined by
Gey = 9ro (8:14: 031o) = V Vás dua,
where di solves the elliptic equation
v (ne Vebe) 0:ps (in distributional sense).
The intrinsic Riemannian distance on M can then written as
dyc(pso (o,) inf S' V0z G(B.)å, P:
and satisfics d px)) 2 W2(Ha, ua) in general (strict unless M is geodesically closed in P2).
C.2.4 Examples and Application to GDEs
Examples include Gaussian families (closed under Wa-geodesics) and general smooth parametric families. For mixture models with finitely many components, one can study the induced Wasserstein metric on parameters, although ambient Wz-geodesies between mixtures typically leave the class.
In this work, interpret GDEs as leamming smooth embeddings of such constrained family of data- tata-generating distributions into Euclidean latent space; empirically the learned latent geometry may approximate the intrinsic geometry induced by the Wasserstein metric.
26

---

D Theory
Throughout, lct (X, d) be a Polish metric space and let B denote its Borel a-algebra. Let P P((2)) denote probability law on X. Given m N, let Sm (X1...., Xm) be i..d sample from P, and let Pm E ōx, denote the empirical measure.
Let Mo(X)) denote the vector space of finite signed Borel measures on X with total mass 0, equipped with the bounded Lipschitz norm || ||B defined below.
We use Pt. P2 to denote two (possibly distinct) probability laws on X, and Si, Sa for independent samples from
Pi, P2 respectively.
For signed measures v, on (X, B) define
(vL(V,
sup 1-111- Lip(
| f f( -1)
We use || [[BL for the corresponding norm || == dBL (v, 0). Since the class of bounded 1-Lipschitz
functions is contained in the class of all bounded measurable functions, we have (v, S || pl/TV.
D.1 Necessity of Distributional Invariance
Our goal is to design encoder architectures that flexibly model unknown data distributions while guaranteeing consistent generation of the underlying law as sample size grows. Since the true distribution P not known in advance, the encoder must be constructed to generalize across all possible P, without leaking spurious information tied the specific realization sample size. If the encoder depends on sample-level such as ordering, multiplicity, or the raw sample size may encode features a generator can exploit, breaking the guarantee that
Motivation
((m)) 4 as m > 00, Sm ppm
This risk arises even under either permutation proportional invariance their own: both permit dependencies that vanish only in expectation and insufficient ensure correct extrapolation with increasing m. For example, encoders based on unnormalized sum aggregations (e.g., DeepSets) will vary with m even when the empirical distribution is unchanged, leading to divergence at inference time.
To formalize this constraint, we appeal to two classical principles from statistical decision theory and invariance: (i) under i..dd sampling. the empirical measure is sufficient statistic for the (nonparametrie) model indexed by the unknown law P. so conditioning invariant under permutations), meaning discards exactly the ancillary degrees of freedom (ordering and other sample-level artifacts) that do not carry information about P. In our setting this motivates enforcing that the encoder depends on the sample only through the empirical distribution: any additional channels (c.g. ordering set-size effects) are unnecessary for identifying P and can be spuriously exploited by a flexible generator, especially when extrapolating across sizes.
it loses information about P; (ii) is also minimal (and the maximal
We define distributional invariance:
Definition (Distributional invariance). A family of encoder maps (Em (mm)22 with &m : X tributionally invariant if there exists a measurable map Φ P(() - Z such that for every m and every Sm. = (,,....,mm))) E x with empirical measure Pm. Em 8.g,
dis-
En = (PP.).
particular, any such family is permutation invariant and consistent across set sizes under proportional duplication, i.e. for every integer 2 1.
(S() Exm (emm(Sm. (Sm...., Sm).
K capics
Empirical measure as a lossless and minimal summary. Fix m E N. Consider the (nonparametric) i.i.d. model pp8m P E P(X)} on X We use the standard (Ncyman- Fisher) notion of sufficiency:
Definition (Sufficiency). A statistic Tm x T is sufficient for the family { pp@m P P(X)) if the conditional distribution of X1m pmm given Te (Xim) does depend on Equivalently, for every bounded measurable f : n + R, there exists a measurable gy TT + such that for all P.
E [f(xm_((X..)) Tm (X1mm)] = (n(TT((1:m) a.s.
Let 6.,. be the symmetric group acting on x by (II,..., Rm) + (EI(I))..., Fr(m)). Call I'm permutation invariant if T' (x) = Tm (xx) for all E 6m.
27

---

Theorem 1 (Empirical measure is sufficient, minimal, and a maximal invariant). Let X1:m = (Xi,..., Xm) pam and Pm. 1x_1xx, Then:
(i) (Sufficiency/ / /'losing nothing"). Pm is sufficient in the sense of Definition Moreover, for any bounded measurable f xxx" + IR, a version of the conditional expectation given by symmetrisation:
Epen (pamm[(Xe (X(mm) Pm
Σ m
a.s.,
m!
which does not depend on P. In particular, the sample ordering is ancillary given P...
(ii) (Maximal invariant / all permutaiooninnvarane summaries factor through Pm). Let (m(()) == {m -1 > Tm X =1 5z; 3iim E 3 be the set of empirical measures with m atoms (counting mulriplicity). If + T is permutaaiio invariant, then there exists measurable Om (m)(X) + such that Tm @m o Pm pointwise on xm (One may extend om all of P ((X) arbitrarily if desired.
(iit) (Minimal sufficient / "not keeping anything unnecessary",, If m is sufficient for {p : P P(x)),
then Pm. is measurable with respect to a (Tm)): there exists measurable hm T + ((X)) such that
Pm (T (X1m) pem -a.. for every P E P(X).
Equivalently, (P)) is the minimal sufficient o-field.
Proof sketch. (i) Let f be bounded mcasurable and define the symmetrisation f(Z1sm) i Enee I [I ((())))))) Then f is permutation invariant, hence depends F1:m only through its multiset, i.c. only through Pm. For any bounded measurable H that is ( mmcas-uaasural H(X H (Xx()))........ Xx(m)) for all π. so by exchangeability.
(l((Xim)H((X))) (mm) =E ml
(( [x(xx))...xxxx.x.))) H(X
Thus f(Ximm is a version of (()( Pm]], and it does not depend on P, establishing sufficiency.
(ii) If x, E d satisfy Pm (x) Pm (v), then y is permutation of x, permutation invariance gives (x) (T) Hence Tm is constant the fibres of the measurable map Pm. : X + Pm (X),, i.e. Tm (m..)-maasuaable. By the Doob-Dynkin lemma, there exists a measurable Φm. Pm)()) + such that Tm @m o Pm.
(iii) If (x) = (y) but (x)(x) * (y). choose an atomic law P that puts positive mass on every point appearing in masses that finite support, the ratio pêm ([=))/ (aal))/sSSmmm)) ({v)) can be changed while keeping T'm Tm(x) (1) = Tm (y). contradicting sufficiency. Therefore T,, (x) Tm (y) = Pm (x) = (g), so Pm is measurable with respect to a (Tm).
Then both sequences have positive probability under pcm By varying the atomic
forcing the conditional distribution of X 1mm given T. depend
Corollary 1 (Necessity of empirical measure dependence for artifact-free encoders). Let Zm = (X be an encoder output.
(i) If Em is permutation invariant, then by Theorem I [ii there exists om such that Z. Dm (Pm):: the encoder can only depend on the data through the empirical distribution.
(ii) If. additionally, the family (Em. mm2) is distributionally invariant in the sense of Definition then there exists single measurable map o: P(X) - Z such that Z.. = ((Pm. for all m.
In this sense, restricting empirical empasreal-measure dependence discards only ancillary sample-level level degrees of freedom (order and size artifacts) and keeps exactly the information relevant to the underlying law P.
measure
Remark (Conncction to predictive sufficiency and scaling). Theorem formalises "losing nothing" (sufficiency) and "not keeping anything unnecessary" (minimality) for the nonparametric i.i.d. model. Definition 7 in Appendix D.3 is an asymptotic, reconstruction- bused analogue restricted to the manifold M: it asks that from low-dimensional coordinate ((P.)) one can reconstruct Pm, (in || ([L).) as - co. Operationally, enforcing duplication invariance removes a particularly dangerous ancillary channel: set size. Without (e.g. sum pooling), a flexible generator can fit finite-m size effects during training and behave impredictably when m changes, even if the empirical distribution is unchanged.
D.2 Complete Large-m Analysis of the Plug-in Loss
Motivation We analyze the statistical properties of plug-in in loss used to train distributional encoders and generators. Our goal is to understand the asymptotic behavior of this loss as the sample size grows, and to establish conditions under which the learned generator recovers the true data distribution. This analysis provides principled foundation for the training objectives used in framework.
28

---

Setting First we establish some notation and definitions.
Definition (Hadamard differentiability). A map T D -
between normed spaces is Hadamard differen- tiable at x E 'D if there exists a continuous linear operator DT, such that for every sequence ht -> in D and t O, T(=+the)-T(e)
[T.[
Definition 4 (Fréchet differentiability).. : y be map between normed vector spaces. We say that T is Fréchet differentiable E D if there exists bounded linear operator A - such that
[T
T(x) A(h)lly O.
lim I
||b||p
The operator A is called the Fréchet derivative of T at Definition (Tangent set at Qu.). Let Qu E P(X) and write Mo(A)) ) for finite signed Borel measures on A' with total mass O. Define the (aa))ttttt space
(Q)) = {N Ma(x)) h < Qo.
dh dQu
E LQ
||l|ie :- dQu dh dQu 1aa
We view (Qo) as a normed linear space via the identification h ++ dh/dQu.
We work in the following general setting:
Assumption (Data Empirical Measure), (X) is a Polish space; P P(X) the true data law. Observations Sm (X,)..., are ii.dd P. The empirical measure is P' Ei=1 ōx,-
Assumption 2 (Encoder regularity). For each probability law P E ((X)) the encoder Φ P(X) + Rd satisfies
(i) Distributional invariance: (5m)mm) ((Pm) depends the sample only via its empirical measure.
(ii) Pathwise (Hadamard) differentiability: é is pathowis differentiable at P and its canonical gradien bpp X T belongs T. (P).
(iii) Asymptotic linearity (AL): there exists a remainder Tin such that
(m(((Pm) - -P(P] mm Σ
wr(X)) + Tms
where (xx=)(p(()] = 0 and rm - 0 in L L
In particular,
Vm (p(Pro)) ((P)} 3 M(O, Es ==:=Vare var=xx-l/vv>x]]]
sup., (o(P.)) -(())) <c.
Assumption (Generator). Let M(X' denote the vector space of finite signed Borel measures X equipped with || BL. and identify P(X) c. -M(X). Assume the generator G Rd. + ((X) admits a local Fréchet expansion at μ p== (P) when viewed as map into M(X): there exists a bounded linear map D₂G TRd Mo(X)) such that
[[ - 0(1) D, = o(llbllz4)
as ||l|| + 0.
Moreover, writing Qo = G(p), the derivative is L
D,T(R*) S Do(Qo) and sup |
D 00.
Finally, the remainder is negligible in the (Q²) tongent norm:
[G(u+t) G(μ) D. =((laa)>))
||A|| 0.
Assumption 4 (Divergence (Hadamard differentiabilityy on an La tangent space)), Let Qo = ((u) and (o(Qo) be as defined above. The discrepancy L: P(x)² + R2 satisfies:
(i) (Hadamard differentiability on Tha (Qu)) for each fixed P, the map Q + L(P, Q) is Hadamard
differentiable at Qo tangentially Do(Qu) (equipped with ||- Ilb.). with continuous linear derivative
D2C(P, Qo) Do(Qo) - RR
In the semiparametric sense of i.c. the unique influence function representing the functional derivative along Mo(X').
29

---

(ii) (Sepurating property) C(P. Q) =0 = P=
(ti) (Weak-continuit)) for each fixed P. Q.) then Qn = P.
We work with the discrepancy L from Assumption The plug-in loss is
mm = (P, (()(Pm)))
and the population loss is
- (PP ((o(P)))
where Pm. is the empirical measure of the sample, Φ is the encoder, and G is the generator. When P is unknown, C(P,-) :) is evaluated via an empirical Monte Carlo estimate based on S the results below describe the additional emror incurred by using (PP..) in place of ((P).
Lemma 1 (Functional Delta Method, [67] Thm. 3.9.4). Let (D, lIb) and (B, 11.) be normed vector spaces. T + be a map that Hadamard differentiable a point E I tangentially to subset Do C B. with continuous linear derivative denoted DT,, By
Suppose:
(a) There exist random elements Zm taking values in DB such that:
Vml( Zm - 2) d Z
for some tight limit Z taking values in Do.
(b) Z is tight and Borel measurable.
Then:
vmm (T()) T(=)) - (Z).
where (Z) is random element of E. In particular, if Z is Gaussian in Do and DT, is continuous and linear, then (Tz Gaussian in E.
Main Result
Theorem 2 (Large-m behaviour of plug-in loss), Assume and
Let == (PP) and C(P, G()(Pm))) Then:
(a) Asymptotic normality of the Encoder.
Vmi( o(P)} S ((0.2.)). Ei (axxrr(vr(X))
(b) Consistency (and mean consistency) of the loss. Writing E(O) = ((P.)( we have
t.. = ((((p)))
E(H)) E".
If. addition, E locally Lipschitz neighbourhood of μ (as a map Bdd B.) and the sequence ((.((P.))) is uniformly integrable, then
-- 2- and hence BLEm)
A sufficient condition for uniform integrability is the following growth bound: there exist >> and C < op such that [E(A)] < (1 + [ for all 9 and sup... El((((P.))) II"" <
Moreover, f is swice continuously differentiable in neighbourhood of p. with bounded Hessian, then second-order Taylor expansion yields
E(fm]] - e" (e,((E(((P.))) - -+]] O( o(m 112).
In common unbiased cases where E(o(P..) - O(m-1) (e.g. sample means and many smooth M -esemmmatros) this simplifies to E(L...] - Er = =0(m²)
(c) Asymptotic normality of the loss. Let E(0) L(P,g(O) and denote its derivative at p by the continuous linear functional
Dlg Rd IR,
Dl,[h] D2C(P. D.9h]), Qo := C(H).
Then
Vm (Cm -) 1 M0, d² De,, Es DÊ
(Identifying De, with a gradient vector V E Rd under the Euclidean inner product yields o² (V,l) (V.l)))
30

---

(d) Consisteney under correct specification. Fix P. 1f (6" "0","")'))))')))"" minimizes the population objective
(0. - [(P, (()(P))
over the model class, and if the model is well-specified in the sense that the minimum value is 0 (equivalently, + DC.
g( (6" (P)) 0), then G* (Pm)) = P in pam probbabilyy
Proof. Step Asymptotic Normality the encoder (a), Assumption 2 liit (asymptotic linearity) gives
(l((PP)-((PP)) 1
(x()) o1)),,
vm
and the classical multivariate CL.T yields the stated convergence.
Am == ((P.) - aa so that, by (a), v A. - NO,,E
Step 2 (consistency and mean consistency). Write {(0) == ((,))()) and A.. := <(Pm) By Assumption 2 jii), A.. = (m 1/²), hence ((Pm.. + A. + μ in probability. Since & is continuous at
(automatic if é is differentiable at po), the continuous mapping thcorem gives E. = El + Am) - E(1) = l" in probability.
For convergence of expectations, assume é is locally Lipschitz u and (e( Pm. m22 is uniformly integrable. Since (Pm.) + yu in probability and continuous ja, we have (( Pm) {(p) in probability,
ie. L. - L" in probability. Uniform integrability then implies B[(..) E" and BIe.. El 0.
If prefers direct bound using local Lipschitz, let U be neighbourhood of μ which (é(0) {(x) L ull. Then
ElC. 01 S B[|e(H + Am) ((x) 1 + Am E UH Flje(( + Am) {()) 1/ Am & U}}-
The first term is < Am and tends to since sup,, Elly m | < co. The second term vanishes by uniform integrability together with P(s + A... U) + O.
If é is C² μ with bounded Hessian, second-order Taylor expansion gives
E[Cm] - E" = D4 [[[[E + O I|²) = De, [[[[]]]]]]] + O(m 0mm=)) = o(m -1/2),
since B[Am] E /// vm o(m - 1/2 under Assumption iii). (If additionally E[Am] = O(m (1mm"),, then the bias is Om )
Step 3: Asymptotic Normality of the loss (c). Let Am == ((P,)) ) a so that v vm + with Z
(0, Es) by part (a).
By Assumption
3 (including the (Q)) remainder control),
V Vm Tm(glo G(u Am) 2o vm D, (1) in (Do(Qo), -11.).
Since Dag IRd + (Qu)) bounded linear, we also have vm D. + D, Do(Qu).
Now apply the functional delta method (Lemma 1 to the map Q + C(P, Q) at Qo. tangentially to Bo(Qo):
vm C(P + Am)) - C(P, = D2C(P, D=C(P,Q0)
Define the continuous linear functional DE,- Bd +R by
De, [h] D2C(P, Qo) [
Then the limit is [2) which is Gaussian with variance 0² = De, Es DeT
Step 4: Consistency under correct specification (d). If (Ø",g") minimises P ((P))()())) and the model is well specified, then ((P, G* (Ø" (P))) 0, (0" (P)) = P Assumption 4 (i). Repeating the expansion from (c) with (ø*. g") shows that
g(P.g (o'.)) Or(m 1/2),
hence C(P. ("P.g"(0' (('.)) + probability. Finally, Assumption 4 (m)) implics that the topology in- duced by E(P,-) is at least as strong as the weak topology: for every 0 there exists > 0 such that C(P, Q) < d entails dar Q) < (Otherwise one could construct sequence (Q..) with L(P, Q.) 0 but dbL Q+) 2 y for all contradicting the assumption.)
31

---

Since ( 0 in probability, for any fixed v > we may choose ó >> O as above and obtain P dBL (g (Pm)), > 7) s ((((P 2 20) m-
Thus (' + P in probability as claimed.
examples, counter-examples, and CLTs The only encoder requirement entering Theorem is [ We now show that it is satisfied by large family of permutation-invarre architectures built from aaymptoticalll-lineer (M/ poolers.
Encoders:
Generic K -llayer pool-concat encoder Fix K E N. Given set of samples S.. = {Ii,..., Lm define recursively
4 p(r:). 0 ("(mm) 1,0 MLP
(1 ( ..").
== L...,,K
and set the encoder output to be another pooler ((Pm) 1 (n(m)
We call permutation permutation-ivvariint invariant functional an asymptotically linear (AL) pooler if it root-m consistent and admits an ceftuence-function expansion; precise details follow.
Definition 6 (Asymplotically-lines pooler). Let P(x) R be fixed statistical functional. A family of symmetric maps (fm)m21 with Tm x +R is an AL pooler for yp at law P each Tm. depends on the sample only through its empirical measure Pm and there exists uhp L (P) such that,
oo,
mm (T (Xi.m)) -(P)
(X()) + o,(1).
Examples include the sample mean, median, trimmed mean, Huber M-estimators, M-quantiles, and studentised 7-estimators with finite variance.
Proposition 4 (CLT for K K-layer AL pool- encoders). Assume
(1 each ( (e=1,..., K 1) is distributionally invariant AL pooler (in the sense of Definition P:
(i)) each ;LP; and the base feature map -RR are C² with bounded derivatives, and weights are frozen as - co.
Then the encoder & is distributionally invoriant, pathwise differentiable, and satisfies the CLT of Assumption with
- [(P)) & N(O,
for some finite covariance matrix Egg.
Sketch. The composition of Lipschitz maps (MLPp) with AL pooters is Hadamard diffcrentiable by repeated application of the delta method (iterating Lemma D (67). Plugging each AL expansion into the chain yields an overall AL expansion whose leading empirical-process term is m 1/2 Em (X.) for some (P) function 4b.. giving the CLT.
Instantiation to common architectures
Corollary (DeepSets, Transformers without positional enc.). Encoder architectures of either type below satisfy Assumption and Proposition
(a) DeepSets / fully-comnected GNN with global mean: 7 and (K++1) are sample means;
(b) Self- celf-aatention block with mean head: (()) are sample means; MLP: includes the softmax-attention update.
Why max-pooling fails The max functional Tmax (m1..) = max; x; is not Hadamard differentiable at continuous laws: its influence function is identically 0 whenever the maximum is attained at a unique point and undefined when it is not. As consequence, the usual /m-sscalig does not yield Gaussian limit for the centered statistic /m (Tmax(Xx..) ) maax (P)); instead, after a different (typically lincar-in-m) rescaling one obtains a non non-Gausse Gaussian extreme-value limit law. Thus Assumption iiii) fails. Using max maxx-poolig inside deep encoder therefore breaks the loss-CLT of Theorem (Softmax pooling with fixed temperature 7 > 0, by contrast, is smooth and can be made into a valid AL pooler.)
The table below summarises the status of common poolers.
32

---

Pooler
AL CLT?
Influence fen. thp in L 2
Sample mean
Huber M M-estimattrr (6 fixed)
Sample median
Top-k max
Softmax (r > fixed)
Smooth Approximation of Non-Regular Statistics The theory developed here establishes that Hadamard difffereniiaiility of the encoder ensures asymptotic normality and consistency, and in Appendix C develop the idea that encoders learn sufficient statistics. But what if the sufficient statistic of interest not Hadamard differentiable? The sample maximum is classic example: it the minimal sufficient statistic for the endpoint of uniform distribution (see Example 2 yet it is not asymptotically normal.
Let ,I...., Unitorm(0, 0). The sample maximum
X(m) == max X1 Xm}
satisfies
mf (m)) 4 Exp(1/0),
so it converges to but asymptotic distribution exponential, not Gaussian. This occurs because the maximum is not smooth functional of the empirical distribution: it fails Hadamard differentiability, so the functional delta method does not apply.
A natural remedy is to approximate the max by smooth, duplication-invaaiant functional. A standard choice is the normalised log-sum-exp:
LSEA (X ...) log
AXx
For fixed this is smooth functional of the empirical measure (under mild moment conditions ensuring the log- moment finite) and is therefore amenable the delta-method theory above. As À + oc, LSEx + max; X so we recover the max in the limit.
Corollary 3 (Smooth approximation suffices for asymptotic normality). Let m(Pm) be non-smooth statis- tic (e.g., maximum), and let TT()) (Pm) be a family of smooth approximations (e-g., LSEx) such that TTA) T(Pm) pointwise À cc. Suppose that for each fixed À the map - (P) is Hadamard differentiable and satisfies Assumption 2 Then for any fixed X, (PP)) admits asymptotically normal plug-in Am
estimators. Allowing
introduces tradeoff between approximation error and /mn-asymptotics.
Thus, even when the true sufficient statistic is not regular, Hadamard differentiahle encoder can still be learned approximate it. This ensures that the asymptotic guarantees from Theorem continue hold, This also highlights why we cannot use max-pon poxolpooiig in the encoder, since breaks the VT CLT.
Generators All neural generators considered in the experiments MLPss and Transformer decoders directly. diffusion/score models when implemented with fixed-step sampler- can be viewed mnite-dimensional compositions of smooth maps from latent codes to synthetic samples, inducing a (locally) smooth dependence of the resulting law on the embedding; this is the modelling assumption captured by Assumption 3
D.3 Embeddings and Predictive Sufticiency
Setting. Let C P(X) be the statistical manifold introduced Section
Here we assume the statistical manifold M is d-dimensional (in the usual lifferential-geomeetn sense), so dim TpM d for every P E M.
For E M observe S.. = (X....., X.) Lid. P and write the empirical measure P.. = m Σm= 5x,-
Throughout use the plug-in in predictor P... Given statistic Tm 6(Pm). where & P(X) + Rd is defined and C² on a neighbourhond of M, let U C Rd be an open set with 6(M) C U such that P()(P..) E U) + for every P E M. Define reconstruction map R U > M that is C on U (in the manifold sense) and set
p = ITT... = ((((..)).
Definition 7 (Predictive sufficiency). The statistic T'.. ((Pm.) is asymptotically predictive sufficient if there exist an open set U Rd with ((M) c U, and reconstruction map R U + -M that C¹ on (in the manifold sense), such that for every P E M,
IF ppm (6(Pm) E UU)- + and Pm R(6(Pa)))
pm
00-300
33

---

We write P == [(((P..))
This notion of sufficiency is reconstruction- based asymptotic analogue of the "no order/size artifacts" principle Section[D) D.1 it asks that from the low-dimensional coordinate ((Pm) ) one can reconstruct the plug-in predictor Pm in the bounded Lipschitz metric (and hence recover all weakly continuous predictive functionals).
Theorem 3 (Embedding < Predictive sufficiency). Assume that á P(X) R is C' neighbourhood of M and satisfies encoder regularity conditions of Assumption 2
Then the following are equivalent.
(i) Smooth embeedding: the restriction ølm M - I2 is injective and its differential døp : TeM Rd is bijective for every P .M.
(ii) Predictive sufficiency: Tmm = ((Pm) ) is asymptotically plug-in in sufficient in the sense of Definition
Proof (sketch). Throughout, || |BLL denotes the bounded-Lipschitz norm signed measures.
ii) = (ii). σi MM a smooth embedding, its image ((M) C TR an emhedded submanifold. By the unverse-function theorem and standard ubbula--reighourrhood constructions, for each / E .M there exists a neighbourhood Vp of (P)) in R and continuous map Rp Vp - M such that Rp(e(Q)) Q for all M with d(Q)) Vp. Using partition of unity may glue these local inverses into single continuous retraction R .M defined on open neighbourhood V of ((M) and satisfying R()(Q)) Q for all Q M.
Encoder regularity (Assumption gives
vm (PPm)) p(P)}
Σ ( + oP(1) in R
Vm
((Pm) (PP) in probahility. Since Pm P in dBL almost surely, have @b Pre) E V with probability tending to one and
P R()(Pm)) pm
R(O(P))
in the bounded Lipschitz topology. by continuity of R. Combining this with Pm + P in |lBL and applying the triangle inequality yields
pm
Pm / I ||mm PL + || PL m-000
0,
which is precisely predictive sufficiency in the sense of Definition 7
(ii) (i) Conversely, assume predictive sufficiency with R e (U. Fix P M. Since Pm - P in [|BL and || Pm R( (lB)))) in probability, we have (()((m)) in probability. Encoder
regularity implies (r.)) + &( P) probability, hence by continuity of R. R(o(P)) = P. Therefore Ro M = idm-
Differentiating the identity map on M and using the chain rule yields
dRR døp = idzp.-
Thus døp is injective for every M, and since dim Tp. M = d dim Rdd it is bijective. Moreover, R) Rla(M) is continuous inverse of dlM. so @lM is smooth embedding.
Remark (Identifiability is automatic). Because each P M already defines a unique predictive distribution. any statistic that plug-in sufficient must be injective; no separate identifiability condition required.
34

---

E
Extensions
E.1 Extension to Multiscale Settings
In many applications, data is naturally organized distributions of samples at a fine scale (e.g.. single cells), grouped into entities at coarser scale (e.g.. patients), which themselves may belong to larger groups (e.g.. hospitals). More generally, we may observe hierarchical data in which each level exhibits internal distributional structure.
multiple scales. For example, may observe
Our framework naturally extends to such multiscalc settings, At each scale we observe set of units indexed by i = 1,...,my (*), Each unit at scale is associated with: a set of samples Stmm = { fj=1, drawn i.i.d. from distribution p(*)) and higher-scale sample x(++1) E x(++1), representing the corresponding entity scale +
(4) mm
The lower-scale distributions p(() are drawn i.dd. from meta-distribution Q over P(x(()), while the higher-scale samples si (s+1) are drawn from p(++1), where p(a+1) Q(++1).
Each lower-scale S defines empirical measure
p(a)
m Σ (a E pm)())))
At each scale learn: an encoder g(+) :m (2 Rd mapping lower-scale empirical distributions into latent space, encoder g(++1), x(++1 + R mapping higher-scale samples into the corresponding latent space, and generators g(*) Id + P(x and g(++1) 1 + ())((+)) at each scale.
To link adjacent scales, we introduce deterministic maps
p(+) R 1 and g(0) 1 - R
which project embeddings upward and downward between latent spaces.
We jointly train to enforce: Approximate identity each scale:
(E(+)) = P("), (g(o)))) ss p(o+1),
and co-embeddinng consistency: the mapped lower-scale embedding f (E()) should align with the higher-scale embedding (x((t(((())) and vice versa via (t))
Formally, we optimize objectives of the form:
L = (f (E()
(2)
o(plot)) (((tt)))))))
(3)
+ (S glot)) (z(*t)))
(4)
+ (gltt
(a)")))
(5)
where d is divergence or distance (e.g., KL divergence, Wasserstein distance) defined by the generative model. One natural approach would be let f(a), (()) both be the identity, forcing the model learn co-embedding across scales. But this may be too rigid and we might prefer more flexilbity practice.
This bi-dircctional coupling ensures that embeddings at adjacent scales are mutually predictive and geometrically aligned, while each scale individually satisfies distributional invariance and approximate identity. The framework naturally generalizes to hierarchies involving more than two scales by recursively composing maps f()) and gh*) across levels.
35

---

Broader impacts
Generative distribution embeddings provide a general framework for modeling data across scales. They are broadly applicable wide variety of problems, including those with direct societal consequences. for example healthcare. In these settings, will be critical consider any potential inequities induced by GDEs, is the case for any modelling approach. Lastly, we acknowledge the environmental impact this paper, which used nontrivial amounts of computational resources, estimated to be about 54kg CO2.
36

---

NeurIPS Paper Checklist
1. Claims
Question: the main claims made the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
The abstract and introduction concretely state the main theoretical and empirical results of the paper, and enumerate the demonstrated applications of our method.
Justification:
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much results can he expected to generalize to other settings.
is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.
Limitations
Question: Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification: We have separate limitations subheading under the Discussion section. We clearly state key limitations of our method (assumption exchangability, etc).
Guidelines:
The answer NA means that paper has no limitation while answer means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create separate "Limitations" section their paper.
paper should point out any strong assumptions and how robust results are to violations of
these assumptions (e.g.. independence assumptions, noiseless settings, model well-specification,
asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect the scope of the claims made, e.g., if the approach was only tested few datasets with few runs. In general, empirical results often depend implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or speech-to- text system might not be used reliably to provide closed captions for online lectures because fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable,, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, worse outcome might that reviewers discover limitations that aren't I acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory assumptions and proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Yes]
Justification: Theoretical results are stated with complete proof and assumptions the Appendix of the paper. Informal versions of theoretical results are provided in the main text.
Guidelines:
The answer NA means that the paper does not include theoretical results.
37

---

All the theorems, formulas, and proofs in the paper should be numbered and cross cross-reference.
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proofs can cither appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors encouraged to provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by
formal proofs provided in appendix or supplemental material.
Theorems and Lemmas the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data provided not)?
Question:
Answer: [Yes]
All experimental details are provided in Appendix. Moreover, experimental results can be reproduced by running the code in the provided (anonymized Github repository). Models can he trained using the appropriate experiment configs in the conf config/en exppe/ieeeeet///////eee directory, and figures from the paper be reproduced by running notebooks the notebooks/ directory.
Justification:
Guidelines:
The answer NA means that the paper does not include experiments.
the paper includes experiments, No answer to this question will not be perceived well by the reviewers: Making the paper reproducible important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For novel architecture, describing the architecture fully might suffice,
example, if the contrihution the contribution is specific model empirical evaluation, it may necessary to either make possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., the case of a large language model), releasing model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily new algorithm, the paper should make it clear how to reproduce that algorithm.
(b) If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution a new model (e.g., a large language model), then there should either be a way to access this model for reproducing the results or a way to reproduce the model (e.g.. with open-source dataset instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, which case authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-soure, models, it may be that access to the model is limited in some way (e.g.. registered users). but it should be possible for other researchers have some path reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, described in supplemental material?
Answer: [Yes]
Justification: Code datasets are made publicly available, code necessary to reproduce results are provided with documentation.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips.cc//nel
quuddess/oodeuubbmisssooooooo... for more details.
38

---

While we encourage the release of code and data, we understand that this might not be possible. "No" is acceptable answer. Papers cannot be rejected simply for not including code, unless this central to the contribution (e.g.. for a new open-soure benchmark).
The instructions should contain the exact command and environment needed the results. See the NeurlPS code and data submission guidelines fattpe://nipen cc/ cc/public/ guides/ https/CooSSuubiisseee ittdeesuuuisssiooooiooy.. for more details.
reproduce
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only a subset of experiments are reproducible, they should state which ones omitted from the script and why.
At submission time, preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code is permitted.
Experimental sefting/details
Question: Does the paper specify all the training and test details (e.g.. data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary understand the results?
Answer: [Yes]
Justification: All experimental details (including train/test splits and model implementation choices) are provided in the Appendix, and can be found in the accompanying (anonymized) Github repository.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to level of detail that is necessary appreciate the results and make of them.
The full details can be provided either with the code, in appendix, supplemental material.
7. Experiment statistical significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate informa- tion about the statistical significance of the experiments?
Answer: [Yes]
Justification: Standard errors are reported where relevant and feasible.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results accompanied by error bars, confidence intervals, or statistical significance lests, at least for the experiments that support the main claims the paper.
The factors variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
It should be clear whether the error bar is the standard deviation the standard error of the mean.
It is OK to report I-sigma error bars, but one should state it. The authors should preferably report 96% CI. if the hypothesis of Normality of errors
a 2-sigma error than state that they not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).
If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments compute resources
Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time execution) needed reproduce the experiments?
Answer: [Yes]
39

---

Justification:
The appendix includes details of the compute resources used for this work. Full internal cluster details will be released after the double-blind period ends.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn make into the paper).
9. Code of ethics
Question: Does research conducted the paper conform, every respect, with NeurIPS Code of Fthics http:///eeiiips.ccpcillll/////aa
Answer: [Yes]
Justification:
We adhere to the code of ethics.
Guidelines:
The answer NA means that the authors have not reviewed the NeurlPS Code of Ethics.
the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive socictal impacts and negative socictal impacts of the work performed?
Answer: [Yes]
Justification:
We discuss potential socictal impacts in the broader impacts section in Appendix of the
paper.
Guidelines:
The answer NA means that there is societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has societal impact or why the paper does address societal impact.
Examples of negative socictal impacts include potential malicious unintended uses (e.g.. disinformation, generating fake profiles, surveillance), fairness considerations (c.g., deploy- ment technologics that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimatc point out that improvement in the quality of generative models could be used generate deepfakes for disinformation. On the other hand, it not needed point out a generic algorithm optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there are negative socictal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition attacks, mechanisms for monitor- ing misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data or models that have high risk for misuse (e.g., pretrained language models, image generators, scraped datascts)?
Question:
Answer: [NA]
Justification: We do not release datasets models with high risk for misuse.
40

---

Guidelines:
The answer NA means that the paper poses no such risks.
high risk misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines restrictions to access the model or implementing safety filters.
Released models that
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but encourage authors take this into account and make best faith effort.
12. Licenses for existing assets
Question: Are the ereators or original owners assets (e.g., code, data, models), used in the paper, properly credited and are license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification: We cite the datasets, code, and models used in the paper.
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset is used and, possible, include a URL.
The name of the license (c.g., ., CC-BY 4.0) should be included for cach asset. For scraped data from particular source (e.g.. website), the copyright terms of service of that source should be provided.
assets are released, the license, copyright information, and terms of use the package should he provided. For popular datasets, paperaai thcode. com/ dataseta has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
If this information is not available online, the authors encouraged to reach out to the asset's creators.
13. New assets
Question: Are new assets introduced the paper well documented and is the documentation provided alongside the assets?
Answer: [Yes]
Justification:
The code provided in the accompanying repository are dell-documented.
Guidelines:
The answer NA means that the paper does not release assets.
Researchers should communicate the details of the dataset/codemmodel as part of their sub- missions via structured templates. This includes details about training. license, limitations, etc.
The paper should discuss whether and how consent obtained from people whose asset used.
At submission time, remember anonymize your (if applicable). You can either create an amonymized URL or include an anonymized zip file.
14. Crowdsourcing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification:
We not perform crowdsourcing experiments or research with human subjects.
Guidelines:
The answer means that paper does not involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper.
41

---

According to the NeurIPS Code of Ethics, workers involved in data collection, curation, other
labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals or equivalent for research with human subjects
Does the paper describe potential risks incurred by study participants, whether such risks were disclosed the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Question:
Answer: [NA]
Justification: No human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval. you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and expect authors adhere the NeurIPS Code of Ethics and the guidelines for their institution.
initial submissions, not include any information that would break anonymity applica- ble), such as institution conducting the review.
16. Declaration of LLL usage
Question: Does the paper describe the usage of LLMs it an important, original, or non-stande aon-standard component of the core methods in this research? Note that if the LLM used only for writing, editing. formatting purposes and does not impact the core methodology, scientific rigorousness, originality of the research, declaration is required.
Answer: [NA]
Justification: LLMs not important component of this work.
Guidelines:
The answer NA that the method development in this research does not involve LLMs as any important, original, or non-standard components.
Please refer our LLM (https: http::/s//muopgccc//////orrrrrer should or should not described.
for what
42


Reviews:
[{"review_id": "tFH5BokUSy", "reviewer": "Reviewer_aHtV", "summary": "Generative Distribution Embeddings (GDE) is a framework that extends autoencoders to learn representations of entire probability distributions rather than individual datapoints. An encoder ingests a set of samples (an empirical distribution) and produces a latent code, while a conditional generator (in place of a traditional decoder) uses that code to generate new samples, aiming to match the original distribution. A key requirement is distributional invariance of the encoder – the representation should depend only on the underlying distribution and not on particular sample realizations. The authors provide a theoretical grounding, showing that GDE encodings behave as predictive sufficient statistics that capture the underlying distribution’s structure while abstracting away sampling noise. Notably, distances in the learned latent space correlate with 2-Wasserstein (W2) distances between distributions, and linear interpolations in latent space correspond to approximate optimal transport paths for simple distribution families. Empirically, GDE is benchmarked on synthetic distribution datasets, outperforming prior methods in reconstructing distributions. The paper further demonstrates GDE’s versatility on six diverse tasks in computational biology – from modeling cell populations and predicting single-cell perturbation responses to designing DNA sequences – showcasing strong performance and meaningful distribution-level embeddings across these domains.", "questions": "**Questions:**\n\n- On line 76, shouldn't the notation be $S_{\\cdot,k}$ to be consistent with the notation introduced earlier?\n- When using a DDPM as a decoder, could you clarify how the training is performed? Specifically, do you backpropagate through the entire diffusion inference procedure?\n- Could you provide a more detailed explanation of how the chosen architectures are modified or constrained to achieve distributional invariance?\n\n**Some suggestions:**\n\n- L70, it would be helpful to clearly define the input and output domains of the encoder and decoder, explicitly stating that: (i) The encoder takes as input a finite set of points (an empirical distribution) and outputs a latent embedding, (ii) The decoder takes this latent embedding as input and outputs samples forming a set of points, effectively reconstructing an empirical distribution. This clarification would facilitate understanding and immediately highlight the architecture’s core functionality.\n\n- Section 3 is somewhat difficult to follow. It appears to describe how to construct the meta-distributions from labeled datasets, covering both discrete and continuous labels. However, the explanation—particularly for continuous labels—should be clarified to make the procedure more transparent and easier to understand.\n\n- Line 78 is somewhat misleading and interrupts the reading flow. It doesn't seem particularly helpful since there is obviously a one-to-one correspondence between the sample set \\(S_m\\) and the empirical distribution \\(P_m\\). It might be clearer and more concise to omit or rephrase this detail.\n\n-For the evaluation on the synthetic data benchmark, I'd recommend also computing the Sinkhorn divergence as a metric for the low-dimensional (i) Gaussian and (ii) Gaussian mixture settings. Given the low dimensionality ($d=5$), directly computing the Sinkhorn divergence in the ambient space is computationally feasible and would provide a richer comparison. Currently, the authors only compute the $W_2$ distance between Gaussian approximations of the empirical distributions (similar to the FID metric), which assesses only reconstruction accuracy for the first two moments.", "limitations": "Yes, the authors have included a dedicated section discussing the limitations of their approach.", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 2, "significance": 3, "originality": 2, "strengths_and_weaknesses": "**Strengths.**\nThe proposed GDE framework is conceptually strong and general, unifying conditional generative modeling with representation learning for distributions. It subsumes or generalizes prior approaches like kernel mean embeddings and a recent “Wasserstein Wormhole” method under one umbrella. The theoretical contributions are substantial: the authors prove that under mild conditions GDE encoders learn asymptotically sufficient summaries of distributions, and they empirically confirm that the latent space geometry mirrors true distributional geometry (with latent Euclidean distances correlating highly with W2 distances and latent interpolation recovering transport geodesics). The experiments are extensive and convincing. On synthetic benchmarks (Gaussians, Gaussian mixtures, and image-defined distributions), GDE achieves lower reconstruction divergences than baselines (e.g. halving the Wasserstein error of a kernel mean embedding + DDPM and vastly outperforming the Wormhole approach). Furthermore, GDE’s practical impact is evidenced by strong results on real-world tasks – for instance, it improves the mutual information in predicting future cell states by over 2 bits compared to a baseline embedding, and it boosts the $R^2$ in zero-shot gene perturbation prediction (versus a mean-only predictor) from ~0.38 to ~0.46 while reducing error. These gains illustrate that GDE’s distribution-level representations are meaningfully more informative than simpler approaches. \n\n**Weaknesses.**\nThe ambition of covering many scenarios comes with trade-offs. Some experimental evaluations feel a bit shallow or limited in baseline comparison; each biology application uses GDE in a novel context, but often the only point of comparison is either a naive approach (e.g. predicting only distribution means) or one specialized baseline (e.g. Wasserstein Wormhole), rather than a broader suite of competing methods. This makes it harder to pinpoint how universally superior GDE is. Additionally, the assumption of distributional invariance in the encoder design, while well-motivated, imposes a non-trivial constraint – not all standard set encoders qualify (e.g. simple DeepSets can be sensitive to sample size), so careful architectural choices (graph neural networks with pooling, attention mechanisms, etc.) are needed. The paper introduces new terminology and a combination of ideas from optimal transport, information geometry, and generative modeling, which means the exposition is quite dense; readers not already familiar with these backgrounds might struggle with clarity in places. Lastly, the method’s scope of novelty could be perceived as somewhat incremental (combining known building blocks) – the true innovation lies in integration and scale, rather than a fundamentally new algorithmic component, and this is something to keep in mind when weighing its contributions. Additionally, the core engineering and methodological components are not always clearly described. For instance, when using a DDPM as the decoder, it appears that the model is trained by backpropagating through the diffusion model's inference process. This effectively makes the training simulation-based and computationally demanding—an important detail that the authors do not explicitly address.", "was_revised": true, "final_justification": "Thank you for your thoughtful responses to my comments. I'm satisfied with the clarifications provided and have decided to increase my score to 5 accordingly.", "rebuttal": "Thank you for this insightful review.\n\n> Some experimental evaluations feel a bit shallow or limited in baseline comparison; each biology application uses GDE in a novel context, but often the only point of comparison is either a naive approach (e.g. predicting only distribution means) or one specialized baseline (e.g. Wasserstein Wormhole), rather than a broader suite of competing methods.\n\nWe appreciate this feedback. While our synthetic benchmarks evaluate a broad range of methods, the reviewer is correct that we have not compared against large panels of methods for each of the applications in Sec. 6. We fully agree that it would be more satisfying to benchmark against several methods in each case. However, due to the scale of the datasets, comparisons to a large suite of competing methods may become impractical.\n\nIn several cases we have compared against a state-of-the-art method. In Sec. 6.2, we use the Wasserstein Wormhole method which is state-of-the-art for cell set representations. In Sec. 6.3, a linear model for predicting perturbation response may appear naive, but in fact achieves state-of-the-art performance for the task [AHA24, WWS+25]. In Sec. 6.7, we include the widely-used ESM model as a baseline. \n\nIn addition to these existing comparisons, we have now included the additional baseline of scVI [LRC+18], a VAE-based method, for our two single-cell transriptomic applications (Sec. 6.2, 6.3). In order to adapt scVI to distributional tasks, we use the mean of sample-wise scVI embeddings. \n\nWe find in both cases that GDEs outperform scVI. For lineage tracing (6.2), scVI achieves 0.86 bits of mutual information -- lower than both Wasserstein Wormhole and GDEs. For perturbation prediction (6.3), scVI achieves $R^2 = 0.42$ and $\\text{MSE}=1.55$, worse than GDEs but better than a pure linear model. We will add these results to the revised manuscript. \n\n[AHA24] Ahlmann-Eltze, C., Huber, W., & Anders, S. (2024). Deep learning-based predictions of gene perturbation effects do not yet outperform simple linear baselines. BioRxiv, 2024-09.\n\n[WWS+25] Wu, Y., Wershof, E., Schmon, S. M., Nassar, M., Osiński, B., Eksi, R., Yan, Z., Stark, R., Zhang, K., & Graepel, T. (2025). PerturBench: Benchmarking machine learning models for cellular perturbation analysis. arXiv.\n\n[LRC+18] Lopez, R., Regier, J., Cole, M. B., Jordan, M. I., & Yosef, N. (2018). Deep generative modeling for single-cell transcriptomics. Nature methods, 15(12), 1053-1058.\n\n> Additionally, the assumption of distributional invariance in the encoder design, while well-motivated, imposes a non-trivial constraint – not all standard set encoders qualify (e.g. simple DeepSets can be sensitive to sample size), so careful architectural choices (graph neural networks with pooling, attention mechanisms, etc.) are needed. \n\nThe reviewer is correct that GDEs require careful choices for encoder architectures. In order to make the constraints clear for practitioners, we will include a \"recipe\" with clear guidelines constructing GDE architectures (see response to reviewer eYba), which shows which pooling functions are permissible under the constraint of distributional invariance.\n\n> The paper introduces new terminology and a combination of ideas from optimal transport, information geometry, and generative modeling, which means the exposition is quite dense; readers not already familiar with these backgrounds might struggle with clarity in places. \n\nWe agree and aim to make the paper as accessible as possible. Currently, in Appendix C (Background), we briefly review notions of sufficiency, Wasserstein spaces, and stastical manifolds. In our revision, we will include additional sections on optimal transport and recent generative modelling approaches.\n\n> Additionally, the core engineering and methodological components are not always clearly described. For instance, when using a DDPM as the decoder, it appears that the model is trained by backpropagating through the diffusion model's inference process. This effectively makes the training simulation-based and computationally demanding.\n\nWe apologize for the lack of clarity -- **we do not backpropagate through the entire diffusion inference procedure**, which, as the reviewer correctly points out, would be computationally  costly.\n\nInstead, we optimize the encoder parameters through the standard conditional diffusion loss, not its full sampling chain. More generally, **the divergence $\\mathfrak{d}$ in the GDE loss function will always just be the standard training objective for the conditional generator.** In many ways this simplicity is the appeal of the approach: we do not need to make any modifications to the generator or its training, just condition on the latent representation from the distribution encoder and backpropogate through both networks.\n\nWe will explicitly specify this training procedure in our revised manuscript (see proposed revision of Sec. 2 in response to reviewer N5oL).\n\n> On line 76, shouldn't the notation be $S_{\\cdot, k}$ to be consistent with the notation introduced earlier?\n\nThank you for highlighting this. Here, we were referring to the multiset constructed by duplicating each element of the sample set $k$ times. In our revised manuscript, we will more clearly refer to this as \n\n$$S_m^{\\\\{k\\\\}} := \\bigcup_{i=1}^k S_m$$\n\n> When using a DDPM as a decoder [...] do you backpropagate through the entire diffusion inference procedure?\n\nWe do not backpropagate through the inference procedure (see response above).\n\n> Could you provide a more detailed explanation of how the chosen architectures are modified or constrained to achieve distributional invariance?\n\nThe key architectural choice to maintain distributional invariance is in the pooling function used by the encoder. In Appendix D.2., we enumerate pooling functions which satisfy the requirements of distributional invariance. In our implementations we have opted for sample mean pooling. However, we show that sample median, and general M/Z estimators will also lead to distributional invariance -- while sum pooling does not (see Appendix Proposition 4). \n\n> L70, it would be helpful to clearly define the input and output domains of the encoder and decoder\n\nThanks for this suggestion. We have implemented it in our revised Sec. 2 (same content rewritten to increase clarity) which can be found in our response to reviewer N5oL.\n\n> Section 3 is somewhat difficult to follow [...] the explanation—particularly for continuous labels—should be clarified\n\nThank you for this suggestion. We will  provide pseudocode for sampling with continuous labels in the revised appendix. Our proposed text is as follows:\n\n---\n\n**Sampling sets from datasets with continuous labels**\n\nTo construct sets from datasets where the label information is continuous, we can sample sets from label neighborhoods (rather than simply grouping by label as in the discrete case). Below, we provide example code for such a sampling approach\n\n```\n# Inputs:\n# D = {(x_k, y_k)}_k : dataset with continuous labels y_k\n# N : number of sets to construct\n# m : set size\n\ndef construct_sets_from_continuous_labels(D, N, m):\n    sets = []\n    for i in range(N)\n\n        # 1. sample target label y_star (e.g. uniformly over label range)\n        y_star = sample_target_label(D)\n\n        # 2. compute similarity weights (using kernel of choice)\n        weights = [exp(-d(y_k, y_star)**2 / (2 * sigma**2)) for _, y_k in D]\n\n        # 3. normalize weights\n        total = sum(weights)\n        weights = [w / total for w in weights]\n\n        # 4. sample m indices according to weights\n        indices = sample_indices(weights, m)\n\n        # 5. form the set\n        S_i = [D[idx][0] for idx in indices]\n        sets.append(S_i)\n    \n    return sets\n```\n\n> Line 78 [...] doesn't seem particularly helpful since there is obviously a one-to-one correspondence between the sample set (S_m) and the empirical distribution (P_m).\n\nWe apologize -- this seems to be a miscommunication which should be resolved by our revised $S^{\\\\{k\\\\}}$ notation above.\n\nWhile in many cases distinct $S_m$ correspond to distinct $P_m$, there are subtle yet important cases where multiple sample sets can result in the same empirical measure. Permutations $\\pi (S)$ and proportional duplications $S^{\\\\{k\\\\}}$ (i.e., copying every data point $k$ times) result in the same empirical measure $P$. These two cases motivate our requirements for distributional invariance.\n\n> For the evaluation on the synthetic data benchmark, I'd recommend also computing the Sinkhorn divergence as a metric for the low-dimensional (i) Gaussian and (ii) Gaussian mixture settings.\n\nThank you for this suggestion. We have now also computed the Sinkhorn divergence in the low-dimensional cases (see below). The results are qualitatively similar, showing that GDEs outperform alternatives.\n\nSinkhorn divergences for synthetic data experiments from Table 1:\n| Method        | MVN  | GMM  |\n|---------------|------|------|\n| GDE           | 0.11 | 0.32 |\n| KME+DDPM      | 0.14 | 0.40 |\n| W2 Wormhole   | 0.21 | 0.83 |"}, {"review_id": "JdXI4tsWJM", "reviewer": "Reviewer_45Ka", "summary": "The paper presents Generative Distribution Embeddings (GDE), a framework for learning representations of probability distributions by combining distribution-invariant encoders with conditional generative models. The authors formalize the concept of distributional invariance, building on empirical process theory and information geometry, and connect GDEs to predictive sufficiency and statistical manifold geometry. The proposed method is benchmarked in low dimensional synthetic data and high dimensional structured domains including single cell omics, DNA methylation, protein sequence modelling and image based phenotyping, each presenting unique challenges for distribution level modelling. The framework leverages modern generative models and suports flexible grouping strategies to enable application to several scientific questions. \nEmpirically, GDEs outperform kernel mean embeddings and Wasserstein Wormholes, and can recover Wasserstein geometry in latent space even in for high dimensional data. The main contributions are the unification of prior approaches, empirical demonstration of geometric properties, and broad applicability.", "questions": "Questions and comments provided in addition to what was mentioned in the strengths and weaknesses.\n- The formalization of \"distributional invariance\" is necessary for the framework, but it is not clear that this property is sufficient for learning useful or interpretable representations in practice, especially when the meta-distribution Q is misspecified or when data are not IID. Can the authors provide theoretical or empirical evidence for the robustness of GDEs to such misspecification, and clarify whether distributional invariance alone is enough for practical utility?\n- The paper positions GDEs as a unifying framework and demonstrates empriical advantages over baselines. Part of the novelty is, indeed, in the synthesis, as well as the breadth of application. I think the authors should do a bit more to clarify the necessity and sufficiency of the formalization: what are the concrete advantages of the GDE abstraction over existing methods (e.g., kernel mean embeddings, Neural Statistician, flow matching) in practical settings?", "limitations": "yes", "rating": 5, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 4, "clarity": 4, "significance": 4, "originality": 4, "strengths_and_weaknesses": "**Strengths:**\n- Rigorous formalization of distributional invariance, with clear connections to foundational work in empirical process theory and information geometry, and grounding the framework in predictive sufficiency, a property supporting that the learned representations capture the essential structure of distributions for downstream inference.\n- Empirical demonstration that GDEs can recover Wasserstein geometry in latent space for simple distributions, which is concrete supporting evidence for the presented theoretical claims.\n- General framework, allowing any distributionally invariant encoder and conditional generator and grouping strategies, and leveraging advantages in modern generative modelling to address the distribution level inference across scientific domains. The success of this framework depends on making sensible, domain informed choices for encoder, generator, and grouping strategies. The paper demonstrates this by benchmarking on several downstream tasks involving high dimensional biological data, adapting architectures and grouping to each domain, and showing strong results in each case. \n- Limitations are explicitly acknowledged and discussed (in particular with what concerns the curse of dimensionalty of Wasserstein metrics that are central to this work), and the manuscript is generally clear.\n\n**Weaknesses:**\n- The evaluation metrics (Wasserstein, Sinkhorn) are less informative in high-dimensional spaces. The paper does mention this, but in the evaluation does not include alternative metrics or qualitative analyses.\n- On interpretability: the method is benchmarked on several computational biology problems. This is not the main focus of the paper, but it would have been valuable to see more analysis of interpretability/biological relevance to close the loop.\n- Furthermore, the impact of grouping strategies and prior selection (meta-distribution Q) is not systematically analyzed. The choices for this component can lead to uninterpretable or misleading representations, so I would have liked to see some discussion on this. Indeed, this could provide further principled guidance on reproducibility to practitioners, which I think could be better addressed by the paper.", "was_revised": true, "final_justification": "i have read the rebuttals and the other reviews of this paper. I am happy to maintain my current score.", "rebuttal": "Thank you for this thoughtful review.\n\n> The evaluation metrics (Wasserstein, Sinkhorn) are less informative in high-dimensional spaces. The paper does mention this, but in the evaluation does not include alternative metrics or qualitative analyses.\n\nThanks for this observation. We now address this issue for our high-dimensional benchmarks (MNIST, FMNIST) by (1) including Frechet distance in pretrained Resnet18 embeddings (akin to the widely-used FID) as an evaluation metric and (2) including visualizations of generated images in our benchmarks. Both additional results align with the results provided in our existing benchmark: GDEs outperform existing methods (Wasserstein Wormhole and KME). We provide the FID results below. While we are not able to provide the generated images (due to new response format), we will include these in the appendix of our revised manuscript. The GDE generated images are clearly visually closer to the ground truth than Wasserstein Wormhole or coupled KME+DDPM.\n\nFrechet distance in Resnet18 representation space for image experiments from Table 2:\n| Method        | MNIST  | FMNIST  |\n|---------------|------|------|\n| GDE           | 103.5 | 132.4 |\n| KME+DDPM      | 110.2 | 146.9 |\n| W2 Wormhole   | 524.05 | 656.1 |\n\n> On interpretability: the method is benchmarked on several computational biology problems. This is not the main focus of the paper, but it would have been valuable to see more analysis of interpretability/biological relevance to close the loop.\n\nWe appreciate this feedback. We have now more carefully interpreted the results of our clonal representation learning model (Sec. 6.2). In particular, we examine the clones for which GDE fate predictions are high quality, while Wasserstein Wormhole predictions are poor quality (as quantified by the difference in point-wise  mutual information, $\\Delta \\text{pMI}$). Of the 10 clones with largest $\\Delta \\text{pMI}$, 8 clones have multiple differentiation fates (e.g., both Neutrophil and Monocyte fates). This observation suggests that performance improvements are due to GDEs more effectively capturing complex fates (combination of two cell types). We will include these results in the appendix of our revised manuscript.\n\nBelow, we also make some informal observations about biological relevance which we will include in Appendix B:\n\n- **Perturbation prediction in optical pooled screening data.** The genes for which GDEs most accurately predict perturbation responses are typically related to cell cycle (e.g., DONSON). This suggests that GDEs have learned representations which capture nuclear features related to abnormal cell cycle progression.\n\n- **Spatiotemporal modeling of viral lineage distributions.** In this application we find that  GDE representations are organized primarily by time and contain a much weaker spatial signal. This is in line with the intuition that there is significantly lower spatial variation in viral protein sequences than variation in time, due to fast transmission.\n    \n> Furthermore, the impact of grouping strategies and prior selection (meta-distribution Q) is not systematically analyzed. The choices for this component can lead to uninterpretable or misleading representations, so I would have liked to see some discussion on this. Indeed, this could provide further principled guidance on reproducibility to practitioners, which I think could be better addressed by the paper. (From below) ... [How does the encoder work] when the meta-distribution Q is misspecified or when data are not IID\n\nWe agree that the choice of metadistribution has a significant impact on the learned representations. We demonstrate this quantitatively in Sec. 5.3, Fig. 5, where  we show that representations can be stretched or concentrated in certain regions depending on choice of metadistribution $Q$.\n\nAs you correctly observe, this introduces a critical implementation choice for practitioners. However, we view this consideration as a key feature of the method -- GDEs can be adapted to learn task-specific geometries by adjusting the metadistribution $Q$.\n\nWe will emphasize this point and provide suggestions for practical choices of metadistribution in the \"recipe\" we are adding to the revised manuscript (text included in response to reviewer eYba). \n\n> The formalization of \"distributional invariance\" is necessary for the framework, but it is not clear that this property is sufficient for learning useful or interpretable representations in practice... Can the authors provide theoretical or empirical evidence for the robustness of GDEs to such misspecification, and clarify whether distributional invariance alone is enough for practical utility?\n\nWe agree that this is an important point to clarify. \n\nThere are two key insights in Thm 1: first as the reviewer notes, we show that the encoder cannot include more information than the empirical distribution; the flip side of the theorem proves that for IID data a distributionally invariant encoder can, in principle, express any property of the underlying distribution. In practice we find that this expressivity is very effective, and across all of our applications (Secs. 6.1–6.7) an empirical-measure based encoder learns representations that are useful and interpretable.\n\nWhen data within a set are not IID (for example due to temporal, spatial, or network dependence) the situation changes. In these cases the empirical distribution is no longer the correct invariant object, so an encoder that simply pools IID samples is not sufficient. Instead, one must adapt the notion of invariance to respect the structure of dependence. For example:\n- for time series or sequential data, one could pool over blockwise empirical distributions or use causal attention mechanisms;\n- for spatial or network data, one could pool over neighborhood empirical distributions or graph orbits.\n\nOnce the correct invariant is used, the same arguments underlying Thm. 1 and Thm. 2 can be extended using tools such as Stein’s method for dependency graphs and martingale CLTs, which allow analogous guarantees under suitable mixing or dependence conditions. \n\nIn short, for IID data Thm. 1 shows that distributional invariance is both necessary and sufficient for universal expressivity, and in practice we see that this simple approach already works well. For more structured dependence, the same principle applies after modifying the invariant, and we view this as an important and exciting avenue for future work. **To clarify these points we will include an additional section in Appendix D that outlines these issues and sketches a general approach to extend the core GDE results to dependent data.**\n\nAll this being said, in our viral lineage application (Sec. 6.2) the data is clearly temporally and spatially dependent, yet performance remains reasonably strong, suggesting that modest violations of IID do not totally degrade results.\n\n> The paper positions GDEs as a unifying framework and demonstrates empriical advantages over baselines. Part of the novelty is, indeed, in the synthesis, as well as the breadth of application. I think the authors should do a bit more to clarify the necessity and sufficiency of the formalization: what are the concrete advantages of the GDE abstraction over existing methods (e.g., kernel mean embeddings, Neural Statistician, flow matching) in practical settings?\n\nGDEs strictly generalize existing approaches like KMEs, Neural Statistician, and the Wasserstein Wormhole, providing a systematic framework for combining different encoders and generators while previous works have focused on particular architectures (like the kernels, VAEs, or transformers). For example, GDEs show how we can use diffusion/flow matching as a distribution generator. This flexibility pays clear dividened in our synthetic benchmark (Sec. 6.1), where we show that GDEs more closely recover data distributions than existing alternatives.\n\nBeyond this generic flexibility and the performance it comes with, a core practical contribution of GDEs is that they enable leveraging existing domain-specific architectures for sample-level modelling into distribution-level representation learning models. This insight is particularly important because while we have many domain-specific modelling approaches at the level individual datapoints (for sequences, images, proteins, etc.) we have relatively few domain-specific approaches for learning representations of distributions.\n\nGDEs allow one to immediately build on domain-specific architectural innovations by connecting conditional generative models to distributionally-invariant encoders (e.g., a sample level encoder with mean-pooling).\n\nA clear illustration of this utility is the protein modeling task (Sec. 6.7): here, we repurpose the ESM and Progen2 architectures, which were not designed for distribution-level tasks, to learn representations of spike protein distributions. This repurposing would not be possible with existing approaches which lack the flexibility to integrate domain specific architectures."}, {"review_id": "QROWZMbJYM", "reviewer": "Reviewer_eYba", "summary": "This paper introduces Generative Distributional Embeddings framework for learning representations for distributions based on an _distributionally invariant_ encoder and conditional generative decoder.", "questions": "- A suggestion: a \"recipe\" for building a distributionally invariant encoder and prefix-based conditional generation would be very useful. I see the details present in the supplementary material but bringing it as its own section (separate from the theoretical results) and a figure/table summarizing encoder/decoder architectures for different experiments would make it more accessible.", "limitations": "yes", "rating": 5, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 4, "clarity": 2, "significance": 3, "originality": 3, "strengths_and_weaknesses": "## Strengths\n\n1. Clear theoretical contribution in terms of \"distributionally invariant\" encoder and interpretation as a asymptotic predictive sufficient statistics.\n2. Positions prior work clearly in context of this framework. \n3. Clear demonstration in several computational biology problems. \n\n## Weaknesses\n\n- I did not check the theoretical results in detail\n- Presentation is dense - I understand showing all the six applications in main text is good but given that this is a strong submission, I would remove some of the experimental section to explain the theoretical contribution more clearly. I read a Wasserstein Wormhole and a couple of other references to understand better.", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for this thoughtful review.\n\n> Presentation is dense - I understand showing all the six applications in main text is good but given that this is a strong submission, I would remove some of the experimental section to explain the theoretical contribution more clearly.\n\nThank you for this suggestion. We are now moving Sec. 6.4. to the Appendix. We will use the extra space to include a revised Sec. 2 with clearer exposition (with the same content), see response to reviewer N5oL for text. \n\n> A suggestion: a \"recipe\" for building a distributionally invariant encoder and prefix-based conditional generation would be very useful. I see the details present in the supplementary material but bringing it as its own section (separate from the theoretical results) and a figure/table summarizing encoder/decoder architectures for different experiments would make it more accessible.\n\nThank you for this suggestion. In our revision, we will add a dedicated section summarizing \"recipes\" for constructing GDEs, including a table that enumerates the encoder and generator pairings used in each of our experiments, to enable easy implementation of GDEs. This addition will be included either in the Appendix or after Section 3, space permitting. See our proposed addition below:\n\n---\n\n## Implementation recipe for GDEs\n\nThe GDE framework is instantiated by pairing a distributionally invariant encoder with a conditional generative model. The following steps outline a general recipe for building GDEs across diverse data domains:\n\n1. **Sample from the metadistribution (construct sets)**  \n   Group raw data into sets $S_i = \\\\{x_{ij}\\\\}_{j=1}^m$, where each set reflects a draw from an unknown latent distribution $P_i$. Groupings can be based on discrete metadata (e.g., text by author, reviews by rating, images by label, cell clones, gene perturbations) or continuous metadata (e.g., time, location, expression quantiles). Sets need not be mutually exclusive, meaning a single data point can belong to multiple sets.\n\n2. **Choose a distributionally invariant encoder**  \n   Select or construct a distributionally invariant encoder $\\mathcal{E}$. This selection generally involves (1) using an architecture for element-wise embeddings and (2) pooling across element-wise embeddings with a sample mean (or other M-estimate). We found that architectures with multiple pooling layers, where each layer's pooled output is concatenated with the element-wise embeddings, were particularly effective. This contrasts with pure DeepSets-style architectures that only pool once at the final layer. For deeper architectures, we have found that including skip-connections improves performance, especially if the generator is also a relatively deep network.\n\n3. **Build a conditional generator**  \n   The generator $\\mathcal{G}$ \"decodes\" from latent space back to the sample space. It should be conditionable on $z = \\mathcal{E}(S)$.\n   \n4. **Train via plug-in loss**  \n   Optimize the generator to minimize the generator loss function $\\ell(P_m, \\mathcal G(\\mathcal E(S_m)))$. This loss should be the standard training objective for the conditional generator. This plug-in loss encourages reconstruction of the true distribution.\n   \nThe encoder-generator pairs used for each application in the paper are shown in Table 3.\n\n### Table 3: Encoder and Generator Architectures by Experiment\n\n| Section | Task | Set Construction | Encoder Architecture | Generator Architecture | Notes |\n|--------|------|--------|-----------------------|--------------------------|-------|\n| 6.1 | MNIST, FMNIST | Same image class | see Table 1 | see Table 1 | Synthetic data benchmark |\n| 6.2 | Lineage-traced scRNA-seq |  Same cell clones |  ResNet-GNN | CVAE |  |\n| 6.3 | Genetic perturbation (scRNA-seq) |  Same perturbation |  ResNet-GNN | CVAE |  |\n| 6.4 | Morphological responses (cell images) |  Same perturbation |  2D Conv-GNN | DDPM (U-Net) |  |\n| 6.5 | Tissue-specific methylation |  Same patient; Same tissue type |  1D Conv-GNN | HyenaDNA |Uses prefix conditioning |\n| 6.6 | Yeast promoter quantile decoding |  Expression quantile (continuous) |  1D Conv-GNN | HyenaDNA | Uses prefix conditioning |\n| 6.7 | Viral protein spatiotemporal modeling |  Same sampling month and location |  ESM + mean pooling | ProGen2 | Uses prefix conditioning |"}, {"review_id": "zyvwno6nCe", "reviewer": "Reviewer_N5oL", "summary": "This paper studies a heirarchical perspective on data analysis and encoder-generator mechanisms in data science. The authors propose *generative distribution embeddings*, or GDEs, which appropriately lift the usual notion of encoder-generator onto the space of distributions. This is akin to many works that are pursuing a similar trend, which the authors discuss in the related work. The premise of this work is the following: the authors work with samples \\{x_{ij}\\}_{j=1}^m \\sim P_i which is a high-dimensional distribution, and and P_i \\sim Q for i = 1,\\ldots, n, where $Q$ is a distribution over probability distributions. The goal is to learn an encoder-generator (E-G) from the samples, which they train using standard divergence-like frameworks, and suggest that essentially any existing scheme falls into their framework. They discuss how their approach is related to learning over statistical manifolds, as they relate the learning of encoders to sufficient statistics. They conclude with a suite of numerical experiments largely focused around computational biology, where their method appears to perform very well compared to existing approaches.", "questions": "- Lines195-199: Such a property (preservation of geodesics of distributions over distributions) is not guaranteed right?\n\n- what is the significance of \"2 bits\" in section 6.2 (rather, maybe just talk about what this is measuring)\n\n- section 6.3: why is the baseline a simple linear model?\n\n- Section 6.4: I might be missing something but it seems like there are no other comparisons? \n\n- at the start of section 2, maybe mention that $ i = 1,\\ldots, n$ is the indexing over $P_i$\n\n- caption of fig2: \"PCs\" is not defined --- I assume this means principal components, but this figure is just not clear in my opinion.\n\n- I also insist on somewhat restructuring the main written portion of the article to increase clarity of the paper.\n\n- line 101: extra \")\" after $\\epsilon_k$ I believe", "limitations": "N/A", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "The main strength of this paper is in the novelty of performing encoding-decoding-type schemes for distributions over distributions. The connections to the Wasserstein space is also quite nice. The figures help guide the reader throughout the paper, and the experiments appear quite thorough, though I am far from an expert.\n\nThe main weakness of the paper is the writing---it is very verbose and seems rushed. In particular, there are many typos and a few places with repetitive sentence structure. For example, lines 189-194 are basically repeats of the same sentence twice. Proposition 3 is included between lines 211 and 212 with basically no context or follow-up; proposition 1 is similar (proposition 2 is motivated from the text).", "was_revised": true, "final_justification": "The authors addressed all of my questions and re-wrote Section 2, which I had the most qualms with.", "rebuttal": "Thank you for this insightful review.\n\n> I also insist on somewhat restructuring the main written portion of the article to increase clarity\n\nThank you for this suggestion. We propose the following changes:\n\n1. Moving Sec. 6.4 to the Appendix to allow more room for exposition.\n2. Replacing Sec. 2 with the revised version below. The content remains the same, but with clearer exposition.\n3. Adding a new section \"Implementation recipe for GDEs\" after Sec. 3., or in the Appendix depending on space (see reviewer eYba response for proposed text).\n\n> section 6.3: why is the baseline a simple linear model?\n\nFor the perturbation response prediction task and dataset considered in Sec. 6.3, simple linear models are state-of-the-art and outperform deep learning-based alternatives such as scGPT and GEARS [AHA24, WWS+25].\n\nWe have now added another baseline method, the VAE-based scVI [LRC+18]. To adapt scVI to distributional tasks, we use the mean of sample-wise scVI embeddings as a distribution embedding. Again, we find that GDEs outperform scVI, which achieves $R^2 = 0.42$ and $\\text{MSE}=1.55$ (compare with 0.46 and 1.50 respectively for GDEs).\n\n[AHA24] Ahlmann-Eltze, C., Huber, W., & Anders, S. (2024). Deep learning-based predictions of gene perturbation effects do not yet outperform simple linear baselines. BioRxiv.\n\n[WWS+25] Wu, Y., et al. (2025). PerturBench: Benchmarking machine learning models for cellular perturbation analysis. arXiv.\n\n[LRC+18] Lopez, R., et al. (2018). Deep generative modeling for single-cell transcriptomics. Nature methods.\n\n> Section 6.4: I might be missing something but it seems like there are no other comparisons?\n\nYou are correct -- in this application we do not compare against any baseline. There are two practical reasons for this: (1) there are, to our knowledge, no existing generative models for predicting the morphological response of novel perturbations in optical pooled screen data (2) the scale of the dataset (30M images) precludes easy training of baseline methods.\n\nOur main goal here is to show that GDEs are compatible with image data and can generate realistic images (see Fig. 7) with relevance to an open problem in biology. We do not claim state-of-the-art performance for this particular task. As such, we will move this section (6.4) to the appendix, and use the space for clearer exposition.\n\nAt the reviewer's suggestion we are training a variant of our GDE implementation using a VAE as the generator for the revised manuscript, and we are happy to consider other baselines.\n\n> what is the significance of \"2 bits\" in section 6.2 (rather, maybe just talk about what this is measuring)\n\nHere, \"bits\" refers to the amount of predictive information ($t_1 \\to t_2$) contained in the representations. We will clarify this and provide intuition. At line 283, we will add:\n\n\"This suggests that GDE representations enable more granular prediction of cell fates. If cell fates were discrete and uniformly likely, 1 bit of mutual information would correspond to resolving 2 fates while 3 bits would correspond to resolving 8 fates.\"\n\n> Lines195-199: Such a property (preservation of geodesics of distributions over distributions) is not guaranteed right?\n\nCorrect, this is an empirical phenomenon that we observe in this particular setting. In fact the geodesics will \"adapt\" to the data and the prior over $\\mathcal M$ as we show in Section 5.3. We will emphasize this in the revised manuscript.\n\n> caption of fig2: \"PCs\" is not defined [...] this figure is just not clear in my opinion.\n\nYes, PCs refers to principal components. We will rewrite the caption:\n\n\"**Figure 2: Concentration of distribution embeddings and plug-in loss.** (Left) Distribution of plug-in GDE loss (with diffusion generator) for MNIST image sets of different sizes sampled from the same distribution $P_i$ over MNIST digit identities. (Right) First two principal components of 64D embeddings of sample sets of different size generated by the same $P_i$\".\n\n> Proposition 3 is included between lines 211 and 212 with basically no context or follow-up; proposition 1 is similar (proposition 2 is motivated from the text).\n\nThank you for highlighting this. We have now entirely rewritten Sec. 2 (see below), addressing this comment.\n\n> There are many typos and a few places with repetitive sentence structure. For example, lines 189-194 are basically repeats of the same sentence twice. \n\nWe apologize. We will remove lines 192-194 and carefully correct typos in the revised manuscript.\n\n> at the start of section 2, maybe mention that $i=1,\\dots, n$ is the indexing over $P_i$\n\nThank you -- we have addressed this in the revised Section 2.\n\n> line 101: extra \")\" after $\\epsilon_k$\n\nThank you for catching this. It will be removed.\n\n## Revised Section 2\n## 2 Setting and methods\n### 2.1 Motivating example\nMany modern datasets consist of several large groups of exchangeable samples: for example, single cells or DNA sequences collected from individual patients. Suppose we observe $n$ such groups,\n\n$$\n  \\mathcal D = \\bigl\\\\{S_{i, m} = \\\\{x_{ij}\\\\}^m_{j=1} \\bigr\\\\}^n_{i=1}, \\qquad x_{ij} \\in \\mathcal X\n$$\n\n\nwhere each group $S_{i,m}$ is drawn i.i.d. from an unknown distribution $P_i \\in \\mathcal{P}(\\mathcal{X})$, and the distributions themselves are drawn from a meta-distribution $Q$:\n\n$$P_i \\stackrel{iid}{\\sim} Q,\\qquad x_{ij} \\stackrel{iid}{\\sim} P_i.$$\n\nThis is a classical hierarchical data generating process, which gives rise to a multiscale problem. The subject of interest here is not only individual cells or DNA sequences $x_{ij}$, but the patients $P_i$. As we explore in Sec. 3 and concretely demonstrate in Sec. 6, this setting is broadly applicable beyond these particular examples.\n\nIn practice, there are often two major challenges in modeling this kind of data. First, unit-level data is often inherently noisy. For example, single-cell data suffers from noise due to molecular undersampling [39]. Our goal is to learn patient embeddings which capture distribution-level signal, rather than the sample-level noise. The second challenge is that groups can contain millions of samples (in the case of DNA sequencing reads per patient, $m$ can be $\\sim10^8$). It is computationally infeasible to train a model on all samples simultaneously, but given the inherent noise at the unit level we would benefit from embedding all available samples at inference time.\n\nIn the remainder of Sec. 2, we will show that both of these practical challenges can be overcome by learning distribution embeddings rather than simply encoding sets of samples. The distributional perspective enables models which distill distribution-level signal, and are able to massively scale at inference time to make use of all available data for precise embeddings.\n\n### 2.2 Learning generative distribution embeddings (GDEs)\nWe address this problem with GDEs, which consist of an encoder $\\mathcal E$ that maps a finite set of samples $S_{i,m}$ to a latent representation, and a conditional generator $\\mathcal G$ that maps the latent representation back to the sample space. Formally, we aim to learn $\\mathcal{E}, \\mathcal{G}$ such that\n\n$$\\mathcal G\\bigl(\\mathcal E(S_{i,m})\\bigr)\\ \\stackrel{m \\to \\infty}{\\longrightarrow}\\ P_i\\tag{1}$$\n\nThe training procedure is shown in Algorithm 1:\n```\nFor each minibatch of sets $S_{i,m}$:\n    Compute latent representation $z_i = \\mathcal E(S_{i,m})$.\n    Compute generator loss $\\ell(P_i, \\mathcal G(z_i))$.\n    Backpropagate gradients through $\\mathcal E$ and $\\mathcal G$.\nEnd for.\n```\n\nThe loss $\\ell$ is the standard training objective for the conditional generator (for example, an evidence lower bound for a VAE or a denoising score-matching objective for a diffusion model); we do not need to backpropagate through the sampling process of $\\mathcal G$.\n\n### 2.3 Distributional invariance uniquely enables achieving Eq. (1)\n\nWe show that achieving Eq. (1) is only possible if the encoder satisfies the following two constraints (formal proofs in Appendix D): \n1. Permutation invariance: reordering the samples in $S_{i,m}$ does not change the embedding.\n2. Proportional invariance: duplicating every sample $K$ times ($\\bigcup_{k=1}^K S_{i, m}$) does not change the embedding.\n\nWe refer to an encoder with these properties as *distributionally invariant*: the encoder must depend only on the empirical distribution. So for some function $\\phi$ we can write:\n$$P_{i, m} = \\frac{1}{m} \\sum_{j=1}^m \\delta_{x_{ij}},\n\\qquad\n\\mathcal E(S_{i,m}) = \\phi(P_{i,m}).$$\n\nWe show formally that distributionally invariant encoders can capture any distributional property and furthermore that any non-distributionally invariant architecture can spuriously encode noise features irrelevant to the distribution:\n\n[Proposition 1 box]\n\nBeyond separating signal and noise, distributional invariance has a second crucial consequence: it enables a *central limit theorem for the embeddings*. As the set size grows, $\\mathcal E(S_{i,m})$ concentrates around its population value with Gaussian fluctuations:\n$$\\sqrt{m}\\ \\bigl(\\mathcal E(S_{i,m})-\\phi(P_{i,m})\\bigr) \\ \\xrightarrow{d}\\  \\mathcal N(0,\\Sigma)$$\n\nThis result, illustrated empirically in Fig. 2, is what makes encoding massive sets possible. \n\nBecause the embeddings converge, the plug-in loss\n$$\\widehat\\ell_m = \\ell\\bigl(P,\\mathcal G(\\mathcal E(S_{i,m}))\\bigr)$$\n\nis an unbiased, asymptotically normal estimator of the population loss. This provides theoretical justification for training GDEs on subsets of larger sample sets: the gradient of $\\widehat\\ell_m$ computed on small sets matches (in expectation) the gradient computed using all samples per set. We prove this in Appendix D.2:\n\n[Proposition 2 box]\n\nViolating distributional invariance (for example, by using sum pooling) causes the embedding to depend on set size and breaks this limit theory causing Eq. (1) to fail.  In contrast, **mean pooling and differentiable M/Z-estimators** satisfy these properties."}]

Author rebuttal:

Meta-review: **Summary:** This paper proposes a new concept of Generative Distribution Embeddings (GDEs), extending autoencoders from single-point embeddings to distributional embeddings. The framework consists of an encoder, enforced by distributional invariance, which maps a sample set into a latent code, and a conditional generator that decodes the latent code back into the distribution. Importantly, GDE distances approximately recover the Wasserstein-2 distance between distributions.

**Strengths:** GDEs introduce a novel paradigm for embedding distributions rather than single datapoints, with theoretical guarantees linking latent geometry to Wasserstein distances. Empirically, the framework outperforms existing baselines across both synthetic benchmarks and diverse biological applications.

**Weaknesses:** The paper’s presentation is weak and could be improved for broader accessibility, as all reviewers noted clarity issues. Sections 1–3 lack sufficient references, and most of the proofs are loose and asymptotic, relying on intuition and high-level arguments with limited rigor.

**Review Process**: Reviewers initially raised concerns about clarity, novelty scope, alternative metrics, qualitative analyses, and interpretability. In the rebuttal, the authors addressed most of these concerns, and three reviewers raised their scores from (4, 4, 5, 3) to (5, 5, 5, 4).

**Recommendation**: This paper presents a novel and well-motivated concept of distributional embeddings, supported by theoretical arguments and validated through experiments. The work is both original and impactful, and with reviewers reaching consensus after rebuttal, I am happy to recommend acceptance. The authors are encouraged to improve the presentation and strengthen the proofs to ensure broader accessibility and clarity.
Decision: Accept (poster)