Flow-based Variational Mutual Information: Fast and Flexible Approximations
Authors: Caleb Dahlke
Jason Pacheco
Venue: ICLR 2025 Poster (2025)
Keywords: Mutual Information
Variational Methods
Normalizing Flows
Bayesian Optimal Experimental Design
Primary area: probabilistic methods (Bayesian methods, variational inference, sampling, UQ, etc.)
TL;DR: We introduce two new flow-based variational Mutual Information estimations that are comparable in expressiveness to critic-based methods, such as MINE or InfoNCE, but can also approximate large MI

Abstract: Mutual Information (MI) is a fundamental measure of dependence between random variables, but its practical application is limited because it is difficult to calculate in many circumstances. Variational methods offer one approach by introducing an approximate distribution to create various bounds on MI, which in turn is an easier optimization problem to solve. In practice, the variational distribution chosen is often a Gaussian, which is convenient but lacks flexibility in modeling complicated distributions. In this paper, we introduce new classes of variational estimators based on Normalizing Flows that extend the previous Gaussian-based variational estimators. Our new estimators maintain many of the same theoretical guarantees while simultaneously enhancing the expressivity of the variational distribution. We experimentally verify that our new methods are effective on large MI problems where discriminative-based estimators, such as MINE and InfoNCE, are fundamentally limited. Furthermore, we compare against a diverse set of benchmarking tests to show that the flow-based estimators often perform as well, if not better, than the discriminative-based counterparts. Finally, we demonstrate how these estimators can be effectively utilized in the Bayesian Optimal Experimental Design setting for online sequential decision making.

Published as a conference paper at ICLR 2025
FLOW-BASED VARIATIONAL MUTUAL INFORMATION: FAST AND FLEXIBLE APPROXIMATIONS
Caleb Dahlke
Jason Pacheco
Department of Mechanical Engincering University of Michigan
Department of Computer Science
University of Arizona
Ann Arbor, MI, USA cdahlke @emmimmihhh..m edu
Tucson, AZ, USA
pacheco jarorjaariiiiii edu
ABSTRACT
Mutual Information (MI) a fundamental measure of dependence hetween random variables, but its practical application is limited because it is difficult to calculate in many circumstances. Variational methods offer one approach by introducing approximate distribution create various bounds on MI, which in turn is an easier optimization problem to solve. In practice, the variational distribution chosen is often Gaussian, which is convenient but lacks flexibility in modeling complicated distributions. In this paper, we introduce new classes of variational estimators based on Normalizing Flows that extend the previous Gaussian-basee variational estimators. Our new estimators maintain many of the same theoretical guarantees while simultaneously enhancing the expressivity of the variational distribution. We experimentally verify that our new methods are effective on large MI problems where diseriminative-bassee estimators, such as MINE and InfoNCE, are fundamentally limited. Furthermore, we compare against a diverse set of benchmarking tests show that the flow-based estimators often perform as well, if not better, than the discriminative-basee counterparts. Finally, we demonstrate how these estimators can be effectively utilized the Bayesian Optimal Experimental Design setting for online sequential decision making.
INTRODUCTION
Mutual Information (MI), fundamental measure of dependence from information theory, has found utility in a variety of fields such as Bayesian optimal experimental design (Lindley 1956 Foster et al. 2019), representation learning (Alemi et al. 2016 Chen et al. 2016), and structure learning (Vinh et al. 2011). Despite having simple interpretation, MI typically lacks a closed form solution making difficult measure to in many practical settings. Straightforward sample-bas based estimates can be used to approximate MI, but these are often inefficient both in terms of computation and sample complexity. Moreover, these so-called nested Monte Carlo (NMC) estimators exhibit large finite sample bias that decays slowly (Zheng et al. 2018 Rainforth et al. 2018).
To resolve the computational issues associated with MI, many so-called discriminative based esti- mators have been proposed. These distribution-fres approximations typically utilize the Donsker- Varadhan lower bound of MI (Donsker & Varadhan 1975), with some examples including MINE (Belghazi et al. 2018) and NWJ (Nguyen et al.) 2010). While these estimators have seen wide adoption in practice, McAllester Stratos (2020) show that they require an exponential number of samples to estimate large MI values. This poor sample complexity presents fundamental limitation in regimes where the achievable information likely to be high.
Generative variational methods, referred to simply variational methods for the remainder of the paper, directly approximate the intractable target distribution with a tractable surrogate have seen promise in practical settings (Barber & Agakov| 2004. Foster et al. 2019 2020). These variational methods make MI estimation more scalable than NMC based approaches and are not subject to the same fundamental limitation of approximating large MI discriminative based estimators. Yet computation of these estimators can be prohibitively costly, particularly in sequential decision making regimes (Foster et al. 2021 [vanova al. 2021) Moreover, the accuracy of variational estimators is

---

Published as a conference paper at ICLR 2025
heavily dependent on the expressibility of the chosen surrogate distribution, which is often limited to a Gaussian approximation to control computation (Dahlke et al. 2023 Pacheco Fisher III 2019
Barber & Agakov 2004).
We address the computational and accuracy aspects of variational MI estimators by introduction of two new estimators based on normalizing flows (Kobyzev et al. 202T). The proposed estimators
flexibly adapt to complex target distributions and allow a straightforward tradeoff between accuracy and computational complexity. The first estimator, Joint Variational Flow (JVF), builds the Joint Variational Gaussian (JVG) estimator proposed in Dahlke et al. (2023). That previous work established that, for Gaussian approximating families. the estimator can be easily computed by moment-matching matching the target distribution. We relax the Gaussian assumption and show that the same efficient updates can be used in a flow-based estimator while simultaneously achieving more flexible density approximation.
Both JVG and JVF assume access to a joint distribution of the random quantities of interest; Gaussian in the first case and in the second more flexible generative distribution given by reversing the normalizing flow. While this assumption of an underlying joint leads to efficient computation, it can restrict the MI approximation. The Neurul Variational Gaussian (NVG) relaxes this assumption using a netrork parameterized Gaussian distribution to approximate the posterior. We extend this estimator to Neural Variational Flow (NVF) to achieve more expressible estimator that incorporates flow-based density estimates into the underlying approximation.
Contributions We summarize the main results of our paper as follows. We introduce 2 families of variational MI estimators (JVF and NVF) that incorporate normalizing flows to flexibly adapt to complex multimodal target distributions. We extend existing results to show that parameters of the underlying flow distribution can be efficiently computed via moment-matching operations. We provide a wide array of benchmarking experiments to highlight the strengths and weaknesses each estimator as well to compare to many of the state of the MI estimators. Finally, we show that our variational estimators flexibly adapt to complex sequential decision making tasks with nonlinear non-Gause Gaussian noise distributions.
PRELIMINARIES
Consider a joint distribution p(X, Y') over latent variable X and observation variable Y. Mutual information (MI) measures the dependence between these random quantities and is given by:
I(X,Y) Hp(o(X)) Hp(x(XIY))
where (x)) = Bo-log (()) is the marginal entropy and (x(p(X()) = E log (X the conditional entropy. The entropy expectations are taken with respect to the joint p(X, Y). Despite its simplicity, MI typically lacks glosed-form solution. Furthermore, nested Monte Carlo estimators of MI have large finite sample bias that is slow to decay (Zheng et al., 2018 Rainforth et al. 2018) making direct sample-based estimates of MI practically infeasible in many settings. To address these issues, variational methods introduce an approximate distribution, g. to estimate the underlying true distribution, p by minimizing Lullback-Leible divergence: min, KL 1| p).
2.1 VARIATIONAL MUTUAL INFORMATION
The MI in Eqn. I requires knowledge of the marginal and posterior distributions. We consider an approach where both distributions are approximated, p(X) se Gmarg (X) and p(X (x)st(X| in the marginal and posterior approximation (Foster et al. Dahlke
Dahlke et al.; 2023):
(X+ Y) = (m (X)) Hppppstt (X]]))))X) v Y(X,Y)
(2)
where Hp(g(-)) = E,l--oo log (()) is the cross-entropy. We then minimize the following upper bound on absolute error (Foster et al., 2019):
Lemma 2.1. For any model p(X, Y) and distributions qmarg (X), Apost (X | Y), the following holds: I| S min Hp (ymaar (X)) min Hpp9po Spss (X Y)) +C
1m+p
gmnnng
Spost
where C = - H (p())) (p(X Y)) does not depend on 9warg or 9post- Further, the RHS is 0
iff "marg (X) p(X) and 9post X (Y) p(X YY) almost surely.
2

---

Published as a conference paper at ICLR 2025
This bound can be made arbitrarily tight by finding good variational approximations of both p(X) and p(X Y) and has the added benefit only assuming that p(X, X Y) can be sampled from. Thus, Imtp can be applied in many settings, such those with implicit (e.g. simulation-biae based) distributions.
2.2
JOINT VARIATIONAL GAUSSIAN (JVG) ESTIMATOR
Dahlke et al. (023) show that under Gaussian approximating distribution the upper bound in Lemma 2.J can be efficiently minimized by moment-matching operations.
Theorem 2.2 (Moment Matching Optimization). Let g(X,Y) be joint Gaussian density:
(X, = N (
x
Pz
3 Mx
May
μr
Σ
Σ
yz
Eyn
then the marginal and posterior optimizing the bound in Lemma 2.1 are given by:
Qmarg (X) = N (X [Az: Eee),
Opost (X (Y) =N(X Mxly Exlu) where
(4)
Pzly Maz EY -My), Eely == Enn Ery Dy Enyz where the mean and covariance matched to the moments of the target density Y(x
(5)
Ep((X,))']== p.
Covp(X, D.
(6)
Furthermore, the value of mmppp can be calculated in closed-form as the mutual information of the moment-matcin Gaussian (Dahlke et al., 2023). We refer to this as the Joint Variational Gaussian
(JVG) estimator:
lJvG == Hp(gmarg(()) - Hplapost (X = log Idet (2Tr)Laa) log det (2Te)eul)) (7)
These results ensure that, for Gaussian approximating q, the approximate MI Impp can be efficiently computed via simple moment-matchine operations. The strength of this approach is that it does not require gradient descent and very fast to compute practice. The drawback of this approach that it can only model linear dependence between random variables.
2.3 NEURAL VARIATIONAL GAUSSIAN (NVG) ESTIMATOR
To capture nonlinear dependence we relax the assumption that X and Y are jointly Gaussian (Foster al. 2019). base Gaussian variational distribution is assumed, but now the mean and variance of the posterior distribution are parameterized by a neural network which is a function of the observation variable y:
(X) = (X Mzz Erz), x(X Y) = (X E(Y))
(8)
where (aY) and Σ(Y) are given by neural network function approximators trained to minimize the bound Lemma 2.1 Marginal moments us and Ege in Eqn. (8) are still learned by moment matching to the true distribution. This results in the Neural Variational Gaussian (NVG) estimator:
1 N
INvG = Hx(g(X)) - H,(g(X(8)) ~ log det (2TeE
Σ log (zala()s), (9)
where (1,,, The relaxed assumption of a Gaussian joint distribution allows for INvG to approximate nonlinear dependencies between X and Y, however this comes at the computational cost of needing to train neural network. Furthermore, INVG is still constrained in its approximation power by assuming a Gaussian for each conditional value as well as the prior.
FLOW-BBSED VARIATIONAL. ESTIMATORS
Previous work considered Iyg and INv as jointly- and conditionally-Gauun Gaussian variational approxi- mations for MI, respectively (Dahlke et al. 2023 Foster et al. 2019). Both of these approaches limited by the expressibility of the Gaussian to an arbitrary distribution To address this, we introduce more flexible class of distributions based normalizing flows. The aim of normalizing flows is to
3

---

Published as a conference paper at ICLR 2025
transform simple base distribution, typically Gaussian, via diffeomorphism to achieve more expressive distribution. Let E IR be a random variable with density g2(2) and let Z = be diffeommrphism. Then the density of X is given by the change of basis formula (Bishop & Nasrabadi 2006):
qx xx(X) o2 (f()) det (Vx](X))
(10)
where V (X) is the Jacobian of f(X) with respect to X. The choice of f(-) can be any diffeomor- phism, however, train effectively we must be able to evaluate and its log determinant efficiently.
There have been many proposed flows (Kobyzev et al. (02T) including planar and radial (Rezende & | 2016), coupling (Kingma et al.] 20166,, and auto-regressive flows (Kingma et al. |016] Pa-
pamakarios et al., 2021). For this work we consider rational quadratic auto-regressive flows (Durkan ct al. 2019) due to their success in related fields (Kobyzev et al. 2021
3.1 JOINT VARIATIONAL FLOW (JVF) ESTIMATOR
Theorem 2.2 states that moment matching yields optimal parameters of variational estimator to minimize Lemma but only applies for a Gaussian approximation. We introduce a new estimator, based on normalizing flows, that adds flexibility the variational estimator but that satisfies the moment matching optimality conditions.
Theorem 3.1 (Moment Matched Flow Distribution). Let Z f(X) and V =YYY) be diffeomor- phisms with Gaussian joint density given by:
(z.v = (1)
Uz
Σ
Σ Expz
22
Eg Exe
(11)
then the marginal and posterior optimizing the bound in Lemma 2. arr given by:
Mare
qmarg (X) N (f(X) Ms, Exx) ddet (x))))
Apost (X Y) = N (f(X) M2 Esjo) deet (Vxf(
where Pz/e = us + En - -e), Σzju = E.s - Egv Σ-EEE Euz
(12)
where the mean and covariance are matched to the moments of the target density p(X, Y): (ox,(((X),, = E.
E((U(()) pt,
(13)
Theorem 3.1 establishes that the optimal flow distribution is moment- to the target distribution, providing simple solution to parameters of the flow distribution. The joint density on X and Y also takes a convenient form given by the change of basis formula in Eqn. (10):
(x,r(X, =zz (J(X), ddet (V= |det (V (V)9())
(14)
-N ((12))
n, Σ Idet (V=J(X)) ddel (V))(Y))
(15)
where the Jacobian terms can be easily computed by construction of the flow. What remains is to train parameters of the flows minimize Lemma [2.]] The following lemma establishes that this bound has a convenient form, which can be optimized via gradient descent.
Lemma 3.2 (Flow Upper Bound), Let p p(X, Y) be an arbitrary target distribution and qx,y (X,Y) be the distribution of the form in Eqn. (14) with moment matched flow density (2,v V). Then the bound Lemma 2.1 is given by:
Imtp S log Idet (2TeE.:)1 1
2 log |det (2 [log |det (X)]( (16)
We see that the marginal and conditional entropy of qx.y can be expressed in terms of the underlying Gaussian entropy of qz,v (., V), which has a closed form resulting the first two terms of Eqn. (16). The last term is the log determinant of the Jacobian, which is easily computed by construction of the normalizing flow. The flows are trained to minimize Eqn. (16) and the parameters μ and are learned from moment matching at each step of the flow training. This results in the Joint Variational Flow (JVF) estimator:
Iv == Hp(9marg(XX) (X)) - Hplgpost (X()) = log Idet (2Tee - log |det (2Tee (17)
2
4

---

Published as a conference paper at ICLR 2025
We see that evaluating the MI does not require the log determinant term introduced by the flow which is the property of MI invariance under invertible transformations (Czyž et al.) 2023). While the resulting MI estimator captures only linear dependence (via Gaussian MI) does so in the flow distribution after application of nonlinear flows. The resulting estimator captures nonlinear dependence in the original X and Y.
3.1.1 STABLE TRAINING ERROR UPPER BOUND
The bound in Lemma 3.2 only sees the log determinant of f(X), and in practice we found that g(Y) tends to overfit during learning. To address this, we introduce new bound which is slight
modification to Lemma 2.1
Lemma 3.3. Let p(X, Y) be any model and yx.y (X,Y) be of the form in Eqn. distribution qz,v (Z, V), then the following bound holds;
14 with base
Live 11 S 2 (Z, V)) 2Epx [log ddet (Vaf(X))]] -2EEp [log lde (V v (Y)) +
(18)
For Gaussian base 92, v Lovv(ZZWV V) (n, Σ) the rightest bound of the form Eqn. (8) is given by the moment- moment-matched flow distribution with Σ = (f g(Y)) and takes the form:
Lvr - I| & log |det (2TeE)| pEx [log |det (V fX))] 2Ep Dlog ddet (V (Vyg(Y))]+C (19)
where C = 21l Y)). This bound is light when 4x.y (X, Y) - p(X, Y) almost surely.
Eqn. 19 contains both the log determinant terms of each flow and in practice enables more stable training. For all experiments in Section 6 we minimize /IvE with respect to this bound. Calculation of the MI estimator remains unchanged from Eqn. (17). The pseudocode for training JVF can be found in Appendix A.. in Algorithm
3.2 NEURAL VARIATIONAL FLOW (NVF) ESTIMATOR
The estimator Isvr assumes that X and Y have joint distribution of the form Eqn. 14. While the underlying joint is non-Gaussian it can still result in restrictions on expressibility. To address this we introduce neural network function approximators the base (flow) distribution. We add flow the marginal, Zmarg fmarg (X), and to the posterior, Zpost = fpest(X (X) with the following distributions:
(z((marz)) = N (1z, Ezs),
Qz (Z Y) = N (O(Y),E(Y)
(20)
where ul-) and E2-) are given by neural network function approximators, Then for two normalizing flows, Fmarg (-) and ()) the corresponding distributions on X are:
(X) =N (fmarg (X) Mg Eex) |det (VSmag(X)) ( Y) =N (fpos(X) p(Y), E(Y)) |det fVo (X))
(21)
(22)
Furthermore, the cross entropy of each distribution is given by:
Hplux(X)) Haz (az (mmrg)) )) Epx [log Idel a. (X)) Hylxxix(X (X (Y) Epx [- log (Spoa (X) ,(Y), E())) [log Idel (Na.sps(()))) (24)
(23)
where Hqz (92( (Zmarg)) is the entropy of the underlying Gaussian. The optimal marginal moments μı and Σgz are found via moment matching, but both flows and the posterior parameters optimized to minimize Lemma 3.3
This results in the Neural Variational Flow (NVF) estimator:
N
1
NVF = log det (2TeE..)1-
log det (V Fmanna (I:))]
2
N
(25)
log_N N (Fpos((x;) p(y:), E(y))
log |det (Vz
where (xx, y:)} ~ p and Σ. = Cov(Smare (X)). The log determinant terms do not cancel this case, since we apply separate transforms to the marginal and posterior, and so MI is not invariant under these transformations. INvr has both the flexibility of a flow-based distribution and nonlinear dependence modeling. INVE is capable of estimating much more complex distributions but comes the cost of needing to train both flows and neural network parameters. The pseudocode for training NVF can be found in Appendix A. in Algorithm 2
5

---

Published as a conference paper at ICLR 2025
True: pix,y)
JVG: g(x.)
NVG: ply]g(xly)
JVF: q(x.y)
NVE: ply)g(xly)
-1
11
D
Figure I: Multi-modal Density A four component Gaussian mixture model with nonlinear dependency is used show the approximating capabilities of each estimator. JVG cannot model the nonlinearity or multi-modality. NVG is able to model the nonlinear dependence between rows, but cannot capture the multi-modality. JVF is able to model the multi-modalit across the system but "hallucinates" extra modes to satisfy its linearity constraint. Finally, NVF models the multi modality and the nonlinear dependency of the GMM.
PRACTICAL APPLICATION OF VARIATIONAL MI ESTIMATORS
Estimators JVG, NVG, JVF, and NVF have different flexibility incorporated improve accuracy in estimating MI. Moving from joint to neural Gaussian increases modeling capability, but comes at the cost of adding a neural network to train and losing access an explicit joint distribution. Although the joint distribution is not necessary for MI estimation, it can be useful to have parameterized joint for applications such as generative models and classification. The second flexibility added to the estimators is flows which enhance the estimators capability of capturing multi multi-modality and nonlinear structures of the distribution but comes at the cost of needing to train two flows.
Figure illlusraate these properties on a four-component Gaussian mixture model. We see that JVG can only match the mean and spread of the data from the joint Gaussian assumption. NVG is able to capture the change of dependency from the top three modes to the bottom but is unable capture the multimodality for any individual y. JVF on the other hand is able to model the multi muuti-moddlity but still constrained to linear dependence of the transformed variable from the joint assumption. This means that JVF "hallucinates" alternate modes to satisfy its linear dependence constraint. Finally, NVF has both capabilities and therefore is capable of capturing both the multi-modality of the distribution and the non-linen non-linearitt of the dependence. It is important to note that for both the neural estimators no explicit joint (X, Y) exists, instead the joint is estimated by p(Y) (Y)((X YY) for illustrative purposes.
The choice of estimator depends on the prior knowledge about the underlying model. JVG is suitable for simple cases, while NVG captures more complex dependencies. JVF appropriate for multi- modal distributions with linear dependencies, and NVF handles both multi-modality and nonlinear dependencies. Table | summarizes the strengths and weaknesses of each estimator to guide the selection of the appropriate method.
Table 1: Flexibility of each estimator. Adding flows increases the variational density flexibility and adding neural parameters increases the dependency flexibility. Adding neural parameters however removes the modeling of an explicit joint distribution.
Estimator
Density
Dependency
Training
Joint
LJVG
Inflexible
Inflexible
None
Explicit
INvg
Inflexible
Flexible
YY)),, IYY)
Implicit
IJVF
Flexible
Inflexible
f(X). (1))g(9>
Explicit
INVF
Flexible
Flexible
(Y),, EYY), prior (X), Spost
Implicit
4.1 BAYESIAN OPTIMAL EXPERIMENTAL DESIGN
common application of MI is Bayesian Optimal Experimental Design (BOED) (Lindley 1956) where the goal is to maximize information through series of decisions, d, about a latent variable. X, given observations. Y. This model consists of an assumed prior on X, p(X), ). and likelihood function, p(Y d). Each decision is quantified with the amount of information it will likely give, which in the context of BOED called expected information gain (EIG)
BIG(d) Tx(X, Y) Hp(x) [p(X)) Hptx.9 (x(XYY,)
(26)

---

Published as a conference paper at ICLR 2025
The optimal decision is the that maximizes EIG: d argmax,EEG(d)) In sequential experi- mental design, a sequence of decisions is made d1...., dr and after each decision, an observation is made ...... Yr. The observations are used to update the prior, p(X), at the next time step using to posterior, p(X Yi...., Yt, di...... ......d, de), to include the observations from the previous decisions.
When making T decisions, the goal is maximize the total expected information gain (TEIG) (Ivanova et al. 2021)
Σ Bexx)pty)Xxdd [oog p(X dr, he-1) p(X XX] (1-11)
TEIGr(D) = (p(X)p(br))))
(27)
where D {d is the set of decisions and h2 = ((Y,, d:)? is the history of previously taken decisions and their corresponding observations. We utilize Iyv. INvG, Isvr, INVF as variational estimations of the EIG at each step in a greedy approach to maximizing TEIG. We learn the corresponding parameters for each model simultaneously with the decision which is approach introduced by Foster et al.) (2020) where they utilized the NVG estimator. We highlight this application for the use in Section show the utility of our estimators in practice.
RELATED WORK
The base MI approximation, mmtp. considered in this work is one of many possible variational approximations. Choosing only one distribution approximate with a variational distribution results in bound of MI. These are considered in alternative work (Foster et al. 2019) Dahlke et al. [023]
Poole al., 2019) but predominantly utilize only Gaussians as the variational distribution. Canonical correlation analysis (CCA) assumes joint Gaussian variational distribution which is closely related to JJyG- (Cheng et al., 2020) proposed an upper bound of MI that utilizes a variational distribution
similar to that of INVG. Recent work from (Butakov et al. 2024) has proposed a similar approach also utilizing normalizing flows which is closely related to Jyve but does not consider an estimator analogous to JNVF. Furthermore, (Dong al. 2025) considers using conditional normalizing flows in the context of sequential BOED.
An alternative approach to approximating the distributions is to instead use a discriminator, The Varadhan representation and the corresponding f-divergence representation lead to tight lower bounds MI, resulting in variety of approaches: DV (Donsker & Varadhan 1975). MINE (Belghazi et al., 2018), and NWJ (Nguyen et al.] 2010). Other methods use noise-contrastive estimation of the density ratio in MI which leads lower-bound of MI, called InfoNCE (van den Oord et al. 2019). Approaches that do not assume a discriminator or a density but instead compare k-nearest neighbors of samples have been considered called KSG (Kraskov et al. 2004). For
Donsker Donsker-Varadhan
thorough review of MI approximations see (Poole et al. 2019).
In the space of BOED multiple approaches have been taken select series of actions that maximizes MI. Multiple greedy approaches have been considered where each design is chosen to maximize instantancous EIG. The approach of simultancously optimizing the variational distributions and decision via gradient decent has been considered (Foster et al. (020). Other approaches, such as LFIRE (Kleinegesse al. 202T), perform ratio estimation of MI with Bayesian optimization for the decisions. Batch optimization is the approach where all decisions are made before testing time and stay fixed regardless of the observed data, such as MINEBED (Kleinegesse & Gutmann 2020). Finally, recent work takes a reinforcement learning approach where a policy-discriminato pair is learned (Foster et al., 2021 Ivanova et al. 2021 Blau et al. 2022] Huan & Marzouk] (0016). The discriminator plays the role of MI estimation for a select sequence of decision and observations while the policy is trained to make informative decisions based upon its history of decisions and observations.
EXPERIMENTS
We consider a wide range of experiments, from synthetic to application based. We start by comparing MI estimation capabilities against many common estimators: DV (Donsker & Varadhan, 1975), MINE (Belghazi et al. 2018), NWJ (Nguyen et al. 2010), InfoNCE (van den Oord et al. 20199. CCA
Murphyy 2023), and KSG Kraskov et al. 2004). The benchmarking tests are large MI estimation
7

---

Published as a conference paper at ICLR 2025
17.4
15
17.4
17.2
17-
MI Value
10
15.0
2000
2900
2000
JVF
NVF
InfoNCE
DV
True
JVG
NVG
MINE
NWJ
500
1000
1500
2000
2500
Gradient Steps
Figure 2: Large MI is difficult for many estimators to approximate. We that the four distribution based estimators, Tyva INva, JIve, and INVE, can accurately predict MI, but the four discriminative-hase methods, MINE, InfoNCE, NWJ, and DV, all converge around O (N))
(McAllester & Stratos. 2020 and a diverse set of 36 individual distributions with known ground truth MI values (Czy² al. 2023). We then look at the sequential decision making setting of location finding. All experiments were run on a high-performance computing cluster with nodes consisting of 2x AMD EPYC 7642 48-core (Rome) CPUs, 512GB of RAM, and NVIDIA V100S GPUs.
6.1 HIGH MUTUAL. INFORMATION EXPERIMENT
We begin with a well-known difficult case of estimating large MI. McAllester & Stratos (2020) showed that any distribution-free high-confidence lower bound on MI, such as the commonly used discriminative-ble based methods, estimated from N samples cannot be larger than O N). We create Gaussian where X,Y R with correlation p .95. The MI of this distribution computed as: T(X,Y) Y) = Dimx In(1 p²) 17.459. We take a large set of 75, 000 samples to train our distribution based estimators as well as a variety of discriminative based estimators for 3000 training steps using batch size of N 256. We utilize 80-20 split of training and testing samples of the total samples. Figure shows the best testing value for each estimator. We see that Jyv and Ivv converge nearly instantaneously to the true MI value and INvG and INVF both converge rapidly to high-quaiin qualityal estimates. For the discriminative-base converge to values drastically lower than the true MI.
be exactly
methods, we see that they learn slower and
6.2 MUTUAL INFORMATION BENCHIARK
Our next experiment a collection benchmarks from Czyz et al. 2023). They construct a diverse family of distributions with known ground-trute MI consisting of Gaussian, Uniform, and Student-T distributions that have MI invariant transformations applied to them (see Appendix B.2 for more details). We use N 1, 000 samples per dimension, with train-test split of 50-50. Figure displays the average MI over 10 runs for each experiment. Due to the large number of experiments we focus on selection here (see Appendix B.2 for all 36). For the majority of experiments we see that incorporating flow improves estimation accuracy, with JVF generally outperforming JVG. In the neural estimator case (NVF vs. NVG) the relative improvement more limited, with a few important exceptions discussed next.
Challenging cases for all methods are the spiral multinormal and the student-t distributions. The spiral multinormal MI is substantially underestimated by all estimators. but we see that flow-based methods (JVF, NVF) typically outperform their foowflow counterparts (JVG, NVG) as well as the citribution-free baselines. For the student-t distributions, all methods produce unstable estimates due to high-variance (due to space we omit standard errors). The high variance for both JVG and NVG comes from the fact that for many cases moments of the student-t distribution are infinite or undefined, limiting effectiveness of moment- noment-matching steps required by JVG and NVG. JVF and NVF can potentially address this issue with the flows, allowing the points to have finite moments, however this is particularly difficult to adequately transform the fat tails of the student-t with sparse samples. [Czy et al. (2023) introduce the inverse hyperbolic sin transformation means of normalizing the
8

---

Published as a conference paper at ICLR 2025
True M1-0,41 0.41
0.41
1.02
1.02
0.42 0.38
0.30
0.38
0.93
0:85
0.73
0.80
0.60
-0.02
InfoNCE 0.40
0.38
0.38
0.97
0.85
-0.15
NWJ -0.0 0.40 0.39
0.31
0.39
0.39
0.32
0.37
0.96
0.90
0.60
D. 74
0.50
0.63
0.84
0.68
0:01
0.09
0.10
DV 0V.0.40 0.38 0.30
0.38
0.40
0.35
0.38
0.96
0.90
0.87
0.92
0.50
0.78
0.85
0.85
-0,00
0.11
0
0.15
CCA - 0.41 0.39 0.19
3.
0.38
0.02
0.41
1.02
0.95 0,97
0.97
0.24
0 84
0.84
0.96
1.48
0.29
0.00
0.00
KSG 0.41 0.42
054-
0.32
0.42
0.41
0.42
0.41
0.96
0.92
0.27 0.24
0. 72
0.24
0.90
0.26
0.20
0.13
0.37
0.19
JVG- 0.42 0.42 0.40
0.19
0.34
0.39
0.02
0.41
1.01
0.95 0.97 0.97
0.24
0.84
0.83
0.96
1.25
0.23
0.00
0.00
JVF 0.42 0.42
0.19
0.42
0.42
0.03
0.41
1.01
1.01
1.04
1.04
0.25
0.84
0.86
0.96
0.90
0.04
0.00
0.00
NVG NVG-0.41 0.43
0.19
0.43
0.41
0.40
0.39 0.99 0.99 0.93 D. 94
0.56
0.
0.95
0.92
2.85
1.75
0.54
0.32
NVF - 0.42 0.40
0.20
0.41
0.41
0.39 0.39
0.98
0.94 0.92
0. 0.96
0.57
0. 0.79
0.96 0.90
0.36
0.39
0.29
Normal CDF @ Bivariate normal x
Bivariate normal x
Uniform x (additive nolse- 0.755)
wiggly Bivaratte normal 1111
Swiss roll 1
Multinrrmal x (dense)
Normal CDF @ Multinomaa x (2-pal))
Multinomal x (2-pair)
Normal CDF @ Multinormal 25 x 25 (2-pain)
Half-cube Multinormal 25 x 25 (2-paln)
Spiral @ Muliinomaal 3 3 (2-pair)
Spiral @ Normal CDF Multinomaal x (2-pai))
Spiral @ Multinormal 25 25 (2-pair)
spiral Normat CDF @ Multinormal 25 25 (2-pair)
Student- 2 (dof=1)
Student-t x (dof=22
Asinh @ Student-t 2 (dof=1)
Asinh @ Student-3 x (dofa2)
Bimodda
Figure 3: MI Benchmark runs against a selection of 19 different scenarios (for space: c.f. Appendix for all 36). Blue indicates an underestimate and red indicates overestimate with the magnitude indicated by saturation. Blank cells indicate that the method produced divergent estimates.
student-t tails, and thus creating finite moments. In this case we see INVF produces good results of the MI, where all other methods seem to fail, suggesting the estimator is useful for fat tailed distributions with appropriate data regularization. Finally, we notice the uniform additive noise as the only case which the discriminative-bassi methods outperform distribution-based methods.
6.3 LOCATION FINDING
We now demonstrate our methods in application to BOED discussed in Section 4.1 on the location finding experiment (Ivanova et al. (002T). In this experiment there are multiple sources with unknown locations, X, producing a signal whose intensity decreases according to the inverse square law. The signals from all sources are summed together to create a total intensity, Y that can be nosily observed from any location, d. Our specific setting contains two sources in two dimensions, for total of four dimensional X e IR. and a one dimensional observation, Y E R!. The goal is to make a sequence of T = 10 decisions di,...,d dr TR.2 that maximize the total expected information gain (Eqn. (27). We consider iDAD (Ivanova et al.) (02T) as strong baseline in this experiment as iDAD performs non-myopic decision making and so should represent an upper bound on what is achievable by the greedy (myopic) methods JVG, JVG, NVG, and NVF. iDAD also requires more total computation time as performs a significant offline training phase that other methods do not.
For our main results we see that incorporating flow uniformly improves MI estimation. When using the correct prior p(X)= = (o, I) the JVG estimator has the lowest MI estimate due assumption of a joint Gaussian distribution on X and Y. Incorporating flow (JVF) shows a sigmificant improvement. The neural estimators show more flexibility but we see similar improvements when incorporating flow (NVF) to the NVG estimator. All greedy methods require less total computation time than the iDAD. Note that neural methods benefit from GPU acceleration that was not applied in our implementation of JVG or JVF. One drawback of ammortized methods, such iDAD, is that they do not adapt to model mismatch during the test phase; a result of the well-known anmortization gap (Cremer al. 20188. We test robustness of each method by using mismatched prior in the test phase p(X) = N (), The relative performance of all methods remains identical to the matched prior setting. But iDAD exhibits 30% reduction in MI estimate whereas the greedy methods, which

---

Published as a conference paper at ICLR 2025
Table 2:
Total Expected Information Gain is reported for two settings of the location finding experiment. matched prior where the training prior matches the true prior and mismatched prior where the testing and training priors differ. We see that JVG. JVF, and NVG all fail to make informative decisions due to their limited flexibility and greedy decision making in both settings. We notice that iDAD performs well in the setting where its training data matches the testing data, hut suffers substantial amortization yiclds informative decisions while being comparatively robust to model mismatch. Total computation time lower for the variational methods when accounting for the offline training required by iDAD.
as the testing values shifted. NVF
Method
Matched Prior
Mismatched Prior
Training Time
Deployment Time (s)
Random
4.62 = 0.24
3.59 = 0.14
iDAD (InfoNCE)
7.54 + 0.19
5.29 + 0.21
8826.683
0.0256 + 0.0016
JVG
3.56 = 0.19
3.59 + 0.15
287 = 35
IVF
4.30 0.19
3.85 + 0.19
1730 = 207
NVG
4.61 = 0.19
3.99 + 0.20
186 =
NVF
5.10 + 0.20
4.748 + 0.18
1253 + 155
JVG
NVG
JVF
NVF
2
2
Figure 4: Location Finding Posterior estimations after single observation. We notice that JVG, NVG both estimate the posterior be Gaussian. The introduction of flow JVF and NVF increases flexibility meaningfully approximate posterior.
perform inference at test time either show reduction in MI (JVG) or smaller relative decrease (7%-13%)
We notice that in this experiment JVG, NVG, and JVF perform significantly worse than NVF. To shed some light on this, we can look the posterior prediction after making single observation (Fig 4 The result of this observation increases the belief that sources are located in radius away from the origin. Neither JVG nor NVG are flexible enough to capture this belief and instead produce Gaussian posteriors. Including flow JVF and NVF have increasesed flexibility to capture the non-Gaussian belief. We see that the additional degrees of freedom in NVF allows for a higher posterior belief in the vicinity of the true target locations, whereas JVF concentrates more away from the targets.
DISCUSSION
We introduced two flow-based variational estimators, JιVF and INVr, which are built on previous Gaussian estimators, IJvG and INVG. The inclusion flows increases the expressibility of the variational estimators. These estimators in general are just as accurate widely used critic-based estimators, but are not limited at their capability of learning large MI values. Furthermore, INve has the capability of estimating wide-tale distribution MI given appropriate data regularization, which notoriously difficult distributions to estimate. Finally. we find that these methods are effective for sequential Bayesian experimental design. In the BOED setting these methods show robustness to model mismatch not exhibited by policy-based approaches while also yielding significantly lower computation time when training accounted for.
Limitations Our proposed flow-based estimators have increased the flexibility of the distribution based variational MI estimators. However, there are still a few cases where we are unable to accurately model MI, such as additive noise in Section 6.2 Furthermore, each method adds increased number of parameters to train which can be costly, specifically at high dimensions. Future work plans to consider ways increase training speed remove the necessity of including neural networks.
10

---

Published as a conference paper at ICLR 2025
REFERENCES
Alexander A Alemi, lan Fischer, Joshua V Dillon, and Kevin Murphy. Deep variational information bottleneck. arXiv preprint 16 2016.
David Barber and Felix Agakov. The IM algorithm: a variational approach to information maximiza- tion, NIPS, 16:201, 2004.
Mohamed Ishmael Belghazi, Aristide Baratin, Sai Rajeshwar, Sherjil Ozair, Yoshua Bengio, Aaron Courville, and Devon Hjelm. Mutual information neural estimation. In International conference on machine learning, pp. 531-540. PMLR, 2018.
Christopher M Bishop and Nasser M Nasrabadi. Pattern recognition and machine learning. Springer New York, 2006.
Tom Blau, Edwin V Bonilla, Iadine Chades, and Amir Dezfouli. Optimizing sequential experimental design with deep reinforcement learning. In International conference on machine learning, pp. 2107-2128. PMLR, 2022.
Ivan Butakov, Alexander Tolmachev, Sofia Malanchuk, Anna Neopryatnaya, and Alexey Frolov. Mutual information estimation via normalizing flows, 2024. URL https::/arxiv..tt
abs 220/.2403..0 02187
Xi Chen, Yan Duan, Rein Houthooft, John Schulman, Ilya Sutskever, and Pieter Abbeel. Info- gan: Interpretable representation learning by information maximizing generative adversarial nets. Advances in neural information processing systems, 29, 2016.
Pengyu Cheng, Weituo Hao, Shuyang Dai, Jiachang Liu,, Zhe Gan, and Lawrence Carin. Club: A contrastive log-ratio upper bound of mutual information, 2020. URL https://arxiv.orr abs/2006 12013
Chris Cremer, Xuechen Li, and David Duvenaud. Inference suboptimality in variational autoencoders.
In International Conference on Machine Learning, pp. 1078-1086. PMLR, 2018.
PawelCzy2, Frederic Grabowski, Julia Vogt, Niko Beerenwinkel, and Alexander Marx. Be- yond normal: On the evaluation of mutual information estimators. In A. Oh, T. Nau- mann, A. Globerson, K. Saenko, M. Hardt, and Levine (eds.), Advances in Neural Information Processing Systems, volume 36, pp. 16957-16990. Curran Associates, Inc., http:/appee.iele/pppp//22222
2023. URL [tttps: wtwpeeeedddiinnn.nnuunee www.nnnuipp... www.uurrss.ccppaaaeeo wwtppperppppeeeeeeeeeeeeeeeee Fi flee/S 2tee36600eee777112226686 ott.::/eeeee..ttooooereeeererrr- Panper-ooofereeeeeeee .ddf
Caleb Dahlke, Sue Zheng, and Jason Pacheco. Fast variational estimation of mutual information for implicit and explicit likelihood models. In Francisco Ruiz, Jennifer Dy, and Jan-Willent van de Meent (eds.), Proceedings of The 26th International Conference on Artificial Intelligence and Statistics, volume 206 of Proceedings of Machine Learning Research, pp. (0262-10278. PMLR, 25-27 Apr 2023.
Jiayuan Dong, Christian Jacobsen, Mehdi Khalloufi, Maryam Akram, Wanjiao Liu, Karthik Du- raisamy, and Xun Huan. Variational bayesian optimal experimental design with normalizing flows. Computer Methods in Applied Mechanics and Engineering, 433: 17447, January 2025. ISSN 0045-7825, doi: 10. 10.10166..ma..224.117220
URL Couhhtt:// doi www.oioorgg/iee owg//0..//00.... 1016/
j. cma. 2024 117457
Monroe D Donsker and SR Srinivasa Varadhan. Asymptotic evaluation of certain markov process expectations for large time, i. Communications on Pure and Applied Mathematics, 28(1):1-47, 1975.
Conor Durkan, Artur Bekasov, lain Murray, and George Papamakarios. Neural spline flows. In H. Wallach, H. Larochelle, A. Beygelzimer, F. Aalché-Buc. E. Fox, and R. Gar- nett (eds.), Advances in Neural Information Processing Systems, volume 32. Curran Asso- ciates, Inc., 2019. ott.p///w.o.o////ooo.......ooao. 201 (

---

Published as a conference paper at ICLR 2025
Adam Foster, Martin Jankowiak, Elias Bingham, Paul Horsfall, Yee Whye Teh, Thomas Rainforth, and Noah Goodman. Variational bayesian optimal experimental design. In Advances in Neural Information Processing Systems 32. pp. 44036-14047 2019.
Adam Foster, Martin Jankowiak, Matthew O' O'Mear, Yee Whye Teh, and Tom Rainforth. uni- fied stochastic gradient approach to designing bayesian-optimal experiments. In International Conference on Artificial Intelligence and Statistics, pp. 9969.2969. PMLR, 2020.
Adam Foster, Desi R. Ivanova, Ilyas Malik, and Tom Rainforth. Deep adaptive design: Amortizing sequential bayesian experimental design, 2021.
Xun Huan and Youssef M Marzouk. Sequential bayesian optimal experimental design via approximate dynamic programming. arXiv preprint arXiv:1604. 08320, 2016.
Desi R. Ivanova, Adam Foster, Steven Kleinegesse, Michael U. Gutmann, and Tom Rainforth, Implicit deep adaptive design: Policy-based experimental design without likelihoods, 2021.
Durk P Kingma, Tim Salimans, Rafal Jozefowicz, Xi Chen, Ilya Sutskever, and Max Welling, Im- proved variational inference with inverse autoregressive flow. In D. Lee, M. Sugiyama, U. Luxburg, Guyon, and R. Garnett (eds.), Advances in Neural Information Processing Systems, volume 29. Curran Associates, Inc., 2016. URL
www. ccpppper Paperr pdf
Files/paper 201 wtlll//flle//eeeeeeeeeee https///p......aeeeee//e/eeeeeaa
Steven Kleinegesse and Michael Gutmann. Bayesian experimental design for implicit models by mutual information neural estimation, 2020.
Steven Kleinegesse, Christopher Drovandi, and Michael U Gutmann. Sequential bayesian exper- imental design for implicit models via mutual information. Bayesian Analysis, 16(3):773-802, 2021.
Ivan Kobyzev, Simon J.D. Prince, and Mareus A. Brubaker. Normalizing flows: An introduction and review of current methods. IEEE Transactions on Pattern Analysis and Machine Intelligence, 43 (11):3964 3979, 2021.
doi: 00.109//PAML222029999959
Alexander Kraskov, Harald Stögbauer, and Peter Grassberger Estimating mutual information. Phys. Rev. E, 69:0661336 13800 Jun 2004. doi: 10. 10.1100/seeeee..6...00 11t.11111/////......... 69. 066138
hhttps://ling aps
D. V. Lindley. On a measure of the information provided by an experiment. The Annals of Mathemat- ical Statistics, 27(4):986-100 1005, December 1956. ISSN 0003-4851.
David McAllester and Karl Stratos. Formal limitations on the measurement of mutual information. In International Conference on Artificial Intelligence and Statistics, pp. 875-884, PMLR, 2020.
Kevin P Murphy. Probabilistic machine learning: Advanced topics. MIT press, 2023.
XuanLong Nguyen, Martin J Wainwright, and Michael Jordan. Estimating divergence functionals and the likelihood ratio by convex risk minimization. IEEE Transactions on Information Theory, 56(11))5847-588) 2010.
Jason Pacheco and John Fisher III. Variational information planning for sequential decision making. In AISTATS, pp. 2028-2036, 2036, 2019.
George Papamakarios, Eric Nalisnick, Danilo Jimenez Rezende, Shakir Mohamed, and Balaji Lakshminarayanan. Normalizing flows for probabilistic modeling and inference. Journal of Machine Learning Research, 22(57):1-64, 2021. URL httLhtt::/ http://mmt...pppppre//s/e
19-1028. html
Ben Poole, Sherjil Ozair, Aaron Van Den Oord, Alex Alemi, and George Tucker. variational bounds of mutual information. In International Conference Machine Learning, pp. 5171-5180.
PMLR, 2019.
Tom Rainforth, Robert Cornish, Hongseok Yang, and Andrew Warrington. On nesting monte carlo estimators. In International Conference Machine Learning, pp. 4264-4273, 2018.
12

---

Published as a conference paper at ICLR 2025
Danilo Jimenez Rezende and Shakir Mohamed. Variational inference with normalizing flows, 2016.
Aaron den Oord, Yazhe Li, and Oriol Vinyals. Representation learning with contrastive predictive coding, 2019.
Nguyen Xuan Vinh, Madhu Chetty, Ross Coppel, and Pramod P Wangikar. Globalmit: learning glob- ally optimal dynamic bayesian network with the mutual information test criterion. Bioinformatics, 27(19):2655-7666 2011.
Sue Zheng, Jason Pacheco, and John Fisher. A robust approach to sequential information theoretic planning. In International Conference on Machine Learning, pp. 5936-5944, 2018.
A APPENDIX
A.1 PSEUDOCODE
We briefly provide an outline of each methods pseudo code. So of the notation is lightly changed in an attempt to highlight parameters of models. For example, for NVE, Fprir (X) is referred as fe (X) to highlight the parameters, Oprior. of the prior flow.
for i=l:k do
Sample {2r J=N ~ P(X);
Sample {y, j= ply kz):
# + Σ j=1 (fo(z,)² Io (y3) (T)
Σ N- Σj= (fo (zy) )T Iy (45) ( (Fo(z,)) 9e (11) up::
Loss log [2Te_ Eラ= log v=fe(z,)) E log (v))(
+ Losss;
- Loss:
end
/ - N Σj=
(2)))"
E N- e/= (fo (v, (fo(2))) "o()))))"" -
LJVF + log 2TEeDial - 1 log 2Te2inIV
Algorithm JVF Pseudocode
for i=l:K do
Sample {zz}j=1:N ~ p(X);
Sample {y5 (yj}i=1s 1 p(y Iz);;
jzz - N - S0ma (xrz);
Σzz + 1 E/=1 frrm (Iz) fo (2,,) -M2H4
Loss + log 12 N Ei= log V (Ir;)]
Σソ (log.N / (frrm (zz) Mo (v:), Eg(y:)) + log V V (zz))));
0prior + Oprior
V W Loss;
Opost + 0post a V ppd Loss;
の VV LLss::
end
uz f- 1
N Li=1 ferm
Σzz & N-1 ΣØ= Sama (zr)) (xz)² - 1zzzzi
IJVF + log 22 "yy log Va.fommn (xz)) +
1 Σ) (log N (Jopm (xz) (w:),, Ev(1)) log V (4;;))
Algorithm 2: NVF Pseudocode
13

---

Published as a conference paper at ICLR 2025
A.2 SECTION 2 PROOFS
Lemma 2.1. For any model p(X, Y) and distributions qmarg, (X), Gpost (X Y), the following holds: 1lmtp- 1| S min Hp Fmmaarg (X)) min H Hplopost (X Y)) +C
Amnng
Spoat
where C' = === Hp(p(X )) does not depend on gmarg or 9post. Further, the RHS is O
iff qmarg (X) p and Spost (X Y) p(X Y) almost surely.
Proof. We recreate the proof from Foster et al.
2019)
1mtp I
H (quarg (X)) Hp Apost
(X Y)) Hp(p(X)) + Hp(p(X Y))
(28)
Hp (x())) + (mmarg X)) + Hp(p(X Y)) - Hp(Apont (X Y))
(29)
( X) Amarg (X)) (T,(p(X Y) Apost (X YY))
(30)
I (x(( Imarg (x)) + (L(p(X IY) 9post (X |Y))
(31)
(x(((X)) Hplamarg (X)) Hp(p(X (Y) + Hp(Apont (X Y))
(32)
Hpp(marrg (X)) + Hp (ppost (X Y)) +
(33)
where - (x))(X)) - (p (Y))
To prove Theorem 2.2 and Theorem 3.1 rely results from Dahlke al. (0023). In Section A.4
we include the related theorems and proofs necessary for the results this paper as a convenience to the reader.
Theorem 2.2 (Moment Matching -Optimization). Let q(X, be a joint Gaussian density:
(X,Y) -N ( x -==
Pvy (
Pz
Exe
Eay Egy
-=
(34)
Zyz
then the marginal and posterior optimizing the bound Lemma 2. arre given by:
Gmarg (X) N (X Mz: Ems) Opost (X YY) N(X Mxllys Ex(y)> where
(35)
Msly == Pz Eny (y- -rr), Exly == Ens Eny E Eyn where the mean and covariance are matched to the moments of the target density p(X,Y):
(36)
(E)((,Y))))= pb Cov,(X,Y) E.
(37)
Proof. It suffices to verify the assumption of the posterior expected statistics being a linear com- bination of joint statistics (Eqn. (9h)) is satisfied. Recall the sufficient statistics of a multivariate Gaussian
X Y vec(X_X vee (T) Lvvee (V)
T(X,Y)
In this case Ti (Y) y and rY( vec(yy?) We now verify that the expected value under Apost (X Y') of each term in the sufficient statistic is linear function of T (Y) and r2(Y)
1.
E Spoee (ahneer(XX1 (XIY)) [x]] Muly = + Exs D (21 (x))
2. y/
E8 E (X ( [y]
3. .r::: T
IP Spont (XIY)
[xa" T
=Ez2p + Helll
=\zx 2) Eng 2 -T + (Mx + Ex 2y mY)) 2n uy
Exe Eny Eg x Maltin +...
Hz(y - H4 EEyy IT + Ezy - (y Hy))2 +...
Egy D (y pr (xr pr 5ー _
14

---

Published as a conference paper at ICLR 2025
4. zy
E Bpoon [xy"] =(Mo + Dy (y 2y)
=(1bo Exy E Hx )y² T + Exz Z yy
5. yy²
E Spoce (XIX) [vN']] = yu
So the statistics are linear functions of 7i (Y) = y and T2(Y) yu² ypnst (X Y) satisfies the conditions of Theorem A.5 and moment matching the joint q(x, X, Y) = [mm, E) yields the optimal Apost (X Y). Furthermore, the joint Gaussian sufficient statistics are
T) [x,Y; vec (eo(XXT), X vec(XXT), vec(YYT)]
and the sufficient statisties of marginal distribution are
T(X)= [X, vec(XXx)]
which simply the first and third sufficient statistic from the joint. Therefore moment matching the joint Gaussian trivially moment matches the marginal, giving the optimal 9marg (X))
Moment matching the joint Gaussian yields the optimal gmarg and Spost With this we can derive the formula for the MI estimate IyvG. We utilize another results from Dahlke
2023
Lemma A.1. Analytic Entropy
Let p(X) be any distribution and ((X) be in the exponential family with constant base measure, h(X)= = C, which analytically moment matched to p(X) and ((X) is empirically moment matched, then
Hp(g(X)) Hx)((X)) Hx()(X)) = HH(((X)))
(38)
Lemma A.I allow us to replace the cross entropy terms with simply the entropy the corresponding Gaussian, allowing for have closed form equation for the JVG estimation
Yva(X,Ys Y)
Hp (omaargg (X)) Hp (post (X
Y))
(39)
1
H Hamerg (Itarg (X)) H spoes (gpost (X (XYY) Y))
log [ log 12
2
2
(40)
A.3 SECTION 3 PROOFS
Theorem 3.1 (Moment Matched Flow Distribution). Let Z = f(X) and V = gYY be diffeomor- phisms with Gaussian joint density given by:
2
4z,v qz,v (Zv( = Vz.
F1=
Σ
(41)
Eve
then the marginal and posterior optimizing the bound in Lemma 2./ are given by:
qmarg (X) N (f(X) P2, Ex.) (z,(()), Apost (X Y) = N (f(X) M2| Esjo) vx)()))
where Pzje == us + EgeE (g pw). Σzjo E.s E.o E Evx
(42)
where the mean and covariance are matched to the moments of the target density p(X,Y* Y): Cor,(((X),,g()) = E.
(43)
Proof. Let 7 f(X) and V =Y(YY be diffeomorphisms with Gaussian joint density, then the joint X and Y can be found by simple change of variable (Bishop & Nasrabadi 2006).
(x.((X,V) = =v ([((2)]
((())
E F2
(44)
15

---

Published as a conference paper at ICLR 2025
We now must show that yx,Y (X, () satisfies the moment matching condition
E (XI)) (T(X, Σ (i(()):(())
(x)r is in the exponential family with base measure M(X, I and T(XY)) = ()) vec(s(X) vee(f(X)gg()))) vec()(Y)g(()))) which are the sufficient statistics we must consider for Theorem A.5 Furthermore, we can derive the posterior
f(x) [(())
N
1.
(xIx(X1V) (x)x
(Y)
[(X)
SN
V de
[(())
(
[(()
14
[x((())
(V())
[10
((())
1
V=f(X)) de
[
(9())
[vf(())
SN
1, dz
(9())
[((X]
[(())
2
Vxf(X))
N (g lAv, Evv)
(f(X) |zz ExIv) [xf(X)l
where PAzlV= = P22 Σzv (y - -v) = uz + Σzv Ez (g() - pv) and Σzl = Σzz
Σzv MV Evz. Using these, we can easily check the linearity condition of each T(X,Y)
1. f(X)
(XXI (f(X) = / N (F(X) (( f(X)dr
f (Z zd=
=Zz = p2 + Σzv EVV pv)
2. 9(Y)
E FXX ((()) = g(V)
3. f(X)((X))
(xII [f())((x)]] / N (1(X) EzIV) (XX(X))
= = - N (Z JUZIV. Ezi) 2zTdz
=ΣZIV + IZIVNZZI
-Ezz Σzo Exy +...
(Uz + Dvv (g(Y) - -uv))(uz + Ezv Ezv - -T
=Enz - Ezv Exy ETy + =zzzz +...
"z(o(Y) E. + (
Ezv 1et (v) S.
5T
16

---

Published as a conference paper at ICLR 2025
4. f(X)g(Y))
(x [f(X)g({)]]
-x (
(uz + Ezv (o) - (vv)
9(Y)²
-wz- -izv Evv g(M)" + Ezv EVbo(())(("
5. (())g())
(x =
Since these are all linear functions of Tı (Y) = g and 72 TY() g())g(Y)) then qxlr YY) satisfies the conditions of Theorem A.3 and moment matching the joint gz,v V) = N(m, E)
yields the optimal 4xlr (X Y). Furthermore, the joint sufficient statistics are
T(X,Y) = [S(X), (V), vec(f(x)s(x)))) vec(f(()g(Y))") veo((
and the sufficient statistics of marginal distribution are
T(X) [f(),, vxo(f(
which are simply the first and third sufficient statistic from the joint. Therefore moment matching the joint trivially moment matches the marginal, giving the optimal (X()
Lemma 3.2 (Flow Upper Bound), Let pl (0XXY Y) be an arbitrary target distribution and qx.y (X, Y) be the distribution of the form in Eqn. (14) with moment matched flow density qz,v (Z, V). Then the bound in Lemma 2.1 is given by:
mm-pp S log E2 + log [2TeDzJo| 2EEx [lo [x](()]]++ +
(45)
Proof. The upper bound in Lemma 2.1 is
Mm+p - I s Hxlux(X)), + Hxlaxix (Y)) +C
This means we need access the cross-entropy terms of the flow distribution
Hp(ax(X)) / 1(x) - Xx
(46)
=H, (Z)) - Ep flg V vxf( = Hzz (2)( - E, [log vxf(x)]
(47)
-2 log (2Te - Ep [log] V xx(()]]
(48)
(X)) = / p(x, Y) log logaxir( Y)dr
Hx(exi((X) (Y))
(49)
/ p(x, Y) log (f g(Y)) (zf((X) de
(50)
= -Hplaziv (Z V)) - Ep flog V [x)]
(51)
Haziv (12lv (Z V)) E, llo / vx/(X)I]
(52)
log |2TeX Ep [log V [X)]]
(53)
2
Here, we recognize that (g v (Z) V)) is a moment matched Gaussian, so Lemma A.1 applies to the cross-entropy. With the cross entropy terms, we can now plug into the upper bound
1m - / < HH (yxxx(X)) + H Hplaxi (x( Y)) +
(54)
1
log 2xee.x1+ log [ 2Ep Dlog vxf(X)] +
(55)
2
17

---

Published as a conference paper at ICLR 2025
The equations for the cross entropies, Equ. 488 and Equ. (53) allow for to explicity write the form of LsvE
Tyr (X, -Hplax(X)) Hp V))
(56)
= log 2TeeDaal - Ep [log V fXx]] - log 2 + Ep [log Vxf(())] (57)
2
log 2
log 2
(58)
2
2
We see that this is simply the MI of the Gaussian q2,v (2.v V) which is the propety of MI invariance under invertible transformations (Czy2 et al., 2023).
Lemma 3.3. Let p(X, Y) be any model and qx.v (x,, of the form in Eqn. (4 with base distribu-
tion qz,v v Z, V), then the following bound holds:
LIF 1 S 2H, (2, px [log Vxf(x)]] pk [log
+C
(59)
For a Gaussian base 92.v(Zz (7.v(Z,V) = NN,, Σ) the tightest bound of the form Eqn. (18) is given by the moment-matched flow distribution with E = (J (X))) and takes the form:
|LVF 1 S log [2TE21]
2Epx [log [x](()]]] - 21ppy [log] (Y)( +
(60)
where C -2H,(p(X,Y)) This bound is tight when qx,y (X, Y) = p(X, Y) almost surely.
Proof. Since JVF assumes and parameterizes a joint distribution (X Y), we have access the symmetric definition of MI
lxv Y) -Hplux Hplaxiy (1))
(61)
-H,lax(X)) Hplax, (X,Y)) H,(qr
(62)
= (Hplor(V) - H [=)] I=)==
(63)
For this proof, we simply use [xvr(X, Y) and (Y, X) to differentiate between the two symmetric forms of JVF but do not necessarily hold this notation constant elsewhere. So we will apply Lemma 2.1 twice
Y) 22 L -(X,
(64)
1LIF((X, 1) I(X, + - X(Y,
(65)
<Halax(X)) H1l(xp (X (Y) + Ci
(66)
H2lay (Y)) + (y x)) + +C2
(67)
where C= + C2, C1 (x,
(68)
- (x Y)), and C2= -H= - x))
which come directly from bound applied to each term. We also were able to combine H, (gx (X)) and
Hy x (y x))) into Hp(px.y(X, Y)) and likewise for Hp (Y)) and (
Y)). Since
this was two applications of Lemmal 2.1 we know that moment matching is optimal for bound and tight
when qx (X) = p(X). (Y) p(Y) ). ( Y) = p(X Y), and xrlx x) = ply r).
Finally, deriving the explicit form of the bound, get
-0x, -XX (XX,) + +C
(69)
-2E, |- (X, +C
(70)
= 21E,, | log q2.v (f (f(X). VV.x(
(71)
-2H,(02,v(Z,, V)) 2E, [log
(72)
=log [2Te_]] - 2Ep [log (VI((XII 2E, flog +
(73)
Where again, we used Lemma A. since gz,v (Z, V) is moment matched Gaussian.
A.4 EXTERNAL SUPPORTING PROOFS
To prove the optimality of moment matching for estimators in Theorem 2.2 and Theorem relied on the generalized results from Dahlke et al. (2023). We include related results from that paper in this section for convenience of understanding the proofs of Theorem 2.2 and Theorem 3.1
18

---

Published as a conference paper at ICLR 2025
Theorem A.2. Let qm (X) be in the exponential famity with statisties T(X). then for any p(X), the
optimal mary is given by moment matching:
R (X) [T(X)= = Px( [T(X)
Proof. Since (X) constant in 4m we have,
argmin Hp(gm(X)) argmin Hp(gm (X)) Hx(X) argmin (L(((())
(74)
9m
4m
4m
It is a known result, as proven in (Bishop & Nasrabadi 2006), that for exponential families, (x) (T)) Bx(x (T) minimizes KI(p(X) ((x))
Lemma A.3. (x Y) takes the form of
(xlx 4x1x (X 9X,Y Y;n)
(75)
(Y;Y:n)
where 9x.Y (X, Y:)) is in the exponential family, then the minimization of H, (x(o)(X(Y)) occurs when
B (x, (T(X,))] - (T(X,Y))
(76)
Proof. The goal is to minimize (X where qp (X is generated from q(X, Y;m) in the exponential family. We will find the minimizing parameters of this distributions. We appeal to the property of exponential families that 0 A(n) = Ex(x,) (T(XY))
än
① On (X,(1x(X11)))
a - on log = = p(xx v) á On ((::n)
log ((n)
(77)
=- f-f --JOX.Y) () 8 (108(1(X, T(X,) - A(1) log (a(Y:; (n)xx:))) drdy 8m
(78)
=-f -/MXX.) T(X,Y)- On A(n) on
log (e(Y::1)) dody
(79)
--Exx [X( E [T(
I (XX,)) ((Y:n) on
(fM(XYY)) (7"T(X', A(1)) de') dedy
(80)
Exx [XX, ExXx.) [T)+
/ p(x,)) ((x:n) (fax' ()(X,))- 2 Am)) de' dadyy
(81)
on
- Exxx [TX, EacA [T(X,
I /MXXY) (FOXIIY (T(X', Y) B, (T(X,)) dz') dædy
(82)
Ex [T(X + Epiry [Bope [T(X, Y)J
(83)
The zero derivative yields ExxX.Y) (T(XYY) Ep(r) (x))) [T(X, Y)] which the stationary condition. It now remains to show that the objective convex in Expanding the form of Hp[g(X Y)) we have the objective,
min -Ep
- A(r) log(g(Y;n)))
(84)
The term p''((X Y) is linear in n,. Convexity of A(r) in is a standard property of the exponential family, however we will show constructive proof that A(1) (og(g(Y;<)) is convex using Hölder's s inequality. Let 77 Am + -)a22 where À E [0. 1] and m, 7/2 in the convex set of valid exponential
19

---

Published as a conference paper at ICLR 2025
family parameters of q then:
A(1) + log(o(Y; :))) A(r) + log () (xxxx
A(i)) dx
(85)
A(v)) + log (exp(-A(1))
T(X,
(86)
los ((x(X.)))sp(
(87)
= log (xxx(
(M Y) TXX, X)), de
(88)
- X) log (/ (xxx
S X log
(89)
X(A(m) log (Y: m)) (1 () log ; m2))
(90)
Thus convexity holds in 1 and the stationary conditions globally optimal.
E pxx, [TXX, Y) Ep(r) (XI)) [TX,
(91)
Theorem A.4. Let q( X, Y) be in the exponential family with sufficient statistics, T(X,Y) (Y)) () where T(X) are the sufficient statistics dependent only on x, +YY) only on y, and 7(X, Y) on both. Further, let the posterior expected statistics be a linear combination of marginal statistics as in,
k
(XIx) [T(X,Y)] - (l1))))
(92)
where TY(Y) is the ith component of T(Y) and ge(T)) are functions of only the parameter Then, the optimal variational distribution, Up, for Ipou defined by joint moment matching: Ep(x, E Y). Eg (XX,))
Proof. From Lemma
know that Ex))) [T Y)] Epry) E (x [T(,, Y)] is the
optimality condition. Let now show that the condition in Eqn. (92) implies that joint moment matching satisfies the optimality condition of Eqn. (76)
Ex(x,) (T(X =xg( (X)
(93)
R (XI)) (T(X, = Ryy { {a(mmm)
(94)
Σ (
= Σ (i()) [x()))((()())
(95)
* Σ (
-E EEp)
= Ep(r) (XI)) (T(X
(96)
So with Lemma Le3 and the assumption of the posterior expected statistics being a linear combination of joint statistics (Eqn. (92) results Eyx,y) [TX, Y)) ( (T, Y)] being the optimal conditions.
Theorem A.5. Let qmarg (X) and q(X,Y) be exponential family distributions where the joint sufficient statistics are T(X,Y) = [r(X), (YY)) D,T(X,Y)) Y)J" and T(X) are the sufficient statistics dependent only T(Y) only y. and 7 (X,Y) Y) both. Further, let q(X,Y) Y) satisfy the linear conditional expectations property
E Expcst ( [(X,Y)) = E.(m)))))
(97)
where y are the natural parameters of q(x, and ge(r)) are arbitrary functions dependent only on the natural parameters. Then, moment matching the joint q(X, Y) and marginal qm (X)
[T(X, Y)] (T(,,)
20

---

Published as a conference paper at ICLR 2025
) [(()) = Bypore (X) [X)]
yield optimal (post Y) x q(X, Y) and (marx(X) that minimize the bound Im+p Lemma
2.7
Notice that the variational distributions Om (X) and Op(X Y) need share common joint q (X, Y). So, let us use different natural parameters, and 1/2, for each (i.e. (X) (XX;; m)) and 9p (X
Proof. We break this down into the two cases of Theorem A.2 and Theorem A.4
(( We optimize the bound in Lemma 2.1 with respect to both natural parameters, beginning with n/:
0 Ōnp1
(-Ep(x,) [log mm) + log q(x Y: C)
PX(X,Y [log ((;m))
(98)
dm
This is exactly the condition in Theorem A.2 which we know is solved by moment matching the marginal. Likewise, for 1722
0 (-Ep(x,)) [log (X::)) log g 127))] +C) 07z
0 p(Y)) [log q(X Y272)] (99)
d7z
The above is the start of the proof for Lemma A.3 in Eqn. (77 and along with Eqn. 92 in Theo- rem A.4 get that moment matching the joint finds the optimal Op- Therefore, the optimization of
Lemma 2.1 simply reduces to moment matching the marginal and the joint.
B EXPERIMENT DETAILS
All code can be found at httpp::///tt htttps://github.com. wcw.b.cco/aaleaaei.... Ftw.oddaaake//oool./...
B. HIGH MUTUAL INFORMATION EXPERIMENT
We consider a synthetic test case where we can create a distribution with high mutual information. Consider joint Gaussian where X, Y E 15
=
x
|
Mx by
E =
Dxx Enyz
Eny Egg
(100)
where = 0 is the zero vector, Σgz = Eyy = I, and Exy Eye = pl. Then I(X,Y)
Dimx In p²) so we are free to choose the value of MI by changing the dimension or by 2 changing the correlation constant. We select p .95 to achieve MI of [(X, Y) 17.459. We train all estimators using batch size of N 256 for total of 3000 steps to ensure all estimators converged.
The critic based methods (DV, MINE, InfONCE, and NWJ) all use neural network with two hidden layers of the structure [16. with ReLU activation functions. Czyż al. (2023) utilized and batch size of N = 256 which we kept the same.
The normalizing flows utilized in JVF and NVF are Rational Quadratic Splines Duurkan al. 2019) which learn 128 knots to parameterize the spline. The spline is bounded between -8 and where linear outside. The knots are learned from neural network of size [8, 8] with ReLU activation functions. To prevent overfitting of the flows, we utilize dropout with rate of .2 and I.2 Regularization with factor of 10-5
The neural parameters used in NVG and NVF are the mean and variance parameterized by two neural networks of size [16, 8] with relu activation functions. The output of the neural network parameterizing the covariance is lower triangular matrix to approaximate the cholesky decomposition of Σ where the diagonal elements have a Softplus activation applied to them to ensure the resulting matrix is semi comi-positive definite. These parameters did not have dropout performed but did have the same L2 Regularization with a factor of 10 5 NVG, JVF, and JVF all utilized a learning rate of .005 and reduced learning rate by factor of 11 test loss if improvement was over 250 testing steps.
21

---

Published as a conference paper at ICLR 2025
Than 0.43
0.2.0.433.0 1.73
0.48
0.41
0.41
0:11
0.22
0.72
9:29
5.41
0:50
1:28
1:02
1,62
1,62
1.00
1.09
1,62
1,50
1,2
(EN6 -2.27 1:75
1220
(HbNCE- -2.0 0:38 3,89
0:28
1.27
0:3
03:
002
0.2
0.30 1,64
0.20
1.22
6.92
0.3
315
CSA-1A: 5.39 2/61
0:29
D.4L
1.90
1.52
1.5
67
033
024
1SG-1AL 244
545
MI
6.39
641
637
323
1.33
639
02
0.0
0.33
02
011
092
02 0.24
024
536
929
035 220
Jon T22 883 LAT
2
2
B U0
1
0.3
Zac soseetotnn
enmene esdees
vopi en
sooe meme on snce
Figure 5: MI Benchmark runs against a selection of 36 different seenarios. Blue indicates an underestimate and red indicates on overestimate with the magnitude indicated by saturation. Blank cells indicate that the method produced divergent estimates.
B.2 MUTUAL INFORMATION BENCHMARK
Czy² et al. (2023) created a large selection of MI benchmark tests with gound truth values via MI invariance properties under certain transformations. We will give a brief description of the experiments but in-depth information can be found in their paper.
Bivariate Normal: simple 1 X dimensional Gaussian with p= .75
Uniform Margins: The Gaussian CDF, (-). is applied to Gaussian variables X and Y, the resulting distributions are marginally uniform but not jointly uniform. Denoted as 'Normal CDF P'. Half-Cube Map: The transformation F(X)= = |z|3//2 is applied to Gaussian variables with the goal of lengthening the tails. Denoted as "Half-cube @ P'.
Asinh Mapping: To shorten tails, the inverse hyperbolic sine function asinh(X) = log X X x2 applied. Denoted as AAsih @ P'.
Wiggly Mapping:
To model non-uniform lengthscales, the mapping F(X) = x Σ ag (wee
applied. Denoted as 'Wiggly @ P' Bimodal Variables: In the inverse CDF of two component GMM is applied to X and Y to create multimodality.
Additive Noise: The random variable X Uniform(O.) 1) is considerd along with noise N Uniform (c., c), then Y = x N. Denoted 'Uniform (additive noise=c).
Uniform(0,) 1) and the Swiss roll embedding is performed on X. X and Y are jointly Gaussian distributed with Cor(X,, Y2)= Cor(X,, Y1) = .8 and 0 correlation everywhere else.
Swiss Roll Embedding: X,
Multivariate Normal (2-pair):
Multivariate Normal (Dense): X and Y are jointly Gaussian distributed with all off diagonal correlations set to 5.
Multivariate Student: A multivariate Student-T distribution with degrees of freedom v. Spiral: The transformation F(X)= exp(vA) ||: [²) where A is a skew-symmetric matrix which "mixes" the dimensions and e is the rate at which the mixing occurs. Denoted "Spiral P..
The dimensions X and Y range from to 25 to give a wide selection of applications. Each estimator is given 1000 samples per dimension with a train test split of 50-50. Each estimator is evaluated on the testing data every 250 gradients steps and early stopping is performed if the estimator test loss has not decreased within the previous 500 steps. All data is pre-processed by centering the data and scaling each dimension by it variance as performed in (Czy2 et al. 2023). The same network structures are used in the Large MI experiment.
B.3 LOCATION FINDING
In the this experiment, we have 2 hidden objects in R² which we wish to learn their locations X (X1, X2). Each source emits a signal that decays with respect to the inverse square law. From
22

---

Published as a conference paper at ICLR 2025
and position, d, we can observe the superposition of all signals together
(xX, d) =6 m | dll
(101)
m || dll
where b is the background noise and m 10-4 controls the maximum signal. We then normally observe the log intensity of the signal from the likelihood
logy y X, ~ (N(oo((((( d))
(102)
where sigma controls the noise in the observation. We condsider two settings for our experi- ments, matched prior and mismatched prior. In the case of the matched prior, every method assumes the prior p(X) = (O, I) which corresponds to the true testing prior where testing æ are sampled from. In the mismatched prior setting, every method still assumes that p(X) NN I) but the true testing z are sampled from (1, I). In practice, we often assume convenient prior but that prior likely does not correspond well with the real world adapt to this mismatch with minimal loss information to still provide informative decision across a sequence of decisions. We notice that in this experiment JVG, NVG, and JVF perform significantly worse than NVF. To shed some light on this, we can look at the posterior prediction after making single observation (Fig The result of this observation increases the belief that sources located in radius away from the origin. Neither JVG nor NVG are flexible enough to capture this belief and instead produce Gaussian posteriors. Including flow JVF and NVF have increasesed flexibility to capture the non-Gaussian belief. We see that the additional degrees of freedom in NVF allows for a higher posterior belief the vicinity of the true target locations, whereas JVF concentrates more away from the targets.
ideal BOED experiment should be able to
23


Reviews:
[{"review_id": "zpUzzGB5Ow", "reviewer": "Reviewer_Gj86", "summary": "This paper introduces two new classes of variational estimators, Joint Variational Flow (JVF) and Neural Variational Flow (NVF) via combining normalizing flows with moment-matching based joint Gaussian estimators. The inclusion of flow has significantly increased the flexibility of the estimator, enabling more accurate approximation for complex interactions. And moment matching strategy allows efficient computation. The author(s) compared the proposed JVF and NVF with popular MI estimator on multiple synthetic datasets and discussed their pros and cons.", "questions": "Can the author(s) clarify how their proposed methods can handle scenarios where the data distribution resides on a lower-dimensional manifold?", "limitations": "", "rating": 5, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "* The introduction of normalizing flow greatly enhanced the flexibility of Gaussian variational MI estimator \n* To mitigate the stability issues observed for the original JVF estimator, author(s) proposed an alternative formulation where the Jacobian of both flows appear in the objective\n* The proposed estimators also give distribution approximation of p(x,y), where people can easily draw samples (for prediction or other simulation-based tasks).\n* Adequate coverage of relevant literature", "weaknesses": "* [Minor] Classification of MI estimators. In the introduction, the author(s) discussed critic-based MI-estimators and variational MI-estimators. However, many of those “critic-based” estimators (e.g., DV, NWJ, InfoNCE, etc.) actually belong to the variational family, see [B. Poole 2019]. These “variational” estimators described by the author(s) should be considered as parametric estimations. \n* [Major] Practical value of the proposed estimators. In practical applications, people often seek to optimize the MI, instead of just estimating its value. For example, one major application of MI optimization is representation learning, where the MI between two related data views are optimized — it is the learned representation that matters, not the estimated MI. “Critic-based” MI estimators are heavily used in such scenarios, given their simple formulation and strong performance for real world complex data. To the best of my knowledge, normalizing flow based approaches rarely moved beyond academic settings. This paper also did not provide any empirical/theoretical evidence that proposed flow-based estimator can deliver competitive performance with popular alternatives such as InfoNCE in representation learning scenario — given ICLR is the International Conference on Learning Representations, the author(s) have to check this box. \n* [Major] Experiments. The experiment section needs to be strengthened. Currently all the experiments are synthetic data, which is inadequate to demonstrate the robustness of the proposed methods to complex real world data. Also while the author(s) mentioned two variants of JVF, there is not empirical comparison between the two. For the location finding experiment, the author(s) should show the comparison posterior distribution of source points after all ten observations.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "SLsnXooNWD", "reviewer": "Reviewer_CMGX", "summary": "This paper introduces a new approach to estimate Mutual Information (MI) using flow-based variational methods. Existing methods like variational and critic-based approaches use restrictive Gaussian approximations and struggle with large MI values.\nThe paper proposes two main methods:\nJoint Variational Flow (JVF): Extends the Joint Variational Gaussian (JVG) estimator by using normalising flows while maintaining efficient moment-matching updates\nNeural Variational Flow (NVF): Builds on the Neural Variational Gaussian (NVG) estimator to provide a more expressive distribution.", "questions": "I would be happy to increase my score if the authors can improve the experiments section by 1.) Testing on at least one more benchmark setting used in the prior works. 2.) Add ablations focusing on the impact of number of sequential decisions (with error bars please).", "limitations": "", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "Originality:\nThe paper creatively combines normalising flows with existing variational MI estimators (JVG and NVG).\nThe approach to stabilising training through the modified bound in Lemma 3.3 is interesting\n\nSignificance:\nThe approach shows robustness to model mismatch which can be important for real-world applications.\nContributions can extend beyond MI estimation to general flow-based density estimation\n\nQuality:\nThe paper provides thorough theoretical development with clear proofs and lemmas.\nHonest discussion of limitations and failure cases\n\nClarity:\nExperiments section has a clear motivation for each of the toy test", "weaknesses": "The experiments although quite interesting, fail to provide sufficient evidence for practical benefits of the proposed approach. The argument for practical benefits hinges on the BOED application. Unfortunately the paper fails to provide convincing experimental evidence for the BOED application. More specifically, the only test is on a location-finding scenario while previous works like Foster et al. (2021) and Ivanova et al. (2021) evaluate on multiple BOED benchmarks:\n\nThe presentation of the experiments is a bit sloppy as there are no error bars or ablation studies on key parameters such as the dimensionality of the problem and the impact of number of sequential decisions\n\nIn my view these type of experiments are more suited for AISTATS than ICLR. For ICLR, I would encourage the authors to provide more realistic benchmarks similar to those used in Foster et al. (2021) and Ivanova et al. (2021).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "S2LZ0B8YJU", "reviewer": "Reviewer_uioR", "summary": "The paper introduces two classes of variational estimators for mutual information that utilize normalizing flows to extend Gaussian-based methods. These estimators are shown to maintain guarantees while improving the flexibility of the variational distribution. Experimental validations demonstrate the efficacy of these methods on mutual information (MI) problems and in applications like Bayesian Optimal Experimental Design (BOED).", "questions": "How do the proposed estimators perform under different levels of model misspecification or with datasets characterized by heavy-tailed noise? Are there particular settings or types of data where the performance of these estimators might degrade significantly? For example, 4th moment bounded heavy tails or 2nd moment bounded heavy tails.", "limitations": "", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "Theoretical Innovation: The introduction of flow-based variational estimators represents a theoretical enhancement over traditional Gaussian-based methods. The authors effectively integrate normalizing flows, which are known for their flexibility in modeling complex distributions.\n\nComprehensive Validation: The paper presents a robust set of experiments that not only compare the proposed methods against existing estimators but also demonstrate their superiority in scenarios where others fail, particularly in large MI estimation problems.", "weaknesses": "Assumptions Limitation: The estimators assume the availability of a joint distribution, which might not be practical or accurate in all real-world scenarios. Although one of the proposed estimators, NVF, relaxes this assumption, the general applicability of the method under different conditions remains underexplored.\n\nRobustness Concerns: The discussion on the robustness of the models, particularly under model misspecifications and with non-Gaussian noise (e.g., Student-T distributions), is not thoroughly addressed. The capability of handling Non-Gaussian (especially heavy-tailed distributions) is not comprehensively validated.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "U9UAsyqa1p", "reviewer": "Reviewer_9DeW", "summary": "This paper focuses on developing a variational bound for the error between a model for mutual information and the true (unknown) mutual information. Through this bound the framework can be used to estimate the mutual information as well as to draw samples from the underlying distribution. A key element employed for achieving model flexibility in modeling the distribution (or samples therefrom) is the use of normalizing flows. This setup is coupled with multiple forms of moment matching, as applied to the approximated first and second order moments of the data, estimated from samples. Neural networks are employed in multiple ways: They are employed within the normalizing flows, and they (for one of the models) are used to model the mean and covariance within conditional distributions (employed within the NVF estimator). An interesting application of this technology is employed within adaptive Bayesian experimental design. Extensive and encouraging experimentation is shown.", "questions": "I enjoyed reading this paper. It was generally well done and well motivated. I have some suggestions for possible improvements.\n\n1. For such a well thought through concept, I think aspects of the writing could be improved. My concern is where to use upper case X (or similar) for a random variable, and lower case x for an instantiation of that random variable. Just below Eq. 19, you write C=2H_p(p(X,Y)). And you use similar uppercase variables in Eq (18) for q_{Z,V}(Z,V). But, in many (most) other places you use lowercase variables, such as in (1), (2) and many other places you use H_p(p(x)) and h_p(p(x|y)).\n\nAt minimum you need to be consistent, and this use of notation is not consistent. I think the notation in (18) and below Eq (19) is better.\n\nIf we take Cover and Thomas, \"Elements of Information Theory,\" (2nd Edition),  as a guide -- as this is a \"bible\" of information theory:\n\nMutual information is I(X;Y) : importantly, (X,Y) are uppercase, because these are RVs, not instantiations of RVs. \nEntropy is H(X): Importantly, X is upper case.\n\nI know that sometimes you may wish to emphasize what distribution you use for these measures, so (like under Eq 19) I suggest you use H(p(X,Y)), but most importantly: use uppercase (X,Y)\n\nAs an example, your Eq 1 should (in my opinion) be H_p(p(X))-H_p(p(X|Y)). This notation is consistent with Cover & Thomas.\n\nIn (2), you do have I(X,Y). But you have H_p(q_{marg}(x)), for example. I think this should be H_q(q_{marg}(X))\n\nThis issue with notation is everywhere in the paper. Usually it is not consistent with Cover & Thomas, and sometimes (rarely) it is.\n\nIn my view, this is a very nice paper. However, the above is very distracting for readers who have a background in information theory (largely your readership). I think you want to write this such that it is consistent with the information theory literature. I suggest you carefully go through and revise. As examples:\n\nq_{marg}(x) -> q_{marg}(X)\nH_p(p(x|y)) -> H_p(p(X|Y))\nq(x,y) -> q(X,Y)\n\netc etc. Above are just a tiny subset.\n\n2. Your bibliography is extensive. However, you missed\n\nPengyu Cheng, Weituo Hao, Shuyang Dai, Jiachang Liu, Zhe Gan, Lawrence Carin, \"CLUB: A Contrastive Log-ratio Upper Bound of Mutual Information,\" ICML 2020 \n\nThat paper also uses sophisticated NNs and addresses a similar problem. Please discuss how your approach relates and compares to this approach.\n\n3. In the Abstract, you write \"In this paper, we introduce two new classes of variational estimators\nthat extend the previous Gaussian-based variational estimators with Normalizing\nFlows.\" This is ambiguous -- did the previous Gaussian-based methods use normalizing flows? No, your method uses normalizing flows. Please consider changing this to: \"In this paper, we introduce new classes of variational estimators based on \nNormalizing Flows that extend the previous Gaussian-based variational estimators.\"\n\n4. Eq (5) is basically the same idea as employed in Gaussian processes to model the distribution of x based on observed y. Maybe you should note this (minor). It could help the reader, particularly GP experts, understand this.", "limitations": "", "rating": 8, "confidence": 5, "soundness": 4, "presentation": 3, "contribution": 4, "strengths": "The paper develops a sophisticated and well-grounded formulation for flexibly estimating the mutual information based on observed samples, and impressive experiments are performed.", "weaknesses": "Generally the paper is well written, but there are places it could be improved. Additionally, while the bibliography is comprehensive, there is some recent work that should be referenced and compared to (theoretically at least, while additional experiments are unnecessary).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)