Mutual Information Estimation via $f$-Divergence and Data Derangements

Authors: nunzio a letizia
nicola  novello
andrea m tonello
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-14 22:17:49

Abstract: Estimating mutual information accurately is pivotal across diverse applications, from machine learning to communications and biology, enabling us to gain insights into the inner mechanisms of complex systems. Yet, dealing with high-dimensional data presents a formidable challenge, due to its size and the presence of intricate relationships. Recently proposed neural methods employing variational lower bounds on the mutual information have gained prominence. However, these approaches suffer from either high bias or high variance, as the sample size and the structure of the loss function directly influence the training process. In this paper, we propose a novel class of discriminative mutual information estimators based on the variational representation of the $f$-divergence. We investigate the impact of the permutation function used to obtain the marginal training samples and present a novel architectural solution based on derangements. The proposed estimator is flexible since it exhibits an excellent bias/variance trade-off. The comparison with state-of-the-art neural estimators, through extensive experimentation within established reference scenarios, shows that our approach offers higher accuracy and lower complexity.

Main paper: 
 introduction: 
 The mutual information (MI) between two multivariate random variables, X and Y , is a fundamental quantity in statistics, representation learning, information theory, communication engineering and biology [1, 2, 3, 4, 5] . It quantifies the statistical dependence between X and Y by measuring the amount of information obtained about X via the observation of Y , and it is defined asEQUATIONUnfortunately, computing I(X; Y ) is challenging since the joint probability density function p XY (x, y) and the marginals p X (x), p Y (y) are usually unknown, especially when dealing with high-dimensional data. Some recent techniques [6, 7] have demonstrated that neural networks can be leveraged as probability density function estimators and, more in general, are capable of modeling the data dependence. Discriminative approaches [8, 9] compare samples from both the joint and marginal distributions to directly compute the density ratio (or the log-density ratio) R(x, y) = p XY (x, y) p X (x)p Y (y) .(EQUATIONWe focus on discriminative MI estimation since it can in principle enjoy some of the properties of implicit generative models, which are able of directly generating data that belongs to the same distribution of the input data without any explicit density estimate. In this direction, the most successful technique is represented by generative adversarial networks (GANs) [10] . The adversarial training pushes the discriminator D(x) towards the optimum value D(x) = p data (x) p data (x) + p gen (x) = 1 1 + pgen(x) p data (x).(3)Therefore, the output of the optimum discriminator is itself a function of the density ratio p gen /p data , where p gen and p data are the distributions of the generated and the collected data, respectively.We generalize the observation of (3) and we propose a family of MI estimators based on the variational lower bound (VLB) of the f -divergence [11, 12] . In particular, we argue that the maximization of any f -divergence VLB can lead to a MI estimator with excellent bias/variance trade-off.Since we typically have access only to joint data points (x, y) ∼ p XY (x, y), another relevant practical aspect is the sampling strategy to obtain data from the product of marginals p X (x)p Y (y), for instance via a shuffling mechanism along N realizations of Y . We analyze the impact that the permutation has on the learning and training process and we propose a derangement training strategy that achieves high performance requiring Ω(N ) operations. Simulation results demonstrate that the proposed approach exhibits improved estimations in a multitude of scenarios.In brief, we can summarize our contributions over the state-of-the-art as follows:• For any f -divergence, we derive a training value function whose maximization leads to a given MI estimator. • We compare different f -divergences and comment on the resulting estimator properties and performance. • We study the impact of data derangement for the learning model and propose a novel derangement training strategy that overcomes the upper bound on the MI estimation [13] ,contrarily to what happens when using a random permutation strategy. • We unify the main discriminative estimators into a publicly available code which can be used to reproduce all the results of this paper. 
related work: 
 Traditional approaches for the MI estimation rely on binning, density and kernel estimation [14, 15] , k-nearest neighbors [16] , and ensemble-based models [17] . Nevertheless, they do not scale to problems involving high-dimensional data as it is the case in modern machine learning applications. Hence, deep neural networks have recently been leveraged to maximize VLBs on the MI [11, 18, 19] . The expressive power of neural networks has shown promising results in this direction although less is known about the effectiveness of such estimators [20] , especially since they suffer from either high bias or high variance.Discriminative approaches usually exploit an energy-based variational family of functions to provide a lower bound on the Kullback-Leibler (KL) divergence. As an example, the Donsker-Varadhan dual representation of the KL divergence [11, 21] produces an estimate of the MI using the bound optimized by the mutual information neural estimator (MINE) [19] . Another VLB based on the KL divergence dual representation introduced in [18] leads to the NWJ estimator (also referred to as f -MINE in [19] ). Both MINE and NWJ suffer from high-variance estimates and to combat such a limitation, the SMILE estimator was introduced in [20] , where the authors proved that the estimate of the partition function is the cause for high-variance in VLB estimators. SMILE is equivalent to MINE in the limit τ → +∞. The MI estimator based on contrastive predictive coding (CPC) [22] provides low variance estimates but it is upper bounded by log N , resulting in a biased estimator. Such upper bound, typical of contrastive learning objectives, has been recently analyzed in the context of skew-divergence estimators [23] .Another estimator based on a classification task is the neural joint entropy estimator (NJEE) proposed in [24] , which estimates the MI as entropies subtraction.Inspired by the f -GAN training objective [25] , in the following, we present a class of discriminative MI estimators based on the f -divergence measure. Conversely to what has been proposed so far in the literature, where f is always constrained to be the generator of the KL divergence, we allow for any choice of f . Different f functions will have different impact on the training and optimization sides, while on the estimation side, the partition function does not need to be computed, leading to low variance estimators.Question: Does the paper discuss the limitations of the work performed by the authors?Answer: [Yes] Justification: In Appendix D.1.4 we discuss a possible limitation of the proposed class of MI estimators. The computational efficiency is discussed in Sec. 6.2 and in Appendix D.1.2. Guidelines:• The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper. • The authors are encouraged to create a separate "Limitations" section in their paper.• The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be. • The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated. • The authors should reflect on the factors that influence the performance of the approach.For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon. • The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size. • If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness. • While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations. 
f -divergence mutual information estimation: 
 The calculation of the MI via a discriminative approach requires the density ratio (2) . From (3), we observe that I(X; Y ) can be estimated using the optimum GAN discriminator D when p data ≡ p X p Y and p gen ≡ p XY . More in general, the authors in [25] extended the variational divergence estimation framework presented in [18] and showed that any f -divergence can be used to train GANs. Inspired by such idea, we now argue that also discriminative MI estimators enjoy similar properties if the variational representation of f -divergence functionals D f (P ||Q) is adopted.In detail, let P and Q be absolutely continuous measures w.r.t. dx and assume they possess densities p and q, then the f -divergence is defined as followsEQUATIONwhere X is a compact domain and the function f : R + → R is convex, lower semicontinuous and satisfies f (1) = 0.In the following, we introduce f -DIME, a class of discriminative mutual information estimators (DIME) based on the variational representation of the f -divergence.Theorem 3.1. Let (X, Y ) ∼ p XY (x, y) be a pair of multivariate random variables. Let σ(•) be a permutation function such that p σ(YEQUATIONthen T (x, y) = arg maxEQUATIONandEQUATIONTheorem 3.1 shows that any value function J f of the form in (5) , seen as the dual representation of a given f -divergence D f , can be maximized to estimate the MI via (7) . It is interesting to notice that the proposed class of estimators does not need any evaluation of the partition term.We propose to parametrize T (x, y) with a deep neural network T θ of parameters θ and solve with gradient ascent and back-propagation to obtain θ = arg max θ J f (T θ ). By doing so, it is possible to guarantee that, at every training iteration n, the convergence of the f -DIME estimator În,fDIME (X; Y ) is controlled by the convergence of T towards the tight bound T while maximizing J f (T ), as stated in the following lemma.Lemma 3.2. Let the discriminator T (•) be with enough capacity, i.e., in the non parametric limit.Consider the problem T = arg maxEQUATIONwhere J f (T ) is defined as in (5) , and the update rule based on the gradient descent methodEQUATIONIf the gradient descent method converges to the global optimum T , the mutual information estimator defined in (7) converges to the real value of the mutual information I(X; Y ).The proof of Lemma 3.2, which is described in the Appendix, provides some theoretical grounding for the behaviour of MI estimators when the training does not converge to the optimal density ratio. Moreover, it also offers insights about the impact of different functions f on the numerical bias.It is important to remark the difference between the classical VLB estimators that follow a discriminative approach and the DIME-like estimators. They both achieve the goal through a discriminator network that outputs a function of the density ratio. However, the former models exploit the variational representation of the MI (or the KL) and, at the equilibrium, use the discriminator output directly in one of the value functions reported in Appendix B. The latter, instead, use the variational representation of any f -divergence to extract the density ratio estimate directly from the discriminator output.In Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof? Answer: [Yes] Justification: All the assumptions are included in the statements of the Lemmas, Theorems, and Corollaries (that are cross-referenced). All the complete proofs are reported in Appendix C. Guidelines:• The answer NA means that the paper does not include theoretical results.• All the theorems, formulas, and proofs in the paper should be numbered and crossreferenced. • All assumptions should be clearly stated or referenced in the statement of any theorems.• The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition. • Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material. • Theorems and Lemmas that the proof relies upon should be properly referenced. 
variance analysis: 
 In this section, we assume that the ground truth density ratio R(x, y) exists and corresponds to the density ratio in (2) . We also assume that the optimum discriminator T (x, y) is known and already obtained (e.g. via a neural network parametrization).We define p M XY (x, y) and p N X (x)p N Y (y) as the empirical distributions corresponding to M i.i.d. samples from the true joint distribution p XY and to N i.i.d. samples from the product of marginals p X p Y , respectively. The randomness of the sampling procedure and the batch sizes M, N influence the variance of variational MI estimators. In the following, we prove that under the previous assumptions, f -DIME exhibits better performance in terms of variance w.r.t. some variational estimators with a discriminative approach, e.g., MINE and NWJ.E p N X p N Y [ R]represents the major issue when dealing with variational MI estimators. Indeed, they comprise the evaluation of two terms (using the given density ratio), and the partition function is the one responsible for the variance growth. The authors in [20] characterized the variance of both MINE and NWJ estimators, in particular, they proved that the variance scales exponentially with the ground truth MI ∀M ∈ NEQUATIONwhereEQUATIONTo reduce the impact of the partition function on the variance, the authors of [20] also proposed to clip the density ratio between e -τ and e τ leading to an estimator (SMILE) with bounded partition variance. However, also the variance of the log-density ratio E p M XY [log R] influences the variance of the variational estimators, since it is clear that VarEQUATIONa result that holds for any type of MI estimator based on a VLB.The great advantage of f -DIME is to avoid the partition function estimation step, significantly reducing the variance of the estimator. Under the same initial assumptions, from (12) we can immediately conclude thatEQUATIONwhereEQUATIONis the Monte Carlo implementation of f -DIME. Hence, the f -DIME class of models has lower variance than any VLB based estimator (MINE, NWJ, SMILE, etc.). Furthermore, we provide in Appendix C two supplementary results. Lemma 4.1 introduces an upper bound on the variance of the f -DIME estimator, a result holding for any type of value function J f . Lemma 4.2, instead, characterizes the variance of the estimator in (14) when X and Y are correlated Gaussian random variables. We found out that the variance is finite and we use this result to verify in the experiments that the variance of f -DIME does not diverge for high values of MI.Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)? Answer: [Yes] Justification: The code is released. In addition, the complete neural architectures are specified in Appendix D.1, together with the optimizer type and hyperparameters. Furthermore, we provide the hardware details in Sec. 6. Finally, the contributions of this paper are the objective functions, which are all reported in Appendix A. Guidelines:• The answer NA means that the paper does not include experiments.• If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not. • If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable. • Depending on the contribution, reproducibility can be accomplished in various ways.For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed. • While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. , with an open-source dataset or instructions for how to construct the dataset). (d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility.In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results. 
derangement strategy: 
 The discriminative approach essentially compares expectations over both joint (x, y) ∼ p XY and marginal (x, y) ∼ p X p Y data points. Practically, we have access only to N realizations of the joint distribution p XY and to obtain N marginal samples of p X p Y from p XY a shuffling mechanism for the realizations of Y is typically deployed. A general result in [13] shows that failing to sample from the correct marginal distribution would lead to an upper bounded MI estimator.We study the structure that the permutation law σ(•) in Theorem 3.1 needs to have when numerically implemented. In particular, we now prove that a naive permutation over the realizations of Y results in an incorrect VLB of the f -divergence, causing the MI estimator to be bounded by log(N ), where N is the batch size. To solve this issue, we propose a derangement strategy.Let the data points (x, y) ∼ p XY be N pairs (x i , y i ), ∀i ∈ {1, . . . , N }. The naive permutation of y, denoted as π(y), leads to N new random pairs (x i , y j ), ∀i and j ∈ {1, • • • , N }. The idea is that a random naive permutation may lead to at least one pair (x k , y k ), with k ∈ {1, . . . , N }, which is actually a sample from the joint distribution. Viceversa, the derangement of y, denoted as σ(y), leads to N new random pairs (x i , y j ) such that i ̸ = j, ∀i and j ∈ {1, • • • , N }. Such pairs (x i , y j ), i ̸ = j can effectively be considered samples from p X (x)p Y (y). An example using these definitions is provided in Appendix D.1.3.The following lemma analyzes the relationship between the Monte Carlo approximations of the VLBs of the f -divergence J f in Theorem 3.1 using π(•) and σ(•) as permutation laws. Lemma 5.1. Let (x i , y i ), ∀i ∈ {1, . . . , N }, be N data points. Let J f (T ) be the value function in (5) . Let J π f (T ) and J σ f (T ) be numerical implementations of J f (T ) using a random permutation and a random derangement of y, respectively. Denote with K the number of points y k , with k ∈ {1, . . . , N }, in the same position after the permutation (i.e., the fixed points). ThenEQUATIONLemma 5.1 practically asserts that the value function J π f (T ) evaluated via a naive permutation of the data is not a valid VLB of the f -divergence, and thus, there is no guarantee on the optimality of the discriminator's output. An interesting mathematical connection can be obtained when studying J π f (T ) as a sort of variational skew-divergence estimator [23] , but this goes beyond the scope of this paper.The following theorem states that in the case of the KL divergence, the maximum of J π f (D) is attained for a value of the discriminator that is not exactly the density ratio (as it should be from (21) , see Appendix A). Theorem 5.2. Let the discriminator D(•) be with enough capacity. Let N be the batch size and f be the generator of the KL divergence. Let J π KL (D) be defined asEQUATION)Denote with K the number of indices in the same position after the permutation (i.e., the fixed points), and with R(x, y) the density ratio in (2). Then, Although Theorem 5.2 is stated for the KL divergence, it can be easily extended to any f -divergence using Theorem 3.1. Notice that if the number of indices in the same position K is equal to 0, we fall back into the derangement strategy and we retrieve the density ratio as output.EQUATION)When we parametrize D with a neural network, we perform multiple training iterations and so we have multiple batches of dimension N . This turns into an average analysis on K. We report in the Appendix (see Lemma 5.4) the proof that, on average, K is equal to 1.From the previous results, it follows immediately that the estimator obtained using a naive permutation strategy is biased and upper bounded by a function of the batch size N . Corollary 5.3 (Permutation bound). Let KL-DIME be the estimator obtained via iterative optimization of J π KL (D), using a batch of size N every training step. Then,EQUATIONWe report in Fig. 1 an example of the difference between the derangement and permutation strategies. The estimate attained by using the permutation mechanism, showed in Fig. 1b , demonstrates Theorem 5.2 and Corollary 5.3, as the upper bound corresponding to log(N ) (with N = 128) is clearly visible.Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material? Answer: [Yes] Justification: We provide the code and a readme file describing how to run it. The code does not require importing data, since it uses samples drawn from probability density functions for the experiments (except from the MNIST datasets, that are automatically downloaded using the scripts provided). Guidelines:• The answer NA means that paper does not include experiments requiring code. Justification: They are written in Appendix D.Guidelines:• The answer NA means that the paper does not include experiments.• The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them. • The full details can be provided either with the code, in appendix, or as supplemental material. 
experimental results: 
 In this section, we firstly describe the architectures of the proposed estimators. Then, we outline the data used to estimate the MI, comment on the performance of the discussed estimators in different scenarios, also analyzing their computational complexity. Finally, we present the outcomes of the self-consistency tests [20] over image datasets.To demonstrate the behavior of the state-of-the-art MI estimators, we consider multiple neural network architectures. The word architecture needs to be intended in a wide-sense, meaning that it represents the neural network architecture and its training strategy. In particular, additionally to the architectures joint [19] and separable [26] , we propose the architecture deranged.The joint architecture concatenates the samples x and y as input of a single neural network. Each training step requires N realizations (x i , y i ) drawn from p XY (x, y), for i ∈ {1, . . . , N } and N (N -1) samples (x i , y j ), ∀i, j ∈ {1, . . . , N }, with i ̸ = j.The separable architecture comprises two neural networks, the former fed in with N realizations of X, the latter with N realizations of Y . The inner product between the outputs of the two networks is exploited to obtain the MI estimate. The proposed deranged architecture feeds a neural network with the concatenation of the samples x and y, similarly to the joint architecture. However, the deranged one obtains the samples of p X (x)p Y (y) by performing a derangement of the realizations y in the batch sampled from p XY (x, y). Such diverse training strategy solves the main problem of the joint architecture: the difficult scalability to large batch sizes. For large values of N , the complexity of the joint architecture is Ω(N2 ), while the complexity of the deranged one is Ω(N ). NJEE utilizes a specific architecture, in the following referred to as ad hoc, comprising 2d -1 neural networks, where d is the dimension of X. I N JEE training procedure is supervised: the input of each neural network does not include the y samples. All the implementation details1 are reported in Appendix D.We benchmark the proposed class of MI estimators on two settings utilized in previous papers [11, 20] . In the first setting (called Gaussian), a multidimensional Gaussian distribution is sampled to obtain x and n samples, independently. Then, y is obtained as linear combination of x and n: y = ρ x + 1 -ρ 2 n, where ρ is the correlation coefficient. In the second setting (referred to as cubic), the nonlinear transformation y → y 3 is applied to the Gaussian samples. The true MI follows a staircase shape, where each step is a multiple of 2 nats. Each neural network is trained for 4k iterations for each stair step, with a batch size of 64 samples (N = 64). The tested estimators are: , in fact, the network's training is guided by the gradient computed using the Jensen-Shannon (JS) divergence (a linear transformation of the GAN divergence). Given the trained network, the clipped objective function proposed in [20] is only used to compute the MI estimate, since when ( 29) is used to train the network, the MI estimate diverges (see Fig. 7 in Appendix D). However, with the proposed class of f -DIME estimators we show that during the estimation phase the partition function (clipped in [20] ) is not necessary to obtain the MI estimate.I N JEE , I SM ILE (τ =We test our estimators over additional complex Gaussian data transformations (half-cube, asinh, and swiss roll mappings, Fig. 3 ) and non-Gaussian distributions (uniform and student distributions, Fig. 4 ) as suggested in [27] . The half-cube mapping is used to lengthen the tails of the Gaussian distributions.The inverse hyperbolic sine (asinh) mapping shortens the tails of the Gaussian distributions. These two transformations are applied to the same scenario of the Gaussian and cubic already present in our paper. The swiss roll mapping increases the dimensionality of the data distribution (from two to three dimensions) and it is usually used to test dimensionality reduction techniques. It considers two Gaussian random variables that are transformed into uniform random variables via the probability integral transform, the same pre-processing approach utilized in [28] to estimate the MI. The swiss roll mapping is applied to the X uniform random variable. The stairs plots are obtained by varying the correlation between the initial Gaussian distributions. The uniform case estimates the MI of the summation of two uniform random variables U (0, 1) and U (-ϵ, ϵ), where we vary the parameter ϵ, modifying the true MI. The student scenario analyzes the case of a multivariate student distribution with dispersion matrix chosen to be the identity matrix and degrees of freedom df . In this scenario, we vary df , implying a variation of the target MI. For the transformed Gaussian scenarios showed in Fig. 3 GAN-DIME attains the best performance in terms of low bias and variance. Among the non-Gaussian settings depicted in Fig. 4 , KL-DIME and GAN-DIME outperform the other methods, exhibiting low bias and exceptionally low variance. of I CP C strongly depends on N , significantly impacting its bias. Therefore, unless the batch size is considerably large, I CP C estimate is not reliable. I N JEE attains higher bias when N increases and, even more severely, when d decreases (see Fig. 2 ).A fundamental characteristic of each algorithm is the computational time. The computational time analysis is developed on a server with CPU "AMD Ryzen Threadripper 3960X 24-Core Processor" and GPU "MSI GeForce RTX 3090 Gaming X Trio 24G, 24GB GDDR6X". Before analyzing the time requirements to complete the 5-step MI staircases, we specify two different ways to implement the derangement of the y realizations in each batch. Random-based. The trivial way to achieve the derangement is to randomly shuffle the y elements of the batch until there are no fixed points (i.e., all the y realizations in the batch are assigned to a different position w.r.t. the starting location). Shift-based. Given N realizations (x i , y i ) drawn from p XY (x, y), for i ∈ {1, . . . , N }, we obtain the deranged samples as (x i , y (i+1)%N ), where "%" is the modulo operator. Although the MI estimates obtained by the two derangement methods are almost indistinguishable, all the results showed in the paper are achieved by using the random-based method. Additionally, we demonstrate the time efficiency of the shift-based approach. The time requirements to complete the 5-step staircase MI when varying the batch size N are reported in the left and center graphics of Fig. 5 . The influence of the MI estimator objective functions in the algorithm's time requirements is marginal, while the architecture type is the impactful component. As discussed in Sec. 6.1, the deranged strategy is remarkably faster than the joint one as N increases. More in general, the architectures deranged and separable are significantly faster w.r.t. the joint and NJEE ones, for a given batch size N and input distribution size d. The need of the separable architecture to use two neural networks implies that when N is significantly large, the deranged implementation is much faster than the separable one. The central graph in Fig. 5 illustrates a detailed representation of the time requirements of these two architectures to complete the 5-step stairs. As N increases, the gap between the time needed by the architectures deranged and separable grows, demonstrating that the former is the fastest. For example, when d = 20 and N = 30k, I GAN -DIM E needs about 55 minutes when using the architecture separable, but only 15 minutes when using the deranged one and less than 9 minutes for the shift-based deranged architecture. I N JEE is evaluated with its own architecture, which is the most computationally demanding, because it trains a number of neural networks equal to 2d -1. Thus, I N JEE can be utilized only in cases where the time availability is orders of magnitude higher than the other approaches considered. The time requirements to complete the 5-step staircase MI when varying the multivariate Gaussian distribution dimension d are reported in the right-side part of Fig. 5 . When d is large, the training of I N JEE fails due to memory requirement problems. For example, our hardware platform does not allow the usage of d > 30.To demonstrate the utility of f -DIME in non-Gaussian scenarios, we investigated the three selfconsistency tests developed by [20] over images datasets using all the estimators previously described, except I N JEE (for dimension constraints). The f -DIME estimators satisfy two out of the three tests, as discriminative approaches tend to be less precise when the MI is high, in accordance with [20] . We report the description of tests and results in Appendix D. 
conclusions: 
 In this paper, we presented f -DIME, a class of discriminative mutual information estimators based on the variational representation of the f -divergence. We proved that any valid choice of the function f leads to a low-variance MI estimator which can be parametrized by a neural network. We also proposed a derangement training strategy that efficiently samples from the product of marginal distributions. The performance of f -DIME is evaluated using three functions f , and it is compared with state-of-the-art estimators. Results demonstrate excellent bias/variance trade-off for different data dimensions and different training parameters.Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments? Guidelines:• The answer NA means that the paper does not include experiments.• The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper. • The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions). • The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.) • The assumptions made should be given (e.g., Normally distributed errors). • It should be clear whether the error bar is the standard deviation or the standard error of the mean. • It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified. • For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates). • If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text. 
experiments compute resources: 
 Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments? Answer: [Yes] Justification: In Sec. 6.2 we report the type of CPU and GPU used. In addition, we develop a full computational time analysis (in Sec. 6.2). Guidelines:• The answer NA means that the paper does not include experiments.• The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage. • The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute. • The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper). 
code of ethics: 
 Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines? 
 
 Appendix: In this section, we provide a concrete list of DIME estimators obtained using three different fdivergences. In particular, we maximize the value function defined in (5)J f (T ) = E (x,y)∼p XY (x,y) T x, y -f * T x, σ(y) ,over T or its transformation. By doing that, and using (7) ,I(X; Y ) = I f DIM E (X; Y ) = E (x,y)∼p XY (x,y) log f * ′ T (x, y) ,we obtain a list of three different MI estimators. The list is used for both commenting on the impact of the f function, referred to as the generator function, and for comparing the estimators discussed in Sec. 2.We consider the cases when f is the generator of: We report below the derived value functions and the mathematical expressions of the proposed estimators.The variational representation of the KL divergence [18] leads to the NWJ estimator in (28) when f (u) = u log(u). However, since we are interested in extracting the density ratio, we apply the transformation T (x) = log(D(x)). In this way, the lower bound introduced in (5) reads as followsEQUATIONwhich has to be maximized over positive discriminators D(•). As remarked before, we do not use J KL during the estimation, rather we define the KL-DIME estimator asEQUATIONdue to the fact thatEQUATIONFollowing a similar approach, it is possible to define f (u) = u log u -(u + 1) log(u + 1) + log 4 and T (x) = log(1 -D(x)). We derive from Theorem 3.1 the GAN-DIME estimator obtained via maximization ofEQUATIONIn fact, at the equilibrium we recover (3), henceEQUATIONThe last example we consider is the generator of the Hellinger distance squared f (u) = ( √ u -1) 2 with the change of variable T (x) = 1 -D(x). After simple manipulations, we obtain the associated value function asEQUATIONwhich is maximized forEQUATIONleading to the HD-DIME estimatorEQUATION)Given that these estimators comprise only one expectation over the joint samples, their variance has different properties compared to the variational ones requiring the partition term such as MINE and NWJ.In this section, we provide a detailed description of the formulas of the MI estimators we summarized in Sec. 2.The Donsker-Varadhan dual representation of the KL divergence [11, 21] produces an estimate of the MI using the bound optimized by the mutual information neural estimator (MINE) [19] EQUATIONwhere θ ∈ Θ parameterizes a family of functions T θ : X × Y → R through the use of a deep neural network. However, the logarithm before the expectation in the second term renders MINE a biased estimator. To avoid biased gradients, the authors in [19] suggested to replace the partition function E p X p Y [e T θ ] with an exponential moving average over mini-data-batches.Another VLB is based on the KL divergence dual representation introduced in [18] (also referred to as f -MINE in [19] )EQUATIONAlthough for a fixed T MINE provides a tighter bound I M IN E ≥ I N W J , the NWJ estimator is unbiased.Both MINE and NWJ suffer from high-variance estimations and to combat such a limitation, the SMILE estimator was introduced in [20] . It is defined asEQUATIONwhere clip(v, l, u) = max(min(v, u), l) and it helps to obtain smoother partition functions estimates. SMILE is equivalent to MINE in the limit τ → +∞.The MI estimator based on contrastive predictive coding (CPC) [22] is defined asEQUATION)where N is the batch size and p XY,N denotes the joint distribution of N i.i.d. random variables sampled from p XY . CPC provides low variance estimates but it is upper bounded by log N , resulting in a biased estimator.The neural joint entropy estimator (NJEE) proposed in [24] is based on a classification task. Let X m be the m-th component of X, with m ≤ d and N the batch size. X k is the vector containing the first k components of X. Let ĤN (X 1 ) be the estimated marginal entropy of the first components in X and let G θm (X m |X m-1 ) be a neural network classifier, where the input is X m-1 and the label used is X m . If CE(•) is the cross-entropy function, then the MI estimator based on NJEE is defined asEQUATIONwhere the first two terms of the RHS constitutes the NJEE entropy estimator. EQUATIONthen T (x, y) = arg maxEQUATIONandEQUATION)Proof. From the hypothesis, the value function can be rewritten asEQUATIONThe thesis follows immediately from Lemma 1 of [18] . Indeed, the f -divergence D f can be expressed in terms of its lower bound via Fenchel convex dualityEQUATIONwhere T : X → R and f * is the Fenchel conjugate of f defined asEQUATIONTherein, it was shown that the bound in (36) is tight for optimal values of T (x) and it takes the following formEQUATIONwhere f ′ is the derivative of f .The MI I(X; Y ) admits the KL divergence representationEQUATIONand since the inverse of the derivative of f is the derivative of the conjugate f * , the density ratio can be rewritten in terms of the optimum discriminator TEQUATIONf -DIME finally reads as followsEQUATIONC.2 Proof of Lemma 3.2 Lemma 3.2. Let the discriminator T (•) be with enough capacity, i.e., in the non parametric limit.Consider the problem T = arg maxEQUATIONwhereEQUATIONand the update rule based on the gradient descent methodEQUATION)If the gradient descent method converges to the global optimum T , the mutual information estimatorEQUATION)converges to the real value of the mutual information I(X; Y ).Proof. For convenience of notation, let the instantaneous MI be the random variable defined asEQUATION)It is straightforward to notice that the MI corresponds to the expected value of i(X; Y ) over the joint distribution p XY . The solution to (42) is given by ( 6) of Theorem 3.1. Let δ (n) = T -T (n) be the displacement between the optimum discriminator T and the obtained oneEQUATIONwhere R (n) (x, y) represents the estimated density ratio at the n-th iteration and it is related with the optimum ratio R(x, y) as followsEQUATIONwhere the last step follows from a first order Taylor expansion in T -δ (n) . Therefore, în,fDIME (X;EQUATION)If the gradient descent method converges towards the optimum solution T , δ (n) → 0 andEQUATIONwhere the RHS is itself a first order Taylor expansion of the instantaneous MI in T . In the asymptotic limit (n → +∞), it holds also for the expected values thatEQUATIONC. EQUATION)where H 2 is the Hellinger distance squared defined asEQUATIONand the infinity norm is defined as||f (x)|| ∞ := sup x∈R |f (x)|.Proof. Consider the variance of R(x, y) when (x, y) ∼ p XY (x, y), thenEQUATION)The power of the log-density ratio is upper bounded as follows (see the approach of Lemma 8.3 in [29] )EQUATIONwhile the mean squared is the ground-truth MI squared, thusEQUATIONFinally, the variance of the mean of M i.i.d. random variables yields the thesisEQUATIONC.4 Proof of Lemma 4.2 Lemma 4.2. Let R be the optimal density ratio and let X ∼ N (0, σ 2 X ) and N ∼ N (0, σ 2 N ) be uncorrelated scalar Gaussian random variables such that Y = X + N . Assume Var p XY [log R] exists. Let p M XY be the empirical distribution of M i.i.d. samples from p XY and let E p M XY denote the sample average over p M XY . Then, under the randomness of the sampling procedure, it holds thatEQUATION)Proof. From the hypothesis, the density ratio can be rewritten as R = p N (y -x)/p Y (y) and the output variance is clearly equal toσ 2 Y = σ 2 X + σ 2 N .Notice that this is equivalent of having correlated random variables X and Y with correlation coefficient ρ, since it is enough to study the case σ X = ρ and σN = 1 -ρ 2 .It is easy to verify via simple calculations thatEQUATION)Similarly,EQUATION)where the last steps use the fact that the Kurtosis of a normal distribution is 3 and that the MI can be expressed as in (59). Finally, the variance of the mean of M i.i.d. random variables yields the thesisEQUATIONIf X and N are multivariate Gaussians with diagonal covariance matrices ρ 2 I d×d and (1 -ρ 2 )I d×d , the results for both the MI and variance in the scalar case are simply multiplied by d.C.5 Proof of Lemma 5.1 Lemma 5.1. Let (x i , y i ), ∀i ∈ {1, . . . , N }, be N data points. Let J f (T ) be the value function in (5) . Let J π f (T ) and J σ f (T ) be numerical implementations of J f (T ) using a random permutation and a random derangement of y, respectively. Denote with K the number of points y k , with k ∈ {1, . . . , N }, in the same position after the permutation (i.e., the fixed points). ThenEQUATIONProof. Define J π f (T ) as the Monte Carlo implementation of J f (T ) when using the permutation function π(•)EQUATIONwhere the pair (x i , y j ) is obtained via a random permutation of the elements of y as j = π(i), ∀i ∈ {1, . . . , N }. Since K is a non-negative integer representing the number of fixed points i = π(i), the value function can be rewritten asEQUATIONwhich can also be expressed asJ π f (T ) = 1 N K i=1 T (x i , y i )+ 1 N N -K i=1 T (x i , y i )- 1 N K i=1 f * T (x i , y i ) - 1 N N -K i=1 f * T (x i , y j̸ =i ) .(65) In (65) it is possible to recognize that the second and last term of the RHS constitutes the numerical implementation of J f (T ) using a derangement strategy on N -K elements, so thatEQUATIONHowever, by definition of Fenchel conjugateEQUATIONsince for t = 1 u -f * (u) ≤ u -(ut -f (t)) = f (1) = 0.(68) Hence, we can conclude thatEQUATIONC.6 Proof of Theorem 5.2 Theorem 5.2. Let the discriminator D(•) be with enough capacity. Let N be the batch size and f be the generator of the KL divergence. Let J π KL (D) be defined asEQUATION)Denote with K the number of indices in the same position after the permutation (i.e., the fixed points), and with R(x, y) the density ratio in (2). Then,EQUATION)Proof. The idea of the proof is to express J π KL (D) via Monte Carlo approximation, in order to rearrange fixed points, and then go back to Lebesgue integration. The value function J KL (D) can be written asEQUATIONSimilarly to (64), we can express J π KL (D) asEQUATION)where K is the number of fixed points of the permutation j = π(i), ∀i ∈ {1, . . . , N }. However, when N → ∞, we can use Lebesgue integration and rewrite (73) asEQUATIONTo maximize J π KL (D), it is enough to take the derivative of the integrand with respect to D and equate it to 0, yielding the following equation in DEQUATIONSolving for D leads to the thesis EQUATIONProof. Theorem 5.2 implies that, when the batch size is much larger than the density ratio (N >> R), then the discriminator's output converges to the density ratio. Indeed,EQUATIONInstead, when the density ratio is much larger than the batch size (R >> N ), then the discriminator's output converges to a constant, in particularEQUATIONHowever, from Lemma 5.4, it is true that K = 1 on average. Therefore, an iterative optimization algorithm leads to an upper-bounded discriminator, sinceEQUATIONwhich implies the thesis.C.8 Proof of Lemma 5.4Lemma 5.4 (see [30] ). The average number of fixed points in a random permutation π(•) is equal to 1.Proof. Let π(•) be a random permutation on {1, . . . , N }. Let the random variable X represent the number of fixed points (i.e., the number of cycles of length 1) of π(•). We define X = X 1 + X 2 + • • • + X N , where X i = 1 when π(i) = i, and 0 otherwise. E[X] is computed by exploiting the linearity property of expectation. Trivially,EQUATIONwhich implies EQUATIONN (N -1) pairs (x, ỹ) ∼ p X p Y N pairs (x, ỹ) ∼ p X p Y N pairs (x, ỹ) ∼ p X p Y Nr. NNs 1 2 1Complexity Ω(N 2 ) Ω(N ) Ω(N )In this section, we show supplementary results for the linear and cubic Gaussian experiments. The implemented neural network architectures are: joint, separable, deranged, and the architecture of NJEE, referred to as ad hoc. See Tab. 1 for a schematic about the architectures.Joint architecture. The joint architecture is a feed-forward fully connected neural network with an input size equal to twice the dimension of the samples distribution (2d), one output neuron, and two hidden layers of 256 neurons each. The activation function utilized in each layer (except from the last one) is ReLU. The number of realizations (x, y) fed as input of the neural network for each training iteration is N 2 , obtained as all the combinations of the samples x and y drawn from p XY (x, y).Separable architecture. The separable architecture comprises two feed-forward neural networks, each one with an input size equal to d, output layer containing 32 neurons and 2 hidden layers with 256 neurons each. The ReLU activation function is used in each layer (except from the last one). The first network is fed in with N realizations of X, while the second one with N realizations of Y .Deranged architecture. The deranged architecture is a feed-forward fully connected neural network with an input size equal to twice the dimension of the samples distribution (2d), one output neuron, and two hidden layers of 256 neurons each. The activation function utilized in each layer (except from the last one) is ReLU. The number of realizations (x, y) the neural network is fed with is 2N for each training iteration: N realizations drawn from p XY (x, y) and N realizations drawn from p X (x)p Y (y) using the derangement procedure described in Sec. 5. The architecture deranged is not used for I CP C because in (30) the summation at the denominator of the argument of the logarithm would require neural network predictions corresponding to the inputs (x i , y j ), ∀i, j ∈ {1, . . . , N } with i ̸ = j.Ad hoc architecture. The NJEE MI estimator comprises 2d -1 feed-forward neural networks. Each neural network is composed by an input layer with size between 1 and 2d -1, an output layer containing N -k neurons, with k ∈ N small, and 2 hidden layers with 256 neurons each. The ReLU activation function is used in each layer (except from the last one). We implemented a Pytorch [31] version of the code produced by the authors of [24] 3 , to unify NJEE with all the other MI estimators.Each neural estimator is trained using Adam optimizer [32] , with learning rate 5 × 10 -4 , β 1 = 0.9, β 2 = 0.999. The batch size is initially set to N = 64.For the Gaussian setting, we sample a 20-dimensional Gaussian distribution to obtain x and n samples, independently. Then, we compute y as linear combination of x and n: y = ρ x + 1 -ρ 2 n, where ρ is the correlation coefficient. For the cubic setting, the nonlinear transformation y → y 3 is applied to the Gaussian samples. During the training procedure, every 4k iterations, the target value of the MI is increased by 2 nats, for 5 times, obtaining a target staircase with 5 steps. The change in target MI is obtained by increasing ρ, that affects the true MI according to EQUATIONAdditional plots reporting the MI estimates obtained from MINE, NWJ, and SMILE with τ = ∞, are outlined in Fig. 6 . The variance attained by these algorithms exponentially increases as the true MI grows, as stated in (10) .We report in Fig. 7 the behavior we obtained for I SM ILE when the training of the neural network is performed by using the cost function in (29) . The training diverges during the first steps when τ = 1 and τ = 5. Differently, when τ = ∞, I SM ILE corresponds to I M IN E (without the moving average improvement), therefore the MI estimate does not diverge. Interestingly, by comparing I SM ILE (τ = ∞) trained with the JS divergence and with the MINE cost function (in Fig. 6 and Fig. 7 , respectively), the variance of the latter case is significantly higher. Hence, the JS maximization trick seems to have an impact in lowering the estimator variance.The class of f -DIME estimators is robust to changes in d and N , as the estimators' variance decreases (see (58) and Fig. 15 ) when N increases and their achieved bias is not significantly influenced by the choice of d. Differently, I N JEE and I CP C are highly affected by variations of those parameters, as described in Fig. 2 and Fig. 10 . More precisely, I CP C is not strongly influenced by a change of d, but the bias significantly increases as the batch size diminishes, since the upper bound lowers. I N JEE achieves a higher bias both when d decreases and when N increases w.r.t. the default values d = 20, N = 64. In addition, when d is large, the training of I N JEE is not feasible, as it requires a lot of time (see Fig. 5 ) and memory (as a consequence of the large number of neural networks utilized) requirements. In addition, Fig. 8 illustrates that the time complexity of the joint architecture is Ω(N 2 ), while the complexity of the deranged architecture is Ω(N ).We show the achieved bias, variance, and mean squared error (MSE) corresponding to the three settings reported in Fig. 2 , 10, and 11 in Fig. 12 , 13, and 14, respectively. The achieved variance is bounded when the estimator used is I KL-DIM E or I CP C . In particular, Figs. The class f -DIME is able to estimate the MI for high-dimensional distributions, as shown in Fig. 17 , where d = 100. In that figure, the estimates behavior is obtained by using the simple architectures described in Sec. D.1 of the Appendix. Thus, the input size of these neural networks (200) is comparable with the number of neurons in the hidden layers (256). Therefore, the estimates could be improved by increasing the number of hidden layers and neurons per layer. The graphs in Fig. 17 illustrate the advantage of the architecture deranged over the separable one.To facilitate the understanding of the role of derangements during training, we provide a practical example in the following.Suppose for simplicity that N = 3. Then, a random permutation of y= [y 1 , y 2 , y 3 ] can be [y 2 , y 3 , y 1 ],where the number of fixed points is K = 0 as no elements remain in the same position after the permutation. However, another permutation of y is [y 1 , y 3 , y 2 ]. In this case, it is evident that y 1 remains in the same initial position, and the number of fixed points is K = 1. A random derangement of y = [y 1 , y 2 , y 3 ], instead, ensures by definition that no element of y ends up in the same initial position, contrarily from a naive random permutation. This idea is essential to avoid having shuffled marginal samples that actually are realizations of the joint distribution. In fact, we proved that a random permutation strategy would lead to a biased estimator (see the permutation bound in Corollary 5.3).It is extremely important to remark that the derangement sampling strategy it is not only applicable to f -divergence based estimators, rather, any discriminative variational estimator should use it to avoid upper bound MI estimates, as it can be observed from Fig. 16 DWe give an insight on how to choose the best estimator in Tab. 6, depending on the desired specifics. We assign qualitative grades to each estimator over different performance indicators. All the indicators names are self-explanatory, except from scalability, which describes the capability of the estimator to obtain precise estimates when d and N vary. The grades ranking is, from highest to lowest: ✓✓, ✓, ∼, ✗. When more than one architecture is available for a specific estimator, the grade is assigned by considering the best architecture within that particular case. Even though the estimator choice could depend on the specific problem, we consider I GAN -DIM E to be the best one. The rationale behind this decision is that I GAN -DIM E achieves the best performance for almost all the indicators and lacks weaknesses. Differently, I CP C estimate is upper bounded, I SM ILE achieves slightly higher bias, and I N JEE is strongly d and N dependent. However, if the considered problem requires the estimation of a low-valued MI, I KL-DIM E is slightly more accurate than I GAN -DIM E .One limitation of this paper is that the set of f -divergences analyzed is restricted to three elements. Thus, probably there exists a more effective f -divergence which is not analyzed in this paper. Another limitation is that f -DIME does not result in neither a lower nor an upper bound on the true MI. Nonetheless, in the following subsection, we show that it is actually possible to obtain a VLB version of the estimator.As discussed above, the f -DIME estimator does not constitute a lower bound on the true MI. This is due to two main reasons that make f -DIME different from the others VLB MI estimators:1. The general value function J f in ( 5) is the dual representation of the more general fdivergence and the KL-divergence is only one special case. Notice that the value J f is a lower bound on the f -divergence;2. We exploit the maximizer of J f , i.e. T , to build the MI estimator at inference time. This is a key component that allows us to get rid of the partition function for MI estimation, and it comes at the expense of not having a lower bound estimator.However, and perhaps quite remarkably, f -DIME can be adjusted to be a lower bound on the MI by adding the partition term (in the SMILE, MINE or NWJ fashion) during inference time. One way to do such adaptation is to use the extracted density ratio inside the expressions of NWJ or MINE, as in the following:EQUATION-E (x,y)∼p X (x)p Y (y) f * ′ T (x, y) + 1,where I f DIM E N W J (X; Y ) is the f -DIME estimator obtained using any f -divergence dual representation of (5) but with the partition term of the NWJ estimator.The benchmark considered for the self-consistency tests is similar to the one applied in prior work [20] . We use the images collected in MNIST [33] and FashionMNIST [34] data sets. Here, we test three properties of MI estimators over images distributions, where the MI is not known, but the estimators consistency can be tested:1. Baseline. X is an image, Y is the same image masked in such a way to show only the top t rows. The value of Î(X; Y ) should be non-decreasing in t, and for t = 0 the estimate should be equal to 0, since X and Y would be independent. Thus, the ratio Î(X; Y )/ Î(X; X) should be monotonically increasing, starting from 0 and converging to 1. 2. Data-processing. X is a pair of identical images, Ȳ is a pair containing the same images masked with two different values of t. We set h(Y ) to be an additional masking of Y of 3 The comparison between the MI estimators for varying values of t is reported in Fig. 18 , 19, and 20.EQUATIONThe behavior of all the estimators is evaluated for various random seeds. These results highlight that almost all the analyzed estimators satisfy the first two tests (I HD-DIM E is slightly unstable), while none of them is capable of fulfilling the additivity criterion. Nevertheless, this does not exclude the existence of an f -divergence capable to satisfy all the tests. Guidelines:• The answer NA means that the abstract and introduction do not include the claims made in the paper. • The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers. • The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings. • It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.Justification: There are no potential harms caused by the research process, and no potential harmful consequences. Guidelines:• The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.• If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics. • The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction). 10. Broader Impacts Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed? Answer: [NA] Justification: Our work is foundational research. Mutual information estimation is used for various machine learning algorithms and there are no direct paths to negative applications. Guidelines:• The answer NA means that there is no societal impact of the work performed.• If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact. • Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations. • The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster. • The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.• The answer NA means that the paper does not release new assets.• Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc. • The paper should discuss whether and how consent was obtained from people whose asset is used. • We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution. • For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.Our implementation can be found at https://github.com/tonellolab/fDIMEhttps://github.com/ermongroup/smile-mi-estimatorhttps://github.com/YuvalShalev/NJEE 


Decision: The authors propose a mutual information estimator using data derangements and explicitly takes permutation invariance into account. Reviewers appreciated the results and how it resolved the exponential complexity of this MI estimation, although there are concerns with the clarity of the paper. Authors, please incorporate the suggestions by the reviewers and make every effort to clarify the writing and presentation.

Reviews: review: paper_summary: paper_summary: This paper investigates the long-standing task of estimating mutual information in high-dimensional data. The authors point out that mutual information estimation methods using variational lower bounds suffer from either high bias or high variance, and proposes an alternative solution leveraging the variational representation of the f-divergence. They further claim that the proposed solution, f-divergence mutual information estimators (f-DIME), exhibits an excellent bias/variance trade-off, higher accuracy and lower complexity.
strength_weakness: strength_weakness: 1.	They benchmark a series of experiments using a wide range of mutual information estimators (Figure 2, 3, 4). However, the empirical results do not seem to support the central claim of “an excellent bias-variance trade-off”. It is clear that CPC is suffering from the log(N) upper bound issue where N is the number of data samples, but the proposed f-DIME method (KL-DIME, HD-DIME, GAN-DIME) do not seem to be consistently better in either bias or variance compared to the competing methods MINE and SMILE.
2.	Minor issue: The methods are not ordered the same way in Figure 2 versus Figure 3 and 4. It would be helpful to order them consistently.
3.	Minor issue: “NWJ” was not defined before first used in line 67.
4. There are non-neural network methods that have good theoretical rates and/or  good experimental results for higher dimensions. They should at least refer to these if not include some comparisons. These include the following:
https://journals.aps.org/pre/abstract/10.1103/PhysRevE.69.066138 (they reference to this but don't compare. This method often performs very well experimentally)
https://doi.org/10.1109/TIT.2021.3100108
https://proceedings.neurips.cc/paper/2015/hash/06138bc5af6023646ede0e1f7c1eac75-Abstract.html
https://www.pnas.org/doi/full/10.1073/pnas.1715593115
questions: questions: 1.	I would ask the authors to address weakness #1. 
2.	I believe it would be helpful to quantify the bias and variance in the empirical results (Figure 2, 3, 4), for better comparison across KL-DIME, HD-DIME, GAN-DIME, NJEE, MINE and SMILE. I just noticed that some relevant information is present in Table 1, 2, 3 and 4 of the Appendix. Would the authors consider presenting such results in the main text instead (for instance, as a nested bar plot)?
3.	In the related works section, the authors describe how some mutual information estimators produce high-biased estimates as they are upper bounded by log(N). Could the authors also explain the root cause(s) of high-variance estimators in the same section?
limitations: limitations: Yes, the authors have adequately addressed the limitations and potential negative societal impact of their work.
paper_summary: paper_summary: This paper proposes a novel discriminative mutual information estimator via the form of f-divergence, which exhibits a bias/variance trade-off.  Experimental results show that the proposed estimator is comparable to existing neural estimators.
strength_weakness: strength_weakness: -	Can the author theoretically show how close the mutual information estimated by using the value function ($J_f(T)$ in Theorem 3.1) is to the true mutual information?
-	The function T(·) in Theorem 3.1 is missing its definition.
-	Eq (10) seems to be a bit of problem, first row: $e^{I(X;Y)} \rightarrow e^{I_{NWJ}^{M,N}}(X;Y)$, 
second row: $e^{I(X;Y)} \rightarrow e^{I_{MINE}^{M,N}}(X;Y)$. 
-	Figure 1 shows that the proposed estimator still has a large variance like the existing methods. High variance does not seem to be well addressed.
-	The experimental utilizes low-dimensional datasets, which raises concerns about the comprehensive evaluation of the proposed estimator's efficacy in high-dimensional contexts. The authors' investigation does not adequately address the potential implications of high-dimensional data on the estimator's performance.
questions: questions: Please refer to Weaknesses
limitations: limitations: N/A
paper_summary: paper_summary: The authors provide a new representation of mutual information (in term of a general $f$-divergence) which has the advantage that the corresponding estimator has a low variance, in stark contrast to the MINE estimator which has a variance which is exponential in the size of the mutual information. The new estimator is tested on multi-dimensional Gaussian distributions.
strength_weakness: strength_weakness: The reviewer is puzzled by the statement of Theorem 3.1 on which the rest of the paper relies.  The issue is the existence of a permutation
$\sigma$ such that $p(\sigma(y)|x) = p(y)$.  Do the author mean to assume the existence of such permutation?  The reviewer may be missing a point but I cannot see why such permutation should exist which makes the permuted random variable $Y$ independent of $X$.  I don't see why this should be true for general data and if it holds true this seems like a very strong assumption because it implies a form independence between $X$ and $Y$.  If this point, which is central to the paper, is not clarified, the paper cannot be considered for publication in the present form.
questions: questions: Clarify the issue of the existence of the permutation in Theorem 1. Prove its existence, or provide examples to which it applies.
limitations: limitations: The limitations have not been addressed properly, especially in regard about the assumptions of Theorem 3.1.
paper_summary: paper_summary: This paper introduces a new method for mutual information estimation using a discriminative training approach based on f-divergences. Notably, the authors address a well-known limitation for this class of estimators which exhibit exponential sample complexity in the strong dependence limit. The authors also demonstrate a simple but surprising observation that the method for drawing marginal samples can have a huge impact on estimates, and show how to correctly handle this with "data derangements".
strength_weakness: strength_weakness: - The theoretical analysis only consider the case when the optimum discriminator is exactly known. It would be nice to have some understanding of the effect of mis-estimation.  Unlike other approaches, we don't get a lower bound on MI, only an estimate. 
- Results section could be improved for clarity. I didn't follow the architecture very well, maybe it can be illustrated? (Appendix D helped - maybe it's just a space issue. A schematic in the appendix would still be helpful though.) 
- The final conclusions in the Variance Analysis could be stated more clearly.
questions: questions: Am I correct that one drawback of this procedure is that it doesn't result in an MI lower bound? (Because the objective that is not optimized is not a MI bound, it just guarantees that at the optimum you get a solution that gives the true MI.)  If so, is it easy to understand why other discriminative approaches do give a lower bound and this one doesn't? 

Why did the derangement sampling strategy affect time so much? Does it change the convergence speed, or am I missing something else?
limitations: limitations: I think these could be made more explicit.

Comments: title: Approaching the end
Satisfactory answers.
Thanks
Answer on benchmarks and stability
Follow-up question on stability / benchmarks
We thank the Reviewer
Response
comment: Dear Reviewer 3r9p,

As the rebuttal period is approaching its end, we would really appreciate your feedback. The other 3 reviewers expressed their satisfaction with our responses, and we would like to know if our answers also clarified your doubts.

Thanks again for your valuable time.
We deeply thank the Reviewer for the positive comment and we will indeed write the clarification of 1 more emphatically in the paper. To ensure accurate reflection of the Reviewer's comment, we kindly request an update of the score to 6 before the rebuttal period ends, as we still do not see this change reflected. Thank you.
The answers are satisfactory to me. I appreciate the clarification of 1 and this should be written more emphatically in the paper. I will raise my score to a 6.
We’d like to thank the reviewer for reading our reply and recognizing the novel contribution on data derangement.

In the second part of the Reviewer’s comment, we believe there still is a misunderstanding about our claim that the proposed MI estimator has low variance and that “the case is theoretically not made”. We politely disagree:

1. We explain that $f$-DIME has lower variance w.r.t. the variational lower bound (VLB) estimators (see lines 134, 152, 242 and Eq. (13)).

2. This statement is theoretically proved in Section 4. Herein, we begin by summarizing what was found in [17]: the estimation of the partition function is responsible for the high variance condition in VLB MI estimators. Then, we explain that $f$-DIME avoids the partition function estimation, thus decreasing the estimator’s variance. This is shown by Eqs. (12-14).

3. Let us try to explain the findings with different terms to ease the understanding. VLB estimators comprise two terms: the first term is the expectation over $p_{XY}$, the second term is the expectation over $p_Xp_Y$. This indeed reflects the notation used, where the variance of VLB estimators is indicated with $Var_{p_{XY}, p_Xp_Y}$. From (11), it is possible to notice that the variance of the VLB estimators is the sum of the variance of two expectations. The variance of the second expectation ($E_{p_Xp_Y}$) is always greater or equal than $0$, which leads to (12). In particular, for high values of true MI, the variance of $f$-DIME ($Var_{p_{XY}}[E_{p_{XY}^M}[\log\hat{R}]]$) is significantly lower than the variance of the VLB estimators, because the second expectation term has exponentially increasing variance with the true MI (demonstrated in [17]).

4. Additionally, we reported two Lemmas in the Appendix that provide: 1) an upper bound on the variance, 2) a characterization of the variance when X and Y are correlated Gaussian random variables.

5. Finally and for completeness, we briefly discuss the bias/variance comparison between $f$-DIME and the other MI estimators analyzed in the paper, to justify why $f$-DIME attains an excellent bias/variance trade-off. SMILE decreases the variance of VLB estimators by clipping the partition term. Therefore, also SMILE attains variance comparable to $f$-DIME, but it is biased. CPC is a low variance estimator, however it is more biased than all the other estimators considered, since it is upper bounded by $\log(N)$. NJEE is low variance because it uses $2d-1$ neural networks to obtain the prediction. For this reason, NJEE is computationally extremely expensive (please see line 290 and subsequent ones). Furthermore, NJEE is biased for various choices of $d$. In conclusion, $f$-DIME attains lower variance w.r.t. the VLB estimators, and in general exhibits an excellent bias/variance trade-off.

6. The experimental results provide further evidence of the findings.

We hope that this provides more clarity and offers the basis to raise the score further. Thank you again.
Thank you for the detailed answers.  I will raise the score to take into account better the contributions of the paper on data derangement strategies which are novel and very interesting.  However am still not convinced by the theoretical part of the paper, and it is not correct to state that the authors hav a low variance estimator of the MI.  It may well be that the estimator proposed is a good substitute but the case is theoretically not made.
Thanks for all the last-minute clarifications. Sorry I didn’t carefully recheck the results when I went back to investigate about the student results. Thank you for the stability tips, and looking forward to the “camera ready” code.
Dear Reviewer,

We believe there might be a misunderstanding here. 
In our paper we reported many tests with the benchmarks from Czyz et al. (we did cite such paper, please check ref. [24]). 
Figs. 3 and 4 especially describe the excellent results attained by $f$-DIME in those benchmarks.  
The reason why the Reviewer’s students did not find them in the arxiv version is that such a version is outdated. 
  
Regarding the experiments for the Czyz et al. benchmarks, the results we obtained were stable. We noticed that the version of the libraries used for the tests may have an impact. We will specify on GitHub the exact versions of the libraries and code we used to obtain the results in Figs. 3,4. Moreover, we will include a docker image in the repository to facilitate reproduction.  
The motivation why the code for the self-consistency tests is not provided yet is partially as the Reviewer wrote, and partially because the results of the self-consistency tests are reported in the Appendix, and it is not required to provide code for results in the Appendix. However, we will also include this part of the code on GitHub. 

Finally, we agree about instability for the self-consistency tests for HD-DIME and some NaNs encountered for KL-DIME (as can be observed from Figs. 17,18,19), while GAN-DIME returned stable results.

Thanks.
I just realized that you didn't cite or include the benchmarks from Czyz et al, "Beyond Normal: On the Evaluation of Mutual Information Estimators", NeurIPS 2023. This seems like an omission, though it doesn't directly probe the things you were focusing on. 

The reason I realized this is the following. I sent the title of your paper to my students and directed them to see if there is an arxiv version (without telling me the authors of course!). They found it and ran your method on our internal benchmarks, including Czyz et al. 

The results were very unstable, and we got NaNs on most of these results. Do you have any thoughts on stability of your results / this benchmark? 
Also, it seems that some of the experiments (self-consistency) were not yet implemented in the code, though I understand cleaning up code often happens closer to final publication. 

(I realize this is is quite late in the cycle, and would appreciate even brief high level thoughts /  speculation on these.)
We sincerely thank the Reviewer for the positive comment and we are glad to read that the Reviewer increased the score to 7.
Thanks for the detailed feedback, and clarification on variance results. I'll raise my score based on that, and partially to balance out the review of dXEo which seems off base.