I MPROVING M UTUAL I NFORMATION E STIMATION WITH A NNEALED AND E NERGY -B ASED B OUNDS
Anonymous authors
Paper under double-blind review
A BSTRACT
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 VAE s and GAN s trained on the MNIST and CIFAR datasets, and showcase signiﬁcant gains over existing bounds in these challenging settings with high ground truth MI .
I NTRODUCTION
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 I NFO NCE (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 inﬁnite-sample limit of the GIWAE framework.
We propose an efﬁcient training procedure using Markov Chain Monte Carlo ( MCMC ) and utilize our Multi-Sample AIS evaluation bounds to show notable improvement over ﬁnite-sample methods for lower bounding MI . In Fig. 1, we summarize the MI bounds discussed in this paper and the relationships between them.
P ROBLEM S ETTING
The mutual information between two random variables x and z with joint distribution p ( x , z ) is
I
[Equation]
where H ( x | z ) denotes the conditional entropy − E p ( 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 difﬁculty 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.
M ULTI -S AMPLE B OUNDS ON M UTUAL I NFORMATION
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.
2. S IMPLE I MPORTANCE S AMPLING B OUNDS
For an arbitrary conditional distribution q θ ( z | x ) , we can obtain lower and upper bounds on log p ( x ) :
[Equation]
where the left-hand side is the familiar Evidence Lower Bound ( ELBO ) with a gap of
D KL [ 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 D KL [ 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
[Equation]
The left hand side of Eq.
3 is the well-known Barber-Agakov ( BA ) bound
(Barber & Agakov, 2003;
Agakov, 2006), which we denote by I BA L ( q θ ) =
I ( x ; z ) − D KL [ p ( z | x ) ||
q θ ( z | x )] .
We refer to the right hand side as the BA upper bound I BA U ( q θ )
=
I ( x ; z ) + D KL [ q θ ( z | x ) || p ( z | x )] .
Relationship with CLUB When the prior is used as the proposal, I BA U ( p ( z )) corresponds to the
Contrastive Log-Ratio Upper Bound ( CLUB )
(Cheng et al., 2020) with known p ( x | z ) .
2. I MPORTANCE W EIGHTED A UTOENCODER
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 q IWAEPROP with K independent samples from the variational distribution q θ ( z | x ) .
The extended state space target p IWAETGT 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 )
[Equation]
The log ratio log p IWAETGT ( x , z
[Equation] [Equation]
K ) )
q IWAEPROP
[Equation] [Equation]
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
[Equation]
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 are D KL [ q IWAEPROP || p IWAETGT ] and D KL [ p IWAETGT || q IWAEPROP ] , 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 I IWAE L ( q θ , K ) and I IWAE U ( q θ , K ) . For example, I IWAE L ( q θ , K ) can be written as
[Equation] 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 I IWAE L ( 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:
[Equation]
2. G ENERALIZED IWAE
In this section, we consider a family of Generalized IWAE ( GIWAE ) lower bounds, which interpolate between I BA L ( q θ ) and I IWAE L ( 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 ) = 1 K ∀ s that speciﬁes which sample z ( s ) ∼ p ( z | x ) is drawn from the posterior
[Equation] Note that marginalization over s leads to the IWAE target distribution p IWAE TGT ( z (1: K ) | x ) . For the GIWAE extended state space proposal, consider a categorical index variable q GIWAEPROP ( s | z (1: K ) , x ) drawn using self-normalized importance sampling ( SNIS ) with weights calculated by a learned energy function T φ q GIWAE PROP ( z (1: K ) , s | x ) =
[Equation]
1: K ) , x ) where q GIWAE PROP ( s |
z (1: K ) , x ) = e T φ ( x , z ( s ) ) K (cid:80) k
=1 e T φ ( x , z ( k ) ) . We provide detailed discussion of the GIWAE probabilistic interpretation in App. C.1. Taking the expectation under p GIWAETGT z (1: K ) , s | x ) of the log importance ratio log p GIWAETGT ( z (1: K ) ,s | x ) q GIWAEPROP ( z (1: K ) ,s | x ) as in App.
A, we obtain an upper bound on log p ( x ) and lower bound on MI as
I GIWAE
[Equation] [Equation]
We see that the GIWAE lower bound decomposes into the sum of two terms, where the ﬁrst 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 I GIWAE L ( q, T φ = c, K ) =
I BA L ( q θ ) for all K .
Similarly, the single sample
I GIWAE L ( 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 I NFO NCE contrastive lower bound (
van den Oord et al., 2018;
Poole et al., 2019), with I
I NFO NCE L ( T φ , K ) =
I GIWAE L ( 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 log K limitations of I NFO NCE (van den Oord et al., 2018). However, using the prior as the proposal in I NFO NCE does allow the energy function to admit a bi-linear implementation T ( x , z ) =
T x ( x ) T z ( 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.
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 ) ,
I GIWAE L ( q θ , T ∗ )
= I IWAE L ( q θ ) .
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
[Equation] [Equation]
While the GIWAE lower bound does not assume access to the full joint distribution p ( x , z ) , Prop. 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 I NFO NCE
Since I NFO NCE and Structured I NFO NCE (SI NFO NCE, 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
[Equation]
M ULTI -S AMPLE AIS B OUNDS FOR E STIMATING M UTUAL I NFORMATION
In the previous sections, we derived probabilistic interpretations of extended state space importance sampling using multiple samples K , as in IWAE .
In this section, we ﬁrst 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 . [Figure]
3. A NNEALED I MPORTANCE S AMPLING B ACKGROUND
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 Z T ( x ) = p ( x ) .
A common choice for intermediate distributions is the geometric mixture path parameterized by { β t } Tt =0 :
[Equation]
where Z t (
[Equation]
In the probabilistic interpretation of AIS , we consider an extended state space proposal q AISPROP ( z 0: T | x ) , obtained by sampling from the initial π 0 ( z | x ) and constructing transitions
T f
( z t +1 | z t , x ) which leave π t ( z | x ) invariant.
The target distribution p AISTGT ( z 0: T | x ) is given by running the reverse transitions
T r ( z t | z t +1 , x ) starting from a target or posterior sample π T ( z | x ) , as shown in Fig.
[Equation]
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 , I AIS L ( π 0 , T ) and I AIS U ( π 0 , T ) , which were suggested for MI estimation in the blog post of Sobolev (2019a).
3. I NDEPENDENT M ULTI -S AMPLE AIS B OUNDS
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 ∼ q AISPROP 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 ∼ p AIS TGT starting from a true posterior sample
z T ∼ p ( z | x ) .
The remaining K − 1 samples are obtained by running forward AIS chains, as visualized in Fig.
[Equation]
where q AISPROP and p AISTGT were deﬁned in Eq.
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)
[Equation]
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 , I IM - AIS L ( π 0 , T ) , decreases (at most) logarithmically in K .
[Equation]
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 = I IM - AIS U
( π 0 , K, T ) − I IM - AIS L ( π 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.
3. C OUPLED R EVERSE M ULTI -S AMPLE AIS B OUNDS
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 deﬁne Independent Reverse Multi-Sample AIS ( IR - AIS ) using the following proposal and target distribution, as shown in Fig. 2.
[Equation]
Note that partition function ratio is Z TGT / Z PROP
= p ( x ) K /p ( x ) K [Equation] = 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 z T ∼ π T ( z | x ) , with remaining transitions denoted p AISTGT ( z 0: T − 1 | z T ) since they are identical to standard AIS in Eq.
[Equation]
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.
[Equation]
Taking the expected log ratio under the proposal and target yields lower and upper bounds on log p ( x ) ,
[Equation]
[Table]
Proposition 3.3.
For a given T , the bias of the Coupled Reverse Multi-Sample AIS upper bound on
MI , I CR - AIS U ( π 0 , K, T ) , decreases (at most) logarithmically in K .
[Equation]
3. D ISCUSSION
Relationship with BDMC While Bidirectional Monte Carlo ( BDMC ) (Grosse et al., 2015; 2016) was the ﬁrst method proposing multi-sample log partition function bounds using AIS chains, our probabilistic interpretations provide novel perspective on BDMC .
Perhaps surprisingly, we ﬁnd 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 log K 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 EUBO IM - AIS ( π 0 , K, T ) ( RHS of Eq. 16) can only improve upon single-sample AIS by log K (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 log K .
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.
MINE-AIS E STIMATION OF M UTUAL I NFORMATION
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 ﬂexible energy-based distribution π θ,φ ( z | x ) as an approximation to the posterior p ( z | x ) (Poole et al., 2019; Arbel et al., 2020)
[Equation]
where Z
[Equation]
[Equation] Plugging π θ,φ ( z | x ) into the BA lower bound, we denote the resulting bound as the Implicit Barber-
Agakov Lower Bound ( IBAL ), since it is often difﬁcult to evaluate explicitly due to the intractable log partition function term.
After simplifying in App.
D.1, we obtain
I (
[Equation] [Equation]
with the gap of the IBAL equal to E p ( x ) [ D KL [ p ( z | x ))
||
π θ,φ ( z | x )]] .
Note that the IBAL generalizes the
Unnormalized Barber-Agakov bound from Poole et al.
(2019), with I UBA ( 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 →∞ I GIWAE L ( q θ , T φ , K ) =
IBAL ( q θ , T φ ) .
Thus, we may view the IBAL as the limiting behavior of the ﬁnite-sample GIWAE bounds as K → ∞ .
While Cor.
2.3 shows that I GIWAE L ( q θ , T ) can improve I BA L ( q θ ) by at most log K nats, we show in
App. E.1 that the IBAL contrastive term is ﬂexible enough to close the entire gap in the BA bound.
[Equation]
However, this ﬂexible contrastive term comes at the cost of tractability, as MINE - AIS involves an intractable partition function log E q θ ( z | x ) [ e T φ ( x , z ) ] compared to the ﬁnite-sum term in GIWAE .
Training of MINE-AIS Although the log partition function log Z θ,φ ( x ) in the IBAL is intractable to evaluate, we only require an unbiased estimator of its gradient for training.
Differentiating Eq.with respect to the parameters θ and φ of the variational and energy function, respectively, we obtain
[Equation]
[Table]
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 ﬁnite-sample GIWAE bounds translates to energy-based training in the limit of K → ∞ .
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 E p ( x , z )
[ T φ ( x , z )] , and will use Multi-Sample AIS to estimate log Z θ,φ ( x ) by annealing between π 0 ( z | x ) =
q θ ( z | x ) and π T ( x , z ) =
q θ ( z | x ) e T φ ( x , z ) .
To preserve a lower bound on I ( x ; z ) , we would like an upper bound on log Z θ,φ ( 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 log Z θ,φ ( x ) by initializing backward chains from the true posterior p ( z | x ) instead of π θ,φ ( z | x ) .
Using a Multi-Sample AIS lower bound on log Z θ,φ ( x ) , we can upper bound the IBAL but cannot guarantee that this is a stochastic lower bound on I ( x ; z ) .
E XPERIMENTS
In this section, we evaluate our proposed Multi-Sample AIS , MINE - AIS , and GIWAE bounds for estimating MI in deep generative models such as VAE s and GAN s trained on MNIST and CIFAR .
5. M ULTI -S AMPLE AIS E STIMATION OF M UTUAL I NFORMATION
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 =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
I NFO NCE.
We see that the improvement of the IWAE lower bound is limited by log K , 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 VAE 2, VAE 10 and GAN 2 on MNIST in cases where the gap is less than log K .
However, using Gaussian q θ ( z | x ) , tight estimation with IWAE is clearly infeasible on VAE 100 and GAN 100.
On VAE 100, for example, at least e 46 total samples are required to improve
I BA L ( q θ ) =
I IWAE L ( 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.
[Figure]
5. E NERGY -B ASED E STIMATION OF M UTUAL I NFORMATION
In this section
[Figure]
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 I GIWAE L ( 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 ﬁnd that the results respect the inequalities shown in
Cor.
C.2, with I BA L ( q θ ) ≤ I GIWAE L ( q θ , T φ ∗ , K ) ≤
I IWAE L ( q θ , K ) ≤ I BA L ( q θ )
+
log K .
For the Linear VAE , the true posterior is in the Gaussian variational family and I BA L ( q θ ) almost exactly matches the analytical MI .
In this case, the contrastive term provides much less than log K 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 log K improvement, and yields a higher overall bound.
This observation is similar to Rainforth et al. (2018), where increasing K in ELBO IWAE ( 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 log K improvement over BA .
MINE-AIS Bounds We can see in Fig.
3a that MINE - AIS improves over BA due to its more ﬂexible variational family, and improves upon GIWAE , IWAE , and (Structured)
I NFO NCE since it is their limiting behavior.
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.
C ONCLUSION
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 log K 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 inﬁnite-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.
Table
of Contents
B IWAE
B.1 Probabilistic Interpretation and Bounds . . . . . . . . . . . . . . . . . . . . . .
B.2 Proof of Proposition 2.1 . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
B.3 Logarithmic Improvement Results for I IWAE L ( q θ , K ) . . . . . . . . . . . . . . .
B.4 Bias Reduction of IWAE Lower Bound on log p ( x ) / Upper Bound on MI . . . .
B.5 Relationship with Structured InfoNCE . . . . . . . . . . . . . . . . . . . . . .
C Generalized IWAE
C.1 Probabilistic Interpretation and C. Does Not Provide Beneﬁt
IWAE . . C.3 Relationship between IWAE and GIWAE Probabilistic Interpretations . . . . . .
C.4 Proof of Proposition 2.2 and Corollary 2.3 . . . . . . . . . . . . . . . . . . . .
D Generalized Mutual Information Neural Estimation ( GMINE )
D.1 Probabilistic Interpretation of IBAL . . . . . . . . . . . . . . . . . . . . . . . .
D.2 Probabilistic Interpretation of Generalized MINE - DV . . . . . . . . . . . . . . . D.3 Probabilistic Interpretation of Generalized MINE - F . . . . . . . . . . . . . . . . E MINE-AIS E.1 Properties of MINE-AIS
[Table] Proof of Proposition 4.1 . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . [Table] [Table]
[Table] [Table] F
[Table]
Known Marginals F.
[Table] BA [Table] [Table]
[Table] [Table] Bound . . . . . . . . . . . . . . . . . . . . . . . . . [Table]
[Table]
Independent
[Table]
AIS
[Table]
[Table]
Probabilistic Interpretation and Bounds . . . . . . . . . . . . . . . . . . . . . . H.
of Proposition 3.1 . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
H.3 Proof of Proposition 3.2 . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
I Reverse IWAE
Bounds . . . . . . . . . . . . . . . . . . . . . .
I. Proof
of
RIWAE . .
L Experimental Details
M.
Intervals
Experiments . . . . . . . . . . . . .
A A G ENERAL A PPROACH FOR D ERIVING E XTENDED S TATE S PACE B OUNDS ON L OG P ARTITION F UNCTIONS
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 q PROP and p TGT chosen such that Z TGT Z PROP = 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
[Equation]
Note that we use ˜ q, ˜ p to indicate unnormalized densities over x and a possibly extended (latent) state space.
For example, p TGT ( z (1: K ) 0: T | x ) = 1 Z TGT ˜
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 q PROP and p TGT , the gap in the lower bound in Eq.
26 is the forward KL divergence,
D KL ( q PROP || p TGT ) , while the gap in the upper bound is the reverse KL divergence, D KL ( p TGT )
|| q PROP ) .
[Equation]
Thus, the bounds in Eq.
26 directly generalize the standard ELBO ( log p ( x ) − D KL [ q ( z | x ) || p
( z | x )] ) and EUBO ( log p ( x ) + D KL [ 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 D KL [ q PROP || p TGT ] in the extended state space to obtain log Z TGT / Z PROP − ELBO ( q PROP , p TGT ) as in Eq.
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 B. P ROBABILISTIC I NTERPRETATION AND B OUNDS
Consider a K -sample proposal distribution q IWAEPROP ( z (1: K ) ) using independent draws from an initial distribution q θ ( z | x ) .
The target distribution is deﬁned 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 ) . [Equation]
[Equation]
Note that the normalizing constant of p IWAETGT ( x , z (1: K ) ) , or (cid:82) p IWAETGT ( x , z
(1: K ) )
d z (1: K ) , equals p ( x ) since q θ ( z | x ) is normalized and (cid:82) p ( x , z )
d z = p ( x ) .
To connect this with the IWAE bound, we show that the log importance ratio reduces to
[Equation]
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 q IWAEPROP , we obtain a lower bound on log p ( x ) as show in App.
A
[Equation]
This corresponds to the following upper bound on MI for known p ( x | z )
[Equation]
Upper Bound on log p ( x ) and Lower Bound on MI Similarly, with expectations under p IWAETGT ( z (1: K ) | x ) , we obtain an upper bound on log p ( x ) .
Since, for independent draws, the uniform mixture p IWAETGT is invariant to index permutations, we may choose the target sample to be z (1) and obtain the upper bound of Sobolev & Vetrov (2019):
[Equation]
Translating this to a lower bound on MI with known p ( x | z ) ,
[Equation]
B. P ROOF OF P ROPOSITION 2.
We ﬁrst recall results that IWAE bounds tighten with increasing K .
Proposition B.1.
For given q θ ( z | x ) , ELBO IWAE ( q θ , K + 1) ≥ ELBO IWAE ( q θ , K ) (Burda et al., 2016) and EUBO IWAE ( q θ , K + 1) ≤ EUBO IWAE ( 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 EUBO IWAE ( q θ , K ) can improve upon the single-sample EUBO :
= log p ( x ) + D KL [ p ( z | x ) ||
q θ ( z | x )] by at most log K nats. [Equation]
[Equation]
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.
[Equation]
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 ﬁrst term is the EUBO and is still an upper bound on log p ( x ) .
However, note that EUBO − log K is no longer guaranteed to be an upper bound on log p ( x ) .
Prop. B.1 implies that EUBO IWAE ( 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:
[Equation]
Proof.
Recall from Sec.
1.1 and Sec. 2.1 that I BA L ( q θ ) )
= − EUBO ( q θ ) − H ( x | z ) , and
I IWAE L ( q θ , K ) =
− EUBO IWAE ( q θ , K ) − H ( x | z ) .
Using Prop.
B.2, EUBO ( q θ , K ) − log K ≤
EUBO IWAE ( q θ , K ) implies that
[Equation]
so that I IWAE L ( q θ , K ) ≤ I BA L ( q θ ) )
+ log K as desired.
I BA L ( q θ ) ≤ I IWAE L ( q θ , K ) and I IWAE U ( q θ , K ) ≤ I BA U ( q θ ) follow from the fact that IWAE bounds tighten with increasing K in Prop.
B.1.
B. L OGARITHMIC I MPROVEMENT R ESULTS FOR I IWAE L ( 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
I
[Equation]
This way of writing I IWAE L ( q θ , K ) provides intuition for the result in Prop. 2.1, since the improvement over I BA L ( q θ ) is simply the contrastive term, which is limited to log K .
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 log K .
For low dimensional VAEs and
GANs, the contrastive term contributes less than log K , 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 log K .
This sheds light on the inherent exponential sample complexity of IWAE , and motives our Multi-Sample AIS methods for MI estimation. [Table] [Table]
B. B IAS R EDUCTION OF IWAE L OWER B OUND ON log p ( x ) / U PPER B OUND 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 { D KL [ p ( z | x ) ||
q θ ( z | x )] } is required for accurate importance sampling estimation.
Maddison et al.
Domke & Sheldon (2018)
ﬁnd that the bias of the IWAE lower bound in the limit of K → ∞ reduces at the rate of O ( 12 K Var [ p ( z | x ) q θ ( z | x ) ]) , although the Var [ p ( x , z ) q θ ( z | x ) ] term is at least exponential in D KL [ p ( z | x ) ||
q θ ( z | x )]
(Song & Ermon, 2019).
B. R ELATIONSHIP WITH S TRUCTURED I NFO NCE
We can recognize the Structured I NFO NCE upper and lower bounds for known p ( x | z ) (Poole et al.
Sec. 2.5) as simply applying the standard IWAE bounds, using the marginal p ( z ) in place of the variational q θ ( z | x )
[Equation]
1 K K (cid:80) k
=1 p ( x | z ( k ) )
(cid:21) ≤ I ( x ; z ) ≤ E p ( x ) p ( z (1: K ) )
(cid:20) log 1 K K (cid:80) k
=1 p ( x | z ( k ) )
(cid:21) − H
x | z ) .
We refer to the lower bound as I S-I NFO NCE L
( K ) and the upper bound as I S-I
NFO NCE U ( K ) .
From
Prop. 2.1, we obtain an alternative proof that the structured I NFO NCE lower bound is upper bounded by log K .
Since I BA ( p ( z )) ≤
I IWAE L ( p ( z ) , K ) ≤ I BA ( p ( z ))
+ log K and the BA bound with a prior proposal equals 0 from Eq.
3, we have that 0 ≤ I S-I NFO NCE ≤ log K .
C G ENERALIZED IWAE C. P ROBABILISTIC I NTERPRETATION AND B OUNDS
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 ) =
1 K ∀ s that speciﬁes 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 ) p GIWAETGT (
[Equation]
, (41) with marginalization over s leading to p IWAETGT ( 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
[Equation] [Equation]
For the GIWAE extended state space proposal distribution, we consider a categorical index variable q GIWAEPROP ( s | z (1: K ) , x ) drawn according SNIS , with weights calculated using the energy function T φ . [Equation]
q
[Equation]
s | z
[Equation]
1: K ) , x ) where q GIWAEPROP ( s | z (1: K ) , x ) =
[Equation]
x
[Equation]
k
[Equation] [Equation]
Similarly to Lawson et al.
and as shown in App.
C.3, the induced SNIS distribution performs variational inference of the posterior p GIWAETGT ( 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 ﬁrst calculate the unnormalized density ratio
[Equation]
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 )
[Equation]
To derive an explicit form for the upper bound on log p ( x )
[Equation]
EUBO GIWAE [Equation] φ [Equation]
 
[Equation]
Translating this to the lower bound on MI from Eq.
9 corresponds to taking I ( x ; z ) =
− E p ( x ) [log p ( x )] − H ( x | z ) ≥ − EUBO GIWAE ( q θ , T φ , K ) − H ( x | z ) .
Writing the conditional entropy term over the index s ,
[Equation]
C.
[Table]
Taking expectations under q GIWAEPROP 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.
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 ) ∼ p GIWAETGT .
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.
C. R ELATIONSHIP BETWEEN IWAE AND GIWAE P ROBABILISTIC I NTERPRETATIONS
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 p GIWAETGT ( s |
z (1: K ) , x ) and the variational distribution q ( s | z (1: K ) , x ) (Eq. 43)
[Equation]
Proof.
To show this result, we consider the posterior over the index variable p ( s |
z
(1: K ) , x ) p GIWAETGT ( s | z (1: K ) , x ) = p GIWAETGT ( z (1: K ) , s, x ) K (cid:80)
[Equation]
z ( s ) | x ) q
( z ( s ) | x )
K
[Equation]
s = p ( z ( s )
| x ) q ( z ( s ) | x )
Now, note that p ( z (1: K ) , s, x ) = (cid:0) K (cid:80) s =1 p
( z (1: K ) , s, x ) (cid:1) · p ( s | z (1: K ) , x ) , so that we can decompose the
KL divergence as
[Equation]
Now, the left hand side equals the gap in the GIWAE upper bound on log p ( x ) in Eq.
(47), while the ﬁrst term on the RHS is the gap in the IWAE upper bound on log p
C. P ROOF OF P ROPOSITION 2.2 AND C OROLLARY 2.
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 ) ,
[Equation]
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. imply
[Equation]
Corollary C.3.
Using the prior q θ ( z | x )
= p ( z ) as in I NFO NCE ,
(a) For the optimal energy function, Cor.
C.2(a) implies T ∗ ( x , z )
= log p ( x | z )
+ c ( x ) , and
[Equation]
[Equation] (b) For an arbitrary energy function T φ ( x , z ) , Cor.
C.2(b) translates to
[Equation]
Proof.
(a) We ﬁrst 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
[Equation]
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. [Equation] ( z ( k ) | x ) (cid:1) · e c ( x )  
[Equation] =
E p ( x [Equation] (cid:20) [Equation]
[Equation]
1) | x ) p (
[Equation] [Equation]
log p (
[Equation]
[Equation]
x
[Equation]
q (
[Equation]
(1) |
[Equation]
) (cid:21)
[Equation]
E p ( x ) p ( z (1) | x )
q θ ( z (2: K )
| x ) (cid:34) log 1 K
K (cid:88) i = p ( z ( k ) , x ) q θ ( z ( k ) | x ) (cid:35) =
−H ( x | z ) + (cid:32)
[Equation]
p ( x ) p (
[Equation]
(1) | x )
[Equation]
θ ( z (2: K ) | x ) (cid:34) log 1 K
K (cid:88) i = p (
[Equation]
( k )
[Equation]
x ) q θ (
[Equation]
( k ) | x ) (cid:35)(cid:33) = I IWAE L ( 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 ) = − E p ( x )
[log p ( x )] .
Combined with the entropy −H ( x | z ) , we obtain the I IWAE lower bound on MI .
See Section 2.2.
(b) We begin by showing that I BA L ( q θ ) ≤ I GIWAE L ( 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 I BA L ( q θ ) =
I GIWAE L ( 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:
I GIWAE L ( q θ , T φ ∗ , K ) ≤ I IWAE L ( q θ , K ) , as the T
φ ∗ is the optimal energy function.
This relationship is visualized in Fig.
3b.
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.and App.
E. See Fig.
5 for a summary of Generalized MINE bounds and their relationships. D. P ROBABILISTIC I NTERPRETATION OF IBAL
For the variational approximation
[Equation]
where
[Equation]
58) we consider the BA lower bound on MI , I BA
[Equation]
contrastive term := IBAL
[Equation]
φ ) so that the gap in the mutual information bound is E p ( x ) (cid:2) D KL [ p ( z | x ) ||
π θ,φ ( z | x )] (cid:3) .
Note that
IBAL ( q θ , T φ ) contains I BA L ( q θ ) as one of its terms.
We visualize the relationship between various energy-based bounds in Fig. 3b and Fig. 5.
D. P ROBABILISTIC I NTERPRETATION OF G ENERALIZED MINE - DV
We can interpret the MINE - DV bound of Belghazi et al.
as arising from a joint energy-based variational approximation π θ,φ ( x , z ) of the full joint distribution p ( x , z ) as follows
[Equation] where [Equation] 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 ) :
= (cid:82) π θ,φ ( x , z ) d z (cid:54) =
p ( x ) due to the contribution of the energy term.
Instead, we have
[Equation]
[Equation]
where Z ( x ) = (cid:82) q ( z | x ) e T φ ( x , z ) d z and Z = (cid:82) p ( x ) q ( z | x ) e T φ ( x , z ) d x d z = E p ( x ) [ Z ( x )] .
Subtracting a joint KL divergence D KL [ p ( x , z ) || π θ,φ ( x , z )] from I ( x ; z ) , we obtain the Generalized
MINE - DV lower bound on MI
[Equation]
[Equation]
[Equation]
By construction,
Eq.
(61) shows that the gap in
Generalized
MINE - DV is
D KL [ p ( x , z ) ||
1 Z 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 I MINE - DV ( T φ )
=
I GMINE - 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 D KL [ p ( x , z ) ||
1 Z 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 I GMINE - DV ( q ∗ θ , T φ ) ≥ I MINE - 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 I GMINE - DV ( q θ , T φ 0 = c )
=
I BA L ( 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
[Equation]
For the optimal energy function associated with a given q θ ( z | x ) , we have I MINE - 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 φ ) ≥ I GMINE - DV ( q θ , T φ ) ), and the difference corresponds to a marginal KL divergence in the x space
[Equation]
where π θ,φ ( x ) is deﬁned in Eq.
As in Poole et al.
(2019), we can also interpret the gap between the IBAL ( q θ , T φ ) and I GMINE - DV ( q θ , T φ ) as an application of Jensen’s inequality, with D KL [ p ( x ) || π θ,φ ( x )]
= log Z − E p ( x ) [log Z ( x )]
= log E p ( x )
[ Z ( x )] − E p ( x ) [log Z ( x )] and log E p ( x )
[ Z ( x )] ≥ E p ( x )
[log Z ( x )] .
D. P ROBABILISTIC I NTERPRETATION OF G ENERALIZED MINE - F
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 deﬁnition 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)
[Equation]
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 D GKL [ r ( z )
|| s ( z )] =
D KL [ r ( z ) ||
s ( z )] .
Finally, if r ( z ) is normalized and ˜ s ( z ) is unnormalized, one can easily conﬁrm that
[Equation]
Generalized EUBO Since the Generalized KL divergence is always nonnegative, we can deﬁne an
Generalized EUBO by adding the Generalized KL divergence between the normalized true posterior and an unnormalized approximate posterior ˜ π
( z | x ) .
[Equation]
[Equation]
where we assume that (cid:82) p ( z | x ) d z = 1 and let Z ˜ π ( x ) :
= (cid:82) ˜
π ( z | x ) d z .
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
[Equation]
[Equation]
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 )
[Equation]
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
[Equation]
z
[Equation]
≥ I GMINE
[Equation]
By construction, we can see that the gap in I GMINE - F ( q θ , T φ ) is equal to the Generalized KL divergence E p
[Equation]
φ
[Equation]
z
[Equation]
x )]] in Eq.
[Equation]
As in the case of MINE - DV , we obtain the standard
MINE - F lower bound I MINE - F ( T φ )
=
I GMINE - 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 I GMINE - 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 I GMINE - F ( q θ , T φ ) to the gap in I IBAL ( q θ , T φ ) .
First, note that the normalized distribution corresponding to ˜ π θ,φ ( z | x ) =
q θ ( z | x ) e T φ ( x , z ) − 1 matches the the energy-based posterior in the IBAL , π θ,φ ( z | x ) =
Z ( x ) q θ ( z | x ) e T φ ( x , z ) .
Thus, using Eq.
(66), we can write
[Equation]
where the ﬁnal term D GKL [ π θ,φ ( z | x ) ||
˜
π θ,φ ( z | x )] ≥ 0 is nonnegative because it is a generalized KL divergence.
After simpliﬁcation, we can also see that D GKL [ π θ,φ ( z | x ) ||
˜
π θ,φ ( z | x )] =
E p ( x ) [ Z ( x ) − 1 − log Z ( x )] ≥ 0 since log u ≤ u − 1 .
We visualize the relationship between
IBAL and Generalized MINE - DV in Fig.
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.
[Equation]
In this case, we have D GKL [ 1 Z p ( x ) q θ ( z | x ) e T φ ( x , z ) || p ( x ) q θ ( z | x ) e T φ ( x , z ) − 1 ] ≥ 0 .
Thus, Generalized
MINE - DV is tighter than Generalized MINE - F (see Fig. 5), which generalizes the ﬁnding in Poole et al.
(2019) that standard MINE - DV is tighter than standard MINE - F .
E MINE-AIS E. P ROPERTIES OF MINE-AIS
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 ∗ [Equation] =
I
[Equation]
x
[Equation]
[Equation]
and Cor.
C.2(a) implies that
[Equation]
(b) Suppose the energy function
T [Equation]
, 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
[Equation]
In particular, the contrastive term in Eq.
22 is upper bounded by D KL [ p ( z | x ) || q θ ( z | x )] .
See Fig. 3b.
Proof.
(a) Since the gap in the IBAL ( q θ , T φ ∗ ) is E p
( x ) (cid:2) D KL [ p ( z | x ) || π θ,φ ( z | x )] (cid:3) =
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
[Equation]
x ,
[Equation] [Equation]
where e c ( x ) is absorbed into the normalization constant.
At this value, the IBAL reduces to I ( x ; z )
[Equation]
(b) The BA bound is a special case of the IBAL ( q θ , T φ ) with constant T φ
= c .
Thus, by learning T φ , we can only improve upon I BA L ( 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 ) .
E. P ROOF OF P ROPOSITION 4.
Proposition 4.1.
For given q θ ( z | x ) and T φ ( x , z ) , lim K →∞ I GIWAE L ( q θ , T φ , K ) =
IBAL ( q θ , T φ ) .
Proof.
Inspecting the form of I GIWAE L ( q θ , T φ , K ) in Eq.
9 to the IBAL ( q θ , T φ ) in Eq.
22, we can see that for a ﬁxed q θ and T φ , both bounds include the same ﬁrst term I BA L ( q θ ) and the same numerator of the contrastive term E p ( x , z ) (cid:2)
e T φ ( x , z ) (cid:3) .
To prove the proposition, we can thus focus on characterizing the limiting behavior of the denominator
[Equation]
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 ﬁxed K and show that each of these sequences of lower and upper bounds converge to log Z θ,φ ( x ) as K → ∞ .
Using the squeeze theorem for sequences, this is sufﬁcient to demonstrate the claim that lim K →∞ log ˆ Z GIWAE ( K, x ) = log
Z θ,φ ( x ) in Eq. 83.
Since the other terms in I GIWAE L ( q θ , T φ , K ) and IBAL ( q θ , T φ ) are identical, this is sufﬁcient to prove the proposition.
Lower Bound: To lower bound log ˆ Z GIWAE ( K, x ) , we rely on the fact that the exponential function e T φ ( x , z ) ≥ 0 to simply ignore the contribution of the p ( z | x ) term.
[Equation]
As K → ∞ , we have log K − 1 K → 0 and lim
[Equation]
LB [Equation] [Equation] 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 p GIWAETGT ( 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
[Equation]
K, x ) = E p ( z (1) | x ) q θ ( z (
[Equation]
) |
[Equation] [Equation]
We consider the term inside the log as an expectation over a Bernoulli random variable with probabilities { 1 K , K − 1 K } . Using the concavity of log to bring this expectation outside, and the linearity of expectation to separate terms, we have
[Equation]
As K → ∞ , we have 1 K → 0 , K − 1 K → 1 , and
[Equation]
z
[Equation]
x ) (cid:2)
[Equation]
( x , z ) (cid:3) .
Finally,
[Equation]
As reasoned above, the convergence of both the upper and lower bound sequences to log Z θ,φ ( x ) implies lim K →∞ log ˆ Z GIWAE ( K, x )
= log
Z θ,φ ( x ) and lim K →∞
I GIWAE L ( q θ , T φ ) =
IBAL ( q θ , T φ ) , as desired.
E. C ONVERGENCE OF GIWAE SNIS D ISTRIBUTION TO MINE-AIS
Proposition E.2.
Let q GIWAEPROP K
( z | x ) denote the marginal SNIS distribution of GIWAE
[Equation]
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.
[Equation]
Proof.
Note that the distribution q GIWAEPROP ( z | x ; K ) is obtained by augmenting the standard (multisample) GIWAE proposal distribution q GIWAEPROP ( 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.
and we consider the distribution q GIWAEPROP ( z | x ; K ) over z after marginalizing the joint q GIWAEPROP ( z , s, z (1: K ) ) over z (1: K ) and s as in Eq.
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 )
[Equation]
Note that marginalization over z will always yield a sample from π θ,φ ( z | x ) .
Since marginalization can only reduce the KL divergence, we have
[Equation]
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 → ∞
[Equation]
Since q GIWAEPROP ( z (1: K ) | x ) = (cid:81)
Kk =1 q θ ( z ( k ) | x ) , it is clear that in the limit as K → ∞ ,
E q GIWAE PROP
[Equation]
x ) . Thus, the difference
[Figure]
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
[Equation]
Similar reasoning as from Eq.
100-Eq.
102 holds for the KL divergence in the reverse direction, with D KL [ p π
[Equation]
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 (cid:20) log 1 K K (cid:80) k
=1 e
T φ ( x , z ( k ) ) (cid:21) → log Z θ,φ ( x ) .
Finally, this implies lim K →∞ D KL [ p π θ,φ TGT ( s, z (1: K ) | x ) || q GIWAEPROP ( s, z (1: K )
| x )]
= 0 and thus lim K →∞ D KL [ π θ,φ ( z | x ) ||
q GIWAEPROP ( z | x ; K )]
= 0 as desired.
E. D ETAILS FOR T RAINING MINE-AIS
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 difﬁcult 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.
E. A PPROXIMATE R EVERSE A NNEALING FOR E VALUATING MINE-AIS
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 log Z θ ( 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 ) =
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:
[Equation]
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
D KL [ p [Equation] z | x ) ||
[Equation] θ ∗ [Equation] φ ∗ ( z | x )] is less that the gap of the BA bound D KL [ p ( z | x ) ||
q θ ∗ ( z | x )] .
Thus we have:
[Equation]
As an illustrative example, Fig.
6a shows a Gaussian base distribution q θ [Equation] ) , a multi modal true posterior distribution p ( z | x ) , a posterior approximation π θ ∗ ,φ ∗ ( z | x ) , whose energy function
T φ parameterized by µ φ , σ φ has been ﬁt to minimize D KL [ p ( z | x ) ||
π θ,φ ( z | x )] .
The standard reverse importance sampling, which take true samples from π θ,φ ( z | x ) , results in an upper bound on log Z θ,φ ( 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 log Z θ,φ ( 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 log Z θ,φ ( 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 log Z θ ( 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 log Z θ,φ ( x ) .
We can also see that both reverse AIS chains provide an upper bound on the log
Z θ,φ ( x ) .
E. M ULTI -S AMPLE AIS C ONVERGENCE R ESULTS FOR MINE-AIS
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 ﬂexible, energy-based posterior approximation, we also compare against the BA [Figure] F A PPLICATIONS TO M UTUAL I NFORMATION E STIMATION WITHOUT K NOWN M ARGINALS
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 q data ( x ) and the representation induced by a stochastic mapping q ψ ( z | x ) parameterized by ψ , [Equation]
[Equation]
where q ψ ( z ) = (cid:82) q data ( x ) q ψ ( z | x ) d x represents the ‘aggregated posterior’ (Makhzani et al., 2015) or induced marginal distribution over z . F. BA L OWER B OUND Using the notation of Eq.
105, we can write the BA lower bound as
[Equation]
negative
[Table]
loss of x
[Equation]
where p θ ( x | z ) is a variational distribution, parameterized by θ , that tries to match to the inverse encoding distribution q ψ ( x | z ) ∝ q data ( x ) q ψ ( z | x ) .
The ﬁrst 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 I BA L ( p θ ( x | z )) with respect to the parameters θ does not depend on the marginal q data ( 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 , H data ( 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 q data ( 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 GAN s (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 classiﬁcation labels (Alemi et al., 2016).
When y is the original data as in unsupervised learning with β - VAE s (Alemi et al., 2018), the distortion or reconstruction loss corresponds to maximizing the BA lower bound on the encoding mutual information in Eq.
F. GIWAE L OWER B OUND
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
[Equation]
Note that we can use a single joint sample from q data ( 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 −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 q GIWAEPROP ( s | z , x 1: K )
=
e T φ ( x ( s ) , z ) / (cid:80)
Kk =1 e T φ ( x ( k ) , z ) .
Finally, we note that optimization over the energy function T φ ( x , z ) improves upon I BA L ( p θ ( x | z )) in
Eq.
(107) by at most log K 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 q data ( 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 q data ( 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 I NFO -NCE lower bounds, it can be used as a drop-in replacement for either of these bounds for optimizing MI .
Contrastive I NFO -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).
F. MINE-AIS / IBAL L OWER B OUND
Recall that MINE - AIS estimation would involve an energy based variational approximation to the inverse encoding distribution q ψ ( x | z ) ,
[Equation]
Note that IBAL lower bound on I ( x ; z ) involves an intractable log partition function term
[Equation]
where the term in the denominator is the partition function
Z
[Equation]
When taking gradients as in Eq.
(24)-(25), we obtain
[Equation]
IBAL
[Equation]
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 ) e T φ ( 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 q data ( 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 log Z ( z ) = log E p θ ( x | z ) (cid:2)
e T φ ( x , z ) (cid:3) .
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 log Z ( 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. G B ACKGROUND ON AIS AIS (Neal, 2001) constructs a sequence of intermediate distributions { π t ( z ) } T t =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 Z T ( 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
Z t ( x ) where Z t ( x ) = (cid:90) π 0 ( z | x ) 1 − β t π T ( x , z ) β t d z .
AIS involves a extended state space proposal distribution q AISPROP ( z 0: T | x ) , obtained by sampling from the initial π 0 ( z | x ) and constructing transitions T f ( z t +1 | z t , x ) which leave π t ( z | x ) invariant.
The target distribution p AISTGT ( z 0: T | x ) is given by running the reverse transitions T r ( z t | z t +1 , x ) starting from a target or posterior sample π T ( z | x ) , as shown in Fig.
[Equation]
Since the forward and reverse kernels obey T f ( z t +1 | z t , x ) π t ( z t | x ) =
T r ( z t | z t +1 , x ) π t ( z t +1 | x ) , we can write the importance weights w :=
p AISTGT ( x , z 0: T ) /q
AISPROP ( z 0: T | x ) in the extended state space as w = p AISTGT ( x , z 0: T ) q
[Equation]
Assuming the initial distribution is normalized Z 0 = 1 , AIS provides the following estimate of the partition function ˆ [Equation] w
[Equation]
where w ( k ) =
p AIS TGT ( x , z ( k ) 0: T ) q AISPROP ( z ( k ) 0: T | x ) and z ( k ) 0: T ∼ q AISPROP ( z 0: T | x ) , (116) which is unbiased since E [ ˆ Z T ( x )] =
E z 0: T ∼ q AISPROP [ w ] =
Z T ( 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. [Equation] The lower bound is an extended state-space ELBO using the proposal and target distributions in
Eq.
(14), with gap equal to D KL [ q AISPROP ( z 0: T | x ) ||
p AISTGT
z 0: T | x )] .
The upper bound is an extended state space EUBO , whose gap is equal to D KL [ p AISTGT ( z 0: T | x ) ||
q AISPROP ( z 0: T | x )] .
These single-chain lower and upper bounds on log p ( x ) translate to upper and lower bounds on MI , I AIS L ( π 0 , T ) and I AIS U ( π 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 I NDEPENDENT M ULTI -S AMPLE AIS H. P ROBABILISTIC I NTERPRETATION AND B OUNDS
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
[Equation]
[Equation] [Equation] The importance ratio log p IM - AIS [Equation]
( x , z ) /q IM - AIS PROP ( x , z ) simpliﬁes to a mixture as in Eq.
30, although we are now averaging the importance weights along AIS chains
[Equation]
t +
[Equation]
Taking the expected log ratio under the proposal and target, respectively, yield lower and upper bounds on log p ( x )
[Equation]
H. P ROOF OF P ROPOSITION 3.
Proposition 3.1 ( Complexity in K ) .
For a given T , the bias of the Independent Multi-Sample AIS lower bound on MI , I IM - AIS L ( π 0 , T ) , decreases (at most) logarithmically in K .
[Equation]
To prove the proposition
[Equation]
we ﬁrst analyze the EUBO IM - AIS ( π 0 , K, T ) log p
[Equation]
so that EUBO IM - AIS ( π 0 , K, T ) ≥ EUBO AIS ( T ) − log K 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 I IM - AIS L ( π 0 , K, T ) =
E p ( x | z ) [log p ( x | z )]
− EUBO IM - AIS ( π 0 , K, T ) and
I AIS L ( π 0 , T ) =
E p ( x | z ) [log p ( x | z )] − EUBO AIS ( π 0 , T ) , we have I IM - AIS L ( π 0 , K, T ) ≤ I AIS L ( π 0 , T )
+
log K as desired.
I AIS L ( π 0 , T ) ≤ I IM - AIS L ( π 0 , K, T ) follows from the fact that multi-sample bounds are tighter than their single sample counterparts.
H. P ROOF OF P ROPOSITION 3.
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 = I IM - AIS U
( π 0 , K, T ) − I IM - AIS L ( π 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 ﬁrst simplify the difference in the single-sample upper and lower bounds as
[Equation]
T,K =1 L
[Equation]
We now assume perfect transitions, which is common in analysis of AIS (Neal, [-
 ;-]2001; Grosse et al., 2013).
Since each forward transition T f ( z t +1 | z t ) leaves π t ( z ) invariant, perfect transitions results in independent exact samples from
[Equation]
so that z 1 ∼ π 0 ( z ) , z 2 ∼ π β 1 ( z ) , z 3 ∼ π β 2 ( z ) , etc.
Similarly, perfect transitions in the reverse direction
T f
( z t | z t +1 ) results in exact samples from z t ∼ π β t ( z ) z (1: K [Equation]
[Equation]
so that z T − 1 ∼ π β T − 1 ( z ) and so on down to z 0 ∼ π 0 ( z ) .
To put this all together, we substitute back into Eq.
126 using a linear annealing schedule β t − β t − 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 z t , with z t ∼ π β t − 1 ( z ) for the forward chain and z t ∼ π β t ( z ) for the reverse chain δ T,K =1 L +
[Equation]
In moving from the third to the fourth line
[Equation]
we see that intermediate terms cancel since each term has the same weight under linear scheduling.
In contrast to Grosse et al.
[Equation] this allows us to derive linear scheduling result which holds for ﬁnite 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 .
I R EVERSE IWAE
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.
I. P ROBABILISTIC I NTERPRETATION AND B OUNDS
Similar to reverse importance sampling, we consider K 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, [Equation]
[Equation]
We visualize this sampling scheme in Fig.
Similarly to Eq.
30, the log-ratio simpliﬁes to
[Equation]
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.
[Equation]
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.
I. P ROOF OF L OGARITHMIC I MPROVEMENT OF RIWAE
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 ELBO RIWAE ( q, K ) can improve upon the single-sample ELBO by at most log K nats. [Equation]
[Equation]
Proof.
[Equation]
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 I NDEPENDENT R EVERSE M ULTI -S AMPLE AIS
In similar fashion to Reverse IWAE , we now use K independent reverse AIS chains to form an extended state space target distribution p ( TGT z (1: K ) 0: T , x ) , and a mixture proposal distribution which contains a single forward AIS chain (see Fig. 2).
[Equation]
As in Eq.
130-131, the unnormalized log ratio becomes
[Equation]
Taking expectations of the log unnormalized density ratio under the proposal and target [Equation] respectively
[Equation]
yield lower and upper bounds on log p ( x ) [Equation] 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
[Equation]
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
[Equation]
Proof.
[Equation]
where, in the third line, we have used the fact that the likelihood ratio is non-negative.
K C OUPLED R EVERSE M ULTI -S AMPLE AIS K. P ROBABILISTIC I NTERPRETATION AND B OUNDS
The Coupled Reverse Multi-Sample AIS extended state space target distribution in Fig.
2 initializes
K backward chains from a single target sample z T ∼ π T ( z | x ) , with remaining transitions denoted p AISTGT ( z 0: T − 1 | z T ) since they are identical to standard AIS in Eq.
[Equation] [Equation]
We can construct log-partition bounds using the unnormalized importance ratio
[Equation]
Taking the expected log ratio under the proposal and target yields lower and upper bounds on log p ( x ) ,
[Equation]
( z 0: T | x ) z (2: K ) 0: T − 1 ∼ p AISTGT ( z 0: T − 1 | z T , x ) (cid:20) log
K
K (cid:88) k
= q AISPROP ( z ( k ) 0: T − 1 , z T | x ) p AISTGT ( x , z ( k ) 0: T − 1 , z T ) (cid:21) ≤ log p ( x ) (146)
EUBO
[Equation]
z
[Equation]
x
[Equation]
[Equation]
x ) (cid:20) log
K
K (cid:88) k
= q AISPROP ( z ( k ) 0: T − 1 , z T | x ) p AISTGT ( x , z ( k ) 0: T − 1 , z T ) (cid:21) ≥ log p ( x ) (147)
K. P ROOF OF P ROPOSITION 3.
Proposition 3.3.
For a given T , the bias of the Coupled Reverse Multi-Sample AIS upper bound on
MI , I CR - AIS U ( π 0 , K, T ) , decreases (at most) logarithmically in K .
[Equation]
Proof.
log
[Equation]
i
[Equation] q [Equation] where, in the third line, we have used the fact that the likelihood ratio is non-negative.
Since I CR - AIS U ( π 0 , K, T ) =
E p ( x | z ) [log p ( x | z )] − ELBO CR - AIS ( π 0 , K, T ) and I AIS L ( π 0 , T ) =
E p ( x | z ) [log p ( x | z )] − ELBO AIS ( π 0 , T ) , then Eq.
148 implies I CR - AIS U
( π 0 , K, T ) ≥ I AIS U
( π 0 , T ) − log K as desired.
L E XPERIMENTAL D ETAILS L. D ATASETS AND M ODELS
We used MNIST (LeCun et al., 1998) and CIFAR -
(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.
L. M ULTI -S AMPLE AIS S ETTINGS
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 deﬁnitions, 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.
but with radius parameter set toinstead 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 CIFAR 10, the radius parameter for the temperature is 6 when the initial distribution is q ( z | x ) and 8 when it is log p ( z ) .
L. R UNTIME C OMPARISON
For MNIST , it took about 20 mins to run IWAE for K = 1 M , and it took about 8 minutes to run
AIS for T = 500 and about 8 hours for T = 30 k .
For CIFAR , it took about 70 mins to run AIS for
T = 500 steps and about nine days for T = 100 k .
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. and an initial distribution q θ ( z | x ) . We plot wall clock time on the x -axis, where increasing runtime reﬂects 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 conﬁrm 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.
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 [Figure] 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:
[Equation]
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 ) β .
L. A NALYTICAL S OLUTION OF THE M UTUAL I NFORMATION ON THE L INEAR VAE
In order to verify our implementations, we have derived the mutual information analytically for
Linear VAEs trained on MNIST .
For simplicity, we assume a ﬁxed 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 + (cid:15) , 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 (cid:15) ∼ N ( 0 , I ) .
It’s easy to show that the conditional likelihood is: p ( x | z )
[Equation]
N ( x | Wz + b , I ) and thus we can solve for the marginal:
[Equation]
Thus the differential entropy of x is:
[Equation]
Where k is the dimension of the observation.
The conditional entropy:
[Equation]
[Equation]
[Equation]
Thus the mutual information:
[Equation]
M A DDITIONAL E XPERIMENTS M. C OMPARISON OF M ULTI -S AMPLE AIS B OUNDS
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 log K 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 log K 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 ) , I IM - AIS L ( 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 log K .
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.
M. C ONFIDENCE I NTERVALS FOR M ULTI -S AMPLE AIS E XPERIMENTS
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% conﬁdence intervals for those experiments.
For MNIST experiments, we computed the conﬁdence interval over 8 batches of 16 data points.
For CIFAR experiments, due to computational constraints, we computed the conﬁdence 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.
[Figure]
[Figure] [Figure]
[Table] [Table]