Gradient Estimation with Discrete Stein Operators
Anonymous Author(s)
AffiliationAddress email
Abstract
Gradient estimation—approximating the gradient of an expectation with respect tothe parameters of a distribution—is central to the solution of many machine learning problems.
However, when the distribution is discrete, most common gradientestimators suffer from excessive variance.
To improve the quality of gradient estimation, we introduce a variance reduction technique based on Stein operators fordiscrete distributions.
We then use this technique to build flexible control variatesfor the REINFORCE leave-one-out estimator.
Our control variates can be adaptedonline to minimize variance and do not require extra evaluations of the target function.
In benchmark generative modeling tasks such as training binary variationalautoencoders, our gradient estimator achieves substantially lower variance thanstate-of-the-art estimators with the same number of function evaluations.
Introduction
Modern machine learning relies heavily on gradient methods to optimize the parameters of a learningsystem.
However, exact gradient computation is often difficult.
For example, in variational inference for training latent variable models [29, 42] and policy gradient algorithms in reinforcementlearning [63] the exact gradient features an intractable sum or integral introduced by an expectationunder an evolving probability distribution.
To make progress, one resorts to estimating the gradientby drawing samples from that distribution [39].
The two main classes of gradient estimators used in machine learning are the pathwise or reparameterization gradient estimators [29, 46, 56] and the REINFORCE or score function estimators [18, 63].
Thepathwise estimators have shown great success in training variational autoencoders [29] but are onlyapplicable to continuous probability distributions.
The REINFORCE estimators are more generalpurpose 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 techniquefor discrete expectations.
Our method is based upon Stein operators [see, e.g., 1, 53], computablefunctionals that generate mean-zero functions under a target distribution and provide a natural way ofdesigning 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 [6, 8, 9, 16, 25, 45, 65].
We then develop a gradientestimation framework—RODEO—that augments REINFORCE estimators with mean-zero CVsgenerated from Stein operators.
Finally, inspired by Double CV [58], we extend our method todevelop CVs for REINFORCE leave-one-out estimators [30, 48] to further reduce the variance.
The benefits of using Stein operators to construct discrete CVs are twofold.
First, the operatorstructure permits us to learn CVs with a flexible functional form such as those parameterized byneural 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.
[Table]
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 functionevaluations, RODEO delivers lower variance and better training objectives than the state-of-the-artgradient estimators DisARM [14, 67], ARMS [13], Double CV [58], and RELAX [20].
Background
We consider the problem of maximizing the objective function E q η [ f ( x )] with respect to the parameters η of a discrete distribution q η ( x ) .
Throughout the paper we assume f ( x ) is a differentiable function of real-valued inputs x ∈ R d but is only evaluated at a discrete subset X d dueto the discreteness of q η . 1 Exact computation of the expectation is typically intractable due tothe complex nature of f ( x ) and q η ( x ) .
The standard workaround is to rewrite the gradient as∇ η E q η [ f ( x )]
=
E q η [ f ( x ) ∇ η log q η ( x )] and employ the Monte Carlo gradient estimator known asthe score function or REINFORCE estimator
[Equation]
[Equation] [Equation]
[Equation]
Here, b is a constant called the “baseline” introduced to reduce the variance of the estimator byreducing the scaling effect of f ( x ( k ) ) .
Since E q η [ ∇ η log q η ( x )] = 0 , the REINFORCE estimator isunbiased for any choice of b , and the term b ∇ η log q η ( x ) is known as a control variate (CV) [41,
Ch. 8].
The optimal baseline can be estimated using additional function evaluations [7, 44, 62].
Asimpler approach is to use moving averages of f from historical evaluations or to train functionapproximators to mimic those values [37].
When K ≥ 2 , a powerful variant of (1) is obtained byreplacing b with the leave-one-out average of function values:
[Equation]
This approach is often called the REINFORCE leave-one-out (RLOO) estimator [30, 47, 48] and wasrecently observed to have very strong performance in training discrete latent variable models [14, 47].
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 sampledependent baseline b ( x ( k ) ) as long as the expectation c ≜ E q η [ b ( x ( k ) )
∇ η log q η ( x ( k ) )] is easilycomputed or has a low-variance unbiased estimate.
In this case, we can correct for any bias introducedby b ( x ( k ) ) by adding in this expectation: 1 K (cid:80) Kk =1 ( f
[Equation]
x ( k ) )
− b ( x ( k ) ))
[Equation]
log q η ( x ( k ) )
+ c. Forexample, b ( x ( k ) ) can be a lower bound or a Taylor expansion of f ( x ( k ) )
Taking this a stepfurther, the Double CV estimator [58] proposed to treat f ( x ( k ) )
− b ( x ( k ) ) as the effective objectivefunction and apply the leave-one-out idea:
[Equation]
The resulting estimator adds two CVs to RLOO: the global CV b k ( x ( k ) ) ∇ η log q η ( x ( k ) ) and thelocal CV b k ( x ( j ) ) .
Intuitively, b k ( x ( j ) ) is aimed at reducing the variance of the LOO average.
Thisis motivated by the fact that replacing 1 K − 1 (cid:80)
j ̸ = k f ( x ( j ) ) with E q η [ f ] in (2) leads to lower variance[58, Prop 1].
To obtain a tractable correction term c , Titsias and Shi [58] adopt a linear design of b :b k ( x ) =
1 K − 1 (cid:80) j ̸ = k f ( x j )( x − µ ) for µ = E q η [ x ] .
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 stronglycorrelated with f .
This may fail to hold for the above linear b k and a highly nonlinear f , especially inapplications like training deep generative models.
In the following two sections, we will introduce a method that allows us to use very flexible CVswhile still maintaining a tractable correction term.
Our method enables online adaptation of CVsto minimize gradient variance (similar to RELAX [21]) but does not assume q η has a continuousreparameterization.
We then apply it to generalize the linear CVs in Double CV to very flexible onessuch as neural networks.
Moreover, we provide an effective CV design based on surrogate functionsthat requires no additional evaluation of f compared to RLOO.
Control Variates from Discrete Stein Operators
At the heart of our new estimator is a technique for generating flexible discrete CVs, that is forgenerating a rich class of functions that have known expectations under a given discrete distributionq .
One way to achieve this is to identify any discrete-time Markov chain ( X ( t ) )
∞ t =0 with stationarydistribution q .
Then, the transition matrix P , defined via P xy =
P ( X ( t +1)
=
y | X ( t ) = x ) , satisfies
P ⊤ q = q
and hence
E q
[( P
[Equation]
I ) h ] = 0 , (4) for any integrable function h .
We overload our notation so that, for any suitable matrix A , ( Ah )( x ) ≜ 87 (cid:80) y
A xy h ( y ) .
In other words, for any integrable h , the function ( P − I ) h is a valid CV as it hasknown mean under q .
Moreover, the linear operator P − I is an example of a Stein operator [1, 53] inthe sense that it generates mean-zero functions under q .
In fact, both Stein et al.
[54] and Dellaportasand Kontoyiannis [12] developed CVs of the form ( P − I )
h based on reversible discrete-time Markovchains 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 ≥ 0 with q as its stationary distribution, then the generator A , defined via( Ah )( x ) =
lim t → 0 E [ h ( X ( t ) )] − h ( x ) t , (5) satisfies A ⊤ q = 0 and hence E q
[ Ah ] = 0 for all integrable h .
Therefore, A is also a Stein operatorsuitable for generating CVs.
Moreover, since any discrete-time chain with transition matrix P can beembedded into a continuous-time chain with transition rate matrix
A
= P −
I , this continuous-timeconstruction is strictly more general [55, 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 kernel of the random-scan Gibbs sampler with stationarydistribution q
[see, e.g., 17] is
[Equation]
(6) where 1 ( · ) is the indicator function.
The associated Stein operator is
A = P − I with
[Equation]
In the binary variable setting, (7) recovers the operator Bresler and Nagaraj [8], Reinert and Ross [45]used to bound distances between the stationary distributions of Glauber dynamics Markov chains.
Zanella Stein operator Zanella [68] studied continuous-time Markov chains with generator
[Equation]
(8) where N x denotes the transition neighborhood of x and κ is a continuous positive function satisfyingκ ( t )
=
tκ (1 /t ) .
Conveniently, the neighborhood structure can be arbitrarily sparse.
Moreover, the κ constraint ensures that the chain satisfies detailed balance (i.e., q ( x ) A xy =
q
( y ) A yx ) and hence that
A ⊤ q = 0 .
Hodgkinson et al.
[24] call the associated operator the Zanella Stein operator :
[Equation]
Important special cases include the minimum probability flow [51] (MPF) Stein operator ( κ ( t ) = √ t )discussed in Barp et al.
[6] and the Barker [4] Stein operator ( κ ( t ) =
t t +1 ).
Birth-death Stein operator
Let e 1 , . .
, e d represent the standard basis vectors on R d .
For anyfinite-cardinality space, we may index the elements of X
= { s 0 , . . .
, s m − 1 } , let idx ( x i ) return theindex of the element that x i represents, and define the increment and decrement operatorsinc i ( x ) =
x
+ e i ( s ( idx ( x i )+1)
mod m − x i ) and dec i ( x ) = x
+ e i ( s ( idx ( x i ) − 1) mod m − x i ) .
Then the product birth-death process [28] on X d with birth rates b i,x = q ( inc i ( x ))
q ( x ) , death ratesd i,x = 1 , and generator
A xy = 1 d (cid:80) di
=1 b i,x 1 ( y = inc i ( x ))
+ d i,x 1 ( y = dec
i ( x )) − ( b i,x + d i,x ) 1 ( y = x ) , has q as a stationary distribution, as A ⊤ q = 0 .
This construction yields the birth-death Stein operator( Ag )( x ) =
1 d (cid:80) di
=1 b i,x ( g ( inc i ( x )) − g ( x )) − d i
[Equation]
An analogous operator is available for countably infinite X [24], and Brown and Xia [9], Eichelsbacherand Reinert [16], Holmes [25] used related operators to characterize convergence to discrete targetdistributions.
Moreover, by substituting h ( x ) = g ( x ) − g ( inc ( x )) in (10), we recover the difference
Stein operator proposed by Yang et al.
[65] without reference to birth-death processes:
[Equation]
Choosing a Stein operator Despite being better-known, the difference and MPF operators oftensuffer from large variance and numerical instability due to their use of unbounded probability ratiosq ( y ) /q
( x ) .
As a result, we recommend the numerically stable Gibbs operator when each componentof x is binary or takes on a small number of values.
When x i takes on a large number of values ( m ),the Gibbs operator suffers from linear scaling with m .
In this case, we recommend the Barker operatorwhere a sparse neighborhood structure can be specified, such as N x = { inc i ( x ) , dec i ( x ) for i ∈ [ d ] } .
The Barker operator is numerically stable as its κ ( q ( y ) q ( x ) )
=
q ( y ) q
( x )+ q ( y ) .
Gradient Estimation with Discrete Stein Operators
Recall that REINFORCE estimates the gradient E q η [ f ( x ) ∇ η log q η ( x )] .
Due to the existence of∇ η log q η ( x ) as a weighting function, we apply a discrete Stein operator to a vector-valued function˜ h : X → R d per dimension to construct the following estimator with a mean-zero CV:
[Equation]
Ideally, we want to choose the ˜ h
such that A ˜ h
i will be strongly correlated with f ( x ) ∇ η i log q η ( x ) toreduce its variance.
The optimal ˜ h
i is given by the solution of Poisson equation
[Equation]
We could learn a separate ˜ h i per dimension to approximate the solution of (13).
However, this posesa 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 ˜ h , we consider an “optimal” Markov chain whichgenerates i.i.d.
sample points from q η .
In this case, ( P − I )
˜ h
i becomes E q η [˜ h
i ] − ˜ h
i , and the optimalsolution is ˜ h
i = f ∇ η i log q η .
Inspired by this, we consider ˜ h of the form
[Equation] [Equation]
where we now only need to learn a scalar-valued function h .
Notably, when h exactly equals fand A
= P − I for any discrete time Markov chain kernel P , our CV adjustment amounts to
Rao-Blackwellization [41, Sec. 8.7], as we end up replacing f ( x ) ∇ η i log q η ( x ) with its conditionalexpectation P ( f ∇ η i log q η )( x ) =
E X t +1 | X t = x
[ f ( X t +1 ) ∇ η i log q η ( X t +1 )] .
This yields a guaranteed 146 variance reduction.
Surrogate function design Based on the above reasoning, we can view h as a surrogate for f .
Weavoid directly setting h
= f because our Stein operators evaluate h at all neighbors of the samplepoints, and f can be very expensive to evaluate [see, e.g., 49].
To avoid this problem, we first observethat ˜ h
(14) can be made sample-specific, i.e., we can use a different ˜ h k for each sample point x ( k ) :
[Equation] [Equation]
We then consider the following choices of h k that are informed about f while being cheap to evaluate:
[Equation] [Equation]
Here y will take on the values of x ( k ) and its neighbors in the Markov chain.
We omit x
( k ) in the sum(16) to ensure that the function h k is independent of x ( k ) and hence that E q η [ 1 K (cid:80)
Kk =
( A ˜ h k )( x ( k ) )]
=0 .
Moreover, this surrogate function design introduces no additional evaluations of f beyond thoserequired for the usual RLOO estimator.
Also, as observed by Titsias and Shi [58], for manyapplications, including VAE training, where f has parameters θ learned through gradient-basedoptimization, {∇ f ( x ( k ) ) }
Kk =1 can be obtained “for free” from the same backpropagation used tocompute ∇ θ f θ ( x ( k ) )
(see Algorithm 1).
RODEO We can further improve the estimator by leveraging discrete Stein operators and the abovesurrogate function design to construct both the global and local CVs in the Double CV framework [58](see Section 2).
We call our final estimator RODEO for R L O O with D iscrete St E in O perators:
[Equation]
where ˜ h ∗ k
( y )
=
h ∗ k
( y ) ∇ η log q η ( y ) and { h k , h ∗
k }
Kk =1 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 1 K − 1 (cid:80)
j ̸ = k f ( x ( j ) ) , while ( A ˜ h ∗ k )( x ( k ) ) acts as a global CV to further reduce the variance of thegradient estimate.
We adopt the aforementioned design of h (16) and ˜ h (14) for the two CVs.
In
Appendix B, we show that the RODEO estimator (17) is unbiased for ∇ η E q η [ f ( x )] when each h k isdefined as in (16) and
[Equation]
Optimization with RODEO
In practice, we let the functions H and H ∗ share a neural network architecture with two output units.
Since the estimator is unbiased, we can optimize theparameters γ of the network in an online fashion to minimize the variance of the estimator(similar to [20]).
Specifically, denoting the RODEO gradient estimate (17) by g γ ( x (1: K ) ) , then∇ γ Tr(Var( g γ ( x (1: K ) )))
=
E [ ∇ γ ∥ g γ ( x (1: K ) ) ∥ 22 ] .
In Algorithm 1, we use an unbiased estimate ofthis hyperparameter gradient, ∇ γ ∥ g γ ( x (1: K ) )
∥ 22 , to update γ after each optimization step of η .
[Table]
Related Work
[Table]
As we have seen in Section
[Table]
there is a long history of designing CVs for REINFORCE estimatorsusing “baselines” [7, 37, 44, 62].
Recent progress is mostly driven by leave-one-out [30, 38, 47,
48]and sample-dependent baselines [20, 22, 42, 58, 60].
REBAR [60] constructs the baseline throughcontinuous relaxation of the discrete distribution [26, 35] and applies reparameterization gradients tothe correction term.
As a result REBAR uses three evaluations of f for each x ( k ) instead of usualsingle 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 andonly 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 (17) is that it incorporates information from neighboring states thanks to the Stein operator while avoiding additional costly f evaluations thanks to thelearned surrogate functions h k .
Estimators based on analytic local expectation [57, 59] and GO gradients [11] also use neighborhood information but only at the cost of many additional target functionevaluations.
In fact, the local expectation gradient [57] can be viewed as a Stein CV adjustment (17)with a Gibbs Stein operator and the target function f used directly instead of the surrogate h .
Thedownside of these approaches is that f must be evaluated Kd times per training step instead of Ktimes as in RODEO, a prohibitive cost when f is expensive and d ≥ 200 as in Section 6.
Prior work has also studied variance reduction methods based on sampling without replacement [31]and antithetic sampling [13–15, 66] for gradient estimation.
DisARM [14, 67] 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 reductionin a variety of learning tasks [2, 5, 36, 40, 50, 52] including gradient estimation [27, 34].
Inparticular, 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].
Incontrast, our approach considers discrete Stein operators for Monte Carlo estimation in discretedistributions with exponentially large state spaces.
Recently, Parmas and Sugiyama [43, App. E.4]used a probability flow perspective to characterize all unbiased gradient estimators satisfying a mildtechnical condition; our estimators fall into this broad class but were not specifically investigated.
Experiments
Python code replicating all experiments can be found in the supplementary material.
6.1 Training Bernoulli VAEs
Following Dong et al.
[14] and Titsias and Shi [58], we conduct experiments on training variationalauto-encoders [29, 46] (VAEs) with Bernoulli latent variables.
VAEs are models with a joint densityp θ ( y, x ) , where x is the latent variable and θ denotes model parameters.
They are typically learnedthrough maximizing the evidence lower bound (ELBO)
E q η ( x | y )
[ f ( x )] for an auxiliary inference network q η ( x | y ) and f ( x ) ≜ log p θ ( y, x ) − log q η ( x | y ) .
In our experiments, p ( x ) and q η ( x | y ) arehigh-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 learnthe parameters η of the inference network.
The VAE architecture and training experimental setupfollows Titsias and Shi [58], and details are given in Appendix D.
The dimensionality of the latentvariable x is d = 200 .
The functions H (16) and H ′ (18) share a neural network architecture with twooutput units and a single hidden layer with 100 units.
For the numerical stability and variance reasonsdiscussed in Section 3, we use the Gibbs Stein operator (7) as a default choice in our experiments,but we revisit this choice in Section 6.3.
We consider the MNIST [33], Fashion-MNIST [64] and Omniglot [32] datasets using their standardtrain, validation, and test splits.
We use both binary and continuous VAE outputs ( y ) as in Titsias and
Shi [58].
In the binary output setting, data are dynamically binarized, and the Bernoulli likelihood isused; in the continuous output setting, data are centered between [ − 1 , 1] , and the Gaussian likelihoodwith learnable diagonal covariance parameters is used.
[Figure]
K =
In the first set of experiments we focus on the most common setting of K = 2 samplepoints and compare our variance reduction method to its counterparts including DisARM [14] and
Double CV [58].
Results for RLOO are omitted since it is consistently outperformed by Double
CV [see 58].
Table 2 shows, on all three datasets, RODEO achieves the best training ELBOs.
In
Figure 1 (left), we plot the gradient variance and average training ELBOs against training steps forall estimators on dynamically binarized MNIST.
RODEO outperforms DisARM and Double CV by alarge margin in gradient variance.
Next, we consider VAEs with Gaussian likelihoods trained on non-binarized datasets.
In this casethe gradient estimates suffer from even higher variance due to the large range of values f ( x ) cantake.
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 =
In the second set of experiments we compare RELAX [20], which uses three evaluationsof f per training step, with RODEO, Double CV, and ARMS [13] for K = 3 .
Figure 1 (right) and
Table 2 demonstrate that RODEO outperforms the three previous methods and generally leads tothe lowest variance.
Although RELAX was often observed to have very strong performance in priorwork [14, 58], our results in Figure 1 suggest that, for dynamically binarized datasets, much largergains 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 log-likelihoodbounds 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 (7) which performs 2 devaluations of the auxiliary functions H and H ′ in (16) and (18).
While this results in some extra cost,the network parameterizing H and H ′ takes only 2-dimensional inputs, produces scalar outputs, andis typically small relative to the VAE model.
As a result, the evaluation of H and H ′ is significantlycheaper than that of f , and its relative contribution to runtime shrinks as the cost of f grows.
6.2 Training hierarchical Bernoulli VAEs
To investigate the performance of RODEO when scaling to hierarchical discrete latent variable models,we follow DisARM [14, 66] to train VAEs with 2/3/4 stochastic layers, each of which consists of
[Figure]
200 Bernoulli variables.
We set K = 2 and compare our estimator with DisARM and Double CV ondynamically binarized MNIST, Fashion-MNIST
[Figure]
and Omniglot.
For each stochastic layer, we usea different CV network which has the same architecture as those in our VAE experiments from theprevious section.
More details are presented in Appendix D.
[Figure]
We plot the results with
-MNIST in Figure 3.
The resultsfor other datasets and for 2 and 3 stochastic layers can be found in Appendix E. RODEO generallyachieves the best training ELBOs.
One difference we noticed when comparing the variance resultswith those obtained from single-stochastic-layer VAEs is that these estimators have very differentbehaviors towards the beginning and the end of training.
For example, in Figure 3, Double CV startswith lower variance than DisARM, but the gap diminishes after around 100K steps, and DisARMstart to perform better as the training proceeds.
In contrast, RODEO has the lowest variance in bothphases for all datasets save Omniglot, where DisARM overtakes it in the long run.
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 we train binary latent VAEs with K = 2 .
Impact of Stein operator Figure 4a explores the impact of RODEO Stein operator choice onnon-binarized MNIST.
As expected, the less stable difference (11) and MPF (9) operators lead tosignificantly higher gradient variances and worse training ELBOs.
In fact, the same operators ledto divergent training on binarized MNIST.
The Barker operator (9) with neighborhoods defined byhaving one differing coordinate yields results very similar to Gibbs (7) as the operators themselvesare very similar for binary variables (notice that
[Equation]
i ) when x − i =
y − i ).
Impact of LOO baseline and global and local CVs Figure 4b explores binarized MNIST performance when RODEO (17) is modified to retain a) only the LOO baseline (this is equivalent tostandard 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
[Figure]
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 (16) andthe remaining RODEO components, we replace the surrogate function c ϕ ( z ) in RELAX [21] with oursurrogates, which only requires a single evaluation of f per x ( 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 canbe attributed to the Stein and double CV components of RODEO.
Conclusions, Limitations, and Future Work
This work tackles the gradient estimation problem for discrete distributions.
We proposed a variancereduction technique which exploits Stein operators to generate control variates for REINFORCEleave-one-out estimators.
Our RODEO estimator does not rely on continuous reparameterization ofthe distribution, requires no additional function evaluations per sample point, and can be adaptedonline to learn very flexible control variates parameterized by neural networks.
One potential drawback of our surrogate function constructions (16) and (18) is the need to evaluatean auxiliary function ( H or H ′ )
at K − 1 locations.
This cost can be comfortably borne when Hand H ′ are much cheaper to evaluate than f , and, in our experiments with the most common samplesizes, the runtime of RODEO was no worse than that of RELAX.
To obtain a more favorable cost forlarge K , one could employ alternative surrogates that require only a constant number of auxiliaryfunction evaluations, e.g.,
[Equation]
The runtime of RODEO could also be improved by varying the Stein operator employed.
Forexample, the Gibbs operator ( Ah )( x ) (7) used in our experiments evaluated its surrogate functionh at d neighboring locations of x .
This evaluation complexity could be reduced by subsamplingneighbors, resulting in a cheaper but still valid Stein operator, or by employing the numerically stable
Barker operator (9) with fewer neighbors.
Either strategy would introduce a speed-variance trade-offworthy of study in follow-up work.
Finally, we have restricted our focus in this work to differentiabletarget functions f .
In future work, this limitation could be overcome by designing effective surrogatefunctions that make no use of derivative information.
Checklist
For all authors...
(a) Do the main claims made in the abstract and introduction accurately reflect the paper’s
contributions and scope?
[Yes]
(b) Did you describe the limitations of your work?
[Yes] See Section 7.
(c) Did you discuss any potential negative societal impacts of your work?
[N/A]
(d) Have you read the ethics review guidelines and ensured that your paper conforms to
them?
[Yes]
If you are including theoretical results...
(a) Did you state the full set of assumptions of all theoretical results?
[Yes]
(b) Did you include complete proofs of all theoretical results?
[Yes]
If you ran experiments...
(a) Did you include the code, data, and instructions needed to reproduce the main experi-
mental results (either in the supplemental material or as a URL)?
[Yes] MIT-licensed
open-source Python code replicating all experiments can be found in the supplementary
material.
(b) Did you specify all the training details (e.g., data splits, hyperparameters, how theywere chosen)?
[Yes] See Appendix D.
(c) Did you report error bars (e.g., with respect to the random seed after running experiments multiple times)?
[Yes] See, e.g., Table 2.
(d) Did you include the total amount of compute and the type of resources used (e.g., typeof GPUs, internal cluster, or cloud provider)?
[Yes] See, e.g., Table 5.
If you are using existing assets (e.g., code, data, models) or curating/releasing new assets...
(a) If your work uses existing assets, did you cite the creators?
[Yes]
(b) Did you mention the license of the assets?
[Yes]
(c) Did you include any new assets either in the supplemental material or as a URL?
[Yes]
MIT-licensed open-source Python code replicating all experiments can be found in the
supplementary material.
(d) Did you discuss whether and how consent was obtained from people whose data you’re
using/curating?
[N/A]
(e) Did you discuss whether the data you are using/curating contains personally identifiable
information or offensive content?
[N/A]
If you used crowdsourcing
or conducted research with human
subjects
(a) Did you include the full text
of instructions given to participants
and
screenshots if
applicable?
[N/A]
(b) Did you describe any potential participant risks, with links to Institutional Review
Board (IRB) approvals, if applicable?
[N/A]
(c) Did you include the estimated hourly wage paid to participants and the total amountspent on participant compensation?
[N/A]
A Sample-dependent Baselines in REBAR and RELAX
We start with the REINFORCE estimator with the sample-dependent baseline b k :
[Equation] [Equation]
REBAR [60] introduces indirect independence on x ( k ) in b k through the continuous reparameterization x
= H ( z ) , z ( k ) ∼ q η ( z | x = x ( k ) ) , where z is a continuous variable and H is an argmax-likethresholding function.
Specifically, b k = f
( σ λ ( z ( k ) )) , where σ λ is a continuous relaxation of Hcontrolled by the parameter λ .
The correction term decomposes into two parts:
[Equation]
Both parts can be estimated with the reparameterization trick [29, 46, 56] which often has lowvariance.
The RELAX [20] estimator generalizes REBAR by noticing that f ( σ λ ( z )) can be replacedwith a free-form differentiable function c ϕ ( z ) .
However, RELAX still relies on parameterizing c ϕ ( z )as f ( σ λ ( z ))
+ r θ ( z ) to achieve strong performance, as noted in Dong et al.
To form modified RELAX in Section 6.3, we replace b k = c ϕ
( z ( k ) ) with b k = h k ( σ λ ( z ( k ) ))
for h kdefined in (16).
B Proof of Unbiasedness of RODEO
Recall our estimator defined in (17) is
[Equation]
To show the unbiasedness, we compute its expectation under q η as
[Equation]
Since the first term is the desired gradient ∇ η E q η [ f ( x )] and the third term is zero, it suffices to showthat the second term also vanishes.
Using the law of total expectations, we find for j ̸ = k ,
E
[Equation]
log q η ( x ( k ) )
| x ( j ) ]] = 0 , which completes the proof.
C Additional Experiments
C.1 Wall clock time comparison with RLOO
Besides necessary target function evaluations, the RODEO estimator comes with the additional costof evaluating the neural network-based H, H ∗ .
Therefore, RODEO is most suited to the problems
[Table]
where evaluating f is very expensive and its cost dominates those of the H
[Table]
H ∗ .
This is often thecase in practice.
For example, state-of-the-art variational autoencoders [e
[Table]
g., 61] are often built onexpensive neural architectures such as deep residual networks (ResNet).
Here, to demonstrate thepractical advantage of our method as the complexity of f grows, we replace the two-layer MLP VAEsused in previous experiments with a ResNet VAE (architecture shown in Table 3), while the neuralnetwork used by H, H ∗ remains a single-layer MLP with 100 hidden units.
We then compare thewall clock performance of RODEO with RLOO.
The latent variables in this experiment remain binaryand have 200 dimensions.
The results are shown in Figure 5.
RODEO achieves better training ELBOs than RLOO in the sameamount of time.
In fact, for this VAE architecture, the per-iteration time of RODEO is 25.2ms , whichis very close to the 23.1ms of RLOO.
This indicates that the cost of f is significant higher than
H, H ∗ .
[Figure]
C.2 Impact of neural network architectures of surrogate functions
We conduct one more ablation study to investigate the impact of neural network architectures usedby H, H ∗ .
Specifically, we replace the single-hidden-layer control variate network used in previousexperiments with a two-hidden-layer MLP (each layer has 100 units) and compare their performanceon binarized MNIST with K = 2 .
We keep other settings the same as in Section 6.1.
The results areplotted in Figure 6.
We do not observe significant difference between the two versions of RODEO.
D Experimental Details
Our implementation is based on the open-source code of DisARM [14] (Apache license) and Double
CV [58] (MIT license).
Our figures display 1M training steps and our tables report performance after1M training steps to replicate the experimental settings of DisARM [14].
[Figure]
D.1 Details of VAE experiments
VAEs are models with a joint density p ( y, x )
= p ( y | x ) p ( x ) , where x denotes the latent variable.x is assigned a uniform factorized Bernoulli prior.
The likelihood p θ ( y | x ) is parameterized by theoutput of a neural network with x as input and parameters θ .
The VAE has two hidden layers with200 units activated by LeakyReLU with the coefficient 0 . 3 .
To optimize the VAE we use Adam withbase learning rate 10 − 4 for non-binarized data and 10 − 3 for dynamically binarized data, except forbinarized Fashion-MNIST we decreased the learning rate to 3 × 10 − 4 because otherwise the trainingis very unstable for all estimators.
We use Adam with the same learning rate 10 − 3 for adapting ourcontrol variate network in all experiments.
The batch size is 100 .
The settings of other estimators arekept the same with Titsias and Shi [58].
In the minimum probability flow (MPF) and Barker Stein operator (9), we choose the neighborhood N x to be the states that differ in only one coordinate from x .
Let y ∈
N
x be an elementin this neighborhood
such that y i ̸
= x i and y − i = x − i .
For the MPF Stein estimator and thedifference Stein estimator (11), the density ratio q ( y ) q ( x ) can be simplified to q ( y i | x − i )
q
( x i | x − i ) .
We further replace it with
[Equation]
i ) q ( x i | x − i )+10 − 3 to alleviate numerical instability.
The Barker Stein estimatordoes not suffer from the numerical issue since the coefficient is bounded in the Bernoulli case:q ( y ) q
( x )+ q ( y )
=
q ( y i | x − i ) q ( x i | x − i )
+ q
( y i | x − i )
=
q ( y i | x − i ) .
In our experiments, we find that the difference Steinestimator is highly unstable and may diverge as the iteration proceeds.
D.2 Details of hierarchical VAE experiments
We optimize the hierarchical VAE using Adam with base learning rate 10 − 4 .
Our control variatenetwork is optimized using Adam with learning rate 10 − 3 .
Settings of training multilayer VAEs arekept the same with Dong et al.
[14], except that we do not optimize the prior distribution of the VAEhidden layer and use a larger batch size 100 .
E Additional Results
In this section, we measure test set performance using 100 test points and the marginal log-likelihoodbound of Burda et al.
[10], which provides a tighter estimate of marginal log likelihood than the
ELBO.
Throughout, we call this the “test log-likelihood bound.”
[Table]
[Figure]
li hood
[Figure]
ound
[Figure]
N on - b i n a r
i ze d
[Figure]
t
Log
[Figure]
i k
[Figure]
hood
B ound
[Figure]
[Table] [Table]
[Figure]
[Figure] [Figure]
[Table]
[Figure]
F a s h i onM N I
S T
[Figure]
t
[Figure] [Figure]
F a s h i onM N I
S T
[Figure] [Figure]
N
[Figure]
S
[Figure]
F a s h i onM N I
S T
[Figure]
m
[Figure]
t
[Figure]