Coupled Gradient Estimators for Discrete Latent Variables
1. Introduction
Training models with discrete latent variables is challenging due to the diﬃculty of estimating the gradients.
The gradient can be expressed as an expectation, however, in all but the simplest settings, it is analytically intractable and we must resort to estimating the gradient using Monte Carlo sampling.
This problem is encountered, for example, in modern variational inference, where we would like to maximize a variational lower bound with respect to the parameters of the variational posterior.
The pathwise gradient estimator , also known as the reparameterization trick, has low variance and has been instrumental to the success of variational autoencoders (Kingma and Welling, 2014; Rezende et al., 2014).
Unfortunately, it can only be used with continuous random variables, and ﬁnding a similarly eﬀective estimator for discrete random variables remains an important open problem.
Score-function estimators (Glynn, 1990; Fu, 2006), also known as REINFORCE (Williams, 1992), have historically been the estimators of choice for models with discrete random variables due to their unbiasedness and few requirements.
As they usually exhibit high variance, previous work has augmented them with variance reduction methods to improve their practicality (Williams, 1992; Ranganath et al., 2014; Mnih and Gregor, 2014).
Motivated by the eﬃciency of the pathwise estimator, recent progress in gradient estimators for discrete variables has primarily been driven by leveraging gradient information.
While the model may only be deﬁned for discrete inputs and hence gradients w.r.t.
the random variables may not be deﬁned, if we can construct a continuous relaxation of the system, then we can compute gradients of the continuous system and use them in an estimator (Gu et al., 2016;
Jang et al., 2017; Maddison et al., 2017; Tucker et al., 2017; Grathwohl et al., 2018).
While such relaxation techniques are appealing because they result in low variance estimators, they do so by taking advantage of gradient information from an appropriate continuous relaxation.
In some cases, constructing a natural continuous relaxation is nontrivial.
In other cases, the computational cost of evaluating the function at the relaxed variable values will be prohibitive, e.g. in conditional computation (Bengio et al., 2013), where discrete variables specify which parts of a large model should be evaluated and using a relaxation would require evaluating the entire model every time.
Recently, Yin et al. (2019) introduced a promising alternative to relaxation-based estimators for discrete latent variables, the Augment-REINFORCE-Swap (ARS) and AugmentREINFORCE-Swap-Merge (ARSM) estimators.
Instead of relaxing the variables, ARS and
ARSM reparameterize them as deterministic transformations of underlying continuous vari- ables.
The estimators leverage coupled samples and a careful construction relying on symmetries of the Dirichlet distribution and exponential racing.
We observe however that the continuous augmentation, which is the ﬁrst step in ARS and ARSM, increases the variance of the REINFORCE estimator.
Inspired by recent work (Dong et al., 2020), we improve both estimators by analytically integrating out unnecessary randomness introduced by the augmentation and reducing the variance of the estimator substantially.
We show that the resulting estimators consistently outperform ARS and ARSM.
However, we ﬁnd that REINFORCE with a leave-one-out-baseline (Kool et al., 2019) greatly outperforms ARS and
ARSM in all cases and is competitive or outperforms our improved estimators.
As it is a simpler estimator to implement, we recommend it in practice.
2. Background
We consider the problem of optimizing
[Equation]
with respect to the parameters θ of a factorial discrete distribution q θ ( z )
[Equation]
where k indexes dimension and α θ,k are the logits of the discrete distribution with C choices 1 .
This situation covers many problems with discrete latent variables, for example, in variational inference f θ ( z ) could be the instantaneous ELBO (Jordan et al., 1999) and q θ ( z ) the variational posterior.
The gradient with respect to θ is
[Equation]
It typically suﬃces to estimate the second term with a single Monte Carlo sample, so for notational clarity, we omit the dependence of f on θ in the following sections.
Monte Carlo estimates of the ﬁrst term can have large variance.
Low-variance, unbiased estimators of the ﬁrst term will be our focus.
2.1. Augment-REINFORCE-Swap (ARS)
Yin et al.
(2019) show that the discrete distribution can be reparameterized by an underlying continuous augmentation: if π ∼ (cid:81) k Dirichlet(1 C ) and
z k : = arg min i π k,i e − α k,i , then z k ∼ Cat( α k ) and
[Equation]
Furthermore, they deﬁne a swapped probability matrix
[Equation]
where π i (cid:11) j k has the entries at indices i and j swapped, and z i (cid:11)
j k
[Equation]
Using these constructions, they show an important identity
[Equation]
which shows that g ARS k,c is an unbiased estimator.
Importantly, if we ﬁx a reference j and compute f ( z c (cid:11) j ) for all c , we can compute g ARS k,c for all k and c with at most C potentially expensive function evaluations.
2.2. Augment-REINFORCE-Swap-Merge (ARSM)
To further improve the estimator, Yin et al.
(2019) suggest averaging over the choice of the reference j resulting in the ARSM estimator
[Equation]
We can compute g ARSM k,c for all k and c with at most C ( C − 1) / 2 + 1 function evaluations.
When C = 2, g ARSM reduces to the ARM estimator for binary variables (Yin and Zhou, 2019).
Notably, both ARS and ARSM only evaluate f at discrete values, so do not require a continuous relaxation.
Both are unbiased and we expect them to have low variance because the learning signal is a diﬀerence of evaluations of f .
Yin et al.
(2019) empirically show that it performs comparably or outperforms previous methods.
3. Methods
ARS and ARSM heavily rely on a continuous reparameterization of the problem, yet the original problem only depends on the discrete values.
Inspired by the ideas in (Dong et al., 2020), we can derive an improved estimator by integrating out the extra randomness.
Starting with ARS, ideally, we would like to compute
[Equation]
taking advantage of independence between dimensions (indexed by k ).
To reduce notational clutter
[Equation]
we omit the dimension index in the following derivation
[Equation]
We have reduced the problem to computing
[Equation]
however, we were unable to compute the expectation analytically.
Instead, we analytically integrate out some dimensions of π and use Monte Carlo sampling to deal with the rest.
First, we know that (cid:80)
i π i
= 1, so one variable is redundant (denote this choice by l ).
Next, we show how to integrate the reference index j (cid:54) =
l (i.e., compute E
π
j | π − j,l ,z 1 (cid:11) j ,...,z C (cid:11) j [ π j ], where π − j,l denotes π excluding its j -th and l -th elements.).
The known values of π − j,l , z 1 (cid:11) j , ..., z C (cid:11) j imply lower and upper bounds on π j .
First, because 1 − (cid:80) i (cid:54) =
l π i =
π l ≥ 0, we conclude that π j ≤ 1 − (cid:80) i (cid:54) =
j,l π i .
To determine the implications of the conﬁgurations z 1 (cid:11) j , ..., z C (cid:11) j , it is helpful to deﬁne additional notation.
Let s c (cid:11) j :
= π c (cid:11) j e − α .
Let’s look at what the value of z m (cid:11) j :=
arg min i s m (cid:11) j
i tells us about π j .
We need to consider two cases:
• z m (cid:11) j
= m : This means that s m (cid:11) j m =
π j e − α m is the smallest entry in s m (cid:11) j : π j e − α m ≤ min i (cid:54)
= m
s m (cid:11) j i which implies that π j ≤ min i (cid:54)
= m
e α m s m (cid:11) j i .
e α m s m (cid:11) j
i contains π l when m = l and i =
j or m (cid:54) =
l and i = l .
When m = l and i =
j , we have that e α l s m (cid:11)
j j
[Equation]
Therefore
[Equation]
A similar computation is required for the case m (cid:54) =
l and i =
l .
• z m (cid:11) j
(cid:54) = m :
This means that π j e − α m is larger than the smallest entry in s m (cid:11) j : [Equation] i which implies that π j ≥ min i e α m s m (cid:11) j i .
As above, we can eliminate π l from the bounds.
Finally, we aggregate the inequalities to compute the lower and upper bounds.
Because π ∼ Dirichlet(1 C ) is a uniform distribution over the simplex, π j | π − j,l , z 1 (cid:11) j , ..., z C (cid:11) j will be uniformly distributed over an interval, which means that it suﬃces to compute the lower and upper bounds to compute the expectation.
In principle, we should be able to apply similar ideas to ARSM.
We plan to implement this in the full version of the paper.
3.1. Leveraging symmetry
We can take advantage of symmetry to further integrate out extraneous randomness.
Let δ k
= 1 z 1 (cid:11)
j k
[Equation]
z C (cid:11) j k .
Then, we have
[Equation]
Now, we claim that inside the expectation E π − k (cid:2) f ( z m (cid:11) j ) (cid:3) is constant with respect to m .
First, we know that z 1 (cid:11) j k [Equation] k inside the expectation and that the dimensions indexed by k are independent.
Because π ∼ (cid:81) k Dirichlet(1 C ), Dirichlet(1 C ) is symmetric, and we are taking an unconditional expectation over the remaining dimensions, the value is invariant to the swapping operation.
As a result, the entire expression vanishes.
Thus, we conclude that
[Equation]
is still an unbiased estimator.
A similar argument holds for g ARSM .
This is complementary to the approach in the previous subsection and can done in combination
[Equation]
where we choose l (cid:54)
=
j uniformly randomly.
This is the estimator we use in our experiments
[Figure]
4. Experiments
[Figure]
We benchmark the proposed gradient estimators by training Variational Auto-Encoders with categorical latent variables.
To facilitate comparison with ARS and ARSM, we use the same architecture as in (Yin et al., 2019).
Brieﬂy, the model has a single layer ofcategorical latent variables which are mapped to Bernoulli logits with an MLP with two hidden layers of 256 and 512 of LeakyReLU units (Xu et al., 2015) with 0 .
2 negative slope.
The encoder mirrors the structure with two hidden layers of 512 and 256 LeakyReLU units.
For ARSM, we reduce the variance only by leveraging the symmetry (Section 3.1) and call the resulting estimator ARSM+.
We train models with 10/5/3/2-dimensional categorical latent variables on dynamically binarized MNIST.
For comparison, we train models with ARM, ARSM, and an n -sample REINFORCE estimator with a leave-one-out baseline (REINFORCE LOO) (Kool et al., 2019).
To match computation, REINFORCE LOO uses C samples for comparing against ARS/ARS+, and uses C ( C − 1) / 2 + 1 samples for
ARSM/ARSM+, where C is the number of categories.
As shown in Figure 1, the proposed estimators, ARS+/ARSM+, signiﬁcantly outperform ARS/ARSM.
Surprisingly, we ﬁnd that both ARS and ARSM underperform the simpler REINFORCE LOO baseline in all cases.
For C = 2, ARS+ 2 and ARSM+ outperform
REINFORCE LOO; however, for C > 2, REINFORCE LOO is superior and the gap increases as C does.
This suggests that partially integrating out the randomness is insuﬃcient to account for the variance introduced by the continuous augmentation.
5. Conclusion
While we achieved our goal of improving ARS and ARSM, along the way we discovered that the methods that we built upon were inferior to a simpler baseline.
In practice, we suggest using REINFORCE LOO as it achieves competitive or superior performance in all of the cases we evaluated.
However, in spite of this negative result, we still believe the techniques and careful benchmarking in this paper will be valuable to the community