Gradient Estimation with Discrete Stein Operators
Authors: Jiaxin Shi
Yuhao Zhou
Jessica Hwang
Michalis Titsias
Lester Mackey
Venue: NeurIPS 2022 Accept (2022)
Keywords: Gradient estimation
Stein's method
Markov chain
score function
REINFORCE
discrete latent variables
VAE
control variates
variance reduction

Abstract: Gradient estimation---approximating the gradient of an  expectation  with respect to the parameters of a distribution---is central to the solution of  many machine learning problems.  However, when the distribution is discrete, most common gradient estimators suffer from excessive variance. To improve the quality of gradient estimation, we introduce a variance reduction technique based on Stein operators for discrete distributions. We then use this technique to build flexible control variates for the REINFORCE leave-one-out estimator.  Our control variates can be adapted online to minimize variance and do not require extra evaluations of the target function. In benchmark generative modeling tasks such as training binary variational autoencoders, our gradient estimator achieves substantially lower variance than state-of-the-art estimators with the same number of function evaluations.

Gradient Estimation with Discrete Stein Operators
Jiaxin Shit
Yuhao Zhou
Stanford University
Tsinghua University yuhaoz.cs cs@gmail.com
jiaxins@stanford edu
Jessica Hwang Stanford University j jhwangdstanford. edu
Michalis K. Titsias DeepMind mtitssias@goggle.cm com
Lester Mackey
Microsoft Research New England 1mackey@nicrosoft.. com
Abstract
Gradient estimation estimation-aaproxxiiinn -pproximating the gradient of an expectation with respect to the parameters of a distribution central to the solution of many machine learn- ing problems. However, when the distribution is discrete, most common gradient estimators suffer from excessive variance. To improve the quality of gradient esti- mation, introduce a variance reduction technique based Stein operators for discrete distributions. We then use this technique to build flexible control variates for the REINFORCE leave-one-out estimator. Our control variates can be adapted online to minimize variance and do not require extra evaluations of the target func- tion. In benchmark generative modeling tasks such training binary variational autoencoders, our gradient estimator achieves substantially lower variance than state-of-te of-th--rrt estimators with the same number of function evaluations.
Introduction
Modern machine learning relies heavily on gradient methods to optimize the parameters of learning system. However, exact gradient computation often difficult. For example, in variational inference for training latent variable models [29, 43], policy gradient algorithms in reinforcement learning [65], and combinatorial optimization [40]. the exact gradient features an intractable sum or integral introduced by expectation under evolving probability distribution. To make progress, one resorts to estimating the gradient by drawing samples from that distribution [39.
The two main classes of gradient estimators used in machine learning are the pathwise or reparameteri- zation gradient estimators [29, 47, 58] and the REINFORCE or score function estimators [18, 65J. 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 new variance reduction technique for discrete expectations. Our method is based upon Stein operators [see, e.g., 1, 55], computable functionals that generate mean-zero functions under target distribution and provide natural way of designing control variates (CVs) for stochastic estimation.
We first provide general recipe for constructing practical Stein operators for discrete distribu- tions (Table 1), generalizing the prior literature [6, 8, 9, 16, 25, 46, 67]. We then develop a gradient estimation Franewor-RRDDOO--n
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 [30, 49] further reduce the variance.
1 Work done while at Microsoft Research New England.
36th Conference on Neural Information Processing Systems (NcurIPS 2022).

---

Table 1: Discrete Stein operators that generate mean-zero functions under q (i.e., (x((AA)(<)]] 0)
Stein Operator
(a))
Gibbs (4)
E Ey_,
1((a)
MPF (6)
EveNs ((W)//(((((()) h(x))
Barker (6)
EveNswwte (()))(() (M(v) h(z))
Difference (8)
E h(dec;(z) (()n) ((x)
[(x)
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 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 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-- gradient estimators DisARM [14, 69), ARMS [13]. Double CV [60], and RELAX [20].
2 Background
We consider the problem of maximizing the objective function [x()] with respect to the pa- rameters 17 of a discrete distribution 9,, (z)).. Throughout the paper we assume f(z) is a differen- tiable function of real-valued inputs x I2 but only evaluated
discrete subset x due to the discreteness of 41. Exact computation of the expectation is typically intractable due to the complex nature of (x) and 9, (z).. The standard workaround is rewrite the gradient as V ,g, [x)) = E,, [f (x) V, log Jn (z)) and employ the Monte Carlo gradient estimator known as the score function or REINFORCE estimator [18, 65]:
1 K
- b) log 9,, (z(1)) for xx11) ,...., (K)
4.6.d. ~
9n-
(REINFORCE)
Here, is a constant called the "baseline" introduced to reduce the variance of the estimator by reducing the scaling effect of (xx(()). Since B,,, V,, (,,
O, the REINFORCE estimator is
unbiased for any choice of b, and the term bV, log 9n (x) is known control variate (CV) [42,
Ch. 8J. The optimal baseline can be estimated using additional function evaluations [7, 45, 64]. A simpler approach is to moving averages of f from historical evaluations to train function approximators to mimic those values [37]. When K 2 2, a powerful variant of REINFORCE is obtained by replacing b with the leave-one-out average of function values:
EK -
X Ejva (((11) V log (x))))
(RLOO)
This approach is often called the REINFORCE leave-one-out (RLOO) estimator [30, 48, 49] and was recently observed to have very strong performance in training discrete latent variable models [14, 48].
All the above methods construct baseline that is independent of the point (R) under consideration, but there are other ways to preserve the unbiasedness of the estimator. We are free to use a sample- dependent baseline ((xx())) as long as the expectation e - E,,, [b(x(())V, log qn (x) is easily computed or has low-variance unbiased estimate. In this case, can correct for any bias introduced by ((x(()) by adding in this expectation: A "E
((x(())) V, log 4y (2x(A)) + For example, (x((k) can be a lower bound or a Taylor expansion of (k) [22, 43). Taking this a step further, the Double CV estimator [60] proposed to treat (xx(+)) - b(z((t)) as the effective objective function and apply the leave-one-out idea:
1/ K (((()())) bu (x)) )
(f(z() by (x())) )] log q, (x((t) +
The resulting estimator adds two CVs to RLOO: the global CV by (x() V n log 4n (x))) and the local CV by (z(i)). Intuitively, by (I)( is aimed at reducing the variance of the LOO average. This is
This assumption holds for many discrete probabilistic models [21] including the binary VAES of Section 6.

---

motivated by the fact that replacing K-i Σ Jes f((()) with Eg [f] in RLOO leads to lower variance [60, Prop 1]. To obtain tractable correction term c, Titsias and Shi [60] adopt linear design of by: (x) = a K-1 Lith (x,))x= ( 14) for μr ag, [x] and a coefficient a.
Although the Double CV framework points to promising new direction for developing better REINFORCE estimators, one only obtains significant reduction in variance when by strongly correlated with f. This may fail to hold for the above linear by and a highly nonlinear especially in applications like training deep generative models.
In the following two sections, we will introduce 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 (, has a continuous reparameterization. We then apply it to generalize the linear CVs in Double CV to very flexible ones such neural networks. Moreover, we provide effective CV design based surrogate functions that requires no additional evaluation f compared to RLOO.
3 Control Variates from Discrete Stein Operators
At the heart of our new estimator is technique for generating flexible discrete CVs, that is for generating a rich class of functions that have known expectations under given discrete distribution q. One way to achieve this is to identify any discrete-time Markov chain (X(')) )2 with stationary
distribution q. Then, the transition matrix P, defined via Pay = (xx(++1) y x xr), satisfies PT q = q and hence
I,((P 1)h] =0
(1)
for any integrable function h. We overload our notation that, for any suitable matrix A, (Z) & E, hty) In other words, for any integrable h. the function (P - - is a valid CV as it has known mean under q. Moreover, the linear operator p- / is an example of Stein operator [1, 55] in the sense that it generates mean-zero functions under In fact, both Stein et al. [56] and Dellaportas and Kontoyiannis [12] developed CVs of the form (P (P-1) r) based on reversible discrete- Markoy chains and linear input functions hzz) T3,
More generally, Barbour [3] and Henderson [23] observed that we identify a continuous-time Markov chain (XX(1) with q as its stationary distribution, then the generator A, defined via
(Ah) lim BMX* IX =x) t-0
satisfies A q - 0 and hence [A]] 0 for all integrable h. Therefore, A is also 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 P - I, this continuous-time
construction is strictly more general [57, Ch. 4].
3.1 Discrete Stein operators
We next present several examples of broadly applicable discrete Stein operators (summarized in Table 1) that can serve as practical defaults.
Gibbs Stein operator The transition kemnel of the random-scal Gibbs sampler with stationary distribution q [see, e.g., 17]
Pay = 1 E (ulk--))y xx-,), where 1(-) is the indicator function. The associated Stein operator is A = P - with
(3)
(Ah)(z) = Ey <(v(l=-))((v) (( flxz).
(4)
In the binary variable setting, (4) recovers the operator Bresler and Nagaraj [8], Reinert and Ross [46]
used to bound distances between the stationary distributions of Glauber dynamics Markov chains.
Zanella Stein operator Zanella [70] studied continuous-timg Markov chains with generator
Ary = k (q())/g(z)) 1 E Nz,, + x) - Ex Agz 1(y = x),
(5)
3

---

where N₂ denotes the transition neighborhood of and is a continuous positive function satisfying
x(t) fx(1/t). Conveniently, the neighborhood structure can be arbitrarily sparse. Moreover, the
constraint ensures that the chain satisfies detailed balance (i.e., ((x) Azy = q(v) Arx) and hence that
AJg 0. Hodgkinson et al. [24] call the associated operator the Zanella Stein operator:
(Ah)(z) = EveN.
Liv
(h(v) 11(2)).
(6)
Important special cases include the minimum probability flow [53] (MPF) Stein operator (x(t) discussed in Barp et al. [6] and the Barker [4] Stein operator (x<()) = (
vt)
Birth-death Stein operator Let ....... Ca represent the standard basis vectors on R For any umite-cardinalig space, we may index the elements of X {se..... Sm-1). let (xx(x;) return the index of the element that represents, and define the increment and decrement operators
(xr) = + es (=((dax((x))+1) mod m x;) and dec,(x) x (i(3(1dx(z))1)) mod xz).
Then the product birth-death process [28] X with birth rates bi.
2(3ner(x)) death rates ((x))
di, 1, and generator
Asy 3 Ei b inc;(x)) + d1 dec,(x)) (bs,z + dr.x) z),
has q as stationary distribution, as A 0. This construction yields the birth-death Stein operator
(Ag)(z) = E (ggn((( - g(x)) - dg.=((()) - g(dec;
An analogous operator is available for countably infinite A' [24], and Brown and Xia [9], Eichelsbacher and Reinert [16], Holmes [25] used related operators to characterize convergence to discrete target
distributions. Moreover, by substituting h(x) = g(x) g(ine (z)) in (7), we recover the difference Stein operator proposed by Yang et al. [67] without reference to birth-death processes:
(Ah)(t) = b h(dec;(x)) (x))
(8)
Choosing Stein operator Despite being better-known, the difference and MPF operators often suffer from large variance and numerical instability due to their use of unbounded probability ratios (x)). As result, recommend the numerically stable Gibbs operator when each component of x is binary or takes on a small number of values. When 2; takes on large number of values (m), the Gibbs operator suffers from linear scaling with m. In this case, we recommend the Barker operator where sparse neighborhood structure can be specified, such as N₁ (inc,(x), dec;(r) for E [d]].
The Barker operator is numerically stable as its K ((v)
o(v)
7(x)
(())+
Gradient Estimation with Discrete Stein Operators
Recall that REINFORCE estimates the gradient E,,, (s(x)) log (x)).. Due to the existence of V, log (z) weighting function, apply discrete Stein operator to vector-valuce function h X > Rd per dimension to construct the following estimator with mean-zero CV:
Eg, (x:) n log (x) + (Aha)
(9)
Ideally, we want to choose the such that Ad, will be strongly correlated with f log (x) reduce its variance. The optimal h, is given by the solution of Poisson equation
to
Egg [ log gml JV ns log 9n = Añ₃-
(10)
We could learn a separate ht per dimension to approximate the solution of (10). However, this poses 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 candidate functional form for h, we draw inspiration from an "optimal" Markov chain in which the current state is ignored entirely and the state is generated independently from gn. In

---

Algorithm 1 Optimizing [E,, with RODEO gradients
input: Objective fe, sample points (1:K) (id. qn, Stein operator A, step sizes ae; S
for T do
1: (fo( (olo()) ( + autooiif((( x(1
2: Compute the surrogates by (x:l)), (g(10)) of (13), (14) for j * & and j. = 1, K.
3: Compute the RODEO gradient estimator on (x(1:k)).
4: 0 a- 0 + o
K
k= Vol(x()))
5: n n+ a.9y (20,((((11)) (x...
6: Update hyperparameters: 7 + 7 - 3.V, (19,
this case, (P- - T) becomes B,, [bd] his and the optimal solution is h, I'Vn log On- Inspired by this, we consider h of the form
h(x) h(x) log 9n, (2),
(11)
where now only need learn scalar-valued function h. Notably, when h exactly equals and = P - T for any discrete time Markov chain kernel P. our CV adjustment amounts to Rao- Noo-Bllccwellizationn [42, Sec. 8.7], end up replacing f(x) nn log 9r, (x) with its conditional expectation P(fV n. log (x) = xxx==l variance reduction.
log (Xx+1)) This yields a guaranteed
Surrogate function design Based on the above reasoning, we can view h as a surrogate for f. We avoid directly setting h = because 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 i (1 1) can be made sample-specific, i.e., we can use a different h2 for each sample point (*));
K # I (f(x())) V, log (x())) + (Ahx))a
(12)
We then consider the following choices of hy that are informed about f while being cheap to evaluate: hv( マー H( (1))), () -
(13)
Here H belongs to a flexible parametric family of functions such as neural networks and chosen have significantly lower cost than f- y will take on the values of (*()) and its neighbors in the Markov chain. We omit (*) in the sum (13) ensure that the function hę is independent of (k) and hence that 1 >k=1 (Ah))
O. Moreover, this surrogate function design introduces no additional evaluations of f beyond those required for the usual RLOO estimator. Also, observed by Titsias and Shi [60], for many applications, including VAE training, where has parameters learned through gradient-based optimization, {V. (zx(((()))) )=== )) can be obtained "for free" from the same backpropagation used to compute (x)) (see Algorithm 1).
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:
RODEO
K (1) x-1 (Ah))))+(AA))))))))))
log +AA (RODEO)
(Ar))
I#K
where h2(()) V, log (2) and (he, hilk are scalar-valued functions. Here, (Ah,) (z01)) is scalar-valued CV introduced to reduce the variance of the leave-one-out base- line RI Ejyx (zx6)), while (x(1) acts as global CV to further reduce the variance of the gradient estimate. We adopt the aforementioned design of h (13) and h (11) for the two CVs. In Appendix B, we show that the RODEO estimator is unbiased for V,l,, [fx) when each hy is defined as in (13) and
(1( R-1 (((E (f((() -
(14)
Optimization with RODEO practice, we let the functions and H* share neural net- work architecture with two output units. Since the estimator is unbiased, we can optimize the
5

---

Table 2: Training binary latent VAEs with K = 2.3 (except for RELAX which uses 3 evaluations) on MNIST, Fushion-MNIST, and Omniglot. We report the average ELBO (#1 standard error) on the training set after 1M steps over 5 independent runs. Test data bounds are reported Table
Bemoulli Likelihoods
Gaussian Likelihoods
MNIST
Fashion-MNIST
Omniglot
MNIST
Fashion- Fashiion-NNST
Ommiglot
DiSARM
102. 75 + 0.08
-237.68 +: 0.13
-116.50 0.04
668
038.03 0.61
182.65 + 0.47
446.61 + 1.16
Double
-102.14 0.06
2397.55
0.16
-116.391
1.18
186.35 0.64
447.65 0.87
RODEO (Ours)
-101.89 1 0.17
237.41 0.09
-115.93 0.06
681.95
+ 0.37
191.81
+ 0.67
451.4 1.11
ARMS
-100.81 0.14
237.03
-115.21
683.55
1.01
193.07
0.34
457.98
1.03
Double CV
-100.94 0.09
237.40
0.11
-115.06
0.12
686.48
0.68
193.93
2 0.20
457.44 : 0.79
RODEO (Ours)
-100.46
0.13
238.88
1 0.12
-115.01 1: 0.05
692.37
1 0.39
196.56 1: 0.42
461.87
0.90
RELAX (3 ovals)
-101.99
+ 0.04
237.74
115.70
688.58
0.52
196.38 2.0.8::00 0.86
462 23 0.63
parameters ?y 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 (x.(1.))) then V., TT,Tr(aa((gg) (ar (x(:xk))) EVV (ll (1()))) In Algorithm I, we use an unbiased estimate of
this hyperparameter gradient. V V (2 (((()))) to update after each optimization step of 1).
5 Related Work
As we have seen in Section 2, there is long history of designing CVs for REINFORCE estimators
using "baselines" [7, 37, 45. 64]. Recent progress is mostly driven by leave-one-out [30, 38, 48, 49] and sample-dependent baselines [20, 22, 43, 60, 62]. 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 (A) instead of usual single evaluation (see Appendix A for 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 small deviation to it. Therefore, RELAX also uses three evaluations of f for each (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 the Stein operator while avoiding additional costly f evaluations thanks the learned surrogate functions h4. 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 / used directly instead of the surrogate h. The downside of these approaches is that must be evaluated K dl times per training step instead of K times as in RODEO, a prohibitive cost when f expensive and Z 200 as in Section
Prior work has also studied variance reduction methods based on sampling without replacement [31] and antithetic sampling [13-15, 68] for gradient estimation. DisARM [14, 69] was shown to outper- form 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, 5, 36, 41, 52, 54] including gradient estimation [27, 34]. 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.
6 Experiments
Python code replicating all experiments can be at /ttot//tt.t//ttttto//////a//i/e/

---

6.1 Training Bernoulli VAEs
Following Dong et al. [14] and Titsias and Shi [60], we conduct experiments on training variational auto-encoders [29, 47] (VAEs) with Bernoulli latent variables. VAES are models with a joint density z), where x is the latent variable and 0 denotes model parameters. They are typically learned through maximizing the evidence lower bound (ELBO) (=)) for an auxiliary inference network (, and f(x) & log po(y. x) log (1 y). In our experiments, p(z) and g, (2ly) are high-dimensional distributions where each dimension of the random variable is an independent Bernoulli. Since exact gradient computations are intractable, we will use gradient estimators to learn the parameters of the inference network. The VAE architecture and training experimental setup follows Titsias and Shi [60], and details are given in Appendix D. The dimensionality of the latent variable z is d 200. The functions H (13) and H (14) share neural network architecture with two output units and a single hidden layer with 100 units. For the numerical stability and variance reasons discussed in Section 3, we use the Gibbs Stein operator (4) as default choice in our experiments, but we revisit this choice in Section 6.3.
We consider the MNIST [33], Fashion-MNISS [66] and Omniglot [32] datasets using their standard train, validation, and test splits. We use both binary and continuous VAE outputs (g) as in Titsias and Shi [60]. In the binary output setting, data are dynamically binarized, and the Bernoulli likelihood is used; in the continuous output setting, data are centered between [-1, 1). and the Gaussian likelihood with learnable diagonal covariance parameters is used.
1o=2
avals
1e=2
evall
evals
evall
Gradient Vertence
1.8
Double CV
DISARM
ROBEO (ours)
1.0
RODEO (ours)
ARMS
Double CV
Tranio SeEa
-004
102
-102
7.4
0.8
1.2
106
-004
1.0
0.6
200K 400K BOOK
B0OK
TM
200K 400K
BOOK B00K 1M
-108
200K
400K
GOOK 500K
1M
-108
200K
4D0K 6O0K 500K 1M
Training Stop
Traiiing Stop
Training Stop
Training Stop
Figure I: Training binary latent VAEs with 2 or 3 evaluations per step on binarized MNIST.
== 2 In the first set of experiments we focus on the most common setting of K = 2 sample points and compare our variance reduction method to its counterparts including DisARM [14] and Double CV [60]. Results for RLOO are omitted since it is consistently outperformed by Double CV [see 60]. Table 2 shows, on all three datasets, RODEO achieves the best training ELBOs. In Figure (left), we plot the gradient variance and average training ELBOs against training steps for all estimators on dynamically binarized MNIST. RODEO outperforms DisARM and Double CV by large margin in gradient variance.
Next, consider VAES with Gaussian likelihoods trained on non-binarized datasets. In this case the gradient estimates suffer from even higher variance due to the large range of values f(z) can take. The results are plotted in Figure 2. We see that RODEO has substantially lower variance than DisARM and Double CV, leading to significant improvements in training ELBOs. In Appendix Figure 7, we find that RODEO also consistently yields the best test set performance in all six settings.
K ==3 3 In the second set of experiments we compare RELAX [20], which uses three evaluations of f per training step, with RODEO, Double CV, and ARMS [13] for K = 3. Figure 1 (right) and Table demonstrate that RODEO outperforms the three previous methods and generally leads to the lowest variance. Although RELAX was often observed to have very strong performance in prior work [14, 60], our results in Figure suggest that, for dynamically binarized datasets, much larger gains can be achieved by using the same number of function evaluations in other estimators.
For the experiments mentioned above, we report final training ELBOs in Table 2, test likeliieldoood bounds in Appendix Table 4, and binarized MNIST average running time in Appendix Table 5. For K = 3. RODEO has the best final performance in 5 out of 6 tasks and runtime nearly identical to RELAX. For K = 2, RODEO has the best final performance for all tasks and is 1.5 times slower than Double CV and DiSARM. We attribute the runtime gap to the Gibbs operator (4) which performs 2d evaluations of the auxiliary functions H and H* in (13) and (14). While this results in some extra cost, the network parameterizing H and H* takes only 2-dimensional inputs, produces scalar outputs, and is typically small relative to the VAE model. As result, the evaluation of H and H' is
7

---

MNIST
Fashior-MNIST
Omniglet
Grodiont Varaanza
2.0
RODEO
Double
DMARM
430K
500K
200K
600K Training Step
200K
600K Training Step
Training Step
875
450
180
Training
180
825
RODEO
140
350
DICARM
120
Training
Stop
Step
Figure 2: Training binary latent VAEs with Gaussian likelihoods, K = 2, and non-binarized datasets.
significantly cheaper than that of f. and its relative contribution to runtime shrinks as the cost of f grows. To demonstrate this, we include wall clock time comparison of method with RLOO in Appendix C.I. In this experiment we replace the two-layer MLP-baaed VAE with ResNet VAE, where the cost of ff is significantly higher than the single-layer MLP of H, H In this case, RODEO and RLOO have very close per-iteration time (0.025s vs. 0.023s). And RODEO achieves better training ELBOs than RLOO for the same amount of time.
Training bierarchical Bernoulli VAEs
To investigate the performance of RODEO when scaling hierarchical discrete latent variable models, we follow DisARM [14, 68] to train VAEs with 2/3/4 stochastic layers, each of which consists of 200 Bernoulli variables. We set K 2 and compare estimator with DisARM and Double CV on dynamically binarized MNIST, Fashion-MNIST. and Omniglot. For each stochastic layer, we use different CV network which has the same architecture those in VAE experiments from the previous section. More details are presented in Appendix D.
6x'00"
Toiniing ELBO
-238
Cradient Vartance
4x10""
Double
RODED (Ours)
8x
-237
2x10"
10"
-238
10"
-239
500K
600K 700K ROOK 500K Training Slep
1M
200K
400K
60CK
500K
11
10K
20K
50K
40K
50K
Training Siep
Training Step
Figure 3: Training hierarchical binary latent VAEs with four stochastic layers on MNISS. In this experiment, the estimators have very different behaviors towards the beginning and the end of training. We show this on the right by zooming into the first 50K steps of the gradient variance plot.
We plot the results on VAEs with four stochastic layers on Fashion for other datasets and for and stochastic layers can be found in Appendix E. RODEO generally achieves the best training ELBOs. One difference noticed when comparing the variance results with those obtained from single-stochastic-lyye behaviors towards the beginning and the end of training. For example, in Figure 3, Double CV starts with lower variance than DisARM, but the gap diminishes after around 100K steps, and DisARM start to perform better the training proceeds. In contrast, RODEO has the lowest variance in both phases for all datasets save Omniglot, where DisARM overtakes it in the long run.
in Figure 3. The results
VAES is that these estimators have very different
6.3 Ablation study
Finally, we conduct an ablation study to gain more insight into the impact of each component of RODEO. In each experiment train binary latent VAEs with K 2.
8

---

2.50
Goddant Vartanoe
Cibla
Difienenee
Gradent Variange
2.25
LO0
100 -gibbal
(00-00000-- loog
Gradient Vaniarce
RODEO
with sunogate
LOO -ooal
1.25
1.50
1-2
200K
400K
GOOK
BO0K
1M
200K
423K
600K
100K
1M
200K
400K
500K
B0OK
1M
Trairing Step
Training Stop
Training Stop
690
-102
Trariing EL80
0.90
-100
Traning ELBO
Trening
106
globoi
Barker
Diflernco
MPF
100
-107
LOO + local
LGO
Only glibal
globel
-106
RELAX with
RODED
590
-100
200K
402K
B00K
800K
1M
200K
400K
800K
BOOK
1M
Training Step
Training Step
Training Stop
Figure 4: Ablation study of impact of RODEO components: (a) Stein operators, (b) LOO baseline and global and local CVs, (c) surrogate functions on binary VAE training performance.
Impact of Stein operator Figure 4a explores the impact of RODEO Stein operator choice on con-binarized MNIST. As expected, the less stable difference (8) and MPF (6) operators lead to significantly higher gradient variances and worse training ELBOs. In fact, the same operators led to divergent training on binarized MNIST. The Barker operator (6) with neighborhoods defined by having one differing coordinate yields results very similar to Gibbs (4) as the operators themselves
very similar for binary variables (notice that {())+((a))
elv) 4(ys [m when r-; y-i).
Impact of LOO baseline and global and local CVs Figure 4b explores binarized MNIST perfor- mance when RODEO is modified to retain a) only the LOO baseline (this is equivalent to standard RLOO), b) only the global CV, c) the LOO baseline the global CV, or d) the LOO baseline the local CV. We observe that the global and local Stein CVs both contribute to variance reduction and have complementary effects. Moreover, remarkably, the global Stein CV alone outperforms RLOO.
Impact of surrogate functions To tease apart the benefits of our new surrogate functions (13) and the remaining RODEO components, we replace the surrogate function c=)=) in RELAX [20] with surrogates, which only requires single evaluation of f per (k) (see Appendix A for more details). Figure 4c compares the performance of RODEO and this modified RELAX on binarized MNIST. Since the surrogate functions are matched, the consistent improvements over modified RELAX be attributed to the Stein and double CV components of RODEO. In Appendix C.2, we also experiment with increasing the complexity of И and H' by using two hidden layers instead of single hidden layer in the MLP. We did not observe significant improvements in variance reduction.
7 Conclusions, Limitations, and Future Work
This work tackles the gradient estimation problem for discrete distributions. We proposed variance reduction technique which exploits Stein operators to generate control variates for REINFORCE leave-one-out estimators. Our RODEO estimator does not rely 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 - 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 Appendix C. 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..
hu(1)= H(( Ejzx f(zl))), x- Enth (f
(y-
The runtime of RODEO could also be improved by varying the Stein operator employed. For example, the Gibbs operator (a) (4) used in our experiments evaluated its surrogate function neighboring locations of æ. 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 -vrriane 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 overcome by designing effective surrogate functions that make no use of derivative information.
Acknowledgments and Disclosure of Funding
We thank Heishiro Kanagawa for suggesting appropriate names for the Barker Stein operator.
References
[1] Andreas Anastasiou, Alessandro Barp, François-Xave Xavier Briol, Bruno Ebner, Robert E Gaunt, Fatemeh Ghaderinezhad, Jackson Gorham, Arthur Gretton, Christophe Ley. Qiang Liu, and Lester Mackey. Stein's method meets statisties: review of some recent developments. arXiv preprint arXiv:2105.034433 105.03481, 2021.
[2] Roland Assaraf and Michel Caffarel. Zero-variance principle for Monte Carlo algorithms. Phys. Rev. Lett.. 83:4682-4685, Dec 1999.
[3] A. D. Barbour. Stein's 's method and Poisson process convergence. J. Appl. Probab., (Special Vol. 25A): A celebration of applied probability.
175 184, 1988. ISSN
[4] Av A Barker. Monte Carlo calculations of the radial distribution functions for proton? electron plasma. Australian Journal of Physics, 18(2):119-134 1965.
[5] Alessandro Barp. Chris Oates, Emilio Porcu, and Mark Girolami. A Riemann-Stein kemnel method. arXiv preprint arXiv:1810.049.. 2018.
[6] Alessandro Barp, Trancois-Xavier Briol, Andrew Duncan, Mark Girolami. and Lester Mackey. Minimum Stein discrepancy estimators. Advances Neural Information Processing Systems, 12:12964-11299 2019.
[7] Yoshua Bengio, Nicholas Léonard, and Aaron Courville. Estimating or propagating gradients through stochastic neurons for conditional computation. arXiv preprint teeXiv//308.35 2013.
[8] Guy Bresler and Dheeraj Nagaraj. Stein's method for stationary distributions of Markov chains and application to Ising models. The Annals of Applied Probubility, 29(5):3230-32666 2019.
[9] Timothy C. Brown and Aihua Xia. Stein's 's Method and Birth-Death Processes. The Annals of Probability, 29(3):1373 1403, 2001.
[10] Yuri Burda, Roger Grosse, and Ruslan Salakhutdinov. Importance weighted autoencoders. International Conference on Learning Representations, 2015.
[11] Yulai Cong, Miaoyun Zhao, Ke Bai, and Lawrence Carin. GO gradient for expectation-basee objectives. In International Conference Learning Representations, 2019.
[12] Petros Dellaportas and loannis Kontoyiannis. Control variates for estimation based on reversible Markov chain Monte Carlo samplers. Journal of the Royal Statistical Society: Series B (Statistical Methodology). 74(1):133-161, 161, 2012.
[13] Aleksandar Dimitriev and Mingyuan Zhou. ARMS: Aniihhetc-RFNNFFRFF--Miii SSamlee gradient for binary variables. In International Conference on Machine Learning, volume 139, pages 2717-2727, 2021.
[14] Zhe Dong, Andriy Mnih, and George Tucker. DisARM: An antithetic gradient estimator for binary latent variables. In Advances in Neural Information Processing Systems, volume 33, pages 18637-18647, 2020.
[15] Zhe Dong, Andriy Mnih and George Tucker. Coupled gradient estimators for discrete latent variables. arXiv preprint 10Xiv:2106.088: 08056, 2021.
10

---

[16] Peter Eichelsbacher and Gesine Reinert. Stein's method for discrete Gibbs measures. The Annals of Applied Probability, 18(4):1588-1615 2008
[17] Charles Geyer. Introduction Markov chain Monte Carlo. In Steve Brooks, Andrew Gelman, Galin Jones, and Xiao-Li Meng, editors, Handbook of Markov Chain Monte Carlo, pages 3-48. CRC Press, 2011.
[18] Peter W. Glynn. Likelihood ratio gradient estimation for stochastic systems. Commun, ACM, 33(10); 75 75-84, October 1990.
[19] Jackson Gorham and Lester Mackey. Measuring sample quality with Stein's 's method. In Advances in Neural Information Processing Systems, pages 226-234, 2015.
[20] W. Grathwohl, D. Choi, Y. Wu, G. Roeder, and D. Duvenaud. Backpropagation through the void: Optimizing control variates for black-box gradient estimation. In International Conference Learning Representations, 2018.
[21] Will Grathwohl, Kevin Swersky, Milad Hashemi. David Duvenaud, and Chris Maddison, Oops I took gradient: Scalable sampling for discrete distributions. In International Conference on Machine Learning. pages 3831-3841, 2021.
[22] Shixiang Gu, Sorgey Levine, Ilya Sutskever, and Andriy Mnih. MuProp: Unbiased backpropagation for stochastic neural networks. In International Conference on Learning Representations, 2016.
[23] Shane G Henderson. Variance reduction via approximating Markov process. PhD thesis, Stanford University, 1997.
[24] Liam Hodgkinson, Robert Salomone, and Fred Roosta. The reproducing stein kernel approach for post-hoc corrected sampling, arXiv preprint arXiv:2001. 09266,, 2020.
[25] Susan Holmes. Stein's 's method birth and death chains. In Persi Diaconis and Susan Holmes, editors. Stein's 's Method: Expository Lectures and Applications, pages 42-65. 2004.
[26] Eric Jang, Shixiang Gu, and Ben Poole. Categorical reparameterizzation with Gumbel-softmax International Conference on Learning Representations, 2017.
[27] Martin Jankowiak and Fritz Obermeyer. Pathwise derivatives beyond the reparameterization trick. In International Conference on Machine Learning, pages 2235-2244. PMIR, 2018.
[28] Samuel Karlin and James McGregor. The classification of birth and death processes. Transactions of the American Mathematical Society, 86(2)3366-400, 1957.
[29] Diederik P. Kingma and Max Welling. Auto-encoding variational Bayes. International Conference on
Learning Representations, 2014.
[30] Kool, V. Hoof, and M. Welling. Buy REINFORCE samples, get a baseline for free! In DeepRLStruccPPed@e ICLR, 2019.
[31] Wouter Kool, Herke van Hoof, and Max Welling. Estimating gradients for discrete random variables by sampling without replacement. In International Conference on Learning Representations, 2020.
[32] Brenden M Lake, Ruslan Salakhutdinov, and Joshua B Tenenbaum. Human-level concept learning through probabilistic program induction. Science, 350(6266):1333
2015.
[33] Yann LeCun, Léon Bottou, Yoshua Bengio, and Patrick Haffner. Gradient-based learning applied to document recognition. Proceedings of the IEEE, 86( 1):2278-2324, 1998.
[34] Hao Liu, Yihao Feng. Yi Mao, Dengyong Zhou, Jian Peng, and Qiang Liu. Action-depedent control variates for policy optimization via Stein's identity. arXiv preprint arXiv: 11188 2017.
[35] Chris Maddison, Andriy Mnih, and Yee Whye Teh. The concrete distribution: A continuous relaxation of discrete random variables. In International Conference on Learning Representations, 2017.
[36] Antonietta Mira. Reza Solgi, and Daniele Imparato. Zero variance Markov chain Monte Carlo for Bayesian estimators. Statistics and Computing, 23(5):653 662, 2013.
[37] Andriy Mnih and Karol Gregor. Neural variational inference and learning in belief networks. In Interna- tional Conference Machine Learning, pages 1791-1799 2014.

---

[38] Andriy Mnih and Danilo Rezende. Variational inference for Monte Carlo objectives. In International Conference on Machine Learning, pages 2188-2196, 2016.
[39] Shakir Mohamed, Mihacla Rosca, Michael Figumov, and Andriy Mnih. Monte Carlo gradient estimation in machine learning. Journal of Machine Learning Research, 21(132):1 -2, 2020.
[40] Mathias Niepert, Pasquale Minervini, and Luca Franceschi. Implicit MLE: backpropagating through discrete exponential family distributions. Advances in Neural Information Processing Systems, 34:1 34:14567- 14579, 2021.
[41] Chris J Oates, Mark Girolami, and Nicolas Chopin. Control functionals for Monte Carlo integration. Journal of the Royal Statistical Society: Series B (Statistical Methodology), 79(3)x695-718, 2017.
[42] Art B. Owen, Monte Carlo theory, methods and examples. 2013.
[43] John Paisley, David M Blei, and Michael Jordan. Variational Bayesian inference with stochastic search. In International Conference Machine Learning, pages 3363-1370 2012.
[44] Paavo Parmas and Masashi Sugiyama. A unified view of likelihood ratio and reparameterization gradients. PMLR, 2021.
In International Conference on Artificial Intelligence and Statistics, pages
[45] Rajesh Ranganath, Sean Gerrish, and David Blei. Black box variational inference. International Conference on Artificial Intelligence and Statistics, page 814 842, 2014.
[46] Gesine Reinert and Nathan Ross. Approximating stationary distributions of fast mixing Glauber dynamics, with applications to exponential random graphs. The Annals of Applied Probability, 29(5):32013322) 2019.
[47] Danilo J. Rezende, Shakir Mohamed, and Daan Wierstra. Stochastic backpropagation and approximate inference deep generative models. In International Conference on Machine Learning, 2014.
[48] Lorenz Richter. Ayman Boustati, Nikolas Nüsken, Francisco Ruiz, and Omer Deniz Akyildiz. VarGrad: low-variance gradient estimator for variational inference. Advances Neural Information Prucessing Systems, volume 33, pages 13481 134422 2020.
[49] Tim Salimans and David A Knowles. On using control variates with stochastic approximation for variational Bayes and its connection to stochastic linear regression. arXiv preprint arXiv:1401.1022 2014.
[50] John Schulman, Nicolas Heess, Theophane Weher, and Pieter Abbeel. Gradicnt estimation using stochastic computation graphs. Advances in Neural Information Processing Systems, 28, 2015.
[51] Noam Shazeer, Azalia Mirhoseini, Krzysztof Maziarz, Andy Davis, Quoe Le,, Geoffrey Hinton, and Jeff Dean. Outrageously large neural networks: The sparsely-gated mixture-of-experts layer. arXiv preprint arXiv:1701.06532 2017.
[52] Shijing Si, Chris Oates, Andrew Duncan, Lawrence Carin, and François- -aviee Briol. Scalable control variates for Monte Carlo methods via stochastic optimization. arXiv preprint arXiv:2006.07487 2020.
[53] Jascha Sohl Sohl-Dicksteen Peter Battaglino, and Michael R DeWeese. Minimum probability flow learning. In International Conference on Machine Learning, pages 905 905-912, 2011.
[54] Leah F South, Chris J Oates, Antonictta Mira, and Christopher Drovandi. Regularised vero-nariance control variates for nigh-dimensional variance reduction. arXiv preprint 1,500,8 2018
[55] Charles Stein. A bound for the error in the normal approximation to the distribution of sum of dependent random variables. In Proceedings of the Sixth Berkeley Symposium on Mathematical Statistics and Probability, Volume 2: Probability Theory. The Regents of the University of California, 1972.
[56] Charles Stein, Persi Diaconis, Susan Holmes, and Gesine Reinert. Use exchangeable pairs the analysis of simulations. Lecture Moter-Monogruph Series, pages 1-26, 2004.
[57] David Stirzaker. Stochastic processes and models. OUP Oxford, 2005.
[58] Michalis K. Titsias and Miguel Lázaro-Gredilla, Doubly stochastic variational Bayes for non-conjugate inference. In International Conference on Machine Learning, 2014.
[59] Michalis K Titsias and Miguel Lázaro-Gredilla Local expectation gradients for black box variational inference. Advances in neural information processing systems, 28:2638 2646, 2015.
12

---

[60] Michalis K Titsias and Jiaxin Shi. Double control variates for gradient estimation in discrete latent variable models. International Conference Artificial Intelligence and Statisticx, 2022.
[61] Seiya Tokui and Issei Sato. Evaluating the variance of likelihood-ratio gradient estimators. In International Conference Machine Learning, pages 3414-3422 2017.
[62] G. Tucker, A. Mnih, C. J. Maddison, and J. Sohl-Dickstein. REBAR: Low-variance, unbiased gradient estimates for discrete latent variable models. In Advances in Neural Information Processing Systems, 2017.
[63] Arash Vahdat and Jan Kautz. NVAE: A deep hierarchical variational autoencoder. Advances in Neural Information Processing Systems, 33:19667-19679, 2020.
[64] Lex Weaver and Nigel Tao. The optimal reward baseline for gradient-based reinforcement learning. In Conference on Uncertainty in Artificial Intelligence, pages 538-545, 2001.
[65] Ronald J Williams. Simple statistical gradient-folooiing algorithms for connectionist reinforcement learning. Machine learning, 8(3-4):229-256, 1992.
[66] Han Xiao, Kashif Rasul, and Roland Vollgraf. Fashion-MNIST: novel image dataset for henchmarking machine learning algorithms. arXiv preprint arXiv:/708.007744 2017.
[67] Jiasen Yang, Qiang Liu, Vinayak Rao, and Jennifer Neville. Goodness-off of-fit fit testing for discrete distributions via Stein discrepancy. In International Conference on Machine Learning, pages 5561-5570. PMLR, 2018.
[68] Mingzhang Yin and Mingyuan Zhou. ARM: Augmen-REINFFRRCF-merg gradient for stochastic binary networks. International Conference on Learning Representations, 2019.
[69] Mingzhang Yin, Nhat Ho, Bowei Yan, Xiaoning Qian, and Mingyuan Zhou. Probabilistic best subset selection gradient-based optimization, 2020.
[70] Giacomo Zanella. Informed proposals for local MCMC in discrete spaces, Journal of the American Statistical Association, 115(530):852 265, 2020.
13


Reviews:
[{"review_id": "tQR9iGuVcV2", "reviewer": "Reviewer_K2dM", "summary": "The paper introduces discrete Stein operators for gradient estimation applied to Bernoulli variational autoencoders. As is widely known, gradients play a crucial role in training machine learning algorithms and taking gradients with respect to a discrete variable is non-trivial. Existing approaches tend to produce gradients with a high variance which then requires control variate approaches to manage the variance of the gradient. In this paper, the authors address this problem with a variance reduction technique based on Stein operators for discrete distributions. \n\nThe authors test the efficacy of their approach on three variations of the binary latent variable VAE. In the experiments, the authors show that their discrete Stein operator is competitive with existing state-of-the-art approaches. ", "questions": "- In the algorithm box $x^{(1:K)} \\sim q_\\eta$ are iid. Is this necessary or can they be generated from a Markov chain?\n- Is requiring $\\nabla_\\theta f(x)$ a restriction of this approach compared to the other approaches discussed in the paper?\n- It feels like the choice of the function $H$ (i.e. eq. 16) is important for the efficiency of the control variate. Can more be said about the choice of the neural architecture in H?\n- What is the computational cost for the algorithms in Table 2?\n- The gradient $g_gamma$ isn't clearly defined.\n", "limitations": "The limitations of the approach are discussed in the conclusion section. There are no obvious negative societal issues.", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Originality - The paper proposes a new approach to variance-reduced gradient estimation for discrete variables using Stein operators. As far as I am aware, this is a novel contribution to the research area. A potential weakness in terms of originality is perhaps that the main benefit of this approach, in terms of its efficiency, is likely to stem from the variance-reduction aspect, which is based on the reinforce leave-one-out (RLOO) estimator and a critical perspective might suggest that this paper is taking RLOO and plugging it into a Stein operator, which may reduce the novelty slightly. \nQuality - The paper is of high quality and well-placed in the NeurIPS community. \nClarity - The paper is very well written and the methodology is clearly presented. There are no obvious weaknesses in the presentation that the authors could aim to improve. \nSignificance - The paper nicely addresses the existing literature in the field which helps to place this work in the wider field. The problem addressed in the paper is significant in the ML community and the contribution of this paper is significant for both its contribution to latent variable VAEs and Stein method. In theory, the paper could be applied to general discrete variables, but in the paper, only VAE-type examples are considered. Is this a limitation of the approach?", "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": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "hnrasw8I1G5", "reviewer": "Reviewer_5uM4", "summary": "In this work, the authors propose a new type of control variate (CV) for discrete distributions. The proposed method is based on Stein operators. \nThe authors further incorporate the double CV into this proposed framework. \nExperiments are shown to support the claim. ", "questions": "See the weaknesses section.", "limitations": "N/A", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n* 1 A new type of control variate (CV) for discrete distributions.\n* 2 The proposed method is sound.\n* 3 The proposed method can be used to train deep NNs. \n* 4 Many empirical results are given to support the claim.\n\nWeaknesses:\n* 1 Discuss whether the proposed method can be applied to a discrete distribution with infinite support such as a Poisson distribution.\n* 2 The authors should discuss the iteration cost (computational budge) of the proposed method. It will be great if the authors discuss the iteration cost of all related methods including baseline methods.\n* 3 The running time instead of the number of training steps should be included in at least one of the plots.\n ", "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": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "Ubrmx-ksHo6", "reviewer": "Reviewer_avva", "summary": "This paper introduces a control variate scheme for REINFORCE-type estimators based on the Stein operators for discrete distributions. Discrete Stein operators generate zero-mean functions under the discrete sampling distribution; therefore, they can be used to construct control variates if a suitable function is chosen. Based on this insight, the authors propose RODEO, a CV scheme based on discrete Stein operators for the REINFORCE leave-one-out (RLOO) estimator. RODEO contains to adaptive function $H$ and $H^*$ that are optimised by minimising the variance of the estimator. RODEO is validated on binary VAE models and is shown to be effective in reducing the variance and inducing faster convergence.", "questions": "* What is the wall-clock time performance of this method compared to RLOO? I am especially interested in comparing its wall-clock performance compared to simply increasing the number of samples in RLOO. I know that this is dependent on the architecture of $H$ and $H^*$, but it is good to have some indication to its practical feasibility.\n* Can this type of estimator benefit from a parallel architecture? For instance, running the model on one GPU and $H$ and $H^*$ on another? $H$ and $H^*$ seem to take $f$ and its gradients as inputs, so this might be challenging at a first glance. Can the authors comment on this point?\n* When is it practically feasible to use this estimator? Are the extra implementation and computation overheads worth it compared to simply increasing the number of MC samples in RLOO or even vanilla REINFORCE? I think it good to add a discussion of its practicality to final section.\n* What is the effect of the architecture of $H$ and $H^*$ on the variance reduction? Is bigger always better? And is there any trade-off in terms of computation vs variance reduction?", "limitations": "Some limitations of this work are addressed in the conclusion and their discussion was, in my opinion, sufficient. There is no potential negative societal impact in my assessment.", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "To the best of my knowledge, the use of Stein operators in gradient estimation in discrete settings is novel. The authors attempt to place their work within the literature, and they make a sufficient argument for the originality of their contribution.\n\nThe development and explanation of the methodology is very clear, even for someone who is not previously familiar with Stein operators. That being said, I found the build-up leadin to equation (14) somewhat confusing. How does a Markov chain generate i.i.d samples? This section is important for developing the estimator, so I think the authors should improve the clarity of this derivation (or expand it in the appendix).\n\nOverall, I think this is a good paper. There are no major concerns with regards to the motivation and the development of the idea. The empirical evaluation can be improved by further comparisons. e.g. vanilla RLOO with different number of samples, and an illustration of the limitation of this method when the number of samples increases (this point was acknowledged in the conclusion). It is lacking a discussion of the practicality of this estimator given its complexity.\n\nIn my assessment, the main strengths of this work are its originality and clarity in the motivation and the development of the methodology, as well as its potential impact on future research. Its main weaknesses are the limited evaluation and discussion on its practicality.\n", "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": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "QziTqiLl7lT", "reviewer": "Reviewer_34qn", "summary": "This paper studies the design of control variates for REINFORCE-type gradient estimators in discrete expectation optimization problems. The paper proposes a general class of control variates using discrete Stein operators, and a practical unbiased gradient estimator RODEO with Stein-inspired control variates. Experiments on Bernoulli VAEs show the proposed method achieves lower gradient variance and better performance compared to a variety of baselines. ", "questions": "There are no major questions the reviewer would like to raise. There are 2 minor questions of interest:\n\n- Equation (14) is built under the intuition of an “optimal” Markov chain, but in practice the Markov chain is clearly not optimal with limited transition neighborhood especially in high dimensions. In the latter case, it would be helpful to analyze its implications for (14). Alternatively, it would be helpful to study the amount of variance reduction given by Rao-Blackwellization when $h=f$ exactly as dimension increases?\n- On page 7, it is mentioned that “RELAX was often observed to have very strong performance in prior work”. Why is it comparatively worse than all baselines in this paper?\n", "limitations": "The authors adequately addressed the limitations of their work and proposed interesting future directions of study. The authors did not discuss any potential negative societal impacts of the work, but as the work is mostly theoretical it is unlikely to have immediate negative societal impacts.    ", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n\n- The paper is well-written and gives a clear introduction to previous variance reduction methodologies of relevance. \n- The theory is well-substantiated and well-explained. The simplifications made in (14) and (16) are justified, and the intuition that $h$ serves as a surrogate for $f$ is clearly built by relating to Rao-Blackwellization. \n- A number of Stein operators are cited, with practical guidelines for choosing suitable operators based on numerical stability. \n- Experiment results are comprehensive and inclusive of the best introduced baseline methods. The comparisons are performed fairly under similar computational costs, and the slightly increased cost for RODEO due to evaluating auxiliary functions is clearly highlighted and justified. Ablation studies clearly demonstrate the effect of different components of the proposed method. \n\nWeaknesses:\n\n- While the proposed methodology is original and novel, similar ideas based on Stein’s methods exist for REINFORCE-type gradients of continuous distributions and as general control variates for evaluating expectations. Nonetheless, this has been clearly detailed in the related works section. \n- The method is evaluated only using VAE with binary latents, which has similar problem structure across different experiments, and may or may not favor the proposed methodology. In particular, the binary state space is relatively simplistic, albeit in a relatively high dimension of 200. \n", "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": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}]

Author rebuttal:

Decision: Accept