Year: 2021

Abstract:
Mutual information (MI) is a fundamental quantity in information theory and machine learning. However, direct estimation of mutual information is intractable, even if the true joint probability density for the variables of interest is known, as it involves estimating a potentially high dimensional log partition function. In this work, we view mutual information estimation from the perspective of importance sampling. Since naive importance sampling with the marginal density as a proposal requires exponential sample complexity in the true mutual information, we propose Multi-Sample Annealed Importance Sampling (AIS) bounds on mutual information to bridge between the marginal and joint distributions. In settings where the full joint distribution is available, we provide lower and upper bounds that can tightly estimate large values of MI in our experiments. In settings where only a single marginal distribution is known, we improve upon existing variational methods by directly optimizing a tighter lower bound on MI, using energy-based training to estimate gradients and multi-sample AIS for evaluation. Our methods are particularly suitable for evaluating MI in deep generative models, since explicit forms for the marginal or joint densities are often available. We evaluate our bounds on estimating the MI of VAEs and GANs trained on the MNIST and CIFAR datasets, and showcase significant gains over existing bounds in these challenging settings with high ground truth MI.

text: Mutual information (MI) is a fundamental quantity in information theory and machine learning. However, direct estimation of mutual information is intractable, even if the true joint probability density for the variables of interest is known, as it involves estimating a potentially high dimensional log partition function. In this work, we view mutual information estimation from the perspective of importance sampling. Since naive importance sampling with the marginal density as a proposal requires exponential sample complexity in the true mutual information, we propose Multi-Sample Annealed Importance Sampling (AIS) bounds on mutual information to bridge between the marginal and joint distributions. In settings where the full joint distribution is available, we provide lower and upper bounds that can tightly estimate large values of MI in our experiments. In settings where only a single marginal distribution is known, we improve upon existing variational methods by directly optimizing a tighter lower bound on MI, using energy-based training to estimate gradients and multi-sample AIS for evaluation. Our methods are particularly suitable for evaluating MI in deep generative models, since explicit forms for the marginal or joint densities are often available. We evaluate our bounds on estimating the MI of VAEs and GANs trained on the MNIST and CIFAR datasets, and showcase significant gains over existing bounds in these challenging settings with high ground truth MI.

heading: 1 INTRODUCTION
text: Mutual information (MI) is among the most general measures of dependence between two random variables. Among many other applications in machine learning, mutual information has been used for both training (Alemi et al., 2016; 2018; Chen et al., 2016; Zhao et al., 2018) and evaluating (Alemi & Fischer, 2018; Huang et al., 2020) generative models. Furthermore, a recent success in neural network function approximation has encouraged a wave of variational or contrastive methods for representation learning (van den Oord et al., 2018) or MI estimation from samples alone (Belghazi et al., 2018; Poole et al., 2019). However, McAllester & Stratos (2020) have shown strong theoretical limitations on any estimator which does not have access to an analytical form of at least one marginal distribution. In light of these limitations, we consider MI estimation in settings where a single marginal or the full joint distribution are available.
In this work, we propose novel Multi-Sample Annealed Importance Sampling (AIS) bounds which can tightly estimate large values of MI when the full joint distribution is known. We motivate these bounds by viewing MI estimation from the perspective of importance sampling, which also suggests novel variational approaches for the case where only a single marginal is known. First, our Generalized IWAE (GIWAE) bounds highlight how multi-sample contrastive estimation can complement variational learning to improve MI lower bounds, and includes INFONCE (van den Oord et al., 2018; Poole et al., 2019) and importance-weighted autoencoder (IWAE) (Burda et al., 2016; Sobolev & Vetrov, 2019) bounds as special cases. We also propose MINE-AIS, which uses an improved energy-based lower bound compared with Mutual Information Neural Estimation (MINE) (Belghazi et al., 2018) and corresponds to the infinite-sample limit of the GIWAE framework. We propose an efficient training procedure using Markov Chain Monte Carlo (MCMC) and utilize our Multi-Sample AIS evaluation bounds to show notable improvement over finite-sample methods for lower bounding MI. In Fig. 1, we summarize the MI bounds discussed in this paper and the relationships between them.

heading: 1.1 PROBLEM SETTING
text: The mutual information between two random variables x and z with joint distribution p(x, z) is
I(x; z) = Ep(x,z) [ log p(x, z)
p(x)p(z)
] = H(x)−H(x|z) = Ep(x,z)[log p(x|z)]− Ep(x)[log p(x)], (1)
whereH(x|z) denotes the conditional entropy −Ep(x,z) log p(x|z). We primarily focus on bounds that assume either a single marginal distribution or the full joint distribution are available. A natural setting where the full joint distribution is available is estimating the mutual information in deep generative models between the latent variables, with a known prior z ∼ p(z), and data x ∼ p(x) simulated from the model (Alemi & Fischer, 2018). 1 Settings where a single marginal is available appear in simulation-based inference (Cranmer et al., 2020), where information about input parameters θ is known and a simulator can generate x for a given θ, but the likelihood p(x|θ) is intractable. While sampling from the posterior p(z|x) for an arbitrary x is often intractable, we can obtain a single posterior sample for x ∼ p(x), z ∼ p(z|x) in cases where samples from the joint distribution of interest p(x)p(z|x) are available. Throughout this paper, we will refer to bounds which involve only a single posterior sample as practical, and those involving multiple posterior samples as impractical.
When the conditional p(x|z) is tractable to sample and evaluate, simple Monte Carlo sampling provides an unbiased, low variance estimate of the conditional entropy term in Eq. 1. In this case, the difficulty of MI estimation reduces to estimating the log partition function, for which importance sampling (IS) based methods are among the most well studied and successful solutions.

heading: 2 MULTI-SAMPLE BOUNDS ON MUTUAL INFORMATION
text: In this section, we review single- and multi-sample MI bounds from the perspective of importance sampling. This provides insight into existing bounds and will provide foundation for our primary contributions in Sec. 3 and 4.

heading: 2.1 SIMPLE IMPORTANCE SAMPLING BOUNDS
text: For an arbitrary conditional distribution qθ(z|x), we can obtain lower and upper bounds on log p(x):
Eqθ(z|x) [ log p(x, z)
qθ(z|x)
] ≤ log p(x) ≤ Ep(z|x) [ log p(x, z)
qθ(z|x)
] (2)
where the left-hand side is the familiar Evidence Lower Bound (ELBO) with a gap of DKL[qθ(z|x)||p(z|x)] = log p(x) − ELBO(qθ). We refer to the right-hand side as the Evidence Upper Bound (EUBO), with gap DKL[p(z|x)||qθ(z|x)] = EUBO(qθ) − log p(x). Throughout this paper, we will derive bounds which correspond to extended-state space versions of Eq. 2 (see App. A).
By plugging the importance sampling lower and upper bounds on log p(x) from Eq. 2 into Eq. 1, we obtain upper and lower bounds on mutual information as follows
IBAL(qθ) = Ep(x,z) [ log
qθ(z|x) p(z) ] ≤ I(x; z) ≤ Ep(x)qθ(z|x) [ log qθ(z|x) p(x, z) ] −H(x|z) = IBAU (qθ) (3)
The left hand side of Eq. 3 is the well-known Barber-Agakov (BA) bound (Barber & Agakov, 2003; Agakov, 2006), which we denote by IBAL(qθ) = I(x; z) −DKL[p(z|x)||qθ(z|x)]. We refer to the right hand side as the BA upper bound IBAU (qθ) = I(x; z) +DKL[qθ(z|x)||p(z|x)].
1An alternative, ‘encoding’ MI between the real data and the latent code is often of interest in representation learning (see App. F), but cannot be directly estimated using our methods due to the unavailability of densities for the data distribution pd(x) or the ‘aggregated posterior’ q(z) = ∫ pd(x)q(z|x)dx.
Relationship with CLUB When the prior is used as the proposal, IBAU (p(z)) corresponds to the Contrastive Log-Ratio Upper Bound (CLUB) (Cheng et al., 2020) with known p(x|z).

heading: 2.2 IMPORTANCE WEIGHTED AUTOENCODER
text: For a given variational distribution qθ(z|x), IWAE bounds (Burda et al., 2016; Sobolev & Vetrov, 2019) improve upon simple importance sampling by considering multi-sample proposal and target distributions in an extended state space (Domke & Sheldon, 2018; Cremer et al., 2017). Consider a proposal qIWAEPROP with K independent samples from the variational distribution qθ(z|x). The extended state space target pIWAETGT is a mixture distribution involving a single sample from the posterior p(z|x) or joint p(x, z) distribution and K − 1 samples from qθ(z|x)
qIWAEPROP (z (1:K)|x) = K∏ k=1 qθ(z (k)|x) pIWAETGT (x, z(1:K)) = 1 K K∑ s=1 p(x, z(s)) K∏ k=1,k 6=s qθ(z (k)|x) . (4)
The log ratio log p IWAE TGT (x,z (1:K))
qIWAEPROP (z (1:K)|x) reduces to familiar ratio appearing in the IWAE objective, while the
ratio of normalization constants corresponds to log p(x). As in App. A, taking expectations under the proposal yields a lower bound, whereas taking expectations under the target yields an upper bound
Eqθ(z(1:K)|x)
[ log 1
K K∑ i=1 p(x, z(k)) qθ(z(k)|x) ] ︸ ︷︷ ︸
ELBOIWAE(qθ,K)
≤ log p(x) ≤ Ep(z(1)|x)qθ(z(2:K)|x)
[ log 1
K K∑ i=1 p(x, z(k)) qθ(z(k)|x) ] ︸ ︷︷ ︸
EUBOIWAE(qθ,K)
. (5)
Note that the indices z(1:K) denote a set of independent draws from a given distribution. As for the ELBO and EUBO, the gap in the lower and upper bounds areDKL[qIWAEPROP ||pIWAETGT ] andDKL[pIWAETGT ||qIWAEPROP ], respectively. See App. B for detailed derivations.
As discussed in Sec. 1.1, lower and upper bounds on the log partition function translate to upper and lower bounds, respectively, on MI when p(x|z) is known. We refer to these bounds arising from IWAE estimation as IIWAEL(qθ,K) and IIWAEU (qθ,K). For example, IIWAEL(qθ,K) can be written as
I(x; z) ≥ IIWAEL(qθ,K) := Ep(x)p(z(1)|x)qθ(z(2:K)|x) [ log
p(x|z(1)) 1 K ∑K i=1 p(x,z(k)) qθ(z(k)|x)
] (6)
While it is well-known that increasing K leads to tighter IWAE bounds (Burda et al., 2016; Sobolev & Vetrov, 2019), we characterize the logarithmic rate of bias reduction for IIWAEL(qθ,K) in the following proposition, with proof in App. B.2.
Proposition 2.1. IWAE bounds on MI improve upon the BA bounds with the following relationships: IBAL(qθ) ≤ IIWAEL(qθ,K) ≤ IBAL(qθ) + logK, IIWAEU (qθ,K) ≤ IBAU (qθ). (7)

heading: 2.3 GENERALIZED IWAE
text: In this section, we consider a family of Generalized IWAE (GIWAE) lower bounds, which interpolate between IBAL(qθ) and IIWAEL(qθ,K) using a contrastive energy function Tφ(x, z). While similar bounds appear in (Lawson et al., 2019; Sobolev, 2019b), we provide a thorough discussion of special cases, and empirical analysis in the context of MI estimation in Sec. 5.2. We also show that our MINE-AIS method in Sec. 4 corresponds to the limiting behavior of GIWAE as K →∞. To derive the GIWAE lower bound, we begin by further extending the state space of the IWAE target distribution in Eq. 4, using a uniform index variable p(s) = 1K ∀s that specifies which sample z(s) ∼ p(z|x) is drawn from the posterior
pGIWAETGT (z (1:K), s|x) = 1
K p(z(s)|x) K∏ k=1,k 6=s q(z(k)|x) . (8)
Note that marginalization over s leads to the IWAE target distribution pIWAETGT (z (1:K)|x). For the GIWAE extended state space proposal, consider a categorical index variable qGIWAEPROP (s|z(1:K),x) drawn using self-normalized importance sampling (SNIS) with weights calculated by a learned energy function Tφ
qGIWAEPROP (z (1:K), s|x) = ( K∏ k=1 q(z(k)|x) ) qGIWAEPROP (s|z(1:K),x) where qGIWAEPROP (s|z(1:K),x) = eTφ(x,z (s))
K∑ k=1 eTφ(x,z (k)) .
We provide detailed discussion of the GIWAE probabilistic interpretation in App. C.1. Taking the
expectation under pGIWAETGT z (1:K), s|x) of the log importance ratio log p
GIWAE TGT (z (1:K),s|x) qGIWAEPROP (z
(1:K),s|x) as in App. A, we obtain an upper bound on log p(x) and lower bound on MI as
IGIWAEL(qθ, Tφ,K) = Ep(x,z) [ log
qθ(z|x) p(z) ] ︸ ︷︷ ︸
IBAL (qθ)
+Ep(x)p(z(1)|x)qθ(z(2:K)|x)
[ log
eTφ(x,z (1))
1 K ∑K i=1 e Tφ(x,z (k)) ] ︸ ︷︷ ︸
contrastive term≤logK
.
(9)
We see that the GIWAE lower bound decomposes into the sum of two terms, where the first is the BA variational lower bound for qθ(z|x) and the second is a contrastive term which distinguishes negative samples from qθ(z|x) and positive samples from p(z|x). Similarly to BA, GIWAE requires access to the analytical form of p(z) to evaluate the bound on MI. However, if the goal is to optimize mutual information in applications such as representation learning, both the BA and GIWAE lower bounds can be used even if no marginal distribution is available. See App. F for more detailed discussion.
Relationship with BA Choosing constant Tφ(x, z) = c means that the second term in GIWAE vanishes, so that we have IGIWAEL(q, Tφ = c,K) = IBAL(qθ) for all K. Similarly, the single sample IGIWAEL(q, Tφ,K = 1) also equals the BA lower bound for all Tφ.
Relationship with InfoNCE When the prior p(z) is used in place of qθ(z|x), we can recognize the second term in Eq. 9 as the INFONCE contrastive lower bound (van den Oord et al., 2018; Poole et al., 2019), with IINFONCEL(Tφ,K) = IGIWAEL(p(z), Tφ,K). From this perspective, the GIWAE lower bound highlights how variational learning can complement contrastive bounds to improve MI estimation beyond the known logK limitations of INFONCE (van den Oord et al., 2018). However, using the prior as the proposal in INFONCE does allow the energy function to admit a bi-linear implementation T (x, z) = Tx(x)Tz(z), which requires only N +K forward passes, rather than NK for GIWAE, where N is the batch size and K is the total number of positive and negative samples.
Relationship with IWAE The following proposition characterizes the relationship between the GIWAE lower bound in Eq. 9 and the IWAE lower bound on MI from Sec. 2.2. See App. C.4 for proofs.
Proposition 2.2. For a given qθ(z|x), the optimal energy function equals the true log importance weights up to an arbitrary constant: T ∗(x, z) = log p(x,z)qθ(z|x) + c(x). With this choice of T ∗(x, z),
IGIWAEL(qθ, T ∗) = IIWAEL(qθ) . (10)
Corollary 2.3. For a given qθ(z|x), let Tφ∗(x, z) denote the energy function that maximizes the GIWAE lower bound when optimization is performed with a neural network parameterized by φ. Then
IBAL(qθ) ≤ IGIWAEL(q, Tφ∗ ,K) ≤ IIWAEL(q,K) ≤ IBAL(qθ) + logK . (11)
While the GIWAE lower bound does not assume access to the full joint distribution p(x, z), Prop. 2.2 suggests that the role of the energy function is to learn the true log importance weights for qθ(z|x). Thus, when p(x|z) is known, IWAE is always preferable to GIWAE. Cor. 2.3 shows that while GIWAE can improve upon BA, this improvement is at most logarithmic with respect to K.
Relationship with Structured INFONCE Since INFONCE and Structured INFONCE (SINFONCE, see App. B.5, Poole et al. (2019)) are special cases of GIWAE and IWAE that use p(z) as the variational distribution, Cor. 2.3 suggests the following relationship
0 ≤ IINFONCEL(Tφ,K) ≤ IS-INFONCEL(K) ≤ logK . (12)

heading: 3 MULTI-SAMPLE AIS BOUNDS FOR ESTIMATING MUTUAL INFORMATION
text: In the previous sections, we derived probabilistic interpretations of extended state space importance sampling using multiple samplesK, as in IWAE. In this section, we first review AIS, which extends the state space using T intermediate distributions. We then show that these approaches are complementary, and derive two practical Multi-Sample AIS methods, which provide tighter bounds by combining insights from IWAE and AIS.

heading: 3.1 ANNEALED IMPORTANCE SAMPLING BACKGROUND
text: AIS (Neal, 2001) constructs a sequence of intermediate distributions {πt(z)}Tt=0, which bridge between a normalized initial distribution π0(z|x) and target distribution πT (z|x) = p(z|x) with unnormalized density πT (x, z) = p(x, z) and normalizing constant ZT (x) = p(x). A common choice for intermediate distributions is the geometric mixture path parameterized by {βt}Tt=0:
πt(z|x) = π0(z|x)1−βt πT (x, z)βt
Zt(x) where Zt(x) =
∫ π0(z|x)1−βt πT (x, z)βtdz. (13)
In the probabilistic interpretation of AIS, we consider an extended state space proposal qAISPROP(z0:T |x), obtained by sampling from the initial π0(z|x) and constructing transitions Tf (zt+1|zt,x) which leave πt(z|x) invariant. The target distribution pAISTGT(z0:T |x) is given by running the reverse transitions Tr(zt|zt+1,x) starting from a target or posterior sample πT (z|x), as shown in Fig. 2.
qAISPROP(z0:T |x) = π0(z0|x) T−1∏ t=0 Tf (zt+1|zt) pAISTGT(x, z0:T ) = πT (x, zT ) T−1∏ t=0 Tr(zt|zt+1,x) (14)
As described in App. H, taking expectations of the log importance weights under the proposal and target yields a lower and upper bound on the log partition function log p(x). These single-chain lower and upper bounds translate to upper and lower bounds on MI, IAISL(π0, T ) and IAISU (π0, T ), which were suggested for MI estimation in the blog post of Sobolev (2019a).

heading: 3.2 INDEPENDENT MULTI-SAMPLE AIS BOUNDS
text: To derive Independent Multi-Sample AIS (IM-AIS), we construct an extended state space proposal by running K independent AIS forward chains z(k)0:T ∼ qAISPROP in parallel. Similarly to the IWAE upper bound (Eq. 5), the extended state space target involves selecting a index s uniformly at random, and running a backward AIS chain z(s)0:T ∼ pAISTGT starting from a true posterior sample zT ∼ p(z|x). The remaining K − 1 samples are obtained by running forward AIS chains, as visualized in Fig. 2
qIM-AISPROP (z (1:K) 0:T |x) = K∏ k=1 qAISPROP(z (k) 0:T |x), p IM-AIS TGT (z (1:K) 0:T ,x) = 1 K K∑ s=1 pAISTGT(x, z (s) 0:T ) K∏ k=1 k 6=s qAISPROP(z (k) 0:T |x). (15)
where qAISPROP and p AIS TGT were defined in Eq. 14. Note that sampling from the extended state space target distribution is practical, as it only requires one sample from the true posterior distribution.
Taking the expectation of the log unnormalized density ratio under the proposal and target yields lower and upper bounds on log p(x), respectively (see App. A)
E z (1:K) 0:T ∼qAISPROP
[ log 1
K K∑ k=1 pAISTGT(x, z (k) 0:T ) qAISPROP(z (k) 0:T |x) ] ︸ ︷︷ ︸
ELBOIM-AIS(π0,K, T )
≤ log p(x) ≤ E z (1) 0:T ∼pAISTGT
z (2:K) 0:T ∼qAISPROP
[ log 1
K K∑ k=1 pAISTGT(x, z (k) 0:T ) qAISPROP(z (k) 0:T |x) ] ︸ ︷︷ ︸
EUBOIM-AIS(π0,K, T )
. (16)
which again have KL divergences as the gap in their bounds. Independent Multi-Sample AIS reduces to IWAE for T = 1, and reduces to single-sample AIS for K = 1. Both the upper and lower bounds become exact as K →∞ or T →∞, and translate to lower and upper bounds on MI as in Sec. 1.1.
Proposition 3.1 (Complexity in K). For a given T , the bias of the Independent Multi-Sample AIS lower bound on MI, IIM-AISL(π0, T ), decreases (at most) logarithmically in K.
IAISL(π0, T ) ≤ IIM-AISL(π0,K, T ) ≤ IAISL(π0, T ) + logK (17)
To characterize the bias reduction for AIS with increasing T , we prove the following proposition.
Proposition 3.2 (Complexity in T ). Assuming perfect transitions and a geometric annealing path with linearly-spaced {βt}Tt=1, the sum of the gaps in the multi-sample AIS sandwich bounds on MI, δT,K = IIM-AISU (π0,K, T )− IIM-AISL(π0,K, T ), reduces linearly with increasing T . This generalizes Thm. 1 of Grosse et al. (2013), which holds as T →∞. See App. H.3 for proof.

heading: 3.3 COUPLED REVERSE MULTI-SAMPLE AIS BOUNDS
text: We can change the role of the forward and backward annealing chains in Independent Multi-Sample AIS to obtain alternative bounds on the log partition function. We define Independent Reverse Multi-Sample AIS (IR-AIS) using the following proposal and target distribution, as shown in Fig. 2.
qIR-AISPROP (x, z (1:K) 0:T ) =
1
K K∑ s=1 qAISPROP(z (s) 0:T |x) K∏ k=1,k 6=s pAISTGT(z (k) 0:T ,x) p IR-AIS TGT (z (1:K) 0:T ,x) = K∏ k=1 pAISTGT(z (k) 0:T ,x)
Note that partition function ratio is ZTGT/ZPROP = p(x)K/p(x)K−1 = p(x). Using these distributions, we derive log p(x) and MI bounds in App. J. However, these will be impractical in most settings since they require multiple true posterior samples (see Sec. 1.1).
To address this, we propose Coupled Reverse Multi-Sample AIS (CR-AIS). The extended state space target distribution in Fig. 2 initializes K backward chains from a single target sample zT ∼ πT (z|x), with remaining transitions denoted pAISTGT(z0:T−1|zT ) since they are identical to standard AIS in Eq. 14.
pCR-AISTGT (z (1:K) 0:T−1, zT ,x) = πT (zT ,x) K∏ k=1 pAISTGT(z (k) 0:T−1|zT ,x). (18)
The extended state space proposal is obtained by selecting an index s uniformly at random and running a single forward AIS chain. We then run K − 1 backward chains, all starting from the last state of the selected forward chain, as visualized in Fig. 2
qCR-AISPROP (z (1:K) 0:T−1, zT |x) =
1
K K∑ s=1 qAISPROP(z (s) 0:T−1, zT |x) K∏ k=1,k 6=s pAISTGT(z (k) 0:T−1|zT ,x) (19)
Taking the expected log ratio under the proposal and target yields lower and upper bounds on log p(x),
−E z (1) 0:T−1,zT∼q AIS PROP(z0:T |x)
z (2:K) 0:T−1∼p AIS TGT(z0:T−1|zT ,x)
[ log 1
K K∑ k=1 qAISPROP(·) pAISTGT(·) ] ︸ ︷︷ ︸
ELBOCR-AIS(π0,K, T )
≤ log p(x) ≤ −E zT∼πT (zT |x)
z (1:K) 0:T−1∼p AIS TGT(z0:T−1|zT ,x)
[ log 1
K K∑ k=1 qAISPROP(·) pAISTGT(·) ] ︸ ︷︷ ︸
EUBOCR-AIS(π0,K, T )
Proposition 3.3. For a given T , the bias of the Coupled Reverse Multi-Sample AIS upper bound on MI, ICR-AISU (π0,K, T ), decreases (at most) logarithmically in K.
IAISU (π0, T )− logK ≤ ICR-AISU (π0,K, T ) ≤ IAISU (π0, T ). (20)

heading: 3.4 DISCUSSION
text: Relationship with BDMC While Bidirectional Monte Carlo (BDMC) (Grosse et al., 2015; 2016) was the first method proposing multi-sample log partition function bounds using AIS chains, our probabilistic interpretations provide novel perspective on BDMC. Perhaps surprisingly, we find that BDMC lower and upper bounds do not correspond to the same extended state space proposal and target distributions. In particular, the BDMC lower bound on log p(x) corresponds to the lower bound of Independent Multi-Sample AIS (Fig. 2 Col. 4, Row 4), while the upper bound of BDMC matches the upper bound of Coupled Reverse Multi-Sample AIS (Fig. 2 Col. 6, Row 3).
Effect of K and T We have shown in Prop. 3.2 that our multi-sample AIS bounds can achieve linear bias reduction with increasing T , although this computation must be done in serial fashion. While increasing K involves parallel computation, its bias reduction is often only logarithmic for both IWAE and Multi-Sample AIS, as we have shown in Prop. 2.1, 3.1, 3.3 and further argue in App. B.2-B.4. Based on these arguments, we recommend increasing K until computation can no longer be parallelized on given hardware and allocating all remaining resources to increasing T . Comparing Multi-Sample AIS Bounds In App. M.1 Fig. 9 , we compare the performance of our various Multi-Sample AIS bounds in order to make recommendations for which bounds to use in practice. For the upper bound on MI, we recommend the Independent Multi-Sample AIS ELBO, which uses independent samples and is not limited to logK improvement compared with the Coupled Reverse Multi-Sample AIS ELBO (Prop. 3.3). The results are less conclusive for the lower bounds on MI. While the MI lower bound obtained from EUBOIM-AIS(π0,K, T ) (RHS of Eq. 16) can only improve upon single-sample AIS by logK (Prop. 3.1), this improvement is easily obtained for low T and may be useful for quickly estimating MI at a similar order of magnitude as logK. The Coupled Reverse AIS lower bound on MI in requires at least moderate values of T to match or marginally improve upon Independent Multi-Sample AIS, suggesting that the preferred MI lower bounds may vary based on the scale of the true MI and the amount of computation available.

heading: 4 MINE-AIS ESTIMATION OF MUTUAL INFORMATION
text: In this section, we present MINE-AIS, which uses our multi-sample AIS bounds from the previous section to bound the mutual information in settings where the true conditional density p(x|z) is unknown. We consider a flexible energy-based distribution πθ,φ(z|x) as an approximation to the posterior p(z|x) (Poole et al., 2019; Arbel et al., 2020)
πθ,φ(z|x) = 1
Zθ,φ(x) qθ(z|x)eTφ(x,z) where Zθ,φ(x) = Eqθ(z|x)
[ eTφ(x,z) ] . (21)
Plugging πθ,φ(z|x) into the BA lower bound, we denote the resulting bound as the Implicit BarberAgakov Lower Bound (IBAL), since it is often difficult to evaluate explicitly due to the intractable log partition function term. After simplifying in App. D.1, we obtain
I(x; z) ≥ IBAL(qθ, Tφ) := IBAL(πθ,φ) = Ep(x,z) [ log
qθ(z|x) p(z) ] ︸ ︷︷ ︸
IBAL(qθ)
+Ep(x,z) [ log
eTφ(x,z) Eqθ(z|x) [ eTφ(x,z) ]]︸ ︷︷ ︸ ≤Ep(x)[DKL[p(z|x)||qθ(z|x)]] , (22)
with the gap of the IBAL equal to Ep(x)[DKL[p(z|x))||πθ,φ(z|x)]]. Note that the IBAL generalizes the Unnormalized Barber-Agakov bound from Poole et al. (2019), with IUBA(Tφ) = IBAL(p(z), Tφ).
Relationship with MINE As discussed in Poole et al. (2019), MINE-DV and MINE-F (Belghazi et al., 2018) utilize lower bounds on IBAL(p(z), Tφ). Instead, MINE-AIS will optimize Eq. 22 directly using an improved energy-based training scheme and evaluate the IBAL using AIS. Relationship with GIWAE We can immediately notice similarities between GIWAE and MINE-AIS, including that their optimal energy functions match and that both bounds include a contrastive term which improves upon the BA lower bound. In fact, we prove the following proposition in App. E.2.
Proposition 4.1. For given qθ(z|x) and Tφ(x, z), lim K→∞ IGIWAEL(qθ, Tφ,K) = IBAL(qθ, Tφ).
Thus, we may view the IBAL as the limiting behavior of the finite-sample GIWAE bounds as K →∞. While Cor. 2.3 shows that IGIWAEL(qθ, T ) can improve IBAL(qθ) by at most logK nats, we show in App. E.1 that the IBAL contrastive term is flexible enough to close the entire gap in the BA bound.
IBAL(qθ) ≤ IBAL(qθ, Tφ∗) ≤ I(x; z) = IBAL(qθ) + Ep(x) [DKL[p(z|x)||qθ(z|x)]] (23) However, this flexible contrastive term comes at the cost of tractability, as MINE-AIS involves an intractable partition function logEqθ(z|x)[eTφ(x,z)] compared to the finite-sum term in GIWAE. Training of MINE-AIS Although the log partition function logZθ,φ(x) in the IBAL is intractable to evaluate, we only require an unbiased estimator of its gradient for training. Differentiating Eq. 22 with respect to the parameters θ and φ of the variational and energy function, respectively, we obtain
∂
∂θ IBAL(qθ, Tφ) = Ep(x,z)
[ ∂
∂θ log qθ(z|x)
] − Ep(x)πθ,φ(z|x) [ ∂
∂θ log qθ(z|x)
] , (24)
∂
∂φ IBAL(qθ, Tφ) = Ep(x,z)
[ ∂
∂φ Tφ(x, z)
] − Ep(x)πθ,φ(z|x) [ ∂
∂φ Tφ(x, z)
] . (25)
We use MCMC with Hamiltonian Monte Carlo (HMC) transition kernels (Neal, 2011) to draw approximate samples from πθ,φ, with additional discussion in App. E.4. Remarkably, using Prop. 4.1 and Eq. 24-25, we observe that direct gradient-based optimization of the contrastive term in finite-sample GIWAE bounds translates to energy-based training in the limit ofK →∞. The goal in both cases is to raise the value of Tφ on positive examples and lower it on negative examples, although, in MINE-AIS, we use only a single, approximate sample from πθ,φ(z|x).
Evaluation of IBAL in MINE-AIS Given an energy function learned using our MINE-AIS training scheme, we still need to evaluate the IBAL lower bound on MI. We have an unbiased, low variance estimate for Ep(x,z) [Tφ(x, z)], and will use Multi-Sample AIS to estimate logZθ,φ(x) by annealing between π0(z|x) = qθ(z|x) and πT (x, z) = qθ(z|x)eTφ(x,z). To preserve a lower bound on I(x; z), we would like an upper bound on logZθ,φ(x). However, this requires true samples from πθ,φ(z|x) to initialize backward AIS chains. Since these samples are unavailable, we argue in App. E.5 that we can preserve an approximate upper bound on logZθ,φ(x) by initializing backward chains from the true posterior p(z|x) instead of πθ,φ(z|x). Using a Multi-Sample AIS lower bound on logZθ,φ(x), we can upper bound the IBAL but cannot guarantee that this is a stochastic lower bound on I(x; z).

heading: 5 EXPERIMENTS
text: In this section, we evaluate our proposed Multi-Sample AIS, MINE-AIS, and GIWAE bounds for estimating MI in deep generative models such as VAEs and GANs trained on MNIST and CIFAR.

heading: 5.1 MULTI-SAMPLE AIS ESTIMATION OF MUTUAL INFORMATION
text: We compare Multi-Sample AIS MI estimation against IWAE, as both methods assume the full joint distribution is available. For the initial distribution of AIS or variational distribution of IWAE, we can use any distribution that is tractable to sample and evaluate. We experiment using both the prior p(z) and a Gaussian q(z|x) trained to optimize the BA lower bound. Tables 1 summarize our results. IWAE As described in Sec. 2.2, IWAE bounds encompass a wide range of MI estimators. The K = 1 bounds with learned qθ(z|x) correspond to BA bounds, while the upper bound with a prior p(z) proposal matches the CLUB. For K > 1 and p(z) as the proposal, we obtain Structured INFONCE.
We see that the improvement of the IWAE lower bound is limited by logK, and thus requires an exponentially large sample size to close the gap from the BA lower bound (K = 1) to the true MI. IWAE can successfully sandwich the MI of VAE2, VAE10 and GAN2 on MNIST in cases where the gap is less than logK. However, using Gaussian qθ(z|x), tight estimation with IWAE is clearly infeasible on VAE100 and GAN100. On VAE100, for example, at least e46 total samples are required to improve IBAL(qθ) = IIWAEL(qθ, 1) = 34.56 to match the ground truth I(x; z) ≈ 80.95 estimated by AIS. Multi-Sample AIS We evaluate our Multi-Sample AIS bounds with K = 48 chains on MNIST and K = 12 on CIFAR, and a varying number of intermediate T . We show results for the Independent Multi-Sample AIS MI lower bound and Coupled Reverse Multi-Sample AIS upper bound in Table 1. Using large enough values of T , Multi-Sample AIS can tightly sandwich large values of ground truth MI in our MNIST-VAE models. Our bounds are more loose in the challenging CIFAR-GAN setting, but these results highlight that increasing T in Multi-Sample AIS is a practical way to reduce bias using additional computation. This is in stark contrast to the exponential sample complexity required for IWAE, as we saw in the previous discussion. We provide runtime details in App. L.3.

heading: 5.2 ENERGY-BASED ESTIMATION OF MUTUAL INFORMATION
text: In this section, we evaluate the family of GIWAE and MINE-AIS bounds, which assume access to a known marginal p(z) but not the conditional p(x|z). We summarize these bounds in Fig. 3b. BA, IWAE and GIWAE Bounds Recall that IGIWAEL(q, Tφ,K) in Eq. 9 consists of the BA lower bound and a contrastive term. In Fig. 3a, we report the contribution of each term for different model complexity and special cases of GIWAE. We find that the results respect the inequalities shown in Cor. C.2, with IBAL(qθ) ≤ IGIWAEL(qθ, Tφ∗ ,K) ≤ IIWAEL(qθ,K) ≤ IBAL(qθ) + logK . 2
For the Linear VAE, the true posterior is in the Gaussian variational family and IBAL(qθ) almost exactly matches the analytical MI. In this case, the contrastive term provides much less than logK improvement for GIWAE and IWAE, since even the optimal energy function cannot distinguish qθ(z|x) and p(z|x). The contrastive term increases for more complex MNIST-VAE posteriors, while the fact that BA term in IWAE is lower for K = 1000 than for K = 100 indicates a tradeoff between the two terms. A worse qθ allows the contrastive term to achieve closer to the full logK improvement, and yields a higher overall bound. This observation is similar to Rainforth et al. (2018), where increasing K in ELBOIWAE(qθ,K) reduces the quality of the posterior approximation in variational inference. Finally, there is essentially no tradeoff between the terms for MNIST-GAN, since the variational family is far enough from the true posterior that both learned GIWAE and optimal IWAE energy functions can achieve logK improvement over BA.
MINE-AIS Bounds We can see in Fig. 3a that MINE-AIS improves over BA due to its more flexible variational family, and improves upon GIWAE, IWAE, and (Structured) INFONCE since it is their limiting behavior. 3 In Fig. 3a, we can compare AIS estimates of MI and the IBAL for MINE-AIS to indicate the quality of the learned energy function. We observe that MINE-AIS lower bounds underestimate the value obtained from Multi-Sample AIS by 1%, 18%, and 28% on linear VAE, MNIST-VAE, and MNIST-GAN, respectively.

heading: 6 CONCLUSION
text: In this work, we have provided a unifying view of mutual information estimation from the perspective of importance sampling. We derived probabilistic interpretations of each bound, which shed light on the limitations of existing estimators, and motivate our novel GIWAE, Multi-Sample AIS, and MINE-AIS bounds. When the conditional is not known, our GIWAE bounds highlights how variational bounds can complement contrastive learning to improve lower bounds on MI beyond the known logK limitations. When the full joint distribution is known, our Multi-Sample AIS bounds are the only practical bounds which have been demonstrated to estimate large values of MI without exponential sample complexity, and thus should be considered the gold standard for MI estimation in these settings. Finally, MINE-AIS extends Multi-Sample AIS evaluation to unknown conditional densities, and can be viewed as the infinite-sample behavior of GIWAE and existing contrastive bounds. Our MINE-AIS and Multi-Sample AIS methods highlight how MCMC techniques can be used to improve mutual information estimation when a single analytic marginal or conditional density is available.
2Note that the BA term in GIWAE and IWAE may not match exactly, since the overall training objectives are different. For example, GIWAE jointly learns qθ and Tφ rather than only optimizing qθ to maximize IBAL(qθ).
3Note that our MINE-AIS results use p(z) as the base distribution, which is the limiting behavior of INFONCE. We might expect further improvement by jointly learning a variational base distribution qθ(z|x) and energy Tφ.

heading: Supplementary Material

heading: Table of Contents
text: A A General Approach for Deriving Extended State Space Bounds on Log Partition
Functions 14
B IWAE 15 B.1 Probabilistic Interpretation and Bounds . . . . . . . . . . . . . . . . . . . . . . 15 B.2 Proof of Proposition 2.1 . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 16 B.3 Logarithmic Improvement Results for IIWAEL(qθ,K) . . . . . . . . . . . . . . . 17 B.4 Bias Reduction of IWAE Lower Bound on log p(x) / Upper Bound on MI . . . . 18 B.5 Relationship with Structured InfoNCE . . . . . . . . . . . . . . . . . . . . . . 18

heading: C Generalized IWAE 18
text: C.1 Probabilistic Interpretation and Bounds . . . . . . . . . . . . . . . . . . . . . . 18 C.2 GIWAE Upper Bound on MI Does Not Provide Benefit Over IWAE . . . . . . . 20 C.3 Relationship between IWAE and GIWAE Probabilistic Interpretations . . . . . . 20 C.4 Proof of Proposition 2.2 and Corollary 2.3 . . . . . . . . . . . . . . . . . . . . 21

heading: D Generalized Mutual Information Neural Estimation (GMINE) 22
text: D.1 Probabilistic Interpretation of IBAL . . . . . . . . . . . . . . . . . . . . . . . . 22 D.2 Probabilistic Interpretation of Generalized MINE-DV . . . . . . . . . . . . . . . 22 D.3 Probabilistic Interpretation of Generalized MINE-F . . . . . . . . . . . . . . . . 24

heading: E MINE-AIS 25
text: E.1 Properties of MINE-AIS . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 25 E.2 Proof of Proposition 4.1 . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 26 E.3 Convergence of GIWAE SNIS Distribution to MINE-AIS . . . . . . . . . . . . 28 E.4 Details for Training MINE-AIS . . . . . . . . . . . . . . . . . . . . . . . . . . 29 E.5 Approximate Reverse Annealing for Evaluating MINE-AIS . . . . . . . . . . . 29 E.6 Multi-Sample AIS Convergence Results for MINE-AIS . . . . . . . . . . . . . 30

heading: F Applications to Mutual Information Estimation without Known Marginals 31
text: F.1 BA Lower Bound . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 31 F.2 GIWAE Lower Bound . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 32 F.3 MINE-AIS / IBAL Lower Bound . . . . . . . . . . . . . . . . . . . . . . . . . 33

heading: G Background on AIS 33
text: H Independent Multi-Sample AIS 34 H.1 Probabilistic Interpretation and Bounds . . . . . . . . . . . . . . . . . . . . . . 34 H.2 Proof of Proposition 3.1 . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 35 H.3 Proof of Proposition 3.2 . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 35

heading: I Reverse IWAE 37
text: I.1 Probabilistic Interpretation and Bounds . . . . . . . . . . . . . . . . . . . . . . 37 I.2 Proof of Logarithmic Improvement of RIWAE . . . . . . . . . . . . . . . . . . 37
J Independent Reverse Multi-Sample AIS 38

heading: K Coupled Reverse Multi-Sample AIS 39
text: K.1 Probabilistic Interpretation and Bounds . . . . . . . . . . . . . . . . . . . . . . 39 K.2 Proof of Proposition 3.3 . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 39

heading: L Experimental Details 40
text: L.1 Datasets and Models . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 40 L.2 Multi-Sample AIS Settings . . . . . . . . . . . . . . . . . . . . . . . . . . . . 40
L.3 Runtime Comparison . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 40 L.4 Initial Distribution of AIS . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 40 L.5 Analytical Solution of the Mutual Information on the Linear VAE . . . . . . . . 41

heading: M Additional Experiments 42
text: M.1 Comparison of Multi-Sample AIS Bounds . . . . . . . . . . . . . . . . . . . . 42 M.2 Confidence Intervals for Multi-Sample AIS Experiments . . . . . . . . . . . . . 42

heading: A A GENERAL APPROACH FOR DERIVING EXTENDED STATE SPACE BOUNDS
text: ON LOG PARTITION FUNCTIONS
We begin by describing a general methodology for deriving lower and upper bounds on the log partition function log p(x) using importance sampling in an extended state space. We employ this approach throughout this Appendix, where the state space may be extended using multiple samples (K > 1) as in IWAE or approximate samples from intermediate densities (T > 1) as in AIS.
Suppose we are given a tractable initial distribution π0(z), which may correspond to a variational distribution qθ(z|x) or prior distribution p(z), and are interested in estimating the normalization constant p(x) for an intractable posterior distribution p(z|x) ∝ p(x, z). In Fig. 4, we summarize various extended state space proposal (second column of Fig. 4) and target distributions (third column of Fig. 4) constructed from these initial and posterior distributions, with qPROP and pTGT chosen such that ZTGTZPROP = log p(x). Each pair of extended state-space proposal and target distributions then provides both upper and lower bounds on the log partition function using importance sampling, by taking the expected log ratio of unnormalized densities under either the proposal or target distribution
EqPROP [ log
p̃TGT q̃PROP
] ≤ log ZTGT
ZPROP ≤ EpTGT
[ log
p̃TGT q̃PROP
] (26)
Note that we use q̃, p̃ to indicate unnormalized densities over x and a possibly extended (latent) state space. For example, pTGT(z (1:K) 0:T |x) = 1 ZTGT p̃TGT(x, z (1:K) 0:T ) represents a state space extended by multiple proposal samples K ≥ 1 as in IWAE and multiple intermediate samples T ≥ 1 as in AIS. For all qPROP and pTGT, the gap in the lower bound in Eq. 26 is the forward KL divergence, DKL(qPROP||pTGT), while the gap in the upper bound is the reverse KL divergence,DKL(pTGT)||qPROP).
EqPROP [ p̃TGT q̃PROP ] = EqPROP [ pTGT qPROP ] + log ZTGT ZPROP = log ZTGT ZPROP
−DKL[qPROP||pTGT]︸ ︷︷ ︸ ELBO(qPROP, pTGT)
≤ log ZTGT ZPROP (27)
EpTGT [ p̃TGT q̃PROP ] = EpTGT [ pTGT qPROP ] + log ZTGT ZPROP = log ZTGT ZPROP
+DKL[pTGT||qPROP]︸ ︷︷ ︸ EUBO(qPROP, pTGT)
≥ log ZTGT ZPROP (28)
Thus, the bounds in Eq. 26 directly generalize the standard ELBO (log p(x)−DKL[q(z|x)||p(z|x)]) and EUBO (log p(x) + DKL[p(z|x)||q(z|x)]), which appear as special cases with K = 1, T = 1, q(z|x) as the proposal distribution, and p(z|x) as the target distribution. We can alternatively derive the bounds in Eq. 26 by expanding the expression for the DKL[qPROP||pTGT] in the extended state space to obtain logZTGT/ZPROP − ELBO(qPROP, pTGT) as in Eq. 27. In Fig. 4, we highlight that the initial or variational distribution (blue circles) and posterior distribution (red circles) distributions may be used multiple times, in either or both of the extended state space proposal and target distributions. Similarly, forward AIS chains (blue circles) starting from the initial distribution and the backward AIS chains (shown in red circles) starting from the target distribution may be used repeatedly in both the proposal or target.
In the next sections, we proceed to derive each of the bounds in Fig. 4 along with probabilistic interpretations of the underlying generative processes.
B IWAE

heading: B.1 PROBABILISTIC INTERPRETATION AND BOUNDS
text: Consider a K-sample proposal distribution qIWAEPROP (z (1:K)) using independent draws from an initial distribution qθ(z|x). The target distribution is defined as a uniform mixture of K components, where each component replaces the initial sample qθ(z(k)|x) in index k with a sample from the target
distribution p(z|x).
qIWAEPROP (z (1:K)|x) = K∏ k=1 qθ(z (k)|x), pIWAETGT (x, z(1:K)) = 1 K K∑ k=1 p(x, z(s)) K∏ k=1,k 6=s qθ(z (s)|x) .
(29)
Note that the normalizing constant of pIWAETGT (x, z (1:K)), or ∫ pIWAETGT (x, z
(1:K))dz(1:K), equals p(x) since qθ(z|x) is normalized and ∫ p(x, z)dz = p(x). To connect this with the IWAE bound, we show that the log importance ratio reduces to
log pIWAETGT (x, z (1:K))
qIWAEPROP (x, z(1:K)) = log
1
K K∑ k=1 p(x, z(s)) K∏ j=1
j 6=k
qθ(z (k)|x)
K∏ k=1 qθ(z(k)|x)
= log
( 1
K K∑ k=1 p(x, z(k)) K∏ j=1
j 6=k
qθ(z (k)|x) · k∏ k=1 qθ(z (k)|x)−1
)
= log 1
K K∑ k=1 p(x, z(k)) qθ(z(k)|x) (30)
As in Eq. (26), we can obtain lower and upper bounds on log p(x) by taking expectations under the proposal and target distributions, respectively.
Lower Bound on log p(x) and Upper Bound on MI First, taking expectations under qIWAEPROP , we obtain a lower bound on log p(x) as show in App. A
ELBOIWAE(qθ,K) := Eqθ(z(1:K)) [ log 1
K K∑ k=1 p(x, z(k)) qθ(z(k)|x) ] (31)
This corresponds to the following upper bound on MI for known p(x|z)
I(x; z) ≤ IIWAEU (qθ,K) = −H(x|z)− ELBOIWAE(qθ,K) (32) = Ep(x,z) [log p(x|z)]− Eqθ(z(1:K)|x) [ log 1
K K∑ k=1 p(x, z(k)) qθ(z(k)|x) ] (33)
Upper Bound on log p(x) and Lower Bound on MI Similarly, with expectations under pIWAETGT (z
(1:K)|x), we obtain an upper bound on log p(x). Since, for independent draws, the uniform mixture pIWAETGT is invariant to index permutations, we may choose the target sample to be z (1) and obtain the upper bound of Sobolev & Vetrov (2019):
EUBOIWAE(qθ,K) := Ep(z(1)|x)qθ(z(2:K)|x)
[ log 1
K K∑ i=1 p(x, z(k)) q(z(k)|x)
] ≤ log p(x) (34)
Translating this to a lower bound on MI with known p(x|z),
I(x; z) ≥ IIWAEL(qθ,K) = −H(x|z)− EUBOIWAE(qθ,K) (35) = Ep(x,z) [log p(x|z)]− Ep(z(1)|x)qθ(z(2:K)|x) [ log 1
K K∑ k=1 p(x, z(k)) qθ(z(k)|x) ] (36)

heading: B.2 PROOF OF PROPOSITION 2.1
text: We first recall results that IWAE bounds tighten with increasing K.
Proposition B.1. For given qθ(z|x), ELBOIWAE(qθ,K + 1) ≥ ELBOIWAE(qθ,K) (Burda et al., 2016) and EUBOIWAE(qθ,K + 1) ≤ EUBOIWAE(qθ,K) (Sobolev & Vetrov, 2019). Proposition B.2. The bias of the IWAE upper bound on log p(x) decreases logarithmically in K. In particular, the K-sample upper bound EUBOIWAE(qθ,K) can improve upon the single-sample EUBO := log p(x) +DKL[p(z|x)||qθ(z|x)] by at most logK nats.
EUBO(qθ)− logK ≤ EUBOIWAE(qθ,K) ≤ EUBO(qθ) (37)
Proof. We include a simple proof of Prop. B.2 here, which provides intuition for the role of the ‘negative’ samples qθ(z(2:K)|x) and forms the basis for similar results later in the paper.
log p(x) ≤ E p(z(1)|x)qθ(z (2:K)|x)
[ log 1
K K∑ k=1 p(x, z(k)) qθ(z(k)|x)
] = E
p(z(1)|x)qθ(z (2:K)|x)
[ log
K∑ k=1 p(x, z(k)) qθ(z(k)|x)
] − logK
≥ E p(z(1)|x)
[ log p(x, z(1))
qθ(z(1)|x)
] − logK (38)
where, in the second line, we have ignored the contribution of the ‘negative’ samples in the summation since the likelihood ratio is non-negative. Using the ‘positive’ sample from the posterior p(z|x), the first term is the EUBO and is still an upper bound on log p(x). However, note that EUBO − logK is no longer guaranteed to be an upper bound on log p(x). Prop. B.1 implies that EUBOIWAE(qθ,K) ≤ EUBO(qθ).
Prop. 2.1 follows directly by translating these log p(x) results to bounds on MI.
Proposition 2.1. IWAE bounds on MI improve upon the BA bounds with the following relationships:
IBAL(qθ) ≤ IIWAEL(qθ,K) ≤ IBAL(qθ) + logK, IIWAEU (qθ,K) ≤ IBAU (qθ). (7)
Proof. Recall from Sec. 1.1 and Sec. 2.1 that IBAL(qθ)) = −EUBO(qθ) − H(x|z), and IIWAEL(qθ,K) = −EUBOIWAE(qθ,K) − H(x|z). Using Prop. B.2, EUBO(qθ,K) − logK ≤ EUBOIWAE(qθ,K) implies that
IIWAEL(qθ,K)− IBAL(qθ) = EUBO(qθ)− EUBOIWAE(qθ,K) ≤ logK (39)
so that IIWAEL(qθ,K) ≤ IBAL(qθ)) + logK as desired.
IBAL(qθ) ≤ IIWAEL(qθ,K) and IIWAEU (qθ,K) ≤ IBAU (qθ) follow from the fact that IWAE bounds tighten with increasing K in Prop. B.1.
B.3 LOGARITHMIC IMPROVEMENT RESULTS FOR IIWAEL(qθ,K)
In Sec. 2.3, we showed that IWAE is a special case of GIWAE, which decomposes into the sum of a variational BA lower bound and a K-sample contrastive term. This suggests that the IWAE lower bound on MI, which arises from an upper bound on log p(x), may be written as
IIWAEL(qθ,K) = Ep(x,z) [ log
qθ(z|x) p(z) ] ︸ ︷︷ ︸
IBA(q)
+Ep(x)p(z(1)|x)qθ(z(2:K)|x) log p(x,z(1))qθ(z(1)|x) 1 K ∑K i=1 p(x,z(k))
qθ(z (k)|x)  ︸ ︷︷ ︸
0≤ contrastive term≤logK
(40)
This way of writing IIWAEL(qθ,K) provides intuition for the result in Prop. 2.1, since the improvement over IBAL(qθ) is simply the contrastive term, which is limited to logK.
Table 2 and Table 3 show the IWAE objective decomposition to the BA term and the contrastive term, on VAEs and GANs trained MNIST and CIFAR10 dataset. We can see that in all the experiments, the contribution of the contrastive term is always less than logK. For low dimensional VAEs and GANs, the contrastive term contributes less than logK, since the posterior is not complicated and the BA term can get close to the ground truth MI. However, for higher dimensional VAE and GAN models, the contribution of the contrastive term is almost exactly logK. This sheds light on the inherent exponential sample complexity of IWAE, and motives our Multi-Sample AIS methods for MI estimation.
B.4 BIAS REDUCTION OF IWAE LOWER BOUND ON log p(x) / UPPER BOUND ON MI
We were unable to show a similar result for the IWAE lower bound on log p(x) and upper bound on I(x; z), although Chatterjee et al. (2018) suggests that exponential sample complexity K ∝ exp{DKL[p(z|x)||qθ(z|x)]} is required for accurate importance sampling estimation. Maddison et al. (2017); Domke & Sheldon (2018) find that the bias of the IWAE lower bound in the limit of K → ∞ reduces at the rate of O( 12KVar[ p(z|x) qθ(z|x) ]), although the Var[ p(x,z) qθ(z|x) ] term is at least exponential in DKL[p(z|x)||qθ(z|x)] (Song & Ermon, 2019).

heading: B.5 RELATIONSHIP WITH STRUCTURED INFONCE
text: We can recognize the Structured INFONCE upper and lower bounds for known p(x|z) (Poole et al. (2019) Sec. 2.5) as simply applying the standard IWAE bounds, using the marginal p(z) in place of the variational qθ(z|x)
Ep(x)p(z(1)|x)p(z(2:K))
[ log
p(x|z(1)) 1 K K∑ k=1 p(x|z(k))
] ≤ I(x; z) ≤ Ep(x)p(z(1:K)) [ log
1
1 K K∑ k=1 p(x|z(k))
] −H(x|z).
We refer to the lower bound as IS-INFONCEL(K) and the upper bound as IS-INFONCEU (K). From Prop. 2.1, we obtain an alternative proof that the structured INFONCE lower bound is upper bounded by logK. Since IBA(p(z)) ≤ IIWAEL(p(z),K) ≤ IBA(p(z)) + logK and the BA bound with a prior proposal equals 0 from Eq. 3, we have that 0 ≤ IS-INFONCE ≤ logK.

heading: C GENERALIZED IWAE

heading: C.1 PROBABILISTIC INTERPRETATION AND BOUNDS
text: To derive a probabilistic interpretation for GIWAE, our starting point is to further extend the state space of the IWAE target distribution in Eq. 4, using a uniform index variable p(s) = 1K ∀s that specifies which sample z(k) is drawn from the posterior p(z|x). This is shown in Fig. 4, and leads to a joint distribution over (z(1:K), s)
pGIWAETGT (z (1:K), s|x) = 1
K p(z(s)|x) K∏ k=1
k 6=s
q(z(k)|x) , (41)
with marginalization over s leading to pIWAETGT (z (1:K)|x) in Eq. 4. The posterior over the index variable s, which infers the ‘positive’ sample drawn from p(z|x) given a set of samples z(1:K), corresponds
to the normalized importance weights
pGIWAETGT (s|x, z(1:K)) = p(x,z(s)) qθ(z(s)|x)∑K k=1 p(x,z(k)) qθ(z(k)|x)
(42)
For the GIWAE extended state space proposal distribution, we consider a categorical index variable qGIWAEPROP (s|z(1:K),x) drawn according SNIS, with weights calculated using the energy function Tφ.
qGIWAEPROP (z (1:K), s|x) = ( K∏ k=1 qGIWAEPROP (z (k)|x) ) q(s|z(1:K),x) whereqGIWAEPROP (s|z(1:K),x) = eTφ(x,z (s))
K∑ k=1 eTφ(x,z (k)) .
(43)
Similarly to Lawson et al. (2019) and as shown in App. C.3, the induced SNIS distribution performs variational inference of the posterior pGIWAETGT (s|x, z(1:K)). As we show in Cor. C.3, the optimal energy function Tφ(x, z) = log p(x,z) q(z|x) + c(x) recovers the IWAE probabilistic interpretation of Domke & Sheldon (2018).
Upper Bound on log p(x) and Lower Bound on MI To derive bounds on the log partition function, we first calculate the unnormalized density ratio
log pGIWAETGT (z (1:K), s,x)
qGIWAEPROP (z(1:K), s|x) = log
1 K p(x, z (s)) K∏ k=1
k 6=s
q(z(k)|x)
q(s|z(1:K),x) K∏ k=1 q(z(k)|x) (44)
= log 1
K
∑ k e T (x,z(k))
eT (x,z(s)) p(x, z(s)) q(z(s)|x) (45)
= log p(x, z(s)) q(z(s)|x) − T (x, z(s)) + log 1 K K∑ k=1 eT (x,z (k)) (46)
Taking the expectation of the unnormalized log ratio under the proposal or target distribution yields a lower or upper bound, respectively, on log p(x)
EqGIWAEPROP
[ log pGIWAETGT (z (1:K), s,x)
qGIWAEPROP (z(1:K), s|x) ] ︸ ︷︷ ︸
ELBOGIWAE(qθ, Tφ,K)
≤ log p(x) ≤ EpGIWAETGT
[ log pGIWAETGT (z (1:K), s,x)
qGIWAEPROP (z(1:K), s|x) ] ︸ ︷︷ ︸
EUBOGIWAE(qθ, Tφ,K)
. (47)
To derive an explicit form for the upper bound on log p(x),
EUBOGIWAE(qθ, Tφ,K) = EpGIWAETGT (z(1:K),s|x)
[ log p(x, z(s))
q(z(s)|x) − T (x, z(s)) + log 1 K K∑ k=1 eT (x,z (k))
]
= Ep(z(s)|x) [ log p(x, z(s))
q(z(s)|x)
] − Ep(z(s)|x) ∏
k 6=s q(z(k)|x) log eT (x,z (s))
1 K K∑ k=1 eT (x,z(k))  (48) Translating this to the lower bound on MI from Eq. 9 corresponds to taking I(x; z) = −Ep(x) [log p(x)]−H(x|z) ≥ −EUBOGIWAE(qθ, Tφ,K)−H(x|z). Writing the conditional entropy
term over the index s,
I(x; z) ≥ Ep(z(s)|x)[log p(x|z(s))] (49)
− ( Ep(z(s)|x) [ log p(x, z(s))
q(z(s)|x)
] − Ep(z(s)|x)q(z(2:K)|x) log eT (x,z (s))
1 K K∑ k=1 eT (x,z(k))
 )
= Ep(x,z(s)) [ log
q(z(s)|x) p(z(s)) ] + Ep(x)p(z(s)|x)qθ(z(2:K)|x) [ log
eTθ(x,z (s))
1 K ∑K i=1 e Tθ(x,z(k))
] . (50)

heading: C.2 GIWAE UPPER BOUND ON MI DOES NOT PROVIDE BENEFIT OVER IWAE
text: Taking expectations under qGIWAEPROP to obtain a lower bound on log p(x) or upper bound on MI, notice that we will obtain expectations under the marginal of z(s), where z(s) is selected via SNIS as in Eq. 43. In translating this to a bound on mutual information, the conditional entropy term H(x|z) takes expectation with respect to p(z|x). In Eq. (50), this term canceled with the expectation over z(s) ∼ pGIWAETGT . However, this will not occur in the case of the GIWAE upper bound on MI. This implies that knowledge of the joint density p(x, z) is required to evaluate both the conditional entropy and log p(x,z (s))
q(z(s)|x) terms. However, if both p(z) and p(x|z) are known, then we will show in Cor. C.2 below that the optimal energy function in GIWAE yields the true importance weights and the resulting bound matches the IWAE upper bound (see Prop. 2.2). We thus conclude that this upper bound is not useful in practice.

heading: C.3 RELATIONSHIP BETWEEN IWAE AND GIWAE PROBABILISTIC INTERPRETATIONS
text: Proposition C.1. For a given qθ(z|x), the IWAE lower bound on MI is tighter than the GIWAE lower bound for any Tφ(x, z), and their difference is the average KL divergence between the normalized importance weights pGIWAETGT (s|z(1:K),x) and the variational distribution q(s|z(1:K),x) (Eq. 43)
IIWAEL(q,K) = IGIWAEL(q, T,K) + Ep(x)pGIWAETGT (z(1:K)|x) [ DKL[p GIWAE TGT (s|z(1:K),x)||qGIWAEPROP (s|z(1:K),x)] ] .
(51)
Proof. To show this result, we consider the posterior over the index variable p(s|z(1:K),x)
pGIWAETGT (s|z(1:K),x) = pGIWAETGT (z (1:K), s,x) K∑ s=1 pGIWAETGT (z(1:K), s,x) =
1 K p(z
(s)|x) ∏K
k=1 k 6=s qθ(z
(k)|x)
1 K K∑ s=1 p(z(s)|x) ∏K k=1 k 6=s qθ(z(k)|x) =
p(z(s)|x) q(z(s)|x)
K∑ s=1 p(z(s)|x) q(z(s)|x)
Now, note that p(z(1:K), s,x) = ( K∑ s=1 p(z(1:K), s,x) ) · p(s|z(1:K),x), so that we can decompose the KL divergence as
DKL[p GIWAE TGT (z (1:K), s,x)||qGIWAEPROP (z(1:K)s,x)] = EpGIWAETGT log K∑ s=1 p(z(1:K), s,x)
K∏ k=1 qθ(z(k)|x) p(s|z(1:K),x) q(s|z(1:K),x)  = DKL[ 1
K K∑ s=1 p(z(s)|x) K∏ k=1
k 6=s
qθ(z (k)|x)|| K∏ k=1 qθ(z (k)|x)] +DKL[p(s|z(1:K),x)||q(s|z(1:K),x)]
= DKL[p IWAE TGT (z (1:K)|x)||qIWAEPROP (z(1:K)|x)] +DKL[pGIWAETGT (s|z(1:K),x)||qGIWAEPROP (s|z(1:K),x)] Now, the left hand side equals the gap in the GIWAE upper bound on log p(x) in Eq. (47), while the first term on the RHS is the gap in the IWAE upper bound on log p(x) in Eq. 36.

heading: C.4 PROOF OF PROPOSITION 2.2 AND COROLLARY 2.3
text: We show the following results, which expand upon Prop. 2.2 and Cor. 2.3 in the main text.
Corollary C.2. The GIWAE lower bound has the following properties:
(a) For a given qθ(z|x), the optimal energy function equals the true log importance weights up to an arbitrary constant: T ∗(x, z) = log p(x,z)qθ(z|x) + c(x). With this choice of T ∗(x, z),
IGIWAEL(qθ, T ∗) = IIWAEL(qθ) (52)
In other words, when the joint distribution p(x, z) is known, the GIWAE bound does not provide any advantage over the IWAE bound.
(b) For a given qθ(z|x), let Tφ∗(x, z) denote the energy function that maximizes the GIWAE lower bound when optimization is performed with a neural network parameterized by φ. Then, Prop. 2.1, and Prop. C.1 imply
IBAL(qθ) ≤ IGIWAEL(qθ, Tφ∗ ,K) ≤ IIWAEL(qθ,K) ≤ IBAL(qθ) + logK . (53)
Corollary C.3. Using the prior qθ(z|x) = p(z) as in INFONCE,
(a) For the optimal energy function, Cor. C.2(a) implies T ∗(x, z) = log p(x|z) + c(x), and
IGIWAEL(p(z), T ∗,K) = IINFONCEL(T ∗,K) = IS-INFONCEL(K) (54)
(b) For an arbitrary energy function Tφ(x, z), Cor. C.2(b) translates to
0 ≤ IINFONCEL(Tφ,K) ≤ IS-INFONCEL(K) ≤ logK (55)
Proof. (a) We first show that the optimal energy function equals the log true importance weights up to a constant c(x), T (x, z) = log p(z|x)q(z|x) + c(x). This implies that learning the energy function using GIWAE becomes unnecessary when this density ratio is available in closed form as in IWAE.
Using Prop. C.1, we can see that the gap in the GIWAE and IWAE bounds, which corresponds to the posterior KL divergence over the index variable s, will equal zero when
qGIWAEPROP (s|z(1:K),x) = pGIWAETGT (s|z(1:K),x) =⇒ eT (x,z
(s))
K∑ k=1 eT (x,z(k)) =
p(z(s)|x) qθ(z(s)|x)
K∑ s=1 p(z(s)|x) qθ(z(s)|x)
(56)
=⇒ T (x, z) = log p(z|x) qθ(z|x) + c(x)
where the last line can include an arbitrary constant c(x) which does not depend on z, since this term will factor out of the denominator in Eq. 56.
For this choice of T ∗(x, z), we now show that the value of the GIWAE objective matches the IWAE lower bound on mutual information.
IGIWAE(q,T ∗) = Ep(x,z(1))
[ log
q(z(1)|x) p(z(1)) ] + Ep(x)p(z(1)|x)qθ(z(2:K)|x) log elog p(z (1)|x)
q(z(1)|x) · ec(x)( 1 K ∑K i=1 e log p(z(k)|x) qθ(z (k)|x) ) · ec(x)  (57)
= Ep(x,z(1))
[ log
q(z(1)|x) p(z(1)) + log p(z(1),x) q(z(1)|x) ] − Ep(x)p(z(1)|x)qθ(z(2:K)|x) [ log 1 K K∑ i=1 p(z(k),x) qθ(z(k)|x) ]
= −H(x|z) + ( −Ep(x)p(z(1)|x)qθ(z(2:K)|x) [ log 1
K K∑ i=1 p(z(k),x) qθ(z(k)|x) ]) = IIWAEL(qθ,K)
The second term is exactly the negative of the IWAE upper bound on log p(x) in Eq. (36), which provides a lower bound of H(x) = −Ep(x) [log p(x)]. Combined with the entropy −H(x|z), we obtain the IIWAE lower bound on MI. See Section 2.2.
(b) We begin by showing that IBAL(qθ) ≤ IGIWAEL(qθ, Tφ∗ ,K). This follows from the fact that there exist φ0 in the pearameter space of the neural network Tφ such that Tφ0 = 0. With this φ0, we would have IBAL(qθ) = IGIWAEL(qθ, Tφ0 ,K). So the optimal φ
∗ in the parameter space, can only improve upon the BA bound. Also note that for any given φ including φ∗, we have: IGIWAEL(qθ, Tφ∗ ,K) ≤ IIWAEL(qθ,K), as the Tφ∗ is the optimal energy function. This relationship is visualized in Fig. 3b.

heading: D GENERALIZED MUTUAL INFORMATION NEURAL ESTIMATION (GMINE)
text: In this section, we provide probabilistic interpretations for the MINE lower bounds on MI from Belghazi et al. (2018), which allows us to derive novel Generalized MINE bounds. In a similar spirit to GIWAE, these bounds use a base variational distribution qθ(z|x) to tighten the MINE-DV and MINE-F bounds and can be evaluated when p(z) is available.
We begin by deriving a probabilistic interpretation for the IBAL lower bound in Sec. 4, which we will show is closely related to the probabilistic interpretations of MINE. Our MINE-AIS method is designed to optimize and evaluate the IBAL lower bound on MI, with detailed discussion in Sec. 4 and App. E. See Fig. 5 for a summary of Generalized MINE bounds and their relationships.

heading: D.1 PROBABILISTIC INTERPRETATION OF IBAL
text: For the variational approximation
πθ,φ(z|x) = 1
Zθ,φ(x) qθ(z|x)eTφ(x,z) where Zθ,φ(x) = Eqθ(z|x)
[ eTφ(x,z) ] , (58)
we consider the BA lower bound on MI,
IBAL(πθ,φ) = I(x, z)− Ep(x) [DKL[p(z|x)||πθ,φ(z|x)]] = Ep(x,z) [ log p(x, z)
p(x)p(z)
] − Ep(x,z) [log p(z|x)] + Ep(x,z) [T (x, z)]− Ep(x,z) [logZθ,φ]
= Ep(x,z) [ log
qθ(z|x) p(z)
] + Ep(x,z) [Tφ(x, z)]− Ep(x) [logZθ,φ(x)]
= Ep(x,z) [ log
qθ(z|x) p(z) ] ︸ ︷︷ ︸
IBAL(qθ)
+Ep(x,z)
[ log
eTφ(x,z) Eqθ(z|x) [ eTφ(x,z) ]]︸ ︷︷ ︸ contrastive term
:= IBAL(qθ, Tφ) so that the gap in the mutual information bound is Ep(x) [ DKL[p(z|x)||πθ,φ(z|x)] ] . Note that IBAL(qθ, Tφ) contains IBAL(qθ) as one of its terms. We visualize the relationship between various energy-based bounds in Fig. 3b and Fig. 5.

heading: D.2 PROBABILISTIC INTERPRETATION OF GENERALIZED MINE-DV
text: We can interpret the MINE-DV bound of Belghazi et al. (2018) as arising from a joint energy-based variational approximation πθ,φ(x, z) of the full joint distribution p(x, z) as follows
πθ,φ(x, z) = 1
Z p(x)qθ(z|x)eTφ(x,z) where Z = Ep(x)qθ(z|x)
[ eTφ(x,z) ] . (59)
where qθ(z|x) is a base variational distribution, and Tφ is an energy function. Note that for the induced marginal of the energy-based approximation, πθ,φ(x) := ∫ πθ,φ(x, z)dz 6= p(x) due to the contribution of the energy term. Instead, we have
πθ,φ(x) = 1
Z p(x)Z(x) (60)
where Z(x) = ∫ q(z|x)eTφ(x,z)dz and Z = ∫ p(x)q(z|x)eTφ(x,z)dxdz = Ep(x)[Z(x)].
Subtracting a joint KL divergence DKL[p(x, z)||πθ,φ(x, z)] from I(x; z), we obtain the Generalized MINE-DV lower bound on MI
I(x; z) ≥ IGMINE-DV(qθ, Tφ) := I(x; z)−DKL[p(x, z)|| 1
Z p(x)qθ(z|x)eTφ(x,z)] (61)
= Ep(x,z) [ log
qθ(z|x) p(z) ] + Ep(x,z) [Tφ(x, z)]− logEp(x)qθ(z|x) [ eTφ(x,z) ] .
(62)
By construction, Eq. (61) shows that the gap in Generalized MINE-DV is DKL[p(x, z)|| 1Z p(x)qθ(z|x)e
Tφ(x,z)], which has a probabilistic interpretation in terms of the approximate joint distribution πθ,φ(x, z).
Relationship with MINE-DV For qθ(z|x) = p(z), we obtain the MINE-DV bound (Belghazi et al., 2018; Poole et al., 2019) as a special case with IMINE-DV(Tφ) = IGMINE-DV(p(z), Tφ). In particular, the joint base distribution in Eq. (59) corresponds to the product of marginals p(x)p(z). Our probabilistic interpretation shows that the gap in MINE-DV corresponds to DKL[p(x, z)|| 1Z p(x)p(z)e Tφ(x,z)].
We expect that the Generalized MINE, with a learned variational distribution qθ(z|x), to obtain tighter bounds than MINE-DV. We can guarantee that IGMINE-DV(q∗θ , Tφ) ≥ IMINE-DV(Tφ) for the optimal q∗θ with a given Tφ, so long as p(z) is in the variational family (i.e., ∃ θ0 such that qθ0(z|x) = p(z)).
Relationship with BA Bound Choosing a constant energy function Tφ0(x, z) = c, we can see that πθ,φ0(x, z) = p(x)qθ(z|x) and IGMINE-DV(qθ, Tφ0 = c) = IBAL(qθ) for a given qθ(z|x).
Optimal Energy Function For a given qθ(z|x), the optimal energy function of Generalized MINEDV corresponds to the log importance weights between the target p(x, z) and the joint base distribution
p(x)qθ(z|x), plus a constant
T ∗(x, z) = log p(x, z)
p(x)qθ(z|x) + c . (63)
For the optimal energy function associated with a given qθ(z|x), we have IMINE-DV(qθ, T ∗) = I(x; z).
Relationship with IBAL We can use our probabilistic interpretation to show that the IBAL is tighter than Generalized MINE-DV (i.e., IBAL(qθ, Tφ) ≥ IGMINE-DV(qθ, Tφ)), and the difference corresponds to a marginal KL divergence in the x space
IBAL(qθ, Tφ) = IGMINE-DV(qθ, Tφ) +DKL[p(x)||πθ,φ(x)] ≥ IGMINE-DV . (64)
where πθ,φ(x) is defined in Eq. (60). As in Poole et al. (2019), we can also interpret the gap between the IBAL(qθ, Tφ) and IGMINE-DV(qθ, Tφ) as an application of Jensen’s inequality, with DKL[p(x)||πθ,φ(x)] = logZ − Ep(x)[logZ(x)] = logEp(x)[Z(x)] − Ep(x)[logZ(x)] and logEp(x)[Z(x)] ≥ Ep(x)[logZ(x)] .

heading: D.3 PROBABILISTIC INTERPRETATION OF GENERALIZED MINE-F
text: The BA lower bound and its corresponding EUBO upper bound on log p(x) are derived using a KL divergence between the true and approximate posterior, which are normalized conditional distributions over z given x. In this section, we interpret the MINE-F bound (Belghazi et al., 2018) as arising from a generalized notion of the KL divergence between possibly unnormalized density functions. In particular, our probabilistic interpretation will involve an unnormalized π̃θ,φ(z|x) which seeks to approximate the true (normalized) posterior p(z|x). Generalized KL Divergence First, we recall the definition of the generalized KL divergence, which takes unnormalized measures as input arguments and corresponds to the limiting behavior of the α-divergence (Minka et al., 2005; Cichocki & Amari, 2010)
DGKL[r̃(z)||s̃(z)] = ∫ r̃(z) log r̃(z)
s̃(z) dz−
∫ r̃(z)dz + ∫ s̃(z)dz . (65)
As in the case of the standard KL divergence, this quantity is nonnegative, convex in either argument, and vanishes for r̃(z) = s̃(z). It is also a member of both the family of Bregman divergences and f -divergences (Amari, 2009). If r(z) and s(z) are normalized, then DGKL[r(z)||s(z)] = DKL[r(z)||s(z)]. Finally, if r(z) is normalized and s̃(z) is unnormalized, one can easily confirm that
DGKL[r(z)||s̃(z)] = DKL[r(z)||s(z)] +DGKL[s(z)||s̃(z)] (66) ≥ DKL[r(z)||s(z)] .
Generalized EUBO Since the Generalized KL divergence is always nonnegative, we can define an Generalized EUBO by adding the Generalized KL divergence between the normalized true posterior and an unnormalized approximate posterior π̃(z|x).
GEUBO(π̃) : = log p(x) +DGKL[p(z|x)||π̃(z|x)] (67) = Ep(z|x) [ log p(x, z)
π̃(z|x)
] − 1 + Zπ̃(x) (68)
where we assume that ∫ p(z|x)dz = 1 and let Zπ̃(x) := ∫ π̃(z|x)dz. Note that for normalized π(z|x), we have EUBO(π) = GEUBO(π). In general, if π is the normalized distribution of π̃, we have EUBO(π) ≤ GEUBO(π̃). Generalized BA Using the Generalized EUBO in place of log p(x), we obtain the following lower bound on MI, which we denote as the Generalized BA lower bound
I(x; z) ≥ IGBAL(π̃) : = I(x; z)− Ep(x) [DGKL[p(z|x)||π̃(z|x)]] = Ep(x,z) [ log
π̃(z|x) p(z)
] + 1− Ep(x) [Zπ̃(x)] . (69)
Generalized MINE-F We now consider an unnormalized, energy-based approximation to the true posterior, involving a base variational distribution qθ(z|x) and a learned energy function Tφ(x, z)
π̃θ,φ(z|x) = qθ(z|x)eTφ(x,z)−1 . (70) Using this unnormalized approximate posterior in the GBA lower bound in Eq. (69), we obtain the Generalized MINE-F (GMINE-F) lower bound I(x; z) ≥ IGMINE-F(qθ, Tφ) := I(x; z)− Ep(x) [DGKL[p(z|x)||π̃θ,φ(z|x)]] (71)
= Ep(x,z) [ log
qθ(z|x) p(z) ] + Ep(x,z) [Tφ(x, z)]− Ep(x)qθ(z|x) [ eTφ(x,z)−1 ] (72)
By construction, we can see that the gap in IGMINE-F(qθ, Tφ) is equal to the Generalized KL divergence Ep(x) [DGKL[p(z|x)||π̃θ,φ(z|x)]] in Eq. (71). As in the case of MINE-DV, we obtain the standard MINE-F lower bound IMINE-F(Tφ) = IGMINE-F(p(z), Tφ) when using the marginal p(z) as the proposal.
Optimal Energy Function The optimal energy function of Generalized MINE-F is T ∗(x, z) = 1 + log p(x,z)p(x)q(z|x) . In this case, we obtain IGMINE-F(qθ, T ∗) = I(x; z) in Eq. (72) (Poole et al., 2019).
Relationship with IBAL We would now like to relate the gap in IGMINE-F(qθ, Tφ) to the gap in IIBAL(qθ, Tφ). First, note that the normalized distribution corresponding to π̃θ,φ(z|x) = qθ(z|x)eTφ(x,z)−1 matches the the energy-based posterior in the IBAL, πθ,φ(z|x) =
1 Z(x)qθ(z|x)e Tφ(x,z). Thus, using Eq. (66), we can write
Ep(x) [DGKL[p(z|x)||π̃θ,φ(z|x)]] = Ep(x) [DKL[p(z|x)||πθ,φ(z|x)]] + Ep(x) [DGKL[πθ,φ(z|x)||π̃θ,φ(z|x)]]
Ep(x) [ DGKL[p(z|x)||qθ(z|x)eTφ(x,z)−1] ] ︸ ︷︷ ︸
gap of GMINE-F
= Ep(x) [ DKL[p(z|x)|| 1
Z(x)qθ(z|x)e Tφ(x,z)] ] ︸ ︷︷ ︸
gap of IBAL
(73)
+ Ep(x) [ DGKL[ 1
Z(x)qθ(z|x)e Tφ(x,z)||qθ(z|x)eTφ(x,z)−1] ] ︸ ︷︷ ︸
≥ 0 where the final term DGKL[πθ,φ(z|x)||π̃θ,φ(z|x)] ≥ 0 is nonnegative because it is a generalized KL divergence. After simplification, we can also see that DGKL[πθ,φ(z|x)||π̃θ,φ(z|x)] = Ep(x) [Z(x)− 1− logZ(x)] ≥ 0 since log u ≤ u − 1. We visualize the relationship between IBAL and Generalized MINE-DV in Fig. 5.
Relationship with Generalized MINE-DV To characterize the gap between Generalized MINEDV and Generalized MINE-F, we can again use the equality from Eq. (66), but this time using divergences over joint distributions.
DGKL[p(x, z)||π̃θ,φ(x, z)] = DKL[p(x, z)||πθ,φ(x, z)] +DGKL[πθ,φ(x, z)||π̃θ,φ(x, z)]
DGKL[p(x, z)||p(x)qθ(z|x)eTφ(x,z)−1]︸ ︷︷ ︸ gap in GMINE-F
= DKL[p(x, z)|| 1
Z p(x)qθ(z|x)eTφ(x,z)]︸ ︷︷ ︸
gap in GMINE-DV
(74)
+DGKL[ 1
Z p(x)qθ(z|x)eTφ(x,z)||p(x)qθ(z|x)eTφ(x,z)−1]︸ ︷︷ ︸
≥ 0 In this case, we have DGKL[ 1Z p(x)qθ(z|x)e
Tφ(x,z)||p(x)qθ(z|x)eTφ(x,z)−1] ≥ 0. Thus, Generalized MINE-DV is tighter than Generalized MINE-F (see Fig. 5), which generalizes the finding in Poole et al. (2019) that standard MINE-DV is tighter than standard MINE-F.

heading: E MINE-AIS

heading: E.1 PROPERTIES OF MINE-AIS
text: Proposition E.1. The lower bound IBAL(qθ, Tφ) used in MINE-AIS has the following properties:
(a) For a given qθ(z|x), the optimal energy function equals the true log importance weights up to an arbitrary constant: T ∗(x, z) = log p(x,z)qθ(z|x) + c(x). With this choice of T ∗(x, z),
IBAL(qθ, T ∗) = I(x; z) (75)
and Cor. C.2(a) implies that
lim K→∞
IGIWAEL(qθ, T ∗,K) = lim
K→∞ IIWAEL(qθ,K) = IBAL(qθ, T
∗) = I(x; z) (76)
(b) Suppose the energy function Tφ(x, z) is parameterized by φ, and that ∃φ0 s.t.∀ (x, z), Tφ0(x, z) = const. For a given qθ(z|x), let Tφ∗(x, z) denote the energy function that maximizes the MINE-AIS lower bound IBAL(qθ, Tφ), then we have
IBAL(qθ) ≤ IBAL(qθ, Tφ∗) ≤ I(x; z) = IBAL(qθ) +DKL[p(z|x)||qθ(z|x)] (77)
In particular, the contrastive term in Eq. 22 is upper bounded by DKL[p(z|x)||qθ(z|x)]. See Fig. 3b.
Proof. (a) Since the gap in the IBAL(qθ, Tφ∗) is Ep(x) [ DKL[p(z|x)||πθ,φ(z|x)] ] = I(x; z) − IBAL(qθ, Tφ∗). This implies that the bound will be optimal iff p(z|x) = πθ,φ(z|x) ∝ p(x, z). We can easily show that the true log importance weights (plus a constant) satisfy this property and thus leads to the optimal energy function T ∗(x, z) = log p(x, z)/qθ(z|x),
πθ,φ(z|x) = 1
logZθ,φ(x) qθ(z|x)e
log p(x,z) qθ(z|x) +c(x) =
ec(x)
logZθ,φ(x) p(x, z) = p(z|x) (78)
where ec(x) is absorbed into the normalization constant. At this value, the IBAL reduces to I(x; z)
IBAL(qθ, T ∗) = Ep(x,z) [ log
qθ(z|x) p(z)
] + Ep(x,z) [Tφ(x, z)]− Ep(x) [logZθ,φ(x)] (79)
= Ep(x,z) [ log
qθ(z|x) p(z)
] + Ep(x,z) [ log p(x, z)
qθ(z|x) + log c(x))
] (80)
− Ep(x) [ logEqθ(z|x) p(x, z)
qθ(z|x) + log c(x)
] (81)
= Ep(x,z) [ log p(x, z)
p(z)p(x)
] = I(x; z) (82)
(b) The BA bound is a special case of the IBAL(qθ, Tφ) with constant Tφ0 = c. Thus, by learning Tφ, we can only improve upon IBAL(qθ). The paramterized family of Tφ may not be expressive enough to match the true log importance weights, so IBAL(qθ, Tφ∗) ≤ IBAL(qθ, T ∗) = I(x; z) by part (a).

heading: E.2 PROOF OF PROPOSITION 4.1
text: Proposition 4.1. For given qθ(z|x) and Tφ(x, z), lim K→∞ IGIWAEL(qθ, Tφ,K) = IBAL(qθ, Tφ).
Proof. Inspecting the form of IGIWAEL(qθ, Tφ,K) in Eq. 9 to the IBAL(qθ, Tφ) in Eq. 22, we can see that for a fixed qθ and Tφ, both bounds include the same first term IBAL(qθ) and the same numerator of the contrastive term Ep(x,z) [ eTφ(x,z) ] . To prove the proposition, we can thus focus on characterizing the limiting behavior of the denominator
lim K→∞ Ep(z(1)|x)qθ(z(2:K)|x)
[ log 1
K K∑ k=1 eTφ(x,z) ] ︸ ︷︷ ︸
:= log ẐGIWAE(K,x)
= logEqθ(z|x) [ eTφ(x,z) ] ︸ ︷︷ ︸
:= logZθ,φ(x)
∀x (83)
where we introduce the notation log ẐGIWAE(K,x) for convenience and the right hand side is the log partition function for the MINE-AIS distribution πθ,φ(z|x). Intuitively, we expect Eq. 83 to hold
since the contribution of the single posterior sample p(z|x) in the GIWAE expectation will vanish as K →∞. More formally, we consider the sequence of values log ẐGIWAE(K,x) as a function of K. We derive lower and upper bounds on the value of log ẐGIWAE(K,x) for fixed K and show that each of these sequences of lower and upper bounds converge to logZθ,φ(x) as K → ∞. Using the squeeze theorem for sequences, this is sufficient to demonstrate the claim that limK→∞ log ẐGIWAE(K,x) = logZθ,φ(x) in Eq. 83. Since the other terms in IGIWAEL(qθ, Tφ,K) and IBAL(qθ, Tφ) are identical, this is sufficient to prove the proposition.
Lower Bound: To lower bound log ẐGIWAE(K,x), we rely on the fact that the exponential function eTφ(x,z) ≥ 0 to simply ignore the contribution of the p(z|x) term.
log ẐGIWAE(K,x) = Ep(z(1)|x)qθ(z(2:K)|x)
[ log
K∑ k=1 eTφ(x,z)
] − logK (84)
≤ Eqθ(z(2:K)|x)
[ log
K∑ k=2 eTφ(x,z)
] − logK (85)
= Eqθ(z(2:K)|x)
[ log 1
K − 1 K∑ k=2 eTφ(x,z)
] + log
K − 1 K
(86)
≤ log 1 K − 1 K∑ k=2 Eqθ(z(2:K)|x) [ eTφ(x,z) ] + log K − 1 K
(87)
= logEqθ(z(2:K)|x) [ eTφ(x,z) ] + log K − 1 K
(88) := LB ( log ẐGIWAE(K,x) ) (89)
As K →∞, we have log K−1K → 0 and
lim K→∞
LB ( log ẐGIWAE(K,x) ) = logEqθ(z|x) [ eTφ(x,z) ] = logZθ,φ(x) (90)
Upper Bound: To upper bound log ẐGIWAE(K,x), we separately consider terms arising from p(z|x) samples and qθ(z|x) samples. Noting that pGIWAETGT (z(1:K), s|x) in Eq. (41) is invariant to the index s, we can assume z(1) ∼ p(z|x) and write
log ẐGIWAE(K,x) = Ep(z(1)|x)qθ(z(2:K)|x)
[ log ( 1
K
∑ eTφ(x,z (1)) + 1
K K∑ k=2 eTφ(x,z (k))
)] (91)
= Ep(z(1)|x)qθ(z(2:K)|x)
[ log ( 1
K eTφ(x,z (1)) + K − 1 K · 1 K − 1 K∑ k=2 eTφ(x,z (k)) )] (92)
We consider the term inside the log as an expectation over a Bernoulli random variable with probabilities { 1K , K−1 K }. Using the concavity of log to bring this expectation outside, and the linearity of expectation to separate terms, we have
log ẐGIWAE(K,x) ≤ 1
K Ep(z(1)
[ Tφ(x, z (1)) ] + K − 1 K Eqθ(z(2:K)|x) [ log 1 K − 1 K∑ k=2 eTφ(x,z (k)) ] (93)
:= UB ( log ẐGIWAE(K,x) ) (94)
As K →∞, we have 1K → 0, K−1 K → 1, and 1 K−1 K∑ k=2 eTφ(x,z (k)) → Eqθ(z|x) [ eTφ(x,z) ] . Finally,
lim K→∞
UB ( log ẐGIWAE(K,x) ) = Eqθ(z|x) [ eTφ(x,z) ] = logZθ,φ(x) (95)
As reasoned above, the convergence of both the upper and lower bound sequences to logZθ,φ(x) implies limK→∞ log ẐGIWAE(K,x) = logZθ,φ(x) and limK→∞ IGIWAEL(qθ, Tφ) = IBAL(qθ, Tφ), as desired.

heading: E.3 CONVERGENCE OF GIWAE SNIS DISTRIBUTION TO MINE-AIS
text: Proposition E.2. Let qGIWAEPROP K(z|x) denote the marginal SNIS distribution of GIWAE
qGIWAEPROP (z|x;K) = K∑ s=1 ∫ qGIWAEPROP (s, z (1:K)|x)δ(z− z(s))dz(1:K) (96)
= K∑ s=1 ∫ δ(z− z(s)) ( K∏ k=1 q(z(k)|x) ) eTφ(x,z (s))
K∑ k=1 eTφ(x,z (k)) dz(1:K) (97)
Then as K → ∞, the KL divergence between the marginal SNIS distribution of the GIWAE and energy-based variational distribution of MINE-AIS goes to zero.
lim K→∞
DKL[q GIWAE PROP (z|x;K)||πθ,φ(z|x)] = 0 and lim K→∞ DKL[πθ,φ(z|x)||qGIWAEPROP (z|x;K)] = 0 (98)
Proof. Note that the distribution qGIWAEPROP (z|x;K) is obtained by augmenting the standard (multisample) GIWAE proposal distribution qGIWAEPROP (z
(1:K), s|x) using an additional variable z = z(s) which copies the sample in index s ∈ [1,K]. This index is chosen according to the SNIS weights induced by the energy function in Eq. (97), and we consider the distribution qGIWAEPROP (z|x;K) over z after marginalizing the joint qGIWAEPROP (z, s, z (1:K)) over z(1:K) and s as in Eq. (96)-97.
To prove the proposition, we consider a mixture target distribution similar to the case of IWAE. However, in this case, we take a single sample from πθ,φ(z|x) instead of p(z|x)
p πθ,φ TGT (z, s, z (1:K),x) = δ(z− z(s)) 1 K πθ,φ(z
(s)|x) K∏
k=1,k 6=s
qθ(z (s)|x) (99)
Note that marginalization over z will always yield a sample from πθ,φ(z|x). Since marginalization can only reduce the KL divergence, we have
DKL[q GIWAE PROP (z|x;K)||πθ,φ(z|x)] ≤ DKL[qGIWAEPROP (z, s, z(1:K)|x)||p πθ,φ TGT (z, s, z (1:K)|x)] (100) = DKL[q GIWAE PROP (s, z (1:K)|x)||pπθ,φTGT (s, z(1:K)|x)] (101)
where the second line follows because z is a deterministic function of z(1:K) and s. Now, we can show that this KL divergence goes to 0 as K →∞
DKL[q GIWAE PROP (s, z (1:K)|x)||pπθ,φTGT (s, z(1:K)|x)] = EqGIWAEPROP log K∏ k=1 q(z(k)|x) e Tφ(x,z (s)) K∑ k=1 eTφ(x,z (k))
1 Kπθ,φ(z
(s)|x) K∏
k=1,k 6=s q(z(k)|x)

= EqGIWAEPROP log q(z(s)|x)eTφ(x,z (s)) q(z(s)|x)eTφ(x,z(s)) logZθ,φ(x)
1 K K∑ k=1 eTφ(x,z (k))  = logZθ,φ(x)− EqGIWAEPROP [ log 1
K K∑ k=1 eTφ(x,z (k))
] (102)
Since qGIWAEPROP (z (1:K)|x) = ∏K k=1 qθ(z (k)|x), it is clear that in the limit as K → ∞,
EqGIWAEPROP [ log 1K K∑ k=1 eTφ(x,z (k)) ] → logEqθ(z|x) [ eTφ(x,z (k)) ] = logZθ,φ(x). Thus, the difference
in Eq. 102, the extended state space KL divergence in Eq. 101, and the marginal SNIS KL divergence in Eq. 100 all go to 0 as K →∞. As desired, lim
K→∞ DKL[q
GIWAE PROP (z|x;K)||πθ,φ(z|x)] = 0.
Similar reasoning as from Eq. 100-Eq. 102 holds for the KL divergence in the reverse direction,
with DKL[p πθ,φ TGT (s, z (1:K)|x)||qGIWAEPROP (s, z(1:K)|x)] = Epπθ,φTGT [ log 1K K∑ k=1 eTφ(x,z (k)) ] − logZθ,φ(x). The left term matches log ẐGIWAE(K,x) in Eq. 83 and the proof of Prop. E.3, apart from the fact that the ‘positive’ sample comes from πθ,φ instead of role of p(z|x). Using identical reasoning as
above, the effect of this sample will vanish as K → ∞. Thus, E p πθ,φ TGT [ log 1K K∑ k=1 eTφ(x,z (k)) ] → logZθ,φ(x). Finally, this implies lim K→∞ DKL[p πθ,φ TGT (s, z
(1:K)|x)||qGIWAEPROP (s, z(1:K)|x)] = 0 and thus lim K→∞ DKL[πθ,φ(z|x)||qGIWAEPROP (z|x;K)] = 0 as desired.

heading: E.4 DETAILS FOR TRAINING MINE-AIS
text: Eq. 24 indicates that in order to maximize the IBAL as a function of θ, φ, we need to increase the value of the energy function Tφ or score function log qθ(z|x) on the real samples of the true joint p(x, z) and decrease the value on fake samples from the approximate p(x)πθ,φ(z|x). However, as is common in training energy-based models, it is difficult to draw samples from πθ,φ.
In order to sample from πθ,φ(z|x), a natural approach is to initialize MCMC chains from a sample of the base distribution qθ(z|x), using HMC transition kernels (Neal, 2011), for example. However, we may require infeasibly long MCMC chains when the base distribution is far from desired energy-based model πθ,φ(z|x). Instead, we can choose to initialize the chains from the true posterior sample z ∼ p(z|x) for a simulated data point x ∼ p(x). This initialization greatly reduces the computational cost and variance of the estimated gradient and enables training energy functions in high dimensional latent spaces, as shown in our experiments. This approach is in spirit similar to Contrastive Divergence learning of energy-based models (Hinton, 2002) where, in order to sample from the model distribution, we start an MCMC chain from the true data distribution to obtain a lower variance gradient estimate.
We use T = 10 MCMC steps in our experiments. After obtaining an approximate sample from πθ,φ(z|x), we simply update the parameters θ using the gradients in Eq. 24.

heading: E.5 APPROXIMATE REVERSE ANNEALING FOR EVALUATING MINE-AIS
text: Here, we argue that we can maintain an approximate upper bound on Zθ∗,φ∗(x) and lower bound on IBAL by sampling from p(z|x) instead of πθ,φ(z|x) in reverse importance sampling or reverse AIS.
Since these samples are unavailable, Prop. E.3 argues that we can preserve an approximate upper bound on logZθ(x) and lower bound on IBAL(qθ, Tφ) by sampling from p(z|x) instead of πθ,φ(z|x) in reverse importance sampling.
Proposition E.3. Suppose Tφ(x, z) is a neural network parameterized θ and πθ,φ(z|x) = 1 Zθ,φ(x)qθ(z|x)e Tφ(x,z). Suppose πθ∗,φ∗(z|x) is the optimal MINE-AIS solution, within a variational family for qθ(z|x), and Tφ parameterized by a neural network. Then, we would have:
Ep(z|x) [ log πθ∗,φ∗(x, z)
qθ∗(z|x)
] ≥ logZθ∗,φ∗(x) (103)
Proof. Note that according to Prop. E.1(b), MINE-AIS optimal solution πθ∗,φ∗(z|x) within a variational family improves upon the BA bound. This means, the gap of the MINE-AIS bound DKL[p(z|x)||πθ∗,φ∗(z|x)] is less that the gap of the BA bound DKL[p(z|x)||qθ∗(z|x)]. Thus we have:
Ep(z|x) [ log πθ∗,φ∗(x, z)
qθ∗(z|x)
] = DKL[p(z|x)||qθ∗(z|x)]−DKL[p(z|x)||πθ∗,φ∗(z|x)] + logZθ∗,φ∗(x)
≥ logZθ∗,φ∗(x) (104)
As an illustrative example, Fig. 6a shows a Gaussian base distribution qθ∗(z|x), a multi modal true posterior distribution p(z|x), a posterior approximation πθ∗,φ∗(z|x), whose energy function Tφ parameterized by µφ, σφ has been fit to minimize DKL[p(z|x)||πθ,φ(z|x)]. The standard reverse importance sampling, which take true samples from πθ,φ(z|x), results in an upper bound on logZθ,φ(x). This is because the initial distribution qθ(z|x) is missing some of the modes of the target distribution πθ,φ(z|x), and when we sample from πθ,φ(z|x), particles from these modes will yield large importance weights since they have low density under qθ(z|x), resulting in an upper bound on logZθ,φ(x). Since MINE-AIS optimizes the reverse KL divergence, the learnt πθ,φ(z|x) will conservatively allocate mass across the support of p(z|x). Prop. E.3 shows that sampling from p(z|x) instead of πθ,φ(z|x) will still yield large importance weights and preserve the upper bound of reverse importance sampling.
While Prop. E.3 shows that we can obtain a single-sample upper bound on logZθ,φ(x), these mathematical guarantees do not immediately translate to Multi-Sample AIS upper bounds, which includes additional transitions. However, we empirically observe in our experiments that this procedure does provide an upper bound on logZθ(x). For example, Fig. 6b shows the performance of the forward AIS starting from qθ(z|x) and reverse AIS starting from both πθ,φ(z|x) and p(z|x). As expected, the forward AIS provides a lower-bound on logZθ,φ(x). We can also see that both reverse AIS chains provide an upper bound on the logZθ,φ(x).

heading: E.6 MULTI-SAMPLE AIS CONVERGENCE RESULTS FOR MINE-AIS
text: In Fig. 7, we compare the performance of our MINE-AIS method, with its learned energy function Tφ, and evaluates Multi-Sample AIS lower and upper bounds on IBAL(qθ, Tφ) ≤ I(x; z) in cases where p(x|z) is unknown. As a comparison, we consider Multi-Sample AIS lower and upper bounds with known p(x|z), since this provides bounds on MI and represents the optimal energy function in MINE-AIS (Prop. E.3).
We also compare the convergence of the Multi-Sample AIS evaluation procedure as a function of T , by plotting our lower and upper bounds as a function of the number of intermediate T used for evaluation of IBAL(p(z), Tφ). While the upper bound on IBAL may not actually lower bound the true MI for low T , we see that our bounds converge to the same estimate for high T .
At convergence, we can compare AIS estimates of MI and the IBAL for MINE-AIS to indicate the quality of the learned energy function. On the other hand, since the IBAL corresponds to the BA lower bound for a flexible, energy-based posterior approximation, we also compare against the BA
lower bound for qθ(z|x) within a Gaussian variational family. We observe that IBAL lower bounds on MI fall between these two baselines, with MINE-AIS underestimating the value obtained from Multi-Sample AIS by 1%, 18%, and 28% on linear VAE, MNIST-VAE, and MNIST-GAN, respectively.

heading: F APPLICATIONS TO MUTUAL INFORMATION ESTIMATION WITHOUT KNOWN MARGINALS
text: While the focus in of our work is evaluating the mutual information I(x; z) in settings where at least a single marginal is available, we are often interested in estimating or optimizing the mutual information where no marginal distribution is available. A natural setting where no marginal distribution is available is representation learning where the goal is to maximize the mutual information between the data distribution qdata(x) and the representation induced by a stochastic mapping qψ(z|x) parameterized by ψ,
I(x; z) = Eqdata(x)qψ(z|x) [ log qψ(x, z)
qdata(x)qψ(z)
] = Eqdata(x)qψ(z|x) [ log
qψ(x|z) qdata(x)
] (105)
where qψ(z) = ∫ qdata(x)qψ(z|x)dx represents the ‘aggregated posterior’ (Makhzani et al., 2015) or induced marginal distribution over z.

heading: F.1 BA LOWER BOUND
text: Using the notation of Eq. 105, we can write the BA lower bound as I(x; z) ≥ IBAL(pθ(x|z)) : = I(x; z)− Eqψ(z) [ DKL[qψ(x|z)||pθ(x|z)] ] ︸ ︷︷ ︸
gap = Eqdata(x)Eqψ(z|x) [ log
pθ(x|z) qdata(x)
] (106)
= Eqdata(x) Eqψ(z|x) [ log pθ(x|z) ] ︸ ︷︷ ︸ negative reconstruction loss of x +Hdata(x)︸ ︷︷ ︸ constant
(107)
where pθ(x|z) is a variational distribution, parameterized by θ, that tries to match to the inverse encoding distribution qψ(x|z) ∝ qdata(x)qψ(z|x). The first term in Eq. 107 can be interpreted as the reconstruction term, and the second term is the entropy of the data distribution, which is constant.
Evaluating MI up to a Constant Note that the gradient of IBAL(pθ(x|z)) with respect to the parameters θ does not depend on the marginal qdata(x), and thus we may still optimize the variational distribution when the data distribution is unknown. We can then use the resulting pθ(x|z) to estimate
the BA lower bound on mutual information up to a constant, Hdata(x). This is useful in comparing the MI induced by two different representations qψ1(z|x) and qψ2(z|x) of the same data distribution.
Optimizing MI with the BA Lower Bound The BA lower bound is also amenable to backpropagation through the parameters of the encoding distribution qψ(z|x), even when analytic marginal densities for qdata(x) or qψ(z) are not available.
Maximization of the BA lower bound appears in various settings, including in representation learning (Alemi et al., 2016; 2018), reinforcement learning (Mohamed & Rezende, 2015), and improving interpretability in GANs (Chen et al., 2016). In particular, in the information bottleneck method (Tishby et al., 2000), we are interested in maximizing the ‘relevant information’ I(z;y) between a learned representation and a target variable y. In supervised learning, y may represent classification labels (Alemi et al., 2016). When y is the original data as in unsupervised learning with β-VAEs (Alemi et al., 2018), the distortion or reconstruction loss corresponds to maximizing the BA lower bound on the encoding mutual information in Eq. (105).

heading: F.2 GIWAE LOWER BOUND
text: In this section, we discuss the applicability of our GIWAE lower bound in mutual information maximization settings.
Rewriting the GIWAE lower bound in Eq. (9) for the setting of estimating I(x; z) in Eq. (105), we have
I(x; z) ≥ IGIWAEL(pθ(x|z), Tφ,K) = Eqdata(x) Eqψ(z|x) [ log pθ(x|z) ] ︸ ︷︷ ︸ negative reconstruction loss of x +Hdata(x)︸ ︷︷ ︸ constant
(108)
+ Eqψ(x,z)pθ(x(2:K)|z)
[ log
eTφ(x,z)
1 K (eTφ(x,z) + ∑K k=2 e Tφ(x (k),z)) ] ︸ ︷︷ ︸
contrastive term≤logK
.
Note that we can use a single joint sample from qdata(x)qψ(z|x) = qψ(z)qψ(x|z) to obtain a positive sample from the inverse encoding distribution x ∼ qψ(x|z) for a particular z ∼ qψ(z), in similar fashion to our ancestral sampling in Sec. 1.1. For a given z, the negative samples correspond to K− 1 samples from the stochastic decoder x(2:K) ∼ pθ(x|z). From the importance sampling perspective, we can view the GIWAE lower bound as corresponding to an upper bound on log qψ(z) for a particular z ∼ qψ(z). The contrastive term in Eq. (108) arises from SNIS sampling of a single x(s) in the extended state space proposal, with qGIWAEPROP (s|z,x1:K) = eTφ(x (s),z)/ ∑K k=1 e Tφ(x (k),z).
Finally, we note that optimization over the energy function Tφ(x, z) improves upon IBAL(pθ(x|z)) in Eq. (107) by at most logK nats using the contrastive term.
Evaluating MI up to a Constant Similar to the BA bound, the GIWAE lower bound can be used to evaluate the mutual information up to the constant data entropy when qdata(x) is unknown. This may be useful in comparing the MI induced by two different representations qψ1(z|x) and qψ2(z|x) of the same data distribution. Note that the reconstruction and contrastive terms in Eq. (108) depend only on samples from qdata(x), so that our ability to take gradients with respect to the parameters of the variational distribution pθ(x|z) or energy function Tφ(x, z) are not affected by the fact that the marginal distribution is unknown.
Optimizing MI with the GIWAE Lower Bound The GIWAE lower bound may also be used for mutual information maximization with respect to the parameters of qψ(z|x), since each term in Eq. (108) is amenable to backpropagation. Since GIWAE generalizes both the BA and INFO-NCE lower bounds, it can be used as a drop-in replacement for either of these bounds for optimizing MI.
Contrastive INFO-NCE bounds appear in representation learning across a variety of settings and modalities, including images (Hjelm et al., 2018; van den Oord et al., 2018; Bachman et al., 2019;
Chen et al., 2020), sequential data (van den Oord et al., 2018), reinforcement learning (Mazoure et al., 2020; van den Oord et al., 2018), or enforcing fairness constraints (Gupta et al., 2021).

heading: F.3 MINE-AIS / IBAL LOWER BOUND
text: Recall that MINE-AIS estimation would involve an energy based variational approximation to the inverse encoding distribution qψ(x|z),
πθ,φ(x|z) = 1
Z(z) pθ(x|z)eTφ(x,z) . (109)
Note that IBAL lower bound on I(x; z) involves an intractable log partition function term Eqψ(z)[logZ(z)] = Eqψ(z) [ logEpθ(x|z) [ eTφ(x,z) ] ] .
I(x; z) ≥ IBAL(pθ(x|z), Tφ) = Eqdata(x) Eqψ(z|x) [ log pθ(x|z) ] ︸ ︷︷ ︸ negative reconstruction loss of x +Hdata(x)︸ ︷︷ ︸ constant
(110)
+ Eqψ(x,z)
[ log
eTφ(x,z) Epθ(x|z) [ eTφ(x,z) ]]︸ ︷︷ ︸ contrastive term≤Eqψ(z)[DKL[qψ(x|z)||pθ(x|z)]]
where the term in the denominator is the partition function Z(z) = Epθ(x|z) [ eTφ(x,z) ] .
When taking gradients as in Eq. (24)-(25), we obtain
∂
∂θ IBAL(pθ, Tφ) = Eqdata(x)qψ(z|x)
[ ∂
∂θ log pθ(x|z)
] − Eqψ(z)πθ,φ(x|z) [ ∂
∂θ log pθ(x|z)
] . (111)
∂
∂φ IBAL(pθ, Tφ) = Eqdata(x)qψ(z|x)
[ ∂
∂φ Tφ(x, z)
] − Eqψ(z)πθ,φ(x|z) [ ∂
∂φ Tφ(x, z)
] . (112)
To obtain approximate negative samples for a given z, we can use MCMC transition kernels which only require access to the unnormalized target density pθ(x|z)eTφ(x,z). This can be evaluated directly since both pθ(x|z) and Tφ(x, z) are tractable.
Evaluating MI up to a Constant Since only samples from qdata(x) are required for the MINE-AIS training procedure in Eq. (111)-(112), we can learn the base variational distribution pθ(x|z) and energy function Tφ(x, z) in cases when the marginal distribution is unknown.
As in the case of GIWAE and BA lower bounds above, we can also evaluate the IBAL lower bound in Eq. (110) up to a constant. As in the main text, this involves using multi-sample AIS techniques to bound the intractable log partition function term logZ(z) = logEpθ(x|z) [ eTφ(x,z) ] .
Optimizing MI with the MINE-AIS Lower Bound If we are interested in optimizing the MINEAIS lower bound with respect to the parameters of a stochastic encoder qψ(z|x), we would need to backpropagate through the multi-sample AIS procedure used to estimate the logZ(z) term. We thus conclude that, among our proposed methods, GIWAE is the most directly applicable in settings of mutual information maximization for representation learning.

heading: G BACKGROUND ON AIS
text: AIS (Neal, 2001) constructs a sequence of intermediate distributions {πt(z)}Tt=0, which bridge between a normalized initial distribution π0(z|x) and target distribution πT (z|x) = p(z|x) with unnormalized density πT (x, z) = p(x, z) and normalizing constant ZT (x) = p(x). A common choice for intermediate distributions is the geometric mixture path parameterized by {βt}Tt=0:
πt(z|x) = π0(z|x)1−βt πT (x, z)βt
Zt(x) where Zt(x) =
∫ π0(z|x)1−βt πT (x, z)βtdz. (113)
AIS involves a extended state space proposal distribution qAISPROP(z0:T |x), obtained by sampling from the initial π0(z|x) and constructing transitions Tf (zt+1|zt,x) which leave πt(z|x) invariant. The target distribution pAISTGT(z0:T |x) is given by running the reverse transitions Tr(zt|zt+1,x) starting from a target or posterior sample πT (z|x), as shown in Fig. 2.
qAISPROP(z0:T |x) = π0(z0|x) T−1∏ t=0 Tf (zt+1|zt) pAISTGT(x, z0:T ) = πT (x, zT ) T−1∏ t=0 Tr(zt|zt+1,x)
(114) Since the forward and reverse kernels obey Tf (zt+1|zt,x)πt(zt|x) = Tr(zt|zt+1,x)πt(zt+1|x), we can write the importance weights w := pAISTGT(x, z0:T )/q AIS PROP(z0:T |x) in the extended state space as
w = pAISTGT(x, z0:T ) qAISPROP(z0:T |x) = πT (x, zT ) π0(z0|x) T−1∏ t=0 πt(x, zt) πt(x, zt+1) = T∏ t=1 πt(x, zt) πt−1(x, zt) = T∏ t=1 ( πT (x, zt) π0(zt|x) )βt−βt−1 (115)
Assuming the initial distribution is normalized Z0 = 1, AIS provides the following estimate of the partition function ẐT (x)
ẐT (x) = 1
K K∑ k=1 w(k) where w(k) = pAISTGT(x, z (k) 0:T ) qAISPROP(z (k) 0:T |x) and z(k)0:T ∼ q AIS PROP(z0:T |x), (116)
which is unbiased since E[ẐT (x)] = Ez0:T∼qAISPROP [w] = ZT (x). Using the general approach in App. A, taking expectations of the log importance weights under the proposal and target yields the following lower and upper bounds on the log partition function.
Ez0:T∼qAISPROP
[ log pAISTGT(x, z0:T )
qAISPROP(z0:T |x) ] ︸ ︷︷ ︸
ELBOAIS(π0, T )
≤ log p(x) ≤ Ez0:T∼pAISTGT
[ log pAISTGT(x, z0:T )
qAISPROP(z0:T |x) ] ︸ ︷︷ ︸
EUBOAIS(π0, T )
(117)
The lower bound is an extended state-space ELBO using the proposal and target distributions in Eq. (14), with gap equal to DKL[qAISPROP(z0:T |x)||pAISTGT(z0:T |x)]. The upper bound is an extended state space EUBO, whose gap is equal to DKL[pAISTGT(z0:T |x)||qAISPROP(z0:T |x)]. These single-chain lower and upper bounds on log p(x) translate to upper and lower bounds on MI, IAISL(π0, T ) and IAISU (π0, T ), which were suggested for MI estimation in the blog post of Sobolev (2019a).
Similar single-sample bounds also appear in variational inference methods (Salimans et al., 2015; Masrani et al., 2019), and backpropagation through these estimators in generative modeling constitutes an active area of research (Geffner & Domke, 2021; Zhang et al., 2021).
H INDEPENDENT MULTI-SAMPLE AIS

heading: H.1 PROBABILISTIC INTERPRETATION AND BOUNDS
text: We can expand the probabilistic interpretation of Independent Multi-Sample AIS from the main text Eq. 120 and Fig. 2 using notation for the full AIS chains
qIM-AISPROP (z (1:K) 0:T |x) = K∏ k=1 qAISPROP(z (k) 0:T |x) (118)
= K∏ k=1 π0(z (k) 0 |x) T−1∏ t=0 Tf (z(k)t+1|z (k) t ) (119)
pIM-AISTGT (z (1:K) 0:T ,x) =
1
K K∑ s=1 pAISTGT(x, z (s) 0:T ) K∏ k=1
k 6=s
qAISPROP(z (k) 0:T |x). (120)
= 1
K K∑ k=1 πT (z (k) T |x) T∏ t=1 Tr(z(k)t−1|z (k) t ) · K∏ j=1
j 6=k
π0(z (k) 0 |x) T−1∏ t=0 Tf (z(k)t+1|z (k) t )
The importance ratio log pIM-AISTGT (x, z)/q IM-AIS PROP (x, z) simplifies to a mixture as in Eq. 30, although we are now averaging the importance weights along AIS chains
log pIM-AISTGT (z (1:K) 0:T ,x)
qIM-AISPROP (z (1:K) 0:T ,x)
= log 1
K K∑ k=1 pAISTGT(x, z (k) 0:T ) qAISPROP(z (k) 0:T |x) = log 1 K K∑ k=1 T−1∏ t=0
( p(x, z
(k) t )
q(z (k) t |x)
)βt+1−βt = log 1
K K∑ k=1 w(k)
(121)
Taking the expected log ratio under the proposal and target, respectively, yield lower and upper bounds on log p(x)
E z (1:K) 0:T ∼qAISPROP
[ log 1
K K∑ k=1 pAISTGT(x, z (k) 0:T ) qAISPROP(z (k) 0:T |x) ] ︸ ︷︷ ︸
ELBOIM-AIS(π0,K, T )
≤ log p(x) ≤ E z (1) 0:T ∼pAISTGT
z (2:K) 0:T ∼q AIS PROP
[ log 1
K K∑ k=1 pAISTGT(x, z (k) 0:T ) qAISPROP(z (k) 0:T |x) ] ︸ ︷︷ ︸
EUBOIM-AIS(π0,K, T )
.
(122)

heading: H.2 PROOF OF PROPOSITION 3.1
text: Proposition 3.1 (Complexity in K). For a given T , the bias of the Independent Multi-Sample AIS lower bound on MI, IIM-AISL(π0, T ), decreases (at most) logarithmically in K.
IAISL(π0, T ) ≤ IIM-AISL(π0,K, T ) ≤ IAISL(π0, T ) + logK (17)
To prove the proposition, we first analyze the EUBOIM-AIS(π0,K, T )
log p(x) ≤ EUBOIM-AIS(π0,K, T ) = EpIM-AISTGT (z(1)0:T |x)qAISPROP(z(2:K)0:T |x)
[ log 1
K K∑ i=1 pAISTGT(z (k) 0:T ,x) qAISPROP(z (k) 0:T |x) ] (123)
= E pIM-AISTGT (z (1) 0:T |x)q AIS PROP(z (2:K) 0:T |x)
[ log
K∑ i=1 pAISTGT(z (k) 0:T ,x) qAISPROP(z (k) 0:T |x)
] − logK
≥ E pIM-AISTGT (z (1) 0:T |x)
[ log pAISTGT(z (k) 0:T ,x)
qAISPROP(z (k) 0:T |x)
] − logK (124)
= EUBOAIS(T )− logK (125)
so that EUBOIM-AIS(π0,K, T ) ≥ EUBOAIS(T ) − logK and the inequality bounds the improvement of the upper bound on log p(x). We use the non-negativity of the likelihood ratio for the inequality in the third line. Since IIM-AISL(π0,K, T ) = Ep(x|z)[log p(x|z)] − EUBOIM-AIS(π0,K, T ) and IAISL(π0, T ) = Ep(x|z)[log p(x|z)]− EUBOAIS(π0, T ), we have IIM-AISL(π0,K, T ) ≤ IAISL(π0, T ) + logK as desired. IAISL(π0, T ) ≤ IIM-AISL(π0,K, T ) follows from the fact that multi-sample bounds are tighter than their single sample counterparts.

heading: H.3 PROOF OF PROPOSITION 3.2
text: Proposition 3.2 (Complexity in T ). Assuming perfect transitions and a geometric annealing path with linearly-spaced {βt}Tt=1, the sum of the gaps in the multi-sample AIS sandwich bounds on MI, δT,K = IIM-AISU (π0,K, T )− IIM-AISL(π0,K, T ), reduces linearly with increasing T .
Proof. We now prove the equality, which expresses the gap between the single sample AIS upper and lower bound for linear scheduling as a 1/T factor multiplied by the symmetrized KL divergence. We
first simplify the difference in the single-sample upper and lower bounds as
δT,K=1L + δ T,K=1 U = EUBOAIS(T,K = 1)− ELBOAIS(T,K = 1)
= Ez0:T∼pAISTGT
[ log qAISPROP(z0:T )
pAISTGT(x, z0:T )
] − Ez0:T∼qAISPROP [ log pAISTGT(x, z0:T )
qAISPROP(z0:T )
]
= Ez0:T∼pAISTGT log πT (zT |x) T−1∏ t=0 Tr(zt|zt+1)
π0(z0|x) T−1∏ t=0 Tf (zt+1|zt)
− Ez0:T∼qAISPROP log πT (zT |x) T−1∏ t=0 Tr(zt|zt+1)
π0(z0|x) T−1∏ t=0 Tf (zt+1|zt)  = Ez0:T∼pAISTGT [ log
T∏ t=1 ( π̃T (x, zt) π̃0(zt) )βt−βt−1] − Ez0:T∼qAISPROP [ log T∏ t=1 ( π̃T (x, zt) π̃0(zt) )βt−βt−1]
= Ez0:T∼pAISTGT [ (βt − βt−1) log π̃T (x, zt)
π̃0(zt|x)
] − Ez0:T∼qAISPROP [ (βt − βt−1) log π̃T (x, zt)
π̃0(zt|x) ] (126)
We now assume perfect transitions, which is common in analysis of AIS (Neal, 2001; Grosse et al., 2013). Since each forward transition Tf (zt+1|zt) leaves πt(z) invariant, perfect transitions results in independent exact samples from
z (1:K) 0:T ∼ q AIS PROP(z (1:K) 0:T |x) =⇒ z (1:K) 0:T ∼ π0(z0) T∏ t=1 πβt−1(zt) (127)
so that z1 ∼ π0(z), z2 ∼ πβ1(z), z3 ∼ πβ2(z), etc. Similarly, perfect transitions in the reverse direction Tf (zt|zt+1) results in exact samples from zt ∼ πβt(z)
z (1:K) 0:T ∼ p AIS TGT(z (1:K) 0:T |x) =⇒ z (1:K) 0:T ∼ πT (zT ) T∏ t=1 πβt−1(zt−1) (128)
so that zT−1 ∼ πβT−1(z) and so on down to z0 ∼ π0(z). To put this all together, we substitute back into Eq. 126 using a linear annealing schedule βt−βt−1 = 1 T ∀t. In particular, since Eq. 127 and Eq. 128 factorize, the expectations over the extended state space simplify to a sum of expectations at each zt, with zt ∼ πβt−1(z) for the forward chain and zt ∼ πβt(z) for the reverse chain
δT,K=1L + δ T,K=1 U = EUBOAIS(T,K = 1)− ELBOAIS(T,K = 1)
= Ez0:T∼pAISTGT [ (βt − βt−1) log π̃T (x, zt)
π̃0(zt|x)
] − Ez0:T∼qAISPROP [ (βt − βt−1) log π̃T (x, zt)
π̃0(zt|x) ] =
T∑ t=1 Eπβt (z) [ (βt − βt−1) log π̃T (x, z) π̃0(z|x) ] − T∑ t=1 Eπβt−1 (z) [ (βt − βt−1) log π̃T (x, z) π̃0(z|x) ] = 1
T EπT (z)
[ log π̃T (x, z)
π̃0(z|x) ] − 1 T Eπ0(z) [ log π̃T (x, z) π̃0(z|x) ] = 1
T
( DKL(π0||πT ) +DKL(πT ||π0) ) .
In moving from the third to the fourth line, we see that intermediate terms cancel since each term has the same weight under linear scheduling. In contrast to Grosse et al. (2013), this allows us to derive linear scheduling result which holds for finite T .
The above proves the proposition for the case of single sample AIS. Using additional chains and averaging inside the logarithm, the gap in the multi-sample AIS bounds is less than or equal to that of the single sample bounds by similar arguments which are used to show IWAE bounds are tighter than the ELBO or EUBO (Burda et al., 2016; Sobolev & Vetrov, 2019). Thus, the gap in the Independent Multi-Sample AIS sandwich bounds also decreases linearly in T .

heading: I REVERSE IWAE
text: In this section, we propose Reverse IWAE (RIWAE), which is an impractical alternative to standard IWAE. However, we use this as the basis for our Independent Reverse and Coupled Reverse MultiSample AIS bounds in Sec. 3.3.

heading: I.1 PROBABILISTIC INTERPRETATION AND BOUNDS
text: Similar to reverse importance sampling, we considerK independent posterior samples in the extended state space proposal distribution. The target distribution for our importance sampling scheme is a mixture of K − 1 posterior distributions and one variational distribution,
qRIWAEPROP (z (1:K)|x) = 1
K K∑ k=1 qθ(z (k)|x) K∏ j=1,j 6=k p(z(j)|x) pRIWAETGT (z(1:K)|x) = K∏ k=1 p(z(k)|x) .
(129) We visualize this sampling scheme in Fig. 4.
Similarly to Eq. 30, the log-ratio simplifies to
log pRIWAETGT (z (1:K),x)
qRIWAEPROP (z(1:K),x) = log
K∏ k=1 p(x, z(k))
1
K K∑ k=1 qθ(z(k)|x) K∏ j=1
j 6=k
p(x, z(k))
(130)
= − log ( 1
K K∑ k=1 qθ(z (k)|x) K∏ j=1
j 6=k
p(x, z(k)) · k∏ k=1 p(x, z(k))−1 )
= − log 1 K K∑ k=1 qθ(z (k)|x) p(x, z(k)) (131)
Taking the expected log ratio under the proposal and target distributions yield lower and upper bounds on log p(x). These translate to upper and lower bounds on log p(x) which are different that standard IWAE bounds in Eq. 5,
−Eqθ(z(1)|x)p(z(2:K)|x)
[ log 1
K K∑ k=1 qθ(z (k)|x) p(x, z(k)) ] ≤ log p(x) ≤ −Ep(z(1:K)|x) [ log 1 K K∑ k=1 qθ(z (k)|x) p(x, z(k)) ] Note these bounds are impractical as they would need more than one true posterior sample from p(z|x). However, we will use them in Sec. 3.3 to derive practical multi-sample AIS bounds.

heading: I.2 PROOF OF LOGARITHMIC IMPROVEMENT OF RIWAE
text: Proposition I.1. The bias of the RIWAE lower bound on log p(x) decreases logarithmically in K. In particular, the K-sample lower bound ELBORIWAE(q,K) can improve upon the single-sample ELBO by at most logK nats.
ELBO(q) ≤ ELBORIWAE(q,K) ≤ ELBO(q) + logK (132)

heading: Proof.
text: log p(x) ≥ ELBORIWAE(qθ,K) = −Eqθ(z(1)|x)p(z(2:K)|x)
[ log 1
K K∑ i=1 qθ(z (k)|x) p(x, z(k))
] (133)
= −Eqθ(z(1)|x)p(z(2:K)|x)
[ log
K∑ i=1 qθ(z (k)|x) p(x, z(k))
] + logK
≤ −Eqθ(z(1)|x) [ log qθ(z (1)|x)
p(x, z(1))
] + logK (134)
= ELBO(qθ) + logK (135)
where, in the third line, we have used the fact that the likelihood ratio is non-negative. Thus, the sum of log likelihood ratios is greater than a single term and the negative sign results in a less-than or equal-to inequality.
The logarithmic improvement on the RIWAE lower bound in Prop. I.1 is analogous to the logarithmic improvement on the IWAE upper bound in Prop. B.2, and similarly translates to bounds on MI.
J INDEPENDENT REVERSE MULTI-SAMPLE AIS
In similar fashion to Reverse IWAE, we now use K independent reverse AIS chains to form an extended state space target distribution p(TGTz (1:K) 0:T ,x), and a mixture proposal distribution which contains a single forward AIS chain (see Fig. 2).
qIR-AISPROP (z (1:K) 0:T ,x) =
1
K K∑ s=1 qAISPROP(z (s) 0:T |x) K∏ k=1,k 6=s pAISTGT(z (k) 0:T |x) p IR-AIS TGT (z (1:K) 0:T |x) = K∏ k=1 pAISTGT(z (k) 0:T |x)
As in Eq. 130-131, the unnormalized log ratio becomes
log pIR-AISTGT (z (1:K) 0:T ,x)
qIR-AISPROP (z (1:K) 0:T ,x)
= log
K∏ k=1 pAISTGT(z (k) 0:T |x)
1 K K∑ k=1 qAISPROP(z (k) 0:T ,x) T∏ t=1 pAISTGT(x, z (k) 0:T )
(136)
= − log ( 1
K K∑ k=1 qAISPROP(z (k) 0:T ,x) K∏ j=1
j 6=k
pAISTGT(x, z (k) 0:T ) · k∏ k=1 pAISTGT(x, z (k) 0:T ) −1 )
= − log 1 K K∑ k=1 qAISPROP(z (k) 0:T ,x) pAISTGT(x, z (k) 0:T )
(137)
Taking expectations of the log unnormalized density ratio under the proposal and target, respectively, yield lower and upper bounds on log p(x)
−E z (1) 0:T ∼qAISPROP
z (2:K) 0:T ∼p AIS TGT
[ log 1
K K∑ k=1 qAISPROP(z (k) 0:T |x) pAISTGT(x, z (k) 0:T ) ] ︸ ︷︷ ︸
ELBOIR-AIS(π0,K, T )
≤ log p(x) ≤ −E z (1:K) 0:T ∼pAISTGT
[ log 1
K K∑ k=1 qAISPROP(z (k) 0:T |x) pAISTGT(x, z (k) 0:T ) ] ︸ ︷︷ ︸
EUBOIR-AIS(π0,K, T )
.
(138)
However, reverse Multi-Sample AIS may be impractical in common settings such as posterior inference or energy-based modeling where the normalization constant is unknown, since is infeasible to obtain more than single sample from the target distribution.
Proposition J.1. For the Independent Reverse Multi-Sample AIS log partition lower bound, increasing K leads to (at most) logarithmic reduction of bias compared with
ELBOIR-AIS(K,T )− ELBOAIS(T ) ≤ logK (139)

heading: Proof.
text: log p(x) ≥ ELBOIR-AIS(K,T ) = −EqAISPROP(z(1)0:T |x)pAISTGT(z(2:K)0:T |x)
[ log 1
K K∑ i=1 qAISPROP(z (k) 0:T |x) pAISTGT(x, z (k) 0:T )
] (140)
= −E qAISPROP(z (1) 0:T |x)p AIS TGT(z (2:K) 0:T |x)
[ log
K∑ i=1 qAISPROP(z (k) 0:T |x) pAISTGT(x, z (k) 0:T )
] + logK
≤ −E qAISPROP(z (1) 0:T |x)
[ log qAISPROP(z (k) 0:T |x)
pAISTGT(x, z (k) 0:T )
] + logK (141)
= ELBOAIS(T ) + logK (142)
where, in the third line, we have used the fact that the likelihood ratio is non-negative.

heading: K COUPLED REVERSE MULTI-SAMPLE AIS

heading: K.1 PROBABILISTIC INTERPRETATION AND BOUNDS
text: The Coupled Reverse Multi-Sample AIS extended state space target distribution in Fig. 2 initializes K backward chains from a single target sample zT ∼ πT (z|x), with remaining transitions denoted pAISTGT(z0:T−1|zT ) since they are identical to standard AIS in Eq. 14.
pCR-AISTGT (z (1:K) 0:T−1, zT ,x) = πT (zT ,x) K∏ k=1 pAISTGT(z (k) 0:T−1|zT ,x). (143)
We can construct log-partition bounds using the unnormalized importance ratio
log pCR-AISTGT (z (1:K) 0:T ,x)
qCR-AISPROP (z (1:K) 0:T |x)
= log
p(zT ,x) K∏ k=1 pAISTGT(z (k) 0:T−1|zT )
1 K K∑ j=1 qAISPROP(z (k) 0:T−1, zT |x) K∏ k=1
j 6=k
pAISTGT(z (k) 0:T−1|zT )
(144)
= − log 1 K K∑ k=1
qAISPROP(z (k) 0:T , zT |x)
pAISTGT(z (k) 0:T−1|zT )p(zT ,x)
(145)
Taking the expected log ratio under the proposal and target yields lower and upper bounds on log p(x),
ELBOCR-AIS(π0,K, T ) := −E z (1) 0:T−1,zT∼q AIS PROP(z0:T |x)
z (2:K) 0:T−1∼p AIS TGT(z0:T−1|zT ,x)
[ log 1
K K∑ k=1 qAISPROP(z (k) 0:T−1, zT |x) pAISTGT(x, z (k) 0:T−1, zT ) ] ≤ log p(x)
(146)
EUBOCR-AIS(π0,K, T ) := −E zT∼πT (zT |x) z (1:K) 0:T−1∼p AIS TGT(z0:T−1|zT ,x)
[ log 1
K K∑ k=1 qAISPROP(z (k) 0:T−1, zT |x) pAISTGT(x, z (k) 0:T−1, zT ) ] ≥ log p(x)
(147)

heading: K.2 PROOF OF PROPOSITION 3.3
text: Proposition 3.3. For a given T , the bias of the Coupled Reverse Multi-Sample AIS upper bound on MI, ICR-AISU (π0,K, T ), decreases (at most) logarithmically in K.
IAISU (π0, T )− logK ≤ ICR-AISU (π0,K, T ) ≤ IAISU (π0, T ). (20)

heading: Proof.
text: log p(x) ≥ELBOCR-AIS(K,T ) = −EqAISPROP(z(1)0:T−1,zT |x)pAISTGT(z(2:K)0:T−1,zT |x)
[ log 1
K K∑ i=1 qAISPROP(z (k) 0:T−1, zT |x) pAISTGT(z (k) 0:T−1|zT )p(zT ,x)
]
= −E qAISPROP(z (1) 0:T−1,zT |x)p AIS TGT(z (2:K) 0:T−1,zT |x)
[ log
K∑ i=1 qAISPROP(z (k) 0:T−1, zT |x) pAISTGT(z (k) 0:T−1|zT )p(zT ,x)
] + logK
≤ −E qAISPROP(z (1) 0:T−1,zT |x)
[ log qAISPROP(z (k) 0:T−1, zT |x)
pAISTGT(z (k) 0:T−1|zT )p(zT ,x)
] + logK
= ELBOAIS(T ) + logK (148)
where, in the third line, we have used the fact that the likelihood ratio is non-negative. Since ICR-AISU (π0,K, T ) = Ep(x|z)[log p(x|z)] − ELBOCR-AIS(π0,K, T ) and IAISL(π0, T ) = Ep(x|z)[log p(x|z)] − ELBOAIS(π0, T ), then Eq. 148 implies ICR-AISU (π0,K, T ) ≥ IAISU (π0, T ) − logK as desired.

heading: L EXPERIMENTAL DETAILS

heading: L.1 DATASETS AND MODELS
text: We used MNIST (LeCun et al., 1998) and CIFAR-10 (Krizhevsky & Hinton, 2009) datasets in our experiments.
Real-Valued MNIST. For the VAE experiments on the real-valued MNIST dataset (Table 1), the encoder’s architecture is 784 − 1024 − 1024 − 1024 − z where z is the latent code size shown in the row header of Table 1. The decoder architecture is the reverse of the encoder architecture. The decoder variance is a global parameter learned during the training.
For the GAN experiments on MNIST (Table 1), we used the same decoder architecture as our VAEs. In order to stabilize the training dynamic, we used the gradient penalty (GP) (Salimans et al., 2016).
The network was trained for 300 epochs with the learning rate of 0.0001 using the Adam optimizer (Kingma & Ba, 2014), and the checkpoint with the best validation loss was used for the evaluation.
CIFAR-10. For the CIFAR-10 experiments (Table 1), we experimented with the DCGAN (Radford et al., 2015) with Gradient Penalty (GP-GAN) (Gulrajani et al., 2017). The number at the end of each model name in Table 1 indicates the latent code size.

heading: L.2 MULTI-SAMPLE AIS SETTINGS
text: We conducted more extensive experimentation for upper bounding the mutual information (lower bounding the partition function log p(x)) and lower bounding the Mutual Information (upper bounding the partition function log p(x)) using coupled reverse multi-sample AIS (See Fig. 4 for the definitions, results are shown in Table 1) For all these experiments, we used 20 leap frog steps for HMC. We used 48 independent chains for each data point. The step size of HMC was adaptive tuned to achieve an average acceptance probability of 65% as detailed in Neal (2001).
For all MNIST experiments in Table 1, we evaluated on a single batch size of 128 simulated data. We used a sigmoid temperature schedule as used in Wu et al. (2016) but with radius parameter set to 10 instead of 4. For all CIFAR experiments in Table 1 we used a single batch of 32 simulated data. It is much harder to get a tight upper and lower bound for CIFAR models, and we discovered that the tightness of the bound is sensitive to the radius parameter, so we used the optimal radius for each setting. For GAN100 models on CIFAR10, the radius parameter for the temperature is 6 when the initial distribution is q(z|x) and 8 when it is log p(z).

heading: L.3 RUNTIME COMPARISON
text: For MNIST, it took about 20 mins to run IWAE for K = 1M , and it took about 8 minutes to run AIS for T = 500 and about 8 hours for T = 30k. For CIFAR, it took about 70 mins to run AIS for T = 500 steps and about nine days for T = 100k.
In Fig. 8, we evaluate the tradeoff between runtime and bound tightness for evaluating the generative MI for VAE and GAN models with 100-dimensional latent codes trained on the MNIST dataset. We compare IWAE and multi-sample AIS evaluation, with the same experimental settings as in App. L.2 and an initial distribution qθ(z|x). We plot wall clock time on the x-axis, where increasing runtime reflects increasing K for IWAE and increasing T for AIS.
As predicted by Prop. 2.1, the improvement of the IWAE lower bound on MI is logarithmic with increasing K, which appears linear with the logarithmic scale of the x-axis. We can also confirm the result in Prop. 3.2 that the gap in the multi-sample AIS lower and upper bounds reduces linearly with increasing T , as each bound appears to improve exponentially in the logarithmic scale of the x-axis.
L.4 INITIAL DISTRIBUTION OF AIS
We can use any distribution that is tractable to sample and evaluate as the initial distribution of AIS. In practice, it is important to use a distribution that is mode-covering for the posterior distribution so
that the AIS particles do not get stuck in a few modes. A typical option for the initial distribution is the prior p(z). In this case, the intermediate distributions will be qβ(z|x) ∝ p(z)p(x|z)β . Another possibility is a qθ(z|x) distribution which has been trained to maximize the importance sampling or BA lower bound on mutual information:
I(x; z) ≥ I(x; z)−DKL[p(z|x)||q(z|x)] = Ep(x,z) [ log
q(z|x) p(z) ] As this bound optimizes a reverse KL divergence, the learned qθ(z|x) will be a mode covering distribution for p(z|x) and thus a good candidate from which to start the AIS chains. In this case, the intermediate distributions for single-sample AIS chains will be qβ(z|x) ∝ q(z|x)(1−β)p(x, z)β .

heading: L.5 ANALYTICAL SOLUTION OF THE MUTUAL INFORMATION ON THE LINEAR VAE
text: In order to verify our implementations, we have derived the mutual information analytically for Linear VAEs trained on MNIST.
For simplicity, we assume a fixed identity covariance matrix I at the output of the conditional likelihood of the linear VAE decoder, i.e., the decoder of the VAE is simply: x = Wz+b+ , where x is the observation, z is the latent code vector z ∼ N (0, I), W is the decoder weight matrix and b is the bias. The observation noise of the decoder is ∼ N (0, I). It’s easy to show that the conditional likelihood is: p(x|z) = N (x|Wz+b, I) and thus we can solve for the marginal:
p(x) = N (x|µx = b,Σx = I + WWᵀ) (149)
Thus the differential entropy of x is:
h(x) = k 2 + k 2 log(2π) + 1 2 log(detΣx) (150)
Where k is the dimension of the observation. The conditional entropy:
h(x|z) = k 2 + k 2 log(2π) + 1 2 log(det I) (151)
= k 2 + k 2 log(2π) (152)
Thus the mutual information:
I(x; z) = h(x)− h(x|z) (153)
= 1
2 log(detΣx) (154)

heading: M ADDITIONAL EXPERIMENTS

heading: M.1 COMPARISON OF MULTI-SAMPLE AIS BOUNDS
text: Fig. 9 compares performance of our various Multi-Sample AIS bounds of Linear VAE, VAE, and GAN models with 20 latent variables trained on MNIST. We perform this comparison with the goal of making recommendations for which Multi-Sample AIS bounds to use in practice.
First, the MI upper bound using the Independent Multi-Sample AIS lower bound on log p(x), which corresponds to the forward direction of BDMC, outperforms its counterparts. This bound uses independent samples and is not limited to only logK improvement, compared with the Coupled Reverse Multi-Sample AIS MI upper bound (Prop. 3.3). We recommend using the Independent Multi-Sample AIS MI upper bound (LHS of Eq. 16) in all cases.
The results are less conclusive for the multi-sample AIS lower bounds on MI, where either the Independent Multi-Sample AIS EUBO or Coupled Reverse AIS EUBO may be preferable for log partition function estimation. Recall that these bounds have different sources of stochasticity that provide improvement over single-chain AIS. The stochasticity in the Independent Multi-Sample AIS EUBO comes from K − 1 independent ‘negative’ forward chains which, by Prop. 3.1, can only lead to logK improvement over single-sample AIS. However, these gains are easily attained for low values of T . For example, with zero intermediate distributions (T = 1), and π0(z) = p(z), IIM-AISL(p(z),K, T = 1) reduces to Structured INFO-NCE. This may be useful for quickly estimating values of MI at a similar order of magnitude as logK.
The stochasticity in the Coupled Reverse AIS EUBO is induced by MCMC transitions in K−1 coupled backward chains. While this does not formally limit the improvement over single-sample AIS, we see in Fig. 9 that at least moderate values of T may be needed to match or marginally improve upon Independent Multi-Sample AIS. These observations suggest that the preferred lower MI bounds may vary based on the scale of the true MI and the amount of computation available.

heading: M.2 CONFIDENCE INTERVALS FOR MULTI-SAMPLE AIS EXPERIMENTS
text: In Table 4 and Table 5, we have conducted more extensive experimentation for upper bounding the mutual information (lower bounding the partition function log p(x)) and lower bounding the Mutual Information (upper bounding the partition function log p(x)) using Coupled Reverse Multi-Sample AIS. We have also computed the 95% confidence intervals for those experiments. For MNIST experiments, we computed the confidence interval over 8 batches of 16 data points. For CIFAR experiments, due to computational constraints, we computed the confidence intervals over 8 batches of 4 data points. Results for MNIST are shown in Table 4 and results for CIFAR are shown in Table 5.

Review summaries:
This work combines the technique of Annealed Importance Sampling (AIS) with Importance Weighted Autoencoder (IWAE) bounds on mutual information to obtain new and significantly tighter bounds. They examine the situations where the joint distribution $p(x, z)$ is known via $p(z)$ and $p(x | z)$, and the situation where only $p(z)$ is known. In the former case, they obtain exact bounds (or asymptotically exact given large enough sample sizes), and in the latter case they find approximate bounds by learning an energy function to approximate $p(x |z)$. Their experiments show the utility of their proposed method. In particular, they obtain impressively tight lower and upper bounds for complex latent variable generative models (VAE and GAN) on complex data such as CIFAR-10.

In this work, the authors proposed some novel estimators for mutual information, namely the GIWAE, Muti-Sample AIS and MINE-AIS bounds. The key idea is to apply importance sampling to mutual information in the model-based setting. In both theoretical and empirical analysis, they prove the bounds are much tighter than the existing ones in the literature. 

This paper investigates annealed importance sampling (AIS) techniques and multi-sample bounds to improve mutual information estimation. Assuming the knowledge of extract likelihood (joint or marginal), this paper shows the mutual information bound can be tightened via exploiting a series of intermediate distributions interpolating between a (convenient) proposal and the joint density. An energy-based perspective was presented and discussed, revealing interesting insights. The effectiveness of this proposal is demonstrated with generative modeling applications on real-world datasets.