Gradient Estimation with Discrete Stein Operators

Modern machine learning relies heavily on gradient methods to optimize the parameters of a learning system. However, exact gradient computation is often difficult. For example, in variational inference for training latent variable models [43, 29], policy gradient algorithms in reinforcement learning [65], and combinatorial optimization [40], the exact gradient features an intractable sum or integral introduced by an expectation under an evolving probability distribution. To make progress, one resorts to estimating the gradient by drawing samples from that distribution [50, 39].

The two main classes of gradient estimators used in machine learning are the pathwise or reparameterization gradient estimators [29, 47, 58] and the REINFORCE or score function estimators [65, 18]. The pathwise estimators have shown great success in training variational autoencoders [29] but are only applicable to continuous probability distributions. The REINFORCE estimators are more general-purpose and easily accommodate discrete distributions but often suffer from excessively high variance.

To improve the quality of REINFORCE estimators, we develop a new variance reduction technique for discrete expectations. Our method is based upon Stein operators [see, e.g., 55, 1], computable functionals that generate mean-zero functions under a target distribution and provide a natural way of designing control variates (CVs) for stochastic estimation.

We first provide a general recipe for constructing practical Stein operators for discrete distributions (Table 1), generalizing the prior literature [9, 25, 16, 8, 46, 67, 6]. We then develop a gradient estimation framework—RODEO—that augments REINFORCE estimators with mean-zero CVs generated from Stein operators. Finally, inspired by Double CV [60], we extend our method to develop CVs for REINFORCE leave-one-out estimators [49, 30] to further reduce the variance.

The benefits of using Stein operators to construct discrete CVs are twofold. First, the operator structure permits us to learn CVs with a flexible functional form such as those parameterized by neural networks. Second, since our operators are derived from Markov chains on the discrete support, they naturally incorporate information from neighboring states of the process for variance reduction.

We evaluate RODEO on 15 benchmark tasks, including training binary variational autoencoders (VAEs) with one or more stochastic layers. In most cases and with the same number of function evaluations, RODEO delivers lower variance and better training objectives than the state-of-the-art gradient estimators DisARM [14, 69], ARMS [13], Double CV [60], and RELAX [20].

We consider the problem of maximizing the objective function $\mathbb{E}_{q_{\eta}}[f(x)]$ with respect to the parameters $\eta$ of a discrete distribution $q_{\eta}(x)$. Throughout the paper we assume $f(x)$ is a differentiable function of real-valued inputs $x\in\mathbb{R}^{d}$ but is only evaluated at a discrete subset $\mathcal{X}^{d}$ due to the discreteness of $q_{\eta}$.¹¹1This assumption holds for many discrete probabilistic models [21] including the binary VAEs of Section 6. Exact computation of the expectation is typically intractable due to the complex nature of $f(x)$ and $q_{\eta}(x)$. The standard workaround is to rewrite the gradient as $\nabla_{\eta}\mathbb{E}_{q_{\eta}}[f(x)]=\mathbb{E}_{q_{\eta}}[f(x)\nabla_{\eta}\log q_{\eta}(x)]$ and employ the Monte Carlo gradient estimator known as the score function or REINFORCE estimator [18, 65]:

Here, $b$ is a constant called the “baseline” introduced to reduce the variance of the estimator by reducing the scaling effect of $f(x^{(k)})$. Since $\mathbb{E}_{q_{\eta}}[\nabla_{\eta}\log q_{\eta}(x)]=0$, the REINFORCE estimator is unbiased for any choice of $b$, and the term $b\nabla_{\eta}\log q_{\eta}(x)$ is known as a control variate (CV) [42, Ch. 8]. The optimal baseline can be estimated using additional function evaluations [64, 7, 45]. A simpler approach is to use moving averages of $f$ from historical evaluations or to train function approximators to mimic those values [37]. When $K\geq 2$, a powerful variant of REINFORCE is obtained by replacing $b$ with the leave-one-out average of function values:

This approach is often called the REINFORCE leave-one-out (RLOO) estimator [49, 30, 48] and was recently observed to have very strong performance in training discrete latent variable models [48, 14].

All the above methods construct a baseline that is independent of the point $x^{(k)}$ under consideration, but there are other ways to preserve the unbiasedness of the estimator. We are free to use a sample-dependent baseline $b(x^{(k)})$ as long as the expectation $c\triangleq\mathbb{E}_{q_{\eta}}[b(x^{(k)})\nabla_{\eta}\log q_{\eta}(x^{(k)})]$ is easily computed or has a low-variance unbiased estimate. In this case, we can correct for any bias introduced by $b(x^{(k)})$ by adding in this expectation: $\frac{1}{K}\sum_{k=1}^{K}(f(x^{(k)})-b(x^{(k)}))\nabla_{\eta}\log q_{\eta}(x^{(k)})+c.$ For example, $b(x^{(k)})$ can be a lower bound or a Taylor expansion of $f(x^{(k)})$ [43, 22]. Taking this a step further, the Double CV estimator [60] proposed to treat $f(x^{(k)})-b(x^{(k)})$ as the effective objective function and apply the leave-one-out idea:

The resulting estimator adds two CVs to RLOO: the global CV $b_{k}(x^{(k)})\nabla_{\eta}\log q_{\eta}(x^{(k)})$ and the local CV $b_{k}(x^{(j)})$. Intuitively, $b_{k}(x^{(j)})$ is aimed at reducing the variance of the LOO average. This is motivated by the fact that replacing $\frac{1}{K-1}\sum_{j\neq k}f(x^{(j)})$ with $\mathbb{E}_{q_{\eta}}[f]$ in RLOO leads to lower variance [60, Prop 1]. To obtain a tractable correction term $c$, Titsias and Shi [60] adopt a linear design of $b_{k}$: $b_{k}(x)=\alpha\cdot\frac{1}{K-1}\sum_{j\neq k}f(x^{(j)})(x-\mu)$ for $\mu=\mathbb{E}_{q_{\eta}}[x]$ and a coefficient $\alpha$.

Although the Double CV framework points to a promising new direction for developing better REINFORCE estimators, one only obtains significant reduction in variance when $b_{k}$ is strongly correlated with $f$. This may fail to hold for the above linear $b_{k}$ and a highly nonlinear $f$, especially in applications like training deep generative models.

In the following two sections, we will introduce a method that allows us to use very flexible CVs while still maintaining a tractable correction term. Our method enables online adaptation of CVs to minimize gradient variance (similar to RELAX [20]) but does not assume $q_{\eta}$ has a continuous reparameterization. We then apply it to generalize the linear CVs in Double CV to very flexible ones such as neural networks. Moreover, we provide an effective CV design based on surrogate functions that requires no additional evaluation of $f$ compared to RLOO.

At the heart of our new estimator is a technique for generating flexible discrete CVs, that is for generating a rich class of functions that have known expectations under a given discrete distribution $q$. One way to achieve this is to identify any discrete-time Markov chain $(X^{(t)})_{t=0}^{\infty}$ with stationary distribution $q$. Then, the transition matrix $P$, defined via $P_{xy}=P(X^{(t+1)}=y|X^{(t)}=x)$, satisfies $P^{\top}q=q$ and hence

for any integrable function $h$. We overload our notation so that, for any suitable matrix $A$, $(Ah)(x)\triangleq\sum_{y}A_{xy}h(y)$. In other words, for any integrable $h$, the function $(P-I)h$ is a valid CV as it has known mean under $q$. Moreover, the linear operator $P-I$ is an example of a Stein operator [55, 1] in the sense that it generates mean-zero functions under $q$. In fact, both Stein et al. [56] and Dellaportas and Kontoyiannis [12] developed CVs of the form $(P-I)h$ based on reversible discrete-time Markov chains and linear input functions $h(x)=x_{i}$.

More generally, Barbour [3] and Henderson [23] observed that if we identify a continuous-time Markov chain $(X^{(t)})_{t\geq 0}$ with $q$ as its stationary distribution, then the generator $A$, defined via

satisfies $A^{\top}q=0$ and hence $\mathbb{E}_{q}[Ah]=0$ for all integrable $h$. Therefore, $A$ is also a Stein operator suitable for generating CVs. Moreover, since any discrete-time chain with transition matrix $P$ can be embedded into a continuous-time chain with transition rate matrix $A=P-I$, this continuous-time construction is strictly more general [57, Ch. 4].

Recall that REINFORCE estimates the gradient $\mathbb{E}_{q_{\eta}}[f(x)\nabla_{\eta}\log q_{\eta}(x)].$ Due to the existence of $\nabla_{\eta}\log q_{\eta}(x)$ as a weighting function, we apply a discrete Stein operator to a vector-valued function $\tilde{h}:\mathcal{X}\to\mathbb{R}^{d}$ per dimension to construct the following estimator with a mean-zero CV:

Ideally, we want to choose the $\tilde{h}$ such that $A\tilde{h}_{i}$ strongly correlates with $f(x)\nabla_{\eta_{i}}\log q_{\eta}(x)$ to reduce its variance. The optimal $\tilde{h}_{i}$ that yields zero variance is given by the solution of Poisson equation

We could learn a separate $\tilde{h}_{i}$ per dimension to approximate the solution of 10. However, this poses a difficult optimization problem that requires simultaneously solving $d$ Poisson equations. Instead, we will incorporate knowledge about the structure of the solution to simplify the optimization.

To determine a candidate functional form for $\tilde{h}$, we draw inspiration from an “optimal” Markov chain in which the current state is ignored entirely and the new state is generated independently from $q_{\eta}$. In this case, $(P-I)\tilde{h}_{i}$ becomes $\mathbb{E}_{q_{\eta}}[\tilde{h}_{i}]-\tilde{h}_{i}$, and the optimal solution is $\tilde{h}_{i}=f\nabla_{\eta_{i}}\log q_{\eta}$. Inspired by this, we consider $\tilde{h}$ of the form

where we now only need to learn a scalar-valued function $h$. Notably, when $h$ exactly equals $f$ and $A=P-I$ for any discrete time Markov chain kernel $P$, our CV adjustment amounts to Rao-Blackwellization [42, Sec. 8.7], as we end up replacing $f(x)\nabla_{\eta_{i}}\log q_{\eta}(x)$ with its conditional expectation $P(f\nabla_{\eta_{i}}\log q_{\eta})(x)=\mathbb{E}_{X_{t+1}|X_{t}=x}[f(X_{t+1})\nabla_{\eta_{i}}\log q_{\eta}(X_{t+1})]$. This yields a guaranteed variance reduction.

Surrogate function design   Based on the above reasoning, we can view $h$ as a surrogate for $f$. We avoid directly setting $h=f$ because our Stein operators evaluate $h$ at all neighbors of the sample points, and $f$ can be very expensive to evaluate [see, e.g., 51]. To avoid this problem, we first observe that $\tilde{h}$ 11 can be made sample-specific, i.e., we can use a different $\tilde{h}_{k}$ for each sample point $x^{(k)}$:

We then consider the following choices of $h_{k}$ that are informed about $f$ while being cheap to evaluate:

Here $H$ belongs to a flexible parametric family of functions such as neural networks and is chosen to have significantly lower cost than $f$. $y$ will take on the values of $x^{(k)}$ and its neighbors in the Markov chain. We omit $x^{(k)}$ in the sum 13 to ensure that the function $h_{k}$ is independent of $x^{(k)}$ and hence that $\mathbb{E}_{q_{\eta}}[\frac{1}{K}\sum_{k=1}^{K}(A\tilde{h}_{k})(x^{(k)})]=0$. Moreover, this surrogate function design introduces no additional evaluations of $f$ beyond those required for the usual RLOO estimator. Also, as observed by Titsias and Shi [60], for many applications, including VAE training, where $f$ has parameters $\theta$ learned through gradient-based optimization, $\{\nabla f(x^{(k)})\}_{k=1}^{K}$ can be obtained “for free” from the same backpropagation used to compute $\nabla_{\theta}f_{\theta}(x^{(k)})$ (see Algorithm 1).

RODEO  We can further improve the estimator by leveraging discrete Stein operators and the above surrogate function design to construct both the global and local CVs in the Double CV framework [60] (see Section 2). We call our final estimator RODEO for RLOO with Discrete StEin Operators:

where $\tilde{h}^{\star}_{k}(y)=h^{\star}_{k}(y)\nabla_{\eta}\log q_{\eta}(y)$ and $\{h_{k},h_{k}^{\star}\}_{k=1}^{K}$ are scalar-valued functions. Here, $(Ah_{j})(x^{(j)})$ is a scalar-valued CV introduced to reduce the variance of the leave-one-out baseline $\frac{1}{K-1}\sum_{j\neq k}f(x^{(j)})$, while $(A\tilde{h}_{k}^{\star})(x^{(k)})$ acts as a global CV to further reduce the variance of the gradient estimate. We adopt the aforementioned design of $h$ 13 and $\tilde{h}$ 11 for the two CVs. In Appendix B, we show that the RODEO estimator is unbiased for $\nabla_{\eta}\mathbb{E}_{q_{\eta}}[f(x)]$ when each $h_{k}$ is defined as in 13 and

Optimization with RODEO   In practice, we let the functions $H$ and $H^{\star}$ share a neural network architecture with two output units. Since the estimator is unbiased, we can optimize the parameters $\gamma$ of the network in an online fashion to minimize the variance of the estimator (similar to [20]). Specifically, if we denote the RODEO gradient estimate by $g_{\gamma}(x^{(1:K)})$, then $\nabla_{\gamma}\mathrm{Tr}(\mathrm{Var}(g_{\gamma}(x^{(1:K)})))=\mathbb{E}[\nabla_{\gamma}\|g_{\gamma}(x^{(1:K)})\|_{2}^{2}].$ In Algorithm 1, we use an unbiased estimate of this hyperparameter gradient, $\nabla_{\gamma}\|g_{\gamma}(x^{(1:K)})\|_{2}^{2}$, to update $\gamma$ after each optimization step of $\eta$.

As we have seen in Section 2, there is a long history of designing CVs for REINFORCE estimators using “baselines” [64, 7, 45, 37]. Recent progress is mostly driven by leave-one-out [49, 38, 30, 48] and sample-dependent baselines [43, 22, 60, 62, 20]. REBAR [62] constructs the baseline through continuous relaxation of the discrete distribution [26, 35] and applies reparameterization gradients to the correction term. As a result REBAR uses three evaluations of $f$ for each $x^{(k)}$ instead of the usual single evaluation (see Appendix A for a detailed explanation). The RELAX [20] estimator generalizes REBAR by noticing that their continuous relaxation can be replaced with a free-form CV. However, in order to get strong performance, RELAX still includes the continuous relaxation in their CV and only adds a small deviation to it. Therefore, RELAX also uses three evaluations of $f$ for each $x^{(k)}$ and is usually considered more expensive than other estimators.

An attractive property of the RODEO estimator is that it incorporates information from neighboring states thanks to the Stein operator while avoiding additional costly $f$ evaluations thanks to the learned surrogate functions $h_{k}$. Estimators based on analytic local expectation [59, 61] and GO gradients [11] also use neighborhood information but only at the cost of many additional target function evaluations. In fact, the local expectation gradient [59] can be viewed as a Stein CV adjustment RODEO with a Gibbs Stein operator and the target function $f$ used directly instead of the surrogate $h$. The downside of these approaches is that $f$ must be evaluated $Kd$ times per training step instead of $K$ times as in RODEO, a prohibitive cost when $f$ is expensive and $d\geq 200$ as in Section 6.

Prior work has also studied variance reduction methods based on sampling without replacement [31] and antithetic sampling [68, 14, 13, 15] for gradient estimation. DisARM [14, 69] was shown to outperform RLOO estimators when $K=2$ and ARMS [13] generalizes DisARM to $K>2$.

Stein operators for continuous distributions have also been leveraged for effective variance reduction in a variety of learning tasks [2, 36, 41, 5, 54, 52] including gradient estimation [34, 27]. In particular, the gradient estimator of Liu et al. [34] is based on the Langevin Stein operator [19] for continuous distributions and coincides with the continuous counterpart of RELAX [20]. In contrast, our approach considers discrete Stein operators for Monte Carlo estimation in discrete distributions with exponentially large state spaces. Recently, Parmas and Sugiyama [44, App. E.4] used a probability flow perspective to characterize all unbiased gradient estimators satisfying a mild technical condition; our estimators fall into this broad class but were not specifically investigated.

Python code replicating all experiments can be found at https://github.com/thjashin/rodeo.

This work tackles the gradient estimation problem for discrete distributions. We proposed a variance reduction technique which exploits Stein operators to generate control variates for REINFORCE leave-one-out estimators. Our RODEO estimator does not rely on continuous reparameterization of the distribution, requires no additional function evaluations per sample point, and can be adapted online to learn very flexible control variates parameterized by neural networks.

One potential drawback of our surrogate function constructions 13 and 14 is the need to evaluate an auxiliary function ($H$ or $H^{*}$) at $K-1$ locations. This cost can be comfortably borne when $H$ and $H^{*}$ are much cheaper to evaluate than $f$, such as in the ResNet VAE example in Section C.1 and many large VAE models used in practice [63]. And, in our experiments with the most common sample sizes, the runtime of RODEO was no worse than that of RELAX. To obtain a more favorable cost for large $K$, one could employ alternative surrogates that require only a constant number of auxiliary function evaluations, e.g.,

The runtime of RODEO could also be improved by varying the Stein operator employed. For example, the Gibbs operator $(Ah)(x)$ 4 used in our experiments evaluated its surrogate function $h$ at $d$ neighboring locations of $x$. This evaluation complexity could be reduced by subsampling neighbors, resulting in a cheaper but still valid Stein operator, or by employing the numerically stable Barker operator 6 with fewer neighbors. Either strategy would introduce a speed-variance trade-off worthy of study in follow-up work. Finally, we have restricted our focus in this work to differentiable target functions $f$. In future work, this limitation could be overcome by designing effective surrogate functions that make no use of derivative information.

We thank Heishiro Kanagawa for suggesting appropriate names for the Barker Stein operator.