Learning to Embed Distributions via Maximum Kernel Entropy
Authors: Oleksii Kachaiev
Stefano Recanatesi
Venue: NeurIPS 2024 poster (2024)
Keywords: distribution regression
kernel methods
Reproducing Kernel Hilbert Spaces (RKHS)
kernel mean embeddings
data-dependent kernel
unsupervised learning
Primary area: other
TL;DR: The paper introduces a novel unsupervised objective for learning data-dependent distribution kernel based on the principal of entropy maximization in measure embeddings space.

Abstract: Empirical data can often be considered as samples from a set of probability distributions. Kernel methods have emerged as a natural approach for learning to classify these distributions. Although numerous kernels between distributions have been proposed, applying kernel methods to distribution regression tasks remains challenging, primarily because selecting a suitable kernel is not straightforward. Surprisingly, the question of learning a data-dependent distribution kernel has received little attention. In this paper, we propose a novel objective for the unsupervised learning of data-dependent distribution kernel, based on the principle of entropy maximization in the space of probability measure embeddings. We examine the theoretical properties of the latent embedding space induced by our objective, demonstrating that its geometric structure is well-suited for solving downstream discriminative tasks. Finally, we demonstrate the performance of the learned kernel across different modalities.

Learning to Embed Distributions via Maximum Kernel Entropy
Oleksii Kachaiev
Dipartimento di Matematica, Università degli Studi di Genova, Genoa, Italy oleksii wew.si..aaaaieveeeeee wwwhaiiewgaamll.cccmmmmmmm
Stefano Recanatesi
Technion Israel Institute of Technology, Haifa, Israel Allen Institute for Neural Dynamics, Seattle, USA
stefano. rwwanateeii@gaailllll com
Abstract
Empirical data can often be considered as samples from set of probability dis- tributions. Kernel methods have emerged natural approach for learning classify these distributions. Although numerous kernels between distributions have been proposed, applying kernel methods to distribution regression tasks remains challenging, primarily because selecting suitable kernel is not straightforward. Surprisingly, the question of learning a data-dependent distribution kernel has received little attention. In this paper, propose novel objective for the un- supervised learning of data-dependent distribution kernel, based on the principle of entropy maximization in the space of probability measure embeddings. We examine the theoretical properties of the latent embedding space induced by our objective. demonstrating that its geometric structure is well-suited for solving downstream discriminative tasks. Finally, we demonstrate the performance of the learned kernel across different modalities.
Introduction
Most discriminative learning methods conventionally assume that each data point is represented as real-valued vector. In practical scenarios, however, data points often manifest as 'set' of features or 'group' of objects. A quintessential example is the task of predicting a health indicator based on multiple blood measurements. In this case, the single data point of a patient has multiple, or distribution of, measurements. One approach to accommodate such cases involves representing each input point as probability distribution. Beyond mere convenience, is more appropriate to model input points distributions when dealing with missing data measurement uncertainty, often encountered when facing the abundance of data, which commonly presents a challenge for data-rich fields such genetics, neuroscience, meteorology, astrophysies, economics.
The task of regressing mapping of probability distributions real-valued response known as distribution regression. Distribution regression has been successfully applied in various fields, such as voting behavior prediction
[3 dark matter halo mass learning [43], human cancer cells
detection [44), brain-age prediction (6), among others [38 33 68) The versatility and effectiveness
of this framework underscore its power solving complex problems [69 29]. Kernel methods have become widely used approach for solving distribution regression tasks by exploiting kemell between distributions referred to as a distribution kernel. Despite the multitude of proposed kernels, the practical application of kernel methods remains challenging due to the nontrivial choice of the appropriate kernel. While some efforts focus on identifying kernels with broad applicability and
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

Framework: train maximize distributional variance H solve multiple downstream tasks.
Classiication
Embeedinng
kamme
f8
10
Classitcatioo
Explot distrbutiinn karnel K over mean umbedcings axa-e opt resentation gesmetry
Z
H
Inout space
Latent 500ce
Figure 1: Learning to embed distributions. (a) Example of multiple distributions over the input
space. (b) he trainable function fo encodes the input dataset into compact latent space, in our case == The first-level embedding kernel &: induces kernel mean embedding map to H. The encoder is optimized to maximize the entropy of the covariance operator embedding the dataset w.r.t. the second-level distribution kernel K between kernel mean embeddings in H. (d) Utilizing learned data-dependent kernel, downstream classification tasks can solved using tools such as Kernel SVM or Kernel Ridge Regression.
favorable statistical properties [54]. others aim to tailor kernels to the geometric characteristics of specific input spaces [6] Remarkably, the question of learning data-dependent kernels has received limited attention. This study is thus driven by fundamental question: What are the underlying principles that facilitate the unsupervised learnin an effective kernel, one that optimally encapsulates the data properties and is well suited for discriminative learning on distributions?
In this work, we leverage key insight: appropriate selection of the distribution kernel enables the embedding of a set of distributions into the space of covariance operators. Building on this claim that quantum entropy maximization of the corresponding covariance
theoretical idea, operator is suitable guiding principle to learn detta-deendent kernel. This, combined with a careful design of kernel parametrization, let devise differentiable optimization objective for learning data-specific distribution embedding kernels from unlabeled datasets (i.e. unsupervised) (Fig. []. We show that the entropy maximization principle facilitates learning of the latent space with geometrical configuration suitable for solving discriminative tasks [3 20). We empirically demonstrate the performance of our method by performing classification tasks in multiple modalities.
In summary, our unsupervised data-dependent distribution kernel learning framework introduces a theoretically grounded alternative the common practice of hand- picking kernels. Such framework could be further leveraged for generalizing existing learning approaches [61] catalyzing the use of distribution-be based representations within the broader scientific community.
Preliminaries
We first introduce the main concepts necessary to formalize our learning framework: kernel mean embeddings and covariance operator embeddings.
2.1 Kernel Embeddings of Distributions
Consider input space &' and dositive-definit (p.d.) kernel k: X' x X' -> R. Let H be the corresponding reproducible kernel Hilbert space (RKHS) induced by such kernel. Consider a probability distribution P E P(X). The kernel mean embedding map embeds the distribution P as function in Hilbert space:
uP PP = (p{)p) = ((x,,))P(())
((x) dP(x),
where o: X + H. is feature map such that d(x) k(x,
Importantly, if the kernel k characteristic [50], the mapping μ P(X') - injective, implying
that all information about the original distribution is preserved in H. This last property underscores
much of power under recent applications of kernel mean embeddings [4] 14). The natural empirical estimator for the kernel mean embedding approximates the true distribution with finite

---

sum of Dirac delta functions:
1 N
Ap
E (()) E H
where ......, ~ N empirical i..d. samples. The estimator has a dimension-free sample complexity error rate of (NN +) [5]]
Additionally, we denote d3 X' x X + R20 a kernel metric induced by a kernel k. (x - lØ(x) - ((x')l
Note that dz a metric on X if feature map @ injective.
Covariance Operators and Entropy
A second way of mapping probability distribution to Hilbert space can be defined by means of covariance operators. For a given feature map <bx)) k(x,;) - Hf and given probability distribution P P)( the covariance operator embedding is defined as:
Jx σ(z) & ((z)d((2))
Σp Σp=
(4)
where @ a tensor product. Ep is a self-adjoint positive semi-definite (p.s.d.) operator acting on mean embedding w.r.t the feature map æ ((x)) Φ(x) and
H. Such operator be therefore, for a universal kernel k, the map P + Ep is injective (see Bach I
Similarly to Eq. 2 the natural empirical estimator is:
Σp
Σ ((x)) 6(x:)
N
where Fi,..., P are N i.dd samples.
For a translation-invariant kernel k(r,, = (fx - x') normalized such that k(r, x) 1, covariance operator Ep is density operator [.. Henceforth, entropy measures can be applied to it, and the quantum Rényi entropy of the order a can be defined as:
S.(Pp)
log tr [{E))]
log
(6)
where {Aily are the eigenvalues of Ep. The Von Neumann entropy can be seen as a special case of Rényi entropy in the limit + 1. However, in our work, we focus primarily on the second-order case of Rényi entropy, i.e. a (see Carlen (9, Müller-Lennert al. [42], Wilde (3), Giraldo et al. [15 for an in-depth overview of the properties and theory of quantum entropies).
3 Unsupervised Distribution Kernel Learning
3.1 Distribution Regression
In this section, discuss the key topics in distribution regression, including problem setup, the notion of a 2-stage sampling process, and the common solutions to regression employing kernel methods.
Distribution regression extends the common regression framework to the setup where covariates given probability distributions available only through samples. Formally, consider the task of finding a regressor f P(X) + y from the dataset of samples I'M {(P,, H M where P. ((X') are distributions provided as set of i.id. empirical samples 21, .......... EN, P. (see Poczos et al. 45). Szabó et al. 53
54) for comprehensive analysis). viable approach to solving this problem is define kernel K P(2)) P(X) R that is universal in P(X). By setting up such a kernel K, we can utilize Bernel-based regression techniques, with SVM often being the
Assuming & be continuous positive definite (p.d.) kemel and X' to be compact.

---

Distribution Entropy
Kernel Norms
Aolipooo 25 bn
Entrpy
sydal
Disebusoorl by (a)Pa)
nemmdrttt buit
varance
04
Gram matrix
15"
0.0
Komd Gram Matrix Egenvatues
10
10
127
10"
10
00
Figure 2: Properties of the entropy the toy example. Entropy and Distributional Variance for 6 distributions on a sphere as a function of their geometrical arrangement parametrized by 7- (b)
Kernel norms that enter the distributional variance bound. The blue shaded area (difference between
blue and red lines) corresponds to the dotted red line in (a) (up to multiplicative factor). (c) Flattening
of Gram matrix eigenvalues function of
preferred method for classification tasks [39 or Kernel Ridge Regression (KRR) when the space y continuous 37
To this end, several kernels have been proposed over time (see details in Sec. 4
One possibility, proposed by Muandet al. [39 is introduce a kernel in the input space, the so called embedding kernel kemb X x X + IR and exploit the induced mean embeddings /emb P(x)) -> Hemb to map input distributions to points
RKHS. Subsequently, to define second level kernel, distribution kernel Kdistr Hemb X Hemb - R between points in the RKHS Hemb The simplest choice for such a distribution kernel is the linear kernel:
Jfxee k(5,, (((,))))(())dd(()))))
K((p,, (Q) ((ppH/)))
standard alternative to the linear kernel a Gaussian kernel with bandwidth parameter 0:
K (RA((UP, (Q) = exp (- (-2 Ippp - pell melliten
)
= = exp - (uP, (-
which was shown to be universal in ((X) [T]. The Gaussian kernel KRAP can be computed from the linear kerne K2 as di.m (Hr, pQ) Ki(P, P) K) Q) (Ki(P,Q) Q)
In practice, we only have access to a finite number samples of each input distribution, thus the distribution kernel is approximated using the natural estimator for the kernel mean embedding. The related excess risk for the regression solution analyzed Szabó al. [33
3.2 Dataset Embedding
Instead of using standard kernels designed to encapsulate the geometry of the input space, consider learning a -ata-dependent kernel, tailored to the specific properties of the dataset. In a similar vein, Yoshikawa et al. [68] proposed learning optimal kernel (or equivalently, feature map) jointly with the classifier to address the text modality. In this work, we focus on an unsupervised problem, aiming learn a data-dependent kernel between probability distributions without access to classification labels.
We first introduce proper parametrization to ensure both expressivity and robustness followed by the definition of the optimization objective. Leveraging the idea 2-level kernel setup, we define the embedding kernel
Kg : X x
X
- IR Kemb (z), fo()))
(9)
where fe is a trainable encoder function fa X' + 2, Z is a latent encoding space, and kemb is kernel defined on the latent space kgub
x Z - The encoder function fo transforms every
2 As pointed out in Meunier et al. 37 under mild conditions, KRR can also be used for classification problems.

---

input probability distribution P E P(X) into distribution over the latent space Pe E P(2) (Fig. Lap. Furthermore, we denote RKHS corresponding to the kernel kemb as Hemb and the kernel mean embedding map as Memb (see Eq. (
PEP(X) L4 POEP(Z)
kanb
/emmb (P) Hemb.
(10)
These transformations define mean embeddings for each input probability distributions through the first level, embedding kernel kemb (Fig. [b). The second level, distribution kernel Kiistr Homb x Hemb + R is defined over the mean embeddings Memb (P'ss We can now consider to embed dataset DM- - [P: an empirical covariance operator (see Eq. 5 i.e.
DM Hembs Kivtr
EDD
Σ PEDM
Kaiste (Pemb (P), -) & K Kaistr (Hemm (P)))
(11)
M
As Ep encapsulates information about the entire dataset, we term it dataset embedding. With the assumption that the dataset DM is sampled i.i.d. from the (unknown) true meta-distribution D, Ep is a natural estimator to approximate the true covariance operator Σp. To simplify the notation we use Ep in place of Dp unless required by the context.
Both the embedding kernel kemb and the distribution kernel training, learning happens by adjusting the parametrization of the latent space encoder fo. Such separation ensures expressivity while conforming all technical requirements for distribution kernel. Throughout the paper, we make the following assumptions on the latent space Z and embedding kernel
Kdistr remain fixed throughout the
Assumption 3.1. Latent space is compact subset of TRdd Kernel kemb : Z characteristic ionslaaiioninvariin kernel kemb(2, f(1l= =²1l²) such that -f' is completely kemt 2)
p.d.
monotone on (0, oo) (see Definition 2.2.4 of Borodachov et and Vz
The choice of a kernel based on Euclidean distance in the latent space makes its definition similar to that in Weinberger and Tesauro (62), though the parametrization of the encoding process differs. Learning kernels by explicitly learning feature maps has been explored in a wide range of settings In contrast, the parametrization proposed in our study applies a known characteristic kernel to learned latent representation. To facilitate our optimization process (which will be explained shortly), we opt for Z = gd-1 (the d-dimensional hypersphere) and Gaussian kemel both for kemb and Kiistr (see Eq. { research.
[66
57 58 67
We retain other suitable choices as potential avenues for future
3.3 Unsupervised Optimization Objective
This dataset level representation depends on the choice of first and second level kernels k, K and, in turn, on the trainable function fo parameterized by the set of parameters In this work, we propose learning the parameters 0 to maximize quantum entropy of the dataset embedding, i.e..
arg max {S2(Ex) == : log [(E)]]]}
(12)
As we will describe in brief, optimizing this target has clear benefits inherited from the underlying geometry of the setup. But, first, we show how to empirically compute S₂ (Ep).. Building upon previous work 1 we exploit the following property of the covariance estimator:
tr ((>b) = (( KD
21 }
(13)
where KD E MM is the distribution kernel matrix, with [Kp]y = (HistrHP,, HP,). This equation follows directly from the fact that Ed and 37 Kp share the same set of eigenvalues. Leveraging this
By the definition of the encoding process, Ps is push-forward measure. For empirical probability distribution q E, 6(2s) P(X) and measurable map f X the push-forward measure f#g E P(Z) defined as Σ, (x,)).
"See the proof of Proposition 6 in Bach [1
5

---

relationship, we can define tractable unsupervised training loss, which we term Maximum Distribution Kernel Entropy (MDKE), with respect to the parameters of the encoder fo:
CMDKE == -=2p(>p)) log (
1 (D) log KDll (14)
log
where the latter relies on the fact that the Frobenius norm || E, (A),, where (i(A) are eigenvalues A.
The MDKE objective is differentiable w.r.t 0 for commonly used kernels, provided that the encoder fe is differentiable as well. While the entropy estimator S2(Ep) is convex in the kernel matrix Kp. the objective whole is generally not convex in O. However, in practice, we show in Sec. [ mini-hatch Stochastic Gradient Descent (SGD) proves to be an effective method for optimizing this objective. The effectiveness of this optimization process is significantly influenced by the parameters of the Gaussian kernels. We elaborate on the methodologies for kernel bandwidth selection in Appendix B.I]
The Frobenius norm formulation in the loss Eq. (1) significantly reduces computational complexity, However, have observed in some of our experiments, it can lead to the collapse of small eigenvalues Kp, particularly near the optimal value of the objective. To address this challenge we introduced regularized version of the loss LMDKE-R that incorporates optional regularization, based on the determinant Kp inspired by the connection with Fekete points (see details in Appendix (B.2).
3.4 Geometrical Interpretation
The optimization objective is specifically designed to minimize the variance within each distribution (nner-distribution variance) while simultaneously maximizing the spread of distributions over compact latent space 2 Sd-1, This shaping of the distributions embeddings in the latent space facilitates easier separation in downstream tasks. In this section we show that the geometry of the optimal (w.Lt. the MDKE loss) configuration of mean embeddings in the RKHS attains describe properties. For doing so we leverage the notion of distributional variance VH (Definition in Muandet et al. [40].
Definition 3.2. For a set of M probability distributions DM, distributional variance (H(DM) of the mean embeddings in the RKHS H given by
M
(n(DM)
1 M Gr|
Gij.
(15)
M
where G is the M M Gram matrix mean embeddings in H. i.c. Giy= (r,,
40
Here we show that the distributional variance Vn can be equally reformulated into two separate contributions:
1 M ||p, 11 - 1|ppl|
(D() = M
(16)
where P denotes mixture distribution with elements of DD being uniformly weighted mixture components (see proof in the Appendix A.Ipp
The relevance distributional variance for MDKE objective established by the following result.
M probability distributions DM, the second-orler Rényi entropy S₂ of the empirical covariance operator embedding En induced by the choice of Gaussian distribution kernel KRRF over points in the RKHS Homb- as defined in Eq. (8) is upper bounded by the distributional variance V Hna (D) as i.e.
Proposition 3.3. For
S2 S V (Hm..(M))
(17)
27
where y is the bandwidth of the distribution kernel KRBF.
The proof of this proposition is provided in Appendix A.3 This result formalizes the fact that our objective increases distributional variance, pushing up the average squared norm of mean embedding

---

of input distributions while minimizing squared norm of the mean embedding of the mixture. We further explore the geometrical implications of such optimization by formalizing connection between the variance of the distribution and the squared norm of the its mean embedding in RKHS.
Proposition 3.4. Under Assumption 3.7 the maximum norm of kernel mean embedding is attained by Dirac distributions {6z)=ez.
This result is trivial due to the fact that the set of mean embeddings is contained in the convex hull of
{Kemb (z, >ee,, and, under Assumption 3.1 Vz E Z ||Kemb (z. >32mm = Remb 2) = 1.
Proposition 3.5. Under Assumption 3.1 uniform distribution U(2) is unique solution of
arg min PEP(Z)
|| A
2)) IJ Riemb
d))((()
(18)
The key intuition here comes from the fact that minimization of the squared norm of mean embedding in RKHS could be seen as minimization of total interaction energy over the given surface where the potential is defined by the kernel k. Thus Proposition [35 is a special case of Theorem 6.2.1 of Borodachov et al. [7]. Similar setup was used w.r.t. Gaussian potential over the unit hypersphere in Proposition from Wang and Isola [6]]. The reformulation in Eq. (16) together with Propositions 3.4 and 3.5 immediately suggests that the framework could be seen as an extension of the Hyperspheri- cal Uniformity Gap [32 to
spaces of probability distributions This extension maintains the goal of reducing variance among input distributions while maximizing the separability between their means. See Appendix C.2 for broader explanation of the connection.
More generally, utilizing the fact that under Assumption 3.1 the kernel metric daam (see Eq. B is monotonically increasing function of Euclidean metric on the latent space, we establish precise connection between the generalized variance (60 and the norm of the mean embedding. Further details can be found in Appendix A..4.
An attempt to directly optimize VH... using SGD resulted in significantly weaker outcomes. While a thorough mathematical explanation necessitates further investigation, we contend that this issue aligns with the recurring challenge reported across various studies regarding direct optimization over Maximum Mean Discrepancy (MMD). We hypothesize that optimal solutions exhibit similar geometric configurations, while the entropy of the covariance operator providing a smoother objective. Nonetheless, VH... retains its value as an insightful and intuitive measure for describing the geometric configuration of the learned system.
3.5 An Illustrative Example
We use a simple example to illustrate the connection between geometrical configurations of embedded distributions and distribution kernel entropy S2 (2(E)) (see Fig. 2). We sample a number of points from 6 different Gaussian distributions and project on a sphere S varying their projected variance y. As 7 decreases, the distributional variance of the overall distribution of Gaussians increases (Fig. 2a For very small y each distribution converges to point (a Dirac distribution). This results in the entropy interpolating between lower and upper bounds, demonstrating how entropy behaves in response changes in distribution variance. Fig. 2b showcases the behavior of two terms comprising distributional variance (Eq. (6):: ): the average kernel norm of the distributions alongside the kernel norm of the mixture. The increase in entropy and variance corresponds to a 'flattening' effect on the spectrum of the distribution kernel matrix. This example provides a simplified picture of how input distributions configurations influence kernel entropy.
3.6 Limitations
Runtime complexity. The applicability of a data data-dependen distribution kernel to solving discrimina- tive tasks relies on the structure of the dataset being well-suited for distribution regression modeling. The model performs best when the number of input distributions is relatively small (e.g., thousands rather than millions), while the number of samples per distribution is large. It is crucial to note that the computational complexity of the proposed method, which is common concern in practical applications, most favorable for the tasks described. A detailed analysis of runtime complexity can be found in Appendix B.3
7

---

Broader impact. We wish to emphasize that the distribution regression framework has emerged as powerful tool for analysis and predictive modeling. especially domains where traditional methods face challenges, including social science, economics, and medical studies. We urge researchers and practitioners applying distribution regression these areas to give special consideration issues such as bias, fairness, data quality, and interpretability, aspects that are currently under-researched in the context of distributional regression, largely due to the relative novelty.
4 Related Work
4.1 Distribution Regression
Distribution regression introduced in Poczos al. [5), while the seminal work of Szabó et al.
provides a comprehensive theoretical analysis of this regression technique. A natural approach to solving distribution regression problems involves using kernels between measures. Notable examples
4
include the Fisher kernel [17]. Bhattacharyya kernel (18), Probability product kernel [19] kernels
based on nonparametrie divergence estimates [52], and Sliced Wasserstein kernels 23
Muandet
et al. (39 proposed leveraging the mean embedding of measures in RKHS, and Szabó et al. [53]
provided theoretical guarantees for learning a ridge regressor from distribution embeddings Hilbert space to the outputs. Distribution kernels have been successfully applied in various kernel based methods, such SVM (39), Ridge Regression (37]. Bayesian Regression (28]. and Gaussian Process Regression 2 They have also been adapted for different modalities like distribution to distribution regression sequential data [29], and more. Some learning paradigms can be considered closely related to distributional classification settings, such as multiple instance learning, where group-level information (i.e., labels) available during training 770 26, 25). For an in-depth exploration of the diverse methodologies employed in distributional regression settings, we invite readers to consult Appendix C.T.
4.2
Matrix Information Theory
Quantum entropy, including Rényi entropy, is powerful metric to describe information in unique way (see Müileer-Lennert et al. [42 for foundational insights). Giraldo et al. [15] designed the measure of entropy using operators in RKHS to mimic Rényi entropy's behavior, offering the advantage of direct estimation from data. Bach [ applied von Neumann entropy of the density matrix to the covariance operator embedding of probability distributions, thereby defining an information-theoree framework utilizing kernel methods. In machine learning, especially within self-supervised learning (SSL) setups, entropy concepts have recently found novel applications. Our study builds most recent developments 449 21 55 by applying quantum Rényi entropy to the covariance operator in RKHS.
5 Experiments
We here demonstrate that our proposed method successfully performs unsupervised learning of data-dependent distribution kernel across different modalities. The experimental setup is divided into two phases: unsupervised pre-training and downstream regression classification using the learned kernel.
For each dataset, we select a hold-out validation subset with balanced classes, while the remainder of the dataset utilized for unsupervised pre-training. We use mini-batch ADAM [22] with a static learning rate of 0.0005. We report mini-batch based (instead of epoch based) training dynamics as our tasks do not require cycling over the entire dataset to converge to the optimal loss value. All experiments use Gaussian kernel both an embedding kernel and distribution kernel, the hyperparameter selection is performed as described in Appendix B.I
Once the samples encoder fe is learned, we employ it to compute distribution kernel Gram matrix, used as an input to the Support Vector Machine (SVM) for solving downstream classification tasks. grid search with 5 splits (70/30) is conducted to optimize the strength of the squared la regularization penalty C, exploring 50 values over the log-spaced range {10-? ..... 10 }. The best estimator then applied to evaluate classification accuracy on the validation subset, which we report.
8

---

Additional experiments exploring the application of data-dependent distribution kernels in domains where distribution regression models are less common, such as image and text, are presented in Appendix D
5.1
Flow Cytometry
Flow Cytometry (FC) is widely used technique for measuring chemical characteristies of mixed cell population. Because population-level properties are described through (randomized) sampling of cells, FC is used as canonical setup of distribution regression. For this study we used a dataset [56] where more than 100.000 cells are measured per each patient (sub- ject). For each cell a total of ten parameters are reported, hence, we treated each subject as an empirical distribution over R.10 We considered downstream classification tasks on two different sets of labels. The first ("Tissue' classification) contains peripheral blood (pB) and bone marrow (BM) samples from = 44 subjects. The second ("Leukemia" classification) presents healthy and leukemia BM cell samples, N 50. Classes were balanced in both cases. We sampled 16 subjects for Tis- sue classification and 20 subjects for Leukemia for training. Unsupervised learning was performed over the en- tire dataset. The encoder fo was parametrized by a 2-layers neural net- work (NN) with ReLU nonlinearity and l2 normalized output (on the unit hypersphere 89). Per each subject we sampled small percentage of cells, and we report performance for the sam- ple size of 200. We repeated each train- ing and testing phase for 100 times to track the variance induced by this ag- gressive subsampling.
Table 1: Distribution regression accuracy on Flow Cytome- try datasets.
MODEL
TISSUE
LEUKEMIA
ACC.
VAR.
ACC.
VAR.
GMM-FV
93.07%
10.308
94.80%
10.18
SWI
87.10%
+0.530
95.07%
10.111
SW2
81.71%
+0.341
95.30 %
+0.224
MMD LINEAR
82.42%
10.840
90.57%
10.208
MMD GAUSSIAN
81.71%
10.574
92.23%
10.216
MMD CAUCHY
81.57%
10.662
93.77%
10.080
MMD IMQ
82.89%
#0.698
91.43%
-0.217
DISTR. VAR.
79.47%
10.011
91.82%
10.007
MDKE RAND
77.50%
#.002
89.50%
+.003
MDKE No REG.
95.30%
#0.002
92.46%
00.02
MDKE REG.
98.89%
10.010
94.57%
10.005
To demonstrate the impact of unsuper- vised pre-training, we compared sev- eral methods across multiple configu-
rations [Ta
a) Kernels on distributions. This group includes Fisher kernels applied to parametric Gaussian Mixture Model (GMM) estimates, as suggested in Krapac et al. [24]. along with Sliced Wasserstein-1 and Sliced Wasserstein-2 kernels Br
b) MMD kernels. Here, we employ Gaussian embedding kernel for mean embeddings, marked as 'MMD' (Maximum Mean Discrepancy) in the table. This category includes various options for the distribution kernel, such as linear, Gaussian, Cauchy, and inverse multiquadrics.
c) Distributional Variance. An ablation study is conducted to demonstrate the results of directly optimizing the distributional variance defined in Eq. (15).
d) MDKE, We explore various configurations of the encoder optimized with the MDKE objective. We report performance for randomly initialized encoder, and for unsupervised pre-trained encoder with and without regularization. Random initialization happens only once, and all subsequent accuracy measurements are taken using the same encoder.
The variance reported for each model is measured across multiple runs to demonstrate the effect of the sampling. Importantly, the optimization of the MDKE objective results in embedding kernels with significantly lower variance compared to con-data-specifee kernels.
5.2 Image and Text Modalities
In the following section, we present additional experiments on learning data- data-dependent distribution kernels in domains not typically considered distributional regression tasks, specifically the image and text domains. While we acknowledge the existence of more powerful domain-specifff models and methods for both modalities, we provide these results to demonstrate the framework's applicability across a wide range of settings under appropriate choice of the representation model. Representing text empirical samples from the finite space of tokens (i.e. words from the dictionary) is quite

---

common, while the choice to model images as histograms over pixel positions is more subtle. We demonstrate that, in both scenarios, unsupervised pre-training of the encoder yields distribution kernel that achieves strong performance on downstream classification tasks, showcasing the versatility of the proposed learning framework in scenarios where distribution regression formulations are uncommon.
and Yashion-MNIST 665 consist of 28 x 28 pixel grayscale images divided into 10 classes. We considered each individual image to be probability distribution (via rescaling pixel intensities so that / (image) 1) over the discrete space of pixel positions (i.e., histogram). Given this support space, the encoder So is a discrete map which we implemented as a table lookup (i.e., embeddings) from pixel indices to the points on the hypersphere 8 Embeddings were initialized by sampling points uniformly. Gradients of the MDKE objective with respect to the embeddings parameters were computed via automatic differentiation using projected gradient steps to ensure that the embeddings remain on the hypersphere. The small size of the support space enables the exact computation of the inner product between kernel mean embeddings of input distributions Eq. 7 and, subsequently, the distribution kernel Eq. (8) during both training and evaluation. This ensured lower variance of the accuracy for the downstream classification. Performing MNIST classification upon pre-trate veetrrannng with our unsupervised encoder significantly improves the baseline (random initialization of latent embeddings) accuracy of 85.0% by reaching plateau at 92.15%. Refer to the detailed analysis of the latent spaces of the trained encoders in Appendix D.I
Images. MNIST 12]
Text. To assess our method's performance in larger discrete support space, we utilized the "20 Newsgroups" [27], a multi-class text classification dataset. We reduced the size of the dataset to 5 classes (resulting in 2, 628 sentences and 38, 969 unique words) by subsampling both training and test subsets. We treated sentences as empirical distributions over words, assuming word sets to be enough for topic classification, despite no positional info. The encoder fo mirrored the setup used in the MNIST case (Appendix D.I. l. with l2 normalized word embeddings on g01, However, while in the MNIST case embeddings computations were performed exactly. here considering the entire embedding kernel Gram matrix is impractical due to its large size. Instead, we optimized embeddings by randomly sampling 20 words per sentence, making the inner product between embeddings a stochastic approximation. This setup is meant to confirm that the optimization of the proposed MDKE objective yields a solution that is robust w.r.t. the excessive risk induced by first-level subsampling. Detailed results can be found in Appendix [.2]
6 Conclusion
In this work, we presented an unsupervised way of learning data-dependent distribution kernel. While previous studies in distribution regression predominantly relied on hand-crafted kernels, our work, in contrast, demonstrates that entropy maximization can serve as a powerful guiding principle for learning adaptable, data-dependent kernel in the space of distributions. Our empirical findings show that this technique can not only serve as a pre-training step to enhance the performance of downstream distribution regression tasks, but also facilitate complex analyses of the input space. The interpretation of the learning dynamics induced by the proposed objective relies on a theoretical link between the quantum entropy of the dataset embedding and distributional variance. This theoretical link, which we have proven, enables us to approach the optimization from a geometrical perspective, providing crucial insights into the flexibility of the learned latent space encoding.
We hope that theoretically grounded way of learning data-dependent kernel for distribution regression tasks will become strong alternative to the common practice of hand-picking kernels. More broadly, our results present a methodology for leveraging the distributional nature of input data along side the novel perspective on the encoding of complex input spaces. This highlights the potential to extend the application of more advanced learning methods, embracing the ever-increasing complexity of data by going beyond more conventional vector-based representations.
We thank the Allen Institute for Brain Science founder, Paul G. Allen, for his vision, encouragement, and support.
Acknowledgments.
10

---

References
[1] Bach. Information theory with kernel methods. IEEE Transactions on Information Theory, 69(2):752-775, 2022.
[2] F. Bachoc, F. Gamboa, J.-M. Loubes, and N. Venet. A gaussian process regression model for distribution inputs. IEEE Transactions on Information Theory, 64(10))6620-6633 6637, 2017.
[3] R. Berman, S. Boucksom, and D. W. Nyström. Fekete points and convergence towards equilib- rium measures on complex manifolds. 2011.
[4] G. Blanchard, G. Lee, and Scott. Generalizing from several related classification tasks to new unlabeled sample. Advances in neural information processing systems, 24, 2011.
[5] O. Boiman, E. Shechtman, and M. Irani. In defense of nearest-neighte based image classifica- tion. In 2008 IEEE conference computer vision and pattern recognition, pages 1-8. IEEE, 2008.
[6] C. Bonet, B. Malézieux, Rakotomamoniy, Drumetz, T. Moreau. M. Kowalski, and N. Courty. Sliced waaserreesnn on symmetric positive definite matrices for m/ceg signals. In International Conference on Machine Learning, pages 2777-2805. PMLR, 2023.
[7] S. V. Borodachov, D. P. Hardin, and E. B. Saff. Discrete energy rectifiable sets. Springer,
2019.
[8] C. Boutsidis, M. Mahoney, and Drineas. An improved approximation algorithm for the column subset selection problem. In Proceedings of the twentieth annual ACM-SIAM symposium on Discrete algorithms, pages 968-9777 SIAM, 2009.
[9] E. Carlen. Trace inequalities and quantum entropy: introductory course. Entropy and the quantum, 529:73-140, 2010.
[10] B. Charlier, Feydy, A. Glaunes, F.-D Collin, and G. Durif. Kernel operations on the gpu, with autodiff, without memory overflows. Journal of Machine Learning Research, 22(74): 2021.
[11] A. Christmann and I. Steinwart. Universal kernels non-standard input spaces, Advances in neural information processing systems, 23, 2010.
[12] L.. Deng. The mnist database of handwritten digit images for machine learning research. IEEE Signal Processing Magazine, 29(6):141-142, 2012.
[13] S. Flaxman, D. Sutherland, Y.-. Wang, and Y. W. Teh. Understanding the 2016 us presidential election using ecological inference and distribution regression with census microdata. arXiv preprint arXiv:1611.03775 2016.
[14] B. Ghojogh. A. Ghodsi, Karray, and M. Crowley. Reproducing kernel hilbert space, mercer's s theorem, eigenfunctions, nyström method, and use of kernels in machine learning: Tutorial and survey. arXiv preprint arXiv:2106.0883
2021.
[15] L. G. S. Giraldo, M. Rao, and C. Principe. Measures of entropy from data using infinitely divisible kernels. IEEE Transactions on Information Theory, 61()):335-548 2014.
[16] K. Grauman and T. Darrell. The pyramid match kerne:: Efficient learning with sets of features.
Journal of Machine Learning Research, 8(4), 2007.
[17] T. Jaakkola and D. Haussler. Exploiting generative models in discriminative classifiers. Advances neural information processing systems, 11, 1998.
[18] T. Jebara and R. Kondor. Bhattacharyya and expected likelihood kernels. In Learning Theory and Kernel Machines: 16th Annual Conference on Learning Theory and 7th Kernel Workshop, COI.T/Kernel 2003, Washington, DCC USA, August 24-27, 2003. Proceedings, pages 57-71. Springer, 2003.
[19] Jebara, Kondor, and Howard. Probability product kernels. The Journal of Machine Learning Research, 5:819-844, 2004.
[20] T. Karvonen, S. Särkkä, and K. Tanaka. Kernel-based interpolation at approximate fekete points. Numerical Algorithms, 87:445-468, 2021.
[21] J. Kim, S. Kang, D. Hwang, J. Shin, and W. Rhee. Vne: An effective method for improving deep representation by manipulating eigenvalue distribution. Proceedings of the IEEF/CVF Conference Computer Vision and Pattern Recognition, pages 3799-3810 3810, 2023.

---

[22] D. P. Kingma and J. Ba. Adam: A method for stochastic optimization, 2017.
[23] S. Kolouri, Y. Zou, and G. K. Rohde. Sliced wasserstein kernels for probability distributions. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, pages 5258-5267, 2016.
[24] Krapac, J. Verbeek, and F. Jurie. Modeling spatial layout with fisher vectors for image categorization. In 2011 International Conference on Computer Vision, pages 1487-1494, IEEE, 2011.
[25] G. Krummenacher, C. S. Ong, and Buhmann. Ellipsoidal multiple instance learning. In International Conference Machine Learning, pages 73-81. PMLR, 2013.
[26] H. Kuck and N. de Freitas. Learning about individuals from group statistics, 2012. URL mttps:://arae http//oww.voo////////////oo wwwooo/g//aa. abs/1b07 1393
[27] K. Lang. Newsweeder: Learning filter netnews. In Proceedings of the Twelfth International Conference on Machine Learning, pages 331-339, 1995.
[28] H. C. L. Law, D. J. Sutherland, D. Sejdinovic, and S. Flaxman. Bayesian approaches to distribution regression. In International Conference Artificial Intelligence and Statistics, pages 1167-1176. PMLR, 2018.
[29] M. Lemercier, C. Salvi, T. Damoulas, E. Bonilla, and Lyons. Distribution regression for sequential data. In International Conference Artificial Intelligence and Statistics, pages 3754-3762. PMLR, 2021.
[30] Leung and J. Malik. Representing and recognizing the visual appearance of materials using phree-dimensions textons. International journal of computer vision, 43:29 44, 2001.
[31] F. Liu, X. Huang, Y. Chen, and J. A. Suykens. Random features for kernel approximation: A survey on algorithms, theory, and beyond. IEEE Transactions on Pattern Analysis and Machine Intelligence, 44(10):7122 1000::1188777, 2021.
[32] W. Liu, L. Yu, A. Weller, and B. Schölkopf. Generalizing and decoupling neural collapse via hyperspherical uniformity gap, 2023.
[33] D. Lopez-Paz, K. Muandet, B. Schölkopf, and I. Tolstikhin. Towards learning theory of cause-effecct inference. In International Conference on Machine Learning, pages 1452-1461.
PMLR, 2015.
[34] S. Lyu. Mercer kernels for object recognition with local features. In 2005 IEEE Computer Society Conference on Computer Vision and Pattern Recognition (CVPR '05), volume 2, pages 223-229. IEEE, 2005.
[35] J. Marzo and J. Ortega-Cerda. Equidistribution of fekete points on the sphere. Constructive Approximation, 32(3):513-521, 2010.
[36] G. Meanti, L. Carratino, L. Rosasco, and A. Rudi. Kernel methods through the roof: handling billions of points efficiently. Advances in Neural Information Processing Systems, 13:14410- 14422, 2020.
[37] D. Meunier, M. Pontil, and C. Ciliberto. Distribution regression with sliced wasserstein kernels. In International Conference on Machine Learning, pages 15501-15523. PMLR, 2022.
[38] Mitrovic, D. Sejdinovic, and Y.-W. Teh. Dr-abe: Approximate bayesian computation with kernel-based distribution regression. In International Conference Machine Learning, pages 1482-1491, PMLR, 2016.
[39] K. Muandet, K. Fukumizu, F. Dinuzzo, and B. Schölkopf. Learning from distributions via support measure machines. Advances in neural information processing systems, 25, 2012.
[40] K. Muandet, D. Balduzzi, and B. Schölkopf. Domain generalization via invariant feature representation. In International conference on machine learning, pages 10-18. PMLR, 2013.
[41] K. Muandet, K. Fukumizu, B. Sriperumbudur, and B. Schölkopf. Kernel mean embedding of distributions: review and beyond. Foundations and Trends@ in Machine Learning, 10(1-2): 1-141, 2017. doi: 10. 10.1561/220000066 URL URL/https: h/tp://oo..oooo orgoo.oorggiii 10 1561/2222000000060
[42] M. Müller-Lennert F. Dupuis, O. Szehr, S. Fehr, and M. Tomamichel. On quantum rényi entropies: A new generalization and some properties. Journal of Mathematical Physics, 54(12), 2013.
12

---

[43] M. Ntampaka, H. Trac, J. Sutherland, S. Fromenteau, B. Póczos, and J. Schneider. Dynamical mass measurements of contaminated galaxy clusters using machine learning. The Astrophysical Journal, 831 (3)(2), 2016.
[44] J. Oliva, B. Póczos, and J. Schneider. Distribution to distribution regression. In International Conference on Machine Learning, pages 1049-1057. PMLR, 2013.
[45] B. Poczos, A. Rinaldo, A. Singh, and L. Wasserman, Distribution-free distribution regression, 2013.
[46] A. Rakotomamonjy, A. Traore, M. Berar, R. Flamary, and N. Courty. Distance measure machines. arXiv preprint orXiv:/8003 002.0, 2018.
[47] A. Rudi, R. Camoriano, and L. Rosasco. Less is more: Nyström computational regularization. Advances in neural information processing systems, 28, 2015.
[48] C.-J. Simon-Gabrie and B. Schölkopf. Kernel distribution embeddings: Universal kernels, characteristic kernels and kernel metrics on distributions. The Journal of Machine Learning Research, 19(1):1708-177)
2018.
[49] O. Skean, A. Dhakal, N. Jacobs, and L. G. S. Giraldo. Frossl: Frobenius norm minimization for self-supervised learning. arXiv preprint arXiv:2310.0290: 2023.
[50] B. K. Sriperumbudur, K. Fukumizu, and G. R. Lanckriet. Universality, characteristic kernels and rkhs embedding of measures. Journal of Machine Learning Research, 12(7), 2011.
[51] B. K. Sriperumbudur, K. Fukumizu, A. Gretton, B. Schölkopf, and G. R. Lanckriet. On the
empirical estimation of integral probability metrics. 2012.
[52] D. J. Sutherland, L. Xiong, B. Póczos, and J. Schneider. Kernels sample sets via nonpara- 2012.
metric divergence estimates. arXiv preprint
[53] Z. Szabó, A. Gretton, B. Póczos, and B. Sriperumbudur. Two-stage sampled learning theory on distributions. In Artificial Intelligence and Statistics, pages 948-957. PMLR, 2015.
[54] Z. Szabó, B. K. Sriperumbudur, B. Póczos, and A. Gretton. Learning theory for distribution regression. The Journal of Machine Learning Research, 17(1):5272-5311, 2016.
[55] Z. Tan, J. Yang. W. Huang. Y. Yuan, and Y. Zhang. Information flow in self-supervised learning. arXiv preprint arXiv:2309. 17281, 2023.
[56] M. C. Thrun, J. Hoffmann, M. Röhnert, M. von Bonin, U. Oelschlägel, C. Brendel, and A. Ultsch. Flow cytometry datasets consisting of peripheral blood and bone marrow samples for the evaluation of explainable artificial intelligence methods. Data in Brief, 43:108382, 2022.
[57] J.-F. Ton, C. Lucian, Y. W. Teh, and D. Sejdinovic. Noise contrastive meta-learning for conditional density estimation using kernel mean embeddings. In International Conference on Artificial Intelligence and Statistics, pages 1099-1107. PMLR, 2021.
[58] J.-F. Ton, D. Sejdinovic, and K. Fukumizu. Meta learning for causal direction. In Proceedings of the AAAI conference on artificial intelligence, volume 35, pages 9897-9905, 2021.
[59] T. Tuytelaars, M. Fritz, K. Saenko, and T. Darrell. The nbnn kernel. In 2011 International Conference Computer Vision, pages 1824-1831. IEEE, 2011.
[60] N. Vakhania, V. Tarieladze, and S. Chobanyan. Probability distributions on Banach spaces, volume 14. Springer Science & Business Media, 2012.
[61] Wang and P. Isola. Understanding contrastive representation learning through alignment and uniformity on the hypersphere, 2022.
[62] K. Q. Weinberger and G. Tesauro. Metric learning for kernel regression. In Artificiaa intelligence and statistics, pages 612-619. PMLR, 2007.
[63] M. M. Wilde. Quantum information theory. Cambridge university press, 2013.
[64] C. Williams and M. Seeger. Using the nyström method to speed up kernel machines. Advances in neural information processing systems, 13, 2000.
[65] H. Xiao, K. Rasul, and R. Vollgraf. Fashion-mmin mnist: novel image dataset for benchmarking
machine learning algorithms, 2017.
[66] I.. Xu, Y. Chen, S. Srinivasan, N. de Freitas, A. Doucet, and A. Gretton. Learning deep features in instrumental variable regression. arXiv preprint arXiv:2010.07154, 2020.
13

---

[67] L. Xu. H. Kanagawa, and A. Gretton. Deep proxy causal learning and its application to confounded bandit policy evaluation. Advances Neural Information Processing Systems, 34: 16264-26275 2021.
[68] Y. Yoshikawa, T. Iwata, and H. Sawada. Latent support measure machines for bag-of-words data classification. Advances in neural information processing systems, 27, 2014.
[69] Z. Yu, D. W. Ho. Z. Shi, and D.-X. Zhou. Robust kemnel-base distribution regression. Inverse Problems, 37110),105014 2021.
[70] Z.-H. Zhou, Y.-Y. Sun, and Y.-F. Li. Multi-instance learning by treating instances as non-i.i.d. samples, 2009. URL URI.https: //arxiv..ivv/////////i org/abs/0807 1997
14

---

A Proofs
In this section, present the proofs for the propositions outlined in study. To ensure clarity, will first restate the setup and introduce necessary concepts.
We define the input space as X. and P(2') represents the space of probability distributions over X. Consider dataset of M probability distributions, denoted as DM = {P, E P( With a p.d. characteristic embedding kernel k X X - I., the corresponding RKHS H. and the feature map á : + H. we define the mean embedding map μ. P(X) + H such that μip = p(P) = Lx ((x) (x). A p.d. maaslation-invarian characteristic distribution kernel K : P(()) P(X) + R is defined using the mean embeddings of corresponding distributions. For simplicity, b4 denotes ((P.).
Additional concepts essential for our proofs include the Gram matrix of mean embeddings G RMX M representing the inner products of mean embeddings in the dataset, i.e., G == (1) Hli--
The kernel matrix Kp I M x M with respect to the distribution kernel K is denoted as Kp = [K,,, Pj)lij
We also recall the definition of distributional variance V H (see Eq. (15):
M
(H(DM) = M [r[
-
ΣE Gij
M i=1
These definitions and notations will be referenced throughout the proofs.
A.1 Kernel Norms Gap and Distributional Variance
For both cases of input distributions being empirical probability distributions or continuous densities, define mixture distribution P, with slight abuse of notation:
P(x) = M 1=1
Σ P((x)
Lemma A.1. For the mixture distribution P and the Gram matrix G, the following relationship holds:
M
|
Gij
(19)
M
Proof. We begin by recalling that the inner product between mean embeddings 14, and in H is given by:
(e. = -xxx M(x,')dP(())
Substituting this into the expression for the Gram matrix, we have:
M
M
1 M M
JJ
ΣΣ<n
Gy
J k(2,z') (2,z
M
M
1 M
|
1 M
X(I,x')d
-
MEP(())
P((x)) d
M
M
1=1
j=1
-J
k(I,x') dP(x)dP(a') dp((')) = ||m l
This completes the proof.
By incorporating Eq. (19) into the definition of distributional variance (see Eq. 15)), and noting that
the trace of G is the sum of squared norms of input distributions (i.e., tr(G) -= [|u,, 1 we obtain:
M
1
Vx(DM)
Σ
11
M
15

---

This result offers new and intuitive perspective on distributional variance, conceptualizing it as the difference between the average squared norm of individual input distributions and the squared norm of the mixture distribution.
A.2 Pairwise Distance and Distributional Variance
Definition A.2. For dataset D of M probability distributions, we define the average pairwise distance between kernel mean embeddings in H as Iu (Dx)), given by:
M ΣΣ
1
(D) =
l14 -1yl
(20)
M
Lemma A.3. For the distributional variance V of the dataset DM, the following relationship holds:
Vn(DM) 2 (D()
(21)
Proof. Starting with the definition of Uu (DM):
M
M
1
(D(
llms - 13l134
M
i=1
j=1
M
1
(p) - M4. pyyn
M
M
M
1
(1 + (1g,10)
2(14 Uy)n)
M
j=1
M
M
1
2M
(ps p3y
M
2
(pin pi)
(par pgy
M
M
diagonal G
Recognizing that M (Mi p4yH tr[G], we can equate expression in the brackets to (D(M leading to:
Vx(DM) = Jx((M).
which concludes the proof.
Lemma A.3 will be instrumental in further proofs, establishing a crucial link between distributional variance in 7l and quantum entropy of the covariance operator embedding Σp.
A.3 Distribution Kernel Entropy Upper-Bound
In this section, we provide the proof for the key theoretical result stated in Proposition 3.3
Consider a dataset D consisting probability distributions {P P(X)); sampled i.i.d. from an unknown meta- meta-distribution D. We assert that the second-ordes Rényi entropy S2 of the empirical covariance operator embedding Ep, induced by the choice Gaussian distribution kernel KREF over points in the RKHS H, is upper-bounded by the distributional variance V (D)):
- S (H(D)
27
16

---

where y is the bandwidth of the distribution kernel KRBF-
Starting from the properties of S2(ED) stated in Eq. 13
M
M
2
M
M
S2(En)= =
kD) M
log
Kiz)
=-log
K3
(22)
M
Applying Jensen's inequality, considering the concavity of the log, to Eq.
22 we have:
M M log K
Sr(Eo)
M
j=1
M
M
log exp (-2 211m - pyll
(23)
M
j=1
M
log exp (->llęss Mey m3)122
M
M M
l1
M
i= j=1
7 (D)
(24)
2y (D)
(25)
Eq. (33
the definition of KRBF from Eq. Eq. (24) is derived from the definition of the
average pairwise distance in Eq. (20). and Eq. 25 follows from Lemma A..3 This completes the
proof of Proposition [33 establishing the upper bound for the second-order quantum Rényi entropy
of the covariance operator embedding En in terms of the distributional variance in RKHS H. [
Remark A.4. It is important to note that S₂ is typically measured using the base-2 logarithm (logg) rather than the natural logarithm. However, the proof remains accurate with a proper re-scaling to account for the change in the logarithmic base.
A.4 Generalized Variance in RKHS
In this section we prove the connection between generalized variance and norm of the kernel mean embedding.
So far established the squared norm of mean embedding is maximized by Dirac points (zero variance) and is minimized by uniform distribution (max variance). We now formalize the intuition that larger squared norm in embedding kernel RHKS corresponds smaller variance in the latent space. We first note that under Assumption 33. the kernel metric dy... (see Eq. (3) induced by the choice of kemb is monotonically increasing function of Fuclidean metric the latent space, thus is representative of the geometry of encoded input distributions. Which let us use a generalized notion of variance here.
Definition A.5. (Vakhania al. (60) Let X be random variable which takes values in a Fréchet space F equipped with seminorm - [|.. And suppose that is gquare-integratable a sense that ||| < co. Let F be a Pettis integral of X (i.e. generalization of the mean). Generalized variance of X w.r.t seminorm || lla is defined as following
Var [X] == E 1ll3.
(26)
Note that for a random variable in RKHS defined as push-forward of a probability distribution P E P(Z) over the latent space, i.e. X = 2 Φ# satisfies conditions of Definition A.5 with F H. ||. la 1-|l and Pettis integral being kernel mean embedding. Denote the described variance as Var,... (P]. We now show that
Proposition A.6. Under Assumption
3.1 for every P P(Z),
Var [P]
|| (P) Hmme
(27)
17

---

Proof.
From the definition of generalized variance we have the following
Ham [P] = Eezn Sttmpllof pl(((2) - pp 12
{lldte - 1p dP(x))
/ (0(x) (r,((x) Hr)) Ham dP(x)
/ (11())))AA..... + 11 2(0(1x); (1)) Hmms) dP(x)
(28)
/ [6(()))) dp(x)+ / ||upll2ku dP(x) -2 / (o(x),p)) dP(x)
= / 1 ddP(x) + |1lp|1bmm. / [IdP(a) - (f (())))
Hanh
1 + || 11 21|1111 11
which concludes the proof.
The proposition demonstrates the motivation behind training the encoder to minimize the variance in the latent space by maximizing the (average) squared norm of the mean embeddings.
B Practical Aspects of Learning
B.1 Kernel Hyperparameter Selection
In practice, the effectiveness of the optimization process significantly influenced by the parameters of the Gaussian kernels. Setting the bandwidth parameter 7 either too low too high can hinder the model's ability to learn effectively. This issuc is a well-known challenge in working with kernel methods, where determining the optimal kernel bandwidth has been an area of extensive study.
In our work, we have employed an empirical approach, which involves adjusting 7 based on the idealized structure of the dataset after its projection onto the latent space. Specifically, we set y such that 1/4 equals 10 times the average distance the nearest neighbor in the set of points sampled uniformly gd-1 (inspired by the experimental approach in Blanchard al. 4 The number of points is chosen to be the number of distributions in the training set for the distribution kernel. For the embedding kernel, it to the number of distributions in a batch multiplied by the number of samples per distribution used for unsupervised training.
B.2 Optional Regularization
matrix the As the theoretical maximum of S. entropy is attained when the spectrum of M KD is uniform, the optimal encoder with respect to LMDKE also tends to maximize the determinant of the kernel Gram matrix Kp. This relationship is intuitive when considering that the determinant product of its eigenvalues. The points that maximize such determinant Ko are theoretically known as Fekete points [35] 3 and, in our case, are relative to the distribution kernel. Fekete points have been shown by Karvonen al. [20] to be an optimal configuration for learning kernel interpolants, making them particularly suitable for downstream tasks framed as kernel regression. As such, encoders optimized under the LMDKE objective facilitate more accurate and robust performance in subsequent regression tasks.
In practice, we found that optimizing the MDKE objective posed certain numerical challenges, particularly due the tendency of too small eigenvalues in the distribution kernel matrix Ki to collapse near the optimal value of the objective when using large batch size for training. To mitigate this issue and prevent undesirable optimization behavior, we have introduced regularization term to the original objective. This term is inspired by the concept of Fekete points configuration, leading to
18

---

the following loss function:
LMKKE R0)
->2[ED) +c. G.((ED))
1 KD
(29)
=log Kollz c. c-log det M
M
Here, serves as a hyperparameter to control the strength of regularization. The regularization term ((pp) is designed to stabilize the optimization process by counteracting the effects of the collapse of small eigenvalues. The empirical evidence supporting the effectiveness of regularization in stabilizing the training dynamies for the sentence representation learning experiment (as detailed in Appendix D.I is showcased in Fig.
B.3 Runtime Complexity
Scalability of the kernel methods is typically key concern when it comes practical applications. The runtime complexity the proposed method could be decomposed into two components:
Computation of the distribution kernel Gram matrix. This involves computing O(N2) inner products between distributions where is a number of distributions. Each inner product involves computing matrix of pairwise distances between samples from each distribution, in the case of computing the Gaussian kernel over points on the hypersphere, the runtime complexity is the complexity of multiplying matrices M where M is number of samples. All computations could be efficiently parallelized.
Solving regression using distribution kernel Gram matrix. The complexity comes from matrix inverse and is typically estimated to be O(N3) with a range of methods proposed to reduce the complexity 8
With this being said, we want to highlight that the runtime complexity of the proposed method is most favorable exactly for those tasks and datasets where distributional regression is an appropriate modeling approach. A significant number of samples per distribution ensures a high accuracy in approximating the kernel between a pair of distributions, and, at the same time, relatively small number of distributions in the dataset alleviates issues related to storing the distribution kernel Gram matrix in memory and performing computations on
Eigenvalues of Gram Matrix (No Reg.)
Eigenvalues Gram Matrix (Reg.)
Step 10k
Step 10k
Step
-10
Step
-10
Step
In
40
-40
500
1000
15000
2000
500
1000
1500
2000
The effect or regularization on the training dynamics. The distribution of the eigen- values of the distribution kernel Gram matrix, calculated for 2,000 sentences sampled from '20 Newsgroups' dataset (details in Appendix D.2p. no regularization leads to the collapse of smaller eigenvalues. (b) The regularization stabilizes the training by preventing eigenvalues from collapsing.
Figure 3:
(.22). is observed throughout the training. (a) Training with
19

---

C Connection to Other Frameworks
C.1
Distributional Regression Landscape
As the distributional regression task differs from common Machine Learning (ML) setups where inputs are given as vectors, the effective practical solution requires unique considerations.
The most obvious approach would be to ignore the fact that inputs are given as distributions and learn classifier on the space of samples, with a proper aggregation of posteriors (e.g. with a simple sum over histograms). This approach, while being simple, has been shown not to yield practically useful results, and was explicitly excluded from reported performance on different tasks by different authors 39 46].
Viable approaches to solving distributional regression could be, approximately, split up into the following categories.
Discriminate generative models. The idea is to fit each input distribution to parametric family, e.g. Gaussian Mixture Model (GMM), to use available closed- fosed-form solutions to compute kernel or similarity or distance between distributions. This group of methods originated from the work on hidden Markov models and it's applications to processing sequence modalities, like text, DNA, proteins, and more. Early work [7] on learning discriminative classifier for generative models leveraged the fact that parametric models forms a Riemannian manifold with local metric given by Fisher information, they derived kernel function termed Fisher kernel suitable for running SVM between generative models. Driving motivation was classification between hidden Markov models, with the developed method being applied to DNA and protein sequence analysis. Following the same modeling approach of analysing DNA sequences with discriminative models between Markov models, Jebara and Kondor [18] proposed to use Bhattacharyya distance between distributions from exponential families to derived so-called Expected Likelihood Kernel. Jebara et al. [19] explored the method of computing kernel between distributions as the integral of the product of pairs of distributions, termed Probability Product Kernel. A new family of kernels was applied to the same setup of discriminating task on text modality, hidden Markov models for biological data. The model was also successfully applied for analysis of linear dynamical systems for time series data. Critical advantage of this work was access to computationally effective way of computing kernel for the distributions without having access to analytical closed-form parametrization (only relying on samples). In Sec. [. we used fitting distributions to GMM with Fisher Kernel applied to learn parameters baseline for the performance the task.
Point clouds. This group includes methods that model each input distribution as a set of points (also known as 'point cloud' 'feature group' in computer vision) and use kernel or similarity function defined on sets. Large portion of the methods in this group arised in computer vision (CV) field when local features extractor were widely employed to pre-process images yielding either per-image histograms or sets of low dimensional vectors. The group includes kernels based on nonparametric divergence estimates, quantized set kernels, and so-called "nearest-neighhor" set kernels to name a few [52] 30 34
5 5 (5). Such kernels were employed to many CV tasks though the successfully application required not only a good kernel but also a high-quality feature extractor. Special attention goes to methods leveraging Wasserstein distance (including both kernel-based and similarity based solutions). Wasserstein distance as a metric-aware discrepancy measure for probability distributions is a natural choice of deriving kernels similarity functions between point clouds (46) While being computationally problematic for large-scale problems, Sliced Wasserstein kernels were succesfully adopted [23 37] they provided both reliable way for the point set comparison with reduced cost leveraging sliced formulation. In Sec. S... we compared performance of Sliced Wasserstein- and Sliced Wasserstein-2 kernels, with the latter yielding significantly lower performance. Such a behavior is consistent with theoretical analysis stating that Sliced Wasserstein-1 1 has favoriable properties when compared kernel based on Wasserstein-2 distance.
16
Kernel Mean Embeddidngs. Muandet et al. 39 proposed to leverage kernel mean embedding of measures in RKHS, so that the distributional regression could now be casted to a regression the corresponding Hilbert space. While the original work leveraged the kernel between the RKHS embeddings to train SVM, Szabó et [53] provided theoretical guarantees for the for learning Ridge regressor from distribution embeddings in the Hilbert space to the outputs. Law et al. (28) applied
20

---

the same approach in Bayesian regression settings, and [2] used it for Gaussian Process Regression. The setup gives a great deal of flexibility by choosing the kernel in RKHS, with a few being tested in practice. Gaussian kernel is common choice, due to its universality [II] Among others, inverse multiquadrie (IMQ) a popular choice, typically paired with random features to improve the runtime complexity of the algorithm. Linear kernel in RKHS, despite not being universal, was reported to produce competitive results in multiple practical setups. In Sec. 5. we compared performance for different kernels in RKHS, namely linear, Gaussian, Cauchy, and IMQ, all defined as functions of Maximum Mean Discrepancy (MMD), which gives as a way to compute the value of the kernel from finite number of samples with high accuracy. While all kernels demonstrated different performance in terms of the test accuracy, the difference reported is not that substantial.
Distribution kernels were successfully applied in different kernel-based methods and in different modalities, like distribution to distribution regression [4], distribution regression for sequential data [29]. and more. In cases where input data is naturally represented as a probability distribution but not immediately applicable for existing distribution regression solutions, pre-processing or encoding of inputs is required. Yoshikawa et al. [68] proposed learning latent space representation for input data to apply distribution regression in the text modality, with the encoder being trained jointly with the classifier. To our best knowledge, methods for learning data-dependent kernel for solving distributional regressions were not previously reported.
C.2 Hyperspherical Uniformity Gap
The link between entropy maximization in the space of distribution, the gap between average norm and norm of the mixture distribution marginalized over the dataset (i.e. 'average' distribution), and properties of distributions that minimize and maximize kernel mean embedding norm (as described in Sec. 3.4 unveils a subtle yet significant connection to the concept of the Hyperspherical Uniformity Gap (HUG), introduced in (32). HUG has been developed to generalize the phenomenon of neural collapse observed in supervised classification settings. In our approach, working directly with samples from input distributions grants explicit access the 'grouping' of points in the input space. Even though the HUG framework setup does not have notion of grouping', we can close the gap by noting that class labels provided with the dataset implicitly create 'groupings' of points, which can be interpreted as empirical samples drawn from latent probability distributions (one for each class).
However, a notable distinction lies in the dimensional aspect: HUG focuses the distribution of points on the surface of a finite-dimensional hypersphere, whereas our work encompasses an infinite-dimension hyperball setting, given the of points in denduelinduced RKHSs. Furthermore, the loss function introduced in our study presents a unified optimization objective that weaves together both the inner-group and inter-group dynamics while being derived from first principles. In the HUG framework these two aspects are addressed separately.
Establishing more formal connection between these frameworks emerges promising direction for future research. Such an endeavor could offer a novel perspective on supervised classification, particularly by conceptualizing class prototypes as probability distributions rather than mere vectors. This exploration might bridge the gap between these distinct approaches, enriching our understanding of classification paradigms in high-dimensional spaces.
D Additional Experiments
D.1 Image Classification Tasks
Performing MNIST classification upon pre-training with our unsupervised encoder significantly improves the baseline (random initialization of latent embeddings) accuracy of 85.0% by reaching a plateau at 92.15%.
To understand the data deteadependency of our encoding procedure, we analyzed the latent spaces of MNIST and Fashion-MNIS datasets. The visualization of pexel-level interaction, computed using Gaussian kernel Gram matrix, reveals complex, dataset-specific interactions (Fig 4a 4cp. Spectral clustering using the kernel Gram matrix provides deeper insight into the pixel interaction landscape. The chart shows clusters of pixels with correlated intensities (Fig. 4d) with the number of clusters set to 10 empirically.
21

---

Pixel level interaction kemel
b
Pixel level clustering colurns
Flattened pixcl space
Pixels rows
MNIST
Flattened pixel spaice
d
Pixels columns
Flattened pixel space
Pixels rows
Fashon MNIST
Flattened pixel space
Figure 4: Unsupervised encoding of Images. Unsupervised learning of image embeddings as support distributions (i.e.. histograms) of pixel intensities. For every pixel position assign a point location on the unit hypersphere and optimize such locations via the covariance operator dataset embedding w.r.t. the MDKE objective. [a) Samples from the MNIST dataset and learned
pixel-to-pixel interaction kernel Gram matrix, (b) Spectral clustering of pixels based on the learned
kernel Gram matrix.
(c) and (d) same
and (b) for Fashion-MNIST dataset.
Learning metrics
b
Keme! Norms
1.0
AVG(nor)) Σ
nom((A/) bylk
Validation
Komal Enbeoding Norm
Karaal artoopy
Anturny
Wardeton encumecy
0.2-
0.8
0.1
Karnel 5,17,)
Distibutonal
2y (PPee 0.4
varlance
0.70
1300C
20000
30000
30000
20000
30000
Training stop
Training stop
Figure 5: Unsupervised encoding of Text. Unsupervised learning of sentences embeddings as empirical distributions of words on the '20 Newsgroup' dataset. Goodness the learned embeddings is evaluated by performing sentence-to-topic classification. (a) Distribution kernel entropy, distribu-
tional variance, and validation aceuracy throughout training. (b) Kernel norms Eq. (T6) throughout
training. Shaded blue area (the difference between the blue and red lines) corresponds to the blue
dotted line in panel (a) (up to multiplicative factor).
D.2 Text Classification Tasks
During the unsupervised pre-training phase, we observed steady decrease in training loss (Fig. 5a despite the small batch size (50 sentences with total of 1, 000 words). Importantly, for this experiment we employed a regularized version of the objective (MDKE-R). We discuss the rationale and empirical evidence supporting the use of this regularization scheme in Appendix B..2 During evaluation, to reduce computational complexity, we sample 2, 000 sentences for the tram and 1, 000 the test split, keeping the classes balanced. The maximum classification accuracy achieved was approximately 89.3%, while the random initialization performance averaged at 37.5% (Fig. 5) The framework's 's high accuracy in downstream classification showcases its prowess learning potent latent representations, when dealing with input distributions with large finite support.
22

---

E Implementation Details
In this section, provide example that illustrates the implementation of the proposed method using the PyTorch framework. All experiments were performed on single machine with 1 GPU and CPUs.
Functions to compute distribution kernel Gram matrix:
def pairwise_kernell (x, gammal):
B, T s.size (0), s. (1)
X_unroll x. r.reshape(B-7 (B+T, -1)
dist (X_unroll[; [;, None, :1 X_unroll (None, :1) **2
dist torch. (dist, dim=2)
F.
dist [None, :.
kernel_size-(1 (T, T),
stride-(T, T)
return corch.exp (- (gamma1/2.) G. G.quuezee()) (0))
def Eestribution,kernel.gge
gammal, gamma2):
Gxy mairwise_kerse (x, (ammal-ganmal)
Gy torch. diag (Gxy)
G Gx[:, None] + Gy [None, :] 2+Gxy
g torch. exp (- (g(amma22/22 * G)
return Gxy,
Distribution kernel entropy estimator and the MKDE loss:
def distribution wtwoooooeeeeeeeeeeeeeeeeee (K)
B size (0)
Cov (1/B) K
return (Cov 2) sum () log2
def mkde_lose Ealoss(encodees (encoder, x, gammal, gamma2):
X. shape is (a.distributions, n_samples, d_input) encoder (X)
# shape is (n_distributions. n_samples, d_latent)
K distribution_kertte kernel_gram (2, gammal, gamma2)
# K. shape is (n.distributions, d.distributions:
loas wdiatributiinnttcc.mm weroonneeeeeeeeeee entropy (K) return loss
The code presented here follows the setup presented in Sec. Sec. Jusing Gaussian kernel both as embedding and distribution kernel. Other kernels could be used by adjusting implementation of both helper functions accordingly. When training on large datasets, Charlier et al. [10 might be used to avoid memory overflow in average pooling.
23

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: the claim outlined in the introduction, and matches theoretical and experimental results presented in the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification: See Sec. 3.6 for details.
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full of assumptions and a complete (and correct) proof?
Answer: [Yes]
Justification: See Sec. 3.4 and Appendix A for details.
4. Experimental Result Reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that it affects the main claims and/or conclusions the paper (regardless whether the code and data are provided or not)?
Answer: [Yes]
Justification: See Sec. and Appendix Dor details.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
Justification: The code for the implementation of the proposed loss is provided in Ap- pendix
Experimental Setting/Details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: See Sec. and Appendix Dor details.
7. Experiment Statistical Significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [Yes]
Justification: We report the variance of accuracy measurements that unique data/- model setup, specifically the resampling variance. Detailed explanations are provided in
Sec. 5
8. Experiments Compute Resources
Question: For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time exccution) needed to reproduce the experiments?
Answer: [Yes]
24

---

Justification: See Appendix E for implementational details.
9. Code Of Ethies
Question: Does the research conducted the paper conform, in every respect, with the NeurlPS Code of Ethics httppp///eerrrp...cpopppo/hhccha
Answer: [Yes]
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [Yes]
Justification: See Sec. 3.6 for details.
Safeguards
Does the paper deseribe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
12. Licenses for existing assets
Are the creators or original owners of assets (e.g.. code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
Justification: Only publicly released datasets were used, with the citations provided for each experiment. See Sec. 5 and Appendix D for details.
13. New Assets
Question: Are new assets introduced in the paper well documented and the documentation provided alongside the assets?
Answer: [NA]
Justification: The paper does not release new assets.
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 details about compensation (if any)?
Answer: [NA]
Justification: The paper does not involve crowdsourcing nor research with human subjects. 15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based the requirements of your country institution) were obtained?
Answer: [NA]
Justification: the paper does not involve crowdsourcing nor research with human subjects.
25


Reviews:
[{"review_id": "G1mrw2AE5D", "reviewer": "Reviewer_1Qk4", "summary": "This paper proposes a new unsupervised way to learn kernels for distribution regression through entropy maximization. In addition, they also propose a geometric interpretation which is very interesting. The learnt kernels are general and have shown on some experimental settings to perform better than standard unlearnt kernels.", "questions": "see above", "limitations": "yes", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 2, "strengths": "- The paper proposes a novel way to learn kernels through entropy maximiazation in the setting of distribution regression.\n- The paper is very well written and easy to follow for someone who has worked the KME before.\n- The paper shows promising results compared to exisiting methods that do not train the kernel (through entropy)", "weaknesses": "- The experimental section does not compare to other kernel distribution regression methods. Could the authors confirm why [1] was omitted in the comparisons?\n- In addition citations to other methods that also learn kernel albeit not with entropy should be cited [2, 3, 4, 5]\n- How were the hyper parameters picked for the fixed kernel such as RBF kernel? Can you please lay out the whole process as well as the corresponding values for each of the hyper parameters?\n\n\n\n[1] Bayesian approaches to distribution regression\n[2] Learning Deep Features in Instrumental Variable Regression\n[3] Noise Contrastive Meta-Learning for Conditional Density Estimation using Kernel Mean Embeddings\n[4] Meta Learning for Causal Direction\n[5] Deep proxy causal learning and its application to confounded bandit policy evaluation", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Dear Reviewer,\n\nThank you very much for your thoughtful review.\n\n* We appreciate the insightful comments from the Reviewer. We will incorporate the Bayesian approaches they highlighted into our extended discussion in Appendix C. We have dedicated this appendix to a thorough review of related literature. The absence of a comparison with the work of [1] in our experimental section is due to the novel nature of our task, which is the unsupervised learning of distribution embeddings. To our knowledge, [1] and similar studies leverage labels in a supervised manner, which contrasts with our unsupervised approach. A direct comparison would necessitate significant modifications to these methods, constituting a separate contribution.\n* In response to the references suggested by the reviewer, we will enrich our manuscript with additional citations and a succinct discussion of [2,3,4,5].\n* We had delineated our hyperparameter selection process in Section B.1 in detail. To enhance clarity, we will insert a second cross-reference to this section. This procedure builds upon and elaborates the methodology established in [6].\n\n[1] Bayesian approaches to distribution regression\n[2] Learning Deep Features in Instrumental Variable Regression\n[3] Noise Contrastive Meta-Learning for Conditional Density Estimation using Kernel Mean Embeddings\n[4] Meta Learning for Causal Direction\n[5] Deep proxy causal learning and its application to confounded bandit policy evaluation\n[6] G. Blanchard, G. Lee, and C. Scott. Generalizing from several related classification tasks to a new unlabeled sample. Advances in Neural Information Processing Systems, 24, 2011.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "mN7L8Xskss", "reviewer": "Reviewer_bRTz", "summary": "The authors propose a method for learning kernel embeddings for distributions via maximizing a Renyi entropy objective. They relate their objective to the distributional variance of the embeddings, and explain why this can lead to good embeddings for downstream tasks. Empirical evaluations show good results on a flow cytometry dataset, but results on image and text classification are weaker.", "questions": "- How is the sampling done in the mini-batch SGD? There are two levels in the two-stage kernel embedding process, one with the individual samples and one with the distributions. How do the authors pick the SGD samples to ensure good optimization progress? And what are the required batch sizes relative to the training dataset size?", "limitations": "Yes", "rating": 4, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 2, "strengths": "- This paper is very clearly written, with good explanations on kernel distribution embedding and covariance operators, and the two-stage embedding scheme used in the estimation. There are also good motivation on the use of entropy maximization as an objective for learning good feature embeddings. \n\n- The authors provide theoretical analysis on the relation between second-order Renyi entropy and distribution variance, and how maximizing the entropy helps increase the variance of the feature embeddings.", "weaknesses": "- The weaker part of this paper is the empirical evaluations. There is only one dataset provided in the main paper on flow cytometry which the proposed method gives improvements over existing methods. There are small-scale experiments on image and text classification in the Appendix. However, even as the results of the learned embeddings are better than randomly initialized embeddings, they are not as good as direct SGD optimization using a cross-entropy loss on MNIST or 20 News-groups. This cast doubt on the effectiveness of the entropy maximization objective for feature learning in text and image problems. \n\n- Another weakness of the current method is scalability. From Equation 14 the objective scales quadratically with the number of distributions. This makes the computation expensive and indeed most of the experiments described in this paper are rather small-scale.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Dear Reviewer,\n\nThank you very much for your thoughtful review.\n\nUpon reading the Reviewer's comments we were confused by what exactly the Reviewer contends when referring to embeddings learned by SGD via Cross-Entropy (CE). Our methodology is inherently unsupervised, whereas CE intrinsically necessitates label information. Consequently, juxtaposing our approach with SGD optimization of the CE loss does not constitute a congruent comparison, given that our optimization goal is unsupervised—learning a data-dependent kernel between distributions devoid of label access—unlike the supervised nature of CE. We invite any additional insights from the Reviewer should there be an alternate perspective.\n\n\nOn the other hand, we agree that our method’s scalability to larger datasets warrants further work. However, we posit that scalability challenges do not preclude the applicability of our method to practical datasets. Our methodology has demonstrated superior performance in real-world contexts, such as clinical studies. The leukemia diagnosis dataset featured in our study, derived from clinical trials, is representative of the complexity typically encountered in medical diagnoses and has been made available to support research in early cancer detection and therapy. This dataset exemplifies common challenges in the same class of ML applications: \na large sample size (in millions of vectors), non-i.i.d. sampling due to a small number of subjects, and challenges in normalizing vector representations because of the nature of the measurements (e.g., different biomarkers). It is our stance that the utility of methods like ours should not be contingent solely on their suitability for the largest available datasets. Our approach is equally applicable to small and medium-sized datasets, as evidenced by our findings.\n\n\n\nTo address specific questions about the sampling procedure, both levels of sampling are performed uniformly to obtain a batch of i.i.d. samples of a given size. We do not have quantitative guarantees for the optimal batch size, and the best hyperparameters were determined using a standard grid search. In addition, in the experiments conducted for this study, our empirical observations suggest that batch size is not a critical determinant of the downstream performance.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "1FDeuXI0oz", "reviewer": "Reviewer_m5U2", "summary": "This work studies the distribution regression problem. The inputs are distributions, and the goal is to learn embeddings for these inputs. The authors propose the following method: First map the distributions to the kernel mean embeddings with respect to a embedding kernel $k_{emb}$, and then conduct a kernel regression on a distribution kernel $K_{distr}$. The authors propose to learn $k_{emb}$ by maximizing the quantum entropy. The effect of this is minimizing the variance within each distribution, while maximizing the spread of distributions.", "questions": "See weaknesses.", "limitations": "Limitations are discussed", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "**Note:** I reviewed this paper at ICML 2024. This review is also based on the difference between the two versions.\n\n1. This paper studies an important problem, which I believe has potential applications in many domains. The paper is well written and easy to read. The intuition behind the proposed method is clearly demonstrated. The proposed approach, as far as I know, is new. Thus, this paper contributes valuable insights and methodologies to the literature.\n\n2. The geometric interpretation, that is minimizing the variance and maximizing the spread, makes sense and is similar to some related methods such as contrastive learning (see Wang & Isola, 2020). \n\n3. The paper is well-structured, offering a comprehensive and clear exposition of the background, methodology, and theoretical aspects of the proposed approach. This clarity enhances the paper's accessibility to readers who might not be intimately familiar with the nuances of kernel methods or distributional regression, thereby broadening its potential impact.", "weaknesses": "Several weaknesses were discussed at the ICML review, which are not fully addressed in the new version:\n\n1. The major weakness is the experiment part, which only considers the flow cytometry task. While this might be an important application (which I am not familiar with), the usefulness of the proposed method in applications is still questionable. Experiments on images and text are moved to Appendix D, while they were in the main body in the previous version. For me I would prefer them to be in the main body. As I pointed out last time: \"One thing that the authors could do to increase the impact of this work is to apply this method to more real tasks. For instance, the authors mentioned voting behavior prediction and dark matter halo mass learning in the introduction, but these data sets are not used in the experiments.\"\n\n2. Other weaknesses in the empirical part pointed out by other reviewers:\n  - Analysis of Important Variables: The experimental setup lacks a thorough analysis of crucial variables, and does not have ablation studies. The impact of hyperparameter selection on the method's performance should be studied and discussed in the paper.\n  - Computational Cost and Scalability: The authors acknowledge this point in the limitations.\n  - Reproducibility: The authors promise to release their code in the main text.\n\n3. One reviewer asked last time why not directly optimizing $V_H$, which is an upper bound of the proposed objective as proved in Eqn. (17). The authors explained that directly optimizing $V_H$ has a lower performance, but did not provide any experiment results. I am not satisfied with this explanation. I suggest the authors add a direct comparison in the paper. This could be an important ablation study.\n\n**Conclusion:** Overall, I am still in favor of accepting this paper. This paper indeed has lots of issues, but it is a preliminary study of an important problem and has some interesting insights, which I believe could inspire future work. My current rating is weak accept, and I am willing to raise it to accept based on the authors' response.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Dear Reviewer,\n\nWe are grateful for your meticulous review and constructive feedback.\n\nSince our initial submission to ICML, the manuscript has undergone significant improvements, particularly in the experimental section, prompted by critiques akin to those presented by the Reviewer.\n\nWe concur that incorporating additional experiments within the main text could be advantageous. However, we opted to relegate the image and text modalities to the Appendix to mitigate potential confusion. \nThe feedback we received indicated a prevalent misunderstanding among readers. They often speculated about the relevance of supervised methods, which are standard for the datasets in question (e.g. MNIST), to our research framework. However, our framework is unsupervised and the quasi entirity of the suggested methods are not comparable to our framework. Therefore, we chose to highlight an experiment where the application of a distributional regression model is more evident.\nThe leukemia diagnosis dataset, derived from clinical trials, best represents the complexities inherent in medical diagnosis and illustrates typical machine learning challenges, such as large sample sizes and non-i.i.d. sampling.\n\nRegarding the other points of critique, we recognize that scalability is a challenge for the practical deployment of kernel methods. We intend to conduct a more thorough analysis of the hyperparameter selection process. Nevertheless, we have shown that our method remains viable for small to medium-sized datasets, where scalability issues are less pronounced. The principal contribution of our paper is the establishment of a theoretically sound objective for unsupervised learning of a data-dependent distributional kernel. This contribution is innovative both as a theoretical construct and in its implications for the application of kernel-based methods to discriminative tasks involving distributions. We anticipate that this novel perspective will stimulate further interest and research in the field. We welcome efforts by other research groups to adapt and scale our methods for use with larger datasets.\n\n\nA brief comment about using $\\mathbb{V}_\\mathcal{H}$ as an optimization objective. We will incorporate an ablation study in the experimental section. Although the outcomes do not surpass those of a random encoder, which is expected as per theoretical considerations discussed in the paper.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "jToZHZVJmo", "reviewer": "Reviewer_pXGE", "summary": "The authors consider the distribution regression problem. The authors propose to learn an embedding function (i.e., embedding support data points into a sphere) and leverages the kernel embedding (into RKHS) and kernel functions (within RKHS space). The authors propose an algorithmic approach to learn such embeddings via maximum kernel entropy in an unsupervised fashion. Empirically, the authors illustrate advantages of the learned kernel with SVM on Tissue and Leukemia classification.", "questions": "The submission is well organized, and easy to follow. The authors also demonstrate detailed interpretations for the proposed approach. The proposed learned kernel approach is interesting and has good empirical results for the distribution regression problem. However, it seems that the experiments are quite weak. \n+ The authors only evaluate on two small datasets (with N <=50 distributions?) and on a classification task only (?). It is better to evaluate on several larger datasets (especially in our current era)! Additionally, it is quite unusual to learn an embedding from only 16 distributions (for Tissue) and 20 distributions (for Leukemia). Could the authors clarify why it only needs a few input distributions to learn such embedding? Or only the number of supports in the distributions matters here?\n+ It is only a plus to illustrate it on a real distribution regression problem.\n+ It is better to report time consumption together with performances, and compare to recent kernel approach for distributions, e.g., based on optimal transport geometry.\n+ It is also better to discuss the parameterization for the embedding and the proposed algorithm to learn it? It seems that the learning algorithm is non-convex?\n+ Please give a proof for results in line 157-158? What is the gradient of L_MDKE w.r.t. \\theta? Please describe S_2(\\Sigma_D) w.r.t. K_D? and does the result in Equ. (13) only hold for the covariance operator \\Sigma_D or also its estimator \\hat\\Sigma_D? Please clarify them.\n\nSome other concerns are as follows:\n\n+ The kernel hyperparameter (in Section B.1) is important and interesting. Could you elaborate it rigorously for the experiments with few distributions?\n+ Could the authors comment on the choice of sphere for the embedding? Why not other spaces? It is also better to use a trivial baseline, e.g., learn embedding by autoencoder separately, and then apply the kernel embedding (into RKHS) and kernel functions (within RKHS) to illustrate the importance of the learned embedding?\n+ It may be interesting to link the distribution regression with optimal transport (OT) geometry besides sliced-Wasserstein kernel [35], e.g., tree-Wasserstein kernel (Le et al., NeurIPS'19), Sobolev transport kernel (Le et al., AISTATS'22)? and other geometry on distributions? Could the authors comment on it?\n+ It seems the proposed approach may be also related to the Log-Hilbert-Schmidt metric between positive definite operators (Ha Quang et al., NIPS'14), especially on application setup? Could the authors comment on it?\n\nBriefly, the proposed method is interesting, and the submission will be stronger in case the authors have a further care for the experiments, e.g., with larger datasets, real distribution regression task.", "limitations": "Yes, the authors have a discussion on the limitations.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "+ The authors propose an interesting framework to learn a kernel based on kernel embedding (into RKHS), parameterized by embedding function (into a sphere) in an unsupervised fashion.\n+ The proposed method is well motivated. The paper is well-organized, easy to follow, with a detailed description for the proposed framework.\n+ The interpretation and illustration of main concepts of the proposed method is also a plus.\n+ The proposed learned kernel with SVM performs well in experiments.", "weaknesses": "+ It seems the experiments are quite weak. The authors only evaluate on two small datasets (with N <=50 distributions?) and on a classification task only (?). It is better to evaluate on larger datasets (especially in our current era)! It is only a plus to illustrate it on a real distribution regression problem.\n+ It is better to report time consumption together with performances.\n\n+ Another point is that the embedding into a sphere is not well motivated (e.g., why sphere? but not other spaces? How is about low-dimensional embedding, as other criteria for embedding?)", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Dear Reviewer,\n\nThank you very much for your thoughtful review.\n\nBefore addressing your specific questions, we want to emphasize that the leukemia diagnosis dataset used in the main part of our paper is a real dataset. It was collected during clinical studies, is sufficiently diverse, and represents the typical medical complexities involved in diagnosis. It has been published to facilitate research in early cancer diagnosis and treatment. From ML perspective, this dataset showcases typical problems in such tasks: large sample size (in millions of vectors), non-i.i.d sampling due to a small number of subjects, and challenges in normalizing vector representations because of the nature of the measurements (e.g., different biomarkers). Appendix D includes experiments with more common ML modalities to demonstrate the applicability of our framework.\n\nTo address your questions:\n\n* Regarding the question of using only a few input distributions. Specifically for the Leukemia dataset, the input sample space is continuous, and the encoder function is a non-linear map (i.e. NN). This complexity makes it difficult to estimate the number of distributions required to learn an encoder that sufficiently separates parts of the support space for downstream tasks. While such a quantitative statement would be beneficial for practical applications, we currently lack these estimates, even for much simpler setups.\n* Regarding time consumption, we report in the paper the computational resources required for our experiments as not significant.\n* The non-convexity of the objective is discussed alongside our choice of optimization algorithm in lines 157-163.\n* The learning objective involves the logarithm of the squared Frobenius norm of the kernel Gram matrix, and the gradient can be easily obtained using any autodiff software. Given that the encoder function is a neural network, deriving the gradient manually may be tedious, and we believe it does not provide additional theoretical value for the discussion in our paper. Computational form of $S_2(\\Sigma_D)$ w.r.t. $K_D$ is given in Eq. 14 on page 5. Eq. 13 is stated in terms of true operator $\\Sigma_D$, the rest relies on the fact that $\\hat{\\Sigma}_D$ is a consistent estimator of $\\Sigma_D$ (with the proof of this fact given in [1]).\n* The hyperparameter selection procedure we applied is defined and analyzed in [2].\n* Regarding the choice of a hypersphere as the latent space: as we stated in lines 141-143, the latent space is a design choice. Other compact spaces with sufficient symmetries and the existence of proper kernels (see Assumption 3.1) would also be valid. The hypersphere is a simple and commonly used latent encoding space in ML, possessing the properties required.\n* While we acknowledge that a specialized kernel suited to the task's geometry might improve solution quality, the study aims to learn the kernel from the dataset (or, more broadly, identify what makes a kernel \"good\" for solving discriminative tasks on distributions).\n* Currently, we do not see an apparent connection to the Log-Hilbert-Schmidt metric. Note that in our work, we do not use covariance operators to define the kernel. The operator embedding is only defined as a learning objective. Once learning is complete, the resulting kernel does not depend on the covariance operators.\n\n[1] Bach, Francis. \"Information theory with kernel methods.\" IEEE Transactions on Information Theory 69.2 (2022): 752-775.\n\n[2] G. Blanchard, G. Lee, and C. Scott. Generalizing from several related classification tasks to new unlabeled sample. Advances in neural information processing systems, 24, 2011", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "uYVK83Pbo9", "reviewer": "Reviewer_Z862", "summary": "This paper focuses on distributional regression (classification) that carries out supervised learning over a collection of datasets, where one instance (subject to one label) is a dataset that can be considered a distinct empirical distribution. The paper proposed a novel learning objective, maximize quantum entropy, to find a suitable embedding function f_{\\theta} for all datasets on a sphere. The proposed quantum entropy is a the lower bound on distributional variance, where optimizing over such objective encourages inner-distribution variance reduction and inter-distribution variance increase for the given kernels, which intuitively lead to well-separability among distributions. Empirical simulation shows the potential of such objective in distributional supervised learning.", "questions": "As I am not an expert in distributional regression, I am unsure about how important regression analysis in this subject is. However, I do not see any empirical evaluation on anything that is technically regression (compared to classification). Moreover, the motivation behind the quantum entropy, as suggested in Figure 2a, is more aligned with classification (separability) than regression. Maybe the author could make a remark on this particular subject (regression), or provide a simple set of simulations on actual regression?", "limitations": "Yes", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The paper is well-written and well-organized, the mathematical notation is on point and easy to follow. I enjoyed reading the paper.\n2. The theoretical results are concise and straightforward, which provides enough motivation and context for the proposed framework.", "weaknesses": "1. My major concern is about the validity of the operation depicted in line 302-303, where the unsupervised pretraining is done with the whole dataset rather than a subset. This seems to be a slight lack of rigor and I hope the author could address how would MDKE perform when the pretraining is only done on the training dataset.\n2. The second weakness is regarding the performance of MDKE in actual regression, for more detailed question, please refer to questions. At this point, the potential use case of MDKE seems rather limited (Distributional classification).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Dear Reviewer,\n\nThank you very much for your thoughtful review.\n\nRegarding the topic of unsupervised learning datasets, the Leukemia experiment conducts unsupervised kernel learning on the entire set of available distributions without accessing the corresponding labels. In the experiments described in Appendix D, unsupervised kernel learning is performed on subsets of the data.\n\nTo address your question about classification versus regression: as it is stated in the abstract, the primary goal of this work was to study **discriminative tasks**. While we believe that minimizing latent variance ultimately facilitates solving regression problems as well, the experimental verification of this hypothesis was not within the scope of the current work.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: This paper addresses the kernel learning problem with the maximum kernel entropy, which is one of the most important (and often neglected) problems in distribution regression. The reviewers unanimously concur on the importance of the problem and the novelty and suitability of the proposed methods. Nevertheless, the reviewers also pointed out several limitations (e.g., small datasets, marginal improvement, etc) in the empirical evaluations of the proposed methods. I encourage the authors to further address these concerns in the camera-ready version of the manuscript, for example, by clarifying the experimental setup and moving some of the results from the appendix to the main paper.
Decision: Accept (poster)