Learning to Embed Distributions via Maximum Kernel Entropy

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.

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

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?

Rating

6

Confidence

4

Soundness

3

Presentation

3

Contribution

3

Limitations

Yes

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.

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. + The proposed method is well motivated. The paper is well-organized, easy to follow, with a detailed description for the proposed framework. + The interpretation and illustration of main concepts of the proposed method is also a plus. + 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. + 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

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.

Rating

6

Confidence

4

Soundness

3

Presentation

3

Contribution

3

Limitations

Yes, the authors have a discussion on the limitations.

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.

Strengths

**Note:** I reviewed this paper at ICML 2024. This review is also based on the difference between the two versions. 1. 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. 2. 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). 3. 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: 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

See weaknesses.

Rating

6

Confidence

3

Soundness

2

Presentation

3

Contribution

2

Limitations

Limitations are discussed

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.

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

- 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?

Rating

4

Confidence

4

Soundness

4

Presentation

3

Contribution

2

Limitations

Yes

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.

Strengths

- The paper proposes a novel way to learn kernels through entropy maximiazation in the setting of distribution regression. - The paper is very well written and easy to follow for someone who has worked the KME before. - 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? - 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

see above

Rating

6

Confidence

3

Soundness

3

Presentation

4

Contribution

2

Limitations

yes

response

I thank the authors for the clarifications and i will keep my score of 6.

Response

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.

Response

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.