I MPROVING M UTUAL I NFORMATION E STIMATION WITH A NNEALED AND E NERGY -B ASED B OUNDS
A BSTRACT
Mutual information ( MI ) is a fundamental quantity in information theory and machine learning.
However, direct estimation of MI 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 present a unifying view of existing MI bounds from the perspective of importance sampling, and propose three novel bounds based on this approach.
Since a tight MI bound without density information requires a sample size exponential in the true MI , we assume either a single marginal or the full joint density information is known.
In settings where the full joint density is available, we propose Multi-Sample
Annealed Importance Sampling ( AIS ) bounds on MI , which we demonstrate can tightly estimate large values of MI in our experiments.
In settings where only a single marginal distribution is known, we propose Generalized IWAE ( GIWAE ) and
MINE - AIS bounds. Our GIWAE bound uniﬁes variational and contrastive bounds in a single framework that generalizes I NFO NCE IWAE and Barber-Agakov bounds. Our MINE - AIS method improves upon existing energy based methods such as
MINE - DV and MINE - F
by directly optimizing a tighter lower bound on MI . MINEAIS uses MCMC sampling to estimate gradients for training and Multi Sample AIS for evaluating the bound.
Our methods are particularly suitable for evaluating MI in deep generative models, since explicit forms of 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 other applications in machine learning, MI 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, successes in neural network function approximation have encouraged a wave of variational or contrastive methods for MI estimation from samples only (Belghazi et al., 2018; van den Oord et al., 2018; Poole et al., 2019).
However, McAllester & Stratos (2020) have shown strong theoretical limitations on any estimator based on direct sampling without an analytic 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 known.
In this work, we view MI estimation from the perspective of importance sampling.
Using a general approach for constructing extended state space bounds on MI , we combine insights from importanceweighted autoencoder ( IWAE )
(Burda et al., 2016; Sobolev & Vetrov, 2019) and annealed importance sampling ( AIS ) (Neal, 2001) to propose Multi-Sample AIS bounds in Sec.
We empirically show that this approach can tightly estimate large values of MI when the full joint distribution is known.
Our importance sampling perspective also suggests improved MI lower bounds that assume access to only joint samples for optimization, but require a single marginal distribution for evaluation.
In
Sec. 2.4, we propose Generalized IWAE ( GIWAE ), which generalizes both IWAE and I NFO NCE (Poole et al., 2019) and highlights how variational learning can complement multi-sample contrastive estimation to improve MI lower bounds.
Finally, in Sec. 4 we propose MINE - AIS , which optimizes a tighter lower bound than MINE (Belghazi et al., 2018) called the Implicit Barber-Agakov Lower bound ( IBAL ).
We demonstrate that the IBAL corresponds to the inﬁnite-sample limit of the GIWAE lower bound, although our proposed energy-based training scheme involves only a single ‘negative’ contrastive sample obtained using Markov Chain Monte Carlo ( MCMC ).
MINE - AIS then uses MultiSample AIS to evaluate the lower bound on MI and shows notable improvement over existing variational methods in the challenging setting of MI estimation for deep generative models.
We summarize the MI bounds discussed in this paper and the relationships between them in Fig.
P ROBLEM S ETTING
The mutual information between two random variables x and z with joint distribution p ( x , z ) is
[Equation]
) [log p ( x )] , (1) 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 MI 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 only a single marginal is known appear, for example, 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 ) in cases where samples from the joint distribution p ( x ) p ( z | x ) are available.
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.
In this case, the difﬁculty of MI estimation reduces to estimating the log partition function log p ( x ) , for which importance sampling ( IS ) based methods are among the most well studied and successful solutions.
U NIFYING M UTUAL I NFORMATION B OUNDS VIA I MPORTANCE S AMPLING
In this section, we present a uniﬁed view of mutual information estimation from the perspective of extended state space importance sampling.
This general approach provides a probabilistic interpretation of many existing MI bounds and will suggest novel extensions in Sec. 3 and Sec. 4.
2. A G ENERAL A PPROACH FOR E XTENDED S TATE S PACE I MPORTANCE S AMPLING B OUNDS
To estimate the log partition function, we construct a proposal q PROP ( z ext | x ) and target p TGT ( x , z ext ) distribution over an extended state space, such that the normalization constant of p TGT ( x , z ext ) is
Z TGT = (cid:82) p TGT ( x , z ext )
d z ext = p ( x ) and the normalization constant of q PROP ( x , z ext ) is Z PROP = 1 .
Taking expectations of the log importance weights log p TGT ( x , z ext ) /q
PROP ( z ext | x ) under the proposal and target, respectively, we obtain lower and upper bounds on the log partition function
[Equation]
These bounds correspond to extended state space versions of the ELBO and EUBO .
In particular, the gap in the lower bound is the forward KL divergence D KL [ q PROP ( z ext | x ) (cid:107) p TGT ( z ext | x )] and the gap in the upper bound equal to the reverse KL divergence D KL [ p TGT ( z ext | x ) (cid:107) q PROP ( z ext | x )] .
See App. A.
B ARBER -A GAKOV L OWER AND U PPER B OUNDS
As a ﬁrst example, consider the standard ELBO ( q θ ) and EUBO ( q θ ) bounds, which are derived from simple importance sampling using a variational distribution q θ ( z | x ) and z ext =
z in Eq.
Plugging these lower and upper bounds on log p ( x ) into Eq.
(1), we obtain upper and lower bounds on MI as
I BA L ( q θ ) : = E p ( x , z )
[Equation]
The left hand side of Eq.
(3) is the well-known Barber-Agakov ( BA ) bound
(Barber & Agakov, 2003), which has a gap of E p ( x ) [ D KL [ p ( z | x ) (cid:107) q θ ( z | x )]] .
We refer to the right hand side as the BA upper bound I BA U ( q θ ) , with a gap of E p ( x ) [ D KL [ q θ ( z | x ) (cid:107) p ( z | x )]] .
In contrast to I BA U ( q θ ) , note that
I BA L ( q θ ) does not require access to the conditional p ( x | z ) to evaluate the bound.
2. I MPORTANCE W EIGHTED A UTOENCODER
The IWAE lower and upper bounds on log p ( x ) (Burda et al., 2016; Sobolev & Vetrov, 2019) improve upon simple importance sampling by extending the state space using multiple samples z ext =
z (1: K ) .
Consider a proposal q IWAEPROP ( z (1: K ) |
x ) with K independent samples from a given variational distribution q θ ( z | x ) .
The extended state space target p IWAETGT ( z (1: K ) | x ) is a mixture distribution involving a single sample from the posterior p ( z | x ) or joint p ( x , z ) and K − 1 samples from q θ ( z | x ) q
[Equation]
The log importance weight log
[Equation]
K )
[Equation]
q IWAEPROP
[Equation] z [Equation] 1: K ) | x ) reduces to the familiar ratio in the IWAE objective, while the normalization constant of p IWAETGT ( x , z (1: K ) ) is p ( x ) .
As in Sec. 2.1, taking expectations under the proposal or target yields a lower or upper bound, respectively,
[Equation]
As for the standard ELBO and EUBO , the gap in the lower and upper bounds are D KL [ q IWAEPROP (cid:107) p IWAETGT ] and D KL [ p IWAE TGT (cid:107) q IWAE PROP ] , respectively.
See App.
B for detailed derivations.
As in Sec. 1.1, with known p ( x | z ) , the lower and upper bounds on log p ( x ) , ELBO IWAE ( x ; q θ , K ) and EUBO IWAE ( x ; q θ , K ) translate to upper and lower bounds on MI , which we denote as I IWAE U ( q θ , K ) and I IWAE L ( q θ , K ) .
While it is well-known that increasing K leads to tighter IWAE bounds (Burda et al., 2016; Sobolev &
Vetrov, 2019), in App.
B.2, we characterize the improvement of ELBO IWAE ( q θ , K ) over ELBO ( q θ ) as a KL divergence.
For EUBO IWAE ( q θ , K ) , we show that the KL divergence measuring its improvement over EUBO ( q θ ) is limited to log K , which implies I IWAE L ( q θ , K ) ≤ I BA L ( q θ )
+
log K .
2. G ENERALIZED IWAE
In this section, we consider a family of Generalized IWAE ( GIWAE ) lower bounds, which improve upon I BA L ( q θ ) using multiple samples and a contrastive critic function T φ ( x , z ) , but do not require access to p ( x | z ) .
While similar bounds appear in (Lawson et al., 2019; Sobolev, 2019), we provide a thorough discussion of special cases, and empirical analysis for MI estimation in Sec. 5.2.
We also show that our IBAL bound in Sec.
4 corresponds to the inﬁnite sample limit of GIWAE .
To derive a probabilistic interpretation for GIWAE , 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 p IWAETGT ( z (1: K ) | x ) in Eq.
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 critic function T φ q GIWAEPROP ( z (1: K ) , s | x
[Equation]
=
[Equation]
k
[Equation]
q ( z ( k ) | x ) (cid:33) q GIWAEPROP ( s | z
(1: K ) , x ) , where q GIWAEPROP ( s | z (1: K ) , x ) =
e
[Equation]
x
[Equation]
s
[Equation]
K
[Equation]
k =
[Equation]
x
[Equation]
k
[Equation]
We discuss the GIWAE probabilistic interpretation in App.
C.1, and ﬁnd that only the GIWAE upper bound on log p ( x ) provides practical beneﬁt in App.
C.2.
The corresponding MI lower bound is
I GIWAE
[Equation]
[Table]
We observe 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 a single positive sample drawn from p ( z | x ) from K − 1 negative samples drawn from q θ ( z | x ) .
Relationship with BA With a constant T φ ( x , z ) = const, the second term in GIWAE vanishes and we have I GIWAE L ( q, T φ = const , K )
=
I BA L ( q θ ) for all K .
For K = 1 , I GIWAE L ( q, T φ , K = 1) also equals the BA lower bound for all T φ .
Similarly to BA , GIWAE requires access to the analytical form of p ( z ) to evaluate the bound on MI .
However, both the BA and GIWAE lower bounds can be used to optimize MI even if no marginal is available.
See App.
N. for detailed discussion.
Relationship with InfoNCE When the prior p ( z ) is used in place of q θ ( z | x ) , we can recognize the second term in Eq.
(7) 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).
Relationship with IWAE The following proposition characterizes the relationship between
I GIWAE L ( q θ , T φ , K ) in Eq.
and I IWAE L ( q θ , K ) from Sec.
See App.
C.5 for proofs.
Proposition 2.1 (Improvement of IWAE over GIWAE ) .
For a given q θ ( z | x ) and any T φ ( x , z ) ,
I IWAE L ( q [Equation] [Equation]
GIWAE [Equation] [Equation] GIWAETGT (
[Equation]
(1: K ) | x ) (cid:104) D KL [Equation] p GIWAETGT ( s | z (1: K ) , x ) (cid:107) q GIWAEPROP ( s | z (1: K ) , x )] (cid:105) .
Corollary 2.2.
For a given q θ ( z | x ) and K > 1 , the optimal critic function is the true log importance weight 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]
[Equation]
Corollary 2.3.
Suppose the critic function T φ ( x , z ) is parameterized by φ , and
∃ φ 0 s.t. ∀
( x , z ) , T φ 0 ( x , z ) = const.
For a given q θ ( z | x ) , let T φ ∗ ( x , z ) denote the critic function that maximizes the GIWAE lower bound.
Using Cor.
2.2, we have
[Equation]
[Equation] While the GIWAE lower bound on MI does not assume access to the full joint distribution, Cor. 2. suggests that the role of the critic 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
I GIWAE L ( q, T φ ∗ , K ) can improve upon BA , this improvement is at most logarithmic in K .
Relationship with Structured I NFO NCE
Since I NFO NCE and Structured I NFO NCE (S- I 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 (see App.
C.6)
[Equation]
[Table]
I LB
Proposal ELBO
I UB z K [Table] [Table]
j
[Table]
[Table]
[Table]
z jT() z 10() z 11() z 1T-1() ...
z 1T() z K0() z K1() z KT-1() ...
z j0() z j1() z jT-1() ...
z T z 10() z 11() z 1T-1() ...
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 bounds using multiple samples K , as in IWAE .
In this section, we ﬁrst review AIS , which extends the state space using T intermediate densities.
We then show that these approaches are complementary, and derive two practical Multi-Sample AIS methods that provide tighter bounds by combining insights from IWAE and AIS .
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 =
π 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
[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 t ( z t | z t − 1 ) which leave π t − 1 ( z | x ) invariant.
The target distribution p AISTGT ( z 0: T | x ) is given by running the reverse transitions ˜
T t ( z t − 1 | z t ) starting from a target or posterior sample π T ( z | x ) , as shown in Fig.
[Equation]
Taking expectations of the log importance weights under the proposal and target again yields a lower and upper bound on the log partition function log p ( x ) (App. E).
These single-chain lower and upper bounds translate to upper and lower bounds on MI , I AIS U ( π 0 , T ) and I AIS L ( π 0 , T ) , which were suggested for MI estimation in the blog post of Sobolev (2019).
To characterize the bias reduction for
AIS with increasing T , we prove the following proposition.
Proposition 3.1 (Complexity in T ) .
Assuming perfect transitions and a geometric annealing path with linearly-spaced { β t } Tt =1 , the sum of the gaps in the AIS sandwich bounds on MI , I AIS U ( π 0 , T ) −
I AIS L ( π 0 , T ) , reduces linearly with increasing T .
See App.
D.1 for a proof.
In our experiments in Sec. 5, we will ﬁnd that this linear bias reduction in
T is crucial for achieving tight MI estimation when both p ( z ) and p ( x | z ) are known.
However, we can further tighten these AIS bounds using multiple annealing chains ( K > 1 ), and we present two practical extended state space approaches in the following sections.
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 AISTGT 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.
As in Sec. 2.1, taking the expectation of the log unnormalized density ratio under the proposal and target yields lower and upper bounds on log p ( x ) , respectively,
[Equation]
which again have KL divergences as the gap in their bounds (see App. E).
Independent Multi-Sample
AIS reduces to IWAE for T = 1 , and reduces to single-sample AIS for K = 1 .
Both upper and lower bounds are tight as K → ∞ or T → ∞ , and translate to lower and upper bounds on MI as in
Sec. 1.1.
In App.
E.3, we show that the Independent Multi-Sample AIS lower bound on MI is limited to logarithmic improvement over single-sample AIS , with I IM - AIS L ( π 0 , K, T ) ≤ I AIS L ( π 0 , T )
+log K .
3. C OUPLED R EVERSE M ULTI -S AMPLE AIS B OUNDS
We can exchange 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.
[Equation]
Note that partition function ratio is Z TGT / Z PROP
= p ( x ) K /p ( x ) K − 1 = p ( x ) .
Using these distributions, we derive log p ( x ) and MI bounds in App.
G.
However, the these bounds 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 ).
As shown in Fig.
2, the extended state space target distribution initializes K backward chains from a single target sample z T ∼ π T ( z | x ) , with the remaining transitions p AISTGT ( z
[Equation]
) matching 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]
We show in App.
H.3 that the Coupled Reverse Multi-Sample AIS upper bound on MI is limited to logarithmic improvement over single-sample AIS , with I CR - AIS U ( π 0 , K, T ) ≥ I AIS U ( π 0 , T ) − log K .
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.1 that 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 , as we show for
I IWAE L ( q θ , K ) (Cor. B.3), I IM - AIS L ( π 0 , K, T ) (App. E.3), and I CR - AIS U ( π 0 , K, T ) (App.
H.3).
Based on these arguments, we recommend increasing K until computation can no longer be parallelized on a given hardware and allocating all remaining resources to increasing T .
Comparing Multi-Sample AIS Bounds In App.
I Fig.
5, we compare performance of our various
Multi-Sample AIS bounds in order to recommend which to use in practice.
For the upper bound on MI , we recommend the Independent Multi-Sample AIS ELBO , or the forward direction of BDMC , since it uses independent samples and is not limited to log K improvement.
The results are less conclusive for the lower bounds on MI .
While the MI lower bound obtained from EUBO IM - AIS ( x ; π 0 , K, T ) ( RHS of Eq.
(14)) can improve upon single-sample AIS by at most log K (App. E.3), this improvement is easily obtained for low T and may be used to quickly estimate MI of a similar magnitude as log K .
The Coupled Reverse AIS lower bound on MI requires moderate values of T to match or marginally improve on Independent Multi-Sample AIS , suggesting that the preferred MI lower bound may differ based on the scale of the true MI and amount of available computation.
MINE-AIS E STIMATION OF M UTUAL I NFORMATION
For settings where the conditional p ( x | z ) is unknown, we propose MINE - AIS , which is inspired by
MINE but optimizes a tighter lower bound on MI (App. J).
Although this bound involves an intractable log partition function, we present a stable, energy-based training scheme and use our Multi-Sample
AIS methods from Sec.
3 to evaluate the bound.
Consider an ﬂ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 θ,φ ( x ) =
E
[Equation]
x
[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.
J.1, we obtain
I ( x
[Equation]
with the gap of the IBAL equal to E p ( x ) [ D KL [ p ( z | x ))
(cid:107) π θ,φ ( 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 φ ) .
Proposition 4.1.
For a given q θ ( z | x ) , the optimal IBAL critic function equals the log importance weights up to a constant T ∗ ( x , z )
= log p ( x , z ) q θ ( z | x ) + c ( x ) .
For this T ∗ , we have IBAL ( q θ , T ∗ ) =
I ( x ; z ) .
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.
L.2.
Proposition 4.2.
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.
L.1 that the IBAL contrastive term is ﬂexible enough to close the entire gap in the BA bound.
However, this ﬂexible contrastive term comes at the cost of tractability, as the IBAL in Eq.
(18) involves a log partition function log E q θ ( z | x ) [ e T φ ( x , z ) ] compared to the ﬁnite-sum term in GIWAE .
Energy-Based Training of IBAL Although the log partition function log Z θ,φ ( x ) in the IBAL is in- tractable to evaluate, we only require an unbiased estimator of its gradient for training.
Differentiating
Eq.
(18) with respect to the parameters θ and φ , respectively, we obtain
[Equation]
[Table]
Eq
19) indicates maximize θ φ , we need to increase of T φ ( x , z ) or log q θ ( z | x ) samples p ( x z ) , lower it As is common in training energy-based models, it is difﬁcult to draw samples from π θ,φ ( z | x ) .
To reduce the cost and variance of the estimated gradient, we initialize chains from true posterior sample z 0 ∼ p ( z | x ) as in contrastive divergence training (Hinton, 2002) and run M of Hamiltonian
Monte Carlo ( HMC ) transition kernels T 1: M ( z | z 0 , x ) (Neal, 2011).
See App.
M.1 for details.
Multi-Sample AIS Evaluation of IBAL After training the critic function using the procedure above, we still need to evaluate the IBAL lower bound on MI .
We can easily upper bound IBAL ( q θ , T φ ) using a Multi-Sample AIS lower bound on log Z θ,φ ( x ) , but this does not ensure a lower bound on
MI .
In order to upper bound log Z θ,φ ( x ) and preserve a lower bound on I ( x ; z ) , Multi-Sample AIS requires true samples from π θ,φ ( z | x ) to initialize backward AIS chains.
Since these samples are unavailable, we instead initialize backward chains from the true posterior p ( z | x ) instead of π θ,φ ( z | x ) . We derive sufﬁcient conditions under which this scheme preserves an upper bound on log Z θ,φ ( x ) in
App. M.2 and provide empirical validation in App.
M.3.
E XPERIMENTS
In this section, we evaluate our proposed MI bounds on VAE s and GAN s trained on MNIST and CIFAR .
M ULTI -S AMPLE AIS E STIMATION OF M UTUAL I NFORMATION
We compare Multi-Sample AIS MI estimation against IWAE , since both methods assume the full joint distribution is available.
For the initial distribution of AIS or variational distribution of IWAE , we experiment using both the prior p ( z ) and a learned Gaussian q θ ( z | x ) .
Table 1 summarizes our results.
IWAE As described in Sec. 2.3, IWAE bounds encompass a wide range of MI estimators.
The K =bounds with learned q θ ( z | x ) correspond to BA bounds, while for K > 1 and p ( z ) as the proposal, we obtain Structured
I NFO NCE.
While the IWAE upper bound on MI , which uses the log p ( x ) lower bound with independent sampling from q θ ( z | x ) , is tight for certain models, we can see that the improvement of the IWAE lower bound on MI is limited by log K .
In particular, we need exponentially large sample size to close the gap from the BA lower bound ( K = 1 ) to the true MI .
For example, on
CIFAR GAN 100, at least e 460 total samples are required to match the lower bound estimated by AIS .
Multi-Sample AIS
We evaluate Multi-Sample AIS 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 for all models and datasets considered.
This is in stark contrast to the exponential sample complexity required for the IWAE MI lower bound, and highlights that increasing T in Multi-Sample AIS is a practical way to reduce bias using additional computation.
We provide runtime details in App.
O.1.3, and provide additional results comparing different Multi-Sample AIS bounds in App.
I Fig.
5. E NERGY -B ASED E STIMATION OF M UTUAL I NFORMATION
In this section, we evaluate the family of GIWAE and MINE - AIS bounds, which assume access to a known marginal p ( z ) but not the conditional p ( x | z ) .
We summarize these bounds in Fig.
3b.
BA, IWAE, and GIWAE Bounds Recalling that I GIWAE L ( q θ , T φ , K ) in Eq.
(7) consists of the BA lower bound and a contrastive term, we report the contribution of each term in Fig.
3a. For a ﬁxed q θ ( z | x ) , Cor.
2.3 shows I BA L ( q θ ) ≤ I GIWAE L ( q θ , T φ ∗ , K ) ≤ I IWAE L ( q θ , K ) ≤ I BA L ( q θ ) +
log K . [Table] [Figure] Although we perform separate optimizations for each entry in Fig.
3a, we ﬁnd that these relationships hold in almost all cases.
We also observe that GIWAE can approach the performance of IWAE , despite the fact that GIWAE uses a learned T φ ( x , z ) instead of the optimal critic in IWAE (Cor. 2.2).
For the Linear VAE , the true posterior is in the Gaussian variational family and I BA L ( q θ ) is close to the analytical MI .
In this case, the contrastive term provides much less than log K improvement for GIWAE and IWAE , since even the optimal critic function cannot distinguish between q θ ( z | x ) and p ( z | x ) .
As K increases, we learn a worse q θ in almost all cases, as measured by a lower BA term.
This allows the contrastive term to achieve closer to its full potential log K , resulting in a higher overall bound.
For more complex VAE and GAN posteriors, there is a reduced tradeoff between the terms since the variational family is far enough from the true posterior (in reverse KL divergence) that either GIWAE or IWAE critic functions can approach log K improvement without signiﬁcantly lowering the BA term.
In all cases, (Structured)
I NFO -NCE bounds saturate to log K .
MINE-AIS Bounds
We report MINE - AIS results using a ﬁxed Gaussian p ( z ) as the base variational distribution and Multi-Sample AIS evaluation of IBAL ( p ( z ) , T φ ) .
We can see in Fig.
3a that MINE - AIS improves over BA due to its ﬂexible, energy-based variational family.
To evaluate the quality of the learned T φ ( x , z ) , we compare the IBAL to the Multi-Sample AIS lower bound, which assumes access to p ( x | z ) and corresponds to the optimal critic in Prop. 4.1.
We ﬁnd that MINE - AIS underestimates the ground truth MI by 11% and 24% on MNIST - VAE and MNIST - GAN , respectively.
We also observe that MINE - AIS outperforms GIWAE (and IWAE ) and by 31% and 49% on MNIST - VAE and MNIST - GAN , respectively.
To investigate whether this is due to its improved critic function or a more costly evaluation procedure, in Fig. 3a we report Multi-Sample AIS evaluation of IBAL ( q θ , T φ ) for q θ , T φ learned by optimizing the GIWAE or I NFO NCE lower bounds with K = 100 . GIWAE and I NFO NCE results only marginally improve, indicating that their critic functions are suboptimal and far from the true log importance weights.
2 We conclude that the improvement of MINE - AIS can primarily be attributed to learning a better critic function using energy-based training.
C ONCLUSION
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 motivated our novel GIWAE , Multi-Sample AIS , and MINE - AIS bounds.
When the conditional is not known, our GIWAE bounds highlight how variational bounds can complement contrastive learning to improve lower bounds on MI beyond known log K limitations.
When the full joint distribution is known, we show that our Multi-Sample AIS bounds can tightly 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 MI estimation when a single analytic marginal or conditional density is available.
A CKNOWLEDGEMENTS
The authors thank Shengyang Sun, Guodong Zhang, Vaden Masrani and Umang Gupta for helpful comments on drafts of this work.
We also thank the anonymous reviewers whose comments greatly improved the presentation and encouraged us to derive several additional propositions.
MG, RG and
AM acknowledge support from the Canada CIFAR AI Chairs program.
Table of Contents
[Table] [Table]
[Table] [Table]
[Table] [Table]
[Table] [Table]
[Table] [Table]
[Table] [Table]
[Table] [Table]
[Table] [Table]
[Table] [Table]
E Independent Multi
[Table]
AIS
F Reverse IWAE
[Table]
[Table]
EUBO . . . . . . . . . . . . . . . . .
[Table]
[Table]
Tighten
Increasing K . . . . . . . . . . . . . .
I Comparison of Multi-Sample AIS Bounds
J Generalized Mutual Information Neural Estimation ( GMINE )
[Table]
[Table] [Table]
[Table] [Table]
Marginals
[Table] [Table]
[Table]
O Experimental Details
[Table] [Table]
[Table] [Table] [Table]
[Table] [Table]
A A G ENERAL A PPROACH FOR D ERIVING E XTENDED S TATE S PACE B OUNDS ON L OG P ARTITION F UNCTIONS
In this section, we give a short proof that the gap in our general extended state space bounds from
Sec.
2.1 corresponds to a forward or reverse KL divergence.
We derive various upper and lower bounds on log p ( x ) using this approach throughout the paper and appendix, and we provide a visual summary in Fig.
First, we consider an extended state space target p TGT ( x , z ext ) and proposal q PROP ( x , z ext ) distributions.
For all cases discussed in this work, we will choose our target and proposal distributions such that log Z TGT ( x ) Z PROP ( x ) = log p ( x ) .
For example, a common construction is to have
Z TGT ( x ) = (cid:82) p TGT ( x , z ext )
d z ext = p ( x ) and Z PROP ( x ) =
(cid:82) q PROP ( x , z ext )
d z ext = 1 .
Our ‘reverse’ importance sampling bounds App.
F-G construct target and proposal such that Z TGT ( x ) =
p ( x ) K and Z PROP ( x ) = p ( x ) K − 1 , which still yields Z TGT ( x ) / Z PROP ( x ) = p ( x ) .
Each pair of extended state-space proposal and target distributions provides both an upper and lower bound on the log partition function.
Taking the expected log ratio of unnormalized densities under either the proposal or target distribution, we have
[Equation] [Equation]
To conﬁrm that these are indeed lower and upper bounds for any q PROP and p TGT , we can show that the gap in the lower bound in Eq.
(21) is the forward KL divergence D KL [ q PROP (cid:107) p TGT ] , and the gap in the upper bound is the reverse KL divergence, D KL [ p TGT (cid:107) q PROP ]
[Equation]
Thus, the bounds in Eq.
(21) directly generalize the standard ELBO ( log p ( x ) − D KL [ q θ ( z | x ) (cid:107) p ( z | x )] ) and EUBO ( log p ( x ) + D KL [ p ( z | x ) (cid:107) q θ ( z | x )] ), which appear as special cases when K = 1 , T = 1 , the proposal distribution is q PROP ( z | x ) =
q θ ( z | x ) , and the target distribution is p TGT ( z | x ) =
p ( z | x ) .
In what follows, our extended state space proposal or target distributions may include q θ ( z | x ) as an initial or base variational distribution, with the posterior p ( z | x ) often appearing within target distributions p TGT ( x , z ext ) .
In Fig.
4, we summarize various extended state space proposal (third column) and target distributions (fourth column).
We emphasize that the base variational distribution q θ ( z | x ) (blue circles) and posterior distribution p ( z | x ) (red circles) 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 posterior 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 as special cases of this general approach, thus interpreting each importance sampling bound in terms of probabilistic inference in an extended state space.
B I MPORTANCE W EIGHTED A UTOENCODER (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]
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.
(21), we can obtain lower and upper bounds on log p ( x ) by taking expectations under the proposal and target distributions, respectively.
Alternative Probabilistic Interpretation We now present an alternative probabilistic interpretation of IWAE , which is similar to Domke & Sheldon (2018) and will be used as the foundation for our
GIWAE bounds in App.
C. Consider the following extended state space target distribution,
[Equation] [Equation]
Note that marginalization over s leads to the IWAE mixture target in Eq.
(4) or Eq.
We consider the extended state space proposal
[Equation]
[Table]
z 1T() z K0() z K1() z KT-1() ...
z KT() z j0() z j1() z jT-1() ...
z jT() z 10() z 11() z 1T-1() ...
z 1T() z K0() z K1() z KT-1() ...
z j0() z j1() z jT-1() ...
z T z 10() z 11() z 1T-1() ...
z K0() z K1() z KT-1() ...
z j0() z j1() z jT-1() ...
z T z 10() z 11() z 1T-1() ... where we have deﬁned q IWAEPROP ( s | z
(1: K ) , x ) =
p ( x , z ( s ) )
q θ ( z ( s ) | x ) / K (cid:80) k
= p ( x , z ( k ) )
q θ ( z ( k ) | x ) .
As desired, the log importance weight match Eq.
[Equation]
Lower Bound on log p ( x ) and Upper Bound on MI Using the general approach in App.
A, taking expectations under q IWAEPROP leads to a lower bound on log p ( x )
[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]
q
[Equation]
Translating this to a lower bound on MI with known p ( x | z ) ,
[Equation]
B. P ROOF OF L OGARITHMIC I MPROVEMENT IN K FOR IWAE EUBO
We ﬁrst recall results that IWAE bounds tighten with increasing K .
Proposition B.1.
For given q θ ( z | x ) , ELBO IWAE ( x ; q θ , K + 1) ≥ ELBO IWAE ( x ; q θ , K ) (Burda et al., 2016) and EUBO IWAE ( x ; q θ , K + 1) ≤ EUBO IWAE ( x ; q θ , K ) (Sobolev & Vetrov, 2019).
Proposition B.2 (Improvement of IWAE and ELBO or EUBO ) .
Let p IWAETGT ( s | x , z (1: K ) )
=
p ( x , z ( s ) )
q θ ( z ( s ) | x ) / (cid:80)
Kk =1 p ( x , z ( k ) )
q θ ( z ( k ) | x ) denote the normalized importance weights and U ( s ) indicate the uniform distribution over K discrete values.
Then, we can characterize the improvement of
ELBO IWAE ( x ; q θ , K ) and EUBO IWAE ( x ; q θ , K ) over ELBO ( x ; q θ ) and EUBO ( x ; q θ ) using KL divergences, as follows ELBO IWAE
[Equation]
Prop.
B.2 demonstrates that the improvement of the IWAE log partition function bounds over its single-sample counterparts is larger for more non-uniform SNIS weights.
Notably, the improvement of EUBO IWAE ( x ; q θ , K ) over the single-sample EUBO ( x ; q θ ) is limited by log K .
Proof.
We ﬁrst note that the single-sample ELBO ( x ; q θ ) and EUBO ( x ; q θ ) are special cases of
ELBO GIWAE ( x ; q θ , T φ , K ) and EUBO GIWAE ( x ; q θ , T φ , K ) (Lemma C.1) with T φ = const.
As a result, the gap between ELBO ( x ; q θ ) and ELBO IWAE ( x ; q θ , K ) , for example, follows as a special case of the gap between ELBO GIWAE ( x ; q θ , T φ , K ) and ELBO IWAE ( x ; q θ , K ) , which we characterize in
Lemma C.2 (App. C.4).
The result in Prop.
B.2 follows directly.
We Improvement of EUBO IWAE :
Note D KL [ p IWAETGT ( s | z (1: K ) , x ) (cid:107)U ( s )]
= log K − H ( p IWAETGT ( s | z (1: K ) , x )) is bounded above by log K since the entropy of a discrete random variable is nonnegative.
Thus, the improvement of EUBO IWAE ( x ; q θ , K ) over EUBO ( x ; q θ ) is limited to log K .
We prove similar results in Prop. 2.1 and Cor.
2.3 (App. C.4), Prop. E.2 (App. E.3) and Prop.
H.2 (App.
H.3).
Improvement of ELBO IWAE : On the other hand, the KL divergence D KL [ U ( s ) (cid:107) p GIWAETGT ( s |
z (1: K ) , x )] is not limited by log K .
However, we do know that the improvement of ELBO IWAE ( x ; q θ , K ) over
ELBO ( x ; q θ ) will be limited by D KL [ q θ ( z | x ) (cid:107) p ( z | x )] , the gap of ELBO ( x ; q θ ) .
Prop.
B.2 shows that both IWAE upper and lower bounds on log p ( x ) improve upon their singlesample counterparts, with the improvement of EUBO IWAE ( x ; q θ , K ) over EUBO ( x ; q θ ) limited by log K .
Cor.
B.3 follows directly by translating these results to bounds on MI .
Corollary B.3.
IWAE bounds on MI improve upon the BA bounds with the following relationships:
[Equation]
Proof.
Recall from Sec.
1.1 and Sec. 2.2 that I BA L ( q θ ) =
E p ( x , z ) [log p ( x | z )] − E p ( x ) [ EUBO ( x ; q θ )] , and I IWAE L ( q θ , K ) =
E p ( x , z ) [log p ( x | z )] − E p ( x ) [ EUBO IWAE ( x ; q θ , K )] .
Using Prop.
B.2, the fact that EUBO ( x ; q θ ) − log K ≤ EUBO IWAE ( x ; q θ , K ) for any x implies that
[Equation]
which results in 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. E XPERIMENTAL R ESULTS SHOWING L OGARITHMIC I MPROVEMENT FOR I IWAE L ( q θ , K )
In Sec. 2.4, 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 IWAE L ( q θ , K ) =
E p
[Equation]
x , z )
[Equation]
This way of writing I IWAE L ( q θ , K ) provides additional intuition for the result in Prop.
B.2 and
Cor.
B.3.
In particular, the improvement of I IWAE L ( q θ , K ) over I BA L ( q θ ) is simply the contrastive term
[Table]
which is limited to log K
[Table]
Table 2 and Table 3 show the IWAE objective decomposition to the BA term and the contrastive term, on VAE s and GAN s trained MNIST and CIFAR -10 dataset.
We can see that in all the experiments, the contribution of the contrastive term is always less than or equal to log K .
For MNIST VAE s and GAN s with two dimensional latent spaces, the contrastive term may contribute notably less than log K .
In these cases, even the optimal critic function T ∗ ( x , z ) = log p ( x , z ) q θ ( z | x ) + c ( x ) , as used in IWAE , has difﬁculty distinguishing posterior and variational samples.
[Table]
However, the contribution of the contrastive term is almost exactly log K for higher dimensional VAE and GAN models, where the posterior p ( z | x ) is more complex and is more easily distinguishable from the variational q θ ( z | x ) .
These highlight the inherent exponential sample complexity of the IWAE lower bound on MI
B. B IAS R EDUCTION IN K FOR IWAE L OWER B OUND ON log p ( x ) / U PPER B OUND ON MI
We have seen in App.
B.2 Prop.
B.2 that the improvement of EUBO IWAE ( x ; q θ , K ) and I IWAE L ( x ; q θ , K ) over the single-sample EUBO ( x ; q θ ) and I BA L ( x ; q θ ) is limited by log K .
This suggests that exponential sample complexity in the gap of the EUBO , K ∝ exp { D KL [ p ( z | x ) (cid:107) q θ ( z | x )] } , is required to obtain a tight lower bound.
The quantity D KL [ U ( s ) (cid:107) p IWAETGT ( s | z (1: K ) , x )] , which measures the improvement of the IWAE lower bound on log p ( x ) over the ELBO , is not explicitly limited by log K .
However, Chatterjee et al.
(2018) suggest that the same exponential sample complexity, K ∝ exp { D KL [ p ( z | x ) (cid:107) q θ ( z | x )] } , is required for accurate importance sampling estimation with proposal q θ ( z | x ) .
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 ) (cid:107) q θ ( z | x )]
(Song & Ermon, 2019).
This exponential sample complexity for exact estimation of log p ( x ) or I ( x , z ) is usually impractical for complex target distributions and limited variational families, where D KL [ p ( z | x ) (cid:107) q θ ( z | x )] may be large.
This motivates our improved, multi-sample AIS proposals in Sec. 3.2, which achieve more favorable (linear) bias reduction by introducing MCMC transition kernels to bridge between q θ ( z | x ) and p ( z | x ) .
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.
(2019) Sec. 2.5) as simply applying the standard IWAE bounds, using the marginal p ( z ) in place of the variational q θ ( z | x )
E p ( x ) p ( z (1) | x ) K (cid:81) k
=2 p ( z ( k )) (cid:34) log p ( x | z (1) )
1 K K (cid:80) k
=1 p ( x | z ( k ) )
(cid:35) ≤ I ( x ; z ) ≤ E p ( x ) K (cid:81) k
=1 p ( z ( k )) (cid:34) log 1 K K (cid:80) k
=1 p ( x | z ( k ) )
(cid:35) − 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
Cor.
B.3, we obtain an alternative proof that the Structured I NFO NCE lower bound is upper bounded by log K .
Since I BA L ( p ( z )) ≤
I IWAE L ( p ( z ) , K ) ≤ I BA L ( 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 L ( K ) ≤ 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 ( x , z (1: K ) , s ) as
[Equation]
with marginalization over s leading to the mixture in Eq.
The posterior over the index variable s , which infers the ‘positive’ sample drawn from p ( z | x ) given a set of samples
z (1: K ) , corresponds to the normalized importance weights
[Equation]
which can be derived from p GIWAETGT ( s
[Equation]
x
[Equation]
z
[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 critic function T φ .
[Equation]
where q GIWAEPROP ( s | z (1: K ) ,
[Equation]
Similarly to Lawson et al.
(2019), the variational SNIS distribution q GIWAEPROP ( s | z (1: K ) , x ) is approximating the true SNIS distribution p GIWAETGT
( s | x , z (1: K ) ) .
As we show in App.
C.4, the optimal critic function is T φ ( x , z )
= log p ( x , z ) q ( z | x ) + c ( x ) , which recovers the IWAE probabilistic interpretation (see
App. B.1).
Log Importance Ratio To derive bounds on the log partition function, we ﬁrst calculate the log unnormalized density ratio
[Equation]
[Equation]
Taking the expectation of the log unnormalized density ratio under the proposal or target distribution yields a lower or upper bound, respectively, on log p ( x )
[Equation]
As in Sec. 2.1 and App.
A, the gap in the lower and upper bounds can be derived as KL divergences in the extended state space.
Upper Bound on log p ( x ) and Lower Bound on MI To derive an explicit form for the GIWAE upper bound on log p ( x ) , we write
[Equation]
where, in the ﬁrst term of the second line, we note that p GIWAETGT ( z (1: K ) , s | x ) speciﬁes that z ( s ) ∼ p ( z | x ) .
Since s ∼ p ( s ) =
1 K is sampled uniformly, we can assume s = 1 in the second term due to permutation invariance.
Translating this into a lower bound on MI , we consider I ( x ; z ) =
− E p ( x ) [log p ( x )] − H ( x | z ) ≥ − E p ( x ) [ EUBO GIWAE ( q θ , T φ , K )] − H ( x | z ) .
Writing the conditional entropy term over the index s ,
I ( x ; z ) ≥ E p
[Equation]
x
[Equation]
z
[Equation]
[log p
[Equation]
x |
[Equation] [Equation]
which matches Eq.
(7) from the main text.
C. GIWAE U PPER B OUND ON MI D OES N OT P ROVIDE B ENEFIT O VER IWAE
To derive an explicit form for the GIWAE lower bound on log p ( x ) ,
[Equation]
Translating this to an upper bound on MI yields
I
[Equation] [Equation]
Thus, in Eq.
(52), knowledge of the full joint density is required to evaluate both the conditional entropy and log p ( x , z ( s ) )
q ( z ( s ) | x ) terms.
If both p ( z ) and p ( x | z ) are known, then we will show in Cor. 2.
below that the optimal critic or negative energy function in GIWAE yields the true importance weights, and the resulting MI or log p ( x ) bounds matches the IWAE bounds.
We thus conclude that the GIWAE upper bound on MI and lower bound on log p ( x ) does not provide any beneﬁt over IWAE in practice.
However, ELBO GIWAE ( x ; q θ , T φ , K ) is still useful for analysis, as our proof of Lemma C.1 below allows us characterize the gap between ELBO IWAE ( x ; q θ , K ) and
ELBO ( x ; q θ ) in Prop.
B.2.
C. ELBO AND EUBO ARE S PECIAL C ASES OF GIWAE L OG P ARTITION F UNCTION B OUNDS
Lemma C.1.
The single-sample ELBO and EUBO are special cases of GIWAE , with
[Equation]
In both cases, the SNIS sampling distribution (Eq. (43) ) is uniform q GIWAEPROP
( z (1: K ) , s | x ) = 1 K
= U ( s ) .
Proof.
We consider the GIWAE probabilistic interpretation (App. C.1) for T φ 0 = const.
We refer to this extended state space proposal as q BAPROP , since it leads to I BA L ( q θ ) and
I BA U ( q θ ) bounds on MI .
[Equation]
where q BAPROP ( s | z (1: K ) , x )
=
[Equation]
φ
[Equation]
x [Equation] Note that the SNIS sampling distribution q BAPROP ( s | z (1: K ) , x ) will be uniform, which matches p ( s )
= 1 K in the GIWAE extended state space target distribution
[Equation]
Now, taking the log unnormalized importance weights, we obtain log
[Equation]
Taking expectations with respect to q BAPROP ( z (1: K ) , s | x ) or p GIWAETGT ( z (1: K ) , s | x ) leads to ELBO ( x ; q θ ) and EUBO ( x ; q θ ) respectively, as in Sec. 2-2.2.
C. P ROOF OF R ELATIONSHIP BETWEEN IWAE AND GIWAE P ROBABILISTIC I NTERPRETATIONS (P ROP . 2.1)
We ﬁrst prove a lemma which relates both the GIWAE lower and upper bounds on log p ( x ) to the respective IWAE bounds.
From this lemma, Prop.
B.2 follows directly and relates the ELBO and EUBO (which are special cases of ELBO GIWAE and EUBO GIWAE ) to ELBO IWAE and EUBO IWAE .
Lemma C.2.
We can characterize the difference between IWAE and GIWAE bounds on log p ( x ) using
KL divergences between their respective SNIS distributions.
[Equation]
Proof.
Recall from Sec.
2.1 or App.
A that the gap of the lower bound ELBO GIWAE ( x ; q θ , T φ , K ) is
D KL [ q GIWAEPROP ( z (1: K ) , s | x ) (cid:107) p GIWAETGT ( z (1: K ) , s, x )] , while the gap of the upper bound
EUBO GIWAE ( x ; q θ , T φ , K ) is D KL [ p GIWAETGT ( z (1: K ) , s | x ) (cid:107) q GIWAEPROP ( z (1: K ) , s | x )] .
We will expand these KL divergences to reveal the relationship between the GIWAE bounds and IWAE bounds.
First, recall from Eq.
(41) that the posterior over the index variable s , or target SNIS distribution, is p GIWAE TGT ( s | z (1: K ) , x ) = p GIWAETGT ( z (1: K ) , s, x ) K (cid:80)
[Equation]
=
[Equation]
= [Equation] ( s ) | x ) . The joint distribution then factorizes as p GIWAETGT ( z (1: K ) , s, x ) =
p GIWAETGT ( z (1: K ) , x ) · p GIWAETGT ( s | z
(1: K ) , x ) .
ELBO Case :
Using this factorization of p GIWAETGT ( z (1: K ) , s, x ) , we can rewrite the gap of
ELBO GIWAE ( x ; q θ , T φ , K ) as follows
[Equation]
where we can recognize the ﬁrst term as the gap in ELBO IWAE ( x ; q θ , K ) .
Noting that ELBO IWAE ( x ; q θ , K ) − ELBO GIWAE ( x ; q θ , T φ , K )
=
D KL [ q GIWAEPROP ( z (1: K ) , s | x ) (cid:107) p GIWAETGT ( z (1: K ) , s | x )] −
D KL [ q IWAEPROP ( z (1: K ) | x ) (cid:107) p IWAETGT ( z (1: K ) | x )] , we obtain Eq.
(57), as desired
[Equation]
EUBO Case :
For ELBO GIWAE ( x ; q θ , T φ , K ) , the derivations follow in a similar fashion to Eq.
Eq.
(60), but using the reverse KL divergence and expectations under p GIWAETGT ( z (1: K ) , s | x ) .
Translating Lemma C.2 to a statement relating IWAE and GIWAE bounds on MI , we obtain the following proposition.
Proposition 2.1 (Improvement of IWAE over GIWAE ) .
For a given q θ ( z | x ) and any T φ ( x , z ) ,
I IWAE
[Equation]
Proof.
The result follows directly from Lemma C.2.
First, note that p IWAETGT ( x , z (1: K ) , s ) =
p GIWAETGT ( x , z (1: K ) , s ) .
Then, using the upper bounds on log p ( x ) and taking outer expectations with respect to p ( x ) , we have
[Equation]
C. P ROOF OF GIWAE O PTIMAL C RITIC F UNCTION AND L OGARITHMIC I MPROVEMENT (C OR . 2.2 AND C OR . 2.3)
We now prove Cor. 2.2 and Cor.
2.3 from the main text, with a corollary stating the results for the special case of I NFO NCE in App.
C.6.
Corollary 2.2.
For a given q θ ( z | x ) and K > 1 , the optimal critic function is the true log importance weight 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]
( q
[Equation]
[Equation]
[Equation]
=
[Equation]
Proof.
Using Prop. 2.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 iff
[Equation]
This condition also ensures that the overall GIWAE proposal p GIWAETGT ( z (1: K ) , s | x ) (Eq.
and target q GIWAEPROP ( z (1: K ) , s | x ) (Eq.
(42)) distributions match.
We will show that any T ( x , z ) which satisﬁes
Eq.
(63) has the form
[Equation]
Let f ( x , z )
= log p ( x , z ) q θ ( z | x ) + g ( x , z ) , which represents an arbitrary choice of critic function.
We will show that g ( x , z ) must be constant with respect to z
[Equation]
where g ( x , z ) =
c ( x ) is required in order to ensure that g ( x , z ( s ) ) =
g ( x , z ( k ) ) for arbitrary choices of z samples.
This form for the optimal critic function T ∗ ( x , z ) in Eq.
(64) implies that learning T φ ( x , z ) in GIWAE becomes unnecessary when the density ratio p ( x , z ) q θ ( z | x ) is available in closed form, as is assumed in the IWAE bound.
For this choice of T ∗ ( x , z ) , the value of the GIWAE objective matches the IWAE lower bound on MI .
[Equation]

[Equation]
The second term in Eq.
(65) is exactly the negative of the IWAE upper bound on log p ( x ) in Eq.
Combined with the entropy − H ( x | z ) , we obtain the I IWAE L ( q θ , K ) lower bound on MI as desired.
Corollary 2.3.
Suppose the critic function T φ ( x , z ) is parameterized by φ , and
∃ φ 0 s.t. ∀
( x , z ) , T φ 0 ( x , z ) = const.
For a given q θ ( z | x ) , let T φ ∗ ( x , z ) denote the critic function that maximizes the GIWAE lower bound.
Using Cor.
2.2, we have
[Equation]
Proof.
We begin by showing that I BA L ( q θ ) ≤ I GIWAE L ( q θ , T φ ∗ , K ) .
This follows from the assumption that there exists φ 0 in the parameter space of the neural network T φ such that T φ 0 = const.
With this φ 0 , we would have I BA L ( q θ ) =
I GIWAE L ( q θ , T φ 0 , K ) as in Lemma C.1.
Thus, the optimal φ ∗ in the parameter space can only improve upon the BA bound.
Next, for any given φ (including φ ∗ ), we have I GIWAE L ( q θ , T φ , K ) ≤ I IWAE L ( q θ , K ) , since IWAE uses the (unconstrained) optimal critic function T ∗ = log p ( x , z ) q θ ( z | x ) + c ( x ) (Cor. 2.2).
The ﬁnal inequality follows from Prop.
B.2, which shows that I IWAE L ( q θ , K ) improves by at most log K over
I BA L ( q θ ) .
These relationships are visualized in Fig.
3b.
C. P ROPERTIES OF I NFO NCE
Corollary C.3.
Using the prior q θ ( z | x )
= p ( z ) as in I NFO NCE , (a)
For the optimal critic function, Cor.
2.2 implies T ∗ ( x , z )
= log p ( x | z )
+ c ( x ) , and
[Equation]
(b) For an arbitrary critic function T φ ( x , z ) , Cor.
2.3 translates to
[Equation]
For I NFO NCE, note that using the prior as the proposal does allow the critic to admit an efﬁcient bi-linear implementation T φ ( x , z ) =
f φ x ( x ) T f φ z ( z ) , which requires only N + K forward passes instead of NK for GIWAE , where N is the batch size and K is the total number of positive and negative samples.
D S INGLE -S AMPLE AIS D. P ROOF OF P ROP . 3.1 (C OMPLEXITY IN T FOR S INGLE -S AMPLE AIS)
In this section, we prove Prop. 3.1, which relates the sum of the gaps in the single-sample AIS upper and lower bounds to the symmetrized KL divergence between the endpoint distributions.
We extend this result to IM - AIS in Cor.
E.1 and CR - AIS in Cor.
H.3.
Proposition 3.1 (Complexity in T ) .
Assuming perfect transitions and a geometric annealing path with linearly-spaced { β t } Tt =1 , the sum of the gaps in the AIS sandwich bounds on MI , I AIS U ( π 0 , T ) −
I AIS L ( π 0 , T ) , reduces linearly with increasing T .
In particular, we will show that
[Equation]
where the left hand side also corresponds to the sum of the gaps D KL [ q AISPROP ( z 0: T | x ) (cid:107) p AISTGT ( z 0: T | x )]
+
D KL [ p AISTGT ( z 0: T | x ) (cid:107) q AISPROP ( z 0: T | x )] .
This result translates to mutual information bounds as in Sec. 1.1.
In contrast to Grosse et al.
Thm. 1, our linear bias reduction result holds for ﬁnite T .
Proof.
For linear scheduling and perfect transitions
[Equation]
we simplify the difference in the single-sample upper and lower bounds as δ
[Equation]
where in (2) , we use the linear annealing schedule β t − β t − 1 = 1 T ∀ t and note that intermediate terms cancel in telescoping fashion.
In (1) , we have used the assumption of perfect transitions ( PT ), which is common in analysis of AIS (Neal, 2001; Grosse et al., 2013).
In this case, the AIS proposal and target distributions have the following factorial form
[Equation]
In other words, for 1 ≤ t ≤ T , perfect transitions results in independent, exact samples from z t ∼ π β t − 1 ( z ) in the forward direction, and z t ∼ π β t ( z ) in the reverse direction.
Using the factorized structure of Eq.
(70) and Eq.
(71), the expectations over the extended state space simplify to a sum of expectations at each z t .
The above proves the proposition for the case of single sample AIS , but should also hold for our tighter multi-sample AIS bounds.
We extend this result to Independent Multi-Sample AIS in App.
E. below, and extend the result to Coupled Reverse Multi-Sample AIS in App.
H.4.
E I NDEPENDENT M ULTI -S AMPLE AIS E. P ROBABILISTIC I NTERPRETATION AND B OUNDS
Our Independent Multi-Sample AIS ( IM - AIS ) are identical to the standard IWAE bounds in App.
B, but using AIS forward q AISPROP ( z 0: T | x ) and backward p AISTGT ( z 0: T | x ) chains of length T instead of the endpoint distributions only q θ ( z | x ) and p ( z | x ) .
In particular, we construct an extended state space proposal by running K independent AIS forward chains z ( k ) 0: T ∼ q AISPROP in parallel.
As in the IWAE upper bound (Eq.
(5)), the extended state space target involves selecting a single index s uniformly at random, and running a backward AIS chain z ( s ) 0: T ∼ p AISTGT 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.
Expanding the log unnormalized density ratio,
[Equation]
which is similar to the IWAE ratio but involves AIS chains.
Taking the expectation under the proposal and target as in App.
A recovers the lower and upper bounds in Eq.
The gap in the lower bound is D KL [ q IM - AIS PROP ( z (1: K ) 0: T | x ) (cid:107) p IM - AIS TGT ( z (1: K ) 0: T | x ))] and the gap in the upper bound is
D KL [ p IM - AIS TGT ( z (1: K ) 0: T | x ))
(cid:107) q IM - AIS PROP ( z (1: K ) {+
 :+} T | x )] .
E. P ROOF OF L INEAR B IAS R EDUCTION IN T FOR IM-AIS
Corollary E.1 (Complexity in T for Independent Multi-Sample AIS Bound) .
Assuming perfect transitions and a geometric annealing path with linearly-spaced { β t } Tt =0 , the sum of the gaps in the
Independent Multi-Sample AIS sandwich bounds on MI ,
I IM - AIS U
( π 0 , T ) − I IM - AIS L ( π 0 , T ) , reduces linearly with increasing T .
Proof.
Using identical proof techniques as for IWAE (Burda et al.
; Sobolev & Vetrov (2019),
Prop. B.1), we can show that our Independent Multi-Sample AIS bounds ELBO IM-AIS ( x ; π 0 , T, K − 1) ≤ ELBO IM-AIS ( x ; π 0 , T, K ) and EUBO IM-AIS ( x ; π 0 , T, K )
≤ EUBO IM-AIS ( x ; π 0 , T, K − 1) improve with increasing K .
Thus, the bias of our multi-sample bounds is less than the bias of the single-sample bounds, so the inequality in Eq.
(68) and linear bias reduction in Prop. 3.1 also hold for Independent Multi-Sample AIS .
We characterize this improvement in Prop.
E.2 below.
E. P ROOF OF L OGARITHMIC I MPROVEMENT OF IM-AIS EUBO
Proposition E.2 (Improvement of Independent Multi-Sample AIS over Single-Sample AIS ) .
Let p IM - AIS TGT ( s | x , z (1: K ) 0: T ) =
[Equation]
AISTGT ( x , z ( s
[Equation]
0: T
[Equation] [Equation]
( z ( k )
[Equation]
T | x ) denote the normalized importance weights over AIS chains, and let U ( s ) indicate the uniform distribution over K discrete val-
ues.
Then, we can characterize the improvement of the Independent Multi-Sample AIS bounds on log p ( x ) , ELBO IM-AIS ( x ; π 0 , T, K ) and EUBO IM-AIS ( x ; π 0 , T, K ) , over the single-sample AIS bounds
ELBO AIS ( x ; π 0 , T ) and EUBO AIS ( x ; π 0 , T ) using KL divergences, as follows
[Equation]
[Table]
Proof.
The result follows directly from App
[Table]
C.4 Prop. 2.1 by viewing Independent Multi-Sample
AIS as IWAE with an AIS proposal as in App.
E.
[Table]
F R EVERSE IWAE
In this section [Table] we propose Reverse IWAE ( RIWAE ), which is an impractical alternative to standard
IWAE .
However, we use this as the basis for our Independent Reverse (App. G) and Coupled Reverse
Multi-Sample AIS bounds (Sec. 3.3 and App.
H).
F. P ROBABILISTIC I NTERPRETATION AND B OUNDS
Similarly to simple reverse importance sampling, 1 p ( x ) =
E p ( z | x ) (cid:2) q θ ( z | x ) p ( x , z ) (cid:3) , we consider K independent posterior samples in an extended state space target distribution.
The proposal distribution for our importance sampling scheme is a mixture of K − 1 posterior distributions and one variational distribution,
[Equation]
Note that we have normalization constants of (cid:82) p RIWAE TGT ( x , z (1: K ) )
d z (1: K ) =
p ( x ) K and (cid:82) q RIWAE PROP ( x , z (1: K ) )
d z (1: K ) =
p ( x ) K − 1 since the transition kernels of the reverse chains do not change the normalization.
We visualize this sampling scheme in Fig.
Similarly to Eq.
(25), the log unnormalized density ratio simpliﬁes to
[Equation]
Taking the expectation 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 require more than one true posterior sample from p ( z | x ) .
However, we use them in Sec.
3.3 and App.
H to derive practical multi-sample reverse AIS bounds.
F.
I MPROVEMENT OF RIWAE OVER ELBO AND EUBO
Proposition F.1 (Improvement of Reverse IWAE over ELBO and EUBO ) .
Let q RIWAEPROP
( s | x , z (1: K ) )
=
q θ ( z ( s ) | x ) p ( x , z (
[Equation]
| x
[Equation]
/ K
[Equation]
k = q θ ( z ( k ) | x ) p ( x , z ( k ) | x ) denote the normalized reverse importance sampling weights, or the posterior over the index variable in Eq.
Let U ( s ) indicate the uniform distribution over K discrete values.
Then, we can characterize the improvement of the Reverse IWAE bounds on log p ( x ) ,
ELBO RIWAE ( x ; q θ , K ) and EUBO RIWAE ( x ; q θ , K ) , over the single-sample bounds ELBO ( x ; q θ ) and
EUBO ( x ; q θ ) using KL divergences, as follows
[Equation]
Proof.
The proof follows similarly as Prop.
B.2 or Lemma C.2.
G I NDEPENDENT R EVERSE M ULTI -S AMPLE AIS G. P ROBABILISTIC I NTERPRETATION AND B OUNDS
In similar fashion to Reverse IWAE
[Equation] we now use K independent reverse AIS chains to form an extended state space target distribution p IR - AIS TGT ( z (1: K ) 0: T , x ) [Equation] and a mixture proposal distribution q IR - AIS PROP ( z (1: K ) 0: T , x ) which includes a single forward AIS chain (see Fig. 2)
[Equation]
Similarly to Eq.
[Equation]
the log unnormalized density ratio becomes
[Equation]
Taking expectations of the log unnormalized density ratio under the proposal and target
[Equation]
respectively, yield lower and upper bounds on log p ( x )
[Equation]
However, Independent Reverse Multi-Sample AIS may be impractical in common settings, since it is infeasible to have access to more than one true posterior sample.
G. P ROOF OF L OGARITHMIC I MPROVEMENT IN K FOR I NDEPENDENT R EVERSE AIS
Proposition G.1 (Improvement of Independent Reverse Multi-Sample AIS over Single-Sample AIS ) .
Let q IR -
[Equation]
PROP
[Equation] [Equation] T ) denote the normalized reverse importance sampling weights over AIS chains
[Equation]
and let U ( s ) indicate the uniform distribution over K discrete values.
Then, we can characterize the improvement of the Independent Reverse Multi-Sample AIS bounds on log p ( x ) , ELBO IR - AIS ( x ; π 0 , K, T ) and EUBO IR - AIS ( x ; π 0 , K, T ) , over the single-sample
ELBO AIS ( x ; π 0 , T ) and ELBO AIS ( x ; π 0 , T ) using KL divergences, as follows
[Equation]
[Table]
Proof.
The proof follows similarly as Prop.
B.2 or Lemma C.2.
H C OUPLED R EVERSE M ULTI -S AMPLE AIS H. 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 ) , which makes the bound useful in practical situations.
We denote the remaining transitions as p AISTGT ( z 0: T − 1 | z T , x ) , since they are identical to standard AIS in Eq.
Thus, the CR - AIS extended state space target distribution is
[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 chain,
[Equation]
See Fig.
2 or Fig.
4 for a graphical model description
[Equation]
We can construct log-partition bounds using the log unnormalized density ratio
[Equation]
Taking the expected log ratio under the proposal and target yields lower and upper bounds on log p ( x )
[Equation]
k = q AISPROP ( z (
[Table] [Table]
1 , z T | x ) p
[Table]
( x , z
[Table] [Table]
[Table]
) (cid:21)
[Table]
p ( x )
EUBO CR - AIS ( x ; π 0 , K, T ) :
= − E
[Table]
T
[Table]
π T
[Table]
z T | x )
z
[Table]
∼ p
[Table]
T , x
[Table]
H. P ROOF THAT CR-AIS B OUNDS T IGHTEN WITH I NCREASING K
In this section, we prove that Coupled Reverse Multi-Sample AIS bounds get tighter with increasing
K .
Our proof provides an alternative perspective to Burda et al.
; Sobolev & Vetrov (2019) for showing the monotonic improvement of IWAE or Independent Multi-Sample AIS with K .
We will also characterize the improvement of CR - AIS bounds over single-sample AIS bounds in Prop.
H.2, as a direct consequence of this lemma.
Lemma H.1 (CR-AIS Bounds Tighten with Increasing K ) .
Coupled Reverse Multi-Sample AIS bounds get tighter with increasing number of samples K .
In other words, for any K > 1 ,
[Equation]
[Equation] [Equation]
Proof.
Our proof will proceed by introducing an additional set of M index variables s 1: M into the K -sample probabilistic interpretation of CR - AIS in Eq.
(83)-(84), with M < K .
We will show that the KL divergence in this joint state space (including s 1: M ) is equal to the gap of the M -sample ELBO CR - AIS ( x ; π 0 , T, M ) or EUBO CR - AIS ( x ; π 0 , T, M ) .
We then show that marginalizing over s 1: M yields the gap of the K -sample ELBO CR - AIS ( x ; π 0 , T, K ) or EUBO CR - AIS ( x ; π 0 , T, K ) .
Since marginalization cannot increase the KL divergence, we will have shown that for any M < K , D KL [ q CR - AIS PROP (
[Equation]
[Equation] [Equation] T
[Equation] [Equation]
, z
[Equation]
|
[Equation]
]] ≤
D
[Equation]
, z T | x [Equation] ]] and thus
ELBO CR - AIS ( x ; π 0 , T, M ) ≤
ELBO CR - AIS ( x ; π 0 , T, K ) .
Identical reasoning holds for the EUBO CR - AIS .
Sub-Sampling Probabilistic Interpretation Let U ( s 1 , ..., s M ) indicate the probability of drawing M < K sample indices uniformly without replacement (i.e. each s m is distinct).
In the CR - AIS target distribution, there is no distinction between the indices { s 1: M } and k (cid:54)∈ { s 1: M } after drawing z
T ∼ π T ( x , z ) .
We can write
[Equation]
which, after marginalization over s 1: M , clearly matches p CR - AIS TGT ( x ,
z
[Equation] [Equation]
K
[Equation]
1 , z T ) .
We also draw s 1: M ∼ U ( s 1 , . . .
, s M ) for the extended state space proposal.
Next, we select an index m uniformly at random from { 1 , . . .
, M } , which is used to to specify which chain z ( s m ) 0: T − 1 is run in the forward direction to obtain z T , as in Eq.
After marginalizing over m , we obtain the following mixture proposal distribution
[Equation]
We will consider marginalizing over s 1: M below, but ﬁrst write the KL divergence in the extended state space which includes s 1: M .
For example, the forward KL divergence matches the gap of the
M -
[Equation]
x ; [Equation] 0 , T, M ) ,
[Equation]
(cid:24)(cid:24)(cid:24)
[Equation]
[Table]
 =
[Equation]
(92) which matches the M -sample probabilistic interpretation of CR - AIS from App.
H.1.
Identical reasoning holds for the case of EUBO CR - AIS ( x ; π
, T, K ) ≤ EUBO CR - AIS ( x ; π 0 , T, M ) using the reverse KL divergence.
Marginalization over s 1: M
We have already seen from Eq.
(89) that the marginal
[Equation]
TGT ( x , s 1: M , z (1: K ) 0: T
[Equation]
1 , z T ) matches the K -sample target distribution p CR - AIS TGT ( x , z (1: K ) 0: T − 1 , z T ) .
We would now like to marginalize over s 1: M in q CR - AIS PROP ( s 1: M , z (1: K ) 0: T − 1 , z T | x ) .
Combining the two product terms in Eq.
[Equation]
[Equation] M [Equation] where in (1) , we can write the marginal U ( s m ) since the terms inside expectation do not explicitly depend on the other indices in s 1: M .
In (2) , we use the fact that the marginal over any s m is uniform.
Thus, we have shown that the marginal distributions match the standard K -sample target and proposal distributions of CR - AIS , with
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
and similar reasoning for the EUBO using the reverse KL divergence.
Conclusion of Proof Since marginalization can not increase the KL divergence, we have
[Figure]
[Equation]
[Equation]
This shows that ELBO CR - AIS ( x ; π 0 , T, K ) is tighter than ELBO CR - AIS ( x ; π 0 , T, M ) , with
ELBO CR - AIS ( x ; π 0 , T, K ) ≥ ELBO CR - AIS ( x ; π 0 , T, M ) .
Identical reasoning holds for the case of
EUBO CR - AIS ( x ; π 0 , T, K ) ≤ EUBO CR - AIS ( x ; π 0 , T, M ) using the reverse KL divergence.
Choosing
M = K − 1 proves that the bounds tighten as K increases, as desired.
Characterizing the Improvement of K -Sample over M -Sample CR-AIS Note that we can explicitly write the gap of the inequality Eq.
(98) using the chain rule for joint probability, for example q CR
[Equation]
The gap in Eq.
(98) also corresponds to the gap between ELBO CR - AIS ( x ; π 0 , T, K ) and
ELBO CR - AIS ( x ; π 0 , T, M ) , so we can write
[Equation]
Below, we analyze the posterior over the index variable for M = 1 as a special case.
This allows us to characterize the gap between ELBO CR - AIS ( x ; π 0 , T, K ) and the single-sample ELBO AIS ( x ; π 0 , T ) =
ELBO CR - AIS ( x ; π 0 , T, K = 1) .
H. P ROOF OF L OGARITHMIC I MPROVEMENT IN K FOR CR-AIS ELBO
Proposition H.2 (Improvement of Coupled Reverse Multi-Sample AIS over Single-Sample AIS ) .
Let
[Equation]
denote the normalized importance weights over the AIS chains used in Coupled Reverse Multi-
Sample
AIS , and let U ( s ) indicate the uniform distribution over K discrete values.
Then, we can characterize the improvement of the Coupled Reverse Multi-Sample AIS bounds on log p ( x ) , ELBO CR - AIS ( x ; π 0 , T, K ) and EUBO CR - AIS ( x ; π 0 , T, K ) , over the single-sample AIS bounds
ELBO AIS ( x ; π 0 , T ) and EUBO AIS ( x ; π 0 , T ) using KL divergences, as follows
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table] [Equation] [Table]
[Equation]
[Table] [Equation] [Table]
[Equation]
[Table]
Proof.
Using M = 1 in Eq.
(99) above
[Equation]
we obtain
[Equation]
From Eq.
(90), we can write the posterior over the index variable as
[Equation]
For the target distribution in Eq.
(89), the posterior p CR - AIS TGT ( s | z (1: K ) 0: T − 1 , z T , x )
=
U ( s ) is uniform.
Thus, we can characterize the improvement of the K -sample CR - AIS ELBO over its single-chain AIS
[Equation]
The proof follows identically for the EUBO .
For the ELBO , the KL divergence with the uniform distribution in the second argument is bounded above by log K .
For the EUBO , the improvement of CR - AIS is limited by the gap in the single-chain
EUBO AIS ( x ; π 0 , T ) , which corresponds to D KL [ p AIS TGT ( z [Equation] T | x ) (cid:107) q AIS PROP ( z 0: T | x )] .
H. P ROOF OF L INEAR B IAS R EDUCTION IN T FOR CR-AIS
Corollary H.
(Complexity in T for Coupled Reverse Multi-Sample AIS Bound) .
Assuming perfect transitions and a geometric annealing path with linearly-spaced { β t } Tt =0 , the sum of the gaps in the
Coupled Reverse Multi-Sample AIS sandwich bounds on MI ,
I CR - AIS U
( π 0 , K, T ) − I CR - AIS L ( π 0 , K, T ) , reduces linearly with increasing T .
Proof.
Since we have shown in Prop.
H.2 that ELBO CR - AIS ( π 0 , K, T ) and EUBO CR - AIS ( π 0 , K, T ) are tighter than ELBO AIS ( π 0 , T ) and EUBO AIS ( π 0 , T ) , the bias in the CR - AIS sandwich bounds
EUBO CR - AIS ( π 0 , K, T ) − ELBO CR - AIS ( π 0 , K, T ) will be less than that of single-sample AIS .
Thus, inequality in Eq.
(68) and linear bias reduction in Prop. 3.1 (see App.
D.1) also hold for CR - AIS , which implies linear bias reduction under perfect transitions and linear scheduling.
I C OMPARISON OF M ULTI -S AMPLE AIS B OUNDS
In Fig.
5, we compare the performance of our various Multi-Sample AIS bounds for MI estimation of a Linear VAE with 10 latent variables and random weights, and VAE and GAN models with 20 latent variables trained on MNIST .
To obtain an upper bound on MI , we recommend using the Independent Multi-Sample AIS ELBO for log partition function estimation. This corresponds to the forward direction of BDMC and achieves the best performance in all cases.
This upper bound uses independent samples and is not limited to log K improvement, in contrast to the Coupled Reverse Multi-Sample AIS upper bound on MI (Prop. H.2).
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 lower bound on MI comes from K − 1 independent negative forward chains which, by App. E.
Prop. E.2, 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 two total AIS distributions, which corresponds to simple importance sampling with π 0 ( z ) =
p ( z ) , the Independent Multi-Sample AIS lower bound on MI reduces to Structured I NFO -NCE and saturates to log K .
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 lower bound on MI is induced by MCMC transitions in K coupled backward chains.
While this does not formally limit the improvement over single-sample AIS , we see in Fig.
5 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 bounds on MI may vary based on the scale of the true MI and the amount of computation available.
[Figure] [Figure]
T for
MI estimation in deep generative models on MNIST .
J G ENERALIZED M UTUAL I NFORMATION N EURAL E STIMATION ( GMINE )
In this section, we provide probabilistic interpretations for the Mutual Information Neural Estimation ( 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 , Generalized MINE uses a base variational distribution q θ ( z | x ) to tighten the MINE - DV or MINE - F bound and can be evaluated when p ( z ) is available.
See Fig.
6 for a summary of Generalized MINE bounds and their relationships.
We discuss probabilistic interpretations in this section, and provide complementary interpretations in terms of conjugate dual representations of the KL divergence in App.
K.
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.
M.
J. P ROBABILISTIC I NTERPRETATION OF IBAL For an energy-based posterior approximation
[Equation]
where Z
[Equation]
[Equation]
104) we consider the BA lower bound on MI ,
[Equation]
x )]
[Equation]
[ T (
[Equation]
z )]
[Equation]
p ( x
[Equation]
z ) [log Z θ
[Equation]
[Equation]
[Equation]
z )
[Equation]
p ( x ) [
[Equation]
Z
[Equation]
( x )]
[Equation]
E p ( x , z ) (cid:34) log e T φ ( x , z )
E q θ ( z | x ) (cid:2) e T φ ( x , z ) (cid:3) (cid:35) (cid:124) (cid:123)(cid:122) (cid:125) contrastive term =: IBAL ( q θ , T φ ) , where the gap in the mutual information bound is E p ( x ) (cid:2) D KL [ p ( z | x ) (cid:107) π θ,φ ( z | x )] (cid:3) .
Note that
IBAL ( q θ , T φ ) includes I BA L ( q θ ) as one of its terms, where we refer to q θ ( z | x ) as the base variational distribution.
We visualize relationships between various energy-based bounds in Fig.
3b and
Fig.
Proposition 4.1.
For a given q θ ( z | x ) , the optimal IBAL critic function equals the log importance weights up to a constant T ∗ ( x , z )
= log p ( x , z ) q θ ( z | x ) + c ( x ) .
For this T ∗ , we have IBAL ( q θ , T ∗ ) =
I ( x ; z ) .
See App.
L.1 for the proof.
J. P ROBABILISTIC I NTERPRETATION OF G ENERALIZED MINE - DV
We can interpret the MINE - DV bound of Belghazi et al.
as arising from an energy-based variational approximation π θ,φ ( x , z ) of the full joint distribution p ( x , z ) as follows
[Equation]
where
[Equation]
q θ ( z | x ) is a base variational distribution, and T φ is a critic or negative energy function.
Note that for the induced marginal π θ,φ ( x ) :
= (cid:82) π θ,φ ( x , z ) d z (cid:54) =
p ( x ) due to the contribution of the critic function.
Instead, we have
[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 the joint KL divergence D KL [ p ( x , z ) (cid:107) π θ,φ ( x , z )] from I ( x ; z ) , we obtain the Generalized
MINE - DV lower bound on MI
[Equation]
By construction,
Eq.
(108) shows that the gap in
Generalized
MINE - DV is
D KL [ p ( x , z ) (cid:107) 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.
(106) 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 ) (cid:107) 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 critic function T φ 0 ( x , z ) =
const, we can see that π θ,φ 0 ( x , z ) =
p ( x ) q θ ( z | x ) and I GMINE - DV ( q θ , T φ 0 = const ) =
I BA L ( q θ ) for a given q θ ( z | x ) .
Optimal Critic Function For a given q θ ( z | x ) , the optimal critic function of Generalized MINE - DV corresponds to the log importance weight between the target p ( x , z ) and the joint base distribution p ( x ) q θ ( z | x ) , plus a constant
[Equation] [Equation]
For the optimal critic 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 , with IBAL ( q θ , T φ ) ≥ I GMINE - DV ( q θ , T φ ) .
Subtracting the gaps in the bounds in Eq.
(105) and Eq.
[Equation]
We see that the difference between IBAL ( q θ , T φ ) and I GMINE - DV ( q θ , T φ ) corresponds to a marginal KL divergence in the x space, 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
[Equation]
since log E p ( x )
[ Z ( x )] ≥ E p ( x )
[log Z ( x )] .
J. 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 (GKL)
First, we state the deﬁnition of the generalized KL divergence, which takes unnormalized measures as input arguments and corresponds to the limiting behavior of the α -divergence (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 ) (cid:107) s
( z )] =
D KL [ r ( z ) (cid:107) s ( z )] .
Finally, if r ( z ) is normalized and ˜ s ( z ) is unnormalized, one can easily conﬁrm that
[Equation]
Generalized EUBO (GEUBO)
Since the Generalized KL divergence is always nonnegative, we deﬁne a Generalized EUBO by adding the Generalized KL divergence between the normalized true posterior and an unnormalized approximate posterior ˜
π
( z | x ) .
[Equation]
where we deﬁne Z ˜ π ( x ) :
= (cid:82) ˜
π ( z | x ) d z .
In general, we have EUBO ( x ; π ) ≤ GEUBO ( x ; ˜ π ) where π is the normalized distribution of ˜ π , with equality if ˜ π is normalized.
Generalized BA (GBA) 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]
Generalized MINE-F (GMINE-F)
We now consider an unnormalized, energy-based approximation to the true posterior, involving a base variational distribution q θ ( z | x ) and a learned critic function
T φ ( x , z )
[Equation]
Using this unnormalized approximate posterior in the GBA lower bound in Eq.
(121), we obtain the Generalized
MINE - F lower bound
[Equation]
By construction, we can see that the gap in I GMINE - F ( q θ , T φ ) is equal to the Generalized KL divergence E p ( x )
[ D GKL [ p ( z | x ) (cid:107) ˜
π θ,φ ( z | x )]] in Eq.
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 Critic Function
The optimal critic function of Generalized MINE - F is T ∗ ( x , z )
= 1 + log p ( x , z ) p ( x ) q
( z | x ) (Poole et al., 2019).
In this case, we obtain I GMINE - F ( q θ , T ∗ )
=
I [Equation] x ; z ) in Eq.
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 ) .
Using Eq.
(118), we have
[Equation]
[Equation]
and substituting in the deﬁnition of each term
[Equation]
where the ﬁnal term D GKL [ π [Equation] [Equation] z |
[Equation]
˜ π
[Equation] z [Equation] x )] ≥ 0 is nonnegative because it is a generalized KL divergence.
Thus, we have
[Equation]
Alternatively, Poole et al.
(2019) use the inequality log u ≤ u − 1 to show that
[Equation]
We visualize the relationship between IBAL and Generalized MINE - F 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.
(118), but this time using divergences over joint distributions.
[Equation]
gap in GMINE - F
[Equation]
In this case, we have D GKL [ 1 Z p ( x ) q θ ( z | x ) e T φ ( x , z ) (cid:107) p ( x ) q θ ( z | x ) e T φ ( x , z ) − 1 ] ≥ 0 .
Thus, Generalized
MINE - DV is tighter than Generalized MINE - F (see Fig. 6), which generalizes the ﬁnding in Poole et al.
(2019) that standard MINE - DV is tighter than standard MINE - F .
K C ONJUGATE D UALITY I NTERPRETATIONS
In this section, we interpret the energy-based MI lower bounds in MINE - AIS , Generalized MINE - DV ,
Generalized MINE - F , GIWAE , IWAE , and I NFO NCE from the perspective of conjugate duality.
In particular, we highlight that the critic or negative energy function in the above bounds arises as a dual variable in the convex conjugate representation of the KL divergence.
In all cases, the KL divergence of interest corresponds to the gap in the BA lower bound
[Equation]
For q θ ( z | x ) =
p ( z ) , the BA lower bound term is 0 and our derivations correspond to taking dual representation of MI directly, e.g. E p ( x ) [ D KL [ p ( z | x ) (cid:107) p ( z )]] , as in Belghazi et al.
Our conjugate duality interpretations are complementary to our probabilistic interpretations, with either approach equally valid for deriving lower bounds and characterizing their gaps.
K. C ONVEX C ONJUGATE B ACKGROUND
The KL divergence to a ﬁxed reference distribution π 0 is a convex function of the distribution in the ﬁrst argument Ω( · ) :
= D KL [ ·(cid:107) π 0 ( z )] .
In later subsections, we will also consider KL divergences over joint distributions, unnormalized density functions, and extended state spaces.
For a given Ω( · ) , we can deﬁne the conjugate function Ω ∗ ( · ) over a dual variable or function T ( z ) , with the following relationships (Boyd & Vandenberghe, 2004)
[Equation]
where inner product notation indicates (cid:104) π ( z ) , T ( z ) (cid:105) :
= (cid:82) π ( z ) T ( z ) d z .
Solving for the optimizing argument in each of Eq.
(130) and Eq.
(131) yields the following dual correspondences
[Equation]
We will proceed to derive closed form expressions for various special cases of Ω and Ω ∗ below.
For an arbitrary p ( z ) and T ( z ) which are not in dual correspondence according to Eq.
(132), we can use Eq.
(131) to derive a lower bound on Ω( p ( z )) known as Fenchel’s inequality
[Equation]
where the gap in the inequality D Ω( · ) (cid:2) p ( z ) , π T ( z ) (cid:3) is the Bregman divergence generated by Ω and π T ( z ) is the dual variable corresponding to T ( z ) using Eq.
K. C ONJUGATE D UALITY I NTERPRETATION OF IBAL
To obtain an alternative derivation of IBAL ( q θ , T φ ) , we consider the conditional KL divergence from a reference q θ ( z | x ) , which is a convex function of its ﬁrst argument [Equation]
[Equation]
To derive the conjugate function Ω ∗ ( T ) , note that we must restrict the optimization to the simplex π ( z | x ) ∈ ∆ , since the standard KL divergence requires a normalized distribution as input [Equation]
[Equation]
where we have solved for the optimizing argument π T ( z | x ) to obtain the conjugate function in
Eq. Eq.
(132) suggests the following dual correspondence between primal and dual variables
[Equation]
We would like to leverage this duality to estimate the KL divergence Ω( p ( z | x )) =
D KL [ p ( z | x ) (cid:107) q θ ( z | x )] from q θ ( z | x ) to the true posterior p ( z | x ) .
In particular, plugging Eq.
(138) into Eq.
(131) suggests the following variational representation
[Equation]
For a suboptimal T π ( x , z ) , which is in dual correspondence with π T ( z | x ) instead of the desired posterior p ( z | x ) , we can use Eq.
(134) to obtain a lower bound on D KL [ p ( z | x ) (cid:107) q θ ( z | x )] .
To characterize the gap in this inequality, one can conﬁrm that the Bregman divergence generated by the
KL divergence in Eq.
(135) is also the KL divergence D D KL [
[Equation]
] [ p, π ] = D KL [ p [Equation] π ] .
Thus, we have
D
[Equation]
Finally, the IBAL ( q θ , T φ ) uses this variational representation of the gap in the BA lower bound, E p ( x ) [ D KL [ p ( z | x ) (cid:107) q θ ( z | x )]] , to obtain a tighter bound on MI .
In particular, for any learned critic function T ( x , z ) , we can use Eq.
(140) to derive the IBAL and its gap
[Equation]
The optimal critic function T ∗ ( x , z ) provides the maximizing argument in Eq.
(139) and is in dual correspondence with the true posterior p ( z | x ) .
In particular, we have π T ∗ ( z | x ) =
p ( z | x ) , resulting in I ( x ; z ) =
IBAL ( q θ , T ∗ ) .
K C ONJUGATE D UALITY I NTERPRETATION OF G ENERALIZED MINE-DV
To obtain a conjugate duality interpretation of (Generalized) MINE - DV , we consider the dual representation of the KL divergence over joint distributions.
Choosing π 0 ( x , z ) =
p ( x ) q θ ( z | x ) as the reference distribution, the KL divergence is a convex function of the ﬁrst argument
[Equation]
This KL divergence is equivalent to the gap in the BA bound E p ( x ) [ D KL [ p ( z | x ) (cid:107) q θ ( z | x )]] after noting the marginal distribution of both p ( x , z ) and p ( x ) q θ ( z | x ) is p ( x ) .
However, the duality associated with Ω( · )
=
D KL [ ·(cid:107) π 0 ( x , z )] holds for general joint distributions, and we will see that the conjugate duality perspective using this divergence leads to looser bound on MI than in App. K.2.
To evaluate the expression Ω( π ( x , z ))
=
sup T ( x , z ) (cid:104) π ( x , z ) , T ( x , z ) (cid:105) − Ω ∗ ( T ( x , z )) , we need to derive the conjugate function Ω ∗ ( T ( x , z )) .
Similarly to Eq.
(136), we constrain the joint distribution to be normalized and obtain
[Equation]
144) where we have used π 0 ( x , z )
= p ( x ) q θ ( z | x ) .
Note that the expectation over p ( x ) now appears inside the log in Ω ∗ ( T ( x , z )) , compared with the conjugate for the conditional KL divergence in Eq.
Solving for the optimizing arguments in Eq.
(131) and Eq.
(143), we have the following relationship between primal and dual variables,
[Equation]
Finally
[Equation]
we use Eq.
(131) to write the dual representation of joint KL divergence as
[Equation]
As in the previous section and Eq.
[Equation]
a suboptimal T ( x , z ) , which is in dual correspondence with π
T ( z | x ) instead of the desired posterior p ( z | x )
[Equation]
yields a lower bound on D KL [ p ( x , z ) (cid:107) p ( x ) q θ ( z | x )] , with the gap equal to a KL divergence
[Equation]
The Generalized MINE - DV bound I GMINE - DV ( q θ , T φ ) uses this variational representation of the gap in the BA lower bound to obtain a tighter bound on MI .
For a learned critic T ( x , z ) , we can use Eq.
(147) to write
[Equation]
z ) (cid:21) + E p ( x ) p ( z | x ) [ T ( x , z )] − log E qθ ( z | x ) (cid:104) e T ( x , z ) (cid:105) (cid:124) (cid:123)(cid:122) (cid:125)
I GMINE - DV ( qθ,Tφ )
+ D KL [ p ( x ) p ( z | x ) (cid:107) π T ( x , z )]
As noted in App.
J.2, the Generalized MINE - DV bound is looser than the IBAL .
K. C ONJUGATE D UALITY I NTERPRETATION OF G ENERALIZED MINE-F
Nguyen et al.
(2010) consider the conjugate duality associated with the family of f -divergences, of which the KL divergence is a special case.
We will show that this dual representation corresponds to taking the conjugate function of the Generalized KL divergence, which can take unnormalized densities as input.
See App.
J.3 for deﬁnition of the Generalized KL divergence, which is convex in either argument since it differs from the standard KL divergence by only a linear term..
For Ω( · ) = D GKL [ ·(cid:107) ˜ q θ ( z | x )] , we write the dual variable using the notation T (cid:48)
( x , z ) and consider
[Equation]
Note that we do not explicitly include a Lagrange multiplier to enforce restriction to normalized distributions.
Solving for the optimizing argument or writing the dual correspondence in Eq.
(132), we obtain
[Equation]
Plugging this ˜ π T ( z | x ) back into Eq.
(149) yields the conjugate function
[Equation]
which leads to a dual representation of the Generalized KL divergence that matches Nguyen et al.
(2010) after plugging into Eq.
[Equation]
We now use the reparameterization T (cid:48)
( x , z ) =
T ( x , z ) − 1 .
Assuming a normalized ˜ q ( z | x ) =
q ( z | x ) and ˜ p ( z | x ) = p ( z | x ) , and noting that D GKL [ p ( z | x ) (cid:107) q θ ( z | x )]
= D KL [ p ( z | x ) (cid:107) q θ ( z | x )] for normalized distributions, we obtain [Equation]
[Equation]
which matches dual representation of the KL divergence found in Belghazi et al.
; Nowozin et al.
See Ruderman et al.
For a suboptimal T π ( x , z ) which is in dual correspondence with π T ( z | x ) instead of the desired posterior p ( z | x ) , we can use Eq.
(134) to obtain a lower bound on D GKL [ p ( z | x ) (cid:107) q θ ( z | x )] =
D KL [ p ( z | x ) (cid:107) q θ ( z | x )] .
To characterize the gap of this lower bound, note that the Bregman divergence generated by the D
GKL divergence is also the D GKL divergence D D GKL [ ·(cid:107) [Equation] [Equation] p
[Equation]
π ]
= D
[Equation]
p
[Equation]
˜ π ] .
Thus, using Eq.
(133), we have
[Equation]
Finally, we obtain the Generalized MINE - F bound by evaluating the dual representation of the
Generalized KL divergence for normalized p ( z | x ) and q θ ( z | x ) .
In particular, I GMINE - F ( q θ , T φ ) uses the critic function T φ to tighten the gap in I BA L ( q θ ) via the dual optimization in Eq. I
[Equation]
[Table]
The optimal critic function is the dual variable of the true posterior p ( z | x ) , which can be found using
Eq.
(150) as T ( x , z ) = 1 + log ˜ p
( z | x ) ˜ q
θ ( z | x ) .
With this optimal critic, the Generalized MINE - F bound is tight.
K. C ONJUGATE D UALITY I NTERPRETATION OF GIWAE, IWAE, AND I NFO NCE
In this section, we use conjugate duality to derive the GIWAE , IWAE , and I NFO -NCE bounds on mutual information and characterize their gaps.
Our approach extends that of Poole et al.
(2019), where the MINE - F dual representation (App. K.4) was used to derive I NFO -NCE.
We provide an alternative derivation using the dual representation associated with the conditional KL divergence and
IBAL in App. K.2.
For either dual representation, GIWAE , IWAE , and I NFO -NCE arise from limiting the family of the critic functions T φ in order to eliminate the intractable log partition function term.
We start from the decomposition of I ( x ; z ) into I BA L ( q θ ) and its gap
[Equation]
We will focus on the dual representation of the normalized, conditional KL divergence Ω( · ) =
D KL [ ·(cid:107) q θ ( z | x )] , as in App. K.2.
Multi-Sample IBAL Consider extending the state space of the posterior or target distribution p ( z | x ) , by using an additional K − 1 samples from a base variational distribution q θ ( z | x ) to construct
[Equation]
where s is an index variable s ∼ U ( s ) drawn uniformly at random from 1 ≤ s ≤ K , which speciﬁes the index of the posterior sample p ( z | x ) .
We similarly expand the state space of the base variational distribution to write
[Equation]
It can be easily veriﬁed that this construction does not change the value of the KL divergence
[Equation]
Consider the convex function
[Equation]
where the primal variable is a distribution in the extended-state space of (cid:0)
z
(1: K ) , s (cid:1) , and the dual variable is a critic function T ( x , z (1: K ) , s ) .
We derive a conjugate optimization in similar fashion to
App. K.2, but now over the extended state space.
For this Ω( · ) , the conjugate function Ω ∗ ( T ) takes a log-mean-exp form analogous to Eq.
We can write the variational representation of Ω( p TGT ) as Ω
[Equation]
For a particular choice of T ( x , z (1: K ) , s ) , we can use Eq.
(161) to obtain a lower bound on the KL divergence D KL [ p ( z | x ) (cid:107) q θ ( z | x )] (as in Eq.
This lower bound translates to the Multi-Sample
IBAL lower bound on MI , I MS - IBAL ( q θ , T ) via Eq.
[Equation]
where Z ( x ; T ) is the normalization constant of the dual distribution π T ( z (1: K ) , s | x ) corresponding to T , π T ( z (1: K ) , s | x ) = Z ( x ; T ) q [Equation] |
[Equation]
As in Eq.
(133), we can write the gap of I MS - IBAL ( q θ , T ) as a Bregman divergence or KL divergence
[Equation]
The optimal K -sample energy function in Eq.
(161) should result in π T ∗ ( z (1: K ) , s | x ) =
p TGT ( z (1: K ) , s | x ) .
Using similar reasoning as in App.
C.5 or App.
L.1, this occurs for
T ∗ ( x , z (1: K ) , s ) = log p ( z ( s ) | x ) q θ
( z ( s ) | x ) + c ( x ) , for which we have I ( x ; z )
=
I MS - IBAL ( q θ , T ∗ ) .
GIWAE is a Multi-Sample IBAL with a Restricted Function Family
Although extending the state space did not change the value of the KL divergence or alter the optimal critic function, it does allow us to consider a restricted class of multi-sample energy functions that yield tractable, low variance estimators.
In particular, GIWAE and I NFO -NCE arise from choosing a restricted family of functions T GIWAE ( x , z (1: K ) , s ) , under which the problematic log Z ( x ; T ) term evaluates to 0 .
This function family is deﬁned as
[Equation] [Equation]
where T GIWAE ( x , z (1: K ) , s ) is speciﬁed by an arbitrary single-sample critic function T φ ( x , z ) .
We can now see that log Z ( x ; T GIWAE )
= 0 ,
[Equation]
using the fact that z (1: K ) ∼ K (cid:81) k
=1 q θ ( z ( k ) | x ) is invariant to re-indexing.
With this simpliﬁcation,
I MS - IBAL ( q θ , T GIWAE ) recovers the GIWAE lower bound on MI
I ( x ; z
[Equation]
- IBAL (
[Equation]
GIWAE )
[Equation] [Equation]
= I GIWAE ( q θ , T φ , K ) .
Finally, we can use Eq.
(164) to recover the probabilistic interpretation of GIWAE and the gap in the lower bound on MI .
As we saw above, the dual distribution π T GIWAE ( z (1: K ) , s | x ) is normalized with
Z ( x ; T GIWAE )
= 1 .
In particular, we can write
[Equation]
which recovers q GIWAEPROP ( z (1: K ) , s | x ) from the probabilistic interpretation of GIWAE in Eq.
(42) or
App.
C.1.
We can write the gap of I GIWAE ( q θ , T φ , K ) =
I MS
[Equation]
IBAL ( q θ , T GIWAE ) as the Bregman divergence or KL divergence as in Eq.
[Equation]
which matches the reverse KL divergence E p ( x ) (cid:2) D KL [ p GIWAETGT ( z (1: K ) , s | x ) (cid:107) q GIWAEPROP ( z (1: K ) , s | x )] (cid:3) derived from the probabilistic approach in App.
C.1.
Recall that I NFO -NCE is a special case of GIWAE with q θ ( z | x ) = p ( z ) (Sec. 2.4).
Conjugate Duality Interpretation of IWAE
We can gain alternative perspective on IWAE from this conjugate duality interpretation.
In particular, IWAE is a special case of GIWAE , where the optimal single-sample critic function T ∗ ( x , z )
= log p ( x , z ) q θ ( z | x ) + c ( x ) (see Sec. 2.4) is used in Eq.
(165) to construct the optimal multi-sample function T IWAE , within the GIWAE restricted multi-sample function family.
Thus, we have I MS - IBAL ( q θ , T IWAE )
=
I IWAE ( q θ , K ) .
Although IWAE uses the optimal critic function, the restriction to the function family in Eq.
(165) is necessary to obtain a tractable bound on the KL divergence D KL [ p ( z | x ) (cid:107) q θ ( z | x )] and mutual information.
Without the restricted function family, the intractable log partition term in Eq.
(161) would require MCMC methods such as AIS to accurately estimate, as we saw for the single-sample
IBAL in Sec. 4.
L P ROPERTIES OF THE IBAL
Our MINE - AIS method in Sec. 4 and App.
M optimizes the Implicit Barber Agakov lower bound ( IBAL ) on MI from Eq.
We ﬁrst recall the probabilistic interpretation of the IBAL bound from App.
J.1.
For a posterior approximation with a learned negative energy function T φ and base variational distribution q θ ( z | x ) ,
[Equation]
where Z θ,φ ( x
[Equation]
=
[Equation]
θ
[Equation]
z |
[Equation]
) we consider the BA lower bound on MI ,
[Equation]
[ p (
[Equation]
(cid:123)(cid:122) (cid:125) contrastive term (176) =
[Equation]
The gap of this lower bound on mutual information is E p ( x ) (cid:2) D KL
[ p ( z | x ) (cid:107) π θ,φ ( z | x )] (cid:3) , as in Sec. 2.2.
L. P ROOFS FOR IBAL O PTIMAL C RITIC F UNCTION (P ROP . 4.1 AND L.1)
Proposition 4.1.
For a given q θ ( z | x ) , the optimal IBAL critic function equals the log importance weights up to a constant T ∗ ( x , z )
= log p ( x , z ) q θ ( z | x ) + c ( x ) .
For this T ∗ , we have IBAL ( q θ , T ∗ ) =
I ( x ; z ) .
Proof.
Recall that the gap in IBAL ( q θ , T φ ∗ ) is E p ( x ) (cid:2) D KL [ p ( z | x ) (cid:107) π θ,φ ( z | x )] (cid:3) =
I ( x ; z ) −
IBAL ( q θ , T φ ∗ ) .
This implies that the bound will be tight 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
[Equation]
where e c ( x ) is absorbed into the normalization constant.
Conversely, using g ( x , z ) which depends on z in T ( x , z ) = log p ( x , z ) q θ ( z | x ) + g ( x , z ) would change the density over z to no longer match p ( z | x ) .
At this value, the IBAL is exactly equal to I ( x ; z )
[Equation]
Proposition L.1.
Suppose the critic function T φ ( x , z ) is parameterized by φ , and
that ∃ φ 0 s.t.
T φ 0 ( x , z ) = const.
For a given q θ ( z | x ) , let T φ ∗ ( x , z ) denote the critic function that maximizes IBAL ( q θ , T φ ) .
Then,
[Equation] [Equation]
In particular, the contrastive term in Eq.
(18) is upper bounded by E p ( x ) [ D KL [ p ( z | x ) (cid:107) q θ ( z | x )]] .
Proof.
The BA bound is a special case of the IBAL ( q θ , T φ ) with constant T φ
= c and the contrastive term equal to 0 .
Since φ 0 is a possible parameterization, we can only improve upon I BA L ( q θ ) by learning T φ .
The parameterized 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 ) using Prop. 4.1.
L. P ROOF OF IBAL AS L IMITING B EHAVIOR OF THE GIWAE O BJECTIVE AS K → ∞ (P ROP . L.2)
Proposition L.2 ( IBAL as Limiting Behavior of GIWAE ) .
For given q θ ( z | x ) and T φ ( x , z ) , we have
[Equation]
Proof.
Comparing the form of I GIWAE L ( q θ , T φ , K ) to the IBAL ( q θ , T φ ) for a ﬁxed q θ and T φ ,
[Equation]
we can see that both bounds include the same ﬁrst term I BA L ( q θ ) and the same numerator of the contrastive term.
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 ( x , K ) for convenience and the right hand side is the log partition function for the IBAL energy-based posterior π θ,φ ( z | x ) .
Intuitively, we expect Eq.
(181) 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 ( x , K ) as a function of K .
We derive lower and upper bounds on the value of log Z GIWAE ( x , K ) for ﬁxed K and show that each of these sequences of lower and upper bounds converge to log Z θ,φ ( x ) in the limit as
K → ∞ .
Using the squeeze theorem for sequences, this is sufﬁcient to demonstrate the claim that lim K →∞ log Z GIWAE ( x , K )
=
log Z θ,φ ( x ) in Eq.
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 on log Z GIWAE ( x , K ) : 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]
where in (1) , we obtain a lower bound by ignoring the contribution of the positive sample from p ( z | x ) .
The ﬁrst term is the k -sample IWAE lower bound with the proposal q θ ( z ) and target π θ,φ ( x , z ) :
E q θ ( z 1: K ) [log 1 K (cid:80) k
π θ,φ ( x , z ( k ) )
q θ ( z ( k ) | x ) ] , and thus converges to log Z θ,φ as K → ∞ .
As K → ∞ , we also have log K − 1 K → 0 , so that the limit of their sum is
[Equation]
Upper Bound on log Z GIWAE ( x , K ) :
To upper bound log Z GIWAE ( x , K ) , 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.
(26) is invariant to the index s , we can assume z (1)
∼ p
( z | x ) and write
[Equation]
Since log( u ) is a continuous function, we know that lim K →∞ log( f ( K ))
=
log lim K →∞ f ( K ) .
We thus reason about the limiting behavior of the terms inside the logarithm in Eq.
As K → ∞ , we have 1 K → 0 , and thus the ﬁrst term inside the log goes to 0 .
For the second term, we also have K −
K → 1 .
Thus, we have
[Equation]
As reasoned above, the convergence of the sequence of both upper and lower bounds to log Z θ,φ ( x ) implies that lim K →∞ log Z GIWAE ( x , K )
=
log Z θ,φ ( x ) .
By the reasoning surrounding Eq.
(181), this implies lim K →∞
I GIWAE L ( q θ , T φ ) =
IBAL ( q θ , T φ ) as desired.
L. C ONVERGENCE OF GIWAE SNIS D ISTRIBUTION TO IBAL E NERGY -B ASED P OSTERIOR
In this section, we consider the marginal SNIS distribution of GIWAE , which is induced by sampling
K times from q θ ( z | x ) and returning the sample in index s with probability q GIWAEPROP ( s | x , z (1: K ) )
∝ e T φ ( x , z ( s ) ) .
As K → ∞ , we show that this distribution converges to the single-sample energy-based posterior approximation underlying the IBAL and MINE - AIS .
A similar observation is made in Sec.
3.2 of Lawson et al.
This result regarding the probabilistic interpretations of GIWAE and
the
IBAL is complementary to the result in Prop.
L.2 regarding the limiting behavior of the bounds.
Proposition L.3.
Deﬁne the marginal SNIS distribution of GIWAE , q GIWAEPROP ( z | x ; K ) , using the following sampling procedure
Sample from s, z (1: K ) ∼ q GIWAEPROP ( s, z (1: K ) | x ) according to Eq.
Return z
= z ( s ) .
Then, as K → ∞ , the KL divergence (in either direction) between the marginal SNIS distribution of
GIWAE and the energy-based variational distribution of IBAL , π θ,φ ( z | x ) ∝ q θ ( z | x [Equation] T φ [Equation]
x , z ) , goes to zero.
[Equation]
and
[Equation]
Proof.
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 the probabilistic interpretations q GIWAEPROP ( s, z (1: K ) | x ) and p GIWAE ,π TGT ( s, z (1: K ) , x ) exactly match the IWAE proposal and target distributions in App.
B.1 where, in the target distribution, the posterior p ( z | x ) has been replaced by π θ,φ ( z | x ) .
Thus, we can analyze the KL divergences
D KL [ q GIWAEPROP ( s, z
(1: K )
| x ) (cid:107) p GIWAE ,π TGT ( s, z (1: K ) | x )] and D KL [ p GIWAE ,π TGT ( [Equation] x
[Equation] (cid:107)
[Equation] GIWAEPROP ( s [Equation] | x )] using techniques from previous work on IWAE .
Following similar arguments as in Domke & Sheldon (2018) (Thm 2) and Cremer et al.
(2017), these extended state space KL divergences upper bound the KL divergence between the marginal SNIS distribution and energy-based target, e.g.
D KL [ q GIWAEPROP ( z | x ; K ) (cid:107) π θ,φ ( z | x )] ≤ D KL [ q GIWAEPROP ( s, z
(1: K )
| x ) (cid:107) p GIWAE ,π TGT ( s, z (1: K ) | x )] .
As K → ∞ ,
D KL [ q GIWAEPROP ( s, z
(1: K )
| x ) (cid:107) p GIWAE ,π TGT ( s, z (1: K ) | x )]
→ 0 since it is an instance of the IWAE gap.
Thus, the KL divergence D KL [ q GIWAEPROP ( z | x ; K ) (cid:107) π θ,φ ( z | x )] also vanishes.
Similar reasoning applies for the reverse KL divergence.
M MINE-AIS M. E NERGY -B ASED T RAINING OF THE IBAL
Eq.
(19) indicates that in order to maximize the IBAL as a function of θ and φ , 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 ) .
Letting T 1: M indicate the composition of T transition steps, the approximate energy function gradient becomes
[Equation]
We can use an identical modiﬁcation for the gradient with respect to θ .
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 one starts an MCMC chain from the true data distribution to obtain a lower variance gradient estimate.
M. M ULTI -S AMPLE AIS E VALUATION OF THE IBAL
After training the variational base distribution q θ ( z | x ) and the critic function T φ ( x , z ) using the MINEAIS training procedure above, we still need to evaluate the IBAL lower bound on MI , IBAL ( q θ , T φ ) .
We can easily upper bound IBAL ( q θ , T φ ) using a Multi-Sample AIS lower bound on log Z θ,φ ( x ) with expectations under the forward sampling procedure q AISPROP ( z 0: T | x ) .
However, an upper bound on
IBAL ( q θ , T φ ) is not guaranteed to preserve a lower bound on MI .
In order to obtain a lower bound on the IBAL , we would need to obtain an upper bound on log Z θ,φ ( x ) = log E
q θ ( z | x ) (cid:2) e T φ ( x , z ) (cid:3) , the log partition function of π θ,φ ( z | x ) .
However, considering this to be the target distribution π T ( z | x ) in Multi-Sample AIS , we would require exact samples from π θ,φ ( z | x ) to guarantee an upper bound on log Z θ,φ ( x ) .
Since these samples are unavailable, we demonstrate conditions under which we can preserve an upper bound on log Z θ,φ ( x ) (and lower bound on IBAL ( q θ , T φ ) ) by sampling from p ( z | x ) instead of π θ,φ ( z | x ) to initialize our backward annealing chains in Prop.
M.1 below. Using the single sample AIS bounds to estimate log Z θ,φ ( x ) , we have the following extended state space proposal and target distributions, [Equation] We emphasize that in p AIS ,π TGT , q AIS ,π PROP , the transition kernels and intermediate densities are based on the critic T φ ( x , z ) and energy-based posterior π θ,φ ( z | x ) whose log partition function we seek to estimate.
Recall from Sec.
2.1 that taking the expected log importance weights under p AIS ,π TGT ( z 0: T | x ) yields an upper bound on log Z θ,φ ( x ) .
However, since it is difﬁcult to draw exact samples from π θ,φ ( z | x ) to initialize backward annealing chains and sample from p AIS ,π TGT ( z 0: T | x ) , we instead consider sampling from the posterior p ( z | x ) . Using the same transition kernels ˜ T t as above, we ﬁrst deﬁne the conditional distribution p AIS ,π TGT ( z 0: T − 1 | x , z T ) of the backward chain for a given z T [Equation] We can then deﬁne the approximate extended state space target distribution which samples z T ∼ p ( z | x )
[Equation]
Using this notation, we can also write p AIS ,π TGT ( z 0: T | x ) =
π
[Equation] z T | x ) p AIS , [Equation] ( z
[Equation]
T − 1 | x , z T ) in Eq.
We now characterize the conditions under which sampling from p APPROXTGT ( z 0: T ) preserves an upper bound on log Z θ,φ ( x ) .
Proposition M.1.
Deﬁne the AIS marginal distribution q AIS ,π PROP
( z T | x ) over the ﬁnal state in the extended state space proposal as q AIS ,π PROP ( z T | x ) := [Equation] q AIS ,π PROP ( z 0: T | x ) d z 0: T − 1 . If we have
[Equation] then initializing the backward AIS chain using z T ∼ p ( z | x ) (i.e., sampling under p APPROXTGT ( z 0: T | x ) ), yields an upper bound on log Z θ,φ ( x ) , [Equation] Proof.
We begin by writing several deﬁnitions, which will allow us to factorize the extended state space proposal q AIS ,π PROP ( z 0: T | x ) in the time-reversed direction. This factorization includes the ﬁnal AIS marginal q AIS ,π PROP ( z T | x ) . Starting from Eq. (192), the forward transitions q AIS ,π PROP ( z 0: T | x ) =
q θ ( z 0 | x ) (cid:81) Tt =1 T t ( z t | z t − 1 ) induce the marginal distributions q AIS ,π PROP ( z t | x ) at each step. The forward transitions and marginals induce a posterior kernel ˜ T qt ( z t − 1 | z t ) , which allows us to rewrite q AIS ,π PROP ( z 0: T | x ) using a reverse factorization [Equation]
where
[Equation]
The posterior reverse transitions ˜ T qt ( z t − 1 | z t ) are intractable in practice, and cannot be simpliﬁed to match the kernels in the target distribution p AIS ,π TGT ( z 0: T − 1 | x , z T ) =
(cid:81) Tt =1 ˜ T t ( z t − 1 | z t ) using the invariance or detailed balance conditions. Doucet et al. (2022) provide a promising approach using score matching to approximate these posterior transitions ˜ T qt ( z t − 1 | z t ) .
Finally, we write the posterior reverse process conditioned on a particular z T as [Equation] With the goal of upper bounding log Z θ,φ ( x ) [Equation]
q
[Equation]
d z , we consider the log importance weights with expectations under the target distribution p APPROXTGT ( z 0: T | x )
= p ( z T | x ) p AIS ,π TGT ( z 0: T − 1 | x , z T ) , as in Eq. Using the above notation, we have [Equation] The intractable KL divergence D KL [ p AIS ,π TGT ( z 0: T − 1 | x , z T ) (cid:107) q AIS ,π PROP ( z 0: T − 1 | x , z T )] compares the reverse kernels in the target distribution ˜ T t ( z t − 1 | z t ) against the posterior ˜ T qt ( z t − 1 | z t ) .
Ignoring this nonnegative term, we can lower bound the expectation on the LHS
E p APPROXTGT
[Equation]
Finally, under the assumption of the proposition that D KL (cid:2) p ( z T | x ) (cid:107) q AIS ,π PROP ( z T | x ) (cid:3) ≥
D KL (cid:2) p ( z T | x ) (cid:107) π θ,φ ( z T | x ) (cid:3) , we have the desired result.
KL Divergence Condition In Prop.
M.1, we have shown that we can preserve an upper bound on log Z θ,φ ( x ) by initializing the reverse chain using a posterior sample, under a condition on the KL divergence, [Equation] While we cannot guarantee this condition, we intuitively expect Eq. (203) to hold in practice since π θ,φ ( z | x ) has been directly trained to match p ( z | x ) .
By contrast, q AIS ,π PROP ( z T | x ) is the ﬁnal state of an AIS procedure, which approximates π θ,φ ( z | x ) and does not have access to information about p ( z | x ) . Burda et al. (2015) use a similar approach for lower bounding the log likelihood in EBM s, but give an example of a Restricted Boltzmann Machines ( RBM ) model (in their Sec. 5) in which Eq.
(203) does not hold. As desired, we ﬁnd in our experiments in Fig. 7 that our approximate reverse annealing procedure underestimates the IBAL in all of the MINE - AIS , GIWAE and I NFO NCE experiments, for all numbers of intermediate distributions T . T = 1 Special Case For T
= 1 , we have q AIS ,π PROP ( z 1 | x ) = q θ ( z 1 | x ) .
In particular, Prop. M.1 will provide an upper bound on log Z θ,φ ( x ) if
[Equation]
This condition is guaranteed under the assumptions of Prop. L.1, where IBAL ( q θ , T φ ∗ ) = I BA L ( π θ,φ ∗ ) improves upon I BA L ( q θ ) for an energy function
T φ ∗ ( x , z ) which has been trained to maximize the IBAL . In Eq.
(204) the KL divergence on [Figure]
which E , z
[Equation]
We can conﬁrm this in Fig. 7, where for T = 1 , the approximate lower bounds on IBAL ( q θ , T φ ) begin from the appropriate BA lower bound.
For example, in the case of IBAL evaluation for GIWAE ( K = 100 ), the light green curve starts from the BA lower bound term reported in the decomposition of the GIWAE ( K = 100 ) lower bound in Fig. 3a.
Using more intermediate distributions ( T > 1 ) for the AIS approximate bound in Prop. M.1, the estimates in Fig.
7 approach the true value of IBAL ( q θ , T φ ) in all cases.
M. MINE-AIS E XPERIMENTS
In Fig.
7, we assess the performance of our MINE - AIS method, which learns an energy function T φ and uses Multi-Sample AIS to evaluate lower and upper bounds on the IBAL ( q θ , T φ ) , 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 these estimators tightly bound the ground truth MI and correspond to the optimal energy function in MINE - AIS (Prop. 4.1).
We train MINE - AIS using IBAL ( p ( z ) , T φ ) for p ( z ) =
N (0 , I ) , and plot our Multi-Sample AIS lower and upper bounds on the IBAL as a function of the number of intermediate AIS distributions T used for evaluation.
Note that our lower bound on IBAL , which uses the approximate reverse annealing procedure from
App. M.2, preserves a lower bound on MI for all values of T .
The upper bound on IBAL may be used to validate the convergence of the Multi-Sample AIS evaluation procedure.
For example, in Fig. 7, we observe that our lower and upper bounds converge to the same estimate for high T , ﬁnding the true value of IBAL ( p ( z ) , T φ ) .
We ﬁnd that MINE - AIS underestimates the ground truth MI by 11% , and 24% on MNIST - VAE and
MNIST - GAN , respectively.
We also compare against the BA lower bound with a Gaussian variational family for q θ ( z | x ) .
Note that the IBAL corresponds to the BA lower bound for a more ﬂexible, energy-based posterior approximation π θ,φ ( z | x ) , and both bounds assume access to only a single known marginal distribution.
We ﬁnd that MINE - AIS outperforms the BA baseline by 2% , 52% , and 90% , on linear VAE , MNIST - VAE , and MNIST - GAN , respectively.
We also use Multi-Sample AIS to evaluate the IBAL corresponding to ( q θ , T φ ) , which are learned by optimizing the GIWAE (with K = 100 ) and I NFO NCE (with K = 100 and q θ = p ( z ) ) lower bounds. As argued in Prop. L.2, the IBAL corresponds to the limiting behavior of GIWAE as K → ∞ . We observe that the AIS evaluation of the IBAL corresponding to GIWAE or I NFO NCE critic functions only marginally improves upon evaluation of the original GIWAE or I NFO NCE lower bounds.
This indicates that the improvement of MINE - AIS over GIWAE or I NFO NCE can be primarily attributed to learning a better critic function using energy-based training. For these evaluation tasks, the approximate reverse annealing procedure from App.
M.2 also underestimates the IBAL for all values of T .
N 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]
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 .
N. BA L OWER B OUND
Using the notation of Eq.
(205), we can write the BA lower bound as
[Equation]
negative reconstruction
[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.
(207) 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), improving interpretability in GAN s (Chen et al., 2016), and variational information bottleneck methods (Tishby et al., 2000; Alemi et al., 2016; 2018).
N. 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.
(7) for the case of estimating I ( x ; z ) in Eq.
(205), we have
[Equation]
[Equation]
Note that we can use a single joint sample from q data ( x ) q ψ ( z | x ) = q ψ ( z ) q ψ
[Equation] x | z ) to obtain a positive sample from the inverse encoding distribution x ∼ q ψ ( x | z ) for a particular z ∼ q ψ ( z ) , in a 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 the “log partition function” log q ψ ( z ) for a particular z ∼ q ψ ( z ) .
The contrastive term in Eq.
(209) 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. (207) 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.
(209) depend on samples from q data ( x ) and not the density.
Thus, 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.
(209) 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 (e.g., see van den Oord et al. N
[Equation]
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 E q ψ [Equation] z ) [log Z
[Equation]
z )] =
[Equation]
q ψ
[Equation]
z )
[Equation]
p
[Equation]
z
[Equation]
[Equation]
where the term in the denominator is the partition function Z ( z )
=
E p θ ( x | z ) (cid:2) e T φ ( x
[Equation]
z ) (cid:3) .
When taking gradients as in Eq. (19)-(20), we obtain
[Equation]
To obtain approximate negative samples from π θ,φ ( x | z ) for a given z , we can use MCMC transition kernels since the unnormalized target density ˜
π θ,φ ( x | z ) =
p θ ( x | z ) e T φ ( x , z ) is tractable.
Evaluating MI up to a Constant Since only samples from q data ( x ) are required for the MINE - AIS training procedure in Eq.
, we can learn the base variational distribution p θ ( x | z ) and energy function T φ ( x , z ) in cases when the marginal distribution q data ( x ) is unknown.
As in the case of GIWAE and BA lower bounds above, we can also evaluate the IBAL lower bound in
Eq.
(211) 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, but computing the gradients are intractable.
We thus conclude that, among our proposed methods, GIWAE is the most directly applicable in settings of MI for representation learning.
O E XPERIMENTAL D ETAILS O. E XPERIMENT D ETAILS OF S EC . 5.
In this section, we provide the experiment details used in Sec. 5.2.
For more details, see the public
GitHub repository, https://github.com/huangsicong/ais_mi_estimation.
O.1. 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 learned scalar.
For the GAN experiments on MNIST (Table 1), we used the same decoder architecture as our VAE s. In order to stabilize the training dynamics, 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 a smaller version of
DCGAN (Radford et al., 2015) (see the public code).
The number at the end of each model name in
Table 1 indicates the latent code size.
O.1. E XPERIMENT D ETAILS OF S EC . 5.
For the AIS temperature schedule, We used sigmoid schedules as used in Wu et al.
The step size of HMC was adaptively tuned to achieve an average acceptance probability of 65% as suggested in Neal (2001).
For all MNIST experiments in Table 1, we evaluated on a single batch size ofsimulated data.
For all CIFAR experiments in Table 1 we used a single batch of 32 simulated data.
All experiments are run on on Tesla P100 or Quadro RTX 6000 or Tesla T4 GPUs.
O.1. R UNTIME C OMPARISON
We benchmarked the runtime on Tesla P100 GPUs.
For MNIST , it took about 35 minutes to run IWAE with K = 1 M , about 8 hours to run AIS with T = 30 K .
For CIFAR , it took about 45 minutes to run
IWAE with K = 1 M , and about 12 hours for the AIS with 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.
O.1.
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 .
[Figure]
E XPERIMENT D ETAILS FOR E NERGY -B ASED B OUNDS (S EC . 5.2) Models and
Data
The linear VAE model has a Gaussian prior z ∼ N (0 , I ) with the dimension of 10.
The dimension of the output x is 100 .
The weights are sampled randomly from a Gaussian distribution with the standard deviation of 1 , and the standard deviation of the Gaussian observation noise at the output is 1 .
The MNIST - VAE 20, is a VAE model that is trained on the real-valued MNIST dataset.
It has a Gaussian prior z ∼ N (0 , I ) with the dimension of 20.
The decoder has one layer of ReLU non-linearity of size 1000 , followed by a linear layer that predicts the mean of a Gaussian observation model with the ﬁxed standard deviation of 0 . 1 .
MNIST - GAN 20 decoder uses one layer of ReLU non-linearity of size 1000 followed by a sigmoid layer that predicts the mean of a Gaussian observation model with the ﬁxed standard deviation of 0 . 1 .
In all the experiments, we use a ﬁxed batch of 100 data points.
BA, IWAE and GIWAE For the BA bound, q θ ( z | x ) is parameterized using a neural network that predicts the mean and log-std of a diagonal Gaussian distribution.
The neural network has two layers of ReLU non-linearity of size 2000 .
The IWAE experiments used the same architecture for q θ ( z | x ) .
The GIWAE experiments, in addition to q θ ( z | x ) with the same architecture, used a critic function
T φ ( x , z ) that is parameterized by a neural network that concatenates ( x , z ) and pass them through two layers of ReLU non-linearity of size 2000 , followed by a linear layer that outputs a scalar value.
The parameters of q θ ( z | x ) and T φ ( x , z ) are trained jointly.
Multi-Sample AIS For the Multi-Sample AIS bounds, for all models, we used up to K =chains (see Fig. 5), and up to T = 50 K intermediate distributions with linear schedule (see Fig.and Fig. 7).
We used HMC as the AIS kernel, with (cid:15) = 0 .
02 and L = 20 leap frog steps.
MINE-AIS For the training of the MINE - AIS , we used a critic function T φ ( x , z ) that is parameterized by a neural network that concatenates ( x , z ) and pass them through three layers of ReLU non-linearity of size 2000 , followed by a linear layer that outputs a scalar value.
We chose the Gaussian prior
N (0 , I ) as the base distribution q θ ( z | x ) .
In order to take a sample from π θ,φ ( z | x ) , we used the
HMC method that is initialized with a true posterior p ( z | x ) sample, with M = 10 iterations each with L = 20 leapfrog steps.
For the step size of HMC , we used (cid:15) = 0 .
for the linear VAE model and (cid:15) = 0 .
02 for the MNIST - VAE 20 and MNIST - GAN 20.
For the evaluation of MINE - AIS , we used
Multi-Sample AIS with the same parameters as the AIS evaluation experiments described above.
O. A NALYTICAL S OLUTION OF THE M UTUAL I NFORMATION ON THE L INEAR MNIST - VAE
In order to verify our implementations, we have derived the MI analytically for the linear VAEs and veriﬁed that it matches the MI estimated by AIS .
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 is easily shown that the conditional likelihood is p ( x | z ) =
N ( x | Wz + b , I ) and thus we can solve for the marginal
[Equation]
The differential entropy of x is:
[Equation]
215) where k is the dimension of the observation.
The conditional entropy is
[Equation]
[Equation]
[Equation]
Thus, the mutual information is
[Equation]
[Equation]
O. C ONFIDENCE I NTERVALS FOR M ULTI -S AMPLE AIS E XPERIMENTS
Table 4 and Table 5 provides the 95% conﬁdence intervals for the AIS results reported in Table 1.
The conﬁdence intervals were computed over conﬁdence interval over 8 batches each with 16 data points for MNIST ; and over 8 batches 4 data points for CIFAR .
[Table]
[Table] [Table]