InfoNet: An Efficient Feed-Forward Neural Estimator for Mutual Information

I thank the authors for their detailed and effective rebuttal. I am largely satisfied with your response. I'll keep my already positive score at this stage.

Below are some further comments:

• On scaling to higher dimensions with slicing. I am glad to see that the authors find the literature [2] helpful. Your new idea of utilizing this method to scale to high-dimensional case is indeed sensible and interesting, and the new experiments on higher-dimensional case and real-world data (e.g. CLIP) are helpful. However, as far as I am concerned, this new method might have two problems. First, it no more estimates the original mutual information, meaning that it may not be so useful in cases where the exact value of mutual information is of interests (for example, formally investigation of the trade-offs in data compression/privacy leakage). Second, the required number of slices may grow exponentially as the dimensionality increases [3]. That being said, I think your new method a very sensible extension, and all these above are only remarks rather than critiques.

• Discussion on flexible/scalable alternatives for measuring statistical dependence. As mentioned in the review, there also exist other parametric techniques that allow scalable/fast computation of statistical dependence. For example, the statistical dependence metrics in [1, 3, 4] can all be computed analytically (similar to solving linear regression/eigen-decomposition), which are also fast to execute compared to neural network training but are still much more powerful than non-parametric/kernel-based methods. Importantly, these methods are also fully differentiable like your method, so they can be applied to representation learning (see e.g. [3]). When I read your paper, I had a sense that why not using these already proven approaches, which are tested on high-dimensional cases (see e.g. [3]) and cases with diverse association patterns (see e.g. [1, 4]). I think it would be highly helpful to include a short discussion between your method and these methods, highlighting that your idea (pre-training) is orthogonal to these methods, and is conceptually more efficient and accurate. I believe this will help readers better understand the merit of your method.

Other comments: Algorithm 1 and Algorithm 2 seem duplicated?

[1] The Randomized Dependence Coefficient, NeurIPS 2017

[2] Sliced mutual information: A scalable measure of statistical dependence, NeurIPS 2021

[3] Scalable Infomin Learning, NeurIPS 2022

[4] Fairness-Aware Learning for Continuous Attributes and Treatments, ICML 2019

Dear Reviewer zWhA,

Thanks for your comprehensive review and the depth of your engagement with our manuscript, and the references you provide give us a lot of inspiration. Your acknowledgment of our method's significance and novelty within the field of mutual information is immensely encouraging. We are also grateful for your constructive critique, which has prompted an expansion of the method's scalability and experimental validation.

In response to your comments and suggestions, we have expanded our experimental scope and provided clarifications on the implementation details in our revised manuscript, which we believe can help reflect the potential and applicability of our proposed approach in a more accurate manner.

We hope that these clarifications can help finalize your assessment and the rating of our paper. Please also let us know if you have any further questions that we need to provide additional clarifications.

W1: ``Limitations (major): while very interesting, the proposed method seems only work for low-dimensional data. This is due to (a) the limitation of simulation-based learning itself, where as the dimensionality grows it becomes more and more difficult to cover diverse cases of distributions; (b) the number of grids in the lookup table grows exponentially with the dimensionality of data, making the discrete representation proposed impractical. That being said, the idea can expect to work well in e.g. 1D/2D cases;''

A: Thank you for your insight into the difficulty of simulation/sampling in high-dimensional space, which may increase the size of the look-up table exponentially.

In the following, we would like to provide an alternative for high-dimensional MI estimation using the pre-trained InfoNet with the sliced technique provided in [1], which projects a high-dimensional signal to a low dimension and allows efficient MI estimation.

We have rigorously evaluated our InfoNet model on high-dimensional random variables compared with MINE, in terms of correlation order accuracy and independence test.

Specifically, we validated the pretrained InfoNet in two tasks:

Task 1: given three jointly sampled random variables, X, Y, and Z from high-dimensional Gaussian, we want to know which one of Y and Z is more correlated with X, where the correct order, i.e., MI(X,Y)>=MI(X,Z) or MI(X,Y)<MI(X,Z) is obtained by the ground-truth MI computed with the analytical solution of MI for Gaussian. The performance metric is the binary classification accuracy.

Task 2: Independence test. Given a pair of random variables X, and Y, either independent or not, we apply InfoNet and MINE to check their independence since both metrics can claim independence between two random variables if the metrics are close to zero. We vary the epsilon as the threshold, and we plot the precision-recall curve to compare their performance using the area under the curve (AUC, the closer to 1, the better).

Results:

A. In terms of the correlation order test (Task 1), we get the following performance:

From the table, we can see that InfoNet performs better in terms of estimating the correlation orders between a triplet of random variables across multiple dimensions. Moreover, InfoNet runs at 0.05s per estimate, while MINE runs at 12.3s per estimate.

B. Results of AUC on independence testing (Task 2)

From the table, we can see that in both 16- and 128-dimensional space, InfoNet with the sliced technique can achieve better AUC for independence testing, which shows the potential of the proposed in high-dimensional applications even only trained on low-dimensional signals. Details of this experiment can be seen in B.4.1 of the revised version.

[1] Sliced mutual information: A scalable measure of statistical dependence, NeurIPS 2021

Dear Reviewer zWhA,

Thanks a lot for your positive response. We are grateful for your insightful comments and appreciation of our new experiments. We will incorporate them into our final version and handle the discussion among the mentioned references carefully and with clarity. Also, thanks for pointing out the duplication of the algorithms. We intentionally keep them here so the change is trackable, but will definitely merge them in the final version.

Also, in case it may cause curiosity on the number of projections for the sliced technique we have used, we provide them for each dimension below:

These numbers (of projections) are typical and can guarantee reasonable performance in our experiments. We also have an ablation on the performance improvement over different numbers of projections in Table 8. Moreover, due to the fact that our network can handle different projections in batch mode, the prediction efficiency is still maintained.

Thanks again for your valuable feedback！

Great, please do include the new results regarding high-dimensional data/real-world data and the new discussion regarding other scalable/efficient alternatives in the final version. Also thanks for the new results regarding #proj, which are instructive.

I will determine my final score after discussing with other reviewers and checking the final manuscript more carefully. As a general opinion, I think your work a quite fine contribution, and I will recommend acceptance.

Dear Reviewer 9kNo,

Thank you for your constructive comments and insightful feedback on our manuscript. We highly appreciate your recognition of the paper's clarity, and that the idea of using a feedforward neural network for mutual information estimation is intriguing. We acknowledge your concerns regarding the dimensionality of the data, the comparison with classical methods, and the generalizability of our approach. We have addressed each of your points in our revised manuscript and believe that our additional experiments and clarifications will shed more light on the efficacy and applicability of our proposed method. We hope that these clarifications can help finalize your assessment and the rating of our paper. Please also let us know if you have any further questions that we need to provide additional clarifications.

W1: The author fails to specify the dimension of the data $(x,y)$. If these are scalars, the task of mutual information estimation becomes considerably simpler. There exists an extensive body of literature on classical (non-neural network) methods addressing this, which warrants consideration.

A: Thank you for highlighting the point of specifying the data dimensionality in mutual information estimation.

Our work primarily focuses on scalar variables in classical mutual information (MI) estimation. However, it can be directly applied to estimating MI between high-dimensional random variables with the technique of sliced mutual information [1]. We have rigorously evaluated our InfoNet model on high-dimensional random variables, in terms of correlation order accuracy and independence test.

Specifically, we validated the proposed InfoNet in two tasks:

Task 1: given three jointly sampled random variables, X, Y, and Z from high-dimensional Gaussian, we want to know which one of Y and Z is more correlated with X, where the correct order, i.e., MI(X,Y)>=MI(X,Z) or MI(X,Y)<MI(X,Z) is obtained by the ground-truth mutual information (MI) computed with the analytical solution of MI for Gaussian. The performance metric is the binary classification accuracy.

Task 2: Independence test. Given a pair of random variables X, and Y, either independent or not, we apply InfoNet and MINE to check their independence since both metrics can claim independence between two random variables if the metrics are close to zero. We vary the epsilon as the threshold, and we plot the precision-recall curve to compare their performance using the area under the curve (AUC, the closer to 1, the better).

Results:

A. In terms of the correlation order test (Task 1), we get the following performance:

From the table, we can see that InfoNet performs better in terms of estimating the correlation orders between a triplet of random variables across multiple dimensions. Moreover, InfoNet runs at 0.05s per estimate, while MINE runs at 12.3s per estimate.

Details of this experiment can be seen in B.3.1 of the revised version.

B. In terms of the independence test, we have the AUCs of MINE-x (where x means the test-time gradient steps) and InfoNet obtained under different dimensions and sequence lengths (number for data points in each sequence):

From the Table, we can observe that the AUCs obtained with InfoNet are consistently higher than the ones obtained with MINE, showing that InfoNet provides a better estimation than MINE of mutual information despite a much shorter running time.

Detailed plots and results can be found in B.3.2 of the revised version.

Our findings, detailed in B.3 in the revised paper, demonstrate that InfoNet, leveraging sliced mutual information, exhibits robust performance at low time complexity in comparison to existing methods for estimating mutual information in high dimension.

[1] Goldfeld Z, Greenewald K. Sliced mutual information: A scalable measure of statistical dependence[J]. Advances in Neural Information Processing Systems, 2021, 34: 17567-17578.

W2: It is a good idea to compare the proposed MI estimator and MCC with ED in terms of the correlation of two random variables (similar to Figure 4 ) and computational speed given the fact that ED is quite fast to use in high dimensional data.

A: Thank you for bringing up the comparison with Energy Distance (ED). Indeed, ED is a noteworthy non-parametric measure known for its computational efficiency, especially in high-dimensional data scenarios, and it is a good measure in measuring the differences between two distributions.

We follow the suggestion and compare InfoNet (operated under 1000 random slices similar to [1]) with ED [2] in terms of computational efficiency and performance. Specifically, we consider two tasks.

Task 1: given three jointly sampled random variables, X, Y, and Z from high-dimensional Gaussian, we want to know which one of Y and Z is more correlated with X, where the correct order, i.e., MI(X,Y)>=MI(X,Z) or MI(X,Y)<MI(X,Z) is obtained by the ground-truth mutual information (MI) computed with the analytical solution of MI for Gaussian. The performance metric is the binary classification accuracy.

Task 2: Independence test. Given a pair of random variables X, and Y, either independent or not, we apply InfoNet and ED to check their independence since both metrics can claim independence between two random variables. We vary the epsilon as the threshold, and we plot the precision-recall curve to compare their performance using the area under the curve (AUC, the closer to 1, the better).

Results:

A. Efficiency: Time compare:

From the Table, we can see that the efficiency of InfoNet is not sensitive to the dimension of the random variable, and InfoNet is more efficient than ED at a higher dimension.

B. In terms of the correlation order test (Task 1), we get the following performance:

From the table, we can see that InfoNet performs better in terms of estimating the correlation orders between a triplet of random variables across multiple dimensions.

Details of this experiment can be seen in B.3.1 of the revised version.

C. In terms of the independence test, we have the AUCs of ED and InfoNet obtained under different dimensions and sequence lengths:

From the Table, we can observe that the AUCs obtained with InfoNet are consistently higher than the ones obtained with ED, showing that InfoNet provides a better metric than ED.

Detailed plots and results can be found in B.3.2 of the revised version.

[1] Goldfeld Z, Greenewald K. Sliced mutual information: A scalable measure of statistical dependence[J]. Advances in Neural Information Processing Systems, 2021, 34: 17567-17578.

[2] Rizzo M L, Székely G J. Energy distance[J]. wiley interdisciplinary reviews: Computational statistics, 2016, 8(1): 27-38.

Dear Reviewer FUUK,

Thank you for your insightful and constructive review of our manuscript. We greatly appreciate the time and effort you have devoted to evaluating our work, and your positive remarks on our proposed feed-forward neural network approach, InfoNet, using the attention-based mechanism for mutual information estimation, as well as the excellent presentation of the work.

We have taken a detailed look into your questions and comments, and have prepared comprehensive responses to each of your points, which we hope will address the issues effectively. We followed your suggestion to explore the comparison with other dependence measures like Energy Distance.

We hope that these clarifications can help finalize your assessment and the rating of our paper. Please also let us know if you have any further questions that we need to provide additional clarifications.

W1 Is speeding up MINE using the attention architecture the only difference between InfoNet and MINE? Fundamentally, these two algorithms are similar and the same objective function using the same Donsker & Varadhan representation of the KL divergence has been used for the MI estimation.

A: Thanks for your comment. Speeding up MINE with the attention architecture is not the only difference. To elaborate, we first provide a brief summary of MINE and InfoNet.

MINE: optimizes a neural network (NN) with the Donsker-Varadhan (DV) objective for a sequence {$x_t,y_t$} ($t=1,..., T$) sampled from a single distribution $p(x,y)$. When the optimization is done, the mutual information (MI) estimation of $p(x,y)$ is obtained. If one needs to estimate the MI of a different distribution $p’(x,y)$ using MINE, the neural network will be optimized again using the DV objective, but the sequence sampled from $p’(x,y)$. Since the optimal parameters of an NN ($\theta^*$) are different for different distributions.

InfoNet: optimizes a neural network (the proposed) also with the DV objective, but for many different distributions, i.e., sequences {$x_t^i,y_t^i$} ($t=1,..., T$,$i=1,..., N$) are sampled from N different distributions $p^i(x,y)$’s. When the optimization is done, it can be used to estimate the MI of an arbitrary distribution via the generalization capability of the proposed InfoNet architecture, without re-training. Equivalently, InfoNet, with a sequence as input, can directly output the optimal set of parameters ($\theta^*$) corresponding to the one obtained by performing a MINE optimization for a distribution (without finetuning InfoNet on this distribution).

With the above in mind, the key differences between MINE and InfoNet are listed below:

1. During training, MINE optimizes the NN for a single distribution, but InfoNet optimizes the proposed NN for many distributions.

2. During training, MINE treats the NN as a (scalar) discriminant whose optimum (defined by the DV formula and the sequence) gives the MI of the distribution represented by the input sequence. However, InfoNet acts as a hyper-network, and InfoNet outputs the optimal discriminant (NN in MINE) corresponding to an arbitrary distribution from the input sequence.

3. During testing, when a new distribution arrives, MINE needs to optimize the NN for estimating its distribution with the sequence sampled from this new distribution. In contrast, InfoNet takes the sampled sequence from this new distribution, and directly outputs the optimal NN (represented as a look-up table) as well as the MI estimation without performing any network update.

Thanks for your effort to include more clarification about your approach in comparison with MINE, and adding more experiments about ED. I have read the comments and the revised version and increased my score from 6 to 8 (There is no 7 in the scores :)).

Dear Reviewer C35n,

Thank you for your careful review and constructive questions. We also appreciate your acknowledgment of the importance of the problem we are tackling, and the efficiency of the InfoNet we are proposing.

We understand that there exists a misunderstanding, and the key differences between MINE and InfoNet need to be clarified. Accordingly, we have revised the paper with all the clarifications derived from these questions, which are answered in detail below.

We hope that these clarifications can help finalize your assessment and the rating of our paper. Please also let us know if you have any further questions that we need to provide additional clarifications.

W1: I am confused with some parts, might be due to writing or organization. Please see the following questions.

A: Thanks for all your questions below. We understand that there is confusion on some parts due to the lack of enough explanations on the key differences between MINE and InfoNet. Next, will first elaborate on the key differences, and then provide clarifications for each of the questions.

Q1: All in all, I think the confusion is due to the lack of enough explanations. Could you please elaborate on the difference between MINE and InfoNet? The same training objective is used. Both can parameterize the function via NNs, even the same NN architecture.

A: First, we provide a brief summary of MINE and InfoNet.

MINE: optimizes a neural network (NN) with the Donsker-Varadhan (DV) objective for a sequence {$x_t,y_t$} ($t=1,..., T$) sampled from a single distribution $p(x,y)$. When the optimization is done, the mutual information (MI) estimation of $p(x,y)$ is obtained. If one needs to estimate the MI of a different distribution $p’(x,y)$ using MINE, the neural network will be optimized again using the DV objective, but the sequence sampled from $p’(x,y)$. Since the optimal parameters of an NN ($\theta^*$) are different for different distributions.

InfoNet: optimizes a neural network (the proposed) also with the DV objective, but for many different distributions, i.e., sequences {$x_t^i,y_t^i$} ($t=1,..., T$,$i=1,..., N$) are sampled from N different distributions $p^i(x,y)$’s. When the optimization is done, it can be used to estimate the MI of an arbitrary distribution via the generalization capability of the proposed InfoNet architecture, without re-training. Equivalently, InfoNet, with a sequence as input, can directly output the optimal set of parameters ($\theta^*$) corresponding to the one obtained by performing a MINE optimization for a distribution (without finetuning InfoNet on this distribution).

With the above in mind, the key differences between MINE and InfoNet are listed below:

1. During training, MINE optimizes the NN for a single distribution, but InfoNet optimizes the proposed NN for many distributions.

2. During training, MINE treats the NN as a (scalar) discriminant whose optimum (defined by the DV formula and the sequence) gives the MI of the distribution represented by the input sequence. However, InfoNet acts as a hyper-network, and InfoNet outputs the optimal discriminant (NN in MINE) corresponding to an arbitrary distribution represented by the input sequence.

3. During testing, when a new distribution arrives, MINE needs to optimize the NN for estimating its MI with the sequence sampled from this new distribution. In contrast, InfoNet takes the sampled sequence from this new distribution, and directly outputs the optimal NN (represented as a look-up table) as well as the MI estimation without performing any network update.

Q2: 'Specifically, we want to explore whether the estimation of mutual information can be performed by a feed-forward prediction of a neural network' -- What does a feed-forward prediction mean? For MINE, we still use NNs to parameterize a function and output a scalar via NNs. Is MINE a feed-forward prediction? Please elaborate it.

A: Thanks for the question. MINE optimizes the NN, which outputs a scalar for a data point in the sequence and accumulates the scalars to get the MI estimate of the distribution represented by the sequence. However, MINE has to perform another optimization for a different distribution (or sequence), with which, we state that MINE is not feed-forward. In other words, MINE is feed-forward for predicting the scalar for a data point, but not feed-forward for estimating the MI of a distribution.

In contrast, after massive training on a spectrum of distributions, when a new distribution is presented to InfoNet in the form of a sequence, InfoNet predicts the optimal discriminant (without changing its parameters), which then computes the scalars and the MI estimate of the new distribution. So, InfoNet is feed-forward for computing the MI of a distribution, no need to perform training for a new distribution. Thus InfoNet is much more efficient than MINE as no gradient ascent is needed with a test distribution.