Approximating mutual information of high-dimensional variables using learned representations

Summary

This paper explores the idea that underlying low-dimensional structure in high-dimensional data can be exploited to approximate mutual information (MI) efficiently and with a reasonable number of samples. The approach learns a low-dimensional embedding of high-dimensional data using a neural network architecture similar to an autoencoder. The MI of the resulting low-dimensional embedding is estimated using a nearest-neighbor approximation. The paper includes extensive experimental evaluation including synthetic Gaussian data, resampled non-Gaussian data, and two examples in the computational biology domain.

Strengths

This is a very nicely-written paper. The problem addressed is important and the ideas and approach are interesting. While the idea is relatively straightforward, the authors provide extensive experimentation to convince the reader that their approach is useful. This reviewer particularly appreciates the inclusion of domain-specific open problems as evidence of the efficacy of the proposed approach.

Weaknesses

The synthetic multivariate Gaussian data evaluation is lacking in some respects. The authors consider a only very specific form of dependence, namely linear dependence that obeys the stochastic process discussed in Sec. 3.1 (second paragraph) and that has low intrinsic dimensionality. Because Gaussians model only linear dependence it is likely that MI estimation based on linear projections (such as Sliced MI) would likely perform well in this setting. Yet the experiments fail to provide comparison to Sliced MI or related methods. A more thorough set of synthetic data experiments, including those that exploit low-dimensional structure, is proposed in (Paweł et al. 2024). The motivation and justification for the proposed approach is a little misleading. In particular, the authors misstate the findings in McAllester and Stratos (2018); this reviewer suspects the intended reference is McAllester and Stratos (2020). In that paper the authors do not demonstrate that statistical efficiency of estimators (MINE and InfoNCE) depend on dimension of the random variables, but rather they show a strong dependence on the value of MI. In particular they show that sample complexity scales exponentially with the value of MI. This finding is not consistent with the claim (L52) "_More generally, it has been shown that no MI estimator can be accurate without making strong assumptions about the distribution..._". In fact the estimators referred to are consistent in the infinite sample limit, but they can require a prohibitive number of samples for accurate estimates. A secondary motivation that is questionable is that the proposed method (L62) makes "_strong, yet reasonable, assumptions about data which enable tractable MI estimation._" In fact the MI of the latent embedding is not tractable. Indeed, the authors use a nonparametric estimator (KSG) to approximate the latent state MI. **Detailed Comments** * Sec. 3.1 : The stochastic process described seems overly complicated. Isn't it equivalent to restrict the covariance matrix for a multivariate Gaussian to a known structure? * Fig. 3c : This figure is not very informative due to the failure of the majority of methods ; perhaps consider a larger epsilon value? * L163 : Change "MIME" to "MINE" * Sec. 3.2 (last paragraph) : Change figure references from Fig. 3 to Fig. 4 * Sec. 4.2 : It isn't totally clear what $X_{6'}$ is referring to as it is not explicitly defined **References** Czyż, Paweł, et al. "Beyond normal: On the evaluation of mutual information estimators." _Advances in Neural Information Processing Systems_ 36 (2024). McAllester, David, and Karl Stratos. "Formal limitations on the measurement of mutual information." _International Conference on Artificial Intelligence and Statistics_. PMLR, 2020.

Questions

See "Weaknesses" section

Rating

7

Confidence

4

Soundness

4

Presentation

4

Contribution

3

Limitations

An obvious limitation of this work is not explicitly discussed. The latent dimension is introduced as a design variable that must be known (estimated) to compute the MI measure. Sensitivity to the choice of latent dimension is not directly addressed.

Summary

The focus of the paper is on approximating mutual information (MI) between multidimensional variables. This problem is challenging as the approximation of the MI suffers from the curse of dimensionality. The authors propose a method that approximates the MI via an embedding in a lower dimensional space. They test their methods on embedding from proteins language model, scRNA-seq data, and toy datasets.

Strengths

The approach is theoretically founded, and I appreciate that the authors also focus on making it practical. For instance, their model is intentionally simple to avoid doing large parameter sweep. The claim of the paper are well validated on toy datasets for which we can vary the dependence and ambient dimension, while having a close form solution of the MI.

Weaknesses

- I would advise reporting the standard deviation in all tables (error bar in figures). - The presentation of the results could be improved, for example a few figures have labels that are not readable (e.g. Fig.7).

Questions

- In Fig.2 , when you estimate 10 MIs, are the samples required to evaluate these MIs seen during training ? - In Fig.4 c), could you also include training time when training is required. - The results on toy datasets really highlight the benefits of the method, but not its limitations. For instance, it would be interesting to show results by varying the latent dimensions of the autoencoder. Especially, what happens if the latent dimension is smaller than $k$ ?

Rating

6

Confidence

3

Soundness

3

Presentation

2

Contribution

2

Limitations

The authors address the limitation of their work at the end of the manuscript.

Summary

This paper proposes latent MI (LMI), a method for estimating the mutual information (MI) between two high-dimensional multivariate random variables. For that, the technique uses the non-parametric MI estimator from [KSG04] on lower-dimensional latent representations that are learned by neural networks such that their MI is close to the one between the original variables. The paper provides some theoretical motivation for the proposed method. Finally, there is an experimental evaluation and comparison with other state-of-the-art methods for MI estimation, together with applications to problems in biology. [KSG04] Kraskov, A., Stögbauer, H., Grassberger, P. (2004). Estimating mutual information. Physical Review E—Statistical, Nonlinear, and Soft Matter Physics, 69(6), 066138.

Strengths

1) Pragmatic presentation of the problem setup and the proposed solution. 2) Interesting approach to MI estimation in (very) high dimensions, an open research subject, exploiting the informative low-dimensional structure of variables, which is a trendy approach (representation learning). 3) Some numerical illustrations of the proposed method, focused on the interpretability of the estimator, which is of paramount utility when it comes to applying mutual information to real-world problems. 4) Most limitations of this work are acknowledged by the authors.

Weaknesses

1) Theoretical justification for the method is rather simplistic, as acknowledged by the authors, and relies on potentially loose approximations such as the data processing inequality. 2) Overall, the proposed method consists in applying an existing estimator to pre-processed input variables, in the form of latent representations. It seems to lack a joint design, which results in two additive and independent sources of error (one from the representation, the other from the estimation itself). 4) Some minor concerns are raised in the 'Questions' field.

Questions

1) While the paper only apply the LMI estimator to problems in biology, have the authors considered other applications? For instance, the MI plays a central role in supervised learning in which the feature vector $X$ is high-dimensional (e.g. a MNIST image) while the label $Y$ consists in a few dimensions. Typically $X$ would be compressed but not $Y$, resulting in quite different latent spaces (different dimension and nature) unlike in the examples introduced in the paper. 2) Could the authors explain the choice to apply the KSG estimator to the latent representations in their LMI method? Have the authors considered combining their method to more modern methods such as MINE [BBR18]? 3) The last paragraph of Section 3.2 refers to "Fig. 3a, 3b" and "Fig 3c" while it should be to Fig. 4. Please make sure the figures are referenced properly throughout the paper. 4) Related to the previous point, please be consistent when referring to figures. It should be "Fig. \#" or "Figure \#" (preferably the former), but both are used. 5) Some relevant references (in my opinion) on MI estimation could be added in the introduction e.g., [NZH19], [GVG15], [MAK20]. [BBR18] Belghazi, M. I., Baratin, A., Rajeswar, S., Ozair, S., Bengio, Y., Courville, A., & Hjelm, R. D. (2018). Mine: mutual information neural estimation. arXiv preprint arXiv:1801.04062. [NZH19] Noshad, M., Zeng, Y., Hero, A. O. (2019, May). Scalable mutual information estimation using dependence graphs. In ICASSP 2019-2019 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP) (pp. 2962-2966). IEEE. [GVG15] Gao, S., Ver Steeg, G., Galstyan, A. (2015, February). Efficient estimation of mutual information for strongly dependent variables. In Artificial intelligence and statistics (pp. 277-286). PMLR. [MAK20] Mukherjee, S., Asnani, H., Kannan, S. (2020, August). CCMI: Classifier based conditional mutual information estimation. In Uncertainty in artificial intelligence (pp. 1083-1093). PMLR.

Rating

7

Confidence

4

Soundness

3

Presentation

3

Contribution

3

Limitations

The authors addressed most of the limitations of their work in the "Limitations" section of their paper. Please however see the 'Weaknesses' section.

Summary

This paper introduces an approach to approximate mutual information (MI), by applying a nonparametric MI estimator to the learned representations. The representations are learned by minimizing a weighted sum of the reconstruction loss and the prediction loss. The authors conducted experiments on both synthetic data and biological data. No higher ratings: The representation learning part has some technical flaws. No lower ratings: The evaluation methods and discussions provide new understandings and suggest some new applications.

Strengths

1. The evaluations are comprehensive with detailed discussions. The demonstrations suggest broader applications of MI in scientific domains. 2. The paper is well organized; presentations are clear and easy to follow.

Weaknesses

Despite the comprehensive evaluations and discussions, the proposed approach has a main technical issue: the learned representations might not capture the dependence between $X$ and $Y$, making the LMI fail. 1. Ideally, the reconstruction loss (autoencoder) alone can preserve the information of $X$, $Y$. However, it can be very inefficient when there are many redundancies in $X, Y$. A counter-example is when $X = (U, W), Y = (V, W)$, and $U$, $V$ have much more information than $W$; 2. The prediction loss highly depends on the $X$, $Y$; in the worst case, the prediction loss does not reveal any information about the dependence. A simple example is when the densities of $X$, $Y$ are symmetric. Two particular cases: (1) when $X$ and $Y$ are uniformly distributed on the unit circle; (2) when $X$ and $Y$ are generated from the mixture Gaussian: $ (\mathcal{N}(0, \begin{bmatrix}1, 1/2; 1/2, 1 \end{bmatrix} + \mathcal{N}(0, \begin{bmatrix}1, -1/2; -1/2, 1 \end{bmatrix})/2$. In both cases, the best predictor is simply zero due to the symmetry. 3. By combining the above examples, one can construct examples where the proposed loss does not provide informative information. 4. However, the above counterexamples, especially for prediction loss, can never appear for Gaussian data. This makes the evaluation of Gaussian synthetic data less convincing.

Questions

None.

Rating

5

Confidence

5

Soundness

2

Presentation

3

Contribution

2

Limitations

The authors have adequately acknowledged the limitations of the proposed method.

Thanks

Thank you for meticulously addressing my points. I have read through your rebuttal. I am convinced that this work deserves publication and am willing to champion the paper.

Thank you for the rebuttal!

I thank the authors for their detailed response and the newly conducted experiments. Specifically, I believe the new analysis in a supervised learning setup offers some insights for applying the proposed method to more general problems. Naturally, this requires a more in-depth study and could be interesting for future work. The analysis could however be added to the appendix, as well as the experiments presented in the response to all reviewers. This work is overall of good quality, at the intersection of exciting topics, and could pave the way to more profound studies.