Learning to Embed Distributions via Maximum Kernel Entropy

Authors: oleksii  kachaiev
stefano  recanatesi
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-14 23:49:16

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.

Main paper: 
 introduction: 
 Most discriminative learning methods conventionally assume that each data point is represented as a real-valued vector. In practical scenarios, however, data points often manifest as a 'set' of features or a '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 a distribution of, measurements. One approach to accommodate such cases involves representing each input point as a probability distribution. Beyond mere convenience, it is more appropriate to model input points as distributions when dealing with missing data or measurement uncertainty, as often encountered when facing the abundance of data, which commonly presents a challenge for data-rich fields such as genetics, neuroscience, meteorology, astrophysics, or economics.The task of regressing a mapping of probability distributions to a real-valued response is known as distribution regression. Distribution regression has been successfully applied in various fields, such as voting behavior prediction [13] , 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 in solving complex problems [69, 29] . Kernel methods have become a widely used approach for solving distribution regression tasks by exploiting a kernel 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 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 a fundamental question: What are the underlying principles that facilitate the unsupervised learning of an effective kernel, one that optimally encapsulates the data properties and is well suited for discriminative learning on distributions?In this work, we leverage a key insight: an appropriate selection of the distribution kernel enables the embedding of a set of distributions into the space of covariance operators. Building on this theoretical idea, we claim that quantum entropy maximization of the corresponding covariance operator is a suitable guiding principle to learn data-dependent kernel. This, combined with a careful design of kernel parametrization, let us to devise a differentiable optimization objective for learning data-specific distribution embedding kernels from unlabeled datasets (i.e. unsupervised) (Fig. 1 ). 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 to the common practice of hand-picking kernels. Such framework could be further leveraged for generalizing existing learning approaches [61, 32] catalyzing the use of distribution-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.Consider an input space X and a positive-definite (p.d.) kernel k : X × X → R. Let H be the corresponding reproducible kernel Hilbert space (RKHS) induced by such kernel. Consider a probability distribution P ∈ P(X ). The kernel mean embedding map embeds the distribution P as a function in Hilbert space:EQUATIONwhere ϕ : X → H is a feature map such that ϕ(x) = k(x, •).Importantly, if the kernel k is characteristic [50] , the mapping µ : P(X ) → H is 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 [41, 48, 14] . The natural empirical estimator for the kernel mean embedding approximates the true distribution with a finite sum of Dirac delta functions:EQUATIONwhere x 1 , . . . , x N ∼ P are N empirical i.i.d. samples. The estimator has a dimension-free sample complexity error rate of O(N -1 2 ) [51] . Additionally, we denote d k : X × X → R ≥0 a kernel metric induced by a kernel k,EQUATIONNote that d k is a metric on X if feature map ϕ is injective.A second way of mapping a probability distribution to a Hilbert space can be defined by means of a covariance operators. For a given feature map ϕ(x) = k(x, •) : X → H 1 , and a given probability distribution P ∈ P(X ), the covariance operator embedding is defined as:EQUATIONwhere ⊗ is a tensor product. Σ P is a self-adjoint positive semi-definite (p.s.d.) operator acting on H. Such operator can be seen as a mean embedding w.r.t. the feature map x → ϕ(x) ⊗ ϕ(x) and therefore, for a universal kernel k, the map P → Σ P is injective (see Bach [1] ).Similarly to Eq. ( 2), the natural empirical estimator is:EQUATIONwhere x 1 , . . . , x N ∼ P are N i.i.d. samples.For a translation-invariant kernel k(x, x ′ ) = ψ(x -x ′ ) normalized such that k(x, x) = 1, the covariance operator Σ P is a density operator [1] . Henceforth, entropy measures can be applied to it, and the quantum Rényi entropy of the order α can be defined as:EQUATION)where {λ i } i are the eigenvalues of Σ P . 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. α = 2 (see Carlen [9] , Müller-Lennert et al. [42] , Wilde [63] , Giraldo et al. [15] for an in-depth overview of the properties and theory of quantum entropies).3 Unsupervised Distribution Kernel Learning 
distribution regression: 
 In this section, we 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 are given as probability distributions available only through samples. Formally, consider the task of finding a regressor f : P(X ) → Y from the dataset of samples Γ M = {(P i , y i )} M i=1 where P i ∈ P(X ) are distributions provided as a set of i.i.d. empirical samples x 1 , . . . , x Ni ∼ P i (see Poczos et al. [45] , Szabó et al. [53, 54] for a comprehensive analysis). A viable approach to solving this problem is to define a kernel K : P(X ) × P(X ) → R that is universal in P(X ). By setting up such a kernel K, we can utilize kernel-based regression techniques, with SVM often being the preferred method for classification tasks [39] or Kernel Ridge Regression (KRR) when the space Y is continuous [37] 2 . To this end, several kernels have been proposed over time (see details in Sec. 4).One possibility, proposed by Muandet et al. [39] , is to introduce a kernel in the input space, the so called embedding kernel k emb : X × X → R and exploit the induced mean embeddings µ emb : P(X ) → H emb to map input distributions to points in RKHS. Subsequently, to define a second level kernel, distribution kernel K distr : H emb × H emb → R between points in the RKHS H emb . The simplest choice for such a distribution kernel is the linear kernel:EQUATIONA standard alternative to the linear kernel is a Gaussian kernel with bandwidth parameter λ > 0:EQUATIONwhich was shown to be universal in P(X ) [11] . The Gaussian kernel K RBF can be computed from the linear kernel K l as d kemb (µP , µ Q ) 2 = K l (P, P ) + K l (Q, Q) -2K l (P, Q).In practice, we only have access to a finite number of 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 is analyzed in Szabó et al. [53] .Instead of using standard kernels designed to encapsulate the geometry of the input space, we consider learning a data-dependent kernel, tailored to the specific properties of the dataset. In a similar vein, Yoshikawa et al. [68] proposed learning an optimal kernel (or equivalently, a feature map) jointly with the classifier to address the text modality. In this work, we focus on an unsupervised problem, aiming to 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 of 2-level kernel setup, we define the embedding kernel as kEQUATIONwhere f θ is a trainable encoder function f θ : X → Z, Z is a latent encoding space, and k emb is a kernel defined on the latent space k emb : Z × Z → R. The encoder function f θ transforms every input probability distribution P ∈ P(X ) into a distribution over the latent space P θ ∈ P(Z) 3(Fig. 1a ). Furthermore, we denote RKHS corresponding to the kernel k emb as H emb and the kernel mean embedding map as µ emb (see Eq. ( 1)).EQUATIONThese transformations define mean embeddings for each input probability distributions through the first level, embedding kernel k emb (Fig. 1b ). The second level, distribution kernel K distr : H emb × H emb → R is defined over the mean embeddings µ emb (P )'s. We can now consider to embed dataset D M = {P i } M i=1 as an empirical covariance operator (see Eq. ( 5)), i.e.EQUATIONAs ΣD encapsulates information about the entire dataset, we term it dataset embedding. With the assumption that the dataset D M is sampled i.i.d. from the (unknown) true meta-distribution D, ΣD is a natural estimator to approximate the true covariance operator Σ D . To simplify the notation we use Σ D in place of Σ D unless required by the context.Both the embedding kernel k emb and the distribution kernel K distr remain fixed throughout the training, learning happens by adjusting the parametrization of the latent space encoder f θ . Such separation ensures expressivity while conforming to all technical requirements for a distribution kernel. Throughout the paper, we make the following assumptions on the latent space Z and embedding kernel k.Assumption 3.1. Latent space Z is a compact subset of R d . Kernel k emb : Z × Z → R is a p.d. characteristic translation-invariant kernel k emb (z, z ′ ) = f (∥z -z ′ ∥ 2 ) such that -f ′ is completely monotone on (0, ∞) (see Definition 2.2.4 of Borodachov et al. [7]) and ∀z ∈ Z : k emb (z, z) = 1.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 [66, 57, 58, 67] . In contrast, the parametrization proposed in our study applies a known characteristic kernel to a learned latent representation. To facilitate our optimization process (which will be explained shortly), we opt for Z = S d-1 (the d-dimensional hypersphere) and Gaussian kernel both for k emb and K distr (see Eq. ( 8)). We retain other suitable choices as potential avenues for future research.This dataset level representation depends on the choice of first and second level kernels k, K and, in turn, on the trainable function f θ parameterized by the set of parameters θ. In this work, we propose learning the parameters θ to maximize quantum entropy of the dataset embedding, i.e.,EQUATION)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 2 (Σ D ). Building upon previous work [1] 4 , we exploit the following property of the covariance estimator:EQUATIONwhereK D ∈ R M ×M is the distribution kernel matrix, with [K D ] ij = K distr (µ Pi , µ Pj ).This equation follows directly from the fact that Σ D and 1 M K D share the same set of eigenvalues. Leveraging this relationship, we can define tractable unsupervised training loss, which we term Maximum Distribution Kernel Entropy (MDKE), with respect to the parameters of the encoder f θ :L MDKE (θ) := -S 2 (Σ D ) = log tr 1 M K D 2 = log M i=1 λ 2 i 1 M K D = log ∥ 1 M K D ∥ 2 F (14)where the latter relies on the fact that the Frobenius norm ∥A∥ 2F = i λ 2 i (A), where λ i (A) are eigenvalues A.The MDKE objective is differentiable w.r.t. θ for commonly used kernels, provided that the encoder f θ is differentiable as well. While the entropy estimator S 2 (Σ D ) is convex in the kernel matrix K D , the objective as a whole is generally not convex in θ. However, in practice, as we show in Sec. 5, mini-batch 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.1.The Frobenius norm formulation in the loss Eq. ( 14) significantly reduces computational complexity. However, as we have observed in some of our experiments, it can lead to the collapse of small eigenvalues of K D , particularly near the optimal value of the objective. To address this challenge we introduced a regularized version of the loss L MDKE-R that incorporates optional regularization, based on the determinant K D inspired by the connection with Fekete points (see details in Appendix B.2).The optimization objective is specifically designed to minimize the variance within each distribution (inner-distribution variance) while simultaneously maximizing the spread of distributions over the compact latent space Z = S d-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.r.t. the MDKE loss) configuration of mean embeddings in the RKHS attains describe properties. For doing so we leverage the notion of distributional variance V H (Definition 1 in Muandet et al. [40] ). Definition 3.2. For a set of M probability distributions D M , distributional variance V H (D M ) of the mean embeddings in the RKHS H is given byEQUATION)where G is the M × M Gram matrix of mean embeddings in H, i.e. G ij = ⟨µ Pi , µ Pj ⟩ H [40] .Here we show that the distributional variance V H can be equally reformulated into two separate contributions:EQUATIONwhere P denotes mixture distribution with elements of D M being uniformly weighted mixture components (see proof in the Appendix A.1).The relevance of distributional variance for MDKE objective is established by the following result. Proposition 3.3. For a set of M probability distributions D M , the second-order Rényi entropy S 2 of the empirical covariance operator embedding ΣD induced by the choice of Gaussian distribution kernel K RBF over points in the RKHS H emb , -as defined in Eq. (8), -is upper bounded by the distributional variance V Hemb (D M ), i.e.,EQUATION)where γ is the bandwidth of the distribution kernel K RBF .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.1, the maximum norm of kernel mean embedding is attained by Dirac distributions {δ z } z∈Z .This result is trivial due to the fact that the set of mean embeddings is contained in the convex hull of {k emb (z, •)} z∈Z , and, under Assumption 3.1, ∀z ∈ Z :EQUATIONThe 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 3.5 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 1 from Wang and Isola [61] . 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 Hyperspherical Uniformity Gap [32] to infinite-dimensional 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 a broader explanation of the connection.More generally, utilizing the fact that under Assumption 3.1, the kernel metric d kemb (see Eq. ( 3)) is a monotonically increasing function of Euclidean metric on the latent space, we establish a 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 V Hemb 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, V Hemb retains its value as an insightful and intuitive measure for describing the geometric configuration of the learned system.We use a simple example to illustrate the connection between geometrical configurations of embedded distributions and distribution kernel entropy S 2 (Σ D ) (see Fig. 2 ). We sample a number of points from 6 different Gaussian distributions and project on a sphere S 2 varying their projected variance γ. As γ decreases, the distributional variance of the overall distribution of Gaussians increases (Fig. 2a ). For very small γ each distribution converges to a point (a Dirac distribution). This results in the entropy interpolating between lower and upper bounds, demonstrating how entropy behaves in response to changes in distribution variance. Fig. 2b showcases the behavior of two terms comprising distributional variance (Eq. ( 16)): 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.Runtime complexity. The applicability of a data-dependent distribution kernel to solving discriminative 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 a common concern in practical applications, is most favorable for the tasks described. A detailed analysis of runtime complexity can be found in Appendix B.3.Broader impact. We wish to emphasize that the distribution regression framework has emerged as a powerful tool for analysis and predictive modeling, especially in domains where traditional methods face challenges, including social science, economics, and medical studies. We urge researchers and practitioners applying distribution regression in these areas to give special consideration to 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 
distribution regression: 
 Distribution regression was introduced in Poczos et al. [45] , while the seminal work of Szabó et al.[54] provides a comprehensive theoretical analysis of this regression technique. A natural approach to solving distribution regression problems involves using kernels between measures. Notable examples include the Fisher kernel [17] , Bhattacharyya kernel [18] , Probability product kernel [19] , kernels based on nonparametric divergence estimates [52] , and Sliced Wasserstein kernels [23, 37] . 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 in Hilbert space to the outputs. Distribution kernels have been successfully applied in various kernel-based methods, such as 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 [44] , 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) is available during training [70, 26, 25] . For an in-depth exploration of the diverse methodologies employed in distributional regression settings, we invite readers to consult Appendix C.1.Quantum entropy, including Rényi entropy, is a powerful metric to describe information in a unique way (see Müller-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 [1] applied von Neumann entropy of the density matrix to the covariance operator embedding of probability distributions, thereby defining an information-theoretic framework utilizing kernel methods. In machine learning, especially within self-supervised learning (SSL) setups, entropy concepts have recently found novel applications. Our study builds on most recent developments [49, 21, 55] by applying quantum Rényi entropy to the covariance operator in RKHS. 
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 is 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 as an embedding kernel and distribution kernel, the hyperparameter selection is performed as described in Appendix B.1.Once the samples encoder f θ 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. A grid search with 5 splits (70/30) is conducted to optimize the strength of the squared l 2 regularization penalty C, exploring 50 values over the log-spaced range {10 -7 , . . . , 10 5 }. The best estimator is then applied to evaluate classification accuracy on the validation subset, which we report.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.Flow Cytometry (FC) is a widely used technique for measuring chemical characteristics of mixed cell population. Because population-level properties are described through (randomized) sampling of cells, FC is used as a canonical setup of distribution regression. For this study we used a dataset [56] where more than 100.000 cells are measured per each patient (subject). 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 N = 44 subjects. The second ('Leukemia' classification) presents healthy and leukemia BM cell samples, N = 50. Classes were balanced in both cases. ). Per each subject we sampled a small percentage of cells, and we report performance for the sample size of 200. We repeated each training and testing phase for 100 times to track the variance induced by this aggressive subsampling.To demonstrate the impact of unsupervised pre-training, we compared several methods across multiple configurations (Table 1 ): 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 [37] . b) MMD kernels. Here, we employ a 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 non-data-specific kernels.In the following section, we present additional experiments on learning 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-specific 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 as 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.Images. MNIST [12] and Fashion-MNIST [65] consist of 28 × 28 pixel grayscale images divided into 10 classes. We considered each individual image to be a probability distribution (via rescaling pixel intensities so that l 1 (image) = 1) over the discrete space of pixel positions (i.e., histogram). Given this support space, the encoder f θ is a discrete map which we implemented as a table lookup (i.e., embeddings) from pixel indices to the points on the hypersphere S 31 . 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 a lower variance of the accuracy for the downstream classification. 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%. Refer to the detailed analysis of the latent spaces of the trained encoders in Appendix D.1.Text. To assess our method's performance in a 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 f θ mirrored the setup used in the MNIST case (Appendix D.1), with l 2 normalized word embeddings on S 31 . 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 D.2. 
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 a 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. This result offers a 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. 
 
 Appendix: Definition A.2. For a dataset D M of M probability distributions, we define the average pairwise distance between kernel mean embeddings in H as J H (D M ), given by:J H (D M ) := 1 M 2 M i=1 M j=1 ||µ i -µ j || 2 H (20)Lemma A.3. For the distributional variance V H of the dataset D M , the following relationship holds:EQUATIONProof. Starting with the definition of J H (D M ):J H (D M ) = 1 M 2 M i=1 M j=1 ||µ i -µ j || 2 H = 1 M 2 M i=1 M j=1 ⟨µ i -µ j , µ i -µ j ⟩ H = 1 M 2 M i=1 M j=1 (⟨µ i , µ i ⟩ H + ⟨µ j , µ j ⟩ H -2⟨µ i , µ j ⟩ H ) = 1 M 2   2M M i=1 ⟨µ i , µ i ⟩ H -2 M i=1 M j=1 ⟨µ i , µ j ⟩ H   = 2       1 M M i=1 ⟨µ i , µ i ⟩ H diagonal of G - 1 M 2 M i=1 M j=1 ⟨µ i , µ j ⟩ H       Recognizing that M i=1 ⟨µ i , µ i ⟩ H = tr[G],we can equate expression in the brackets to V H (D M ), leading to:V H (D M ) ≡ 1 2 • J H (D M ),which concludes the proof.Lemma A.3 will be instrumental in further proofs, establishing a crucial link between distributional variance in H and quantum entropy of the covariance operator embedding Σ D .In this section, we provide the proof for the key theoretical result stated in Proposition 3.3.Consider a dataset D consisting of probability distributions {P ∈ P(X )} i sampled i.i.d. from an unknown meta-distribution D. We assert that the second-order Rényi entropy S 2 of the empirical covariance operator embedding Σ D , induced by the choice of Gaussian distribution kernel K RBF over points in the RKHS H, is upper-bounded by the distributional variance V H (D):1 2γ S 2 ( ΣD ) ≤ V H (D)Proof. From the definition of generalized variance we have the followingVarEQUATIONwhich 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.In practice, the effectiveness of the optimization process is significantly influenced by the parameters of the Gaussian kernels. Setting the bandwidth parameter γ either too low or too high can hinder the model's ability to learn effectively. This issue 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 γ based on the idealized structure of the dataset after its projection onto the latent space. Specifically, we set γ such that 1/γ equals 10 times the average distance to the nearest neighbor in the set of points sampled uniformly on S d-1 (inspired by the experimental approach in Blanchard et 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 is set to the number of distributions in a batch multiplied by the number of samples per distribution used for unsupervised training.As the theoretical maximum of S 2 entropy is attained when the spectrum of 1 M K D is uniform, the optimal encoder with respect to L MDKE also tends to maximize the determinant of the kernel Gram matrix K D . This relationship is intuitive when considering that the determinant of a matrix is the product of its eigenvalues. The points that maximize such determinant K D 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 et 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 L MDKE 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 to the tendency of too small eigenvalues in the distribution kernel matrix K D 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 a regularization term to the original objective. This term is inspired by the concept of Fekete points configuration, leading to the following loss function:EQUATIONHere, ϵ serves as a hyperparameter to control the strength of regularization. The regularization term Ω(Σ D ) 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 dynamics for the sentence representation learning experiment (as detailed in Appendix D.1) is showcased in Fig. 3 .Scalability of the kernel methods is typically a key concern when it comes to practical applications. The runtime complexity of the proposed method could be decomposed into two components:• Computation of the distribution kernel Gram matrix. This involves computing O(N 2 ) inner products between distributions where N 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 R M ×d where M is a 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(N 3 ) with a range of methods proposed to reduce the complexity [8, 47, 64, 36, 31] .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, a relatively small number of distributions in the dataset alleviates issues related to storing the distribution kernel Gram matrix in memory and performing computations on it 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 a parametric family, e.g. Gaussian Mixture Model (GMM), to use available closed-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 [17] 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. 5.1 we used fitting distributions to GMM with Fisher Kernel applied to learn parameters as a baseline for the performance on the task.Point clouds. This group includes methods that model each input distribution as a set of points (also known as 'point cloud' or '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-neighbor' set kernels to name a few [52, 30, 34, 16, 5, 5, 59] . 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 or similarity functions between point clouds [46] . While being computationally problematic for large-scale problems, Sliced Wasserstein kernels were succesfully adopted [23, 37] as they provided both reliable way for the point set comparison with a reduced cost leveraging sliced formulation. In Sec. 5.1 we compared performance of Sliced Wasserstein-1 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 has favoriable properties when compared to kernel based on Wasserstein-2 distance. 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 al. [53] provided theoretical guarantees for the for learning Ridge regressor from distribution embeddings in the Hilbert space to the outputs. Law et al. [28] During the unsupervised pre-training phase, we observed a steady decrease in training loss (Fig. 5a ) despite the small batch size (50 sentences with a 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 train and 1, 000 for 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 high accuracy in downstream classification showcases its prowess in learning potent latent representations, even when dealing with input distributions with large finite support.Assuming k to be a continuous positive definite (p.d.) kernel and X to be compact.As pointed out in Meunier et al.[37], under mild conditions, KRR can also be used for classification problems.By the definition of the encoding process, P θ is a push-forward measure. For an empirical probability distribution q = i δ(xi) ∈ P(X ) and a measurable map f : X → Z, the push-forward measure f #q ∈ P(Z) is defined as i δ(f (xi)).See the proof of Proposition 6 in Bach[1]. 


Decision: 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.

Reviews: review: paper_summary: paper_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.
strength_weakness: strength_weakness: - The experimental section does not compare to other kernel distribution regression methods. Could the authors confirm why [1] was omitted in the comparisons?
- In addition citations to other methods that also learn kernel albeit not with entropy should be cited [2, 3, 4, 5]
- 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?



[1] Bayesian approaches to distribution regression
[2] Learning Deep Features in Instrumental Variable Regression
[3] Noise Contrastive Meta-Learning for Conditional Density Estimation using Kernel Mean Embeddings
[4] Meta Learning for Causal Direction
[5] Deep proxy causal learning and its application to confounded bandit policy evaluation
questions: questions: see above
limitations: limitations: yes
paper_summary: paper_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.
strength_weakness: strength_weakness: - 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. 

- 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.
questions: 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: limitations: Yes
paper_summary: paper_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.
strength_weakness: strength_weakness: Several weaknesses were discussed at the ICML review, which are not fully addressed in the new version:

1. 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."

2. Other weaknesses in the empirical part pointed out by other reviewers:
  - 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.
  - Computational Cost and Scalability: The authors acknowledge this point in the limitations.
  - Reproducibility: The authors promise to release their code in the main text.

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

**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.
questions: questions: See weaknesses.
limitations: limitations: Limitations are discussed
paper_summary: paper_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.
strength_weakness: strength_weakness: + 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.
+ It is better to report time consumption together with performances.

+ 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?)
questions: 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. 
+ 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?
+ It is only a plus to illustrate it on a real distribution regression problem.
+ 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.
+ 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?
+ 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.

Some other concerns are as follows:

+ The kernel hyperparameter (in Section B.1) is important and interesting. Could you elaborate it rigorously for the experiments with few distributions?
+ 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?
+ 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?
+ 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?

Briefly, 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: limitations: Yes, the authors have a discussion on the limitations.
paper_summary: paper_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.
strength_weakness: strength_weakness: 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.
2. 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).
questions: 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: limitations: Yes

Comments: title: Response
Response
response
comment: Thanks for engaging on technical details. We really appreciate and enjoy discussing this with you:

**(1) Time consumption**

Mini-batch SGD is a common approach for tackling non-convex problems, typically able to overcome local minima due to the stochasticity of gradient estimates. Since it is usually not feasible to predict the performance of an SGD solution for a specific problem at hand, we propose it as an efficient method for solving the described optimization objective based on empirical evaluations. Note that the objective is convex w.r.t. to the kernel Gram matrix. The use of a neural network as the encoder turns it non-convex w.r.t. to parametrization of the NN. However, in our experience, SGD remains a suitable choice, often yielding satisfactory results in such scenarios. While we obviously recognize the value of quantitative statistical guarantees, our implementation of SGD aligns with standard practices in the literature for deep NNs and we expect the properties of our implementation to be mostly inherited by studies of mini-batch SGD convergence in deep-NNs.

Regarding computational time, our process is relatively efficient, typically requiring only hours on a single machine. This negates the need for extensive parallelization across multiple GPUs or machines, underscoring our method’s flexibility and accessibility.

**(2) Data-size is too small**

We tried to explain at best the dataset but, as for any dataset, real insight comes only by playing with it. This is actually not a very small dataset: it has millions of samples! Each distribution has 10^5-10^6 samples on which the optimization is performed. What keeps the dataset compact is the dimension of each sample ranging from 20 to 50 based on the selected criteria. While the dataset’s dimensionality is smaller than what you might find in image datasets, the number of samples—which we refer to as the “distribution part”—is quite extensive. The apparent simplicity of the downstream classification task we use to properly assess quality of the distribution kernel learned is the result of specific modeling and parameterization choices made. We hope it demonstrates the power of the techniques applied.

We use 70/30 splits for training, details are provided in the opening paragraphs of the Section 5.

**(3) Gradient formulation (for minibatch SGD)**

We agree with the Reviewer that the gradient is the most important building block of the learning procedure. We do not expect any substantial improvement from replacing autodiff with a theoretically derived custom gradient iteration. Moreover, by establishing a general framework, we aim to keep it flexible: any differentiable encoder can be incorporated without changes to the outer training loop, leveraging the power of modern autodiff software. We acknowledge that we might be wrong in such intuition and we are curious to know if the Reviewer suspects a different outcome. However, for this reason, while it is possible to manually derive a custom gradient, we didn't devote too much time to it.
Thank you for the response. I have no more questions as all my previous questions were rather minor. I will keep my score of 6.
I thank the authors for the explanation in the rebuttal.

I have some quick questions as follows:

**(1)Time consumption**
- It is not clear to me why the mini-batch SGD is efficient for the proposed non-convex learning problem yet? There is no analysis for it (e.g., convergence, affects of initialization, stopping condition, step-size, etc)? Could you elaborate why the computational time is not significant? 
- I agree that in Section 5, the authors show the advantages on performances for the proposed method. However, it is not clear about other aspects, e.g., algorithms to solve the non-convex learning problems? 

**(2) Data-size is too small**
- I agree that the used datasets are real-world.
- However, their size is just too small! The applications use kernel SVM (i.e., non-linear classification), how many samples do the authors use for training such classifiers (split 70/30 from a very small datasets with less than 50 samples?) Could the authors elaborate it more rigorously.

**(3) Gradient formulation (for minibatch SGD)**
- the gradient is essentially the most important part for SGD approach.
- How does the usage of autodiff affect the convergence and learning procedure?
- Is it possible to derive the formulation of the gradient?
I thank the authors for the response. I don't have any further questions at this point. I will discuss with my fellow reviewers and the AC, and notify the authors if I change my rating.
Thank you for your valuable feedback. Your clarification has greatly helped us understand your perspective.

We recognize that your emphasis on state-of-the-art (SOTA) performance is aimed at advancing current methodologies. However, we believe this focus may overlook the foundational nature of our contribution. Our work seeks to establish a new approach in unsupervised kernel learning, which we believe holds significant potential for future research and development. If we understand correctly, you are requesting a comparison between unsupervised and self-supervised approaches (e.g., SimCLR, MoCo) within a domain-free setup like ours.

This is indeed a fundamental question, but one that other subfields (such as vision and language) have addressed at more mature stages. Our contribution is the first to propose unsupervised learning of full distribution embeddings. While we appreciate your feedback and find it very helpful, we feel that rejecting our contribution based on these remarks may be asking too much from a single paper.

Here are some specific comments on your remarks:

1. **Domain-Specific vs. Domain-Free Approaches**: Our algorithm is designed to be domain-free, which we believe offers a significant advantage over domain-specific self-supervised methods. This flexibility allows our approach to be applied across various domains without requiring domain-specific adjustments.

2. **State-of-the-Art (SOTA) Focus**: While we acknowledge the importance of achieving SOTA results, our primary goal in this paper is to introduce a novel framework for unsupervised learning of whole distribution embeddings. We believe that establishing this new framework is a crucial first step that can pave the way for future improvements and comparisons.

Thank you once again for your insightful comments. We look forward to addressing them in our future work.
Yes, when I say SGD optimization with CE loss I do mean supervised learning with labels. In computer vision there is a line of work on unsupervised representation learning with contrastive methods such as SimCLR or MoCo, and the performance of these methods is on par with supervised learning on ImageNet. So it is not completely unfair to compare the performance of unsupervised representation learning with their supervised counterpart. The same is true in NLP, where unsupervised representation learning (although via pretraining on large amount of texts) beats supervised learning with limited data. It could be difficult to require a general unsupervised learning method like the author's proposal to do better than domain-specific approaches used in vision or NLP, but it also limits the applicability of the proposed algorithm on these type of data since there are strong alternatives.
I thank the authors for the clarifications and i will keep my score of 6.