Improving Mutual Information Estimation with Annealed and Energy-Based Bounds
Authors: Rob Brekelmans
Sicong Huang
Marzyeh Ghassemi
Greg Ver Steeg
Roger Baker Grosse
Alireza Makhzani
Venue: ICLR 2022 Poster (2022)
Keywords: mutual information estimation
annealed importance sampling
energy-based models

Abstract: 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 unifies variational and contrastive bounds in a single framework that generalizes InfoNCE, 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. MINE-AIS 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 VAEs and GANs trained on the MNIST and CIFAR datasets, and showcase significant gains over existing bounds in these challenging settings with high ground truth MI.

Published as a conference paper at ICLR 2022
IMPROVING MUTUAL INFORMATION ESTIMATION WITH ANNEALED AND ENERGY-BASED BOUNDS
Rob Brekelmans
Sicong Huang*
Marzyeh Ghassemi
Greg Ver Steeg!
Roger Grosse 2,3
Alireza Makhzani
Information Sciences Institute, University of Southern California
Vector Institute University of Toronto MITT BECS/ / IMES / CSAIL
ABSTRACT
Mutual information (MI) 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 known, as involves estimating potentially high-dimensional log partition function. In this work, we present a unifying view existing MI bounds from the perspective of importance sampling, and propose three novel bounds based on this approach. Since tight MI bound without density information requires a sample size exponential in the true MI, we assume either single marginal 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 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 GIWAR bound unifies variational and contrastive bounds in a single framework that generalizes INFONCE, IWAE, and Barber-Agako bounds. Our MINE-AIS method improves upon existing energy-based methods such MINE-DV and MINE-F by directly optimizing a tighter lower bound on MI. MINE- AIS 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 joint densities are often available. We evaluate our bounds on estimating the MT of VAES and GANS trained on the MNIST and CIFAR datasets, and showcase significant gains over existing bounds in these challenging settings with high ground truth MI.
INTRODUCTION
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 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 al., 2018; van den Oord al., 2018; Poole al.. 2019). However, McAllester Stratos (2020) have shown strong theoretical limitations on any estimator based on direct sampling without an analytic form of least one marginal distribution. 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 general approach for constructing extended state space bounds on MI, we combine insights from importance- weighted autoencoder (IWAE) (Burda et al., 2016; Sobolev Vetrov, 2019) and annealed importance sampling (AIS) (Neal, 2001) to propose Multi-Sample Sample AIS bounds in Sec. 3. We empirically show that this approach can tightly estimate large values of MI when the full joint distribution is known.
"Equal Contribution. brekelma @ edu; (huang, makhzani @esstorontooedd edu. See arXiv for extended paper.

---

Published as a conference paper at ICLR 2022
[(x B)
Schematic of MI bounds discussed in this paper. Green shading indicates our contri- butions, while columns and gold labels indicate single- or multi-sample bounds. Blue arrows in- dicate special cases using the indicated proposal distribution. Relationships based on learned critic functions are indicated by red arrows. obtain only lower bounds on MI with unknown p(x(z), but both upper and lower bounds with known p(x|z), All bounds require a known marginal p(z) for evaluation, apart from (Structured) INFO-NCE.
Figure 1:
aaggl-anapple
traltionmpple
suit samplle
singlo dingloioneele
BA
CIWAE
TWAR dhale tesponal
AIS
pr (IdoNNC
sorgg-leed
S S
MINE-AIS
Mulli Samplee AIS
and ocalaras
Lowe, Dound on MI
Lowe and Upper Dound on XII
Our importance sampling perspective also suggests improved MT 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 INFONCE (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-Agakon Lower bound (IBAL). We demonstrate that the IBAL corresponds to the infinite-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 Multi- Sample 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. 1.
1.1
PROBLEM SETTING
The mutual information between two random variables x and z with joint distribution p(x, z) is
log p(x, x) p(x)p(z)
1(x; z) Ep(x,)
H(x) (x Ep(x.))loopp((ee)
Epco Eooo]log (I)
where Hx((z) denotes the conditional entropy - -Fpp(x.)) p(xz)) We primarily focus 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)." Settings where only single marginal 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 given 6, but the likelihood p(x(8) (0) 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))()(x) are available. We will refer to bounds which involve only single posterior sample as practical, and those involving multiple posterior samples as impractical.
When the conditional p(x /2 is tractable to sample and evaluate, simple Monte Carlo sampling provides an unbiased, low variance estimate of the conditional entropy term Eq. (1). In this case, the difficulty 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.
UNIFYING MUTUAL INFORMATION BOUNDS VIA IMPORTANCE SAMPLING
In this section, we present a unified view of mutual information estimation from the perspective of ex- tended 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.1 A GENERAL APPROACH FOR EXTENDED STATE SPACE IMPORTANCE SAMPLING BOUNDS
To estimate the log partition function, we construct a proposal qpzar (Zea) xx) and target Pror (x, Zest) distribution over an extended state space, such that the normalization constant of Pror (x, Zexi) is
ZTar = f Pror (x, Zea) ddze = p(x) and the normalization constant of 4exop (x, Zexi) is Zprop = 1. Taking expectations of the log importance weights log Prey (x, Zea) qroor (Zexi (Zexi) (x) under the proposal
An alternative, "encoding" MI hetween the real data and the latent code is often of interest (see App. N), but cannot he directly estimated using our methods due the unavailability of pa(x) or g(z) = J (a(x))(zlxxxdx
2

---

Published as a conference paper at ICLR 2022
and target, respectively, obtain lower and upper bounds the log partition function
Prat Zext)
Pror (x, Zeet
yn (Reulx)
log
I log p(x) S Prar(teulx)) log
(2)
Qexor Zest x)
Qenor ( (Zex) X
ELBO x: Ivxor: Proc)
EUBO (x; Tenor: Pret)
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 DKL 14rRop DKL.4FrOP! (Zext xx) PlarrllPrraaaaaa x)] and the gap in the upper bound equal to the reverse KIL divergence DxL [Pro (exx xx) APROP Zeei x)]. See App.
2.2 BARBER-AGAKOV AGAKOV LOWER AND UPPER BOUNDS
As a first example, consider the standard ELBO (ge) and BUBO( bounds, which are derived from simple importance sampling using a variational distribution ()(z)x)
x) and Zextt = z in Eq. (2). Plugging these lower and upper bounds on p(x) into Eq. (1), we obtain upper and lower bounds on MI as
I (qs) == Ep(x..) log 2o (z)x))
log [0(x - H(x[2) =: Inny (go). (3) p(x, 2)
S (xx:2) S (p(x))o(=xx)
p(z)
The left hand side of Eq. (3) is the well-known Barber-Agakov (BA) bound (Barber & Agakov, 2003), which has a gap of Eppx) [KI. (e(z]p(({(x) I1 [/]x]]]_ We refer to the right hand side as the BA upper bound Luve (go), with gap of Bp(x) [KL [qo (z) [|p[z|x)]]] In contrast to Laxe (go), note that TAA (9o) does not require access the conditional p(x[z) to evaluate the bound.
2.3 IMPORTANCE WEIGHTED AUTOENCODER
The IWAE lower and upper bounds on p() (Burda et al., 2016; Sobolev & Vetrov, 2019) improve upon simple importance sampling by extending the state space using multiple samples Zext = (I:K), Consider proposal GRAA (z( (z(1:) x) with K independent samples from a given variational distribution
qe (z) The extended state space target Pror= (z (1::k) xx) is mixture distribution involving a single
sample from the posterior p(z|x) or joint p(x; 2) and K samples from 40 (z)xx
K
K Σ
K
TWAE 4rROP
(1:R)
II 1100 (x
Pror (x, 2 (1:3)
p(x,2'
II k==1, Eopte
(x (x).
(4)
K
Piar (1)00)
The log importance weight log
(x(0 reduces to the familiar ratio in the IWAE objective, while
the normalization constant of prer (x, z (l::K) is p(x). As in Sec. 2.1, taking expectations under the proposal or target yields a lower or upper bound, respectively,
1 K
p(x, z (A))
1 K
p(x, 2 (5)
R
log
S log log)(x) S F
log
(5)
(x(
(2)
[x]
p(2())))) (x) ff (a((A))))))
(2(2
ELBO (x: ge,
EUBO (X; qo, K)
As for the standard ELBO and EUBO, the gap in the lower and upper bounds are DKL [g | p and DKL [ [|
[WAR]
[]waF], respectively. See App. B for detailed derivations. As in Sec. 1.1, with known ((x the lower and upper bounds on log p(x). ELBO qo, K) and HUBO qo, K) translate to upper and lower bounds on MI, which we denote as IIAAE (go. K) and IIAAE (g0,
While it is well-known that increasing K leads tighter IWAE bounds (Burda et al., 2016; Sobolev Vetrov, 2019), in App. B.2, we characterize the improvement of ELBO (qo. K) over EL))(gg) as a KI. divergence. For EUBO K), we show that the KI. divergence measuring its improvement over (g)) is limited to log K, which implies TIWAR, (go, K) < TRAy (go) + log K.
2.4 GENERALIZED IWAE
In this section, we consider family of Generalized IWAE (GIWAE) lower bounds, which improve upon IaA, (go) 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 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 infinite sample limit of GIWAE
To derive a probabilistic interpretation for GIWAF, we begin by further extending the state space of the IWAE target distribution in Eq. (4), using uniform index variable p(s) K which sample (+) ~ p(z/x) is drawn from the posterior
Vs that specifies
3

---

Published as a conference paper at ICLR 2022
K II q(2 (x). 1=1,
PIGT
(1:K) .0/x)
p(z
p(2(1))))
(6)
K
Note that marginalization over leads the IWAE target P (1:K) x) in Eq. (4). For the GIWAE extended state space proposal, consider categorical index variable s (s)z (s|z(1:)), x) drawn using self-normalized importance sampling (SNIS), with weights calculated by learned critic function To
TT(xxxl1)
ULWAL (1:K
s|xx =
II q(2 (4)
(x))
CIWAR (s|z (1:K) x), where
GENAL (s (1:K) X)
Σ eTo(xa(t)
We discuss the GIWAE probabilistic interpretation in App. C.I, and find that only the GIWAE upper bound on log p(x) provides practical benefit in App. C.2. The corresponding MI lower bound is
Yo (z/xx) p(z)
e (T(xxa!"))
IOIWAEL (yo, To, K)
E
log
log
p(x.z)
p( (x(2()
K > (o
Tang (ge)
contrastive term S log
We observe that the GIWAE lower bound decomposes into the sum two terms, where the first the BA variational lower bound for go(2)x and the second is contrastive term which distinguishes single positive sample drawn from p(z) x) from
negative samples drawn from go (z|x).
Relationship with BA With a constant T (x z) const, the second term in GIWAE vanishes and we have /GIAAE (g, Ta = const, K) = TBA (go for all K. For K = 1, /CIWAE (q. Top 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 optimize MI even if marginal is available. See App. N. for detailed discussion.
1)
Relationship with InfoNCE When the prior pzz) is used in place of (z(x), the second term in Eq. (7) as the INFONCE contrastive lower bound (van den Oord et al., 2018; Poole et al., 2019), with (INFNNCE, (T', K) [GIWAEz (p(z), Te, 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 INFONCE (van den Oord et al., 2018).
recognize
Relationship with IWAE The following proposition characterizes the relationship between IOIWAE (go. Tg, K) in Eq. (7) and IIAAE (go, K) from Sec. 2.3. See App. C.5 for proofs.
Proposition 2.1 (Improvement of TWAE over GIWAE). For given (o(z)x) and any T 2)).
IIAAR (96, K) LawAAR (go, To, K) + E (p(0)ar""(
Dki [Kallptrr (o000"(())2 (1:K), .x) (allar (s) (1:K)
x)]
Corollary 2.2. For a given qo (2)x) and K > 1, the optimal critic function is the true log importance weight up to arbitrary constant: T (x. log p(x.z)
[x(x]x] + c(x). With this choice of T (x. z).
IGIWAE, (go, T K) JIWAE (go. K).
Corollary 2.3. Suppose the critic funncioon Tx(x,z) is parameterized by G. and 3 s.t. Vl (x,2), 2), Top (x, z) = const. For a given qo xx))), let T (x,2) denote the critic function that maximizes the GIWAE lower bound. Using Cor. 2.2, we have
TaA (go) S TGIWAES (go, Tar, K) TIWAE (g. K) s Jax (go) + log K.
(9)
While the GIWAE lower bound on MI does assume access to the full joint distribution, Cor. 2.2 suggests that the role of the critic function is to learn the true log importance weights for qa(z) x). Thus, when p(x, z) known, IWAE is always preferable GIWAE Cor. 2.3 shows that while IGIWAE, To- K) can improve upon BA, this improvement is at most logarithmic in K.
Relationship with Structured INFONCE Since INFONCE and Structured INFONCE (S- INFONCE, 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)
v 4ISINCEE (T,, K) ISASSONCE, (K) < log K.
(10)
4

---

Published as a conference paper at ICLR 2022
IWAE
Singlo-Sumple A15
Indepeendend Multth Stutti-Smmple 18
Independentt Reverse Miati-Sannn
Coupled Reverse Sulti-Sample
Praetiral?
Target BX30 Tta
Proposal
EL.0
Aua
Figure 2: Extended state-space probabilistic interpretations of multi suulti-sample AIS bounds. Forward chains are colored in blue, and backward chains are colored in red. Note that ELBOS and EUBOS are obtained by taking the expectation of the log importance weight log Prer (-) (/Qror(() under either the proposal or target distribution, and can then be translated to MI bounds.
MULTI-SAMPLE AIS BOUNDS FOR ESTIMATING MUTUAL INFORMATION
In the previous sections, derived probabilistic interpretations of extended state space importance sampling bounds using multiple samples K, as in IWAE. In this section, we first review ATS, which extends the state space using T' intermediate densities. We then show that these approaches are complementary, and derive two practical Multi- Mulpi-Sample AIS methods that provide tighter bounds by combining insights from IWAE and AIS.
3.1 ANNEALED IMPORTANCE SAMPLING BACKGROUND
AIS (Neal, 2001) constructs a sequence of intermediate distributions {z)(Z)) It=0- which bridge between a normalized initial distribution To xx) and target distribution my (z(x) p(z[x) with
unnormalized density T (x(x, z) = p(x, z) and normalizing constant ZT(x) xx)). A common
choice for intermediate distributions is the geometric mixture path parameterized by {8.)7
no [x)2)))= x)" At TT 2x)
where Zx(x) = / [o(z)))²= 1-2 TI (x, z) dz
<z|x/x =
(11)
(x)
In the probabilistic interpretation of AIS, we consider an extended state space proposal TRROP (Z [x). obtained by sampling from the initial No (z]x) and constructing transitions Tz (=1)) which leave Te- 1 (z(x) invariant. The target distribution pror (Zo.T(x) [x) (X) is given by running the reverse transitions Te (z (z,) starting from a target or posterior sample Ty (z|x), (x). , as shown in Fig. 2,
A18 (Zoer) (x) (Zoer == no (2o(x) II Te (z:
AIS (x, zo.r) )= = (x(xx 27) II T((z(-1/z)) (12)
Taking expectations of the log importance weights under the proposal and target again yields 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, Jaise (wo, T) and LAI (mo, T), which were suggested for MT 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 7). Assuming perfect transitions and geometric annealing path with linearly-spaced { (Br)= the sum of the gaps in the AIS sandwich bounds on MI, LAssE (To, T') Lasz (To, T',, reduces linearly with increasing T.
See App. D.I for a proof. our experiments in Sec. 5, we will find that this linear bias reduction in T's is crucial for achieving tight MI estimation when both p(z) and p(x|z) are known. However, can further tighten these AIS bounds using multiple anncaling chains (K > 1), and we present two practical extended state space approaches in the following sections.
3.2 INDEPENDENT MULTI-SAMPLE AIS BOUNDS
To derive Independent Multi- Multi-Sample AIS (IM-AIS), we construct an extended state space proposal by (k)
running K' independent AIS forward chains Zp:T" ~ T in parallel. Similarly to the IWAE upper bound (Eq. (5)), the extended state space target involves selecting index uniformly random, and running a backward AIS chain oo:T ~ picy starting from a true posterior sample Zy ~ p(z(xx. The remaining K samples obtained by running forward AIS chains, as visualized in Fig.
5

---

Published as a conference paper at ICLR 2022
(1:K) (uu
(x)
(1) sp ( (oar [x). k=1
II
PTOt
2022"
(1:K))
K pras
K Σ
II Expppe
K
(e) 2o. x),
(13)
where OPROP and Prer were defined in Eq. (12). 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 p(x), respectively,
AIS X. OoT7
1 K
Proo (x, Zzy Zo
(R)
log K
S logp(x) SE
log
(14)
(10))
(E) Queop (z u. x))
PIST (2:K) -ugt
Tenar 20. (x)
(k)
ELBO To, K, T)
EUBOD (x; To; K, T)
which again have KL divergences 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 K -> or T - co, and translate to lower and upper bounds MI in Sec. 1.1. In App. E.3, we show that the Independent Multi-Sample AIS lower bound on MI is limited logarithmic improvement over single-sample AIS, with L IIM-AS (TO, K, K) < LAIS (To, T) + log K.
3.3 COUPLED REVERSE MULPIISAMMLE AIS BOUNDS
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 define Independent Reverse Multi-Sami Sample AIS (IR-AIS) using the following proposal and target distribution, as shown in Fig. 2.
K
IR-IIS (x, Zo:T
(1:K)) =
Σ
AIS
ZooTT
x)
(k) AIS (x, Zo:T ) II 1=1,8
(1.K) IR-AIS (x, ZOOTT
II
AIS (x, Zo:T)
Note that partition function ratio is ZTor/ Zrror = p(x)* /p(x)x(1= p(x). Using these distribu- tions, 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 See. 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 single target sample 2r Ty TT(22xx (2)x,, with the remaining transitions y11 (Zo:T (TT) matching standard AIS in Eq. (12).
CR-AIS
Zo.T ëT, x) = NT (Zr, x)
II Pier (z xr,
(15)
The extended state space proposal is obtained by selecting an index s uniformly at random and running a single forward AIS chain. We then run K 1 backward chains, all starting from the last state of the selected forward chain, as visualized in Fig. 2
K
K II k= 1,
OPROP 2o
FR-AIS (1:K)
x)
OFROP
P [T)
(K)
(16)
K
s=1
Taking the expected log ratio under the proposal and target yields lower and upper bounds on log p(x),
-2 (1)
Log
AIS P
log p(x) -B
log
Drie (
Apen (50 (o0))
0.2 1 A01 Foorr-11 (x)
(3:A)
ELBOcK.AIS No. K,T)
EUBOCK (x; mo, K,T)
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 JeR-AIS (To, K,T) Z TAIsa (To, T) - log K.
3.4 DISCUSSION
Relationship with BDMC While Bidirectional Monte Carlo (BDMC) (Grosse et al., 2015; 2016) was the first method proposing multi-sample log partition function bounds using AIS chains, our probabilistic interpretations provide novel perspective BDMC. Perhaps surprisingly, we find that BDMC lower and upper bounds do not correspond to the same extended state space proposal and target distributions. In particular, the BDMC lower bound on log,p(x 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. Row 3).

---

Published as a conference paper at ICLR 2022
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, we show for IwwAR (go, K) (Cor. B.3), JIA,IS, (To. K, T) (App. E.3), and AIS (To, K, T') (App. H.3). Based on these arguments, a given hardware and allocating all remaining resources to increasing T"
recommend increasing K until computation can longer be parallelized on
Comparing Multi-Sample AIS Bounds In App. Fig. 5, we compare performance of our various Multi- 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 BUBO To, K, T) (RIIS 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 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.
4 MINE-AIS ESTIMATION OF MUTUAL INFORMATION
For settings where the conditional p(x l)) is unknown, we propose MINE-AIS, which is inspired by MINE but optimizes 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 Multi-Sample AIS methods from Sec. 3 to evaluate the bound. Consider an flexible, energy-based distribution To (z) x) as an approximation to the posterior p(z x) (Poole et al., 2019; Arbel et al., 2020)
(z)x (x) 2o (xxx2)
where Z = Eggo (eTo
(17)
Plugging Te Te, into the BA lower bound, we denote the resulting bound as the Implicit Barber- Agakov Lower bound (IBAL.), since it often difficult to evaluate explicitly due to the intractable log partition function term. After simplifying in App. J.1, we obtain
go (z|x)
Tx(x.)
I(x: z) 2 Jane (TT Epcx..) log
+ Ep(xa)
log
IBAL To), (18)
p(z)
E
Egg x) [eTa(x
Lung (go)
SPpck) (o)))
with the gap of the IBAL equal to Ep(x) D [Io (z|x)]] Note that the IBAL generalizes the Unnormalized Barber- Aaaker-Agakov bound from Poole et al. (2019), with (UB)(TT)) = (pAL(p(z), Ts).
Proposition 4.1. For a given (z(x)), the optimal IBAL critic function equals the log importance weights up to a constant T (x, z) = log pixzz
[o(2 +c(x). For this T*, we have IBAL(q,, = [(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 go(z)x) and (x, z), lim IGIAAE (go, To. K) = IBAL To).
K-+00
Thus, we may view the IBAL as the limiting behavior of the samile-sampl GIWAE bounds as K -> co. While Cor. 2.3 shows that Tawar, (ge, T') can improve I,,, (go) by at most log K nats, we show in App. L.I that the IBAL contrastive term is flexible enough to close the entire gap in the BA bound. However, this flexible contrastive term comes at the cost of tractability, as the IBAL in Eq. (18) involves log partition function log qo (x1x) [e (xx2)) compared to the finite-sum term in GIWAF
Energy-Based Training of IBAL Although the log partition function log (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 Φb respectively, we obtain
8 IBAL (go, TT) Bp(x).)
log qo (z) (x)
8 20
log 9o po(
(19)
20
a0
0 [gA TT) = pp(x)
To (x, x)
To (x, x)
(20)
03
D6
7

---

Published as a conference paper at ICLR 2022
Method
Propreul
MNIST-VAEIO
MNIST-YAEI10
MNNISTGAANI
MINST-GGANI0
Method
Proposel
CIFAR-GANI0
CIFAR-GAN100
AIS
M(6)
10.00,1920.85
[0.00, 786.12)
10,00,80 861.35)
AIS
100000 4035635 75)
(0.00, (853110.50)
(T=1)
g(4)p)
(21.0G, 63.00)
(34.49, 362.13)
[3.07, 314.723
[2.61, 513.33)
T=1
9/3.00
[17.300
4036779 22)
(20.17, 2378257.90) (0)
A1S
p(9)
39.09)
(79.9099.177) (5,17)
(21,57, 22.47)
(25.66 27.55)
AIS
[(2)
(20.52,
33080 (K0)
(104.51, 63290.40)
(T-500)
(34.16 34.29)
(80.19, 82.34)
(21,60, 23,06)
(25.58, 29.53)
T-500
(08.16, 196.15)
(145.19, 2786, 53)
AIS
p(0)
(44,21,34.21)
(50.78, 80.84)
(21.97, 22.02)
(26.47, 26.52)
((0)
(71.87, 73.98)
(480.26, 488 07)
(T-30K)
o(4)3)
(34.21, 34,21)
(80.77, 80.80)
(22.01, 22.01)
(26.53, 26.54)
T-100K
9(400)
(72.85, 73.54)
(479.27, 484.84)
IWAE
((9)
(0,00,8827.68
(0,00,11501999 (92)))
(0.00, 1630.00)
[0.00, 1740.29)
((2)
(0.000 77655005.50) 5.0)
(0.00, 90161
(K=1)
g(a)s)
(25.20,35.34)
(5.66)
[4.23, 67-47)
[3:23,260.8] 260.87)
(17.45, 77.52)
(20.00, 6346.85)
IWAE
p(a)
(0.91, 1197,7))
(230.19)
[6.91, 400.50)
(6,91,494799)
IWAE
((8)
10.91. 2064170.73)
(0.91, (856714455)
(K=IK)
(31 (31.60,85 34.24)
(61 (5..344))
(11.14, 62.73)
[10.14, 201.
K=1K
9/9/8)
(23.58, 74.00)
(26.98, 5283 13)
IWAR
pla)
(13,82,370:
(13. 82, 2247.73)
(15.81, 81.01)
[13.82, 114.01)
IWAE K-IM
p(a)
[13.82, 710511 (63)
(17.82, 19008044.00)
(K-MM)
(34.10, 34.22)
(08.35, 88.39)
(17.76630.880
(16.98, 68.04)
(30.73, 73,36)
(38.81, 0271..0)
Table 1: MI Estimation with IWAE (with varying K) and Multi-Sample ATS (with varying T) on MNIST (left) and CIFAR (right). Tight estimates, with gap of less than 2 nats, are in bold.
Eq. (19) indicates that to maximize the IBAL as a function of f and Φ. we need to increase the value of (x, z) or log (z)x) on samples from p(x, z), and lower on samples from p(x)TTe (z x). As is common in training energy-based models, it is difficult to draw samples from To (z(x). To reduce the cost and variance of the estimated gradient, we initialize chains from true posterior sample Zu ~ p(z[x) as in contrastive divergence training (Hinton, 2002) and run M steps of Hamiltonian Monte Carlo (HMC) transition kernels TIM (z) (o)zo.)) (Neal, 2011). See App. M. 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 (BAL( T)) using a Multi-Sample AIS lower bound on log Zg,0(x) but this does not ensure a lower bound on MI. In order to upper bound log (x,+(x) and preserve lower bound on 7(x: z), Multi-Sample AIS requires true samples from 00.0 x) to initialize backward AIS chains. Since these samples are unavailable, we instead initialize backward chains from the true posterior p( p(z(x) instead of Te (z x). We derive sufficient conditions under which this scheme preserves an upper bound on log Za. (x) in App. M.2 and provide empirical validation in App. M.3.
5 EXPERIMENTS
In this section, we evaluate our proposed MI bounds on VABS and GANS trained on MNIST and CIFAR.
5.1 SULTI-SAMPLE AIS ESTIMATION OF MUTUAL INFORMATION
We compare Multi Sample AIS MI estimation against IWAE, since both methods assume the full joint distribution available. For the initial distribution of AIS or variational distribution of IWAF, we experiment using both the prior p(z) and learned Gaussian go (z(x). Table summarizes our results.
IWAE As described in Sec. 2.3, IWAE bounds encompass wide range of MI estimators. The K= bounds with learned go (z)x) correspond to BA bounds, while for K > 1 and p(z) as the proposal, we obtain Structured INFONCE. While the IWAE upper bound on MI, which uses the log p(x) lower bound with independent sampling from qo (z) x). is tight for certain models,
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 close the gap from the BA lower bound (K = 1) to the true MI. For example, on CIFAR GAN 1 100, at least e 460 total samples are required to match the lower bound estimated by ATS.
Multi-Sample AIS We evaluate Multi Multi--Smmpl AIS with K 48 chains MNIST and K = 12 on CIFAR, and a varying number of intermediate T. We show results for the Independent 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 in stark contrast to the exponential sample complexity required for the IWAE M1 lower bound, and highlights that increasing T' in Multi-Sampi 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.
5.2 ENERGY-BASE
ESTIMATION OF MUTUAL INFORMATION
In this section, evaluate the family of GIWAE and MINE-AIS bounds, which assume access to 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 IGIWAE (go. To, 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 fixed go go(2|x). Cor, 2.3 shows BA, (go) S TaIAAE (go, To., Tor,K K) S JIWAR, (go, K) S TBA (go) log K.
8

---

Published as a conference paper at ICLR 2022
Iepur Used
Bound
Model
Lincar VAEIO
MNISST VAE20
MNIST GANZO
1x;;2)
3 www.lbieeiillleeeee
p(x.4)
Analytical
23.25
NIA
NI
IBAL goo
AIS Bound on Tree MI
(25.25, 25:23)
(65.11,65.177
153. 12, 53.50)
ptx, x))
TWAE LB (K-
1000)
20.53 + 2.66 = 23.19 [9|
36 21
6.00
05.11
20:97
+ 6.91 = 27.88
(MINE
Samples
LWA (K- 100)
21.01 + 1,50
23.14|
38,
4.61
63.47
20.56
4.00 25.06
Ploa bith and Mettattt (00
Structured InfoNCE LD (K- 1000)
6.91
6.91
0.91
E Ppccpp(abe) (xx)c aa (x)))]
Stractored TnfoNCF 1R (K 100)
4.61
4.61
4.61
AIS Bound IRAL (MINE-AIS)
(25.15, 23.15)
(57.72,57.74)
(40.70, (0.79)
Samples
AIS Hound IBAL (GWAL A
(24.7, 22.77
(40.17,444.0)
(28.01, 28.62)
log K
Luwas (oeeelgggo (go, R.
AIS Bound IBAL (lubNNCE K 1009
(11.38, 11,30)
(6.18,6.18)
(7.42) 7.42)
Geseralized TWAR (K 1000)
0.38 22.00
(57.23
6.00
43.76
20.50
6.77 27/222
Generalized [WAC 1000
22,48
0.39 22.87
37.50
41.90
20.68
+ 4.57 25.25
Lamaur 190. Top K)
Barber Barber-Agaloo (XC
22.09
37:99
21.42
2. Daallttr to (600 (e66ll
Jvine Samppes
InfoNCE (K 1000)
6.01
6.91
6.01
Inf-NCE L.B (K' 103)
4.61
4.61
4.81
Ina.
(b) Figure 3: (a) Comparison of energy-basee bounds (GIWAE and MINE-AIS) with other MI bounds.
(a)
(b) Visualizing the gaps of various energy- lower bounds and their relationships.
Although we perform separate optimizations for each entry in Fig. 3a, we find 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 learned T (x. z) instead of the optimal critic in IWAE (Cor. 2.2).
For the Linear VAF, the true posterior is in the Gaussian variational family and TA, (10) 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 (o (z x) and p(z]x). As K increases, we learn a worse qu in almost all cases, as measured by lower BA term. This allows the contrastive term to achieve closer to its full potential log K, resulting in higher overall bound. For more complex VAE and GAN posteriors, there a reduced tradeoff between the terms since the variational family is far enough from the true posterior (in reverse KL divergence) that either GIWAF or IWAE critic functions can approach log K improvement without significantly lowering the BA term. In all cases, (Structured) INFO-NCE bounds saturate to log K.
MINE-AIS Bounds We report MINE-AIS results using fixed Gaussian p(z) as the base variational distribution and Multi-Sample AIS evaluation of IBAL.(p(2) , Ta) We can see in Fig. 3a that MINE-AIS improves over BA due to its flexible, energy- based variational family. To evaluate the quality of the learned To (x, z). we compare the IBAL to the Multi-Sample AIS lower bound, which assumes access to p(x (z) and corresponds the optimal critic in Prop. 4.1. We find that MINE-AIS underestimates the ground truth MI by 11 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 more costly evaluation procedure, in Fig. 3a we report Multi-Sample AIS evaluation of IBAL TT) for qo. To learned by optimizing the GIWAE or INFONCE lower bounds with K = 100. GIWAE and INFONCE 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.
6 CONCLUSION
We have provided a unifying view of mutual information estimation from the perspective of impor- tance 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 Salti-Sammpe 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 infinite-sampi behavior of GIWAE and existing contrastive bounds. Our MINE-AIS and Saltie AIS methods highlight how MCMC techniques can be used to improve MI estimation when single analytic marginal or conditional density is available.
Note that IWAE and Structured INFONCE use the true log importance weights 7" (xx z) = log p(x.2)
(=)x) +c(x)
(Cor. 2.2). For this optimal critic, Ig,1,(go, T*) T(x;; z), go (Prop. 4.1) and AIS will sandwich the true MI.
9

---

Published as a conference paper at ICLR 2022
ACKNOWLEDGEMENTS
The authors thank Shengyang Sun, Guodong Zhang, Vaden Masrani, and Umang Gupta for helpful comments 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 Al Chairs program.
REFERENCES
Alexander Alemi, Ben Poole, Ian Fischer, Joshua Dillon, Rif A Saurous, and Kevin Murphy. Fixing broken elbo. In International Conference on Machine Learning, pp. 159-168, 2018.
Alexander A Alemi and lan Fischer. Gilbo: one metric to measure them all. In Proceedings of the 32nd International Conference Neural Information Processing Systems, pp. 7037-7046, 2018.
Alexander A Alemi, lan Fischer, Joshua V Dillon, and Kevin Murphy. Deep variational information bottleneck. arXiv preprint 1612.00414 2016.
Shun-Ichi Amari. Aipha-divergence is unique, belonging to both f-divergence and bregman divergence classes. IEEE Transactions Information Theory, 55(11):4925-499) 4931, 2009.
Michael Arbel, Liang Zhou, and Arthur Gretton. Generalized energy based models. In International Conference on Learning Representations, 2020.
David Barber and Felix Agakov. The im algorithm: a variational approach to information maxi- mization. In Proceedings of the 16th International Conference on Neural Information Processing Systems, pp. 201-208, 2003.
Mohamed Ishmael Belghazi, Aristide Baratin, Sai Rajeshwar, Sherjil Ozair, Yoshua Bengio, Aaron Courville, and Devon Hjelm. Mutual information neural estimation. In International Conference on Machine Learning, pp. 531-540. PMLR, 2018.
Stephen Boyd and Lieven Vandenberghe. Convex optimization. Cambridge university press, 2004.
Yuri Burda, Roger Grosse, and Ruslan Salakhutdinov. Accurate and conservative estimates mrf log-likelihood using reverse annealing. In Artificial Intelligence and Statistics, pp. 102-110. PMLR, 2015.
Yuri Burda, Roger B Grosse, and Ruslan Salakhutdinov. Importance weighted autoencoders. In International Conference Learning Representations, 2016.
Sourav Chatterjee, Persi Diaconis, al. The sample size required in importance sampling. The Annals of Applied Probability, 2892): 2018.
Xi Chen, Yan Duan, Rein Houthooft, John Schulman, Ilya Sutskever, and Picter Abbeel. Infogan: Interpretable representation learning by information maximizing generative adversarial nets. In Proceedings of the 30th International Conference on Neural Information Processing Systems, pp. 2180-2188, 2016.
Andrzej Cichocki and Shun-ichi Amari. Families of alpha-beta-and gamma-divergences: Flexible and robust measures of similarities. Entropy, 1236):1532-1568, 2010.
Kyle Cranmer, Johann Brehmer, and Gilles Louppe. The frontier of simulation based inference. Proceedings of the National Academy of Sciences, 117448):0055-30002 2020.
Chris Cremer, Quaid Moris, and David Duvenaud. Reinterpreting importance autoencoders. arXiv preprint arXiv: /ed:::2040000666 2017.
Justin Domke and Daniel R Sheldon. Importance weighting and variational inference. In Advances in neural information processing systems, pp. 4470 4479, 2018.
Arnaud Doucet, Will Sussman Grathwohl, Alexander G de G Matthews, and Heiko Strathmann. Annealed importance sampling meets score matching. In ICLR Workshop on Deep Generative Models for Highly Structured Data, 2022.
10

---

Published as a conference paper at ICLR 2022
Roger B Grosse, Chris J Maddison, and Ruslan Salakhutdinov. Annealing between distributions by averaging moments. In NIPS, pp. 2769-2777. Citeseer, 2013.
Roger Grosse, Zoubin Ghahramani, and Ryan P Adams. Sandwiching the marginal likelihood using bidirectional monte carlo. arXiv preprint arXiv:1511.02554 2015.
Roger B Grosse, Siddharth Ancha, and Daniel M Roy. Measuring the reliability of meme inference with bidirectional monte carlo. In Proceedings of the 30th International Conference on Neural Information Processing Systems, pp. 2469-2467, 2016.
Geoffrey E Hinton. Training products of experts by minimizing contrastive divergence. Neural computation, 14(8):1771-1800, 2002.
Sicong Huang, Alireza Makhzani, Yanshuai Cao, and Roger Grosse. Evaluating lossy compression rates of deep generative models. In International Conference on Machine Learning. PMLR, 2020.
Diederik Kingma and Jimmy Ba. Adam: A method for stochastic optimization. CoRR,
abs/1412.6980, 2014.
Alex Krizhevsky and Geoffrey Hinton. Learning multiple layers of features from tiny images. Technical report, Citeseer, 2009.
Dieterich Lawson, George Tucker, Bo Dai, and Rajesh Ranganath. Energy-insping models: learning with sampler-induced distributions. In Proceedings of the 33rd International Conference on Neural Information Processing Systems, pp. 8501-8513, 2019.
Yann LeCun, Léon Bottou, Yoshua Bengio, Patrick Haffner, et al. Gradient-based learning applied to document recognition. Proceedings of the IEEE, 86(11):2278-25
1998.
Chris J Maddison, Dieterich Lawson, George Tucker, Nicolas Heess, Mohammad Norouzi, Andriy Mnih, Arnaud Doucet, and Yee Whye Teh. Filtering variational objectives. In Proceedings of the 31st International Conference Neural Information Processing Systems, pp. 6576-6586 66866, 2017.
Alireza Makhzani, Jonathon Shlens, Navdeep Jaitly, Ian Goodfellow, and Brendan Frey. Adversarial autoencoders. arXiv preprint arXiv: w:Xi::551.00644 1.05644, 2015.
David McAllester and Karl Stratos. Formal limitations on the measurement of mutual information. In International Conference on Artificial Intelligence and Statisticx, pp. 875-884. PMLR, 2020.
Shakir Mohamed and Danilo Jimenez Rezende. Variational information maximisation for intrinsically motivated reinforcement learning. Advances in Neural Information Processing Systems, 28:2125- 2133, 2015.
Radford M Neal. Annealed importance sampling. Statistics and computing, 11(2):125-139, 2001.
Radford M Neal. Meme using hamiltonian dynamics. Handbook of Markov Chain Monte Carlo, 2011.
XuanLong Nguyen, Martin J Wainwright, and Michael I Jordan. Estimating divergence functionals and the likelihood ratio by convex risk minimization. IEEE Transactions on Information Theory, 56(1 (B((1))5847-5611 2010.
Sebastian Nowozin, Botond Cseke, and Ryota Tomioka. f-gan: Training generative neural samplers using variational divergence minimization. In Proceedings of the 30th International Conference on Neural Information Processing Systems, pp. 271-279, 2016.
Ben Poole, Sherjil Ozair, Aaron Van Den Oord, Alex Alemi, and George Tucker. On variational bounds of mutual information. In International Conference on Machine Learning, 2019.
Alec Radford, Luke Metz, and Soumith Chintala. Unsupervised representation learning with deep convolutional generative adversarial networks. arXiv preprint arXiv:1511.064:: 2015.
Avraham Ruderman, Mark D Reid, Darío Marcia-Gareía, and James Petterson. Tighter variational
representations of f-divergences via restriction to probability measures. In Proceedings of the 29th International Coference International Conference Machine Learning, pp. 1155-1162, 2012.
11

---

Published as a conference paper at ICLR 2022
Tim Salimans, lan Goodfellow, Wojeiech Zaremba, Vicki Cheung, Alec Radford, and Xi Chen. Improved techniques for training gans. In Advances in Neural Information Processing Systems, pp. 2234-2242, 2016.
Artem Sobolev.
Thoughts on mutual information estimation: More estima-
tors. Blog post, 2019. URL http://artem, www..ew..eo.e.................a. /posts/ html.
Artem Sobolev and Dmitry Vetrov, Importance weighted hierarchical variational inference. In Advances in Neural Information Processing Systems, volume 32, 2019.
Jiaming Song and Stefano Ermon. Understanding the limitations of variational mutual information estimators. In International Conference on Learning Representations, 2019.
Naftali Tishby, Fernando C Pereira, and William Bialek. The information bottleneck method. arXiv preprint physics/00040577, 2000.
Aaron van Oord, Yazhe Li, and Oriol Vinyals. Representation learning with contrastive predictive coding. arXiv preprint arXiv: 1807.03748, 2018.
Yuhuai Wu, Yuri Burda, Ruslan Salakhutdinov, and Roger Grosse. On the quantitative analysis of decoder-based generative models. arXiv preprint arXiv:1611.04233 2016.
Shengjia Zhao, Jiaming Song, and Stefano Ermon. The information autoencoding family: A la- grangian perspective on latent variable generative models. Proc. 34th Conference on Uncertainty in Artificial Intelligence, 2018
12

---

Published as a conference paper at ICLR 2022
Supplementary Material
Table of Contents
A General Approach for Deriving Extended State Space Bounds on Log Partition Functions
14
B Importance Weighted Autoencoder (IWAE)
15
B.I Probabilistic Interpretation and Bounds
15
Proof Logarithmic Improvement in K for IWAE EUBO
Experimental Results showing Logarithmic Improvement for AwAn (ga. K)
B.4 Bias Reduction in K for IWAE Lower Bound log p(x) Upper Bound MIl
19
B.5
Relationship with Structured InfoNCE
19
C
Generalized IWAE
19
C.I Probabilistic Interpretation and Bounds
19
C.2 GIWAE Upper Bound on MI Docs Not Provide Benefit Over IWAE
21
ELBO EUBO are Special Cases GIWAE Log Partition Function Bounds
22
C.4 Proof of Relationship between IWAE and GIWAE Probabilistic Interpretations
(Prop. 2.1)
Proof GIWAE Optimal Critic Function Logarithmic Improvement (Cor. and Cor. 2.3)
23
C.6 Properties of InfoNCE
25
D Single-Sample AIS
25
D.1 Proof of Prop. 3.1 (Complexity in T for Single-Sample AIS)
25
E Independent Multi-Sample AIS
26
E.1 Probabilistic Interpretation and Bounds
26
E.2 Proof Linear Bias Reduction in T for IM-AIS
26
Proof Logarithmic Improvement of IM-AIS EUBO
27
F
Reverse IWAE
27
F.1
Probabilistic Interpretation and Bounds
27
F.2
Improvement of RIWAE over ELBO and EUBO
28
G Independent Reverse Multi-Sample AIS
28
G.1 Probabilistic Interpretation and Bounds
28
G.2 Proof of Logarithmic Improvement in K for Independent Reverse AIS
29
Coupled Reverse Multi-Sample AIS
29
H..1 Probabilistic Interpretation and Bounds
29
H.2 Proof that CR-AIS Bounds Tighten with Increasing K
30
H.3 Proof of Logarithmic Improvement in K for CR-AIS ELBO
32
H.4 Proof of Linear Bias Reduction in T for CR-AIS
33
Comparison of Multi-Sample AIS Bounds
33
Generalized Mutual Information Neural Estimation (GMINE)
34
J.1 Probabilistic Interpretation of TBAI.
35
J.2 Probabilistic Interpretation of Generalized MINE-DV
35
J.3 Probabilistic Interpretation of Generalized MINE-F
37
K Conjugate Duality Interpretations
39
K.1 Convex Conjugate Background
39
Conjugate Duality Interpretation of IBAL
39
Conjugate Duality Interpretation of Generalized MINE-DV
40
K.4 Conjugate Duality Interpretation of Generalized MINE-F
41
13

---

Published as a conference paper at ICLR 2022
K.5 Conjugate Duality Interpretation of GIWAE, IWAE, and InfoNCE
42
L Properties of the IBAL
45
L.1 Proofs for IBAL Optimal Critic Function (Prop. 4.1 and L.1)
45
Proof IBAL as Limiting Behavior the GIWAE Ohjective as K + oc (Prop. I.22)
46
L.3 Convergence of GIWAE SNIS Distribution to IBAL Energy-Base Based Posterior
47
M MINE-AIS
48
M.1 Energy-Based Training of the IBAL
48
M.2 Multi-Sample AIS Evaluation of the IBAL
48
M.3 MINE-AIS Experiments...
51
N Applications to Mutual Information Estimation without Known Marginals
52
N.1 BA Lower Bound
52
N.2 GIWAE Lower Bound
53
N.3 MINE-AIS/ / IBAL Lower Bound
53
Experimental Details
54
Experiment Details Sec. 5.2
54
0.2 Experiment Details for Energy- Baaegy-Baeedd Bounds (Sec. 5.2)
55
0.3 Analytical Solution of the Mutual Information on the Linear MNIST-VAE
56
0.4 Confidence Intervals for Multi SullleSSmple AIS Experiments
56
A A GENERAL APPROACH FOR DERIVING EXTENDED STATE SPACE BOUNDS ON LOG PARTITION FUNCTIONS
In this section, give short proof that the gap in 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 pxx) using this approach throughout the paper and appendix, and provide visual summary in Fig. 4.
First, consider extended state space target Prer (x, Zeu) and proposal [xRop Rexi) dis- tributions. For all cases discussed in this work, we will choose our target and proposal dis- tributions such that log Znce(x)
log 1xx). For example, common construction is to have (xar(x) = PPer (x, Zeai )deea = pxx) and [x)or(x) = AeRoe (x, dex = 1. Our 'reverse' importance sampling bounds App. F-G construct target and proposal such that Zrar = p(x) and (x) p(x)* which still yields Zror (x) Z (x) p(x).
Zun(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
Pror(X, (X,, Zext
Zrer(x) Epror P (Roolx)
Prer (x, Zext)
E Engour (x)
log
< log ZyROr(x)
log
(21)
QPROP (x, Zext)
QEROP (x, Zext)
To confirm that these are indeed lower and upper bounds for any qraor and Prar- we can show that the gap in the lower bound Eq. (21) is the forward KL. divergence DKL 1qrzoe lProol] and the gap the upper bound is the reverse KL divergence, (rrrPrror [[Roor]]
Prer (x, Zexr)
(x)r(x) log Zpnop (x)
Zxor(x)
I B (Roal
DKL. [Jenoe DKL(genop (Zesi xx) Pror (Zesa xx] log Zrrop (x)
(22)
OPROP (x, Zuzs).
ELBO (pRoe-PrOT)
Puet (x, Zoa)) Orsop Zeex)
loog Zun (x)
Z (x)
F. E (Read x
DKL. [Pim ( (x)alg (zea x)]
log
(23)
ZpRoP (x)
ZpRor (x)
EUBO(x; Qunors Prui)
14

---

Published as a conference paper at ICLR 2022
Thus, the bounds in Eq. (21) directly generalize the standard ELBO (log p(x) DKL <>)ppppz x)]
and EUBO (log p(x) + D [x( (2(x))) which appear as special cases when = T.
the proposal distribution is GRROP (z/x) go and the target distribution is Prar Prar(z|X) (z|x) p(z x).
In what follows, our extended state space proposal or target distributions may include [x(2|x) as
an initial base variational distribution, with the posterior p(z|x) often appearing within target distributions Prer(X, Zex)).
In Fig. 4, we summarize various extended state space proposal (third column) and target distributions (fourth column). We emphasize that the base variational distribution (a(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 ATS 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. as special cases of this general approach, thus interpreting each importance sampling bound in terms of probabilistic inference in extended state space.
B IMPORTANCE WEIGHTED AUTOENCODER (IWAE)
B.1 PROBABILISTIC INTERPRETATION AND BOUNDS
Consider a K K-sample proposal distribution PREOP (z (2:))) using independent draws from an initial distribution to (z)x). The target distribution is defined as uniform mixture of K components, where
each component replaces the initial sample qe (z( (*) (x) in index k with sample from the target distribution p(z x).
K
IWAE (x, z (1:K))
Σ p(x, z ())
II go (x))
IWAE (z (1:K)
(x) =
II
(z (s) lx),
(24)
Note that the normalizing constant of p1a*=(x (x, 2 (2:K)), or "pteaExx ((2:K) (2:2 equals p(x)
since go (z) x) is normalized and f p(x, z)d2 = p(x). To connect this with the TWAF bound, we show that the log importance ratio reduces
1
K p(x, (0)
Σ
II (x)x())))
PTGT X 2 (1:K))
log
IWAE log
(2( x)
K 11 (0(z(()(x)
1 K
p(x, (8))
log
(25)
4o (x( (x)
As in Eq. (21), we can obtain lower and upper bounds log p(x) by taking expectations under the proposal and target distributions, respectively.
Alternative Probabilistic Interpretation We now present alternative probabilistic interpreta- tion of IWAF, which is similar to Domke & Sheldon (2018) and will be used as the foundation for our GIWAF bounds in App. C. Consider the following extended state space target distribution,
K
1 K p(x, z (o)
Proy (x, (I:K),
II (x))
xx)
(26)
Note that marginalization over s leads to the IWAE mixture target in Eq. (4) or Eq. (24). We consider the extended state space proposal
Greop (z) (1:K)
,//x)
[I
(z (1)
1x)
JWAE (s]z
(27)
km
15

---

Published as a conference paper at ICLR 2022
Practicality
Extended State-Space Proposal
Extended State- Space Target Upper Bound on log p(x) (EUBO) Lower Bound 1(x; 2)
Lower Bound on log p(x) (ELBO) Upper Bound [(x;
Simple
IWAE
x-
PN
Gemeralized IWAE
Reverse IWAE
Single-Sample
AIS
Independent
Multi- Multi-Sample AIS
#
+
Independent
Reverse
Sample
+4
Coupled Reverse
Multi-Sample AIS
Comparison of the probabilistic extended space interpretations of different multi- sample bounds. Forward chains of AIS and variational distributions in is / IWAE are colored in blue. Backward chains in ATS or posterior distributions in IWAF are colored in red. Note that lower bound on log p(x), translates to upper bound on MI, and vice versa.
Figure 4:
where we have defined (.WAR (s22(())) (1:K)()))) ;x)
p(xc,x (1)
(x(2()) (x)
(x)))))x))
As desired, the log importance weight match Eq. (25)
p(x, (())
II
qo (z
log
DEAT
x, z (1:K) .)
log
kyda
log
p(x, x (2))
(28)
C (1:K), s|x)
)(x))()))
To (z) (x)
(1) (x)
(p(ni=) xx
(())
16

---

Published as a conference paper at ICLR 2022
Lower Bound log p(x) and Upper Bound on MI Using the general approach in App. A, taking expectations under IWAR leads to a lower bound on log p(x)
K p(x, 2 (11))
ELBO Qo. K) = E
log
(29)
Qa (x((0))
K
go (2((1)))
k=1
This corresponds the following upper bound on MI for known p(x(2)
I [xx;z)= Iiaar (90, K) (pp(x.)) [log((x/2)] Ep(x) [ELBO (x; yo, K)J
(30)
log 1 p(x, x K (z)
= (x(x)) [ogp(x/z)] - p(x) fi (x(1)))
(31)
Upper Bound on logp(x) and Lower Bound on MI Similarly, with expectations under Prote (z( (z(I:)(px, we obtain an upper bound on logp(x), Since, for independent draws, the uni- form mixture Prat is invariant to index permutations, we may choose the target sample to be z11) and obtain the upper bound of Sobolev Vetrov (2019):
K p(x, (*))
EUBO (x; qo: K)
E p(a())) (x) II (x(2)))
log (((())/x)
log
(32)
Translating this to lower bound on MI with known p(x(2),
T(x; 2) 2 JIWAE, (go, K) Ex(xx.)) [logp(x(z)) Byco) (EUBO (x; (x; qo, K]]
(33)
K
p(x, (k))
(p(x) [log((x(z]] p(x)p(a(11p)) Ii co(a()))))) k
log
2 K
(34)
g0 (z(k) xx)
k=
B.2 PROOF OF LOGARITIMMIC IMPROVEMENT IN K FOR IWAE EUBO
We first recall results that IWAE bounds tighten with increasing K.
Proposition B.1. For given (z(x). ELBO (x: qo. K + 1) 2 ELBO qa. K) (Burda et al., 2016) and EUBO)WAE(X; qo, + 1) S EUBOJWAE (x; qo, K) (Sobolev & Vetrov, 2019).
Proposition B.2 (Improvement of IWAE and ELBO or EUBO). Let prer (IAAE(ssx (1:K)
(() (x(x())))) EK qa(x (xxx denote the normalized importance weights and U(s) indicate the uniform distribution over K discrete values. Then, we can characterize the improvement of
((x.x())
ELBOIWAE (x: qe, K) and EUBOIWAE(X; qo, Ko, K) over ELBO(x; qo) and EUBO(x; qe) using KL diver- gences, as follows
ELBOIWAE q0, K) ELBO (x; qa) + Expur (2(1 (2(1-)) (x())
DKL(U(() (Ip WAA (s/z (1)2( x)]
(35)
0 s of uniform from SNIS weighty S DxL Dxi (z)x [|(()
(1:00, EEBOO x: qo) E (2(1))
DKL [piWA (s/2 (2)2((: x) [{{}]]
(36)
0 KL of SNIS weights from uniform S log K
Prop. B.2 demonstrates that the improvement of the IWAE log partition function bounds its single-sample counterparts is larger for more non-uniform SNIS weights. Notably, the improvement of EUBO xx; ge, K) over the single-sample EUBO(x; qe) is limited by log K.
Proof. We first note that the single-sample ELBO(x x; ga) and EUBO(x; qa) are special cases of ELBO (x; qo, To, K) and EUBO ge. To. K) (Lemma C.1) with Te = const. As a re- sult, the gap between ELBO(x; qe) and ELBO xx; qa K). for example, follows special case of the gap between ELBO (x: qo. Ta. K) and ELBOOWAr(X; qo. K). which we characterize in Lemma C.2 (App. C.4). The result in Prop. B.2 follows directly.
17

---

Published as a conference paper at ICLR 2022
We now justify the range of the KL divergences in Eq. (35) and Eq. (36) referenced in the underbraces.
Improvement of IUBO Note DxL (y (s|z(1:K), [)][[[(s]] = log K - H (ptare (s (1)2(1)),, x)) is bounded above by log K since the entropy of a discrete random variable is nonnegative.
Thus, the improvement of EUBO (x: qo. K) over EUBO(X; qo) 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 ELRO,WAE On the other hand, the KI. divergence DKL (o,()((())))) ([)))]lp (s) (z)z( x)] is not limited by log K. However, we do know that the improvement of ELBO (x; qo, K) over ELBO qe) will be limited by DKL [go (z|x) [[p(z|x)], the gap of ELBO(x: qo).
Prop. B.2 shows that both IWAE upper and lower bounds on log p(x) improve upon their single- sample counterparts, with the improvement of EUBOIWAE (x; qo. K) over EUBO(x; qo) 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:
IBAL (go) S /TAAE (qo, K) S IBAL (qo) + log K,
IIWAE (ge, K) S IBAU (ge).
(37)
Proof. Recall from Sec. 1.1 and Sec. 2.2 that TAA, (go) Ppxx ((((2)] Exyx) [RUBO {o)], and TIWAEL (go, K) 2x [log p(x|z)] p()) [EUBOIWAE (x; qo, K)]. Using Prop. B.2, the fact that EUBO(X; yo) log K S EUBO (x; qo, K) for any x implies that
IIWAE, (go, K)- IBAI (go) = Ex(x) [EUBO(X; go) - EU qo, K)] S log K,
(38)
which results in TIWAE (go. K) S IBA (go) + log K, as desired.
IAA (g0) S IIWAE (go, K) and IIWAE (go, K) S Ia (90) follow from the fact that IWAE bounds tighten with increasing K in Prop. B.1.
B.3 EXPERIMENTAL RESULTS SHOWING LOGARITIMIC IMPROVEMENT FOR LIWAR (qo, K)
In Sec. 2.4, we showed that IWAE is a special case of GIWAE, which decomposes into the sum of variational BA lower bound and K -ample 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
((xxx(1))
go (z|xx) p(z)
No (
IWARz (go, K) = Ep(x,)) log
+E
log
(39)
s(x))(e()))))
fi (a()(()
((x)
K
In)
D. contrasive temn Slog
This way of writing JIWAE, (go. K) provides additional intuition for the result in Prop. B.2 and Cor. B.3, In particular, the improvement of LIWAE, (ge, K) over Luzz (go) is simply the contrastive term, which is limited to log K.
Mathed log K
Priginal
VAF2
VAEIO
VAR100
GAN2
GANI6
GANI0I
IWAE
2(a)
0+--- 0
0
0
0+0-1 D
0.400
0- 0-
siels
8.0
0
25.20
01.44
8.3+0=
8:53
4.23++0 4.28
3/3
0.223
IWAE
p(4)
0.46.11
681
0.4 6.01
0.6601
6.91
0.4688
6.01
046691 6.91
(K-IK)
9.0 25.20
31.00
44.44
51.44
7.82
10.74
+
11.14
3.23 6.01
10.14
(K-IM)
1382
p(9)
18.76
0.09.00
000
9.09
25.30
8.50
44.54
13,82
1881 58.85
1882
6.02 10.81
1077
4.23
13.08
17.70
817
18.80
15.90
16.98
Table 2: Decomposition of TwAaz (ge, K) into BA term and contrastive term (< log K) on MT estimation for VAE and GAN models trained on MNIST.
Table 2 and Table show the IWAE objective decomposition to the BA term and the contrastive term, on VAES and GANS 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 VAES and GANS with two dimensional latent spaces, the contrastive term may contribute notably less than log K. In these cases, even the optimal critic function 7" (x( z) log p(x.) as used in IWAF, has difficulty distinguishing posterior and variational samples.
(x +c(x),
18

---

Published as a conference paper at ICLR 2022
Method
log
Proposal
GANS
GAN10
GANTOO
IWAR
p(R)
0. 0 -0-0
0.4
0.400 0
K=l
9(9
14.53 0 14.53
17.45 4 0 17:45
20.00 4 20.00
IWAR
((2)
6.91
K-lk
9i9e x)
14.53 + 6:00
21.43
16.68 6.9 23.08
20.07 6.9.1 26.88
IWAE
K=IM
13.82
p(2)
13.82 = 13.82
0++ 13.82 = 13.82
0.4 13.82 13.82
o(9 x)
14.53 13.81
28.34
16.92 + 13.81 = 30.73 20.00 13.81 33.81
Table 3: Decomposition of AwAcz (ge, K) into BA term and contrastive term (< log K) on MI estimation for GAN models trained on CIFAR-10.
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; more complex and more easily distinguishable from the variational Q (z x). These results highlight the inherent exponential sample complexity of the IWAE lower bound on MI.
B.4 BIAS REDUCTION IN K FOR IWAE LOWER BOUND ON log p(x) UPPER BOUND ON MI
We have seen in App. B.2 Prop. B.2 that the improvement of EUBO (x; Qo. K) and LIAAz (x; qa, K) over the single-sample EUBO(x; x; go) and IuAz (x; qo) is limited by log K. This suggests that expo- nential sample complexity in the gap of the EUBO, K x exp (xxxx/// is required to obtain a tight lower bound.
The quantity DRL [x) P (s| x)], which measures the improvement of the IWAE lower bound on log pxx) over the ELBO, is not explicitly limited by log K. However, Chatterjee et al. (2018) suggest that the same exponential sample complexity, K ε exp x) go (z x []) is required for accurate importance sampling estimation with proposal qo (z) x). Maddison et al. (2017) 1); Domke &
Sheldon (2018) find that the bias of the IWAE lower bound the limit of K + oo reduces at the rate Varl p(xx x)
of O 2K go (x)x) 1). although the Varl p(x.x)
term is at least exponential in (xL[p(z)x)ll (glp
(x))x)
(Song& & Ermon, 2019).
This exponential sample complexity for exact estimation of log p(x) or T(x, 2) is usually impractical for complex target distributions and limited variational families, where DRL. [p(z, 4o (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 (z(x)) and p(z]x).
B.5 RELATIONSIIIP WITH STRUCTURED INFONCE
We can recognize the Structured INFONCE upper and lower bounds for known p(x|z) (Poole et al. (2019) Sec. 2.5) simply applying the standard IWAE bounds, using the marginal p(z) in place of the variational go (z/x)
log p(x)2(1))
1
E (())p(x(1) xy (()))))
S ((x;2) S. p(x) p(alk))
log
(x(2))
Σ pixtali)) k=
# Σ p(x] kmm
We refer to the lower bound as IS-INFONCE, (K) and the upper bound as Is IS-INFONCE, (K). From Cor. B.3, we obtain an alternative proof that the Structured INFONCE lower bound is upper bounded by log K. Since Inaz (p(z)) < IwaRz (p(z), K) < Inna (p(z)) log K and the BA bound with a prior proposal equals O from Eq. (3), we have that 0 < /S-INFONC,, (K) S log K.
C GENERALIZED IWAE
C.I PROBABILISTIC INTERPRETATION AND BOUNDS
To derive a probabilistic interpretation for GIWAE, our starting point to further extend the state space of the IWAE target distribution in Eq. (4), using uniform index variable p(s) K Vss that specifies which sample (8) is drawn from the posterior p(z|x). This is shown in Fig. 4, and leads to
19

---

Published as a conference paper at ICLR 2022
joint distribution over (x, z (11:k),)) ,8)
K
ProT (X, z (I:K) ,s
1 K
p(x, z )
[((
q(z (* xx)
(40)
with marginalization over leading to the mixture in Eq. (4). The posterior over the index variable s, which infers the 'positive' sample drawn from p(z|x) given a of samples z (1::)), corresponds to the normalized importance weights
p(x.0 (
Prot
PIOTAR P z (1:K)
Go (x(()))
p(x, p(x,a((k))
(41)
Σ k= do (2(F))x)
which can be derived from Prot
(s|x, P (1l:
DIWAR
Prot (x. (1:46)
* p(xxx ) Il (0(2 lx)
pONAA (pc (1))))
Σラ= p(x26)) II (z(z [x)
For the GIWAE extended state space proposal distribution, we consider categorical index variable APROP (a/= (s/z(1:K) x) drawn according SNIS, with weights calculated using the critic function To.
K II (1(o(())))))
QEROP
GIWAE
(1:R)
.8/x)
SPROP (s/2(1) (1:K) xx).
(42)
(s/z) (1:K) x) (x.,(on))
where
TEROP
(43)
Tolxxa(t))
Similarly to Lawson et al. (2019), the variational SNIS distribution SPROP (s|z(1:K), ,x) is approxi- mating the true SNIS distribution Prer (s x, (1::)). As we show in App. C.4, the optimal critic function is Tx(x, z) log App. B.1).
g(*[x] + c(x), which recovers the IWAF probabilistic interpretation (see
p(x=)
Log Importance Ratio To derive bounds on the log partition function, we first calculate the log unnormalized density ratio
K
1
p (())
II
(2) x)
GIWAE (z (1(:1),,,,) 8, log log Orsor (z (2(:k),
Pier
k
(44)
K q(2(4) (x)
((1/2(1:33, x)
k=1
1 Σx Txx,0
p(x, Z ((t))
log K (xx,2(())
(45)
((2(0) [x)
x(x, log
T'xx 2()) ) log
1
Tixa
(46)
{(((+))
K
Taking the expectation of the log unnormalized density ratio under the proposal target distribution yields a lower or upper bound, respectively, on log p(x)
GIWAF
Zz (1:K), s,x)
CIWAA (2 (1:K) 8, log TVROP
E
log
PTGT
< log p(x) < E,graa
(47)
GIWAF
(1:K), s(x)
[2(1KK),
GIROP
ELBO To, K)
EUBO (x; go, To, K)
As in Sec. 2.1 and App. A, the gap in the lower and upper bounds can derived as KL divergences in the extended state space.
20

---

Published as a conference paper at ICLR 2022
Upper Bound on log p(x) and Lower Bound on MI upper bound on log p(x). we write
To derive an explicit form for the GIWAE
EUBO (go, To. K) = B yours ( (l)x))
log pxx, 2(x,2!0) g(2(0))xx (x)
T(x, (11) + log
1
T (xxx)(1))
K k=
p(x, z) Rp()/x) log
e )x((x(011))
(z(xx)
B pp(((1)))))
IT ((2(0)px)
log
(48)
kw2
(x,2(1))
K
where, in the first of the second line, we note that piat (z (l))))),)))))))))))))))))) s|xx specifies that z(*) p(z) x). is sampled uniformly, we can assume = in the second term due to
Since ~ p(s) K
permutation invariance.
Translating this into lower bound on MI, we consider ((x:2) = - Pp(x) [lo]] Hx > -Ep(x) [EUBOGWAE To, K)] - H (x|z). Writing the conditional entropy term over the index 8,
(x: 2) 2 (log
(49)
log p(x, z) (z x)
log e (x,x
(p(x))(2))))
E p(x)p(x()))) 11 p(((1))))) l0
K
= -ppx..) log g( g(2|x) x) p(z)
loog
(xx (x),
+ E
(50)
p(x)p(a())x))
go (x) km
Di e (x(,((())))
K
which matches Eq. (7) from the main text.
C.2 GIWAE UPPER BOUND ON MI DOES NOT PROVIDE BENEFIT OVER IWAE
To derive an explicit form for the GIWAE lower bound on logp(x)
log p(x, z(²)) ((x()))x)
(x, z))) +log Σ (x.a(()))
ELBOutWA (x; go, To, K)
E area (0100)
(51)
Translating this to an upper bound on MI yields
K
To (o(x2()))
log p(x,2(1) 2(11)
1xx: 2) S (xcet)))x)) (ogp(((2((1)
(52)
of (x))
K
(x
g( (x)
km
log Tox,allly
F. mvan (n (2(1:))) (lox)
K Σ Tonaatti
Thus, in Eq. (52), knowledge of the full joint density is required to evaluate both the conditional entropy and log (x2 terms. If both p(z) and pxx|z) are known, then we will show in Cor. 2.2 below that the optimal critic or negative energy function in GIWAE yields the true importance weights, log p(x) bounds matches the IWAE bounds.
p(x,z(1)
and the resulting
We thus conclude that the GIWAE upper bound on MI and lower bound on p(x) does not provide any benefit over IWAE in practice. However, ELBO (x; qa. Tg. K) is still useful for analysis, as our proof of Lemma C. below allows us characterize the gap between ELBO, (x; go. K) and EI,BO qo) in Prop. B.2.
21

---

Published as a conference paper at ICLR 2022
C.3 ELBO AND EUBO ARE SPECIAL CASES OF GIWAE LOG PARTITION FUNCTION BOUNDS
Lemma C.I. The single-sample ELBO and EUBO are special cases of GIWAE, with
(LBO (x; qo) ELBOGIWAE (x; q0, Too const, K),
EXBO q6) = EUBOGIWAE (x; q0. Too = const, K).
In both cases, the SNIS sampling distribution (Eq. (43)) is uniform CIINAE (z (1:K), s|x) = U(s).
Proof. We consider the GIWAF probabilistic interpretation (App. C.1) for Too = const. We refer to this extended state space proposal as SpRor= since it leads to TEa, (go) and TAN (go) bounds on MI.
K Arrop (1:K) .s|x) = II
go (2(21) km
- (pao0(s|2 (1/2(1: x).
(53)
(x
where
UA (s 1z (1:K)
Σ @ (x.(()))))) = U(s) K
(54)
Note that the SNIS sampling distribution ArooP (s|z(1.K), (s in the GIWAE extended state space target distribution
x) will be uniform, which matches p(s) = K
1 Prot (g (1:K) ,|xx) k p(((1) x) II (x)(2())
(55)
Now, taking the log unnormalized importance weights, we obtain
R p(((())
xx)
go
(x)
Piut (z. (1:K), ,//x)
llog p(x, p(x,z!)) x))
log
log
(56)
RA (2( (I:K), s(x)
K
go [0(2)=)|x)
(2(())))
1=1
Taking expectations with respect to PPRO (z (1:K), and EUBO(x; qo) respectively, as in Sec. 2-2.2.
, s(x) or (2:K), s/xx leads ELBO go) p
(1:K),
C.4 PROOF OF RELATIONSHIP BETWEEN IWAE AND GIWAE PROBABILISTIC INTERPRETATIONS (PROP. 2.1)
We first prove 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 ELBOqiaE and EUBO to ELBOJWAE and EUBO|WAE-
Lemma C.2. We can charucterize the difference between IWAE and GIWAE bounds on log p(x) using KL divergences between their respective SNIS distributions.
BLBOrAAn (x: do, K)= BLBOGIWAR X: do. To. K) (2 (x)
Det. Orpor (s) (1:K), xx] pier ( (1:K) x)]
(57)
EUBO]WAn(x; go. K)= EUBO go. To. K)
Dxr, [Pror (1:K) x)) || some (1:8)
x3]
(58)
Proof. Recall from Sec. 2.1 or App. A that the gap of the lower bound ELBBO (x; qo, Tep K) is DKL. COWAE
[ (2 (1:K), 8) xx) GIWAE (2) (1:K), x)). while the gap of the upper bound EUBO (x; qo, To, K) is DKL lx GIwAF (z(1:K) s|xx)]99888 (lg (Z:KK), 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 8, or target SNIS distribution, is
1 p(x, a (())
1
go (x()))
(())
(s/z Pror (I (1:K)
x)
(z (I:K) 8.
g(x(=))x) xx)
Σ
K
Pier
UIWAR
(z (21 x.))
K
p(x, ((0) II
Co (x(1))
Σ
K
p(xx(())
kyp
22

---

Published as a conference paper at ICLR 2022
The joint distribution then factorizes PIGTAR (z s. pior PioraF(z(1:33)
(z(1:k),
(z(1:K), x) p (s|z(1:)) (1))
ELBO Case: Using this factorization of prot (z , ELBOG (x; qo, To, K) as follows
rewrite the gap of
Dxt (9nce (x (x(1:)), xo (z(1:))) s|x)]
(99)
amas (o (a(::)))
log
Treor
Prac (xx ALLRY px) POERA (n/on 21
GEWAR
(2 (1.K (x) (ela) (1:K) x))
(
x II go (2 (x)
DRi
(x))
K pla
ga ga(w (x) + souaa Sehro (el (1/1 .x)llmm (elz
x)]
= DKL pas (ze
(1:13) 1x)llptor ras (x (1:K) (p)) (xc] + E Exance DKL (ereor (s/22 (1=K) Prar P (1/2 xx1]
(60)
where we can recognize the first term that ELBO (x; qo, K) ELBOawar (x; To, K) = Dxt. 9esor (z (2:K), s/xx) (ppior (1:K)
the gap in ELBO (x; go, K).
Noting
s|x)]
Dxi [ypnor (x (I:K) xx) saar (1:K) [x]] we obtain Eq. (57), as desired
ELBO (x; qo, K) ELBOaaan qo, Tøp K) (2(1))))))))))
D Ia (1 (5/2(280) x) ( (1)2 xi)
EUBO Case: For ELBO (x; qu. TV. K), the derivations follow in similar fashion to Eq. (59)-
Eq. (60), but using the reverse KL divergence and expectations under poWAE (
(I:KK), s(xx.
Translating Lemma C.2 to a statement relating IWAE and GIWAE bounds on MI, we obtain the following proposition.
Proposition 2.1 (Improvement IWAE over GIWAE). For given (z(x)x) and any Tx(x,, z).
TWAEz (go; K) = IOIAAR (go, To: K) + P. (pxe)ssaa | Logme (s|z (1:K) x))lqrum
orNaR (s) (s/z (1:K) x)]
Proof. The result follows directly from Lemma C.2. First, note that per
P (x, z (I:K), s) =
DIWAR (x, (1:K), s)) Then, using the upper bounds on p(x) and taking outer expectations with respect p(x), we have
TIWAR (20. K) Tamas, (20, To, K) - -H(x/z) By(x)) [EUBOIWAE go, K)]
(61)
Hx(22)) [pp()) [EUROGIWAE qu, To, K)]
B (2 xx)
DKL IWAE
(s/z (s|z (1:K)
x) GIWAE (s (1:K)
x)]
(62)
C.5 PROOF OF GIWAE OPTIMAL CRITIC FUNCTION AND LOGARITHMIC IMPROVEMENT (COR. 2.2 AND COR. 2.3)
We now prove Cor. 2.2 and Cor. 2.3 from the main text, with corollary stating the results for the special case of INFONCE in App. C.6.
Corollary 2.2. For a given (z)xx) and > the optimal critic function is the true log importance
weight up to an arbitrary constant: T (x, z) log p(x.x)
(2]x] + c(x). With this choice of T (x, z),
IGIWAE (qo. T K) IIAAE (go. K).
Proof. Using Prop. 2.1, we can see that the gap the GIWAE and IWAE bounds, which corresponds to the posterior KL divergence over the index variable will equal zero iff
((x,(1))
()x
GoRMA (s/z (1:K), x) =Pror
GIWAE (s|2 (1:K) x)
20 (0(2
(63)
((x((())
Σ
K
p(xx(())
(x)
23

---

Published as a conference paper at ICLR 2022
This condition also ensures that the overall GIWAE proposal piGrAs (z(= s/xx) (Eq. (40)) and target ( s|x) (Eq. (42)) distributions match. We will show that any T(x, z) which satisfies Eq. (63) has the form
p(x, z)
T (x, = log (z(z)x)
c(xx).
(64)
Let fxx, z) = log p(xxx) + g(x, 2). which represents an arbitrary choice of critic function. We will show that g(x, z) must be constant with respect to
(x(x)x)
log p(x r(m, +x(xx
p(x22 (*)
go (2 (x)
log p(x 90(xl0)pp) (x(xx
K
p(x.a(1))
(z(110)) x)
log pixa(e))
K
p(x, (k))
p(x, (+))
log pix +z(x,x(1))
e (x(
(0(xx (1))
(o(u())
4o (x(x1))))
2o [o(())) x)
p(x, ( p(x, 2 K
K
pxx, kkk))
p(x, 2())
c0(x,((1))
colxaal)
Am 1 (0( x) (0(z(>)|x)
(z)) (x) qo (z(x)
g(x, 2) c(x)
where g(x, z) = c(x) is required in order to ensure that g(x, (n) = g(x, z ((1)) for arbitrary choices of z samples.
This form for the optimal critic function (x, z) in Eq. (64) implies that learning To (x, z) in GIWAE becomes unnecessary when the density ratio p(x.z) is available in closed form, as assumed in the go (m|x)
IWAE bound.
For this choice of (x, 2). the value of the GIWAF objective matches the TWAF lower bound on MT.
log ((x,a(1))
log g(2)x) p(x)
11x(0)x0) (x)
I Tawas (go, T" K) Bx(xx))
log
((x,a()))
9 (x0)))))0)) k
log pi ((0)
(o(ta()) (x0 (0(x))
log glaxx)
+ log p(x, z)
p(x, z 2(1) K
= -p((xx))
log
(65)
p(z)
a(x/x}]
())) ff k=
(x(2()))2))
1=1 qe (z(()) xx)
K
p (
(x(())) +
-I8
log
((x)p(2(0))))
K
2o (z)(N) xx)
= TIAAR (go, K).
The second term in Eq. (65) is exactly the negative of the IWAE upper bound on log in Eq. (34). Combined with the entropy -H(z)), we obtain the TWAAE (40, K)) lower bound on desired.
Corollary 2.3. Suppose the critic function Tx(x, 2) is parameterized by G. and 3 s.t. V (x,z), To, (x, const. For a given qe (z)x), let To. (x,z)) denote the critic function that maximizes the GIWAE lower bound. Using Cor. 2.2, we have
IBAL (go) I IGIAAE (go, To., K) S LIWAR (q, K) I IBAz (go) log K.
(9)
Proof. We begin by showing that JAA (go) S TGIWAE (go, K.) This follows from the assump- tion that there exists do in the parameter space of the neural network To such that T const. With this Φo. we would have JBA, (ge) = TGIWAE, (49. Ton K) as in Lemma C.I. Thus, the optimal Φ" in the parameter space can only improve upon the BA bound.
Next, for any given & (including Φ²), we have IGIWAE (yo, To, K) S IIAAR (4o, K), since IWAE uses the (unconstrained) optimal critic function T log p(xxx) + c(x) (Cor. 2.2). The final ga [2|x]
inequality follows from Prop. B.2, which shows that IIWAR, (go, K) improves by at most log K over TBA, (go). These relationships are visualized in Fig. 3b.
24

---

Published as a conference paper at ICLR 2022
C.6 PROPERTIES OF INFONCE
Corollary C.3. Using the prior qa(z|x) p(z) as in INFONCE,
(a) For the optimal critic function, Cor. 2.2 implies T T* (x, z) lxgp(x(z) + c(x), and
IGIWAE LG(WAELpp(z).... T K) = INNONCE, (T" K) = IS-INIONCEL (K).
(66)
(b) For an arbitrary critic function To z). Cor. 2.3 translates to
ANTONCE, (Tg. K) S IS-INIONCE_ (K) S log K.
(67)
For INFONCE, note that using the prior as the proposal does allow the critic to admit an efficient bi-linear implementation Tx(x, z) = for (x)" for (z), which requires only + K forward passes
instead of N for GIWAE, where N is the batch size and K is the total number of positive and negative samples.
D SINGLE-SAMPIE AIS
D.I PROOF OF PROP. 3.1 (COMPLEXITY IN T' FOR SINGLE-SAMPLI AIS)
In this section, we prove Prop. 3.1, which relates the sum of the gaps in the single-sampi AIS upper and lower bounds the symmetrized KL divergence between the endpoint distributions. We extend this result to IM-AIS in Cor. E. and CR-AIS in Cor. H.3.
Proposition 3.1 (Complexity in T). Assuming perfect transitions and a geometric annealing path with linearly-spaced (BoYE=) the sum of the gaps in the ATS sandwich bounds on MI, TAISy (To, T) Ass (mo, T), reduces linearly with increasing T.
In particular, we will show that
EUBO To. T) ELBOAIs(X; no, To. T) S T (2) [m)( + DRr. [xr[[=[[2]][ (x|xx)
(68)
where the left hand side also corresponds to the sum of the gaps DKL [xkoor (Zo:T x)||ptr (Zo.T x)] +
DKL. [ (Zo:r) xx) f (Zo:r [x)]. This result translates to mutual information bounds in Sec. 1.1.
In contrast to Grosse et al. (2013) Thm. 1, our linear bias reduction result holds for finite T.
Proof. For linear scheduling and perfect transitions, we simplify the difference in the single-sample upper and lower bounds
67. K= + 6T k=1
EUBOAIS (x: mo, 1')
ELBOAIS (x; no, T")
AIS
(x, zo.r)
Piot (x, 20:T) ToRop (Zo:T (x))
log
Beo.r
log
(FROP (zor (x)
क.. (2rlxx (x) [-1
II Ti(za- (2()
TT(Zr) x) II
E P
log
E B Sraer
log
1=1
(z)l)) II Ti (z.-1)
To Tu(Zu|x)
II Te(ze (zc-1)
T
ix, 2e) #u(z)|x))
Pi Pi-Pe-1
T
ir (x, 2e) #u(z,(x)
P
log
II 1=1
E Exg A
log
II 1=1
B
Σ (Br 1) log ir fr 2i) #o cio(zo1x
Exp
Σ (Br Br 1)
For (x,, 2c) #u(z+|
(69)
(1)
T
E Bxg (x)
(B. (Bi- Bo-1)
ūr 2)
Σ
T
B p, (x)
(8. Bc-1) log #r x) #0
#o (z/x)
1m]
[-]
(
#r (x,2) 2)
io (x, 2)
Eur((a)
log
T Ta)(())
log
(z(x
To (2|xx
(Dea (oo | (Iu(FTr/Foo))
7
25

---

Published as a conference paper at ICLR 2022
where in (2), the linear annealing schedule Bt 3 vt 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
ZO:r
(1:K)
To (2o)
II TT. (ze);
(70)
T
ZO:T= ~ PIGr (Zo:T
(I:K)
x)
(PT)
mr (T))
II
TT (zt-1).
(71)
In other words, for 1 S t S T, perfect transitions results in independent, exact samples from 22 T& (z) in the forward direction, and TH, (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 each Ze.
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 Independent Multi-Sample AIS in App. E.2 below, and extend the result to Coupled Reverse Multi-Sample AIS in App. H.4.
E INDEPENDENT MULTI-SAMPLE AIS
E.1 PROBABILISTIC INTERPRETATION AND BOUNDS
Our Independent Multi-Sample AIS (IM-AIS) are identical the standard IWAE bounds in App. B, but using AIS forward ypgor (zarl|x and backward yorr (Za:r(x) chains of length instead of the endpoint distributions only go(z)x) and p(z]x).
In particular, we construct an extended state space proposal by running K independent AIS forward (k)
chains Zo:T ~ drgor in parallel. As in the IWAE upper bound (Eq, (5)), the extended state space target involves selecting single index uniformly random, and running backward AIS chain (+)
Znrr plar starting from true posterior sample Zr ~ p(z]x). The remaining K samples are
obtained by running forward AIS chains, as visualized in Fig.
IM-AIS
(ZigeT
(1:K) (x) == II
Orzor (Zo:r x))
IM-AIS (x, Zor
(1:K)
Pror (x, oor)
II
Oraos ( x).
where FpROP and pror were defined in Eq. (12). Expanding the log unnormalized density ratio,
K
Σ
Pior
(x,
(s)
π
OFROP
(k) Zorr (x)
DMM-AIS (x, (1:KK) log
log
kyf
(72)
DM-AIS (Zu:T (x)
(1:K)
K
II
qeas (2 O.T x)
(x, Zor))
(k)
llog
(73)
K
(k) TRRP Zurr (x)
which is similar to the IWAE ratio but involves AIS chains. Taking the expectation under the proposal and target in App. A recovers the lower and upper bounds in Eq, (14). The gap in (1:K),
the lower bound is DKL P (zo.T [x) Il (2
(Zo.T (x))) and the gap in the upper bound is [x))]
(1:K)
DRL. DKI.I
IM--AI (20TT
(1:K)
[x)) (k)))lugM IM-AIS (Z0:T |x)).
(1:K)
E.2 PROOF OF LINEAR BIAS REDUCTION IN T FOR IM-AIS
Corollary E.1 (Complexity in T for Independent Multi Salti-Sammle AIS Bound). Assuming perfect transitions and geometric annealing path with linearly-spaced 1T the sum of the gaps in the Independent Multi Aulti-Sample AIS sandwich bounds on MI, LIM-AISU (To. T) IIM (To. T), reduces linearly with increasing T.
26

---

Published as a conference paper at ICLR 2022
Proof. Using identical proof techniques as for IWAE (Burda al. (2016); Sobolev & Vetrov (2019), Prop. B.1), we can show that our Independent Aulli--Smplle AIS bounds ELBOIM-AIS (x: To. T, 1) I ELBO (x; mp. T, K) and EUBOMM - (xs To. T. K) I EUBOIM AIS (x; To. T, K 1) im- prove with increasing K. Thus, the bias of our multi-sample bounds less than the bias of the single-sample bounds, 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.3 PROOF LOGARITIMIIC IMPROVEMENT IM-AIS EUBO
Proposition E.2 (Improvement of Independent Multi-Sample ATS over Single-Sample AIS).
Let Prer IM-AIS (s|x, (1:K)) = Pirr (x. R..7
Prer (x. (8)
EE=
denote the normalized importance
(
weights over ATS chains, and let U(s) indicate the uniform distribution over K discrete val- ues. Then, we characterize the improvement of the Independent Multi-Sample AIS bounds on p(x), ELBOIM-AIS (x; To. T. and EUBOIM-AIS (x; To. T..T.KK K) over the single-sample AIS bounds ELBOAIS x;; no, T) and EUBO,Is(X; To, T) using KL divergences, as follows
ELBOIM (x: No, K) ELBOAIs (x; No, T)
CUK) D.
(x) DEL
IM-Ais (s/zarr x)]
uniform from SNIS weights S DEL (Ra:2) (x) prar (Morr [x]]
(74)
EUBOIM-AIS no, K) FUUBOAIS To, 7)
meat (1:K) (*0TT" (x) DKL Piar (s|2 (1-K) [x)]([{s]]
S KL of SNIS weights from uniform S log
(75)
Proof. The result follows directly from App. C.4 Prop. 2.1 by viewing Independent Multi-Sampa AIS as IWAE with an AIS proposal as App. E.1.
F REVERSE IWAE
In this section, we propose Reverse IWAE (RIWAE), which is an impractical alternative to standard IWAE. However, we use this as the basis for our Independent Reverse (App. G) and Coupled Reverse Multi-Sample ATS bounds (Sec. 3.3 and App. H).
F.1 PROBABILISTIC INTERPRETATION AND BOUNDS
consider K indepen- dent posterior samples in extended state space target distribution. The proposal distribution for our importance sampling scheme is mixture of K posterior distributions and one variational distribution,
Similarly to simple reverse importance sampling, pixj Ep(x p(xx))
[x(x)
K Σ (z x) II z()
JUWAR: Arwor
(2:K)
Pior
(2:K)
p(x, 2 (K)).
(76)
K
1=1
Note that we have normalization constants of 1 4pz0r (x, (1::)))d2 = p(x)K=1 K- since the transition kernels of the reverse chains do not change the normalization.
RIWAE (x, (1: (1222)
p(x)* and
27

---

Published as a conference paper at ICLR 2022
We visualize this sampling scheme in Fig. 4. Similarly to Eq. (25), the log unnormalized density ratio simplifies to
K
11 A=1
p(x,
(1)))
PIGT (x, (1:KK)
log RRWAE (x, (1:)K)
log
(77)
1
K
K
2 40 (2(0) (x)
II 2(1))
K
Right
qo (k( x)
log
(78)
p x, z (A:)
Taking the expectation under the proposal and target distributions yield lower and upper bounds log p(x). These translate to upper and lower bounds on p(x) which are different that standard IWAE bounds in Eq. (5),
K (0
(1) xx) Xx, 2(k))
K p (R))
xx)
log
S log p(x)
log K
(a (x)
(x)) 2=-2
K
II pxx)
k=
Note these bounds are impractical as they would require more than posterior sample from x). However, we use them in Sec. 3.3 and App. to derive practical multi-sample reverse AIS bounds.
F.2 IMPROVEMENT RIWAE OVER ELBO AND EUBO
Proposition F.1 (Improvement of Reverse IWAE over ELBO and BUBO). Let Grror (s|x,
( zz,,z!! (1::K))
20(2(0) bc) p(x, k=
(x) denote the normalized reverse importance sampling weights, or the pos-
p(x,
terior over the index variable in Eq. (76). 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).
ELBORIWAR (x; qo, K) and EUBORRWAL (x; go, K), over the magle-sample bounds ELBO(x; qo) and
EUBO(x; qo) using divergences, as follows
EL ERRORE (x; go, K) ELBO go) + E aanA (1:KK) x))
DKL ( RIWAL:
(s|x, (2:K) ((((()))
0 I KL of uniform SNIS weightt & log K
EU EUBORRWA (x: K') EUBO(X; go) ) E (u
DKL [(((s) RIWAN (s|x, 2 (1:80)]]
0 < KL of SNIS weights from uniform D Dwl.pp2/xx)llaxxzxxs
Proof. The proof follows similarly as Prop. B.2 or Lemma C.2.
G INDEPENDENT REVERSE MULTI-SAMPLE AIS
G.1 PROBABILISTIC INTERPRETATION AND BOUNDS
In similar fashion Reverse IWAE, we now use K independent reverse AIS chains form an extended
state space target distribution pl-Ais (zorr
(1:K) x), and a mixture proposal distribution g (2
(1:K)
which includes single forward AIS chain (see Fig.
(1:K)
1
(s) Zo x)
II
AIS (x, (8)
4PKOP
(x, Z
AEROP
K
PTGr
IR-AIS
(x,
2p.T
(1:K)
Prer (x,
AIS
2o.T)
(k)
(79)
1=1
28

---

Published as a conference paper at ICLR 2022
Similarly to Eq. (77)-(78), the log unnormalized density ratio becomes
I1 pAI (x; Zo:T)
(k)
PR-AAIS (1:K)
PTGT DTT x)
log
log
(80)
IR-AIS
0:T
(1:K)
x)
PROP OoT x)
(k)
prar (x, 2o.T)
1=1
K
AIS
( x)
(k)
log
(81)
K
AIS
X
ZorT)
(A)
k=
Taking expectations of the log unnormalized density ratio under the proposal and target, respectively, yield lower and upper bounds on log p(x)
1 log
APROP O: (x)
(R)
AeROP 20: x))
(1)
(1) SoR
K
Pior (x, ZO:T)
(k)
< log pogp(x) S EE (1:0) fer
log
Prer (x, Zo:T
(1)
#o.T
ELBO (X; To. K)
EUBO (x; No, K,T)
(82)
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.2 PROOF OF LOGARITHMIC IMPROVEMENT K FOR INDEPENDENT REVERSE AIS
Proposition G.1 (Improvement of Independent Reverse Multi-Sample AIS over Single-Sample AIS).
Let grtor | Zo:T
(1:K))
Frice (x)
E
Treor (7 (N)
denote the normalized reverse importance
(x,
k=1 Prot (x)
sampling weights over AIS chains, and let ((s) indicate the uniform distribution over K discrete values. Then, we can characterize the improvement of the Independent Reverse Multi-Sample AIS bounds log p(x). ELBOIR-AIS (x; Tu, K, K,T) and EUBOIR-AIS (x; Tu, K. T), over the aingl-sample ELBO AIS (x; No; T) and ELBO AIS X2 Toy T') using KL divergences, follows
ELBO (x; No. K,T) ELBOAIS (x: no. To, T) + E
(1:K) (#::T
DEL.(erar TR-AI (s|x, (1:K)
[2]
S KL. of uniform from SNIS weights S log K
EUBOIRAAS (x; no. K,T) -BBBO.n(X; no, T)
E Play
(1: (For
DKL. (=)))))))oo 1R-A15 x, (1-
KL weights from aniform DEc. (pits (Rour (xc) so (Roer [x]]
Proof. The proof follows similarly as Prop. B.2 or Lemma C.2.
H COUPLED REVERSE MULTI-SAMPLE AIS
H. PROBABILISTIC INTERPRETATION AND BOUNDS
The Coupled Reverse Multi Sample AIS extended state space target distribution in Fig. 2 initializes K backward chains from a single target sample Zr ~ Ty (z|x), which makes the bound useful in practical situations. We denote the remaining transitions as AIS (Zn:T- |zr, x), since they identical to standard AIS in Eq. (12). Thus, the CR-AIS extended state space target distribution is
(1:K). CR-AIS (x, 10:T-11 ZT) = n. (x, ZT)
II 1=
AIS Prer (Z.TT /ITTTIIZT,X Ix.
(83)
The extended state space proposal is obtained by selecting an index s uniformly at random and running single forward AIS chain. We then K - backward chains, all starting from the last
29

---

Published as a conference paper at ICLR 2022
state of the selected forward chain,
K
OPROP D1: 1
CR-AIS (1:K)
xx)x)=
=
APROP (2
AIS (+)
0:7 2r x)
II
Prer (Z 2r, x).
AIS
(84)
K
See Fig. 2 or Fig. 4 for graphical model description.
We can construct log-pun pog-patttion bounds using the log unnormalized density ratio
K II
CR. AIS (1:K)
[rr, x)
Ais
2r, x)
log
Prat 0.T 1 ZT,
log
k=1
(85)
(1:K)
K
K II
(PROP (Z0TT11:24 17
AIS
0:7-1 2T ZT(x) x)
prar Zo.T ZT)
K
Rm
K
PPROP U:T U:-' Zr x)
(k)
log
(86)
K
k=
Piar
(z
(k)
n. (TT x)
Taking the expected log ratio under the proposal and target yields lower and upper bounds on log p(x).
K
(k) ArROP
2u.T
ELBOcR.AIs (x: No, K,
log
< logg((x)
Sraoe rlx)
K
piu (x, 1: 2r)
(+)
(2:K) Por (#0:T (RT.x))
Rml
AiS
(Z Po: Zr xx)
EUBOCR-ARS TO, K,T) =
By (erle)))))))) (1:K) uppor (Roon- (xm.x)
log
K
pier (x, 20.7 2T)
> log p(x).
H.2 PROOF THAT CR-AIS BOUNDS TIGHTEN WITH INCREASING
In this section, we prove that Coupled Reverse Multi-Sample ATS bounds get tighter with increasing K. Our proof provides an alternative perspective to Burda et al. (2016); 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 direct consequence of this lemma.
Lemma H.1 (CR-AIS Bounds Tighten with Increasing K). Coupled Reverse Multi Saatii-Sampeeeee AIS bounds get tighter with increasing number of samples K. other words, for any K > 1,
ELBOce AIs (x; To, T, K 1) S ALBOcr (x; To, T, K) < logp(x);
(87)
EUBOCR (x; no K 1) EUBOcR (x; mo. T, K) log p(x)
(88)
Proof. Our proof will proceed by introducing an additional set of M index variables S1:M into the K-sample probabilistic interpretation of CR-AIS in Eq. (83)-(84), with M K.
We will show that the KI. divergence in this joint state space (including 81:M) is equal to the gap of the M-sample ELBOCR-AIS (x; no. T, M) or EUBOCR-AIS (x; no. M,M. We then show that marginalizing over 81.M yields the gap of the K-sample ELBOCR-AIS (x; no, T, K) or EUBOCR-AIS (x; To, T, K). Since marginalization cannot increase the KL divergence, we
will have shown that for any M < K, DKL (z) (Z) z7 [x) (I (ZorT-11 ZT|x]] s (1:M)
(1:K)
(1:K),
DKL. (oC--AS (ZOTT 2r 2r(x) [DC-AAIS (1:Aa) 2r [x]x and thus ELBOCR (x; ALBOORRAisX,, No, T, M)
S
ELBOCR-AIS (x; Wo. T. K). Identical reasoning holds for the EUBOCR-AIS-
Sub-Sampling Probabilistic Interpretation (si, 8M) indicate the probability of drawing M < K sample indices uniformly without replacement (i.e. each 8m is distinct). In the CR-AIS target distribution, there is no distinction between the indices {$IM } and k & {81:22 after drawing
Zr (x,, z). We can write
Pret
CR-AI8 (x, 81:M, 2 Zτ) = U($1 ..., 835) TT (x, ZT)
Prer (7 BT, x),
(89)
30

---

Published as a conference paper at ICLR 2022
which, after marginalization over 81.M, clearly matches prer (x, 11 ZT).
(I:K)
We also draw $1:M U(s............................. SM) for the extended state space proposal. Next, we select index m uniformly at random from . which is used to to specify which chain 2o:T-1 is run in the forward direction obtain zy, in Eq. (84). After marginalizing over m, we obtain the following mixture proposal distribution
CR-AIS (1:K) GEROP (81:M: -1 zr (x) = _(..........................
F (2 D:T 2r(xx) (x) II
Prar (o:TT |ZT,
AIS (e4)
M
m==1
jmm
( 2o kr, (r, x).
AIR
(90)
We will consider marginalizing over 9I.M below, but first write the divergence in the extended state space which includes $1:M- For example, the forward KL divergence matches the gap of the M M-sample ELBOcR-AIS (x; To, T, M).
(1:K) DKL|9pRoP ($1:M; Zo:T ZT| (r(x)) PTaT CR-AIS ($1.M 11: 2T [x)]
CR-AIS
(1:K)
(91)
Als (tm)) Σ
- 27/3)
II
(n))
Ii
mamm
log
Tr (x ET)
piu (0 (xx,x)
(K)
ARI-TM
DRL. [grRor (0 Zr xx) 11 CR-AIS (1:M)
CR-AIS
2 2r|xx]
(92)
which matches the M M-sample probabilistic interpretation of CR-AIS from App. H.1. Identical reason- ing holds for the case of BUBOCR-AIS (x; To. T, K) < EUBOCR (x; To. T, M) using the reverse KL divergence.
Marginalization over 81:M We have already seen from Eq. (89) that the marginal (I:K)
En Prer (x, 81:a5, 2o 11 Zr) matches the K-sample target distribution pier (x, 20:7 Zr).
(1:K)
(1:K) We would now like marginalize over s1:25 in ooRoP ($1.M 11 1. zr xx.. Combining the two product terms in Eq, (90),
CR-AIS (SI.M, Bor-1: Zr Zr Eepcoras)) Σ
1 M
Σ
(
(#m)
2r (x(x)
II
Piet (Zu.T ZT,, x)
(93)
Arzor
( ( 2r x)) II
Prer (2 12T, x))
(94)
(1)
E((co)
Oraor ( 2r (x)
Pien (Zo.T IT,, x))
(95)
de Aeylaam
Grocw
AIS
(2 (1)
2o.r 1s 2r (x) II
Pin (2 2r,, x)
AIS (K)
(96)
M
K
A=1
K
Denap
11 Zr
Pien (2 - L2,
(97)
31

---

Published as a conference paper at ICLR 2022
where in (1), we can write the marginal (sm)) since the terms inside expectation do not explicitly depend on the other indices in $1:M. In (2), we use the fact that the marginal over any sm is uniform.
Thus, we have shown that the marginal distributions match the standard K K-sample target and proposal distributions of CR-AIS, with
DEL
(1:K (51:M oo zrbblll
E "
(fL
(1.))) TO By [x]
Dxt. (eaoe
(1.K) (x - By [x) (yy
iooe (7. By bx]]
(1-K)
II
#1:M
pppe BO0KR mg.
and similar reasoning for the EUBO using the reverse KL divergence.
Conclusion of Proof
Since marginalization can not increase the KL divergence, we have
Dia.lergor
CReals (BeTT zy (0 (20:7 [r]x]]
Dia (ybor (an:ar, "Dr [x)] CR-A (BI:M; -0:7
[x]]
zap of Ans No.
X)
pap of tLbiiee AIS (aiing 21, M] (Bg
(98)
This shows that ELBOCR-AIS To, T, K) is tighter than ELBOCR-AIS (x; no. T,,,, with EIBBO AIS X: To, T, K) 2 ELBOCR (x; To, T. M). Identical reasoning holds for the case of
EUBOcR AIS X: no, T, K) S EUBOCR-ArS (x; To, T, M) using the reverse KI. divergence. Choosing
= K - proves that the bounds tighten as K increases, as desired.
Characterizing the Improvement of K- K-Saplle 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 CR-AIS
SPROP example CR-AIS ( 20 zr|x)
(1:K)
OpRor (Zo:7-1 (x
(1:K)
2r (xrr(x)xx) TFRO ($1.M) (0:T 2T, (T,x). x).
(1:K)
The gap in Eq. (98) also corresponds ELBOORR-II (x; No, T, M), so we can write
the gap between ELBOCR (x; no. T, K) and
ELB x: mu, T, K) - ELBOCR AIS AIs No. M)
(I:K) DKLllgooo ($1:M / (1:K) 1. Zr. x) | CR--IS ($1.M Z -1 x),
(99)
Below, we analyze the posterior over the index variable for M as special case. This allows us to characterize the gap between ELBOCR-AIS(N) mo, T, K) and the single-sample ELBO (x; To. T)) ELBBO (x; no, T, K = 1).
H.3 PROOF OF LOGARITHMIC IMPROVEMENT IN K FOR CR-AIS ELBO
Proposition H.2 (Improvement of Coupled Reverse Multi-Sample AIS over Single- AIS). Let
Teaop Oz (r(x)
p4ts (x, 20.7 (TT)
(a)
Grace (()
(I:R) Zg.TT MT)
(100)
GrROP 20: - Zr x)
AIS (R)
piul (x, z (ZT)
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, can characterize the improvement of the Coupled Reverse Multi-Sample AIS bounds on log p(x), ELBOCR-AIS (x; no. T, K) and EUBOCR-AIS (x: 50, T, K), over the single-sample AIS bounds EISBOAIS (x; no, T) and EUBO (x; To, T) using KI. divergences, as follows
ELBOCR AIS (x; no, T, K) ELBOAIS T")
(3:K) (4
(1:K) (x) DKL cu-Ai (s|x, 20. ZT) [[))
0 S KL. of uniform from SNIS weights S log
EUBOCR-AIS (x; T, K)) (U; - B cas (2 Por 1 zr(x)
DKL [(((s)](qao0"" groc ( (s)xx, Z 27
(1-K)
0 KL. of SNIS weights from unitform S. DKL, (RIlPAN (Ma.r) x) 1 (No xx)
Proof. Using M 1 in Eq. (99) above, we obtain
EL ELBBOCR-AIS (x; no; T, K) ELBOAIS (x; no; T, =)
DRL. AI (1)2 (1:K)
(1:X) (s|z).--1, Zr, x) | CRIAIS (1) 2T, xi]
(101)
32

---

Published as a conference paper at ICLR 2022
From Eq. (90), write the posterior the index variable as
Orace (2 4:7-1 -1
CR-AIS 4PROP
(I:K) (s 0:7 27, 2T,
Pror (x. xx)
(102)
K
Qrmrr (2 1 (x)x)
(6)
Piot (x.2 (2r)
(1:K) For the target distribution in Eq. (89), the posterior Pror (s|Z we can characterize the improvement the K sammpl CR-AIS ELBO over its single-chain AIS
ZT x) = ((s) is uniform. Thus,
ELBOrrR-Is (x; no, T, K) (X: no, T)
Qasor OOrT 1
(1:K)
DRI (Qveor (s|xx T0 [(((s))
(1:K)
(103)
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 x; no, T), which corresponds to DKL [pAss (Zoer |x) }lllAt (Zo:r|x)
H.4 PROOF OF LINEAR BIAS REDUCTION IN T FOR CR-AIS
Corollary H.3 (Complexity in T' for Coupled Reverse Multi-Sample AIS Bound). Assuming perfect transitions and a geometric annealing path with finerll-spaced { {T the sum of the gaps the Coupled Reverse Multi-San Multi-Sample ATS sandwich bounds on MI, ICR-AIS (To, K, T) ICR-AIS, (To, K, T), reduces linearly with increasing T.
Proof. Since we have shown in Prop. H.2 that ELBOCR-AIr (To K, T) and EUBOCR AIS (To. K, T) tighter than ELBO (To, T) and EUBO AIS ( (To T), the bias in the CR-AIS sandwich bounds EUBOCR-AIS (To, K,T) ELBOCR-AIS 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 COMPARISON OF MULTI-SAMPLE AIS BOUNDS
In Fig. 5, we compare the performance of our various Multi-Sampi AIS bounds for MI estimation of a Linear VAE with 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 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 the Coupled Reverse Multi-Sample AIS upper bound MI (Prop. H.2).
The results are less conclusive for the Multi-Sample AIS lower bounds on MI, where either the Independent Multi-Sample 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 independent negative forward chains which, by App. E.3 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 no (z) = p(z), the Independent Multi-Sample AIS lower bound on MI reduces to Structured INFO-NCE and saturates log K. This may be useful for quickly estimating values of MI at a similar order magnitude as log K.
EUBO or Coupled Reverse AIS EUBO may be preferable for log
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, in Fig. 5 that at least moderate values of T may be needed to match 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 and the amount of computation available.
33

---

Published as a conference paper at ICLR 2022
LiWe Boird
1400
Uhear VAZ S00
1100
Uvear VAZ 1010000
Nort Boo
Shadet oo. lure Bockd
Soalia horeedd Love Bond
Granc Porcoa Lare Bar
314
(a) Linear VAE (AE(==) 11)
(b) Linear VAE TIK=10 100)
(c) Linear VAE (K- 10000)
1805
49SPACI1-2
PMESANE D-
W9STRNCIe- 2008
515
565
Couyted Poreoot Now
K35
Etvemsed Lol Board
(d) VAE 1)
VAE (K~ 100)
(f) VAB (K 1000)
304
304
200
24
(g) GAN 11
(h) GAN (K- = 100)
(i) GAN (K- = 1100))
Figure 5: Comparing Multi-Sample AIS sandwich bounds and evaluating the effect of K and T for MI estimation in deep generative models on MNIST.
J GENERALIZED MUTUAL INFORMATION NEURAL ESTIMATION (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 similar spirit to GIWAE, Generalized MINE uses a base variational distribution qe (2(x) to tighten the MINE-DV or MINE-F bound and can be evaluated when p(z) is available.
See Fig. 6 for 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.
34

---

Published as a conference paper at ICLR 2022
20(2|x)
I(x; z) = Bysoo)
log
Byso) DDx [exeppx)))lll
p(z)
1s Da. p(xtx) F Fuool
2(x)
g0(2/xc) p(z)
IBAL(go. Ts) Bxoca)
log
Bp(xa) (x,(xx Eyool log [Bette o² Fa(aa)
Piu ((s.e)) (20000000000
JaNNNE- (gs. To) = Bynna) log g0(x)xx)
+ Bpona) [Tx(x, 2)] log Eptxinntts)) [e Tolnn))
p(z)
2 D Wolibiee wes Tolon))
Zix)
Doxi pp0jaelaa Fanaal po (ospe
2 Dou (xal FArk--1
g0(2|x)
JGMINE (96. TT)
Bo(ax)
log
Bptxa) [T((x, Bpp0ogoto)) [e"
p(z)
(el2/x)
log
p(z)
Figure 6: Generalized Energy Based Bounds. Arrows indicate the gaps in each MI lower bound or its relationship to other lower bounds. DeaL (-ll-) represents the generalized KL divergence between two unnormalized densities (see App. J.3). All bounds are written in terms of a base variational distribution (2 x), which may be chosen to be the marginal p(z) as in MINE-DV and MINE-F.
J.1 PROBABILISTIC INTERPRETATION IBAL For an energy-based posterior approximation
To. (2)x) Zo. (o(2) (z) (x(x.))
where E0,0(x) Ego(e(x) le [eTo(xa))
(104)
we consider the BA lower bound on MI, JRA, (To.o) = T(x, z)
Bp(x) [DLL (To
(105)
p(x, z) p(x)p(z))
E Ex(x.x)
log
Ep(x.2) [logp(z|x) + Ep(x..) [T(x, 2)]
-Ep(xx) Dlog Zn,ol
log Go (o/(xx)
(x)(x.)
+ (x(()2) [T)(X - Ep(x)) [log Z
p(z)
log go go(z|x)
(x((x))
Ep(x22)
log [x (x
p(z)
TeAL (go)
contrastive term
(go To):
where the gap in the mutual information bound is Ep(x) [Xx]]]//)))(/// To. [qo TT) includes Lax, (ge) as one of its terms, where we refer to (z(x)) as the base varia- tional distribution. We visualize relationships between various energy-basee bounds Fig. 3b and Fig. 6.
Note that
Proposition 4.1. For a given (z(z)x), the optimal IBAL critic function equals the log importance weights up to constant T*x(x, z) = log p(x.z)
(2) + c(x). For this T*, we have IBAL(qe. T*) ) = 1 (x; z).
See App. L.1 for the proof.
J.2 PROBABILISTIC INTERPRETATION OF GENERALIZED MINE-DV
We can interpret the MINE-DV bound of Belghazi et al. (2018) as arising from an energy-based variational approximation a (x, z) of the full joint distribution p(x, z) as follows
(x, = x((x) (z) (e)x)) (o(x.))
1
where Z - (p(x))o((1x)
To(xx2)
(106)
35

---

Published as a conference paper at ICLR 2022
qo (z/x)) is base variational distribution, and To is critic negative energy function.
Note that for the induced marginal o (x) = ne T z)dz * p(x) due to the contribution of the critic function. Instead, we have
Te (x) 2 2 ((x),
(107)
where 2(x) - g(x)x)e de and = p( ((x))((z)x) (o dxda = (x)))
Subtracting the joint divergence DxL (xL(p(x, |T z)] from (x;z). obtain the Generalized MINE-DV lower bound on MI
1(x: z) 2 1(x:2) z) DKL p(xx, z)
(108)
-xxx) log (0(2|x)
+ pxxx2) (Tx(x, log (xxx0(())
(T((xa))
(109)
p(z)
DO (go, To).
By construction, Eq. (108) shows that the gap in Generalized MINE-DV is DKL[p(x, z) ( which has probabilistic interpretation in terms of the approximate joint distribution (x,
Relationship with MINE-DV For (z = p(z), we obtain the MINE-DV bound (Belghazi et al., 2018; Poole al., 2019) as special case with MINE- ny (Tş) ICMINNE (p(z), To). .). In particular, the joint base distribution in Eq. (106) corresponds to the product of marginals p (((). Our probabilistic interpretation shows that the gap in MINE-DV corresponds to DKL DKl/p(x, z) 2 p(x (x()ee)
We expect that the Generalized MINE, with a learned variational distribution go g0(z[x), obtain tighter bounds than MINE-DV. We can guarantee that LosINE-ov (4g. TT) > MINNF- (T))) for the optimal up with given Tø, so long as pzz) the variational family (i.e., 30, such that yo, (z) x) = p(z)).
Relationship with BA Bound Choosing a constant critic function To (x, 2) const, we can see that 70,00 (x, z) = p(x)go(2/xx and TOMINE. ny (go, Too = const) JAA (46) for given (z(x)))
Optimal Critic Function For a given go (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)qe (z)x), plus constant
p(x, z)
T (x, 2) = loo (xx))x(Zz)
+ c.
(110)
For the optimal critic function associated with given (z(x))). we have DN (qo. 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 [BAI.(qo, Ta) JOVINE.DV (go, Tş). Subtracting the gaps in the bounds in Eq. (105) and Eq. (108),
[go Tτ) TOMINE-DV (go, Ta) + DKL. [p(x 2) To,0 (/,0 2)) Ex(x)) [Dii. [p(z)] (p/z///))))) [|TT)) (2)x)]
gap in OMINE-DV To)
gap TRAI (QA. T&)
log p(x)
= ToMINE-UV (ye, Tτ) + W. p((x))
(111)
(x)
-TMMINE-D ov (40, Tτ) + DKL [l
(112)
2 LOMINE-NY (go. To).
(113)
We see that the difference between IBAL(qo. TT) and I (go. TT) corresponds to a marginal KL divergence in the space, where (x) is defined in Eq.. (107).
36

---

Published as a conference paper at ICLR 2022
As in Poole et al. (2019), also interpret the gap between the IBAL(90. To) and IOMINE-DV (go. TT) as an application of Jensen's inequality, with
Dx = DKL p(xx) z p(xx) (xx)
(114)
log z pxxx) flog (xx)
(115)
= log Ep(x) [xxx) Ex(x)) log 2(x),,
(116)
since log (p(x) 2 Ep(x) log [xx]
J.3 PROBABILISTIC INTERPRETATION OF GENERALIZED MINE-F
The BA lower bound and its corresponding EUBO upper bound on log,(x) are derived using a KL divergence between the true and approximate posterior, which normalized conditional distributions over z given x. In this section, we interpret the MINE-F bound (Belghazi et al., 2018) as arising from generalized notion of the divergence between possibly unnormalized density functions. In particular, our probabilistic interpretation will involve an unnormalized (z)x(z which seeks to approximate the (normalized) posterior p(z)x). x).
Generalized KL Divergence (GKL) First, we state the definition of the generalized KI. divergence,
which takes unnormalized measures as input arguments and corresponds to the limiting behavior of the a-divergence (Cichocki & Amari (2010))
F(z) dz
f F(z)dzz ʃ 8(z)dz.
DOKL (F() = F(z) log 5(z)
(117)
As in the case the standard KI. divergence, this quantity is nonnegative, convex in either argu- ment, and vanishes for f(z) $(z). It is also a member of both the family of Bregman diver- gences and f-divergences (Amari, 2009). r(2) and s(z) are normalized, then (rx),[((z)][s)) = (xL(r())]<<(2)) Finally, if r()) is normalized and ((z) is unnormalized, one can easily confirm that
DeKL (x()((())(3())) DxL (xL(r())(((())) + (skr[s(z)]
(118)
> Dx((r())]((())
Generalized EUBO (GEUBO) Since the Generalized KL divergence is always nonnegative, define a Generalized EUBO by adding the Generalized KL divergence between the normalized true posterior and unnormalized approximate posterior 7 (z)(x).
GEUBO(x; 7) logp(x) DAKL [p(z)x)]|
(119)
p(x, 2)
= Eppea)
log #(2|x)
(x),,,
(120)
where we define (x(x) = A (z|x)ddz. In general, have EUBO(x; #) GEUBO(X; f) where the normalized distribution स, with equality if is normalized.
Generalized BA (GBA) Using the Generalized EUBO in place of log p(x), obtain the following lower bound on MI, which we denote as the Generalized BA lower bound
I(x: 2 I(x: z)
Ep(x)
[DaIL. feee..p(2ee (z) (z|x)]
(z/
Ep(x,x)
log
+ Bp(x)) [Z (x))
(121)
p(z)
Lours (T)).
Generalized MINE-F (GMINE-F) We now consider unnormalized, energy-based approximation to the true posterior, involving a base variational distribution qe (e(2)x) x)) and a learned critic function To (x,))) no. (2)x)) == go (oo
(122)
37

---

Published as a conference paper at ICLR 2022
Using this unnormalized approximate posterior in the GBA lower bound in Eq. (121), obtain the Generalized MINE-F lower bound
T(x; x) 2 T(x;z) z)
p(x) [Daki (ptx///p)/)))]ffooo1mm
(123)
log [0(2)))
p(z) (p(xx) [Ta(x.. x)] (p(x000(a [o(xx))1]]
p)(xx)
(124)
= IoMNNE ( TT).
(125)
By construction, we can see that the gap in (oMIND-r(qe (T) is equal to the Generalized KL divergence Ep(x) [DakL [p(z)x) oo (z)x)]] in Eq. (123). As in the case of MINE-DV, obtain the standard MINE-F lower bound I (Ta) JoMINe. (p(z). TT) when using the marginal p(z) the proposal.
Optimal Critic Function The optimal critic function of Generalized MINE-F is T (x, z) = 1
log p(x,)) (Poole et al., 2019). In this case, obtain T (go, T*) T(x: in Eq. (124).
p())((()))
Relationship with IBAL We would now like to relate the gap in LOMINE (go To) to the gap in (gs, To). First, note that the normalized distribution corresponding to Tg. (z(x) = go [z(2x)ee (o (x.))1 matches the the energy-based posterior in the IBAL, To. (z)x() = x go (z)x)) Using Eq. (118), we have
px(x) [Deku (p<x//p(z//))) no. (2]x)]] = pp()) [Dx (xx x]]] Ep(s) (aas.
and substituting in the definition of each term
DGKL p(z|x)
Ep(x)
DKL p(z|x)
2(x) 20 (2(22))) Fo(o.o)) Poto
(126)
gap of GMINE-F
gap of IBAL
+ Ep(x)
DOKL
go (z(x))) (olox))
[o(2)x) (o(x))))
[(x)
where the final term DOKL (o [z divergence. Thus, we have
> 0 is nonnegative because it is generalized KL
[BAL(qa, Ta) I (go, T + Ep(x) [DGKI. [Ta (z) [x 2 LONINE. (1o, To). (127)
Alternatively, Poole et al. (2019) use the inequality log u S 1 to show that
DAKL. [ (z|x) To (z|x)] = Ex(x) [ - log 2(x)) 2
We visualize the relationship between IBAL and Generalized MINE-F in Fig. 6.
Relationship with Generalized MINE-DV DV and Generalized MINE-F, we can again use the equality from Eq. (118), but this time using divergences over joint distributions.
To characterize the gap between Generalized MINE-
DCKL [p(x, (x, DRL (x(x, 2) (x, 2)] + DEKL [x),((X,2)) 2) (x, z)]
DGKL [x(x, z)
/0/l| l (zxx = DKL p(x, z))
1 1 (z) ((x.2)
(128)
gap in GMINE-F
gap GMINE-DV
DOKL
p(x) (z) (z)x)e (xxz) p(x)go (x.c))
Z
In this case, have DOKL ( (x..x) ( ||x()) (z) (z)x) (xxx)) 0. Thus, Generalized
MINE-DV is tighter than Generalized MINE-F (see Fig. 6), which generalizes the finding in Poole et al. (2019) that standard MINE-DV is tighter than standard MINE-F.
38

---

Published as a conference paper at ICLR 2022
K CONJUGATE DUALITY INTERPRETATIONS
In this section, we interpret the energy-based MI lower bounds MINE-AIS, Generalized MINE-DV, Generalized MINE-F, GIWAE, IWAE, and INFONCE from the perspective of conjugate duality. In particular, we highlight that the critic or negative energy function the above bounds arises as dual variable in the convex conjugate representation of the divergence. In all cases, the KL divergence of interest corresponds to the gap in the BA lower bound
yo (z|x)
T(x:z) ())())) log
Ep(x) [PKL. [2 (20
(129)
p(z)
For (0(2)x) = p(z), the BA lower bound term and our derivations correspond to taking dual representation of MI directly, e.g. Ep(x) (Dxi.p((z) [p(x))] in Belghazi et al. (2018).
Our conjugate duality interpretations are complementary to our probabilistic interpretations, with either approach equally valid for deriving lower bounds and characterizing their gaps.
K.1 CONVEX CONJUGATE BACKGROUND
The KL divergence to fixed reference distribution To is a convex function of the distribution in the first argument (2(-) == DKL [-]] [z)).. In later subsections, we will also consider KL divergences over joint distributions, unnormalized density functions, and extended state spaces.
For given (2(),, we can define the conjugate function (-) over a dual variable function T(z), with the following relationships (Boyd & Vandenberghe, 2004)
(T(((z) = sup (T(z), [()) ((((()),
(130)
=(x)
((T(2)) = sup [r(z), (T()))))
(131)
T(x)
where inner product notation indicates (z), T(z) == f π (z) Solving for the optimizing argument each of Eq. (130) and Eq. (131) yields the following dual correspondences
(2) = V (T(z)), (z) = V (())))))
(132)
We will proceed to derive closed form expressions for various special cases of Ω and S2" below.
For an arbitrary p(z) and T(z) which are not in dual correspondence according Eq. (132), we can use Eq. (131) to derive lower bound on ((p(z)) known as Fenchel's inequality
((p(z)) = (p(z), T()) - 22" (T(z) + Dac) [p(z),(())))))))))))
(133)
2 (p(z), T(x)) - 22 (T(z)).
(134)
where the gap in the inequality Dnc)) [p(z), TT TT(Zz) the Bregman divergence generated by Ω and Tr (z) is the dual variable corresponding to T(z) using Eq. (132).²
K.2 CONIUGATE DUALITY INTERPRETATION IBAL
To obtain alternative derivation of [BAL( To), we consider the conditional KL divergence from a reference qe (z)x), which a convex function its first argument
2(.) D (xlalll
(135)
To derive the conjugate function Ω (T), note that we must restrict the optimization to the simplex क (z)x) x) A, since the standard KL divergence requires a normalized distribution as input
(T(xx)))) = sup / x(2xx)
T(2)x) T(x, 2)dz - [(T(2)))) - AC/ (2)x))dz-1)
(136)
log (x)()) T(xx)
(137)
'This follows directly from the definition of the Bregman divergence by using To = V(i(g) and the identity in Eq. (130), Dalp. a(p) l(g) (Vi(g),p q) 9(p) 2(g) (T (T..p)
a(p) + (g) (Ta.P).
39

---

Published as a conference paper at ICLR 2022
where have solved for the optimizing argument (z/x) to obtain the conjugate function in
Eq. (137). Eq. (132) suggests the following dual correspondence between primal and dual variables
Tx(x, z) = log (z|xx)
(T(2/xx 2(x; T) d0 (x)x)e=((x)
c(x).
(138)
(o(2)x)
We would like to leverage this duality to estimate the KI. divergence ((p(zlx))
DxI [x_lpp(()))llgxxll from qo [z|zx]] to the true posterior p(z|x). In particular, plugging Eq. (138) into Eq. (131) suggests the following variational representation
sup / plzix) T(x, 2)dz log (xp(( T(x.x)
(x.aa
(139)
For suboptimal T, (x, 2), which is in dual correspondence with (T|x)) instead of the desired posterior p(z)x), we can use Eq. (134) to obtain a lower bound on DRi. [p(z) (l)) To characterize the gap in this inequality, one can confirm that the Bregman divergence generated by the KL divergence in Eq. (135) is also the KL divergence D π = DxI. [P]TT Thus, we have
[xllp(zzxx)l<o((xx]]]] = p(z T(x, z)dz log (eo
T(xx)]
+ Dx.
(140)
Finally, the (go TΦ) uses this variational representation of the gap in the BA lower bound, Ep(x) [DxL (p(z](() x) qo x)]]. to obtain a tighter bound on MI. In particular, for any learned critic function T(x, 2), we can use Eq. (140) derive the IBAL and its gap
log d (2)x)
I(x;2) Ex(x.))
Bo0 [Dar. (o
(141)
p(z)
lany (ite)
= Epptne)
log
(2/2x) p(z)
Epeo) (Epeaio (T 2) - log (ml()))) ((x.)))] +Bpoo [Px.lp(alx]]/x><xxxx]]]]
The optimal critic function T (xx 2) provides the maximizing argument in Eq. (139) and is in dual correspondence with the true posterior p(z|x). In particular, we have Ty (z|x) = p(z|x), resulting in 2x: = [gal(ge, T*).
K.3 CONJUGATE DUALITY INTERPRETATION OF GENERALIZED MINE-DV
To obtain a conjugate duality interpretation of (Generalized) MINE-DV, we consider the dual rep- resentation of the KL divergence over joint distributions. Choosing (x, z) - (z(x) as the reference distribution, the KL divergence is a convex function of the first argument
I2(-) = DKL (z)
(142)
This divergence equivalent the gap the BA bound Epcx) [Dku (px//////)))ll/ after noting the marginal distribution of both p(x, z) and p(x)qo(z(x) is p(x).
However, the duality associated with 22(-) DxL [: (o 2)] 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 Q(T((x, z)) = supr(x.)) (T(x, z), Tx,, z))
Ω' (T(x, z)), we need to derive the conjugate function Ω (T(x, 2)). Similarly to Eq. (136), we constrain the joint distribution be normalized and obtain
/ x(x,2) 2) T(x, z z))xdz Q(T((X 2)) - Acʃ T(x, 2)dzdx -1) -1) (143)
Ω (T(x, == sup -(x.x)
= log R (x1))
T(x.x))
(144)
where we have used mi mn(x, 2) = p(x) (z]x). Note that the expectation over p(x) now appears inside the log in Q2* (T z)), compared with the conjugate for the conditional KL divergence in Eq. (137).
Solving for the optimizing arguments in Eq. (131) and Eq. (143), we have the following relationship between primal and dual variables.
T=(X,z) log T(x,2)
Tr T 2) == Z(T) D ((x))x(2)x)) T(xx)
+ +c,
(145)
p(x)go
40

---

Published as a conference paper at ICLR 2022
Finally, we use Eq. (131) to write the dual representation of joint KL divergence
/ p(xx2) 2) T(x; 2)dxdz log l. (2))) ((x2))) T(x,2)
sup
(146)
As in the previous section and Eq. (134), suboptimal T'x, z), which is in dual correspondence with Ty (z|x) instead of the desired posterior p(z[x), yields lower bound on D z) (o (z)x).
with the gap equal to KI. divergence
DEL (p(x, 2) ())x(())))))) (z/x)] f p(x; 2xx 2) log Ppp00o [(x + DEL [p((px क (x,z))
(147)
The Generalized MINE-DY DV bound JAMINE-DV (qo T²) uses this variational representation of the gap in
the BA lower bound to obtain tighter bound on MI. For learned critic T(x, 2), we can use Eq. (147) to write
go(2)
x(x:x)
Fp(x.))
Ing.
+ + (ex]
(148)
n(x)
log
(2(xx) p(x)
(p(x)p(ale)) (T(x,, x)] log Eugia x)) ((xx)) (aa.ppxx)pdakkllllyyx.....
oo Top
As noted in App. J.2, the Generalized MINE-DV bound is looser than the IBAL.
K.4 CONJUGATE DUALITY INTERPRETATION OF GENERALIZED MINE-F
Nguyen et al. (2010) consider the conjugate duality associated with the family of f-divergences, of which the KL divergence is special case. We will show that this dual representation corresponds to taking the conjugate function of the Generalized KI. divergence, which can take unnormalized densities as input. See App. J.3 for definition of the Generalized KI. divergence, which is convex in either argument since it differs from the standard KI. divergence by only a linear term..
For 2(-) DGKI. [ [z)()) we write the dual variable using the notation T' T'(x, z) and consider
22" (T' (x,x)) 2)) = sup #(xlx)
[x2)))T((x, 2)d2 DCKL (2 0o (2|x)]
(149)
Note that we do not explicitly include Lagrange multiplier to enforce restriction to normalized distributions. Solving for the optimizing argument or writing the dual correspondence in Eq. (132), we obtain
=zz
(z(z)x) = [e(xx)ee" (x.2)
[x, z) = log
(150)
[x(2
Plugging this (2(x)) back into Eq. (149) yields the conjugate function
Ω (T'(xxx, = ā( ((z)x)e" (xxx) (x.))) t ((zlx)dz,
(151)
which leads to dual representation of the Generalized KL divergence that matches Nguyen et al. (2010) after plugging into Eq. (131)
/ j(z xiax)T"(( (xx z)dz ((z/x)e (x..) dz (lzix)da.
DOKL ((lz))
sup r(xx)
(152)
We now use the reparameterization T' 2) T(x,2) Assuming normalized q(z)x)
q(z|x) and p(z]x)= = p(z(x), and noting that Doxxlp(alx)lloo(xxx) = -xx.p(zzix)laoo(zxx)) for normalized distributions, we obtain
/
/ (0(x)x))) (xx2)) dz.
sup T(xx)
(153)
which matches dual representation of the KL divergence found in Belghazi et al. (2018); Nowozin al. (2016). See Ruderman et al. (2012) for further discussion.
41

---

Published as a conference paper at ICLR 2022
For suboptimal T5 (x, z) which is in dual correspondence with Tr (z|x) instead of the desired posterior p(z|x), can Eq. (134) obtain lower bound DGRL. [p(z)x) qo (z) x)] (xl[p(z) ())]]g()) To characterize the gap this lower bound, note that the Bregman divergence generated by the DKKL divergence is also the DGXL divergence D Dexal (ilxsllllil [ = DGKL Thus, using Eq. (133), we have
(xxll(e)x)(llozzxx))) f ((2))))T(xxx / ( dz + DOKL D -
(154)
Finally, we obtain the Generalized MINE-F bound by evaluating the dual representation of the Generalized KL divergence for normalized p(z]x) and (z(x)x). In particular, IOMINE- TT) uses the critic function Tę to tighten the gap in BA, (go) via the dual optimization in Eq. (154),
(xx; 2) = Ep(xe log p(x)
+Bp(se) [Diu (pa.(p(a/>)lleo0aaaa
(155)
Fane (16)
Ep(xa)))
logt (2(x(
Epix)
(aiaale) [T(,, log Bag(nts)) [(x(x))) +Epts) [Dona (xx))
p(x)
foumsu Log. (g))
(156)
The optimal critic function is the dual variable of the true posterior p(z[xx) which can be found using p(xxx)
Eq. (150) as Txx, z) + log (x tight.
With this optimal critic, the Generalized MINE-F bound
K.S CONJUGATE DUALITY INTERPRETATION OF GIWAE, IWAE, AND INFONCE
In this section, we use conjugate duality to derive the GIWAE, IWAE, and INFO-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 INFO-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, GIWAF, IWAE, and INFO-NCE arise from limiting the family of the critic functions To order climinate the intractable log partition function term.
We start from the decomposition of / (x; z) into Lune. (qa) and its gap
log 4o (2|x)) x) p(z)
1xx;2) Ep(x..)
+Epix) [DRI]p(2)))]l
(157)
Jasy (go)
We will focus on the dual representation of the normalized, conditional KL divergence 22(-) =
Dxul Dxul:ll9ozz/xx (z(x)) 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 samples from a base variational distribution (z x) to construct
K
Pror (z (1:K) s|xx (())((( x) II (z (x)
(158)
1=1
where s is an index variable U(s) drawn uniformly at random from S s S K, which specifies the index of the posterior sample p(z) x). We similarly expand the state space of the base variational distribution to write
AeROP (2:0)(((1:K)) (1:K)
,(x)) = U(s))
II
(z (k) lx).
(159)
k=
42

---

Published as a conference paper at ICLR 2022
It can be easily verified that this construction does not change the value of the KL divergence
Pran (2):K) s|x)
[Ru(Pror(z" (1:K) (l (1:K) s/x]] = (1
log
4prOP (x(1:K), s/xx
R (rur(a(1:8), log (4STp((z()) x) Ie 20 (z(20))
2AST II (z(x()))xx)
k=1
log p(z(=)/x (x)
(0(2(0) (x)
= Ballix) log p(z[x)
go (z|xx)
Consider the convex function
2(- = DKL 11 (0 (z(::K) s|x)]]
(160)
where the primal variable is distribution in the extended-state space of (2(::), ,)) and the dual variable a critic function T(x, g(1:k),s) We derive a conjugate optimization in similar fashion to App. K.2, but now over the extended state space. For this (2(.), the conjugate function (7) takes log-mean-exp form analogous to Eq. (137)*. We can write the variational representation of (p
as
((pre)) = sup
Prar (z (1:K)
s (1:K) .s)d log onsco (x110)) ((x,0110),00 .6)
(161)
For particular choice of T(x, z( (I:K), divergence DRi. D IBAL lower bound on MI, MMS- IRAI. (go; T) via Eq. (157).
s). we can use Eq. (161) to obtain a lower bound on the KL (as in Eq. (134)). This lower bound translates to the Multi-Sample
x(x;))= 2 pp(x))) log
96 (2/x²) p(x)
+ Pooe
[
T(x, (1:00
[0]]
log
wrear (1: (1: abld
Tor (TR)
Z
=: Lass-ma (qo, T).
(162)
where 2(x; T) is the normalization constant of the dual distribution क. 7z (11::K,s)x)) corresponding to T,
Tz s|xx) Z(x: T) GrRor (z (I:K)
s/x)e T(xx)e ((x,x(::8)
(163)
As in Eq. (133), we can write the gap of IBA (go. T) as a Bregman divergence or KL divergence
1(x: z) [gs-mar 7) + Eptx) DKL. [pror (1:K), (1.-)), s|))||=T(( ,|/)]]
(164)
The optimal K- K-sample energy function in Eq. (161) should result in ky- (2:1)) s/x) = Pror (z (l:K s/x). Using similar reasoning as in App. C.5 or App. L.1, this occurs for pxx)" c(x), for which we have [(x; = IMs-mar (qo. T*).
T T*(x z (1:K), 5) log (x()))
GIWAE is a Multi-Sample IBAL with 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 INFO-NCE arise from choosing a restricted family of functions Tawas (x, (1::K), s), under which the problematic log [(x; term evaluates to 0. This function family is defined as
We consider the conjugate with restriction normalized distributions in App. K.2 Eq. (136).
43

---

Published as a conference paper at ICLR 2022
Taiase (x, (1::), ,8) log (T((x,l)))
(165)
K
Σ
K
k=1
where TaIaaF (x, (11:K), 8) is specified by an arbitrary single-sample critic function (x( 2). We can now see that log Z TarwAR) = O,
(x. ()))
log Z(x; Terma))
urnurr (1:10 (1(n)
log E 20(s) vnta( bx)
Folaralt)
using the fact that (1:K)
II
go (k) xx) is invariant to re-indexing. With this simplification,
k=1
IMS-IBAL (go, TAIAAB) recovers the GIWAE lower bound on MI
[(x; 2) 2 IMS-IHAL (go, TaIWAR)
(166)
log (0(2(x) p(z)
(T((xa(0))
[x(x,x)
+
log
(167)
* #(x)p((())
11 go(2(1) (c) Aipia k=1
log go [z(x)) p(z)
ef((xx
E
log
(168)
p(x)p(z((1(p)) ft (x(
K
(x,((11())))
K
= Iawas TT. K).
Finally, we can use Eq. (164) to recover the probabilistic interpretation GIWAE and the gap in the lower bound on MI. As we saw above, the dual distribution Tamas (z (1:3,, x) is normalized with Z(x TaIWAR) = 1. In particular, we can write
Tonaa (1:K) ss(x)
ARROP (z (1:K) ,8 ,s/x)e (1.:)
(169)
2(x: TAIAE))
K
(To(xx(1))
20 (0(2(11
(170)
[(x; TaImse)
1==1
Σ (T((x
K
To (2)
(171)
K
Toe (x.,(())
which recovers FPRor (2 (211:), App. C.I. We write the gap of Tawau (go, To. K) las- TaIaa) the Bregman diver- gence or KL. divergence as in Eq. (133)
s|x) from the probabilistic interpretation of GIWAR in Eq. (42) or
[(x; 2) = LMS-MALL (go, Tawas) + Ep(x) [DkL prer (2(:)) s s/xx | (1:3)
(2 .s/x)]]
(172)
which matches the reverse KL. divergence Fx(x) [DIL. DKKIlpSCWAn (I (1:K) s/ x/x)
(1 GIWAR (x(1:)) s|x)]] de-
rived from the probabilistic approach App. C. Recall that INFO-NCE is special case of GIWAF with (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 special of GIWAE, where the optimal p(xx)
single-sample critic function T*(x, z) log [2|x] c(x) (see Sec. 2.4) is used in Eq. (165) to construct the optimal sample function TIWAs within the GIWAE restricted multi-sample function family. Thus, have J (go, TIwAs) JIwae (go, K).
44

---

Published as a conference paper at ICLR 2022
Although IWAE the optimal critic function, the restriction the function family in Eq. (165) necessary to obtain a tractable bound on the KL divergence DxL|p(zz (au[pp///)]ll (//o)///////// and mutual information. Without the restricted function family, the intractable log partition term in Eq. (161) would require MCMC methods such as ATS to accurately estimate, as we saw for the single-sample IBAL in Sec. 4.
L PROPERTIES OF THE IBAL
Our MINE-AI5 AIS method in Sec. 4 and App. M optimizes the Implicit Barber Agakov lower bound (IBAL) on MI from Eq. (18). We first recall the probabilistic interpretation the IBAL bound from App. J.1. For posterior approximation with learned negative energy function Tę and base variational distribution qa (2)x),
(z)x(
(x) to (z x)e
1
where Zu.a = Eno(a(x) (eFo(xz))
(173)
we consider the BA lower bound on MI,
1(x, 2) 2 Jnnn (To,e)
(174)
1(x, z)
Epix) ( (2|x)]
(175)
log (z(2))
(x(x2.)
(x(x))
+ p(x.))
log E (x)() [x_6(
(176)
p(z)
Lunz (go)
contrastive term
IBAI-(go: To)-
(177)
The gap of this lower bound on mutual information is Ep(x) D (xL.p(x/x)]l as in Sec. 2.2.
L.1 PROOFS FOR IBAL OPTIMAL CRITIC FUNCTION (PROP. 4.1 L.1)
Proposition 4.1. For given go (z|x),, the optimal IBAL critic function equals the log importance pix.2]
weights up to constant T (xx, = log (u(x)) +c(x). For this T*, we have IBAL(qe, T*) ) = I(x; z).
Proof. Recall that the gap in IBAL TΦ-) Rxp()) [[xx.p/aaxs/(|xoo<(x)))) [BAL TT., This implies that the bound will be tight iff p(z|x) (z) x p(x,z). z). We can easily show that the true log importance weights (plus constant) satisfy this property
T(x;2)
ex(x)
(z/x( = Za qo (z)
(z/x x)e
dpte +c(x)
Zx. p(p p(x, p(z|x),
(178)
where ex(x) is absorbed into the normalization constant. Conversely, using g(x, z) which depends on z in T(x. z) log
p(x.z) gu (x|x) + g(x, z) would change the density over no longer match p(z[x).
At this value, the IBAL is exactly equal to [(x; z)
qo p(x)
[BAL(go. T p(x.))
E p(xxx) [x,, Eptx) [log Zo.o(x))
Bp(xa))
(x(xxx)
p(x,x) By(xa)) log (0(x)x)
log c(x) Eppo log (ng
p(x,2) go (2(x) + log loge(x)
p(z)
log p(x, x) p(x)p(x))
Ep(xx)
I(x;2).
Proposition L.I. Suppose the critic function Tx(x,z) z) is parameterized by b. and that 30 s.t. Tx,(X, const. For a given {o(z)x), let To. (x, denote the critic function that maximizes IBAL(go. To). Then,
IBAL (qa) S IBAL(go, Tø-) S I(x; IBAL (go) Ep(x) [Dxx.p(eaxx))loo(<xx]]]]] (179) In particular, the contrastive term in Eq. (18) upper bounded by Fxpx)) D (p(llx()))) [2((x)))]]]
45

---

Published as a conference paper at ICLR 2022
Proof. The BA bound is special case of the [9AL To) with constant Tou and the contrastive term equal to O. Since Φo is a possible parameterization, we can only improve upon IBA (go) by learning T.. The parameterized family of a may not be expressive enough to match the true log importance weights, so [go TΦ-) S IBAT.(go, T*) = [xx; z) using Prop. 4.1.
L.2 PROOF OF IBAL AS LIMITING BEHAVIOR OF THE GIWAE OBJECTIVE AS K + 00 (PROP. L.2)
Proposition L.2 (IBAL as Limiting Behavior of GIWAE). For given go and To (x, 2), we have IGIAAE (go. To. K) 1BAL(qo,
lim
(180)
K K0000
Proof. Comparing the form of lawAE (go. To. K) to the IBAL( Tω) for a fixed go and To,
(o(xxx(1)) R EE e (To(teaa())
log g(<|x)
Tamms (go: To, K) Brp(o..))
log
p(z)
((x,a(1))
Af (x0k)))
K
k=
IBAL[qo, TT) Ex(x)))
log
[0(2)x) p(z)
+ Bp(x,))
log T. (x)x))
(x..)))
(e (x(())
we can see that both bounds include the same first term Lone (g0) and the same numerator of the contrastive term. To prove the proposition, we can thus focus on characterizing the limiting behavior of the denominator
K log
VX : lim E
c Tx(x,)
log E
e [o(x22)
(181)
K K00
p()
II go (z(())))))
K
go (x/x)
R=1
log ZarwAr (x, K)
log (x,
where we introduce the notation log ZaraR (x, K) for convenience and the right hand side is the log partition function for the IBAL energy- posterior (z.o(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 + 0o.
More formally, we consider the sequence of values log ZoIwAR K) derive lower and upper bounds on the value of log Zawas (x, K) for fixed K and show that each of these sequences of lower and upper bounds converge to log Zo, in the limit as K + co. Using the squeeze theorem for sequences, this is sufficient to demonstrate the claim that x log ZaIaaE K) log Zo,a(x) in Eq. (181). Since the other terms in I TAIWAE, (40, Tep K) and [B, are identical, this is sufficient to prove the proposition.
function of K. We
Lower Bound on log ZxIaRR(x K): We rely on the fact that the exponential function TT((x,) > to simply ignore the contribution of the p(z|x) term.
E To(xxa(t))
log ZeawAR (x, K) = E
log
llog K
(182)
pxx) II (x)) (x)
1=1
K
log Σ Tolxa(lt))
- log K
(183)
do (xx())))))
k=2
K Σ k=2
E
log
Tolxa(t))
+ llog
-
(184)
(x(0)))
K
K
Lx;
(185)
where in (1), we obtain a lower bound by ignoring the contribution of the positive sample from p(z)x). The first term is the k-sample IWAE lower bound with the proposal go(z) and target (x, z):
(x,2)))))))
( (oaax)/oogg flog K Σx 20(2 x)] and thus converges to log Ea,s as K - co. As K + we also
46

---

Published as a conference paper at ICLR 2022
have log K-1 - 0, so that the limit of their sum is K
lim LB(x; K) = log g (x)x) K-000
[ = log (x,
(186)
Upper Bound on log ZoIaAR (x; K): To upper bound log Zaiar (x, K). separately consider terms arising from p(z]x) samples and go (2(x))) samples. Noting that PlOrA (anrA((((lgg (z(1:K), in Eq. (26) is invariant to the index
assume z p(z x) and write
log Zarwne (x, K)=
E
log πr
p(x(1))))) no (x6
K
K
e Foleaalht
E
log
p(a(1))
(2(
K
K
K
km
(1)
S log
K 18 (ppc00110)))
(
+ E
K-1 1
To (o(x)
11 (0(2)())2))
K
K
log
1
18 (((((11/2))
(
+ K-1
1
Σ
E Bxca(t) x) [ (eToleatinn)
K
K
K
log
1 E Paccossps) K
K K-1 +
20,0(x))
(187)
K
UB(x; K).
(188)
Since log(u) is continuous function, we know that lim (og(f(KK) log lim f(K). We thus reason about the limiting behavior of the terms inside the logarithm in Eq. (187). As K -> oc, we have O, and thus the first term inside the log goes to O. For the second term, also have + 1. Thus, we have
lim UB(X;K) = logg E (xxxx)
[T(xx))) log Zx.((x).
(189)
As reasoned above, the convergence of the sequence of both upper and lower bounds to log Zo. (x) implies that 1imo log Zaraa (x, (x, K)= log Za (x). By the reasoning surrounding Eq. (181), this
implies lim JawaAR (4o, Ta) = Iqo To) as desired.
L.3 CONVERGENCE OF GIWAE SNIS DISTRIBUTION TO IBAL ENERGY-BASHI POSTERIOR
In this section, we consider the marginal SNIS distribution of GIWAF, which is induced by sampling
K times from go (z/x) and returning the sample in index s with probability GIWAE fs|x, z (1:8)) x
(T((x,(")), As K oc, show that this distribution converges to the single-sample energy-based posterior approximation underlying the IBAL and MINE-AIS. similar observation is made in Sec. 3.2 of Lawson et al. (2019). 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. Define the marginal SNIS distribution of GIWAE, GIWAF (z) (z)x: K), using the follow- ing sampling procedure
1. Sample from 8, z (1:K)
GIWAF
(1:K() according to Eq. (42)
2. Return ())
Then, as K oo, the KL divergence (in either direction) between the marginal SNIS distribution of GIWAF and the energy-based variational distribution of IBAI, (z)x) x qo (z(2)) (e (x.22 zero.
goes to
lim DKI. | (2|x; TT To, (z) xx]
GIWAE
K)] O. and lim DKL (T, 11 1-x
GIWAE (z)
47

---

Published as a conference paper at ICLR 2022
Proof. To prove the proposition, consider a mixture target distribution similar to the case of IWAE. However, in this case, we take a single sample from To,o (z)xx instead of p(z|x)
K
OIWAB Prot
(s, (k:z(1
1 K no (z
xx
II
go (z
(*) x).
(190)
GIWAR, Note that the probabilistic interpretations (s 2z (1: xx) and Prat (s. z ((::k), 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 To (z)x). Thus, we can analyze the KI. divergences
Dri. [xr_[uxxp== (1: / ]lPteer (OWAE, (8, (1: [x)] and DxI. [Prer (s, zz ((::K) x)] VPROPE (8, (1 |x)] using techniques from previous work IWAE. Following similar arguments
in Domke & Shel- don (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. DEL CAENAE (z|x; K" go. [z(x))
DKL [q (s, 21:: |x) ||ar | (s, ((x:)(x)) [x)]. As K cc, (s, (1: (x) ||xx|lprrr |llpoet (s. z (1:K) x)] + since an instance of the IWAE gap. Thus,
DKL 19pp0r the KL divergence DxL D (z|x: To o x)] also vanishes. Similar reasoning applies for the reverse KI. divergence.
M MINE-AIS
M.1 ENERGY- ENERGY-BASED TRAINING OF THE IBAL
Eq. (19) indicates that in order to maximize the IBAL as a function of f and ø, we need to increase the value of the energy function To or score function log (o(z)x) on the real samples of the true joint p(x,2 2) and decrease the value on fake samples from the approximate p(X))TO,d (z/x). However, as is common in training energy-based models, is difficult to draw samples from TT..d-
In order to sample from e (z) x). natural approach is to initialize MCMC chains from sample of the base distribution qe (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 a (z(x). Instead, we can choose to initialize the chains from the true posterior sample p(z[x) for a simulated data point transition steps, the approximate energy function gradient becomes
p(x). Letting TEM indicate the composition of T
0 [go To) R Ep(x.2)
To (x, 2)) - E (4/2,))) Os
To To (xx Oib
(191)
We can use an identical modification for the gradient with respect to A.
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 in spirit similar Contrastive Divergence learning of energy- models (Hinton, 2002) where one starts an MCMC chain from the true data distribution to obtain lower variance gradient estimate.
M.2 MULTI-SAMPLE AIS EVALUATION OF THE IBAL
After training the variational base distribution qe (z|x) and the critic function [x(x, z) using the MINE- AIS training procedure above, we still need to evaluate the IBAL lower bound MI, IBAL Ts). We can easily upper bound IBAL(qe T&) using a Multi-Sample AIS lower bound on log Zg. a (xx) with expectations under the forward sampling procedure AIS [x): However, an upper bound on IBAL(qe, 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 Zo,o(x) (x) log (e)x) [e=fo(x.x), the log partition function of Ta.s (z|x). However, consid- ering this to be the target distribution (z)( x) in Multi-Sample AIS, would require exact samples from To, (z(x) to guarantee an upper bound on log (x). Since these samples are unavailable,
demonstrate conditions under which we can preserve upper bound on log Zx, (and lower
48

---

Published as a conference paper at ICLR 2022
bound IBAL(9g. T&)) by sampling from p(z/x) instead of Tz. to initialize backward annealing chains in Prop. M.1 below.
Using the single sample AIS bounds estimate log Z (x), we have the following extended state space proposal and target distributions,
T
AB. (Zo:T|x) = (0 x) (z(2-11
GRROP (Zo.T)x) = qo(zo x) II Te(zelze-1).
(192)
We emphasize that in Pror
MS.* AIS, the transition kernels and intermediate densities are based on the critic To(x, 2) and energy-basee posterior (z(x( whose log partition function seek to estimate. Recall from Sec. 2.1 that taking the expected log importance weights under Prer" (Zo:T) x) yields an upper bound on log F0,((x), However, since it is difficult to draw exact samples from To. (z/x) x) to initialize backward annealing chains and sample from Preí AIS.T (zo:r) (x), we instead consider sampling from the posterior p(z] p(z|x). Using the same transition kernels T, as above, we first define the conditional
distribution PiGr (Zo:r-1 Xx Zr) of the backward chain for a given 27
T (Zo:T-I 27) = II T2 (7 (2.)
AIS.
(193)
We can then define the approximate extended state space target distribution which samples zy p(z|x)
PTor
APPROX
AIS.T (Zorl)x) = p(zr x) II Tz p(2rp((())pppp) (Zo.r-1|x, 2r).
(194)
AIS,T Using this notation, also write Prer (zo:r (x) no, (2r x Prei AI8,5
(Zo:r-1|x, ZT) in Eq. (192).
We now characterize the conditions under which sampling from Prar (Zo:T) preserves an upper bound on log (x).
APPROx
Proposition M.1. Define the AIS marginal distribution dFRor (ZT over the final state in the MS..
extended state space proposal as QeRoP (2T(x) ) == AtS,T (Zo:T )d) If we have
Dxı.p(zz [x) (eeoe AI8, (Zr[x)) (xI.pp(zT (mx
(195)
then initializing the backward AIS chain using zr ~ p(z|x) (i.e., sampling under ApprOX (Zo:T|x)) (x)), yields upper bound log (x).
p v (x, zu:r) AIS.7 (Zo.r (x)
Epenpoo (mo
log
log (x).
(196)
Proof. We begin by writing several definitions, which will allow factorize the extended state space proposal AIS (Zg:r(x x) in the time-reversed direction. This factorization includes the final AIS marginal GRROP (Zr x).
Starting from Eq. (192), the forward transitions AIS. (zo:r xx) = qe xx) I[ T((z,|z,1) induce the marginal distributions AIS.T (2/(x) at each step. The forward transitions and marginals induce posterior kernel ,." (z.-14 (e-1/2,) (z.), which allows to rewrite AI8. (Zoer x) using a reverse factorization
A15.5 GEROP (Z).T|x))
AIS. OFROP
(zT (x)
II [=1
T. (z Zc).
(197)
where (zt-11122)
OrRor
A18.7
(z4-1 1 |x (2x A18.
(198)
(z, [x)
The posterior reverse transitions 7.7 (Z(-1/zz)) are intractable in practice, and cannot be simplified
to match the kernels in the target distribution Pret ( (T) I Te 1/2)) using the
49

---

Published as a conference paper at ICLR 2022
invariance or detailed balance conditions. Doucet et al. (2022) provide promising approach using score matching to approximate these posterior transitions T, (Z)
Finally, we write the posterior reverse process conditioned on a particular 27 as
ATS,T OFROP
(Z.TT1
TV (z (z)))
(199)
t=1
With the goal of upper bounding log (x) (z)e dz, we consider the log importance weights with expectations under the target distribution APpROx (zo:r|x) = p(zr|x] x x)por A (Zo:T-1|x, ZT), as in Eq. (196). Using the above notation, we have
AIS,T (x, Zo:r) log
A18, (Zo:7 x log (x, + Bepg log
Bppano (x))))))) Plot
(200)
ATS,7 (zo:r) x)
Venor (Z0:7 x)
= llog (x)x(x) + E Epprpion (s0.r xc) log mi (zr xx x)pr (Zo: |x, 2T) p(zr xx)
(201)
A18. (zq) x
Iyrur (Zo:T |x, zr)
AIS.
p(zr ((zr|x) (x)
= llog (x.+(x) + DEL Lplz xx)le (zrl [x]] - DxL
(202)
+ Ex(er(e) [DLL pror Als." (Zo:T- X, Zr)
AI5.5 (Zo:T - X, 2r]]
The intractable KL divergence DkL (PAGr (Zo.T-I|x, ZT) [ AIS,5 (Zo:T - xx, ZT) compares the re-
AIS,T
verse kernels in the target distribution Te -z Zi) against the posterior T, (7(z-11 (z.) Ignoring this nonnegative term, can lower bound the expectation on the LHS
EAAPRO
log
Prar (x, Zo:T)
11 (Za.T log (x) DKL (xlzrl))]lg AIs. (zr[x) DEL [x(zr [x)]
Finally, under the assumption of the proposition that DXL [p(zr [x] AIS. (zr (zr(x]]
DXL (x(rr/x)lloo, 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 by initializing the reverse chain using posterior sample, under a condition the KL divergence,
DEL [p(zr ([f/zrx)llqgma ( AIS. (x))x) 2 DKL. bsl.
(203)
While we cannot guarantee this condition, we intuitively expect Eq. (203) to hold in practice since T (z|x) has been directly trained to match p(z|x). By contrast, 4rzor (zT (x) is the final 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 EBMS, but give an example of a Restricted Boltzmann Machines (RBM) model (in their Sec. 5) in which Eq. (203) does not hold.
As desired, find in our experiments in Fig. 7 that approximate reverse annealing procedure underestimates the IBAL in all of the MINE-AIS, GIWAE and INFONCE experiments, for all numbers of intermediate distributions T.
For T' 1, we have TeROP (zi(x) = qo (zi xx.. In particular, Prop. M. will provide an upper bound on log Ze (x) if
T 1 Special Case
D
2 (xx.
(204)
This condition is guaranteed under the assumptions of Prop. L.1, where IBAL(ge. To-) IBA, (To. improves upon Iasz (qo) for an energy function T (x, z) which has been trained to maximize the IBAL.. In Eq. (204) the KI. divergence the left-hand side corresponds to the gap in the BA lower
50

---

Published as a conference paper at ICLR 2022
27
126
VIGT OAX
26
18
15
2F ONL SOACTTAA
222
33
22
Sesterid Dishadore
(a) Linear VAE
(b) MNIST VAE
(c) MNIST GAN
Figure Estimating the mutual information of deep generative models with AIS and MINE-AIS.
bound, which is larger than the gap in the IBAL on the right-hand side. We can further show that the lower bound on IBAL (qo, Ts) resulting from Prop. M.I with T 1 the BA lower bound.
go [x(x))
[gA TT) = E)(x,x)
log
+ lE px(xxx) [x, z)] Ep(x)) llog [x))
p(z)
log go [z|x)
Pror (x,, Zu:r
AIS,
2 (x.x)
+ px(xx) (ovxxx)/T((xx) (xx z)) -Eparmox
log AIS.7 (Zo:T) x)
p(z)
log 20(2)xx
log 4o (z) TT (x(22)
Ep(x,x)
+ 1 Ep(x.)) (T 2)] - Ep(xx)
p(z)
2o (0(2
(0(z)x)
Ex(xxx)
log
+ E (T (x, 2]
E E [Ttx
p(z)
LBAZ (go)
We can confirm this in Fig. 7, where for 7 1. the approximate lower bounds on IBAL Tω) begin from the appropriate BA lower bound. For example, in the case TBAL evaluation for GIWAF (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' I) for the AIS approximate bound in Prop. M.1, the estimates in Fig. 7 approach the true value of [g, Tτ) in all cases.
M.3 MINE-AIS EXPERIMENTS
In Fig. 7, we assess the performance of our MINE-AIS method, which learns an energy function To and uses Multi-Sample AIS to evaluate lower and upper bounds on the [gAL(ge, TT), in cases where p(x|z) is unknown. As 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 (pAL(p(z). T') for p(z) = N (0, f), 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 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, finding the true value of (p(L(p(z), T.).
We find 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 Gaussian variational family for qo (z)x).. Note that the IBAL corresponds to the BA lower bound for a more flexible, energy-based posterior approximation o (z/xx, and both bounds assume access only a single known marginal distribution. We find that MINE-AIS outperforms the BA baseline by 2%, 52%, and 90%, linear VAF, MNIST-VAF, and MNIST-GAN, respectively.
51

---

Published as a conference paper at ICLR 2022
We also Multi-Sample AIS to evaluate the IBAL corresponding (ge. To), which are learned by optimizing the GIWAE (with K = 100) and INFONCE (with = 100 and qe = p(z)) lower bounds. As argued in Prop. L.2, the IBAL corresponds to the limiting behavior of GIWAE K - DC. We observe that the AIS evaluation of the TRAL corresponding to GIWAR or INFONCE critic functions only marginally improves upon evaluation of the original GIWAR INFONCE lower bounds. This indicates that the improvement of MINE-AIS over GIWAE or INFONCE 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 T.
N APPLICATIONS TO MUTUAL INFORMATION ESTIMATION WITHOUT KNOWN MARGINALS
While the focus in our work is evaluating the mutual information [(x; z) in settings where at least single marginal is available, we are often interested in estimating optimizing the mutual information where no marginal distribution available. A natural setting where no marginal distri- bution available is representation learning where the goal is to maximize the mutual information between the data distribution Gdata (x) and the representation induced by a stochastic mapping de ge (z|x) parameterized by 45
[(x; z) Eyan (2(x) log Loo (x, 2)
log To (x (z)
qe (x)g (z/x)
(205)
Idma qu.
Idata (x)
where (z) I ydata (xnn(())yv (z)x)((()))))))) represents the "aggregated posterior" (Makhzani et al., 2015) or induced marginal distribution over z.
N.1 BA LOWER BOUND
Using the notation of Eq. (205), we can write the BA lower bound as
[xx;2)= 2 1(x;z)= (gy()))
DRi Iaa (x|2) [lpo 2)]
gap
Po (x zz) log
Regar (x) Bu (2)x)
(206)
Idinta (x)
-Bgan(x) (x) En (x)
log pa (x(z)
+ H (x)
(207)
negative reconstruction
Innz (x)))
(208)
where po (x|z) is a variational distribution, parameterized by 0, that tries to match to the inverse encoding distribution qu (x(z) x Qdam (x) (z|x). The first 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 Constant Note that the gradient of Laz (Po z)) with respect to the parameters 0 does not depend on the marginal (x), and thus we may still optimize the variational distribution when the data distribution unknown. We can then use the resulting po (x|2) to estimate the BA lower bound on mutual information up constant, H (x). This is useful in comparing the MI induced by two different representations qya (z|x) and Tw2 (z) x) of the same data distribution.
Optimizing MI with the BA Lower Bound The BA lower bound is also amenable to backprop- agation through the parameters of the encoding distribution gx even when analytic marginal densities for Qdaa (x) or qu (z) not available.
Maximization of the BAA. lower bound appears in various settings, including in representation learning (Alemi et al., 2016; 2018), reinforcement learning (Mohamed & Rezende, 2015), improving inter- pretability in GANS (Chen et al., 2016), and variational information bottleneck methods (Tishby et al., 2000; Alemi al., 2016; 2018).
52

---

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 ls (z)x), since each term in Eq. (209) is amenable to backpropagation. Since GIWAE generalizes both the BA and INFO-NCE Published as a conference paper at ICLR 2022
N.2 GIWAE LOWER BOUND
In this section, discuss the applicability of GIWAE lower bound in mutual information maximization settings. Rewriting the GIWAE lower bound in Eq. (7) for the case estimating [xx; z) in Eq. (205), we have
5 (x.0))
[xx x) (x (el (ala)) (x]le) Hxmi( + (8.0)) ft pg(x(())p)
log
* (e (o (K.00) + Ex To(n (0))
nagarise
(209)
Famasy (pe (x/2). TT. K)
Note that we can use a single joint sample from qdma( [z]x)- = qu (z)gg (x|z) to obtain a positive
sample from the inverse encoding distribution x (x/z))) for particular (z(z,, in similar fashion to our ancestral sampling in Sec. 1.1. For given z, the negative samples correspond to K samples from the stochastic decoder x ((2:K) (x(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 (z) ) for a particular z ~ qo(z)) The contrastive
term in Eq. (209) arises from SNIS sampling of a single x(*) in the extended state space proposal, with UpRoP g x (2:K))
(x(x
Σm= / efo(((t)))))
Finally, we note that optimization over the energy function (x, z) improves upon TAA (po (x(2)) 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 quau (x) is unknown, This may be useful in comparing the MI induced by two different representations qe, (z|x) and Tv., (z(x) of the same data distribution. Note that the reconstruction and contrastive terms Eq. (209) depend on samples from Gdata (x) and not the density. Thus, our ability to take gradients with respect to the parameters the variational distribution Pe (x|z) or energy function (x z) are not affected by the fact that the marginal distribution is unknown.
lower bounds, can be used as drop-in replacement for either these bounds for optimizing MI (e.g., see van den Oord et al. (2018)).
N.3 MINE-AIS/ / IBAL LOWER BOUND
Recall that MINE-AIS estimation would involve an energy based variational approximation to the inverse encoding distribution (xlz)
T0,0 = 2(z) (e
(210)
Note that IBAL lower bound on ((x;2) involves an intractable log partition function term (2) [log E(z)] = (x) log Exo x) [oxxx)])]].
(x,x))
1(x;2)= 2 Exsn(()) By (x|x)]
Hama(x) +
Egg log Ep (xc m)
(211)
(x
negative reconssruction of
contrastvve term B (1)) DDa (xc (x(l)))
=: (x
where the term in the denominator is the partition function Z(z) (x(()))) [eFo(x.2]]
53

---

Published as a conference paper at ICLR 2022
When taking gradients in Eq. (19)-(20), obtain
0
0
(pa, To) - B gx (x/x))
log po (x(())
Egg (w) (xc(e)
log po(x
(212)
00
00
30
0 O6
0
IBAL ( TT) = (x) (x)bx)
To (x, 2))
(x) (x/(2)
T T z)
(213)
06
To obtain approximate negative samples from na, (x/z) for a given
MCMC transition
kernels since the unnormalized target density a (x(z) Po (x (ox is tractable.
Evaluating MI up to a Constant Since only samples from qama (x) are required for the MINE-AIS training procedure in Eq. (212)-(213), we can learn the base variational distribution Po (x and energy function Ta (x, z) cases when the marginal distribution gama unknown.
As in the case of GIWAF and BA lower bounds above, we can also evaluate the TBAL 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)) = log (x(()x)) [x.)]
Optimizing MI with the MINE-AIS Lower Bound If we are interested in optimizing the MINE- AIS lower bound with respect to the parameters of stochastic encoder qe (z x), would need to backpropagate through the Solti-Sample AIS procedure used to estimate the log ((z) term, but computing the gradients intractable. We thus conclude that, among our proposed methods, GIWAE the most directly applicable in settings of MI for representation learning.
O EXPERIMENTAL DETAILS
O.1 EXPERIMENT DETAILS OF SEC. 5.2
In this section, we provide the experiment details used in Sec. 5.2. For more details, see the public GitHub repository,
O.1.1 DATASETS AND MODELS
We used MNIST (LeCun et al., 1998) and CIFAR-10 (Krizhevsky & Hinton, 2009) datasets in our experiments.
Real-Valued MNIST For the VAE experiments on the real-valued MNIST dataset (Table 1), the encoder's architecture is 784 1024 1024 1024 2, where 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 VAES. In order stabilize the training dynamics, 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 opti- mizer (Kingma & Ba, 2014), and the checkpoint with the best validation loss was used for the evaluation.
For the CIFAR-10 experiments (Table 1), experimented with 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.
CIFAR-10
O.1.2 EXPERIMENT DETAILS OF SEC. 5.2
For the AIS temperature schedule, We used sigmoid schedules used in Wu et al. (2016). The step size HMC was adaptively tuned to achieve an average acceptance probability 65% as suggested in Neal (2001). For all MNIST experiments in Table 1, evaluated simulated data. For all CIFAR experiments in Table we used single batch of 32 simulated data. All Tesla P100 Quadro RTX 6000 Tesla T4 GPUs.
single batch size of 128
experiments
54

---

Published as a conference paper at ICLR 2022
0.1.3 RUNTIME COMPARISON
We benchmarked the runtime on Tesla P100 GPUs. For MNIST, it took about 35 minutes to run IWAE with K LM, about hours run AIS with T = 30K. For CIFAR, it took about 45 minutes to run IWAE with K = 1M, and about 12 hours for the AIS with T = 100K.
In Fig. 8, we evaluate the tradeoff between runtime and bound tightness for evaluating the generative MI for VAE and GAN models with 100-dimensional latent codes trained on the MNIST dataset. We compare IWAE and multi-sample AIS evaluation, with the same experimental settings as in App. O.1.2 and an initial distribution qo [z) x). We plot wall clock time on the z-axis, where increasing runtime reflects increasing K for IWAE and increasing T for AIS.
202
TI4E Lower Bound
IMAE Lower Sound
175
Semblle AS Lawas Bound
tipper Boind
175
Mulli Sumpie Lower Roard
HMAF Itppr found
Upper
M1 Exteraaees
125
Estimmtos
135
75
20
32
25
25-
0.00
320
0.00
0.3
Miniees
Minutes
(a) VAF-100
(b) GAN-100
Figure 8: Runtime vs. Bound Tightness for VAE (left) and GAN (right) models on MNIST.
O.2 EXPERIMENT DETAILS FOR ENERGY-BASEI BOUNDS (SEC. 5.2)
Models and Data The linear VAE model has Gaussian prior z ~ NO, 7) with the dimension of 10. The dimension of the output x is 100. The weights are sampled randomly from Gaussian distribution with the standard deviation of 1, and the standard deviation of the Gaussian observation noise at the output is 1. The MAF22, is a VAE model that is trained on the real-valued MNIST dataset. It has Gaussian prior ~ (O, I) with the dimension of 20. The decoder has one layer of ReLU non-line non-lincarity of size 1000, followed by a lincar layer that predicts the mean of a Gaussian observation model with the fixed standard deviation of 0.1. MNIST-GAN20 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 fixed standard deviation of 0.1. In all the experiments, we use fixed batch of 100 data points.
BA, IWAE and GIWAE For the BA bound, (z(x))) is parameterized using 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 qe (z) x). The GIWAE experiments, in addition to go (z|x) with the same architecture, used a critic function (xx, 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 scalar value. The parameters of go (z(x)) and T7 (x, 2) are trained jointly.
Multi-Sample AIS For the Multi-Sample AIS bounds, for all models, we used up to K = 1000 chains (see Fig. 5), and up to T = 50K intermediate distributions with linear schedule (see Fig. and Fig. 7). We used HMC the AIS kernel, with 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
(O I) as the base distribution go (z)x). In order to take sample from no.s (z,
used the
HMC method that is initialized with a true posterior p(z|x) sample, with M = 10 iterations each with T. 20 leapfrog steps. For the step size of HMC, we used 0.05 for the linear VAE model
55

---

Published as a conference paper at ICLR 2022
and c = 0.02 for the MNIST- VAE20 and MNIST-GAN20. For the evaluation of MINE-AIS, used Multi-Sample AIS with the same parameters as the AIS evaluation experiments described above.
O.3 ANALYTICAL SOLUTION OF THE MUTUAL INFORMATION ON THE LINEAR MNIST-VAL
In order verify our implementations, and verified that it matches the MI estimated by AIS. For simplicity. we assume fixed identity covariance matrix decoder of the VAE is simply:
have derived the analytically for the linear VAEs
the output of the conditional likelihood of the linear VAE decoder, i.e. the Wz+ b + e, where x is the observation, z is the latent code vector 2 (O, I), W the decoder weight matrix and b is the bias. The observation noise of the decoder (O, I). It is easily shown that the conditional likelihood is p(x|z) (xx + b, I) and
thus we can solve for the marginal
p(x) NXxx b, Σx = I+ WWT).
(214)
The differential entropy of x is:
k k + log(2r) + log (det Ex),
H(x)
(215)
2 where k is the dimension of the observation. The conditional entropy is
2
2
k + log(2z) 2 log(det I)
H(x
(216)
2
k
2 + log(2x).
(217)
Thus, the mutual information is
1(x: H(x) (x(2z)
(218)
2 log(det Ex)-
(219)
0.4 CONFIDENCE INTERVALS FOR MULTI-SAMPLE AIS EXPERIMENTS
Table and Table 5 provides the 95% confidence intervals for the AIS results reported Table The confidence intervals were computed over confidence interval over batches each with 16 data points for MNIST; and over 8 batches 4 data points for CIFAR.
56

---

Published as a conference paper at ICLR 2022
Method
Proposal
VAE10
VAE100
GAN10
GAN100
UB
(1849.03,
2010.66)
(5564..3, 6096.51)
(745 68, 826.55)
(795.81, 826.99999))
AIS
x(x)
LB
(0.00.0.00
(0.00, 0.00)
(0.00, 0.00)
(0.00, 0.000
T=1
UB
(59.34, 66, 66)
(345.45, 378.80)
(293 60, 335.88 83)
(482. 40, 0022440,444.2 544.26)
x(2/x)
LB
(60.09, 66.61)
(32/46, 36.52 52)
(3941
(2.40,2.80
UB
(38. 61. 39.57)
(92.32, 98. 03)
(21.91, 23.04)
(26. 80, (2.222
p(x)
AIS
LB
(33.69, 34.41)
(77.77, 82 22.02)
(21.18, 21.96)
(25.36,26, 26.36)
T-500
UB
(33. 613.94.33 34. 34.63)
(79.96, 94.71)
(22.55 23.58)
(28.82, 22,30.24
((2/x))
LB
(33.80,34.63
(78.01, 82:37)
(21.19, 22.00)
(25.09, 26,06)
UB
(33 (33.82.34.0)
(78.62, 83.07)
(21.54, 22.50)
(25.97. 27.07)
p(x)
AIS
LB
(33.85 34.56)
(78.51, 83.05)
(21.50,22.0
125.96,26.90
T=30K
4(2/x)
UB
(33.85, 34.58)
(78.56, 83. (3.03)
(21.54,22 22.47)
(26.01, 27 (07
LB
(33.86, 34.56)
(78.52, 02)
(21.55, 22.55,22.45 22.48)
(26.02, 27.03)
UB
(3628.
86, 4026. 29)
(10705.98% 12297.86)
(1556.00 00, 1704.00)
(1680.50, 1800. 28)
p(z)
IWAE
LB
(0.00, 0.00)
(0.00,
(0.00, 0.00)
(0.00,0.00
K=1
g(2(x)
UB
(34.82, 35.86)
(92.82, 98.44)
(38.80) (8..80,76.14)
(243.34,278 278.41)
LB
(24.54, 25.86)
(41.53, 47 .4)
(3.85,
(..94,3.52)
p(z)
CB
(1132.54,
1262.97)
(3987.51 (47.5114480..0 4480.86)
(430-40, 463..40,63..00 463.19)
(462.95 95, 526.51)
IWAE
LB
(6.91. 6.91)
(6.91,6.9) 6.91)
(6.91, 6.91)
(6.91, 6.91)
K=IK
((2/x))
C
(33.87. 34.61)
(83.15, 87.46)
(34.11, 71.36)
(182.82. 219.55)
LB
(31.08. 32.33)
148.44,54.48
(10.75, 11.52)
(9.84, 10.43)
CB
(352.91,
400.87)
(2078.30, 2417.17)
(71.93, 91.09
(100.27, 127.75)
p(z)
IWAE K-IM
LB
(13.82, 13.82)
(13.82, 13.82)
(13.79, 13.82)
(13.82, 13.82)
g(2(x)
UB
(33.87 333.87.34.57
(81.42, 85 35) (55.35, 61.36)
(28.51,33 33.26)
(62. -1, 63.68) (16.65 17.32)
LB
(33.67, 34.53)
(17.31, 18.21)
Table 4: Confidence intervals of AIS and IWAE estimates MI on MNIST. UB stands for Upper Bound, and LB stands for Lower Bound.
Model
Proposal
GANIO
GANI00
UB
(3593559. 12, 4477712.38)
(4038603.0g 09, 5668217.91)
p(z)
AIS
(0.00, 0.00)
0..00,0.00)
(T=1)
G(2/x)
UB
(241342 62, 566015 82)
(1874938.5) 58, 2881676.42)
LB
(15.70, 89)
(18.41, 21.93)
UB
(24683. 41496.70)
(25453.02, 101 101127.79)
p(z)
AIS
LB
128.41, 30.63)
(98.37, 110.65)
(T=500)
UB
(116 02, 156. 28)
(1603. 13, 3969.92)
((2/x)
LB
(45.45, 50.88)
(134.35, 156.02)
UB
(72.17, 75.80) (70.19, 73.55)
(479.04, 497.10)
p(z)
ATS (T=100K)
L.B
(470.05, 490.47)
UB
(71.87, 75. 22)
(475.07, 494.61)
o(x/x)
L.B
(71.35, 74.75)
(468.64, 489.89)
UB
(6019901. 11,, 91199511.9999 (0.00, 0.00)
(7339-111.95 12492792.05)
p(x)
IWAE
LB
0.00)
(K=1)
UB
(75.00, 80.03)
(-181 (-1810.70. 125044.11
((z|x)
LB
(16.28, 18. 62)
(18.39, 21.60
p(x)
UB
(1673333.00 .0,, 2 2415008.47)
(2181660.84 84, 3531363.7880 16)
IWAE
LB
(6.91, 6.91)
(6.91, 6.91)
(K=IK)
UB
(72.31, 75.68)
(- 12, 12438.39) (25.32, 28.
o(z|x)
LB
(22.18, 24.99)
UB
(582880.8 62, 838142 63)
(1479829. 44, 2327879.55 (13.82, 13.82)
p(x)
IWAE (K=IM)
LB
(13.82, 13. 82)
UB
773, 75.08) (29.86, 31.60)
(- 1883.22, 22, 12426.34) (32.21, 35.42)
g(z/x)
L.B
Table 5: Confidence intervals of AIS and IWAE estimates of M1 on CIFAR. UB stands for Upper Bound, and LB stands for Lower Bound.
57


Reviews:
[{"review_id": "qUpyFBTBzZ0", "reviewer": "Reviewer_eNs9", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "First I identify myself as Reviewer cNMQ from the last round of peer review of this paper. I see the author(s) have carefully revised their manuscript to address the concerns from the reviews. In particular, the current draft has made clear the distinction between parametric and non-parametric MI estimation, both are important research topics and they are fitted for different application scenarios. Not only that, the presentation has been improved. Better visualization of the bounds has been provided to help understand the relationship between different estimators, and there are many in-place remarks and discussions on the connection to existing bounds. I believe this is a piece of solid work and enthusiastically recommend acceptance. \n\nStrength \n* Importance sampling is a key technique in statistical inference and its application to MI estimation constructs a promising direction for research. Such motivated, this paper has derived a family of multi-sample novel bounds with annealed proposals.\n* The connection between the InfoNCE bound and the generalized multi-sample IWAE MI bound is quite fascinating. The author(s) have made an interesting point in the statement that the naive choice of proposal p(z) has led to the well-known log-K cap of InfoNCE bound, which I believe is critical for improving MI estimation. The current presentation is underselling its significance, so I would encourage the author(s) to enrich the discussions and give more concrete examples for easier understanding.\n* Clarity of presentation. I am glad to see the significant improvements made compared to the last submission. \n\n\nBelow are some minor suggestions that might help the author(s) to further improve the quality of this manuscript. \n\nTechnical points:\n\nA. While the importance sampling perspective for MI estimation is novel, this is not new in the broader context of variational inference. The following works [1-3] are highly relevant: they all follow an annealed path to tighten the bound. The discussions on these works are currently missing from the manuscript, which I think would significantly boost the significance of this work by making strong connections. \n\n[1] T Salimans, et al. Markov Chain Monte Carlo and Variational Inference: Bridging the Gap. ICML 2015\n[2] L Maaloe, et al. Auxiliary Deep Generative Models. ICML 2016\n[3] V Masrani, et al. The Thermodynamic Variational Objective. NeurIPS 2019\n\nB. Another major missing piece is the discussion of whether this strategy may inadvertently harm learning. This point is related to the seminal work of [4], which analyzed a particular scenario that tight multi-sample VI bounds actually harm inference. Intuitively, the inference arm is the approximate posterior (i.e., proposal distribution), and since tighter bounds are less sensitive to the quality of approximate posterior, the algorithm is less incentivized to improve it. Related to my technical point A, [3] has provided some perspectives on additional tradeoffs involved. Nonetheless, a careful analysis is beyond the scope of this study, so I only encourage the author(s) to give some thoughts and make some remarks in the paper. \n\n[4] T Rainforth. Tighter Variational Bounds are Not Necessarily Better. ICML 2018\n\nC. The derivation of GIWAE is one of the key contributions of this paper, but there is no intuition provided in the main text on how this is derived. The author(s) simply say \"we define GIWAE lower bound\". At least the probabilistic interpretation should appear in the main paper. Also, the author(s) should mention the disadvantage relative to the InfoNCE estimator. In particular, InfoNCE admits an efficient bi-linear implementation [5], because the z is sampled from the prior. I think the difference between IWAE and GIWAE should be further emphasized, say using a table comparing what is needed for both of them, and changing the order of the first two sentences of Sec 2.3 \n\n[5] T Chen, et al. A Simple Framework for Contrastive Learning of Visual Representations. ICML 2020\n\nD. Sorry to see the section \"Energy-based Bounds\" in the original writing moved to Sec 4, personally, I think it can be presented earlier.  And I do not fully appreciate the importance of Sec 3.3, could the author(s) elaborate?\n\nE. Selecting K, T is a joint optimization problem, but the discussion in Sec 3.4 is singled threaded: it only talks about suggestions when either K ot T is fixed. \n\nNon-technical points:\n\nF. Please add runtime analysis, and highlight the tradeoff between computation/memory and bound tightness. \n\nG. Please give the explicit expression of I_{IWAE} in the paper. \n\nH. Please use the booktab package to format the Tables. Follow the practice of minimalism, remove unnecessary horizontal lines and all vertical lines in the tables. Actually, for Table 1 a figure might be more intuitive. \n\nI. Figure 2 (b) is excellent. I think it is better to present it earlier in the text, right where the bounds are discussed. ", "summary_of_the_paper": "This paper investigates annealed importance sampling (AIS) techniques and multi-sample bounds to improve mutual information estimation. Assuming the knowledge of extract likelihood (joint or marginal), this paper shows the mutual information bound can be tightened via exploiting a series of intermediate distributions interpolating between a (convenient) proposal and the joint density. An energy-based perspective was presented and discussed, revealing interesting insights. The effectiveness of this proposal is demonstrated with generative modeling applications on real-world datasets.", "summary_of_the_review": "This is an interesting work and my major concerns from the last round of paper reviews have been adequately addressed. As such I recommend acceptance. ", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "l7k2X3jVNR", "reviewer": "Reviewer_YniD", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Strengths:\n* The authors present a comprehensive framework for mutual information bounds, called GIWAE, that unifies many existing approaches.\n* Combining AIS with IWAE is elegant, especially when one can define a sensible approximation for $p(z | x)$ in the form of either a prior $p(z)$ or approximate posterior $q(z | x)$ as with a VAE. The MINE-AIS extension is also intriguing due to its potential to extend the proposed techniques far beyond the VAE and GAN scenarios considered in this work.\n* The experimental results are quite impressive and convincingly demonstrate the utility of the proposed method.\n* The authors provide a comprehensive presentation of their method and an exhaustive review and comparison with other methods.\n\nWeaknesses:\n* The MINE-AIS technique is not fully explored. It would be interesting to see MINE-AIS applied in a scenario where the IWAE-type AIS bounds are not applicable (i.e. when $p( x | z )$ is unknown) but this is not investigated.\n* The core of the proposed method is previously suggested in a blog post by Sobolev. I believe this is not a significant weakness because the impressive effectiveness of this approach was not previously demonstrated.\n\nOther Comments:\n* I am unsure of why a single $z$ is used in the numerator but multiple $z$ are used in the denominator of (7) and other IWAE-type equations. A brief explanation would aid my understanding.", "summary_of_the_paper": "This work combines the technique of Annealed Importance Sampling (AIS) with Importance Weighted Autoencoder (IWAE) bounds on mutual information to obtain new and significantly tighter bounds. They examine the situations where the joint distribution $p(x, z)$ is known via $p(z)$ and $p(x | z)$, and the situation where only $p(z)$ is known. In the former case, they obtain exact bounds (or asymptotically exact given large enough sample sizes), and in the latter case they find approximate bounds by learning an energy function to approximate $p(x |z)$. Their experiments show the utility of their proposed method. In particular, they obtain impressively tight lower and upper bounds for complex latent variable generative models (VAE and GAN) on complex data such as CIFAR-10.", "summary_of_the_review": "Overall, this is an impressive work both in terms of its comprehensive mathematical presentation and its high-quality experimental results. Personally, I find the experimental results to be the most interesting aspect, because it opens the possibility of accurate MI estimation in complex scenarios studied in contemporary generative modeling. I recommend this work be accepted.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works."}, "extra_text": {}}, {"review_id": "0QW0fwUXeyg", "reviewer": "Reviewer_78Uf", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Strength: \n(1) This work unifies the view of mutual information from the perspective of importance sampling. This gives us another angle of thinking.\n(2) Good summary of the prior literature\n(3) Clear theoretical analysis.\n(4) The experiments provide evidence on the effectiveness of the proposals.\n\nWeakness:\n(1) This paper is more of a theoretical flavor and there is no experiment to show this method is practically useful. For example, is it useful to improve the prediction accuracy (using the learned latent representation z)?\n(2) How about the efficiency of the proposed strategies. (e.g. information about the running time). In the experiments, the author(s) have used tens of thousands of annealing steps, which may not be practically feasible. Also, it is helpful to visualize the trade-offs in computing time and bound sharpness.\n(3) I am curious about the stability of the training process. \n(4) The method has limited applicability in real-world settings, Generative modeling for feature engineering is not used that often, mostly because it is so hard to assume the knowledge of a likelihood model for a machine learning problem.\n\nMinor Comments:\n(1) Make clear distinction between empirical estimators and theoretical estimators. For example, I_{BA} is theoretical and I_{IWAE} is empirical. Please add hat over the empirical estimators.\n(2) In Figure 1, Z_{T-1} is bigger than other circles. Is there any different meanings?\n(3) Table 1 is messy, tidy it up. There too much data and at the first glance, no idea about where I should focus on. Also, there is inconsistency in notations: in the first column, T=30K, K=1K. Should change the capital letter K to lower case letter k.\n\n", "summary_of_the_paper": "In this work, the authors proposed some novel estimators for mutual information, namely the GIWAE, Muti-Sample AIS and MINE-AIS bounds. The key idea is to apply importance sampling to mutual information in the model-based setting. In both theoretical and empirical analysis, they prove the bounds are much tighter than the existing ones in the literature. ", "summary_of_the_review": "The reviewer is convinced by the theoretical analysis of this paper and the novel perspectives it proposed. \nIt makes some valuable contribution to our community. However, I am concerned about the practicability and efficiency of this method. ", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper investigates a tighter bound for mutual information and proposes some novel estimators of MI from the importance sampling perspective. The proposed approach provides a unifying framework for mutual information bounds that can deduce many existing approaches. The theoretical and experimental analyses well justify the proposed approach and shows the bounds are much tighter than the existing ones.  
Overall, this paper is well written. The relevant literatures are exhaustively reviewed and well compared with the proposed method. The experimental results show remarkable superiority of the proposed method. The proposed framework would shed light on the literatures and open up a new direction of the relevant researches.  
For those reasons, I would like to recommend this paper to be accepted by ICLR2022 conference.
Decision: Accept (Poster)