Data-Efficient Variational Mutual Information Estimation via Bayesian Self-Consistency

Authors: desi r ivanova
marvin  schmitt
stefan  radev
Venue: NeurIPS_cc_2024_Workshop_BDU
Published: 2024-09-06 03:46:35

Abstract: Mutual information (MI) is a central quantity of interest in information theory and machine learning, but estimating it accurately and efficiently remains challenging.
In this paper, we propose a novel approach that exploits Bayesian self-consistency to improve the data efficiency of variational MI estimators.
Our method incorporates a principled variance penalty that encourages consistency in marginal likelihood estimates, ultimately leading to more accurate MI estimation and posterior approximation with fewer gradient steps.
We demonstrate the effectiveness of our method on two tasks: (1) MI estimation for correlated Gaussian distributions; and (2) Bayesian experimental design for the Michaelis-Menten model.
Our results demonstrate that our self-consistent estimator converges faster whilst producing higher quality MI and posterior estimates compared to baselines.

Main paper: 
 introduction: 
 Estimating dependencies between random variables is a fundamental problem in science and engineering. Mutual information (MI) is a general information-theoretic metric to quantify pairwise relationships that manifest themselves beyond the covariance matrix [15] . Additionally, estimating and optimizing MI is fundamental to many machine learning problems and tasks, such as representation learning [5, 28] , understanding neural networks [26] , Bayesian optimization and active learning [10, 14, 18] , and Bayesian experimental design [17, 23] , to name just a few.Despite its favorable theoretical properties, estimating MI involves nesting of expectations, rendering the problem doubly intractable [22, 29] . As a result, its optimization is extremely challenging, especially in high-dimensional settings [27] . Recent advances in representation learning have proposed computationally tractable variational surrogate objectives, parameterized by neural networks [2, 20, 21, 28] . These approaches promise improved scalability not only in terms of dimensionality but also in handling complex data distributions and large-scale datasets. However, as McAllester & Stratos [19] show, MI estimation becomes challenging not only with increasing dimensionality, but also with the magnitude of the MI itself. They prove that any distribution-free estimator requires an exponential number of samples, which clearly does not scale in a computationally feasible way.In many applications of probabilistic modeling, however, we have access to at least one closed-form distribution. In Bayesian analysis, this could be the prior distribution over a set of parameters θ which expresses our knowledge about plausible (latent) parameter values before observing data. In such cases, the limitations of distribution-free estimators [19, 25] no longer apply, opening up opportunities for more efficient estimators. In line with this research stream, we investigate a novel self-consistency penalty that has the potential to improve the sample efficiency and estimation quality of variational MI estimators. By encouraging self-consistency [16, 24] in the probabilistic joint model p(x, y), even when parts of the model are approximated with surrogate distributions, we demonstrate that we can further enhance the sample efficiency of variational MI estimators. This penalty is designed as a simple plug-in for existing variational estimators, requiring minimal modifications to current implementations. The trade-off for this improved efficiency is a slight increase in computational cost during training, though importantly, it does not require additional samples from real-world observations or expensive simulators or extra computation during inference time. 
background: 
 Mutual information (MI) is a fundamental concept in information theory, measuring the amount of information shared between two random variables. For random variables X and Y , the MI is the Kullback-Leibler (KL) divergence between their joint probability and the product of their marginals:EQUATIONEstimating MI in practice is computationally costly, particularly when dealing with high-dimensional or continuous variables, with various estimation techniques developed to address these challenges.If one of the conditional distributions, for instance, p(y | x), is available in closed-form, we can approximate the marginal distribution p(y) = E p(x) [p(y | x)] using L samples x l ∼ p(x). This gives us the Nested Monte Carlo (NMC) estimator (also known as double-loop Monte Carlo or DLMC):EQUATION)The concavity of log(•) means that NMC is biased, providing an upper bound on the MI. To obtain a corresponding lower bound and thus "sandwich" the true MI, we can modify the NMC estimator by including the probability of the "positive" sample from the numerator to the "negative" (or contrastive) samples in the denominator. This leads to the Prior Contrastive Estimator (PCE [9] ):EQUATIONBoth bounds are asymptotically consistent, converging to the the true MI in the limit as L → ∞.Optimal convergence efficiency is achieved by setting L ∝ √ N , where N is the number of samples we use to approximate the outer expectation. This, at best, gives us an overall convergence rate of O(C -1/3 ), where C = N L is the total computational budget, which is considerably worse than O(C -1/2 ) for conventional Monte Carlo. One way to improve the efficiency of these basic estimators is to introduce a variational proposal distribution q v (x | y) and apply importance sampling to the x values, potentially leading to a more accurate approximation of the marginal p(y):EQUATIONSubstituting Eq. 4 into the denominator of Eq. 2 yields the Variational Nested Monte Carlo (VNMC) estimator, which we denote by U VNMC (q v ; L). Similarly, adjusting for the positive sample and substituting in Eq. 3, we obtain the Variational Prior Contrastive Estimator (VPCE), which we denote as L VPCE (q v ; L). For clarity, the estimators are given in Appendix A.These variational estimators preserve the upper and lower bound properties of their non-variational counterparts whilst offering an improved estimation efficiency. Specifically, in addition to asymptotic consistency as L → ∞, VNMC and VPCE can also achieve finite sample unbiasedness (with L = 1) for the optimal proposal, which equals the posterior q * v (x ℓ | y) = p(x ℓ | y). 
method: 
 By rearranging Bayes' rule, we can express the marginal distribution (i.e., the evidence in Bayesian inference) of y as p(y) = p(x) p(y | x) p(x | y) -1 . An important observation that follows directly is the self-consistency of the marginal p(y) when evaluated based on different values of x. In other words, for any set of values x 1 , . . . x K we have: However, if we replace the posterior p(x | y) with an approximation q(x | y), Eq. 5 will not hold exactly, with different x k yielding different p(y) values. This variance is undesirable, as it is a direct result of approximation error. Thus, a high-quality approximate posterior q(x | y) minimizes the variance of p(y). Recently, [24] used this insight to propose a variance penalty that improves the sample efficiency of neural posterior estimation in amortized Bayesian inference,EQUATIONEQUATIONwhere q(x) is an easy-to-sample proposal distribution whose choice has non-trivial consequences for the empirical behavior of Eq. 6. Building on these results, we explore the utility of self-consistency in the context of MI estimation and optimization. 
empirical evaluation: 
 In this section, we report preliminary results validating our self-consistent estimation algorithm.Our experiments focus on two key areas: mutual information estimation and Bayesian optimal experimental design. The main baseline we compare against is the VPCE estimator. We use identical neural network architectures and hyperparameters for all methods to ensure a fair comparison. To provide a controlled setting for these initial investigations, we assume access to an analytic likelihood function throughout our experiments. We note that a likelihood function can be learned within the same optimization scheme [24] , and we plan to extend our approach to such scenarios in future work.The main focus on two key aspects of MI estimators: accuracy and data efficiency, particularly in high MI regimes. Whilst quick and accurate MI estimation is desirable in the context of optimization, obtaining high-quality amortized posteriors is also of paramount interest due to their practical utility in downstream tasks, such as marginal likelihood estimation or cross-validation. To this end, we use the squared maximum mean discrepancy (MMD; [11] ) distance between a "groundtruth" reference posterior p(x | y), estimated with HMC [3, 12] , and our approximation q(x | y):EQUATIONwhere κ(•, •) is a positive definite kernel. Further evaluation metrics encompass (1) the relative error in estimated mutual information; and (2) the estimation error of the log posterior density p(x | y).This experiment serves as a proof-of-concept in a controlled setting where closed-form values for all quantities of interest are known, allowing for a principled evaluation. We follow a similar experimental setup to that in Poole et al. [21] and consider D-dimensional correlated Gaussian random vectors X and Y with some correlation ρ.Setup We evaluate both VPCE and our self-consistent VPCE across various dimensionalities D ∈ {2, 4, . . . , 20} and ρ = 0.98, resulting in MIs between 3.0 and 36. We train the estimators with batch size of 16 and set the number of negative samples to L = 4. We use K = 8 selfconsistency samples to estimate the variance in Eq. 6. We track the estimated mutual information and posterior quality over the course of 1000 gradient steps. We train a simple Gaussian proposal with a learnt mean and standard deviation:q v (• | y) = N (• | α + β y, σ).To ensure robustness, we repeat the training and evaluation procedures five times with different seeds and report the averages. Results As shown in Figure 1 , the VPCE baseline requires about 800 gradient steps to achieve accurate MI estimates. In contrast, our SC-VPCE achieves the same performance with only 600 gradient steps, which amounts to a 25% relative increase in sampling efficiency. Furthermore, since we have access to the true posterior in closed form for this analytical experiment, we can compute the KL divergence exactly, providing a precise measure of posterior quality. Figure 2 illustrates that our self-consistent VPCE yields essentially perfect approximations of the posterior distribution (measured by KL divergence) while being approximately 50% more sample-efficient than the VPCE baseline.Originally rooted in statistics [4, 17] , Bayesian experimental design (BED), has recently seen renewed interest, particularly in the engineering and machine learning communities [1, 8, 13, 23] . The objective in BED is to choose experiments that yield the most informative data about underlying model parameters θ. Denoting the experiment parameters by ξ and the outcome of the experiment by y, the optimal design ξ * maximizes the mutual information between y and θ, i.e. ξ * = arg max ξ I(θ; y | ξ).We consider a static experimental design problem, where the goal is to learn the parameters in the Michaelis-Menten model [6] . Michaelis-Menten is a non-linear model that is widely used in biology and chemistry to describe the relationship between the (scaled) concentration of a substrate ξ ∈ [0, 1] and the rate of an enzymatic reaction θ = (θ 1 , θ 2 ) ∈ R 2 . The outcome y of an experiment with design ξ is modelled probabilistically as y | ξ, θ ∼ N (f (ξ; θ), σ 2 ) with f (ξ; θ) = θ 1 (ξS) b /θ b 2 + (ξS) b , where S and b are a fixed scaling factors. The prior on the parameters of interest is p(θ 1 ) = N (θ 1 ; 0.5, 0.1 2 ), p(θ 2 ) = N (θ 2 ; 0.5, 0.1 2 ). Setup We use the SC-VPCE lower bound to learn 10 designs ξ 1 , . . . , ξ 10 using a batch size of N = 32, L = 32 contrastive samples, and vary the number of self-consistency samples K ∈ [64, 32, 16, 0], with K = 0 corresponding to the VPCE baseline. We use b = 6, S = 400, and observation noise σ = 10.0.We choose an expressive neural-based proposal distribution, namely a neural spline flow (NSF) [7] that learns an invertible transformation from a Gaussian base distribution to the target posterior distribution p(x | y). The posterior density is tractable via the change-of-variables formula [16] . The NSF q ν is parameterized by a total of 5 980 learnable parameters ν, spread across 2 neural transformation blocks with 32 hidden features. MMD computations are performed over 50 simulations, using the learned design.As the number of gradient steps increases from 500 to 2000, we observe a general trend of decreasing MMD values, indicating improved posterior quality (see Table 1 ). Our method consistently achieves lower MMD values, indicating superior posterior estimates. We observe that the MI values remain relatively stable, suggesting that the designs quickly converge to their optimum. 
conclusion: 
 In this paper, we presented a method to integrate self-consistency losses into variational MI estimators. Compared to baselines, our data-efficient estimator requires fewer gradient steps to converge and produces better posterior distributions, which can be useful for downstream tasks. In future work we plan to reduce the required modeling assumptions by extending our method to learned likelihoods. 
 
 Appendix: The variational nested Monte Carlo (VNMC) upper bound is given by: U VNMC (q v ; L) := E x0,y0, The corresponding variational PCE lower bound is given by: L VPCE (q v ; L) := E x0,y0,  


Decision: Both reviews are positive, though the first is suspiciously so and includes language that is otherwise atypical in a review (“innovative”, “underline the practical significance”, ..). Nonetheless, the second review is also positive, so I recommend acceptance.

Reviews: review: title: title: Evaluation of Data-Efficient Variational Mutual Information Estimation via Bayesian Self-Consistency: Insights on Methodology and Application
main_review: main_review: The manuscript introduces a novel approach to estimating mutual information (MI) by harnessing Bayesian self-consistency, representing a considerable advancement in the landscape of information theory and machine learning. The authors rigorously derive their variational estimator, incorporating a principled variance penalty that encourages consistency in marginal likelihood estimates. This technical innovation effectively addresses the challenges associated with MI estimation, particularly in high-dimensional spaces where conventional methods exhibit significant sample inefficiency. The rigorous empirical evaluation, demonstrated through tasks such as MI estimation for correlated Gaussian distributions and Bayesian experimental design with the Michaelis-Menten model, provides compelling evidence that the proposed self-consistent estimator converges faster and yields more accurate MI and posterior approximations compared to existing baselines.

In terms of clarity, the manuscript is generally well-structured, facilitating a logical progression from foundational concepts to detailed methodological exposition. However, the complexity of the mathematical derivations and algorithms, particularly those surrounding Nested Monte Carlo (NMC) estimators and Variational Prior Contrastive Estimators (VPCE), can pose challenges for readers without a strong foundation in probabilistic modeling and Bayesian inference. Increased elaboration on crucial equations and the underlying assumptions would enhance understanding and accessibility. Furthermore, a more thorough explanation of experimental settings and performance metrics—such as the squared maximum mean discrepancy (MMD) and KL divergence—would provide deeper insights into the validity of the results.

The originality of this research is evident in its innovative application of self-consistency penalties to variational MI estimators, which not only enhances sample efficiency but also maintains the desirable properties of asymptotic consistency. The empirical results underline the practical significance of the findings, especially given the increasing dimensionality of datasets common in machine learning tasks. While the focus on two specific experimental setups demonstrates the method's efficacy, expanding the empirical evaluation to include a wider variety of applications would strengthen the claims regarding generalizability. Overall, this manuscript demonstrates substantial technical merit and presents a strong contribution to the field, warranting recognition for its innovative approach and comprehensive validation of results.
title: title: Review for "Data-Efficient Variational Mutual Information Estimation via Bayesian Self-Consistency"
main_review: main_review: The paper proposes a bound for use in mutual information maximization that can be an alternative to VNMC and VPCE. 


Major comments: 
- The self consistency explanation for the derivation of the bound (although the basic idea is proposed in a cited paper and not here, I am going to criticise it here) is informal, perhaps too informal (admittedly a subjective value judgement). For example, it leads to statements like that in line 76-77: " Thus, a high-quality approximate posterior q(x | y) minimizes the variance of p(y). " p(y) is either a number (scalar) or a PDF, function of y, not a random variable, so it does not have a variance. The way I can make sense of this self consistency is that it is a more vague way of saying that we would like low-variance estimates of p(y) like with the estimator in Eq. (4), i.e. an importance sampling estimator. And, well, that we already know from standard MC theory - yes, in Eq. (4) we want low variance, amounting to the q being close to the true posterior ( and variance = chi^2 divergence ). Therefore the jump to Eq. (6) is too loosely motivated. 
- It should also be mentioned that the proposed loss of Eq (6) appears before for instance in Richter, L. et al (2020) (see Eq. (5)), although in a different context (gradient estimation). 
- The full paper should say how you plan to maximize and / or estimate Eq 6, conditions under which your (gradient?) estimators have finite variance, and how this compares with VNMC/VPCE. 
- An interesting question is that now you have two distributions to choose/optimize: q(x) and q_v(x | y), while e.g. VNMC/VPCE only have the latter. How can we compare apples to apples here? You should elaborate on this. What should q(x) be? if you do a simple MC estimate of Eq (6) sampling from q(x) , by IS theory (see e.g. Owen's "Monte Carlo theory, methods and examples" (2013) ) the optimal q(x) to estimate Eq (6) is a weird function. Or maybe you want a q(x) that minimizes the variance of a *gradient* estimator of Eq (6). All things that should be elaborated upon. And how to intertwine with the optimization of q_v(x|y) which supposedly should just try to get close to p(x|y).
- In the metrics, "the estimation error of the log posterior density" - what do you mean precisely ? I am guessing you mean for some samples / values x_1, etc. you calculate  approximation(x1) vs true_density(x1) and see how far they are. This is a weird metric, if so. Better to have things like estimates of divergences. Also the MMD is a bit arbitrary / fishy since it depends on an arbitrary kernel.  Error in mutual info estimation makes more sense, I would suggest plotting boxplots of estimates vs true value (if not closed form, obtained with a very long run of a consistent method).  Although also not a fully relevant metric if the aim is to *maximize* MI as opposed to *estimate it*, in which case some notion of convergence speed to the minimum should be there. 

I have to say that I personally do not see what the "self consistency" perspective brings to arriving to Eq 6, which maybe could have a more formal motivation (it could borrow from Richter, L. et al 2020 but not necessarily). 

Refs. 

- Richter, L., Boustati, A., Nüsken, N., Ruiz, F. and Akyildiz, O.D., 2020. VarGrad: a low-variance gradient estimator for variational inference. Advances in Neural Information Processing Systems, 33, pp.13481-13492.