The Representation Jensen-Shannon Divergence

Thank you for your valuable feedback and insightful comments on our paper. We appreciate the time and effort you have dedicated to reviewing our work. Next, we address the main concerns raised.

Questions

1. We would like to clarify that the kernel used to measure RJSD does not need to be normalized to be a probability density, and that the only restriction is that $\kappa(x,x) = 1$. In definition 1. we normalize the kernel to be a probability density, to be able to compare the covariance operator entropy $S(C_\mathbb{P})$ to Shannon's entropy $H(p)$. Notice that the result in Theorem 1. shows that the cross-entropy $H(p,\hat{p}) = S(C_\mathbb{P}) + \log(h)$, where $h$ is the normalization constant. However, $S(C_\mathbb{P})$ is independent of the normalization constant and does not require the kernel to be a density function to be computed. Additionally, notice that in Theorem 4 (Theorem 5 in the revised document), we establish the conditions under which $D_{JS}(\mathbb{P},\mathbb{Q}) = D_{JS}^\phi(C_\mathbb{P},C_\mathbb{Q})$, and notice, that in this case, the normalization constant is not involved in this equivalence. In fact, one of the main benefits of the covariance operator entropy is that it avoids the intermediate step of density estimation.

2. The RJSD formulation has several advantages. First, in GANs, the problem of mode collapse is usually attributed to using a classifier to measure the divergence since the generator can learn to create a few realistic samples to fool the discriminator and trivially minimize the loss. Besides, since the discriminator handles each example independently, there is a lack of coordination between the gradients, and consequently, the generator is not penalized by producing similar samples. On the contrary, our framework processes the covariance matrices of entire batches of real and generated data rather than individual samples. This leads to coordinated gradients to either increase or decrease the divergence between the distributions. Second, RJSD avoids entropy regularizers to avoid mode collapse. Specifically, RJSD explicitly matches the entropy of the generated samples and the training distribution. This feature penalizes generative distributions with lower entropy than the real distribution.

3. In section 2.2 we mention that the cost of the kernel-based estimator is $\mathcal{O}(N^3)$, and the covariance-based is $\mathcal{O}(D^3)$. The cost of the kernel-based estimator is not really prohibitive, since we are able to compute eigendecompositions of up to 20-30K samples easily on modern computers. However, the covariance-based estimator is independent of sample size, which makes this estimator even more appealing, since even for larger amounts of data, we can estimate this quantity faster. In future work, we aim to test the proposed estimators in large scale datasets.

We hope to have addressed your main concerns. If you have any further suggestions or questions, please feel free to let us know.

Thank you for your valuable feedback and insightful comments on our paper. We appreciate the time and effort you have dedicated to reviewing our work. Next, we address the main concerns raised.

Major

1. Issues about confusing notation. We appreciate your thorough review and feedback on the clarity of notation and terminology used in our paper. Your observations regarding the potential challenges for readers without a background in quantum mechanics/quantum information theory are valuable. In response to your concerns, we revised the manuscript to enhance the clarity of notation and provide more detailed explanations of relevant operations such as the tensor product $\otimes$, nuclear norm $\lVert \cdot \rVert_*$, Hilbert-Schmidt norm $\lVert \cdot \rVert_{HS}$, among others. Regarding the bra-ket notation, we changed it to a more familiar notation in Machine learning, using a more standard inner product instead avoiding any confusion among the readers. We also defined properly entropy and cross-entropy and added a footnote to define a density matrix.

2. Theorems 1 and 2 are important for establishing a connection between kernel-based entropy and Shannon's differential entropy. The convergence and consistency of RJSD are addressed later in the document where we leverage the results in Bach, (2022). We added the following: "Leveraging the convergence results in Bach, (2022)[Proposition 7] of the empirical estimator $S(\mathbf{K_X})$ to $S(C_{\mathbb{P}})$, we can show that $D_{\scriptscriptstyle JS}^{\kappa}(\mathbf{X},\mathbf{Y})$ converges to the population quantity $D_{\scriptscriptstyle JS}^\phi(C_{\mathbb{P}}, C_{\mathbb{Q}})$ at a rate $\mathcal{O}\left(\frac{1}{\sqrt{N}}\right)$, assuming $N=M$. Details of this rate are given in Appendix A.8. Additionally, a direct consequence of Theorem 2 is that, under the same assumptions of the theorem, $D_{\scriptscriptstyle JS}^{\kappa}(\mathbf{X},\mathbf{Y})$ converges to $D_{\scriptscriptstyle JS}(\mathbb{P},\mathbb{Q})$ as $N\rightarrow \infty$ with probability one." We acknowledge that this might not have been expressed clearly, and we enhanced the clarity of our paper by explicitly stating these two different results.

3. This concern is related to the previous point. We acknowledge the potential for misunderstanding and appreciate the reviewer's feedback. This has been addressed in the revised manuscript.

4. Having shown the supporting theoretical results of the proposed RJSD estimator (Points 2 and 3), we appreciate your feedback regarding the clarity of the performance improvement over other metrics. In response, we would like to highlight Theorem 5 (Theorem 6 in the revised manuscript) in our paper, where we establish that RJSD serves as an upper bound on MMD. This result implies that if MMD can effectively distinguish between two distributions, then RJSD is expected to have similar or superior discriminative power. We argue in the paper that this theoretical relationship suggests that RJSD should be at least as good as MMD in identifying discrepancies between distributions. While it is true that the observed performance gain may not bet substantial, deep kernel MMD is a state-of-the-art approach. In this context, even marginal improvements in performance on the evaluated datasets are noteworthy. To address your concern, in the revised paper, we elaborate further on why RJSD has a superior test power. We agree that these additional explanation will convey more clearly the advantages offered by our proposed method.

5. We appreciate your attention to the presentation of the test power results in Section 5.3 and your concern about the significance levels. To address this, we would like to clarify that all tests were conducted with a significance level of $\alpha = 0.05$. Detailed information about the testing procedure was provided Appendix C.3. Additionally, we included figures (3 to 6 in the appendix) displaying the standard deviations of test power for all experiments, offering a more comprehensive view of the test performance. Furthermore, we have made the code used to generate these results available on an anonymous GitHub repository to ensure transparency and reproducibility. This allows interested readers to replicate our experiments and verify the reported results. We believe these additional details and resources strengthen the robustness of our findings.

Thank you for your valuable feedback and insightful comments on our paper. We appreciate the time and effort you have dedicated to reviewing our work.

• "The paper would have been much stronger if the performance improvement over MMD had been shown more clearly. As it stands, there is a rather small performance gain, and it is not completely clear why." (Q3) We appreciate your feedback regarding the clarity of the performance improvement over MMD in our paper. We agree that providing a more explicit demonstration of this improvement is crucial and it would make the article stronger. In response, we would like to highlight Theorem 5 in our paper, where we establish that RJSD serves as an upper bound on MMD. This result implies that if MMD can effectively distinguish between two distributions, then RJSD can achieve at least similar discriminative power. We argue in the paper that this theoretical relationship suggests that RJSD should be at least as good as MMD in identifying discrepancies between distributions. While it is true that the observed performance gain may not be substantial, deep kernel MMD is a state-of-the-art approach. In this context, even marginal improvements in performance on the evaluated datasets are noteworthy. To address your concern, in the revised document, we elaborate further on why RJSD is expected to have a superior test power. We agree that these additional explanation will convey more clearly the advantages offered by our proposed method.

• (Q1) The condition $\kappa(x,x) = 1$ ensures that the covariance operator will have unit trace. Note that from $\kappa(x,x)=1$, we obtain that $\lVert \phi(x)\rVert = 1$. Therefore, $\text{Tr}(\phi(x)\otimes \phi(x)) = \lVert \phi(x) \rVert^2$ for all $x \in \mathcal{X}$. On the other hand, in the results connecting kernel-based entropy with Shannon's differential entropy, as $\gamma \rightarrow \infty$, $\kappa(x,x)$ is still one, since for the Gaussian kernel $\kappa(x,x) = \exp(-\gamma\lVert x - x \rVert^2) = \exp(-\gamma \cdot 0) = 1$ regardless the value of $\gamma$. To avoid confusion, we explicitly define the Gaussian kernel as above in the revised manuscript.

• (Q2) We would like to clarify that our generative formulation is indeed a generative adversarial network (GAN). The formulation (16) is an adversarial game (min max problem), between the generator $G_\theta$ trying to maximize RJSD and the discriminator $\phi_\omega \circ f_\omega$ aiming to minimize RJSD. We mention, however, that we differ from the standard GAN formulation because instead of using a classifier as the discriminator, we use a Fourier Features network $\phi_\omega \circ f_\omega$ as the discriminator. Notice that in vanilla GANs, the loss function implicitly minimizes the Jensen-Shannon divergence; however, we achieve a similar goal by explicitly minimizing RJSD, without using a classifier.

We believe that the proposed enhancements will improve the clarity and quality of the paper. If you have any further suggestions or questions, please feel free to let us know.

Thank you for your valuable feedback and insightful comments on our paper. We appreciate the time and effort you have dedicated to reviewing our work.

• "A hard question seems to be why RJSD outperforms the "Learning Deep Kernels for Non-Parametric Two-Sample Test" experimentally or establishing the relationship between its power and power of MMD-like algorithms." One important result in our paper is given in Theorem 5 (Theorem 6 in the revised manuscript). There, we show that RJSD is an upper bound on MMD, particularly, $D_{JS}^\phi(C_{\mathbb{P}},C_{\mathbb{Q}})\geq \frac{1}{8}\text{MMD}_{\kappa^2}(\mathbb{P},\mathbb{Q})$. This result implies that if MMD can effectively distinguish between two distributions, then RJSD can achieve at least similar discriminative power. We argue in the paper that this theoretical relationship suggests that RJSD should be at least as good as MMD in identifying discrepancies between distributions. In the revised manuscript, we elaborate further on this theoretical property.

We appreciate your suggestions of comparing our method against "Interpretable Distribution Features with Maximum Testing Power," and we will intend to perform additional experiments in the future work.

Thank you for your valuable feedback and insightful comments on our paper. We appreciate the time and effort you have dedicated to reviewing our work. Next, we address the main concerns expressed in the revision.

• "However the authors do not present any statistical properties of the estimator". We would like to clarify that we do introduce statistical properties of our estimator. The estimator of RJSD is based on the kernel-based estimator of Shannon’s entropy whose statistical properties are presented in Theorems 1 and 2 in the original paper. Specifically, we show in Theorem 1, that our entropy estimator can be related to a plug-in Parzen density estimator, with convergence and consistency extensively studied (Dimitriev & Tarasenko ,(1974)) as we point out in the document. Moreover, we show in Theorem 2 a similar convergence result when using Fourier Features. Therefore, the statistical properties of RJSD are a direct consequence of these results. Notice, that these results apply, for both estimators, using Kernel matrices or the explicit covariance matrices. We acknowledge the potential for misunderstanding and appreciate the reviewer's feedback. To address this concern, in the updated draft we incorporated an additional theorem (New Theorem 2) explicitly outlining the convergence of the kernel-based and covariance-based entropy estimators to Shannon's differential entropy. Additionally, we added a paragraph after Eqn. 13 to discuss the convergence of the kernel-based RJSD estimator to the population quantity.

• "The RFF based estimator is perhaps discussed too soon. I would have liked to read the properties and performance of the kernel matrix based estimator when $O(N^3)$ is not really prohibitive." We appreciate your suggestion to elaborate further the properties of the kernel matrix-based estimator. It's important to note that the introduction of the Fourier features approach is not solely aimed at reducing computational cost by approximating the kernel-based estimator. Instead, the Fourier features serve a dual purpose in our methodology. Firstly, they facilitate the construction of empirical covariance matrices, crucial for estimating RJSD. Secondly, these features enable a parameterization of the representation space, leading to a kernel-learning strategy. We treat the Fourier features as learnable parameters within a neural network (Fourier Feature network), optimizing them to maximize divergence and enhance the discriminatory power between sets of samples. Consequently, the Fourier features approach offers a more versatile estimator that extends beyond addressing computational cost. In the revised manuscript, we provide a clearer exposition of these two functionalities to enhance understanding. Also, we changed the order of the RJSD estimators, presenting first the kernel-based estimator as you suggested.

• "The authors say the estimator is "differentiable" in the introduction. Is this for the RFF estimator or the kernel matrix estimator. The differentiability for the kernel based estimator is not clear, and would be good to clarify." We would like to clarify that when we mention the estimator as 'differentiable,' we are referring to both kernel-based estimator and the covariance-based estimator using Fourier Features. The kernel-based estimator is indeed differentiable, and this was introduced in Sanchez Giraldo & Principe (2013); Yu et al. (2021). On the other hand, Sriperumbudur & Szabó (2015) show that the Fourier Features are twice differentiable and therefore the kernel or matrix produced by this approach. Therefore, the result in Sanchez Giraldo & Principe (2013); Yu et al. (2021) holds for this estimator as well. We acknowledge that this may have been unclear in our original presentation. In the revised manuscript, we provided additional clarification on the differentiability of the kernel-based estimator, ensuring a more explicit and transparent discussion.

• "The writing has a lot of room for improvement. The following notational issues bothered (and confused) me while reading this paper." Regarding the notational issues, we appreciate the thorough revision of the notation inconsistencies or confusing symbols, and we addressed all the concerns raised in the revised manuscript. Next, we reply individually to some of them that require further clarification.