U-Statistics for Importance-Weighted Variational Inference
Abstract
We propose the use of U-statistics to estimate the (gradients of) the importance-weighted evidence lower bound (IW-ELBO), a variational objective that uses multiple samples from a proposal distribution to lower-bound the log-likelihood.
We propose a complete U-statistic estimator, which has variance that is never higher than the standard IW-ELBO estimator, and, under certain conditions, the lowest variance of any unbiased estimator.
However, it requires evaluating the objective on a large number of subsets of samples from the proposal distribution, which can be computationally expensive.
We propose to use incomplete Ustatistics as practical alternatives.
We ﬁnd empirically that both methods reduce estimator variance for its gradients with little computational cost, and lead to faster optimization.
1. INTRODUCTION
An important recent development in variational inference (VI) is the use of ideas from
Monte Carlo sampling to obtain tighter variational bounds (Burda et al., 2016; Maddison et al., 2017; Le et al., 2018; Naesseth et al., 2018; Domke and Sheldon, 2019).
These lead to better approximations of the log-likelihood for learning (Burda et al., 2016) and better approximate posterior distributions (Cremer et al., 2017; Domke and Sheldon, 2018;
Naesseth et al., 2018; Domke and Sheldon, 2019).
Assume a target distribution p ( z, x ) =
p ( z ) p ( x | z ) where x is observed and z is latent.
VI uses the following evidence lower bound (ELBO), given approximating distribution q , to approximate ln p ( x ) (Saul et al., 1996; Blei et al., 2017):
L = E Z ∼ q (cid:104)
[Equation]
Burda et al.
ﬁrst showed that a tighter bound can be obtained by using the average of m importance weights within the logarithm.
The importance-weighted ELBO ( IW-ELBO ) is L m =
E Z 1: m (cid:104)
ln 1 m (cid:80) mi =1 p ( Z i ,x ) q ( Z i ) (cid:105) ≤
ln p ( x ), where the expectation is over Z 1 , . .
, Z m drawn independently from q .
We expect Jensen’s inequality to provide a tighter bound because the distribution of this sample average is more concentrated around p ( x ) than the distribution of one estimate.
Indeed, L m ≥ L m (cid:48) for m > m (cid:48) and approaches ln p ( x ) as m → ∞ (Burda et al., 2016).
In practice, the IW-ELBO is estimated by sampling.
It is convenient to deﬁne the logweight random variables W i =
ln p ( Z i ) − ln q ( Z i ) for Z
i ∼ q and rewrite the IW-ELBO as
[Figure]
sonar
[Figure]
c o m p l e t e - U r ando m s ub s .
pe r m u t ed 0. [Figure] r a t i o [Figure]
of gradients’ covari
[Figure]
ance.
L m =
E W 1: m (cid:104)
ln 1 m (cid:80) mi =
e W i (cid:105) .
Then, for example, with m = 2, we can draw two log-weights
W 1 and W 2 and estimate the IW-ELBO as ˆ L 2 = ln
( e W 1 + e W 2 ) − ln 2.
It is standard to use multiple replicates to reduce variance (Rainforth et al., 2018; Domke and
Sheldon, 2018).
For example, we could draw four log-weights instead of two and estimate the IW-ELBO as ˆ L 4 , 2 = 12 ln( e W 1 + e W 2 ) + 12 ln( e W 3 + e W 4 ) − ln 2.
[-the IW-ELBO as ˆ L 4 , 2 = 12 ln( e W 
-]
However, using the same n log-weights, we could instead compute
[Equation]
which uses all possible pairs of log-weights instead of two disjoint pairs.
This is an instance of a complete U-statistic, a family of estimators introduce by Hoeﬀding (1948) following the work of Halmos (1946).
It is clear that both estimators are unbiased estimators of L 2 .
Additionally, it was proven by Halmos that the complete U-statistics (and hence ˆ L U 4 , 2 ) are optimal in the sense of having the smallest variance among all unbiased estimators (of L 2 in our case).
These two properties suggest that the complete U-statistics are good candidates for the task of optimizing the IW-ELBO, thus motivating our work.
The contributions of this paper are: i) we propose the complete U-statistic IW-ELBO estimator, which is the unbiased estimator of L m with the smallest variance
[Section 2], and describe how to construct an analogous estimator for gradient estimation that is compatible with diﬀerent base gradient estimators [Section 2.1]; ii) to reduce the computational burden of evaluating the objective on (cid:0) nm (cid:1) sets for the complete U-statistic, we construct particular incomplete U-statistics
(Blom, 1976) and prove they achieve most of the variance reduction of complete U-statstics at a modest cost [Section 3]; iii) we show empirically that our approaches consistently reduce the trace of the gradients’ covariance compared to the standard IW-ELBO estimator [Section 4], which leads to faster learning [Fig. 1( a )].
2. U-STATISTIC ESTIMATORS
We now formalize the examples above by introducing the standard IW-ELBO estimator and its corresponding complete U-statistic IW-ELBO estimator , and applying the theory of
U-statistics to relate their variances.
The theory of U-statistics was developed in a seminal work by Hoeﬀding (1948) and extends the theory of unbiased estimation introduced by
Halmos (1946).
Detailed background can be found in the original works, or see (Lee, 1990; van der Vaart, 2000).
Our goal is to estimate the IW-ELBO L m and its gradient with respect to parameters of the log-weight distribution.
We focus here on IW-ELBO estimation, but an analogous development holds for reparameterization gradients (see Section 2.1).
We consider estimators that are functions of n independent log-weights W =
( W 1 , . . .
, W n ), where we assume for convenience that n = rm for an integer number of replicates r ≥ 1.
Estimator
The standard IW-ELBO estimator is
[Equation]
This estimator partitions the n log-weights into r disjoint sets and computes the average of the unbiased estimates from each set.
In contrast, estimators based on U-statistics will use overlapping sets of log-weights.
To set this up, we introduce notation for
size- m
subsets of the indices from 1 to n .
Let (cid:74) n (cid:75) =
{ 1 , . . . , n } ⊆ N , let (cid:0) Am (cid:1) denote the set of all subsets of
A with exactly m elements, and, for s ∈ (cid:0) (cid:74) n (cid:75) m (cid:1) , let s i be the i th smallest index in s .
Estimator
The complete U-statistic IW-ELBO estimator is
[Equation]
In words, ˆ L Un,m takes the average of ln (cid:0) 1 m (cid:80) mi =1 e W si (cid:1) over all distinct subsets s of m indices.
Contrast this to ˆ L n,m , which takes the same average but over r disjoint subsets of m indices.
It follows from linearity of expectation that both ˆ L U n,m and ˆ L n,m are unbiased estimators for L m .
The relationship to general U-statistics studied by Halmos (1946) and Hoeﬀding (1948) is as follows.
Given n random variables W 1 , . . .
, W n drawn iid from a distribution F , and an unbiased estimator of the form θ ( F ) =
E h ( W 1 , . . .
, W m ) for a kernel h that is symmetric in its m ≤ n arguments, the complete U-statistic is (cid:0) nm (cid:1) − 1 (cid:80) s ∈ ( (cid:74) n (cid:75) m ) h
( W s 1 , . . .
, W s m ).
Our estimator is a complete U-statistic for h ( W 1 , . . .
, W m )
=
ln (cid:0) 1 m (cid:80) mi =1 e W i (cid:1) .
A result of Halmos (1946) implies that, subject to certain conditions, ˆ L Un,m has the smallest variance of any unbiased estimator of the IW-ELBO.
The technical conditions are needed to deﬁne the class of “unbiased estimators” as ones that are unbiased for all log-weight distributions in a non-trivial class.
We observe in practice that there is a gap between the two variances that leads to practical gains for the complete U-statistic estimator in real VI problems.
2.1. Gradients
Variance reduction via U-statistics also applies to gradient estimators.
For example, assume the log-weights are reparameterizable as W i =
W ( (cid:15)
i , θ ) for (cid:15) i drawn iid from a ﬁxed base distribution.
By linearity, the reparameterization partial derivative estimator ∂∂θ
j ˆ L Un,m (cid:0) W ( (cid:15) 1 , θ ) , . .
, W ( (cid:15) n , θ ) (cid:1) is itself a complete U-statistic corresponding to the kernel h j ( (cid:15) 1 , . . .
, (cid:15) m ; θ ) =
∂∂θ j ln 1 m (cid:80) mi =1 e W ( (cid:15)
i ,θ ) , which is “base” reparameterization gradient estimator.
More broadly, we can form a complete U-statistic gradient estimator (cid:0) nm (cid:1)
[Equation]
( (cid:15)
[Equation]
, (cid:15) s m ; θ ) from any base gradient estimator, for example, the doubly reparameterized gradient estimator h j ( (cid:15) 1 , . .
, (cid:15) m ; θ ) =
[Equation]
(Tucker et al., 2018), or the score function estimator.
Because the theory of this section derived from general U-statistic theory, for each base estimator, it holds an exact analog of Prop 9 in the Appendix, this time using h j .
The complete U-statistic estimator will have variance that is never higher, and usually lower, than the base estimator averaged over r disjoint samples.
As a concrete result, this implies that E (cid:107)∇ θ ˆ L Un,m (cid:107) 22 ≤ E (cid:107)∇ θ ˆ L n,m (cid:107) 22 .
The same situation will apply in the next section: results about variance of U-statistic estimators for the IW-ELBO also apply to U-statistic estimators of gradients.
Analogs of Proposition 7 also hold for gradient estimation, but should be taken lightly.
Gradient estimators often use properties of the underlying sampling distribution of the log-weights—e.g., that the distribution can be reparameterized, or that log p ( Z ; θ ) is differentiable with respect to θ (for the score function estimator)—that make the conditions unlikely to hold.
Indeed, it is precisely by using these properties that diﬀerent approaches are often able to reduce variance.
2.2. Computational Complexity
There are two main factors to consider for the computational complexity of an IW-ELBO estimator:
The cost to compute n log-weights W i =
ln p
( Z i , x ) − ln q ( Z i ) for i ∈ (cid:74) n (cid:75) , and 2) the cost to compute the estimator given the log-weights.
A problem with ˆ L Un,m is that it averages (cid:0) nm (cid:1) distinct subsets of indices in 2), which is expensive.
It should be noted that the individual operations to compute ˆ L Un,m are very simple, while, for many probabilistic models, computing each log-weight is expensive, so, for modest m and n , the computation may still be dominated by Step 1).
However, for large enough m and n , Step 2) is impractical.
3. INCOMPLETE U-STATISTIC ESTIMATORS
In practice, we can achieve most of the variance reduction of the complete U-statistic IWELBO estimator with only modest computational cost by averaging over only k (cid:28) (cid:0) nm (cid:1) subsets of indices selected in some way.
Such an estimator is called an incomplete U-statistic .
Incomplete U-statistics were introduced and studied by Blom (1976).
Note that the standard
IW-ELBO estimator ˆ L n,m is itself an incomplete U-statistic, where the k = r = nm index sets are disjoint.
We can improve on this by selecting k > r sets.
Given S a collection of size- m subsets of (cid:74) n (cid:75) , a general incomplete U-statistic for the IWELBO is deﬁned as ˆ L S ( W ) = 1 |S| (cid:80) s ∈S ln (cid:16)
1 m (cid:80) mi =1 e W si (cid:17) .
We will allow S to be a multi-set, that is, the same subset may appear more than once.
Estimator 3 (Random subsets)
The random-subset incomplete-U-statistic estimator for the IW-ELBO is the estimator ˆ L S k where S k is a set of k subsets ( s i ) ki
=1 drawn uniformly at random (with replacement) from (cid:0) (cid:74) n (cid:75) m (cid:1) .
Estimator 4 (Permuted block)
The permuted-block incomplete-U-statistic estimator for the IW-ELBO is the estimator ˆ L S (cid:96) Π with the collection S (cid:96) Π deﬁned as follows.
Let π denote a permutation of (cid:74) n (cid:75) .
Deﬁne S π as the collection obtained by permuting indices according to π and then dividing them into r disjoint sets of size m .
That is,
[Equation]
Now, let S (cid:96)
Π
= (cid:85) π ∈ Π S π where Π is a collection of (cid:96) random permutations and (cid:85) denotes union as a multiset.
The total number of sets in S (cid:96) Π is k = r(cid:96) .
Both incomplete-U-statistic estimators can achieve variance reduction in practice for a large enough number of sets k , but the permuted block estimator has an advantage: its variance with k subsets is never more than that of the random subset estimator with k subsets, and never more than the variance of the standard IW-ELBO estimator (and usually smaller).
On the other hand, the variance of the random subset estimator is more than that of the standard estimator unless k ≥ k 0 for some threshold k 0 > r .
Proposition 5 Given m and n = rm , the variances of the estimators satisfy the following partial ordering:
[Equation]
Moreover, if (cid:96) > 1 and Var[ ˆ L Un,m ] <
Var[ ˆ L n,m ] , then (a) is strict; if r > 1 , then (b) is strict.
(Note that ˆ L S (cid:96) Π and ˆ L S r(cid:96)
both use k = r(cid:96) subsets.)
A remarkable property of the permuted-block estimator is that we can choose the number of permutations (cid:96) based on how much of the available variance reduction we want to achieve.
Say we would like to achieve 95% of the variance reduction; then it suﬃces to set (cid:96) = 20.
The following Proposition formalizes this result.
Proposition 6 Given m and n = rm , for (cid:96) ∈ N the permuted-block estimator achieves a (1 − 1 /(cid:96) ) fraction of the variance reduction provided by the complete U-statistic IW-ELBO estimator, i.e.,
[Equation]
4. EXPERIMENTS
In this section we analyze empirically the methods proposed in this paper for IW-VI (i.e., optimizing the IW-ELBO with respect to the parameters of the proposal distribution q ).
The general setup is shared across experiments: we experiment with Bayesian logistic regression using a diagonal Gaussian prior and a posterior approximated with a Gaussian distribution with full covariance matrix.
The models were optimized using stochastic gradient descent with ﬁxed learning rate.
We will set n = 16, except for the analysis of the required time [see below].
We ran the experiments with m ∈ [2 , 4 , 8] and we report the results for m = 4.
We set (cid:96) =
20 for the permuted block estimator, and for the random subsets estimator, we put k = 20 nm .
We show the results on the a1a ( d = 120), ionosphere ( d = 34) and sonar ( d = 60) datasets, but additional ﬁgures can be found in the Appendix.
Ultimately, our goal is to provide a more eﬃcient optimization method.
To empirically show this, we optimized each alternative over a range of diﬀerent learning rates.
In Fig.
1( a ) we show the envelope of the objective for all methods, i.e., at every iteration and method, we take the maximum value of the objective across all learning rates.
The Figure shows that the presented methods outperform the standard IW-ELBO estimator.
The reason is that those methods allow for higher learning rates.
We conjecture that this phenomenon is due to the total variance of the gradients, i.e., the trace of the covariance matrix of the gradients.
We estimated this quantity for each estimator taking 200 samples every 200 iterations.
We present in Fig.
1( b ) the ratio to the estimation of the total variance of the gradients for the standard IW-ELBO estimator using the a1a dataset.
The ratio is approximately 60% for all methods, with the random subsets estimator showing the highest ratio, and the complete U-statistic the lowest.
Moreover, the permuted estimator with (cid:96) = 20 achieves approximately 95 .
58% of the variance reduction, as predicted by Proposition 6.
To complete, we provide in Table 1 the times required to complete 1 k iterations of the optimization.
Here we used n = 24 and m = 12, which makes it a challenging setting for the complete U-statistic IW-ELBO estimator because there are (cid:0) 2412 (cid:1)
= 2 , 704 , 156 sets.
As expected, the only outlier method is the complete U-statistic.
The diﬀerences in the required times by the other methods to the time required by the standard IW-ELBO estimator are not statistically signiﬁcant and can be attributed to noise in the environment. [Table] 5. Citations and Bibliography Acknowledgments
Acknowledgements JB would like to thank Tom´as Geﬀner for earlier discussions about this project.
This material is based upon work supported by the National Science Foundation under
Grant No. 1908577.
Appendix A. Additional results
Proposition 7 Let E F [ · ] and Var F [ · ] denote expectation and variance with respect to logweights W 1 , . . .
, W n drawn independently from distribution F , and denote the IW-ELBO with log-weight distribution F by L m ( F ) =
E F (cid:2)
ln 1 m (cid:0)(cid:80) mi =
e W i (cid:1)(cid:3) .
Let F denote the set of distributions supported on a ﬁnite subset of R .
Suppose Φ is any estimator such that
E F [Φ( W 1: n )]
=
L m ( F ) for all F ∈ F .
Then,
[Equation]
whenever the latter quantity is deﬁned, for any distribution F on the real numbers (up to conditions of measurability and integrability).
Proof The result is a direct application of Theorem 5 of Halmos (1946).
For IW-ELBO estimation, the conditions are rather mild: we expect an IW-ELBO estimator to work for generic log-weight distributions.
For gradient estimation, we take the conclusion lightly, because gradient estimators often use speciﬁc properties of the underlying distributions; see Section 2.1.
A.1. Variance Comparison
How much variance reduction is possible with the complete U-statistic IW-ELBO estimator?
This section shows the variance of the ˆ L Un,m is never more than that of ˆ L n,m , and is strictly less under certain conditions (that occur in practice), using classical bounds on U-statistic variance due to Hoeﬀding (1948).
Since ˆ L Un,m is a sum of terms, one for each s ∈ (cid:0) (cid:74) n (cid:75) m (cid:1) , its variance depends on the covariances between pairs of terms for index sets s and s (cid:48) , which in turn depends on how many indices are shared by s and s (cid:48) .
This motivates the following deﬁnition:
Deﬁnition 8 Let W 1 , . . .
, W 2 m be i.i.d.
log-weights.
For 0 ≤ c ≤ m , take s , s (cid:48) ∈ (cid:0) (cid:74) n (cid:75) m (cid:1) with | s ∩ s (cid:48)
| = c .
Deﬁne
[Equation]
, which does not depend on the particular choices of s and s (cid:48) .
For example, when m =
we have ζ 0 = 0,
[Equation]
and
[Equation]
Then, due to Hoeﬀding’s classical result,
Proposition 9 Let ˆ L n,m and ˆ L Un,m be as in Estimators 1 and 2 for n = rm with r ∈ N .
Then
[Equation]
Moreover, for a ﬁxed m , the quantity n Var[ ˆ L Un,m ] tends to its lower bound m 2 ζ 1 as n increases.
Proof The inequalities and asymptotic statement follow directly from Theorem 5.2 of Hoeﬀding (1948).
The equality follows from the deﬁnition of ζ m .
We can now proceed to prove Proposition 5 and 6 from the paper.
Proposition 5 Given m and n = rm , the variances of the estimators satisfy the following partial ordering:
[Equation]
Moreover, if (cid:96) > 1 and Var[ ˆ L Un,m ] <
Var[ ˆ L n,m ] , then (a) is strict; if r > 1 , then (b) is strict.
(Note that ˆ L S (cid:96) Π and ˆ L S r(cid:96)
both use k = r(cid:96) subsets.)
Proof The complete U-statistic is at the lhs of (1) because all the estimators are unbiased, and ˆ L Un,m is the one with smallest variance
[cf.
Prop. 7].
By Def. 8, if s and s (cid:48) are uniformly drawn from (cid:0) (cid:74) n (cid:75) m (cid:1) , we have
[Equation]
Observe now that, from the deﬁnition of ˆ L S (cid:96) Π and for a given permutation π , all sets in
S π will be independent.
Hence, all dependencies between diﬀerent sets are due to relations between permutations, i.e., each of the (cid:96)r terms will have a dependency with the ( (cid:96) − 1) r terms not in the same permutation.
Therefore, it follows from (2) that the total variance of ˆ L S (cid:96) Π is
[Equation]
i.e., a convex combination of 1 r ζ m =
Var[ ˆ L n,m ] and Var[ ˆ L Un,m ].
Hence, (a) holds.
By a similar argument, the total variance of ˆ L S r(cid:96) is
[Equation]
Then, (b) holds because
[Equation]
Proposition 6 Given m and n = rm , for (cid:96) ∈ N the permuted-block estimator achieves a (1 − 1 /(cid:96) ) fraction of the variance reduction provided by the complete U-statistic IW-ELBO estimator, i.e.,
[Equation]
Proof Indeed, this follows directly from Eq.
Appendix B. Experiments B.1. Variance of Objective
The use of U-statistics was motivated because they provide a framework to potentially reduce the variance of the objective (and gradients) within IW-VI, and we conﬁrm this empirically.
We performed IW-VI using the complete U-statistic ˆ L Un,m , and every 200 iterations we estimated the variance of the objective for the alternatives presented: the standard IWElBO estimator, the complete U-statistic IW-ELBO, the permuted-block estimator with (cid:96) = 20, and the random subsets estimator with k = 20 nm , i.e., a number of sets equal to the permuted version.
Recall that, according to Prop. 6, (cid:96) =
20 implies that the permuted estimator achieves a 95% of the variance reduction provided by the complete-U-statistic
IW-ELBO.
In Figure 2, we show the ratio of the variance of the alternatives to that of the standard IW-ELBO estimator for the a1a dataset.
This Figure conﬁrms that it is possible to reduce the variance of the objective by considering the U-statistics.
Moreover, the estimators can be ordered by their variances, with the complete U-statistic IW-ELBO estimator showing the smallest variance, the permuted-block estimator in a middle ground, and ﬁnally the random subsets estimator being the one with the smallest variance reduction.
[Figure]
B.2. Additional Figures
In this section we augment the Figures shown in the paper with two additional datasets, australian and mushrooms , and for m ∈ [2 , 4 , 8].
In all cases we used n = 16 independently sampled log-weights W 1 , . . .
, W 16 .
We trained all models for 10 k iterations, but we show results for all iterations only for the largest models, i.e., a1a ( d = 120) and mushrooms ( d = 112), and the rest, i.e, australian ( d = 14), ionosphere ( d = 34) and sonar ( d = 60), we show the results for the ﬁrst 1 k iterations.
Additionally, we reported in Fig.
4 the objective achieved by each method after 2k and 10k iterations as a function of the learning rate [cf.
Domke and Sheldon (2018)].
In each case, we evaluated the learned approximating distribution using the standard IW-ELBO estimator with n (cid:48)
= 1000 ∗ m .
[Figure] [Figure]
i
[Figure]