Double Control Variates for Gradient Estimation in Discrete Latent Variable Models
Abstract
Stochastic gradient-based optimisation for discrete latent variable models is challenging due to the high variance of gradients.
We introduce a variance reduction technique for score function estimators that makes use of double control variates .
These control variates act on top of a main control variate, and try to further reduce the variance of the overall estimator.
We develop a double control variate for the REINFORCE leave-one-out estimator using Taylor expansions.
For training discrete latent variable models, such as variational autoencoders with binary latent variables, our approach adds no extra computational cost compared to standard training with the REINFORCE leave-one-out estimator.
We apply our method to challenging high-dimensional toy examples and training variational autoencoders with binary latent variables.
We show that our estimator can have lower variance compared to other state-of-the-art estimators.
1. Background
Several problems in machine learning, such as variational inference and reinforcement learning, require the optimisation of an intractable expectation of an objective function f ( x ) under a distribution q η ( x ) with tunable parameters η .
Here f ( x ) is a diﬀerentiable objective function.
x is a D -dimensional vector.
Since f ( x ) can have a complex non-linear form, E q η ( x ) [ f ( x )] and its exact gradients are generally intractable.
Several techniques apply stochastic optimisation based on unbiased Monte Carlo gradients by sampling from q η ( x ).
Pathwise or reparametrization gradients (Glasserman, 2003) have been shown to be eﬀective for machine learning problems (Kingma and Welling, 2014; Rezende et al., 2014;
Titsias and L´azaro-Gredilla, 2014), but they are only applicable to continuous distributions.
A very general class of gradient estimators that apply to both continuous and discrete variables is the score function or REINFORCE estimator (Glynn, 1990; Williams, 1992;
Carbonetto et al., 2009; Paisley et al., 2012; Ranganath et al., 2014; Mnih and Gregor, 2014).
However, these estimators suﬀer from high variance and reducing the variance remains an important open problem.
Variance reduction techniques for REINFORCE estimators range from simple baselines (Ranganath et al., 2014; Mnih and Gregor, 2014) and Rao-blackwellization (Titsias and L´azaro-Gredilla, 2015; Tokui and Sato, 2017) to more advanced gradient-based control variates (Tucker et al., 2017; Grathwohl et al., 2018;
Gu et al., 2016) and coupled sampling (Yin and Zhou, 2019; Dong et al., 2020; Yin et al., 2020;
Dimitriev and Zhou, 2021).
The score function estimator with a baseline b is given by
[Equation]
[Figure] Another variance reduction method that has become prominent recently is the REINFORCE leave-one-out estimator (RLOO) (Salimans and Knowles, 2014; Kool et al., 2019;
Richter et al., 2020).
Given K ≥ 2 samples, it takes advantage of multiple evaluations of f to avoid learning the baseline b :
[Equation]
where each (cid:16) 1 K − 1 (cid:80) j (cid:54) =
k f ( x j ) (cid:17) ∇ η log q η ( x k ) acts as a sample-speciﬁc control variate.
Despite its simplicity, this estimator performs very strongly for training discrete latent variables models (Dong et al., 2020;
Richter et al., 2020).
Presumably this is because the leave-one-out stochastic baselines can automatically adapt to the non-stationarity of the f ( x ).
Speciﬁcally, f ( x ) := f θ ( x ) often contain additional model parameters θ updated at each optimization step 1 , as for instance in variational autoencoders (VAEs) (Kingma and Welling, 2014;
Rezende et al., 2014).
Although θ is changing, the sample-speciﬁc baseline
[Equation]
always remains an unbiased estimate of E q η ( x ) [ f θ ( x )].
However, RLOO is still limited in how much variance reduction it can achieve.
Proposition 1 Consider the estimator R ∗ ( η ) = 1 K (cid:80)
Kk =
( f ( x k ) − E f ) ∇ η log q η ( x k ) , where
E f = E q η ( x ) [ f ( x )] is a constant baseline across all samples.
Then, V ar ( RLOO ) ≥ V ar ( R ∗ ) .
According to Prop. 1, the performance of RLOO is bounded by R ∗ which uses the mean
E f (usually intractable in practice) as a constant baseline.
Therefore, there is scope to further reduce the variance of this estimator.
In this work, we focus on the RLOO estimator and enhance it by adding extra control variates.
We refer to the added baselines as double control variates since they co-exist with the main RLOO baseline, and are designed to have a complementary eﬀect by reducing the variance of the initial RLOO estimator.
For training latent variable models with discrete variables, our proposed estimator runs roughly at the same speed as the RLOO estimator.
2. Double Control Variates for REINFORCE LOO
We construct these new control variates along two directions:
(a) Since the main baseline
[Equation]
is stochastic and thus has variance, we can try to reduce the variance by adding a control variate for each stochastic term f ( x j ).
(b) We want to add a diﬀerent type of control variate that depends on x k which may have a complementary eﬀect to the main RLOO baseline.
In the remaining of Section 2 we use s ( x ) :=
∇ η log q η ( x ) to denote the score function.
To accomplish both (a) and (b) simultaneously we start with the unbiased estimator 1 K
[-Kk =
-][Equation]
+ [Equation]
− α E q [Equation]
( x ) [ b ( x ) s ( x )] , where we introduced a control variate b ( x k ), that depends on the current sample x k and has analytic global correction
E q η ( x ) [ b ( x ) s ( x )].
Then, to create a double control variate estimator we treat f ( x ) + αb ( x ) as the “new eﬀective objective function” and apply the leave-one-out procedure to it:
[Equation]
The scalar α is a regression coeﬃcient that can be further optimised to reduce the variance.
In the above estimator we have highlighted with blue the ﬁrst appearance b ( x k ), which can be thought of as a baseline paired with the value f ( x k ), and with red the second appearances b ( x j ) paired with the remaining values f ( x j ).
Intuitively, b ( x k ) can be considered as targeting to reduce the variance of f ( x k ) and b ( x j ) the variance of f ( x j ).
In Sections 2.1 and 2.2 we describe two approaches to specify b ( x ).
For training latent variable models such as VAEs, the second will be the most practical since it adds no extra cost.
The ﬁrst method helps to introduce the idea and is based on a mean ﬁeld argument.
2.1. Mean Field Approach
To specify b ( x ) we can construct an approximation of f ( x ) that correlates well with the exact value f ( x ).
While any surrogate of f ( x ) with a tractable global correction could work, next we focus on the case when f ( x ) is diﬀerentiable w.r.t.
the input x and we use a ﬁrst order
Taylor expansion around the mean µ = E q η ( x ) [ x ], so that f ( x ) ≈ f ( µ )+ ∇ f ( µ ) (cid:62) ( x − µ ).
Since any constant term in b ( x ) cancels out in (2), the constant f ( µ ) in the Taylor approximation can be dropped, yielding the double control variate b ( x ) = ∇ f ( µ ) (cid:62) ( x − µ ) .
By substituting this function in Eq.
(2) we obtain the general estimator
[Equation]
where E q η ( x ) [ s ( x )( x − µ ) (cid:62) ] will typically have an analytical form.
For binary latent variables x ∈ { 0 , 1 } d and a factorised Bernoulli distribution of the form q η ( x ) = (cid:81) di =1 µ x i i (1 − µ i )
− x i , µ i = σ ( η i ) .
E q η ( x ) [ s ( x )( x − µ ) (cid:62) ]
= diag( µ ◦ (1 − µ )) and the global correction term simpliﬁes to − αµ ◦ (1 − µ ) ◦ ∇ f ( µ ), where ◦ denotes element-wise vector product.
2.2. An Estimator without Extra Gradient Evaluations
The estimator in Eq.
(3) requires a backpropagation operation to compute the gradient ∇ f ( µ ), which adds extra computational cost compared to standard RLOO.
Next, we wish to develop an alternative estimator that avoids this extra cost for certain problems.
For many applications, such as VAEs, the function f ( x ) depends on model parameters θ (typically diﬀerent than η ) that we update at each optimisation iteration by computing the gradients {∇ θ f ( x j ) } Kj =1 .
Then, from the same backpropagation operations is easy to also return the gradients w.r.t.
the latent vectors, i.e. to compute {∇ f ( x j ) } Kj =1 .
We would like to utilize these latter gradients to deﬁne the double control variate b ( x ).
Starting from b ( x ) = ∇ f ( µ ) (cid:62) ( x − µ ), we ﬁrst want to modify b ( x k ) by replacing ∇ f ( µ ) with some new gradient computed from {∇ f ( x j ) } Kj =1 .
We cannot use the full average because this will lead to ( 1 K (cid:80)
Kj
=1 ∇ f ( x j )) (cid:62) ( x k − µ ) which has an intractable global correction due to the intractable term E q η ( x k ) [ ∇ f ( x k ) (cid:62) ( x k − µ ) ∇ η log q η ( x k )].
However, we can use the leave-one-out gradient, i.e. by leaving out ∇ f ( x k ), which gives
[Equation]
This has a tractable correction term E q η ( x k ) [ b k ( x 1: K ) ∇ log q η ( x k )].
Proposition 2 For b k ( x 1: K ) from (4) we obtain the following unbiased gradient estimator
[Equation]
The proof of unbiasedness is given in the appendix.
We choose the regression coeﬃcient α by minimizing the total variance.
If g ( α ) denotes the stochastic gradient and ¯ g =
E [ g ( α )] the exact gradient where the latter does not depend on α , the total variance is
Tr[ E ( g ( α ) − ¯ g )( g ( α ) − ¯ g ) (cid:62) ]
=
E [ || g ( α ) || 2 ] + const .
Thus, in practice at each optimisation iteration we can perform a gradient step towards minimizing the empirical variance
|| g
( α )
|| 2 .
3. Experiments 3.1. Toy Learning Problem
We consider a generalization of the artiﬁcial problem considered by Tucker et al.
(2017), where the goal is to maximize E ( η ) =
E q η ( x ) [ D − 1 (cid:80) Di
=1 ( x i − p 0 ) 2 ] , where q η ( x ) = (cid:81) Di =1 σ ( η i ) x
i (1 − [Table] σ ( η i ))
1 − x i , p 0 = 0 .
499 and the optimal solution is σ ( η i )
=
for all i = 1 , . . .
, D .
While
Tucker et al.
(2017) considered D = 1, here we additionally consider a more diﬃcult highdimensional case with D = 200.
We compare three methods: (i) RLOO, (ii) DisARM and our proposed double control variates estimator (Double CV) from Eq.
We use K =samples for all methods.
Also we include in the comparison R ∗ which is tractable in this toy example.
Fig.
1 compares the methods in terms of variance, the objective function and the average value of the D probabilities σ ( η i ).
Fig.
6 shows further comparison for the D =case, as in Tucker et al.
We observe that Double CV gradients have smaller variance which results in much faster optimisation convergence.
3.2. Variational Autoencoders with Binary Latent Variables 3.2.1. Experimental setup
We consider training nonlinear variational autoencoders (Kingma and Welling, 2014; Rezende et al., 2014) with binary latent variables.
We conduct separate experiments for binary output data y ∈ { 0 , 1 } d and continuous data y ∈ R d .
For binary data we use the standard Bernoulli likelihood.
For continuous data we centered data between [ − 1 , 1] and consider a Gaussian likelihood of the form p θ ( y | x ) =
N ( y | m θ ( x ) , Σ), where m θ ( x ) is a decoder mean function that depends on the latent variable x and Σ is a learnable diagonal covariance matrix.
We consider the datasets MNIST, Fashion-MNIST and Omniglot.
For all three datasets we use both the dynamically binarized versions and their original continuous versions.
More details and the results for linear VAEs are included in the appendix.
We compared the following estimators: RLOO, DisARM and the proposed Double CV method where all three estimators use K samples.
We experimented with K = 2 and K = 4.
For K = 4 we also compare to the state-of-the-art ARMS estimator recently proposed by
Dimitriev and Zhou (2021).
Besides, we include in the comparison the RELAX estimator that combines concrete relaxation (Tucker et al., 2017) with a learned control variate (Grathwohl et al., 2018).
We point out that RLOO, DisARM, Double CV, and ARMS (when K = 4) have roughly the same running time on a P100 GPU while RELAX is computationally more expensive and is roughly twice slower than the other four estimators with K = 4 (see
Table 3).
Also note that RELAX is less generally applicable since it assumes the existence of a concrete relaxation for x . [Table] [Figure] G r ad i en t V a r i
an c e
[Figure]
3.2.2. Results
Table 1 shows the training ELBO for binarized and continuous datasets when training the
VAE by diﬀerent estimators with K = 2.
We can observe that Double CV consistently outperforms RLOO in all experiments, while having approximately the same running time.
Double CV also outperforms DisARM in all cases for both Bernoulli and Gaussian likelihoods.
Furthermore, Fig.
3 plots the gradient variance and the training ELBO for the binarized datasets as a function of the training steps.
Similarly, Fig. 4 shows the corresponding results for the non-binarized (continuous) datasets where a Gaussian likelihood is used.
We observe that the Double CV estimator can have lower variance than RLOO and DisARM.
Also, while RELAX performs better than the other methods it is less generally applicable and more expensive.
For K = 4, the ﬁnal training ELBO values are reported in Table 2 and the variances of the diﬀerent estimators are plotted in Fig.
We can observe that Double CV consistently has lower variance than other estimators and it outperforms ARMS in terms of training
ELBO in most cases.
It also signiﬁcantly outperforms RELAX.
Note that, even with K = 4,
Double CV is still nearly twice faster than RELAX.
Appendix A. Related Work
Our proposed gradient estimators follow the general form of unbiased REINFORCE estimators (Williams, 1992; Glynn, 1990; Carbonetto et al., 2009;
Paisley et al., 2012; Ranganath et al., 2014; Mnih and Gregor, 2014), which unlike reparametrization or pathwise gradients (Kingma and Welling, 2014;
Rezende et al., 2014; Titsias and L´azaro-Gredilla, 2014), are applicable also to discrete latent variables.
The double control variates we develop build on top of the RLOO estimator (Kool et al., 2019; Salimans and Knowles, 2014; Richter et al., 2020); see also the VIMCO method of Mnih and Rezende (2016) who also used a leave-one-out procedure.
RLOO was shown to be a competitive estimator for challenging problems such as training VAEs with binary or categorical latent variables (Dong et al., 2020;
Richter et al., 2020; Dong et al., 2021).
As shown by our experiments, our enhancement of RLOO with double control variates leads to further variance reduction, and without increasing the computational cost when training VAEs.
In our current framework, the double control variates are constructed by using the gradients of the objective function f θ ( x ).
These gradients are also used by other unbiased gradient techniques based on control variates, such as the MuProp estimator (Gu et al., 2016), the concrete relaxation methods REBAR (Tucker et al., 2017) and RELAX (Grathwohl et al., 2018).
Our method diﬀers signiﬁcantly since our control variates act on top of the sample-speciﬁc RLOO baseline 1 K − 1 (cid:80) j (cid:54) =
k f θ ( x j ), i.e., they try to have complementary eﬀect to this existing control variate.
This means that our estimators preserve RLOO’s property of capturing the non-stationarity of f θ ( x ), since the leave-one-out baseline always tracks the expected value E [ f θ ( x )] as θ evolves.
In contrast, previous gradient-based estimators use stand-alone global control variates.
For instance, the baseline in MuProp (Gu et al., 2016) is constructed using only f θ ( µ ) and x k , which can be a poor tracker of the expected value
E [ f θ ( x )].
Unlike MuProp, REBAR (Tucker et al., 2017) and RELAX (Grathwohl et al., 2018) are much more eﬀective, however they are more expensive than our method — they require diﬀerentiating f θ three times, while our method can work with just two, and they are less generally applicable since they assume a concrete relaxation for x .
Other recent REINFORCE type of estimators for discrete latent variables are based on coupled sampling (Owen, 2013), such as antithetic sampling (Yin and Zhou, 2019;
Dong et al., 2020; Yin et al., 2020; Dimitriev and Zhou, 2021).
For instance, the recent DisARM estimator independently proposed by Dong et al.
and Yin et al.
(2020) was shown to give state-of-the-art results for binary latent-variable models with K = 2 antithetic samples.
Appendix B. Proofs B.1. Proof of Proposition
The RLOO estimator can be written as
[Equation]
(6) where R ∗ is the REINFORCE estimator with baseline E f and E is a residual term of zero mean.
To prove the Proposition we will use V ar (RLOO)
=
V ar (R ∗ + E ) =
V ar (R ∗ )
+
V ar ( E ) + 2 Cov
(R ∗ , E ).
Then, it suﬃces to show that Cov (R ∗ , E ) = 0.
We have
[Equation]
where we used f − k
[Equation]
=
[Equation] [Equation]
(cid:54) =
k (cid:48) f
( x j ) for short.
For all terms in the double sum such that k = k
(cid:48)
the expectation
[Equation]
because the zero-mean random variable E f − f − k is independent from the remaining product (since it does not contain the sample x k ).
For all cross terms k (cid:54) =
k (cid:48) the whole product ( f ( x k ) − E f )( E f − f − k (cid:48) ) ∇ η log q η ( x k ) does not contain the sample x k (cid:48) .
Therefore this product is independent from ∇ η log q η ( x k (cid:48) ) and thus each cross term is zero because of the score function property E [ ∇ η log q η ( x k (cid:48) )]
= 0.
This shows that
Cov (R ∗ , E ) = 0 which completes the proof.
B.2. Proof of Prop.
The estimator can be written as
[Equation]
where b k ( x 1: K ) = (cid:16) 1 K − 1 (cid:80) j
(cid:54) =
k ∇ f ( x j ) (cid:17) (cid:62) ( x k − µ ) and b j ( x 1: K ) = (cid:16) 1 K − 1 (cid:80) m (cid:54) =
j ∇ f ( x m ) (cid:17) (cid:62) ( x j − µ ).
It suﬃces to show that the expectation of the second line is minus the correction term at the third line.
The expectation of each term b j ( x 1: K ) ∇ η log q η ( x k ) for j (cid:54) =
k is zero because the zero-mean term x j − µ is always independent from the rest terms in the product.
Then, we need to examine only the expectation of
[Equation]
Then observe that the expectation of ∇ η log q η ( x k ) × ( x k − µ ) (cid:62) is the same for every sample x k , so the above reduces to
[Equation]
µ ) (cid:62) ]
K ( K − 1)
K (cid:88) k
= (cid:88) j (cid:54) =
k ∇ f ( x j ) from which the result follows since (cid:80)
Kk =1 (cid:80)
j (cid:54) =
k ∇ f ( x j ) =
( K − 1) (cid:80)
Kk =1 ∇ f ( x k ).
B.3. The “half” Double Control Variate Estimators
One question is whether we need both b ( x k ) and b ( x j ) or we could keep one of them, i.e. to use an “ b ( x k ) only” or “ b ( x j ) only” estimator.
It is straightforward to express these latter unbiased estimators, as follows.
The “ b ( x k ) only” estimator is given by
[Equation]
and the “ b ( x j ) only” by
[Equation]
It is easy to show that both estimators are unbiased.
However, in practice these estimators can be much less eﬀective in terms of variance reduction than their Double CV combination.
In Figure 5 we apply these two estimators to the toy learning problem with D = 10.
Both estimators are signiﬁcantly outperformed by the full Double CV estimator.
Notably, the “ b ( x k ) only” estimator could outperform R ∗ since it uses a baseline that depends on the current sample x k , while “ b ( x j ) only” reduces the variance of the RLOO control variate but remains bounded by R ∗ .
[Figure]
Appendix C. Additional Results and Experimental Details C.1. Toy Experiment with D =
For completeness, we include the results of a simpler version of the toy experiment described in Section 3.1, where we set D = 1.
This is the setting used in several previous works (Tucker et al., 2017; Grathwohl et al., 2018;
Yin and Zhou, 2019; Dong et al., 2020).
The variances of the gradient estimators and the training curves of σ ( η ) are plotted in Fig.
[Figure] [Figure]
[Equation]
[Figure]
[Equation]
[Figure]
C.2. Training Binary Latent VAEs C.2.1. Experimental Details
We follow the VAE models used in Yin and Zhou (2019); Dong et al.
The VAE model uses fully connected neural networks with two hidden layers of 200 LeakyReLU activation units with the coeﬃcient 0 .
All models are trained using Adam (Kingma and
Ba, 2014) with learning rate 10 − 3 for the binarized data, while for the continuous data we used smaller learning rate 10 − 4 .
In all experiments α was trained with learning rate 10 − 3 .
For all experiments we use a uniform factorized Bernoulli prior over the D =dimensional latent variable x .
The model was trained by maximizing the ELBO using an amortised factorised variational Bernoulli distribution.
C.2.2. Time comparison
In Fig.
We report the per-step running time of RLOO, Double CV, DisARM, ARMS estimators when K = 4 and compare to RELAX.
RELAX is almost twice slower.
[Table]
C.2.3. Full results of training ELBOs
Here we include the full results of ﬁnal training ELBOs from the experiment in Section 3.2.
Table 4 and Table 5 extend Table 1 to include the linear VAE results trained under the same setting.
Table 6 and Table 7 extend Table 2 to include the linear VAE results trained under the same setting.
The linear VAE has 200 dimensional latent variable x and use a single fully-connected layer to produce the logits (for Bernoulli likelihoods) or the mean (for
Gaussian likelihoods) of the distribution of y .
C.2.4. Additional figures for nonlinear VAEs
In Fig.
8 we plot the average training ELBOs as a function of training steps from the K =experiment in Section 3.2.
C.2.5. Additional figures for linear VAEs
We plot the gradient variance and average training ELBOs of training linear VAEs in
Figures 9,10,11, and 12.
[Table] [Table]
[Table] [Table]
[Figure] [Figure]
[Figure] [Figure]
ad
[Figure]
en t V a r i
an c e
[Figure] [Figure]
t V a r
[Figure]
an c
[Figure]
r a i n i
[Figure]
L
[Figure]
O
[Figure]