Mutual Information Estimation via Normalizing Flows

We thank Reviewer 2t4W for the work! We are glad to receive helpful criticism of our article. We further provide answers to the main points raised in the review.

Weaknesses:

1. It is true that some parts of our work are dense with theoretical results. As our manuscript has reached the page limit, we hope to utilize the additional content page in the final version of the article to add more complementary explanation and discussion.

2. We agree that our article would benefit from the proposed amendments. Some examples of density plots will be added to the article; they are also available in the supplementary PDF (Figure 5); please, see the global reply. An additional comparison with MINE and DINE-Gaussian has been also conducted; please, refer to the same PDF, Figure 6. In our setup, DINE-Gaussian did not yield reasonable results in the low-dimensional case, severely underestimating the MI.

3. We believe that the overview provided in the introduction is mostly sufficient to familiarize the reader with the related work. Nevertheless, we will additionally elaborate on this topic to explain the relation of our method to the previous estimators in more details.

Detailed comments:

1. Unfortunately, the absolute values are required for the following reason. In general, $D_\text{KL}(p_{\xi,\eta}||q_{\xi,\eta})$ and $D_\text{KL}(p_\xi \otimes p_\eta||q_\xi \otimes q_\eta)$ are incomparable. For example, if $\xi=\eta$ under $p$ and $q$, then $D_\text{KL}(p_{\xi,\eta}||q_{\xi,\eta}) = D_\text{KL}(p_\xi||q_\xi) < 2 \cdot D_\text{KL}(p_\xi||q_\xi) = D_\text{KL}(p_\xi \otimes \eta||q_\xi \otimes q_\eta)$. Conversely, if $p_\xi \otimes p_\eta = q_\xi \otimes q_\eta$ but $p_{\xi,\eta} \neq q_{\xi,\eta}$, we have $D_\text{KL}(p_{\xi,\eta}||q_{\xi,\eta}) > 0 = D_\text{KL}(p_\xi \otimes \eta||q_\xi \otimes q_\eta)$. Thus, the proposed bound is tight in both ways.

2. Thank you, the typo will be corrected.

3. Thank you again, the text will be changed.

4. Likelihood maximization is equivalent to $D_\text{KL}$ minimization (see [31]), which tightens the bound in Corollary 3.2 and Corollary 4.3. We will add this clarification to the manuscript.

5. To derive the bounds in question, one also have to utilize the monotonicity property of the KLD ($D_\text{KL}(p_{X,Y}||q_{X,Y}) \geq D_\text{KL}(p_X||q_X)$). This is reflected in our proofs of the bounds. Please, refer to lines 780-781 (Appendix, proof of Corollary 3.2).

6. Thank your for the suggestion! The pictures will be added to the article. We will also add them to the supplementary PDF of the general reply.

   Please, also note that we do not call synthetic images "incompressible"; this term is used to describe the datasets for experiments in Figure 1. These datasets have no low-dimensional latent structure.

7. It is true that adding more examples of applications will make the motivation of our work more clear. This will be done in the next revision.

We again sincerely thank Reviewer 2t4W for carefully reading our article and providing us with useful comments and helpful criticism. We will be glad to address any further concerns if they arise.

Dear Reviewer 2t4W,

Once again, thank you very much for your detailed review and the time you spent. As the end of the discussion period is approaching, we would like to ask if the concerns you raised have been addressed. We hope that you find our responses useful and would love to engage with you further if there are any remaining questions.

We understand that the discussion period is short, and we sincerely appreciate your time and help!

[a] Kingma, D. P., Salimans, T., Jozefowicz, R., Chen, X., Sutskever, I., and Welling, M. Improved variational inference with inverse autoregressive flow. In Advances in neural information processing systems, pp. 4743–4751, 2016.

[b] Atanov, Andrei, Alexandra Volokhova, Arsenii Ashukha, Ivan Sosnovik and Dmitry P. Vetrov. Semi-Conditional Normalizing Flows for Semi-Supervised Learning. ArXiv abs/1905.00505, 2019.

[c] Z. Goldfeld, K. Greenewald, J. Niles-Weed, and Y. Polyanskiy. Convergence of smoothed empirical measures with applications to entropy estimation. IEEE Transactions on Information Theory, 66(7):4368–4391, 2020.

[d] David McAllester and Karl Stratos. Formal limitations on the measurement of mutual information. Proceedings of the Twenty Third International Conference on Artificial Intelligence and Statistics, 2020

Dear Reviewer m9Tu, thank you sincerely for your profound and comprehensive review! In the following text, we provide responses to your questions. We hope that all your concerns are properly addressed.

Weaknesses: for 1 and 2 please refer to the questions section.

3. In order to assess the quality of our method and compare it to other approaches, synthetic datasets had to be utilized, as they are the only option with available ground truth MI. To improve the tests, we tried to reproduce the main challenges of realistic data: high dimensionality and latent structure. It is, however, true that our work would still benefit from experiments on realistic data.

   To partially address this issue, we employ our method to conduct a nonlinear canonical correlation analysis of the handwritten digits dataset. We describe the experiment and provide the results in the general reply. Please, note that this experiment was conducted under the time constraints of the rebuttal process.

Questions:

1. We believe that having three models instead of one should not be a significant problem if such approach allows for less learnable parameters to be used. In contrast to a conventional approach, where three normalizing flows are used to estimate $h(X)$, $h(Y)$ and $h(X,Y)$ separately, our refined method should be as expensive as learning only two flows for modeling $p_X$ and $p_Y$ (as we employ a Cartesian product of two flows).

   The general method additionally requires learning (possibly complex) $q$, but this approach still should not be as expensive as training a whole separate flow to model the joint distribution $p_{X,Y}$ from scratch. That is because the joint distribution should be already partially disentangled and simplified by $f_X \times f_Y$.

   In [22], several clever tricks are employed to learn only one model. More specifically, manipulating the conditioning parameter $c$ in MINDE-C or diffusion speed modulators $\alpha,\beta$ in MINDE-J allows for using a unique score network to model all the required distributions. One is able to apply similar tricks to normalizing flows by using autoregressive flows [a,b], but such approach is incompatible with our idea, as it can not be implemented via a Cartesian product of two flows. However, this amortization technique can be applied to $q$.

   We also note that all three models are learned in a single training loop with a shared loss function. Thus, this is basically a single model consisting of three blocks.

   Finally, the complexity of $q$ can be gradually increased (e.g., more components are added to a Gaussian mixture) to ensure that the most simple model is used. This procedure should be fairly cheap, as it does not require retraining $f_X$ and $f_Y$ every time the complexity of $q$ is increased. This can be interpreted as a gradual unfreezing of some parts of the model.

   Summary: we believe that our method is as expensive as modeling just $p_{X,Y}$ in the worst case scenario. The same should be true for both methods from [22].

2. In this work, we decided to focus on the refined approach for several reasons: mathematical elegance, better error bounds and relatively decent performance. Nevertheless, we also consider the question of sample complexity of the general method to be very important and interesting. Although we admit that a proper investigation of this topic requires conducting an additional set of experiments, we still would like to provide some related discussion.

   It is indeed known that any consistent MI estimator requires an exponentially (in dimension) large number of samples [b,c]. This fact is related to the high sample complexity of the PDF estimation task. One may expect $\mathcal{N}$-MIENF to be less prone to the curse of dimensionality, as it is a fairly restricted model. This, of course, comes at a cost of less expressive power. On the other hand, increasing the complexity of $q$ in the general method makes the model less regularized and more expressive, which increases the effects of the curse of dimensionality, but also improves the quality of the estimate for large datasets. Thus, one can expect the sample complexity to increase as a consequence of using more expressive $q$.

   Finally, the sample complexity also depends on $f_X$ and $f_Y$. More expressive flows may require large numbers of samples to avoid overfitting. In the case of small, but high-dimensional datasets NN-based MI estimators tend to "memorize" pairs of samples from $p_{X,Y}$, yielding extremely high MI estimates. We demonstrate this behaviour in the supplementary PDF, see the global reply.

   Summary: using less expressive $q$, $f_X$ and $f_Y$ may serve as a regularization, thus partially decreasing the sample complexity at a cost of lower accuracy for large datasets.

3. Due to the restrictive nature of normalizing flows, they are relatively expensive and less stable to train. They are also inferior to diffusion models in terms of generative capabilities. We believe that all these nuances may make it harder to apply our method to real data of very high dimensionality and complex structure. We consider combining our binormalization method with diffusion models to alleviate these difficulties.

4. Yes, our method is perfectly suitable for PMI estimation. Moreover, it is even possible to calculate PMI via a closed-form expression while using less restricted $q$ (e.g. Gaussian mixture). This allows us to achieve both consistency and simplicity of the PMI estimator.

We hope that we were able to answer the raised concerns. These answers will be incorporated into the manuscript. We again sincerely thank Reviewer m9Tu for carefully reading our article and providing us with useful comments, helpful criticism and interesting questions! We will be glad to answer any further questions if they arise.

We would like to deeply thank Reviewer 4miK for reading the article and providing us with a profound review! We further provide answers to the main points raised in the review.

Weaknesses:

1. We wanted to stress out that the biggest possible gap between the lower bound and the true value of MI is indeed achievable. Thus, the bound in Corollary 4.4 is tight. The same example from Remark 4.5 also shows that it is insufficient to marginally Gaussianize a pair random vectors, as the MI estimation error may become arbitrary high. We will explain this remark more thoroughly in the next revision.

2. This appears to be a small mistake in the legend, as the results are provided for tridiag-$\mathcal{N}$-MIENF. Choosing the tridiagonal version over the basic was indeed motivated by the near-identical performance (Figure 1), as well as the mathematical equivalence of the methods (see Statement 4.13). To avoid clutter, we plot only one method in Figure 2, as the results of both methods on synthetic images are also very close to each other. The label will be changed to represent the right method.

3. In order to assess the quality of our method and compare it to other approaches, synthetic datasets had to be utilized, as they are the only option with available ground truth MI. To improve the tests, we tried to reproduce the main challenges of realistic data: high dimensionality and latent structure.

   It is true that our work would benefit from experiments in the setting of explainable AI and an information-theoretic analysis of DNNs. However, we decided to leave these experiments for the future, and focus on refining and evaluating the method. We believe that the resulting estimator is still a notable contribution to the field.

   Nevertheless, to partially address this issue, we employ our method to conduct a nonlinear canonical correlation analysis of the handwritten digits dataset. We describe the experiment and provide the results in the general reply. Please, note that this experiment was conducted under the time constraints of the rebuttal process.

Questions:

We do not require $f_X$ and $f_Y$ to be block-diagonal. Instead, we only utilize the fact that $f=f_X \times f_Y$ has a block-diagonal Jacobian, as it is a Cartesian product of two mappings. This is just a small simplification, which is not of great importance.

Limitations:

This is a valid point, which we will address in the next revision. The quick answer is that exact likelihood calculation is not required in our method. However, optimization of the exact likelihood (instead of a lower bound) should tighten the bound from Corollary 3.2 better, which in turn should yield better MI estimates. We also consider extending our idea to non-likelihood based methods as a part of the future work. This discussion will be added to the limitations section.

We again sincerely thank Reviewer 4miK for carefully reading our article and providing us with helpful comments. We will be glad to answer any further questions if they arise.

Dear Reviewer cdSA, thank you for reading and reviewing our article carefully! In the following text, we provide responses to your questions. We hope that all the concerns are properly addressed.

Weaknesses:

1. We agree that experiments with real datasets will further improve our work and complement the current results acquired on synthetic data. Unfortunately, such data does not allow for a direct assessment of the quality of the MI estimators. That is why we focus on synthetic tests, while still trying to reproduce the main challenges of realistic data: high dimensionality and latent structure.

   To partially address this issue, we employ our method to conduct a nonlinear canonical correlation analysis of the handwritten digits dataset. We describe the experiment and provide the results in the general reply. Please, note that this experiment was conducted under the time constraints of the rebuttal process.

2. Two diagrams of both MI estimatiors will be added to the article, as well as some examples of the synthetic data. We have also added these figures to the global reply PDF.

Questions:

1. Note that we store $w_j$ and compute $\hat I = \sum_j e^w_j$ and $\rho_j = \sqrt{1 - \exp(- 2 \cdot e^{w_j})}$. As $-2 \cdot e^{w_j} < 0$, $\exp(-2\cdot e^{w_j}) \in (0;1)$, which partially secures the method from numerical instabilities. Some other issues might arise from the gradient calculation and backpropagation, but we did not observe them. Moreover, the training process for tridiag-$\mathcal{N}$-MIENF was relatively stable, with the loss function plot being quite smooth and almost monotonic.

   One can also use other smooth mappings from $\mathbb{R}$ to $(0;+\infty)$ to parametrize $\rho_j$. For example, if the softplus function is employed, $\hat I = \sum_j \log(1 + e^{w_j})$ and $\rho_j = \sqrt{1 - 1/\sqrt{1 + e^{w_j}}}$, which eliminates one of the exponential. We will also reflect this approach in the final revision of the work.

2. The typo will be corrected, thank you!

We hope that we were able to answer all the raised concerns. We again thank Reviewer cdSA for the work! If there are any questions left, we will be glad to answer them.