Learning to Embed Distributions via Maximum Kernel Entropy

Dear Reviewer,

Thank you very much for your thoughtful review.

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

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

Dear Reviewer,

Thank you very much for your thoughtful review.

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

To address your questions:

• 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.
• Regarding time consumption, we report in the paper the computational resources required for our experiments as not significant.
• The non-convexity of the objective is discussed alongside our choice of optimization algorithm in lines 157-163.
• 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]).
• The hyperparameter selection procedure we applied is defined and analyzed in [2].
• 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.
• 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).
• 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.

[1] Bach, Francis. "Information theory with kernel methods." IEEE Transactions on Information Theory 69.2 (2022): 752-775.

[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

Dear Reviewer,

We are grateful for your meticulous review and constructive feedback.

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

We 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. The 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. The 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.

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

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

Dear Reviewer,

Thank you very much for your thoughtful review.

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

On 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: a 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.

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

Dear Reviewer,

Thank you very much for your thoughtful review.

• 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.
• 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].
• 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].

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